Update app.py

app.py

@@ -1,106 +1,90 @@
-import spaces
 import os
-import json
 import time
-import torch
-from PIL import Image
-from tqdm import tqdm
+from datetime import datetime
+import uuid
 import gradio as gr
+from PIL import Image
 
-from safetensors.torch import save_file
-from src.pipeline import FluxPipeline
-from src.transformer_flux import FluxTransformer2DModel
-from src.lora_helper import set_single_lora, set_multi_lora, unset_lora
+# Define folders
+gallery_folder = "./gallery"
+task_folder = "./tasks"
 
-# Initialize the image processor
-base_path = "black-forest-labs/FLUX.1-dev"    
-lora_base_path = "./models"
+os.makedirs(gallery_folder, exist_ok=True)
+os.makedirs(task_folder, exist_ok=True)
 
-pipe = FluxPipeline.from_pretrained(base_path, torch_dtype=torch.bfloat16)
-transformer = FluxTransformer2DModel.from_pretrained(base_path, subfolder="transformer", torch_dtype=torch.bfloat16)
-pipe.transformer = transformer
-# 移除 pipe.to("cuda")，默认使用CPU
+# Function to save image and metadata
+def save_image_to_gallery(image, prompt, seed):
+    timestamp = datetime.now().strftime("%Y%m%d-%H%M%S")
+    filename = f"{timestamp}_seed{seed}.png"
+    metadata_filename = filename.replace(".png", ".txt")
+    
+    image_path = os.path.join(gallery_folder, filename)
+    metadata_path = os.path.join(gallery_folder, metadata_filename)
 
-def clear_cache(transformer):
-    for name, attn_processor in transformer.attn_processors.items():
-        attn_processor.bank_kv.clear()
+    image.save(image_path)
+    with open(metadata_path, "w") as f:
+        f.write(f"Prompt: {prompt}\nSeed: {seed}\nTime: {timestamp}")
 
-# Define the Gradio interface
-@spaces.GPU()  # 改为 @spaces.CPU() 或直接移除，因为免费层没有GPU
-def single_condition_generate_image(prompt, spatial_img, height, width, seed, control_type):
-    # Set the control type
-    if control_type == "Ghibli":
-        lora_path = os.path.join(lora_base_path, "Ghibli.safetensors")
-    set_single_lora(pipe.transformer, lora_path, lora_weights=[1], cond_size=512, device="cpu")
-    
-    # Process the image
-    spatial_imgs = [spatial_img] if spatial_img else []
-    image = pipe(
-        prompt,
-        height=int(height),
-        width=int(width),
-        guidance_scale=3.5,
-        num_inference_steps=15,  # 减少步数以适应CPU
-        max_sequence_length=512,
-        generator=torch.Generator("cpu").manual_seed(seed), 
-        subject_images=[],
-        spatial_images=spatial_imgs,
-        cond_size=512,
-    ).images[0]
-    clear_cache(pipe.transformer)
+# Modified generation function with saving
+def single_condition_generate_image_with_save(prompt, spatial_img, height, width, seed, control_type):
+    image = single_condition_generate_image(prompt, spatial_img, height, width, seed, control_type)
+    save_image_to_gallery(image, prompt, seed)
     return image
 
-# Define the Gradio interface components
-control_types = ["Ghibli"]
+# Get gallery images
+def load_gallery():
+    images = []
+    for fname in sorted(os.listdir(gallery_folder), reverse=True):
+        if fname.endswith(".png"):
+            images.append(os.path.join(gallery_folder, fname))
+    return images
 
-# Example data (调整分辨率以适应CPU)
-single_examples = [
-    ["Ghibli Studio style, Charming hand-drawn anime-style illustration", Image.open("./test_imgs/00.png"), 512, 512, 5, "Ghibli"],
-    ["Ghibli Studio style, Charming hand-drawn anime-style illustration", Image.open("./test_imgs/02.png"), 512, 512, 42, "Ghibli"],
-    ["Ghibli Studio style, Charming hand-drawn anime-style illustration", Image.open("./test_imgs/03.png"), 512, 512, 1, "Ghibli"],
-    ["Ghibli Studio style, Charming hand-drawn anime-style illustration", Image.open("./test_imgs/04.png"), 512, 512, 1, "Ghibli"],
-    ["Ghibli Studio style, Charming hand-drawn anime-style illustration", Image.open("./test_imgs/06.png"), 512, 512, 1, "Ghibli"],
-    ["Ghibli Studio style, Charming hand-drawn anime-style illustration", Image.open("./test_imgs/07.png"), 512, 512, 1, "Ghibli"],
-    ["Ghibli Studio style, Charming hand-drawn anime-style illustration", Image.open("./test_imgs/08.png"), 512, 512, 1, "Ghibli"],
-    ["Ghibli Studio style, Charming hand-drawn anime-style illustration", Image.open("./test_imgs/09.png"), 512, 512, 1, "Ghibli"],
-]
+# Save user task
+def submit_task(prompt, image):
+    task_id = str(uuid.uuid4())
+    task_img_path = os.path.join(task_folder, f"{task_id}.png")
+    task_txt_path = os.path.join(task_folder, f"{task_id}.txt")
+    image.save(task_img_path)
+    with open(task_txt_path, "w") as f:
+        f.write(prompt)
+    return f"✅ Task submitted successfully with ID: {task_id}"
 
-# Create the Gradio Blocks interface
+# ---- GUI ----
 with gr.Blocks() as demo:
     gr.Markdown("# Ghibli Studio Control Image Generation with EasyControl")
-    gr.Markdown("The model is trained on **only 100 real Asian faces** paired with **GPT-4o-generated Ghibli-style counterparts**, and it preserves facial features while applying the iconic anime aesthetic.")
-    gr.Markdown("Generate images using EasyControl with Ghibli control LoRAs.（Running on CPU due to free tier limitations; expect slower performance and lower resolution.）")
-    
-    gr.Markdown("**[Attention!!]**：The recommended prompts for using Ghibli Control LoRA should include the trigger words: `Ghibli Studio style, Charming hand-drawn anime-style illustration`")
-    gr.Markdown("😊😊If you like this demo, please give us a star (github: [EasyControl](https://github.com/Xiaojiu-z/EasyControl))")
 
     with gr.Tab("Ghibli Condition Generation"):
         with gr.Row():
             with gr.Column():
                 prompt = gr.Textbox(label="Prompt", value="Ghibli Studio style, Charming hand-drawn anime-style illustration")
                 spatial_img = gr.Image(label="Ghibli Image", type="pil")
-                height = gr.Slider(minimum=256, maximum=512, step=64, label="Height", value=512)  # 限制最大分辨率
-                width = gr.Slider(minimum=256, maximum=512, step=64, label="Width", value=512)    # 限制最大分辨率
+                height = gr.Slider(minimum=256, maximum=512, step=64, label="Height", value=512)
+                width = gr.Slider(minimum=256, maximum=512, step=64, label="Width", value=512)
                 seed = gr.Number(label="Seed", value=42)
-                control_type = gr.Dropdown(choices=control_types, label="Control Type")
-                single_generate_btn = gr.Button("Generate Image")
+                control_type = gr.Dropdown(choices=["Ghibli"], label="Control Type")
+                generate_btn = gr.Button("Generate Image")
             with gr.Column():
-                single_output_image = gr.Image(label="Generated Image")
+                output_img = gr.Image(label="Generated Image")
 
-        gr.Examples(
-            examples=single_examples,
+        generate_btn.click(
+            single_condition_generate_image_with_save,
             inputs=[prompt, spatial_img, height, width, seed, control_type],
-            outputs=single_output_image,
-            fn=single_condition_generate_image,
-            cache_examples=False,
-            label="Single Condition Examples"
+            outputs=output_img
         )
 
-    single_generate_btn.click(
-        single_condition_generate_image,
-        inputs=[prompt, spatial_img, height, width, seed, control_type],
-        outputs=single_output_image
-    )
+    with gr.Tab("Gallery"):
+        gallery_gallery = gr.Gallery(label="Previous Generations", show_label=False, elem_id="gallery").style(grid=3)
+
+        gallery_btn = gr.Button("🔄 Refresh Gallery")
+        gallery_btn.click(fn=load_gallery, outputs=gallery_gallery)
+
+    with gr.Tab("Tasks"):
+        with gr.Row():
+            task_prompt = gr.Textbox(label="Your Prompt")
+            task_image = gr.Image(label="Your Image", type="pil")
+        task_submit_btn = gr.Button("Submit Task")
+        task_submit_output = gr.Textbox(label="Submission Result")
+
+        task_submit_btn.click(fn=submit_task, inputs=[task_prompt, task_image], outputs=task_submit_output)
 
-# Launch the Gradio app
-demo.queue().launch()
+demo.queue().launch()