Update app.py

app.py

@@ -1,5 +1,4 @@
 import spaces
-
 import gradio as gr
 import cv2
 import numpy as np
@@ -42,7 +41,6 @@ pipe = StableDiffusionControlNetInpaintPipeline.from_pretrained(
     torch_dtype=torch.float16,
 ).to(device)
 
-
 # Set the K_EULER scheduler
 pipe.scheduler = EulerDiscreteScheduler.from_config(pipe.scheduler.config)
 
@@ -80,12 +78,15 @@ def resize_image(image, max_size=1536):
         new_size = (int(w * scale), int(h * scale))
         image = cv2.resize(image, new_size, interpolation=cv2.INTER_AREA)
     return image
-    
+
 @spaces.GPU(duration=60)
 # Function to inpaint the hair area using ControlNet
 def inpaint_hair(image, prompt):
-    # Resize the input image if necessary
-    image = resize_image(image)
+    # Only resize the input image if it's larger than 1536 in any dimension
+    h, w = image.shape[:2]
+    if max(h, w) > 1536:
+        image = resize_image(image)
+    
     # Segment hair to get the mask
     mask = segment_hair(image)
     # Convert to PIL image for the inpainting pipeline
@@ -102,16 +103,18 @@ def inpaint_hair(image, prompt):
     # Generate inpainted image
     generator = torch.Generator(device).manual_seed(42)
     negative_prompt = "lowres, bad quality, poor quality"
-    output = pipe(
-        prompt=prompt,
-        negative_prompt=negative_prompt,
-        image=image_pil,
-        mask_image=mask_pil,
-        control_image=inpaint_condition,
-        num_inference_steps=25,
-        guidance_scale=7.5,
-        generator=generator
-    ).images[0]
+    
+    with torch.cuda.amp.autocast():  # Enable automatic mixed precision
+        output = pipe(
+            prompt=prompt,
+            negative_prompt=negative_prompt,
+            image=image_pil,
+            mask_image=mask_pil,
+            control_image=inpaint_condition,
+            num_inference_steps=25,
+            guidance_scale=7.5,
+            generator=generator
+        ).images[0]
 
     return np.array(output)
 
@@ -129,4 +132,4 @@ iface = gr.Interface(
 )
 
 if __name__ == "__main__":
-    iface.launch()
+    iface.launch()