Update app.py

app.py

@@ -12,177 +12,101 @@ from src.pipeline import FluxPipeline
 from src.transformer_flux import FluxTransformer2DModel
 from src.lora_helper import set_single_lora, set_multi_lora, unset_lora
 
-class ImageProcessor:
-    def __init__(self, path):
-        device = "cuda"
-        self.pipe = FluxPipeline.from_pretrained(path, torch_dtype=torch.bfloat16, device=device)
-        transformer = FluxTransformer2DModel.from_pretrained(path, subfolder="transformer", torch_dtype=torch.bfloat16, device=device)
-        self.pipe.transformer = transformer
-        self.pipe.to(device)
-        
-    def clear_cache(self, transformer):
-        for name, attn_processor in transformer.attn_processors.items():
-            attn_processor.bank_kv.clear()
-            
-    @spaces.GPU()
-    def process_image(self, prompt='', subject_imgs=[], spatial_imgs=[], height=768, width=768, output_path=None, seed=42):
-        image = self.pipe(
-            prompt,
-            height=int(height),
-            width=int(width),
-            guidance_scale=3.5,
-            num_inference_steps=25,
-            max_sequence_length=512,
-            generator=torch.Generator("cpu").manual_seed(seed), 
-            subject_images=subject_imgs,
-            spatial_images=spatial_imgs,
-            cond_size=512,
-        ).images[0]
-        self.clear_cache(self.pipe.transformer)
-        if output_path:
-            image.save(output_path)
-        return image
+from huggingface_hub import hf_hub_download
+hf_hub_download(repo_id="Xiaojiu-Z/EasyControl", filename="models/Ghibli.safetensors", local_dir="./checkpoints/models/")
 
 # Initialize the image processor
 base_path = "black-forest-labs/FLUX.1-dev"    
-lora_base_path = "EasyControl/models"
-style_lora_base_path = "Shakker-Labs"
-processor = ImageProcessor(base_path)
+lora_base_path = "checkpoints/models/models"
 
-# Define the Gradio interface
-def single_condition_generate_image(prompt, subject_img, spatial_img, height, width, seed, control_type, style_lora=None):
-    # Set the control type
-    if control_type == "subject":
-        lora_path = os.path.join(lora_base_path, "subject.safetensors")
-    elif control_type == "depth":
-        lora_path = os.path.join(lora_base_path, "depth.safetensors")
-    elif control_type == "seg":
-        lora_path = os.path.join(lora_base_path, "seg.safetensors")
-    elif control_type == "pose":
-        lora_path = os.path.join(lora_base_path, "pose.safetensors")
-    elif control_type == "inpainting":
-        lora_path = os.path.join(lora_base_path, "inpainting.safetensors")
-    elif control_type == "hedsketch":
-        lora_path = os.path.join(lora_base_path, "hedsketch.safetensors")
-    elif control_type == "canny":
-        lora_path = os.path.join(lora_base_path, "canny.safetensors")
-    set_single_lora(processor.pipe.transformer, lora_path, lora_weights=[1], cond_size=512)
-    
-    # Set the style LoRA
-    if style_lora=="None":
-        pass
-    else:
-        if style_lora == "Simple_Sketch":
-            processor.pipe.unload_lora_weights()
-            style_lora_path = os.path.join(style_lora_base_path, "FLUX.1-dev-LoRA-Children-Simple-Sketch")
-            processor.pipe.load_lora_weights(style_lora_path, weight_name="FLUX-dev-lora-children-simple-sketch.safetensors")
-        if style_lora == "Text_Poster":
-            processor.pipe.unload_lora_weights()
-            style_lora_path = os.path.join(style_lora_base_path, "FLUX.1-dev-LoRA-Text-Poster")
-            processor.pipe.load_lora_weights(style_lora_path, weight_name="FLUX-dev-lora-Text-Poster.safetensors")
-        if style_lora == "Vector_Style":
-            processor.pipe.unload_lora_weights()
-            style_lora_path = os.path.join(style_lora_base_path, "FLUX.1-dev-LoRA-Vector-Journey")
-            processor.pipe.load_lora_weights(style_lora_path, weight_name="FLUX-dev-lora-Vector-Journey.safetensors")
+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")
 
