Update app.py

app.py

@@ -17,13 +17,22 @@ pipe = FluxControlPipeline.from_pretrained(
 )
 processor = DepthPreprocessor.from_pretrained("LiheYoung/depth-anything-large-hf")
 
+def cleanup_memory():
+    """Clean up GPU memory"""
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
+        torch.cuda.ipc_collect()
+
 @spaces.GPU
 def load_lora(lora_path):
     if not lora_path.strip():
         return "Please provide a valid LoRA path"
     try:
+        cleanup_memory()
+        
         # Move to GPU within the wrapped function
         pipe.to("cuda")
+        pipe.enable_model_cpu_offload()
         
         # Unload any existing LoRA weights first
         try:
@@ -35,25 +44,38 @@ def load_lora(lora_path):
         pipe.load_lora_weights(lora_path)
         return f"Successfully loaded LoRA weights from {lora_path}"
     except Exception as e:
+        cleanup_memory()
         return f"Error loading LoRA weights: {str(e)}"
 
 @spaces.GPU
 def unload_lora():
     try:
+        cleanup_memory()
         pipe.to("cuda")
         pipe.unload_lora_weights()
         return "Successfully unloaded LoRA weights"
     except Exception as e:
+        cleanup_memory()
         return f"Error unloading LoRA weights: {str(e)}"
 
+def round_to_multiple(number, multiple):
+    """Round a number to the nearest multiple"""
+    return multiple * round(number / multiple)
+
 @spaces.GPU
 def infer(control_image, prompt, seed=42, randomize_seed=False, width=1024, height=1024, 
           guidance_scale=3.5, num_inference_steps=28, progress=gr.Progress(track_tqdm=True)):
     
-    if randomize_seed:
-        seed = random.randint(0, MAX_SEED)
-    
     try:
+        cleanup_memory()
+        
+        if randomize_seed:
+            seed = random.randint(0, MAX_SEED)
+        
+        # Ensure dimensions are divisible by 16
+        width = round_to_multiple(width, 16)
+        height = round_to_multiple(height, 16)
+        
         # Move pipeline to GPU within the wrapped function
         pipe.to("cuda")
         
@@ -61,18 +83,21 @@ def infer(control_image, prompt, seed=42, randomize_seed=False, width=1024, heig
         control_image = processor(control_image)[0].convert("RGB")
         
         # 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("cuda").manual_seed(seed),
-        ).images[0]
+        with torch.inference_mode():
+            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("cuda").manual_seed(seed),
+            ).images[0]
         
+        cleanup_memory()
         return image, seed
     except Exception as e:
+        cleanup_memory()
         return None, f"Error during inference: {str(e)}"
 
 css="""
@@ -131,7 +156,7 @@ with gr.Blocks(css=css) as demo:
                     label="Width",
                     minimum=256,
                     maximum=MAX_IMAGE_SIZE,
-                    step=32,
+                    step=16,  # Changed to 16 to ensure divisibility
                     value=1024,
                 )
                 
@@ -139,7 +164,7 @@ with gr.Blocks(css=css) as demo:
                     label="Height",
                     minimum=256,
                     maximum=MAX_IMAGE_SIZE,
-                    step=32,
+                    step=16,  # Changed to 16 to ensure divisibility
                     value=1024,
                 )