Update app.py

app.py

@@ -12,8 +12,6 @@ from diffusers import StableDiffusionControlNetPipeline, ControlNetModel, UniPCM
 # -----------------------------------------------------------------------------
 # 1. Configuration & Registry
 # -----------------------------------------------------------------------------
-# This dictionary serves as the "Registry" for valid models and their specific
-# trigger words. This decouples configuration from logic.
 LORA_REGISTRY = {
     "None (Base SD1.5)": {
         "repo": None,
@@ -22,9 +20,9 @@ LORA_REGISTRY = {
     },
     "Lego Style": {
         "repo": "lordjia/lelo-lego-lora-for-xl-sd1-5", 
-        "trigger": "LEGO Creator, LEGO MiniFig, ", # Combined triggers for general usage
+        "trigger": "LEGO Creator, LEGO MiniFig, ", 
         "weight": 0.8,
-        "file": "Lego_XL_v2.1.safetensors" # Note: Auto-resolution usually handles this, but explicitly noted for context
+        "file": "Lego_XL_v2.1.safetensors"
     },
     "Claymation Style": {
         "repo": "DoctorDiffusion/doctor-diffusion-s-claymation-style-lora",
@@ -46,8 +44,7 @@ print("Initializing Inference Pipeline...")
 device = "cuda" if torch.cuda.is_available() else "cpu"
 dtype = torch.float16 if device == "cuda" else torch.float32
 
-# Load ControlNet (Canny Edge Detection)
-# We use the standard lllyasviel checkpoint which is the gold standard for SD1.5
+# Load ControlNet
 controlnet = ControlNetModel.from_pretrained(
     "lllyasviel/sd-controlnet-canny",
     torch_dtype=dtype,
@@ -55,7 +52,6 @@ controlnet = ControlNetModel.from_pretrained(
 )
 
 # Load Base Stable Diffusion 1.5
-# We use the official RunwayML checkpoint
 pipe = StableDiffusionControlNetPipeline.from_pretrained(
     "stable-diffusion-v1-5/stable-diffusion-v1-5",
     controlnet=controlnet,
@@ -63,12 +59,10 @@ pipe = StableDiffusionControlNetPipeline.from_pretrained(
     use_safetensors=True
 )
 
-# Optimization: Use UniPC Scheduler for fast convergence (20-30 steps)
 pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config)
 
-# Optimization: Offload model to CPU when not in use to save VRAM
-# Crucial for running on constrained hardware (e.g., free tier Spaces)
-pipe.enable_model_cpu_offload()
+if device == "cuda":
+    pipe.to(device)
 
 print("Base Pipeline Loaded Successfully.")
 
@@ -77,23 +71,14 @@ print("Base Pipeline Loaded Successfully.")
 # -----------------------------------------------------------------------------
 
 def get_canny_image(image, low_threshold=100, high_threshold=200):
-    """
-    Converts a PIL image into a Canny edge map.
-    The map is converted to RGB (3-channel) to match ControlNet input requirements.
-    """
     image_array = np.array(image)
-    
-    # Canny edge detection via OpenCV
     canny_edges = cv2.Canny(image_array, low_threshold, high_threshold)
-    
-    # Replicate the single channel to 3 channels (RGB)
     canny_edges = canny_edges[:, :, None]
     canny_edges = np.concatenate([canny_edges, canny_edges, canny_edges], axis=2)
-    
     return Image.fromarray(canny_edges)
 
 # -----------------------------------------------------------------------------
-# 4. Inference Logic (The "Middleware")
+# 4. Inference Logic
 # -----------------------------------------------------------------------------
 
 @spaces.GPU(duration=120)
@@ -110,52 +95,58 @@ def generate_controlled_image(
         raise gr.Error("Validation Error: Please upload an image first!")
 
     # 1. Preprocess Image
-    # Resizing to 512x512 is standard for SD1.5 to avoid duplication artifacts
     width, height = 512, 512
     input_image = input_image.resize((width, height))
     canny_image = get_canny_image(input_image)
 
     # 2. Manage LoRA State
-    # We must explicitly unload previous weights to prevent style contamination
+    pipe.unload_lora_weights()
+    
+    style_config = LORA_REGISTRY[lora_selection]
+    repo_id = style_config["repo"]
+    trigger_text = style_config["trigger"]
+    lora_weight = style_config["weight"]
+    
+    final_prompt = f"{trigger_text}{prompt}"
+    
     try:
-        pipe.unload_lora_weights()
-        
-        style_config = LORA_REGISTRY[lora_selection]
-        repo_id = style_config["repo"]
-        trigger_text = style_config["trigger"]
-        
-        # Modify prompt with trigger words
-        final_prompt = f"{trigger_text}{prompt}"
-        
         if repo_id:
             print(f"Loading LoRA: {repo_id}")
             pipe.load_lora_weights(repo_id)
-            
-            # Note: In more complex setups with multiple adapters, we would use 
-            # pipe.set_adapters() and fuse_lora(), but for single-style swap,
-            # load/unload is sufficient and memory-safe.
+            pipe.fuse_lora(lora_scale=lora_weight)
+            print("LoRA fused successfully.")
             
     except Exception as e:
         print(f"LoRA Load Error: {e}")
-        # Fallback to base model if LoRA fails, but warn user via prompt
-        final_prompt = prompt 
-        gr.Warning(f"Failed to load LoRA {lora_selection}. Using base model.")
+        gr.Warning(f"Failed to load LoRA {lora_selection}. Using base model. Error: {str(e)}")
 
-    # 3. Deterministic Generation
-    # Using a manual seed ensures reproducibility
+    # 3. Generation
     generator = torch.Generator(device).manual_seed(int(seed))
 
     print(f"Generating with Prompt: {final_prompt}")
     
-    output_image = pipe(
-        prompt=final_prompt,
-        negative_prompt=negative_prompt,
-        image=canny_image,
-        num_inference_steps=int(steps),
-        controlnet_conditioning_scale=float(controlnet_conditioning_scale),
-        guidance_scale=7.5, # Standard CFG scale
-        generator=generator,
-    ).images
+    try:
+        output_image = pipe(
+            prompt=final_prompt,
+            negative_prompt=negative_prompt,
+            image=canny_image,
+            num_inference_steps=int(steps),
+            controlnet_conditioning_scale=float(controlnet_conditioning_scale),
+            guidance_scale=7.5,
+            generator=generator,
+        ).images
+    except Exception as e:
+        pipe.unfuse_lora()
+        pipe.unload_lora_weights()
+        raise e
+
+    # 4. Cleanup
+    if repo_id:
+        print("Unfusing LoRA...")
+        pipe.unfuse_lora()
+        pipe.unload_lora_weights()
+    
+    torch.cuda.empty_cache()
 
     return canny_image, output_image
 
@@ -168,10 +159,7 @@ css = """
 .guide-text {font-size: 1.1em; color: #4a5568;}
 """
 
-# Example Data (Using Public Domain / CC0 URLs for reproducibility)
-# Nested list format:
-# Example Data
-# Format: [Image URL, Prompt, Negative Prompt, LoRA Selection, ControlNet Scale, Steps, Seed]
+# Example Data (Using resolve URLs)
 examples = [
     [
         "https://huggingface.co/takuma104/controlnet_dev/resolve/main/gen_compare/control_images/converted/control_bird_canny.png",
@@ -231,8 +219,6 @@ examples = [
     ]
 ]
 
-   
-
 with gr.Blocks(theme=gr.themes.Soft(), css=css) as demo:
     
     with gr.Column(elem_id="col-container"):
@@ -297,7 +283,7 @@ with gr.Blocks(theme=gr.themes.Soft(), css=css) as demo:
             inputs=[input_image, prompt, negative_prompt, lora_selection, controlnet_conditioning_scale, steps, seed],
             outputs=[output_canny, output_result],
             fn=generate_controlled_image,
-            cache_examples=True # Pre-compute examples for instant display
+            cache_examples=False # CRITICAL FIX: Set to False to prevent async loop errors
         )
 
     # Event Wiring
@@ -316,4 +302,4 @@ with gr.Blocks(theme=gr.themes.Soft(), css=css) as demo:
     )
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch()