Update app.py

app.py

@@ -24,6 +24,19 @@ def set_timesteps_patched(self, num_inference_steps: int, device = None):
     self._begin_index = None
     self.sigmas = self.sigmas.to("cpu")  # to avoid too much CPU/GPU communication
 
+def resize_image(image, resolution):
+    original_width, original_height = image.size
+    
+    if original_width > original_height:
+        new_width = resolution
+        new_height = int((resolution / original_width) * original_height)
+    else:
+        new_height = resolution
+        new_width = int((resolution / original_height) * original_width)
+    
+    resized_img = image.resize((new_width, new_height), Image.ANTIALIAS)
+    return resized_img
+
 EDMEulerScheduler.set_timesteps = set_timesteps_patched
 
 pipe_edit = CosStableDiffusionXLInstructPix2PixPipeline.from_single_file(
@@ -37,14 +50,16 @@ pipe_normal.scheduler = EDMEulerScheduler(sigma_min=0.002, sigma_max=120.0, sigm
 pipe_normal.to("cuda")
 
 @spaces.GPU
-def run_normal(prompt, negative_prompt="", guidance_scale=7, progress=gr.Progress(track_tqdm=True)):
-    return pipe_normal(prompt, negative_prompt=negative_prompt, guidance_scale=guidance_scale, num_inference_steps=20).images[0]
+def run_normal(prompt, negative_prompt="", guidance_scale=7, steps=20, progress=gr.Progress(track_tqdm=True)):
+    return pipe_normal(prompt, negative_prompt=negative_prompt, guidance_scale=guidance_scale, num_inference_steps=steps).images[0]
 
 @spaces.GPU
-def run_edit(image, prompt, negative_prompt="", guidance_scale=7, progress=gr.Progress(track_tqdm=True)):
-    resolution = 1024
-    image.resize((resolution, resolution))
-    return pipe_edit(prompt=prompt,image=image,height=resolution,width=resolution,negative_prompt=negative_prompt, guidance_scale=guidance_scale,num_inference_steps=20).images[0]
+def run_edit(image, prompt, negative_prompt="", guidance_scale=7, steps=20, progress=gr.Progress(track_tqdm=True)):
+    image = resize_image(image, 1024)
+    print("Image resized to ", image.size)
+    width, height = image.size
+    #image.resize((resolution, resolution))
+    return pipe_edit(prompt=prompt,image=image,height=height,width=width,negative_prompt=negative_prompt, guidance_scale=guidance_scale,num_inference_steps=steps).images[0]
 css = '''
 .gradio-container{
 max-width: 768px !important;
@@ -67,6 +82,7 @@ with gr.Blocks(css=css) as demo:
           with gr.Accordion("Advanced Settings", open=False):
             negative_prompt_edit = gr.Textbox(label="Negative Prompt")
             guidance_scale_edit = gr.Number(label="Guidance Scale", value=7)
+            steps_edit = gr.Slider(label="Steps", minimum=10, maximum=50, value=20)
       gr.Examples(examples=edit_examples, fn=run_edit, inputs=[image_edit, prompt_edit], outputs=[output_edit], cache_examples=True)
     with gr.Tab("CosXL"):
       with gr.Group():
@@ -77,6 +93,7 @@ with gr.Blocks(css=css) as demo:
           with gr.Accordion("Advanced Settings", open=False):
             negative_prompt_normal = gr.Textbox(label="Negative Prompt")
             guidance_scale_normal = gr.Number(label="Guidance Scale", value=7)
+            steps_normal = gr.Slider(label="Steps", minimum=10, maximum=50, value=20)
       gr.Examples(examples=normal_examples, fn=run_normal, inputs=[prompt_normal], outputs=[output_normal], cache_examples=True) 
     
     gr.on(
@@ -85,7 +102,7 @@ with gr.Blocks(css=css) as demo:
             prompt_normal.submit
         ],
         fn=run_normal,
-        inputs=[prompt_normal, negative_prompt_normal, guidance_scale_normal],
+        inputs=[prompt_normal, negative_prompt_normal, guidance_scale_normal, steps_normal],
         outputs=[output_normal],
     )
     gr.on(
@@ -94,7 +111,7 @@ with gr.Blocks(css=css) as demo:
             prompt_edit.submit
         ],
         fn=run_edit,
-        inputs=[image_edit, prompt_edit, negative_prompt_edit, guidance_scale_edit],
+        inputs=[image_edit, prompt_edit, negative_prompt_edit, guidance_scale_edit, steps_edit],
         outputs=[output_edit]
     )
 if __name__ == "__main__":