Update app.py

app.py

@@ -26,24 +26,25 @@ MAX_SEED = np.iinfo(np.int32).max
 examples = json.loads(open("examples.json").read())
 
 aspect_ratios = {
-    "Big 1:1": (1080, 1080),
-    "Big 16:9": (1920, 1080),
-    "Big 9:16": (1080, 1920),
-    "Big 4:3": (1440, 1080),
-    "Big 3:4": (1080, 1440),
-
-    "Small 1:1": (720, 720),
-    "Small 16:9": (1280, 720),
-    "Small 9:16": (720, 1280),
-    "Small 4:3": (960, 720),
-    "Small 3:4": (720, 960),
+    "1080, aspect 1:1": (1080, 1080),
+    "1080, aspect 16:9": (1920, 1080),
+    "1080, aspect 9:16": (1080, 1920),
+    "1080, aspect 4:3": (1440, 1080),
+    "1080, aspect 3:4": (1080, 1440),
+
+    "720, aspect 1:1": (720, 720),
+    "720, aspect 16:9": (1280, 720),
+    "720, aspect 9:16": (720, 1280),
+    "720, aspect 4:3": (960, 720),
+    "720, aspect 3:4": (720, 960),
+
+    "480, aspect 1:1": (480, 480),
+    "480, aspect 16:9": (854, 480),
+    "480, aspect 9:16": (480, 854),
+    "480, aspect 4:3": (640, 480),
+    "480, aspect 3:4": (480, 640),
 }
 
-def get_image_size(aspect_ratio):
-
-
-
-    return aspect_ratios[aspect_ratio]
 
 def polish_prompt_en(original_prompt):
     
@@ -66,13 +67,13 @@ def polish_prompt_en(original_prompt):
     return polished_prompt
 
 
-@spaces.GPU(duration=45)
+@spaces.GPU(duration=90)
 def infer(
     prompt,
     negative_prompt=" ",
     seed=42,
     randomize_seed=False,
-    aspect_ratio="16:9",
+    aspect_ratio="480, aspect 3:4",
     guidance_scale=4,
     num_inference_steps=50,
     progress=gr.Progress(track_tqdm=True),
@@ -81,7 +82,8 @@ def infer(
 
     if randomize_seed:
         seed = random.randint(0, MAX_SEED)
-    width, height = get_image_size(aspect_ratio)
+        
+    width, height = aspect_ratios[aspect_ratio]
     
     print("Generating for prompt:", prompt)
     image = pipe(
@@ -89,9 +91,9 @@ def infer(
         negative_prompt=negative_prompt,
         width=width,
         height=height,
-        num_inference_steps=50,
-        true_cfg_scale=4.0,
-        generator=torch.Generator(device="cuda").manual_seed(42)
+        num_inference_steps=num_inference_steps,
+        true_cfg_scale=guidance_scale,
+        generator=torch.Generator(device="cuda").manual_seed(seed)
     ).images[0]
 
     #image.save("example.png")
@@ -125,8 +127,9 @@ with gr.Blocks(css=css) as demo:
             run_button = gr.Button("Run", scale=0, variant="primary")
 
         result = gr.Image(label="Result", show_label=False)
+        seed_output = gr.Textbox(label="Used seed", lines=1)
 
-        with gr.Accordion("Advanced Settings", open=False):
+        with gr.Accordion("Advanced Settings", open=True):
             negative_prompt = gr.Text(
                 label="Negative prompt",
                 max_lines=1,
@@ -145,14 +148,13 @@ with gr.Blocks(css=css) as demo:
             randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
 
             with gr.Row():
-                aspect_ratio = gr.Radio(
-                    label="Image size",
+                aspect_ratio = gr.Dropdown(
+                    label="Image size (aprox.)",
                     choices=list(aspect_ratios.keys()),
-                    value="Small 4:3",
+                    value="480, aspect 16:9",
                 )
 
 
-
             with gr.Row():
                 guidance_scale = gr.Slider(
                     label="Guidance scale",
@@ -183,7 +185,7 @@ with gr.Blocks(css=css) as demo:
             guidance_scale,
             num_inference_steps,
         ],
-        outputs=[result, seed],
+        outputs=[result, seed_output],
     )
 
 if __name__ == "__main__":