Update app.py

app.py

@@ -201,7 +201,7 @@ def randomize_seed_fn(seed: int, randomize_seed: bool) -> int:
 
 @spaces.GPU
 def run(
-    image: PIL.Image.Image,
+    image: dict,
     prompt: str,
     negative_prompt: str,
     style_name: str = DEFAULT_STYLE_NAME,
@@ -213,59 +213,65 @@ def run(
     use_canny: bool = False,
     progress=gr.Progress(track_tqdm=True),
 ) -> PIL.Image.Image:
-    width, height  = image['composite'].size
-    ratio = np.sqrt(1024. * 1024. / (width * height))
-    new_width, new_height = int(width * ratio), int(height * ratio)
-    image = image['composite'].resize((new_width, new_height))
+    # Get the composite image from the EditorValue dict
+    composite_image = image['composite']
+    width, height = composite_image.size
+    
+    # Calculate new dimensions to fit within 1024x1024 while maintaining aspect ratio
+    max_size = 1024
+    ratio = min(max_size / width, max_size / height)
+    new_width = int(width * ratio)
+    new_height = int(height * ratio)
+    
+    # Resize the image
+    resized_image = composite_image.resize((new_width, new_height), Image.LANCZOS)
     
     if use_canny:
-        controlnet_img = np.array(image)
+        controlnet_img = np.array(resized_image)
         controlnet_img = cv2.Canny(controlnet_img, 100, 200)
         controlnet_img = HWC3(controlnet_img)
         image = Image.fromarray(controlnet_img)
-    
     elif not use_hed:
-          controlnet_img = image
+        controlnet_img = resized_image
     else:
-        controlnet_img = processor(image, scribble=False)
-      # following is some processing to simulate human sketch draw, different threshold can generate different width of lines
+        controlnet_img = processor(resized_image, scribble=False)
+        # Process controlnet_img as before...
         controlnet_img = np.array(controlnet_img)
         controlnet_img = nms(controlnet_img, 127, 3)
         controlnet_img = cv2.GaussianBlur(controlnet_img, (0, 0), 3)
-
-        # higher threshold, thiner line
         random_val = int(round(random.uniform(0.01, 0.10), 2) * 255)
         controlnet_img[controlnet_img > random_val] = 255
         controlnet_img[controlnet_img < 255] = 0
         image = Image.fromarray(controlnet_img)
     
-    
     prompt, negative_prompt = apply_style(style_name, prompt, negative_prompt)
 
     generator = torch.Generator(device=device).manual_seed(seed)
+    
     if use_canny:
         out = pipe_canny(
-        prompt=prompt,
-        negative_prompt=negative_prompt,
-        image=image,
-        num_inference_steps=num_steps,
-        generator=generator,
-        controlnet_conditioning_scale=controlnet_conditioning_scale,
-        guidance_scale=guidance_scale,
-        width=new_width,
-        height=new_height,
-    ).images[0]
+            prompt=prompt,
+            negative_prompt=negative_prompt,
+            image=image,
+            num_inference_steps=num_steps,
+            generator=generator,
+            controlnet_conditioning_scale=controlnet_conditioning_scale,
+            guidance_scale=guidance_scale,
+            width=new_width,
+            height=new_height,
+        ).images[0]
     else:
         out = pipe(
-        prompt=prompt,
-        negative_prompt=negative_prompt,
-        image=image,
-        num_inference_steps=num_steps,
-        generator=generator,
-        controlnet_conditioning_scale=controlnet_conditioning_scale,
-        guidance_scale=guidance_scale,
-        width=new_width,
-        height=new_height,).images[0]
+            prompt=prompt,
+            negative_prompt=negative_prompt,
+            image=image,
+            num_inference_steps=num_steps,
+            generator=generator,
+            controlnet_conditioning_scale=controlnet_conditioning_scale,
+            guidance_scale=guidance_scale,
+            width=new_width,
+            height=new_height,
+        ).images[0]
 
     return (controlnet_img, out)
 
@@ -281,7 +287,7 @@ with gr.Blocks(css="style.css", js=js_func) as demo:
     with gr.Row():
         with gr.Column():
             with gr.Group():
-                image = gr.ImageEditor(type="pil", image_mode="L", crop_size=(512, 512))
+                image = gr.ImageEditor(type="pil",label="Sketch your image or upload one", crop_size="1:1", width=1024, height=1024,)
                 prompt = gr.Textbox(label="Prompt")
                 style = gr.Dropdown(label="Style", choices=STYLE_NAMES, value=DEFAULT_STYLE_NAME)
                 use_hed = gr.Checkbox(label="use HED detector", value=False, info="check this box if you upload an image and want to turn it to a sketch")