-    # Process the image
-    subject_imgs = [subject_img] if subject_img else []
-    spatial_imgs = [spatial_img] if spatial_img else []
-    image = processor.process_image(prompt=prompt, subject_imgs=subject_imgs, spatial_imgs=spatial_imgs, height=height, width=width, seed=seed)
-    return image
+def clear_cache(transformer):
+    for name, attn_processor in transformer.attn_processors.items():
+        attn_processor.bank_kv.clear()
 
 # Define the Gradio interface
-def multi_condition_generate_image(prompt, subject_img, spatial_img, height, width, seed):
-    subject_path = os.path.join(lora_base_path, "subject.safetensors")
-    inpainting_path = os.path.join(lora_base_path, "inpainting.safetensors")
-    set_multi_lora(processor.pipe.transformer, [subject_path, inpainting_path], lora_weights=[[1],[1]],cond_size=512)
+@spaces.GPU()
+def dual_condition_generate_image(prompt, spatial_img, height, width, seed, control_type, zero_steps):
+    # 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)
 
     # Process the image
-    subject_imgs = [subject_img] if subject_img else []
     spatial_imgs = [spatial_img] if spatial_img else []
-    image = processor.process_image(prompt=prompt, subject_imgs=subject_imgs, spatial_imgs=spatial_imgs, height=height, width=width, seed=seed)
-    return image
 
-# Define the Gradio interface components
-control_types = ["subject", "depth", "pose", "inpainting", "hedsketch", "seg", "canny"]
-style_loras = ["Simple_Sketch", "Text_Poster", "Vector_Style", "None"]
-
-# Example data
-single_examples = [
-    ["A SKS in the library", Image.open("./test_imgs/subject1.png"), None, 1024, 1024, 5, "subject", None],
-    ["In a picturesque village, a narrow cobblestone street with rustic stone buildings, colorful blinds, and lush green spaces, a cartoon man drawn with simple lines and solid colors stands in the foreground, wearing a red shirt, beige work pants, and brown shoes, carrying a strap on his shoulder. The scene features warm and enticing colors, a pleasant fusion of nature and architecture, and the camera's perspective on the street clearly shows the charming and quaint environment., Integrating elements of reality and cartoon.", None, Image.open("./test_imgs/spatial1.png"), 1024, 1024, 1, "pose", "Vector_Style"],
-]
-multi_examples = [
-    ["A SKS on the car", Image.open("./test_imgs/subject2.png"), Image.open("./test_imgs/spatial2.png"), 1024, 1024, 7],
-]
+    # Image with use_zero_init=True
+    image_true = pipe(
+        prompt,
+        height=int(height),
+        width=int(width),
+        guidance_scale=3.5,
+        num_inference_steps=25,
+        max_sequence_length=512,
+        generator=torch.Generator("cpu").manual_seed(seed), 
+        subject_images=[],
+        spatial_images=spatial_imgs,
+        cond_size=512,
+        use_zero_init=True,
+        zero_steps=int(zero_steps)
+    ).images[0]
+    clear_cache(pipe.transformer)
+
+    # Image with use_zero_init=False
+    image_false = pipe(
+        prompt,
+        height=int(height),
+        width=int(width),
+        guidance_scale=3.5,
+        num_inference_steps=25,
+        max_sequence_length=512,
+        generator=torch.Generator("cpu").manual_seed(seed), 
+        subject_images=[],
+        spatial_images=spatial_imgs,
+        cond_size=512,
+        use_zero_init=False
+    ).images[0]
+    clear_cache(pipe.transformer)
+
+    return image_true, image_false
 
+# Define the Gradio interface components
+control_types = ["Ghibli"]
 
 # Create the Gradio Blocks interface
 with gr.Blocks() as demo:
