Update app.py

app.py

@@ -4,21 +4,38 @@ import spaces
 import torch
 import random
 from peft import PeftModel
-
 from diffusers import FluxControlPipeline, FluxTransformer2DModel
 from image_gen_aux import DepthPreprocessor
 
 MAX_SEED = np.iinfo(np.int32).max
 MAX_IMAGE_SIZE = 2048
 
-pipe = FluxControlPipeline.from_pretrained("black-forest-labs/FLUX.1-Depth-dev", torch_dtype=torch.bfloat16).to("cuda")
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+# Initialize the pipeline and move it to GPU
+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")
 
 def load_lora(lora_path):
     if not lora_path.strip():
         return "Please provide a valid LoRA path"
     try:
+        # Unload any existing LoRA weights first
+        try:
+            pipe.unload_lora_weights()
+        except:
+            pass
+            
+        # Load new LoRA weights and move to the same device
         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)}"
@@ -26,6 +43,8 @@ def load_lora(lora_path):
 def unload_lora():
     try:
         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)}"
@@ -36,18 +55,25 @@ def infer(control_image, prompt, seed=42, randomize_seed=False, width=1024, heig
     
     if randomize_seed:
         seed = random.randint(0, MAX_SEED)
+    
+    try:
+        # Process control image
+        control_image = processor(control_image)[0].convert("RGB")
         
-    control_image = processor(control_image)[0].convert("RGB")
-    image = pipe(
-        prompt=prompt,
-        control_image=control_image,
-        height=height,
-        width=width,
-        num_inference_steps=num_inference_steps,
-        guidance_scale=guidance_scale,
-        generator=torch.Generator().manual_seed(seed),
-    ).images[0]
-    return image, seed
+        # Generate image
+        image = pipe(
+            prompt=prompt,
+            control_image=control_image,
+            height=height,
+            width=width,
+            num_inference_steps=num_inference_steps,
+            guidance_scale=guidance_scale,
+            generator=torch.Generator(device=device).manual_seed(seed),
+        ).images[0]
+        
+        return image, seed
+    except Exception as e:
+        return None, f"Error during inference: {str(e)}"
 
 css="""
 #col-container {