Update app.py

app.py

@@ -10,41 +10,38 @@ from image_gen_aux import DepthPreprocessor
 MAX_SEED = np.iinfo(np.int32).max
 MAX_IMAGE_SIZE = 2048
 
-device = "cuda" if torch.cuda.is_available() else "cpu"
-
-# Initialize the pipeline and move it to GPU
+# Initialize models without moving to CUDA yet
 pipe = FluxControlPipeline.from_pretrained(
     "black-forest-labs/FLUX.1-Depth-dev", 
     torch_dtype=torch.bfloat16
-).to(device)
-
+)
 processor = DepthPreprocessor.from_pretrained("LiheYoung/depth-anything-large-hf")
 
+@spaces.GPU
 def load_lora(lora_path):
     if not lora_path.strip():
         return "Please provide a valid LoRA path"
     try:
+        # Move to GPU within the wrapped function
+        pipe.to("cuda")
+        
         # Unload any existing LoRA weights first
         try:
             pipe.unload_lora_weights()
         except:
             pass
             
-        # Load new LoRA weights and move to the same device
+        # Load new LoRA weights
         pipe.load_lora_weights(lora_path)
-        
-        # Ensure all model components are on the correct device
-        pipe.to(device)
-        
         return f"Successfully loaded LoRA weights from {lora_path}"
     except Exception as e:
         return f"Error loading LoRA weights: {str(e)}"
 
+@spaces.GPU
 def unload_lora():
     try:
+        pipe.to("cuda")
         pipe.unload_lora_weights()
-        # Ensure model is on correct device after unloading
-        pipe.to(device)
         return "Successfully unloaded LoRA weights"
     except Exception as e:
         return f"Error unloading LoRA weights: {str(e)}"
@@ -57,6 +54,9 @@ def infer(control_image, prompt, seed=42, randomize_seed=False, width=1024, heig
         seed = random.randint(0, MAX_SEED)
     
     try:
+        # Move pipeline to GPU within the wrapped function
+        pipe.to("cuda")
+        
         # Process control image
         control_image = processor(control_image)[0].convert("RGB")
         
@@ -68,7 +68,7 @@ def infer(control_image, prompt, seed=42, randomize_seed=False, width=1024, heig
             width=width,
             num_inference_steps=num_inference_steps,
             guidance_scale=guidance_scale,
-            generator=torch.Generator(device=device).manual_seed(seed),
+            generator=torch.Generator("cuda").manual_seed(seed),
         ).images[0]
         
         return image, seed
@@ -113,6 +113,7 @@ with gr.Blocks(css=css) as demo:
             run_button = gr.Button("Run", scale=0)
         
         result = gr.Image(label="Result", show_label=False)
+        error_message = gr.Textbox(label="Error", visible=False)
         
         with gr.Accordion("Advanced Settings", open=False):
             seed = gr.Slider(