Update optimized.py

optimized.py

@@ -8,6 +8,7 @@ from diffusers import FluxControlNetModel, FluxControlNetPipeline, AutoencoderKL
 import gradio as gr
 from accelerate import dispatch_model, infer_auto_device_map
 from PIL import Image
+import gc
 # Corrected and optimized FluxControlNet implementation
 
 def self_attention_slicing(module, slice_size=3):
@@ -83,6 +84,8 @@ pipe.enable_attention_slicing(1)
 print(f"VRAM used: {torch.cuda.memory_allocated()/1e9:.2f}GB")
 @spaces.GPU
 def generate_image(prompt, scale, steps, control_image, controlnet_conditioning_scale, guidance_scale):
+    print(f"Memory Usage: {torch.cuda.memory_summary(device=None, abbreviated=False)}")
+
     # Load control image
     control_image = load_image(control_image)
     w, h = control_image.size
@@ -100,8 +103,11 @@ def generate_image(prompt, scale, steps, control_image, controlnet_conditioning_
     ).images[0]
     print(f"VRAM used: {torch.cuda.memory_allocated()/1e9:.2f}GB")
     # Aggressive memory cleanup
-    # torch.cuda.empty_cache()
+    torch.cuda.empty_cache()
     # torch.cuda.ipc_collect()
+    
+    del variables
+    gc.collect()
     print(f"VRAM used: {torch.cuda.memory_allocated()/1e9:.2f}GB")
     return image
 # Create Gradio interface
@@ -122,6 +128,7 @@ iface = gr.Interface(
     description="Generate images using the FluxControlNetPipeline. Upload a control image and enter a prompt to create an image.",
 )
 print(f"Memory Usage: {torch.cuda.memory_summary(device=None, abbreviated=False)}")
-
+del variables
+gc.collect()
 # Launch the app
 iface.launch()