Update app.py

app.py

@@ -3,7 +3,6 @@ import numpy as np
 import spaces
 import torch
 import random
-import gc
 from peft import PeftModel
 from diffusers import FluxControlPipeline, FluxTransformer2DModel
 from image_gen_aux import DepthPreprocessor
@@ -11,107 +10,69 @@ from image_gen_aux import DepthPreprocessor
 MAX_SEED = np.iinfo(np.int32).max
 MAX_IMAGE_SIZE = 2048
 
-def init_pipeline():
-    """Initialize pipeline with memory-efficient settings"""
-    pipe = FluxControlPipeline.from_pretrained(
-        "black-forest-labs/FLUX.1-Depth-dev",
-        torch_dtype=torch.bfloat16,
-        low_cpu_mem_usage=True,
-        use_safetensors=True
-    )
-    return pipe
-
-# Initialize models without moving to CUDA
-pipe = init_pipeline()
+# Initialize models without moving to CUDA yet
+pipe = FluxControlPipeline.from_pretrained(
+    "black-forest-labs/FLUX.1-Depth-dev", 
+    torch_dtype=torch.bfloat16
+)
 processor = DepthPreprocessor.from_pretrained("LiheYoung/depth-anything-large-hf")
 
-def cleanup_memory():
-    """Aggressive memory cleanup"""
-    if torch.cuda.is_available():
-        with torch.cuda.device('cuda'):
-            torch.cuda.empty_cache()
-            torch.cuda.ipc_collect()
-    gc.collect()
-
-def reinit_pipeline():
-    """Reinitialize the pipeline if needed"""
-    global pipe
-    cleanup_memory()
-    pipe = init_pipeline()
-    cleanup_memory()
-
 @spaces.GPU
 def load_lora(lora_path):
-    global pipe
     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")
         
-        # Reinitialize pipeline
-        reinit_pipeline()
-        
-        # Enable sequential CPU offload
-        pipe.enable_sequential_cpu_offload()
-        
-        # Load LoRA weights
+        # Unload any existing LoRA weights first
+        try:
+            pipe.unload_lora_weights()
+        except:
+            pass
+            
+        # Load new LoRA weights
         pipe.load_lora_weights(lora_path)
-        
-        cleanup_memory()
         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():
-    global pipe
     try:
-        cleanup_memory()
-        reinit_pipeline()
-        pipe.enable_sequential_cpu_offload()
+        pipe.to("cuda")
         pipe.unload_lora_weights()
-        cleanup_memory()
         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):
-    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")
         
         # Process control image
         control_image = processor(control_image)[0].convert("RGB")
         
-        # Generate image with memory optimization
-        with torch.inference_mode(), torch.cuda.amp.autocast():
-            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]
+        # 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]
         
-        cleanup_memory()
         return image, seed
     except Exception as e:
-        cleanup_memory()
         return None, f"Error during inference: {str(e)}"
 
 css="""
@@ -122,6 +83,7 @@ css="""
 """
 
 with gr.Blocks(css=css) as demo:
+    
     with gr.Column(elem_id="col-container"):
         gr.Markdown(f"""# FLUX.1 Depth [dev] with LoRA Support
 12B param rectified flow transformer structural conditioning tuned, guidance-distilled from [FLUX.1 [pro]](https://blackforestlabs.ai/)  
@@ -169,7 +131,7 @@ with gr.Blocks(css=css) as demo:
                     label="Width",
                     minimum=256,
                     maximum=MAX_IMAGE_SIZE,
-                    step=16,
+                    step=32,
                     value=1024,
                 )
                 
@@ -177,7 +139,7 @@ with gr.Blocks(css=css) as demo:
                     label="Height",
                     minimum=256,
                     maximum=MAX_IMAGE_SIZE,
-                    step=16,
+                    step=32,
                     value=1024,
                 )