Update app.py

app.py

@@ -1,5 +1,6 @@
 import os
 import torch
+import gc
 import gradio as gr
 import numpy as np
 from PIL import Image
@@ -13,17 +14,18 @@ from diffusers import FluxControlNetPipeline, FluxControlNetModel
 
 # Source: https://github.com/XLabs-AI/x-flux.git
 name = "flux-dev"
-device = torch.device("cuda")
-offload = False
-is_schnell = name == "flux-schnell"
+
+# Load the model on CPU
+device_cpu = torch.device("cpu")
+device_gpu = torch.device("cuda")
 
 base_model = 'black-forest-labs/FLUX.1-dev'
 controlnet_model = 'InstantX/FLUX.1-dev-Controlnet-Union'
 
-# Load the new ControlNet model and pipeline
-controlnet = FluxControlNetModel.from_pretrained(controlnet_model, torch_dtype=torch.bfloat16)
+# Load the ControlNet model and pipeline on CPU
+controlnet = FluxControlNetModel.from_pretrained(controlnet_model, torch_dtype=torch.bfloat16, device_map="cpu")
 pipe = FluxControlNetPipeline.from_pretrained(base_model, controlnet=controlnet, torch_dtype=torch.bfloat16)
-pipe.to(device)
+pipe.to(device_cpu)  # Keep on CPU initially
 
 controlnet_conditioning_scale = 0.5
 
@@ -68,15 +70,23 @@ def preprocess_image(image, target_width, target_height, crop=True):
     
     return image
 
+def clear_cuda_memory():
+    gc.collect()
+    torch.cuda.empty_cache()
+    torch.cuda.ipc_collect()
+
 @spaces.GPU(duration=120)
 def generate_image(prompt, control_image, control_mode, num_steps=50, guidance=4, width=512, height=512, seed=42, random_seed=False):
+    clear_cuda_memory()
+    
     if random_seed:
         seed = np.random.randint(0, 10000)
     
     if not os.path.isdir("./controlnet_results/"):
         os.makedirs("./controlnet_results/")
 
-    torch_device = torch.device("cuda")
+    # Move model to GPU for inference
+    pipe.to(device_gpu)
 
     control_image = preprocess_image(control_image, width, height)
     
@@ -93,6 +103,9 @@ def generate_image(prompt, control_image, control_mode, num_steps=50, guidance=4
             guidance_scale=guidance,
         ).images[0]
     
+    # Move model back to CPU after inference
+    pipe.to(device_cpu)
+
     return [control_image, image]  # Return both images for slider
 
 interface = gr.Interface(