-    gr.Markdown("# Image Generation with EasyControl")
-    gr.Markdown("Generate images using EasyControl with different control types and style LoRAs.")
+    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.（Due to hardware constraints, only low-resolution images can be generated. For high-resolution (1024+), please set up your own environment.）")
+
+    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("Single Condition Generation"):
+    with gr.Tab("Ghibli Condition Generation"):
         with gr.Row():
             with gr.Column():
-                prompt = gr.Textbox(label="Prompt")
-                subject_img = gr.Image(label="Subject Image", type="pil")  # 上传图像文件
-                spatial_img = gr.Image(label="Spatial Image", type="pil")  # 上传图像文件
-                height = gr.Slider(minimum=256, maximum=1536, step=64, label="Height", value=768)
-                width = gr.Slider(minimum=256, maximum=1536, step=64, label="Width", value=768)
+                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=1024, step=64, label="Height", value=768)
+                width = gr.Slider(minimum=256, maximum=1024, step=64, label="Width", value=768)
                 seed = gr.Number(label="Seed", value=42)
+                zero_steps = gr.Number(label="Zero Init Steps", value=1)
                 control_type = gr.Dropdown(choices=control_types, label="Control Type")
-                style_lora = gr.Dropdown(choices=style_loras, label="Style LoRA")
                 single_generate_btn = gr.Button("Generate Image")
             with gr.Column():
-                single_output_image = gr.Image(label="Generated Image")
-
-        # Add examples for Single Condition Generation
-        gr.Examples(
-            examples=single_examples,
-            inputs=[prompt, subject_img, spatial_img, height, width, seed, control_type, style_lora],
-            outputs=single_output_image,
-            fn=single_condition_generate_image,
-            cache_examples=False,  # 缓存示例结果以加快加载速度
-            label="Single Condition Examples"
-        )
-
-
-    with gr.Tab("Multi-Condition Generation"):
-        with gr.Row():
-            with gr.Column():
-                multi_prompt = gr.Textbox(label="Prompt")
-                multi_subject_img = gr.Image(label="Subject Image", type="pil")  # 上传图像文件
-                multi_spatial_img = gr.Image(label="Spatial Image", type="pil")  # 上传图像文件
-                multi_height = gr.Slider(minimum=256, maximum=1536, step=64, label="Height", value=768)
-                multi_width = gr.Slider(minimum=256, maximum=1536, step=64, label="Width", value=768)
-                multi_seed = gr.Number(label="Seed", value=42)
-                multi_generate_btn = gr.Button("Generate Image")
-            with gr.Column():
-                multi_output_image = gr.Image(label="Generated Image")
-                
-        # Add examples for Multi-Condition Generation
-        gr.Examples(
-            examples=multi_examples,
-            inputs=[multi_prompt, multi_subject_img, multi_spatial_img, multi_height, multi_width, multi_seed],
-            outputs=multi_output_image,
-            fn=multi_condition_generate_image,
-            cache_examples=False,  # 缓存示例结果以加快加载速度
-            label="Multi-Condition Examples"
-        )
-
+                image_with_zero_init = gr.Image(label="Image CFG-Zero*")
+                image_without_zero_init = gr.Image(label="Image CFG")
 
     # Link the buttons to the functions
     single_generate_btn.click(
-        single_condition_generate_image,
-        inputs=[prompt, subject_img, spatial_img, height, width, seed, control_type, style_lora],
-        outputs=single_output_image
-    )
-    multi_generate_btn.click(
-        multi_condition_generate_image,
-        inputs=[multi_prompt, multi_subject_img, multi_spatial_img, multi_height, multi_width, multi_seed],
-        outputs=multi_output_image
+        dual_condition_generate_image,
+        inputs=[prompt, spatial_img, height, width, seed, control_type, zero_steps],
+        outputs=[image_with_zero_init, image_without_zero_init]
     )
 
 # Launch the Gradio app
-demo.queue().launch()
+demo.queue().launch()