add height/width

app.py

@@ -1,3 +1,4 @@
+import random
 import spaces
 import torch
 import gradio as gr
@@ -16,11 +17,14 @@ pipe = StableDiffusionDMMPipeline.from_pretrained(
 )
 pipe.to("cuda")
 
+
 @spaces.GPU
 def generate(prompt: str,
              negative_prompt: str,
              model_id: int,
              seed: int = 1234,
+             height: int = 512,
+             width: int = 512,
              all: bool = True):
     if all:
         outputs = []
@@ -28,8 +32,8 @@ def generate(prompt: str,
             output = pipe(
                 prompt=prompt,
                 negative_prompt=negative_prompt,
-                width=512,
-                height=512,
+                width=width,
+                height=height,
                 num_inference_steps=25,
                 guidance_scale=7,
                 model_id=i,
@@ -41,8 +45,8 @@ def generate(prompt: str,
         output = pipe(
             prompt=prompt,
             negative_prompt=negative_prompt,
-            width=512,
-            height=512,
+            width=width,
+            height=height,
             num_inference_steps=25,
             guidance_scale=7,
             model_id=int(model_id),
@@ -57,20 +61,27 @@ def main():
         gr.Markdown("# DMM")
         with gr.Row():
             with gr.Column():
+                with gr.Column():
+                    model_id = gr.Dropdown(list(range(8)), label="Model Index")
+                    all_check = gr.Checkbox(label="All (ignore the selection above)")
                 prompt = gr.Textbox("portrait photo of a girl, long golden hair, flowers, best quality", label="Prompt")
                 negative_prompt = gr.Textbox("worst quality,low quality,normal quality,lowres,watermark,nsfw", label="Negative Prompt")
-                seed = gr.Number(1234, label="Seed", precision=0)
-                with gr.Column():
-                    model_id = gr.Slider(label="Model Index", minimum=0, maximum=7, step=1)
-                    all_check = gr.Checkbox(label="All")
-                btn = gr.Button("Submit", variant="primary")
+                with gr.Row():
+                    seed = gr.Number(0, label="Seed", precision=0, scale=3)
+                    update_seed_btn = gr.Button("🎲", scale=1)
+                with gr.Row():
+                    height = gr.Number(512, step=8, label="Height (suggest 512~768)")
+                    width = gr.Number(512, step=8, label="Width")
+                submit_btn = gr.Button("Submit", variant="primary")
             output = gr.Gallery(label="images")
     
-        btn.click(generate,
-                  inputs=[prompt, negative_prompt, model_id, seed, all_check],
-                  outputs=[output])
+        submit_btn.click(generate,
+                         inputs=[prompt, negative_prompt, model_id, seed, height, width, all_check],
+                         outputs=[output])
+        update_seed_btn.click(lambda: random.randint(0, 1000000),
+                              outputs=[seed])
 
-    demo.launch()
+    demo.launch(server_name="0.0.0.0", server_port=7860)
 
 
 if __name__ == "__main__":