Update app.py

app.py

@@ -20,7 +20,7 @@ pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.conf
 
 # Function to generate an image
 @spaces.GPU()  # Adjust the duration as needed
-def generate_image(prompt, negative_prompt, use_defaults, width, height, guidance_scale, num_inference_steps, seed, randomize_seed, progress=gr.Progress()):
+def generate_image(prompt, negative_prompt, use_defaults, resolution, guidance_scale, num_inference_steps, seed, randomize_seed, progress=gr.Progress()):
     pipe.to('cuda')  # Move the model to GPU when the function is called
     
     if randomize_seed:
@@ -34,6 +34,7 @@ def generate_image(prompt, negative_prompt, use_defaults, width, height, guidanc
         progress(step / num_inference_steps)
         return
     
+    width, height = map(int, resolution.split('x'))
     image = pipe(
         prompt, 
         negative_prompt=negative_prompt,
@@ -52,12 +53,12 @@ def generate_image(prompt, negative_prompt, use_defaults, width, height, guidanc
     return image, seed
 
 # Define Gradio interface
-def interface_fn(prompt, negative_prompt, use_defaults, width, height, guidance_scale, num_inference_steps, seed, randomize_seed, progress=gr.Progress()):
-    image, seed = generate_image(prompt, negative_prompt, use_defaults, width, height, guidance_scale, num_inference_steps, seed, randomize_seed, progress)
+def interface_fn(prompt, negative_prompt, use_defaults, resolution, guidance_scale, num_inference_steps, seed, randomize_seed, progress=gr.Progress()):
+    image, seed = generate_image(prompt, negative_prompt, use_defaults, resolution, guidance_scale, num_inference_steps, seed, randomize_seed, progress)
     return image, seed, gr.update(value=seed)
 
 def reset_inputs():
-    return gr.update(value=''), gr.update(value=''), gr.update(value=True), gr.update(value=832), gr.update(value=1216), gr.update(value=7), gr.update(value=28), gr.update(value=0), gr.update(value=False)
+    return gr.update(value=''), gr.update(value=''), gr.update(value=True), gr.update(value='1024x1024'), gr.update(value=7), gr.update(value=28), gr.update(value=0), gr.update(value=False)
 
 with gr.Blocks(title="UrangDiffusion 1.0 Demo", theme="NoCrypt/miku@1.2.1") as demo:
     gr.HTML(
@@ -68,8 +69,14 @@ with gr.Blocks(title="UrangDiffusion 1.0 Demo", theme="NoCrypt/miku@1.2.1") as d
             prompt_input = gr.Textbox(lines=2, placeholder="Enter prompt here", label="Prompt")
             negative_prompt_input = gr.Textbox(lines=2, placeholder="Enter negative prompt here", label="Negative Prompt")
             use_defaults_input = gr.Checkbox(label="Use Default Quality Tags and Negative Prompt", value=True)
-            width_input = gr.Slider(minimum=256, maximum=2048, step=32, label="Width", value=832)
-            height_input = gr.Slider(minimum=256, maximum=2048, step=32, label="Height", value=1216)
+            resolution_input = gr.Radio(
+                choices=[
+                    "1024x1024", "1152x896", "896x1152", "1216x832", "832x1216",
+                    "1344x768", "768x1344", "1536x640", "640x1536"
+                ],
+                label="Resolution",
+                value="896x1152"
+            )
             guidance_scale_input = gr.Slider(minimum=1, maximum=20, step=0.5, label="Guidance Scale", value=7)
             num_inference_steps_input = gr.Slider(minimum=1, maximum=100, step=1, label="Number of Inference Steps", value=28)
             seed_input = gr.Slider(minimum=0, maximum=99999999, step=1, label="Seed", value=0, interactive=True)
@@ -83,7 +90,7 @@ with gr.Blocks(title="UrangDiffusion 1.0 Demo", theme="NoCrypt/miku@1.2.1") as d
     generate_button.click(
         interface_fn,
         inputs=[
-            prompt_input, negative_prompt_input, use_defaults_input, width_input, height_input, guidance_scale_input, num_inference_steps_input, seed_input, randomize_seed_input
+            prompt_input, negative_prompt_input, use_defaults_input, resolution_input, guidance_scale_input, num_inference_steps_input, seed_input, randomize_seed_input
         ],
         outputs=[output_image, seed_input]
     )
@@ -92,7 +99,7 @@ with gr.Blocks(title="UrangDiffusion 1.0 Demo", theme="NoCrypt/miku@1.2.1") as d
         reset_inputs,
         inputs=[],
         outputs=[
-            prompt_input, negative_prompt_input, use_defaults_input, width_input, height_input, guidance_scale_input, num_inference_steps_input, seed_input, randomize_seed_input
+            prompt_input, negative_prompt_input, use_defaults_input, resolution_input, guidance_scale_input, num_inference_steps_input, seed_input, randomize_seed_input
         ]
     )