Update app.py

app.py

@@ -21,21 +21,18 @@ torch.backends.cudnn.deterministic = True
 
 # Initialize the models
 base_model = "runwayml/stable-diffusion-v1-5"
-dtype = torch.float16
 
 # Load the custom UNet
 unet = UNet2DConditionModelEx.from_pretrained(
     base_model,
-    subfolder="unet",
-    torch_dtype=dtype
+    subfolder="unet"
 )
 
 unet = unet.add_extra_conditions("ow-gbi-control-lora")
 
 pipe = StableDiffusionControlLoraV3Pipeline.from_pretrained(
     base_model, 
-    unet=unet,
-    torch_dtype=dtype
+    unet=unet
 )
 
 pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config)
@@ -45,9 +42,6 @@ pipe.load_lora_weights(
     weight_name="40kHalf.safetensors"
 )
 
-# Move to CPU initially
-pipe = pipe.to("cpu")
-
 def get_random_condition_image():
     conditions_dir = Path("conditions")
     if conditions_dir.exists():
@@ -63,16 +57,16 @@ def get_canny_image(image, low_threshold=100, high_threshold=200):
     elif isinstance(image, str):
         image = np.array(Image.open(image))
     
-    if len(image.shape) == 2:  # If grayscale
+    if len(image.shape) == 2:
         image = np.stack([image] * 3, axis=-1)
-    elif image.shape[2] == 4:  # If RGBA
+    elif image.shape[2] == 4:
         image = image[..., :3]
     
     canny_image = cv2.Canny(image, low_threshold, high_threshold)
     canny_image = np.stack([canny_image] * 3, axis=-1)
     return Image.fromarray(canny_image)
 
-@spaces.GPU(duration=300)  # Increased duration to 5 minutes
+@spaces.GPU(duration=300)
 def generate_image(input_image, prompt, negative_prompt, guidance_scale, steps, low_threshold, high_threshold, seed, progress=gr.Progress()):
     if input_image is None:
         raise gr.Error("Please provide an input image!")
@@ -89,35 +83,22 @@ def generate_image(input_image, prompt, negative_prompt, guidance_scale, steps,
         progress(0.1, desc="Processing input image...")
         canny_image = get_canny_image(input_image, low_threshold, high_threshold)
         
-        progress(0.2, desc="Moving model to device...")
-        # Move pipeline to GPU
-        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-        pipe.to(device)
-        
         progress(0.3, desc="Generating image...")
         with torch.no_grad():
             image = pipe(
                 prompt=prompt,
                 negative_prompt=negative_prompt,
-                num_inference_steps=steps,
-                guidance_scale=guidance_scale,
+                num_inference_steps=int(steps),  # Convert to int
+                guidance_scale=float(guidance_scale),  # Convert to float
                 image=canny_image,
                 extra_condition_scale=1.0,
                 generator=generator
             ).images[0]
         
-        progress(0.9, desc="Moving model back to CPU...")
-        # Move back to CPU to free up GPU memory
-        pipe.to("cpu")
-        torch.cuda.empty_cache()
-        
         progress(1.0, desc="Done!")
         return canny_image, image
     
     except Exception as e:
-        # Move back to CPU in case of error
-        pipe.to("cpu")
-        torch.cuda.empty_cache()
         raise gr.Error(f"An error occurred: {str(e)}")
 
 def random_image_click():
@@ -204,13 +185,11 @@ with gr.Blocks() as demo:
         cache_examples=True
     )
 
-    # Handle the random image button
     random_image_btn.click(
         fn=random_image_click,
         outputs=input_image
     )
 
-    # Handle the generate button
     generate.click(
         fn=generate_image,
         inputs=[
@@ -226,5 +205,5 @@ with gr.Blocks() as demo:
         outputs=[canny_output, result]
     )
 
-demo.queue()  # Enable queuing for better handling of concurrent requests
-demo.launch()
+demo.queue()
+demo.launch(share=True)