Update app.py

app.py

@@ -9,12 +9,12 @@ from src_inference.pipeline import FluxPipeline
 from src_inference.lora_helper import set_single_lora
 
 BASE_PATH = "black-forest-labs/FLUX.1-dev"
-LOCAL_LORA_DIR = "./LoRAs"          
-CUSTOM_LORA_DIR = "./Custom_LoRAs"  
+LOCAL_LORA_DIR = "./LoRAs"
+CUSTOM_LORA_DIR = "./Custom_LoRAs"
 os.makedirs(LOCAL_LORA_DIR, exist_ok=True)
 os.makedirs(CUSTOM_LORA_DIR, exist_ok=True)
 
-# ------------------ DEVICE SETUP (✅ supports CPU-only Spaces) ------------------ #
+# ------------------ DEVICE SETUP ------------------ #
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 dtype = torch.bfloat16 if device.type == "cuda" else torch.float32
 print(f"🚀 Running on device: {device}")
@@ -37,32 +37,16 @@ pipe = FluxPipeline.from_pretrained(
 set_single_lora(pipe.transformer, omni_consistency_path,
                 lora_weights=[1], cond_size=512)
 
-# ------------------ Download LoRA Styles ------------------ #
-def download_all_loras():
-    lora_names = [
-        "3D_Chibi", "American_Cartoon", "Chinese_Ink", "Clay_Toy",
-        "Fabric", "Ghibli", "Irasutoya", "Jojo", "LEGO", "Line",
-        "Macaron", "Oil_Painting", "Origami", "Paper_Cutting",
-        "Picasso", "Pixel", "Poly", "Pop_Art", "Rick_Morty",
-        "Snoopy", "Van_Gogh", "Vector"
-    ]
-    for name in lora_names:
-        hf_hub_download(
-            repo_id="showlab/OmniConsistency",
-            filename=f"LoRAs/{name}_rank128_bf16.safetensors",
-            local_dir=LOCAL_LORA_DIR,
-        )
-download_all_loras()
-
+# ------------------ Util ------------------ #
 def clear_cache(transformer):
     for _, attn_processor in transformer.attn_processors.items():
         attn_processor.bank_kv.clear()
 
-# ------------------ Generation Function ------------------ #
-@spaces.GPU()  # Will fallback silently if GPU not available
+# ------------------ Generation ------------------ #
+@spaces.GPU()
 def generate_image(
-    lora_name,        
-    custom_repo_id,   
+    lora_name,
+    custom_repo_id,
     prompt,
     uploaded_image,
     width, height,
@@ -73,6 +57,7 @@ def generate_image(
     width, height = int(width), int(height)
     generator = torch.Generator("cpu").manual_seed(seed)
 
+    # Custom LoRA path
     if custom_repo_id and custom_repo_id.strip():
         repo_id = custom_repo_id.strip()
         try:
@@ -91,8 +76,12 @@ def generate_image(
         except Exception as e:
             raise gr.Error(f"Load custom LoRA failed: {e}")
     else:
-        lora_path = os.path.join(
-            f"{LOCAL_LORA_DIR}/LoRAs", f"{lora_name}_rank128_bf16.safetensors"
+        # Built-in LoRA: download only the one selected
+        lora_filename = f"LoRAs/{lora_name}_rank128_bf16.safetensors"
+        lora_path = hf_hub_download(
+            repo_id="showlab/OmniConsistency",
+            filename=lora_filename,
+            local_dir=LOCAL_LORA_DIR
         )
 
     pipe.unload_lora_weights()
@@ -104,13 +93,14 @@ def generate_image(
     except Exception as e:
         raise gr.Error(f"Load LoRA failed: {e}")
 
-    spatial_image  = [uploaded_image.convert("RGB")]
+    spatial_image = [uploaded_image.convert("RGB")]
     subject_images = []
+
     start = time.time()
     out_img = pipe(
         prompt,
         height=(height // 8) * 8,
-        width=(width  // 8) * 8,
+        width=(width // 8) * 8,
         guidance_scale=guidance_scale,
         num_inference_steps=num_inference_steps,
         max_sequence_length=512,
@@ -124,7 +114,7 @@ def generate_image(
     clear_cache(pipe.transformer)
     return uploaded_image, out_img
 
-# ------------------ UI Interface ------------------ #
+# ------------------ Gradio UI ------------------ #
 def create_interface():
     demo_lora_names = [
         "3D_Chibi", "American_Cartoon", "Chinese_Ink", "Clay_Toy",
@@ -135,22 +125,22 @@ def create_interface():
     ]
 
     def update_trigger_word(lora_name, prompt):
-      for name in demo_lora_names:
-        trigger = " ".join(name.split("_")) + " style,"
-        prompt = prompt.replace(trigger, "")
-      new_trigger = " ".join(lora_name.split("_"))+ " style,"
-      return new_trigger + prompt
+        for name in demo_lora_names:
+            trigger = " ".join(name.split("_")) + " style,"
+            prompt = prompt.replace(trigger, "")
+        new_trigger = " ".join(lora_name.split("_")) + " style,"
+        return new_trigger + prompt
 
     examples = [
-        ["3D_Chibi", "", "3D Chibi style, Two smiling colleagues enthusiastically high-five in front of a whiteboard filled with technical notes about multimodal learning, reflecting a moment of success and collaboration at OpenAI.", 
+        ["3D_Chibi", "", "3D Chibi style, Two smiling colleagues high-five at a whiteboard filled with technical notes.", 
         Image.open("./test_imgs/00.png"), 680, 1024, 3.5, 24, 42],
-        ["Clay_Toy", "", "Clay Toy style, Three team members from OpenAI are gathered around a laptop in a cozy, festive setting, with holiday decorations in the background; one waves cheerfully while the others engage in light conversation, reflecting a relaxed and collaborative atmosphere.", 
+        ["Clay_Toy", "", "Clay Toy style, A holiday-themed OpenAI team photo full of smiles and warmth.", 
         Image.open("./test_imgs/01.png"), 560, 1024, 3.5, 24, 42],
-        ["American_Cartoon", "", "American Cartoon style, In a dramatic and comedic moment from a classic Chinese film, an intense elder with a white beard and red hat grips a younger man, declaring something with fervor, while the subtitle at the bottom reads 'I want them all' — capturing both tension and humor.",  
+        ["American_Cartoon", "", "American Cartoon style, A dramatic subtitle moment from a classic film.",  
         Image.open("./test_imgs/02.png"), 568, 1024, 3.5, 24, 42],
-        ["Origami", "", "Origami style, A thrilled fan wearing a Portugal football kit poses energetically with a smiling Cristiano Ronaldo, who gives a thumbs-up, as they stand side by side in a casual, cheerful moment—capturing the excitement of meeting a football legend.", 
+        ["Origami", "", "Origami style, A Portugal football fan posing with Cristiano Ronaldo.", 
         Image.open("./test_imgs/03.png"), 768, 672, 3.5, 24, 42],
-        ["Vector", "", "Vector style, A man glances admiringly at a passing woman, while his girlfriend looks at him in disbelief, perfectly capturing the theme of shifting attention and misplaced priorities in a humorous, relatable way.", 
+        ["Vector", "", "Vector style, The distracted boyfriend meme reimagined.", 
         Image.open("./test_imgs/04.png"), 512, 1024, 3.5, 24, 42]
     ]
 
@@ -169,33 +159,32 @@ def create_interface():
 
         with gr.Row():
             with gr.Column(scale=1):
-               image_input = gr.Image(type="pil", label="Upload Image")
-               prompt_box = gr.Textbox(label="Prompt",
-                                        value="3D Chibi style,",
-                                        info="Remember to include the necessary trigger words if you're using a custom LoRA."
+                image_input = gr.Image(type="pil", label="Upload Image")
+                prompt_box = gr.Textbox(
+                    label="Prompt",
+                    value="3D Chibi style,",
+                    info="Include a style like 'Ghibli style,' in your prompt for better results."
                 )
-               lora_dropdown = gr.Dropdown(
+                lora_dropdown = gr.Dropdown(
                     demo_lora_names, label="Select built-in LoRA")
-               custom_repo_box = gr.Textbox(
+                custom_repo_box = gr.Textbox(
                     label="Enter Custom LoRA",
-                    placeholder="LoRA Hugging Face path (e.g., 'username/repo_name')",
-                    info="If you want to use a custom LoRA, enter its Hugging Face repo ID here and built-in LoRA will be Overridden. Leave empty to use built-in LoRAs. [Check the list of FLUX LoRAs](https://huggingface.co/models?other=base_model:adapter:black-forest-labs/FLUX.1-dev)"
+                    placeholder="e.g. username/repo_name",
+                    info="Overrides built-in LoRA if provided."
                 )
-               gen_btn = gr.Button("Generate")
+                gen_btn = gr.Button("Generate")
+
             with gr.Column(scale=1):
                 output_image = gr.ImageSlider(label="Generated Image")
+
         with gr.Accordion("Advanced Options", open=False):
-          height_box = gr.Textbox(value="1024", label="Height")
-          width_box  = gr.Textbox(value="1024", label="Width")
-          guidance_slider = gr.Slider(
-              0.1, 20, value=3.5, step=0.1, label="Guidance Scale")
-          steps_slider = gr.Slider(
-              1, 50, value=25, step=1, label="Inference Steps")
-          seed_slider = gr.Slider(
-              1, 2_147_483_647, value=42, step=1, label="Seed")
-
-        lora_dropdown.select(fn=update_trigger_word, inputs=[lora_dropdown,prompt_box], 
-                             outputs=prompt_box)
+            height_box = gr.Textbox(value="1024", label="Height")
+            width_box = gr.Textbox(value="1024", label="Width")
+            guidance_slider = gr.Slider(0.1, 20, value=3.5, step=0.1, label="Guidance Scale")
+            steps_slider = gr.Slider(1, 50, value=25, step=1, label="Inference Steps")
+            seed_slider = gr.Slider(1, 2_147_483_647, value=42, step=1, label="Seed")
+
+        lora_dropdown.select(fn=update_trigger_word, inputs=[lora_dropdown, prompt_box], outputs=prompt_box)
 
         gr.Examples(
             examples=examples,
@@ -213,8 +202,10 @@ def create_interface():
                     width_box, height_box, guidance_slider, steps_slider, seed_slider],
             outputs=output_image
         )
+
     return demo
 
+# ------------------ Run ------------------ #
 if __name__ == "__main__":
     demo = create_interface()
     demo.launch(ssr_mode=False)