Update app.py

app.py

@@ -37,7 +37,7 @@ LORA_REGISTRY = {
 }
 
 # -----------------------------------------------------------------------------
-# Model Initialization (CPU only, ZeroGPU handles device transfer)
+# Model Initialization
 # -----------------------------------------------------------------------------
 print("Initializing SDXL Pipeline on CPU...")
 
@@ -62,7 +62,7 @@ pipe = StableDiffusionXLControlNetPipeline.from_pretrained(
 
 pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config)
 
-print("Pipeline loaded. ZeroGPU will handle device management.")
+print("Pipeline loaded successfully.")
 
 # -----------------------------------------------------------------------------
 # Helper Functions
@@ -92,6 +92,10 @@ def generate_controlled_image(
     if input_image is None:
         raise gr.Error("Please upload an image first!")
 
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    
+    pipe.to(device)
+    
     width, height = 1024, 1024
     input_image = input_image.resize((width, height))
     canny_image = get_canny_image(input_image)
@@ -101,7 +105,6 @@ def generate_controlled_image(
     style_config = LORA_REGISTRY[lora_selection]
     repo_id = style_config["repo"]
     trigger_text = style_config["trigger"]
-    lora_weight = style_config["weight"]
     lora_file = style_config.get("file", None)
     
     final_prompt = f"{trigger_text}{prompt}"
@@ -118,7 +121,7 @@ def generate_controlled_image(
             print(f"LoRA Load Error: {e}")
             gr.Warning(f"Failed to load LoRA. Using base model.")
 
-    generator = torch.Generator("cuda").manual_seed(int(seed))
+    generator = torch.Generator(device).manual_seed(int(seed))
 
     print(f"Generating: {final_prompt[:100]}...")
     
@@ -138,7 +141,9 @@ def generate_controlled_image(
         raise e
 
     pipe.unload_lora_weights()
-    torch.cuda.empty_cache()
+    
+    if device == "cuda":
+        torch.cuda.empty_cache()
 
     return canny_image, output_image