Update app.py

app.py

@@ -27,21 +27,6 @@ device = torch.device("cuda")
 offload = False
 is_schnell = name == "flux-schnell"
 
-model, ae, t5, clip, controlnet = None, None, None, None, None
-
-def load_models():
-    global model, ae, t5, clip, controlnet
-    t5 = load_t5(device, max_length=256 if is_schnell else 512)
-    clip = load_clip(device)
-    model = load_flow_model(name, device=device)
-    ae = load_ae(name, device=device)
-    controlnet = load_controlnet(name, device).to(device).to(torch.bfloat16)
-
-    checkpoint = load_safetensors(model_path)
-    controlnet.load_state_dict(checkpoint, strict=False)
-
-load_models()
-
 def preprocess_image(image, target_width, target_height, crop=True):
     if crop:
         image = c_crop(image)  # Crop the image to square
@@ -78,11 +63,16 @@ def generate_image(prompt, control_image, num_steps=50, guidance=4, width=512, h
 
     torch_device = torch.device("cuda")
     
-    model.to(torch_device)
-    t5.to(torch_device)
-    clip.to(torch_device)
-    ae.to(torch_device)
-    controlnet.to(torch_device)
+    torch.cuda.empty_cache()  # Clear GPU cache
+    
+    model = load_flow_model(name, device=torch_device)
+    t5 = load_t5(torch_device, max_length=256 if is_schnell else 512)
+    clip = load_clip(torch_device)
+    ae = load_ae(name, device=torch_device)
+    controlnet = load_controlnet(name, torch_device).to(torch_device).to(torch.bfloat16)
+
+    checkpoint = load_safetensors(model_path)
+    controlnet.load_state_dict(checkpoint, strict=False)
 
     width = 16 * width // 16
     height = 16 * height // 16
@@ -116,8 +106,8 @@ interface = gr.Interface(
         gr.Image(type="pil", label="Control Image"),
         gr.Slider(step=1, minimum=1, maximum=64, value=28, label="Num Steps"),
         gr.Slider(minimum=0.1, maximum=10, value=4, label="Guidance"),
-        gr.Slider(minimum=128, maximum=2048, step=128, value=1024, label="Width"),
-        gr.Slider(minimum=128, maximum=2048, step=128, value=1024, label="Height"),
+        gr.Slider(minimum=128, maximum=1024, step=128, value=512, label="Width"),
+        gr.Slider(minimum=128, maximum=1024, step=128, value=512, label="Height"),
         gr.Number(value=42, label="Seed"),
         gr.Checkbox(label="Random Seed")
     ],
@@ -127,5 +117,4 @@ interface = gr.Interface(
 )
 
 if __name__ == "__main__":
-    interface.launch()
-    
+    interface.launch()