Update optimized.py

optimized.py

@@ -83,38 +83,25 @@ 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):
-    # Clean up input handling
+    # Load control image
     control_image = load_image(control_image)
     w, h = control_image.size
-    scale = min(scale, 2.0)  # Cap scale factor
-    
-    # Size calculation with safety limits
-    max_dim = 1536  # Set based on your VRAM
-    target_w = min(int(w * scale), max_dim)
-    target_h = min(int(h * scale), max_dim)
-    
-    control_image = control_image.resize(
-        (target_w, target_h),
-        PIL.Image.BICUBIC
-    )
-    
-    # Generation with memory-friendly parameters
-    with torch.autocast("cuda"):  # Mixed precision
-        image = pipe(
-            prompt=prompt,
-            control_image=control_image,
-            controlnet_conditioning_scale=controlnet_conditioning_scale,
-            num_inference_steps=steps,
-            guidance_scale=guidance_scale,
-            height=target_h,
-            width=target_w,
-            output_type="pil",  # Avoid extra latent decoding steps
-            generator=torch.Generator(device="cuda").manual_seed(0)
-        ).images[0]
+    # Upscale x1
+    control_image = control_image.resize((int(w * scale), int(h * scale)))
+    print("Size to: " + str(control_image.size[0]) + ", " + str(control_image.size[1]))
+    image = pipe(
+        prompt=prompt,
+        control_image=control_image,
+        controlnet_conditioning_scale=controlnet_conditioning_scale,
+        num_inference_steps=steps,
+        guidance_scale=guidance_scale,
+        height=control_image.size[1],
+        width=control_image.size[0]
+    ).images[0]
     print(f"VRAM used: {torch.cuda.memory_allocated()/1e9:.2f}GB")
     # Aggressive memory cleanup
-    torch.cuda.empty_cache()
-    torch.cuda.ipc_collect()
+    # torch.cuda.empty_cache()
+    # torch.cuda.ipc_collect()
     print(f"VRAM used: {torch.cuda.memory_allocated()/1e9:.2f}GB")
     return image
 # Create Gradio interface