Update app.py

app.py

@@ -43,10 +43,11 @@ LORA_REGISTRY = {
 # -----------------------------------------------------------------------------
 print("Initializing SDXL Inference Pipeline...")
 
+# On ZeroGPU, we initialize standard variables, but we rely on the decorator for device placement
 device = "cuda" if torch.cuda.is_available() else "cpu"
-dtype = torch.float16 if device == "cuda" else torch.float32
+dtype = torch.float16
 
-# 1. Load VAE (Critical for SDXL fp16 stability)
+# 1. Load VAE (Critical for SDXL fp16 stability to avoid NaNs)
 vae = AutoencoderKL.from_pretrained(
     "madebyollin/sdxl-vae-fp16-fix", 
     torch_dtype=dtype
@@ -71,11 +72,12 @@ pipe = StableDiffusionXLControlNetPipeline.from_pretrained(
 # Optimization
 pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config)
 
+# For ZeroGPU/Spaces, enable_model_cpu_offload is the standard way to handle SDXL
+# This registers hooks that automatically move layers to GPU when the @spaces.GPU function is called
 try:
     pipe.enable_model_cpu_offload()
 except Exception as e:
-    print(f"Warning: CPU offload failed, moving to device manually. {e}")
-    pipe.to(device)
+    print(f"Offload warning: {e}")
 
 print("SDXL Pipeline Loaded Successfully.")
 
@@ -112,9 +114,7 @@ def generate_controlled_image(
     input_image = input_image.resize((width, height))
     canny_image = get_canny_image(input_image)
 
-    # 2. Manage LoRA State
-    pipe.unload_lora_weights()
-    
+    # 2. Configuration
     style_config = LORA_REGISTRY[lora_selection]
     repo_id = style_config["repo"]
     trigger_text = style_config["trigger"]
@@ -123,9 +123,16 @@ def generate_controlled_image(
     
     final_prompt = f"{trigger_text}{prompt}"
     
+    # 3. LoRA & Generation Block
+    # We use a try/finally block to ensure LoRA is ALWAYS unloaded, 
+    # preventing state corruption on the shared GPU.
     try:
+        # A. Load LoRA
         if repo_id:
             print(f"Loading LoRA: {repo_id}")
+            # Ensure we are in a clean state before loading
+            pipe.unload_lora_weights()
+            
             if lora_file:
                 pipe.load_lora_weights(repo_id, weight_name=lora_file)
             else:
@@ -133,38 +140,36 @@ def generate_controlled_image(
                 
             pipe.fuse_lora(lora_scale=lora_weight)
             print("LoRA fused successfully.")
-            
-    except Exception as e:
-        print(f"LoRA Load Error: {e}")
-        gr.Warning(f"Failed to load LoRA {lora_selection}. Using base model. Error: {str(e)}")
-
-    # 3. Generation
-    generator = torch.Generator(device).manual_seed(int(seed))
 
-    print(f"Generating with Prompt: {final_prompt}")
-    
-    try:
+        # B. Generate
+        generator = torch.Generator("cuda").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.0, # SDXL usually prefers slightly lower CFG than SD1.5
+            guidance_scale=7.0,
             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()
+        
+    finally:
+        # C. Cleanup (Always run this)
+        if repo_id:
+            print("Cleaning up LoRA weights...")
+            try:
+                pipe.unfuse_lora()
+                pipe.unload_lora_weights()
+            except Exception as cleanup_error:
+                print(f"Cleanup warning: {cleanup_error}")
+        
+        # Explicit cache clearing for ZeroGPU shared environment
+        torch.cuda.empty_cache()
 
     return canny_image, output_image
 
@@ -177,7 +182,7 @@ css = """
 .guide-text {font-size: 1.1em; color: #4a5568;}
 """
 
-# Example Data (Updated for SDXL context)
+# Example Data
 examples = [
     [
         "https://huggingface.co/takuma104/controlnet_dev/resolve/main/gen_compare/control_images/converted/control_bird_canny.png",
@@ -272,7 +277,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=False # Keep False for stability
+            cache_examples=False # Must be False for ZeroGPU async compatibility
         )
 
     # Event Wiring