Update app.py

app.py

@@ -1,91 +1,106 @@
+import spaces
 import os
+import json
 import time
-from datetime import datetime
-import uuid
-import gradio as gr
+import torch
 from PIL import Image
+from tqdm import tqdm
+import gradio as gr
 
-# Define folders
-gallery_folder = "./gallery"
-task_folder = "./tasks"
-
-os.makedirs(gallery_folder, exist_ok=True)
-os.makedirs(task_folder, exist_ok=True)
-
-# 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")
+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
+
+# Initialize the image processor
+base_path = "black-forest-labs/FLUX.1-dev"    
+lora_base_path = "./models"
+
+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
+
+def clear_cache(transformer):
+    for name, attn_processor in transformer.attn_processors.items():
+        attn_processor.bank_kv.clear()
+
+# 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")
     
-    image_path = os.path.join(gallery_folder, filename)
-    metadata_path = os.path.join(gallery_folder, metadata_filename)
-
-    image.save(image_path)
-    with open(metadata_path, "w") as f:
-        f.write(f"Prompt: {prompt}\nSeed: {seed}\nTime: {timestamp}")
-
-# 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)
+    # 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)
     return image
 
-# 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
-
-# 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}"
-
-# ---- GUI ----
+# Define the Gradio interface components
+control_types = ["Ghibli"]
+
+# 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"],
+]
+
+# Create the Gradio Blocks interface
 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=["Ghibli"], label="Control Type")
-                generate_btn = gr.Button("Generate Image")
+                control_type = gr.Dropdown(choices=control_types, label="Control Type")
+                single_generate_btn = gr.Button("Generate Image")
             with gr.Column():
-                output_img = gr.Image(label="Generated Image")
+                single_output_image = gr.Image(label="Generated Image")
 
-        generate_btn.click(
-            single_condition_generate_image_with_save,
+        gr.Examples(
+            examples=single_examples,
             inputs=[prompt, spatial_img, height, width, seed, control_type],
-            outputs=output_img
+            outputs=single_output_image,
+            fn=single_condition_generate_image,
+            cache_examples=False,
+            label="Single Condition Examples"
         )
 
-    with gr.Tab("Gallery"):
-        gallery_gallery = gr.Gallery(label="Previous Generations", show_label=False, elem_id="gallery", columns=[3], rows=[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)
+    single_generate_btn.click(
+        single_condition_generate_image,
+        inputs=[prompt, spatial_img, height, width, seed, control_type],
+        outputs=single_output_image
+    )
 
-demo.queue().launch()
+# Launch the Gradio app
+demo.queue().launch()