Update app.py

app.py

@@ -64,6 +64,7 @@ default_negative_prompt = "色调艳丽, 过曝, 静态, 细节模糊不清, 字
 
 def get_duration(
     prompt,
+    img_size,
     negative_prompt,
     guidance_scale,
     guidance_scale_2,
@@ -77,6 +78,7 @@ def get_duration(
 @spaces.GPU(duration=get_duration)
 def generate_image(
     prompt,
+    img_size,
     negative_prompt=default_negative_prompt,
     guidance_scale = 3.5,
     guidance_scale_2 = 4,
@@ -93,6 +95,7 @@ def generate_image(
     
     Args:
         prompt (str): Text prompt describing the desired image.
+        img_size (str): size and orientation of output image, choices are square(1024), portrait(1024,576), landscape(576,1024).
         negative_prompt (str, optional): Negative prompt to avoid unwanted elements. 
             Defaults to default_negative_prompt (contains unwanted visual artifacts).
         guidance_scale (float, optional): Controls adherence to the prompt. Higher values = more adherence.
@@ -123,11 +126,21 @@ def generate_image(
    
     current_seed = random.randint(0, MAX_SEED) if randomize_seed else int(seed)
 
+    if img_size == "square":
+        height=1024
+        width=1024
+    elif img_size == "portrait":
+        height=1024
+        width=576
+    else: #landscape
+        height=576
+        width=1024
+    
     out_img = pipe(
         prompt=prompt,
         negative_prompt=negative_prompt,
-        height=1024,
-        width=1024,
+        height=height,
+        width=width,
         num_frames=1,
         guidance_scale=float(guidance_scale),
         guidance_scale_2=float(guidance_scale_2),
@@ -154,6 +167,7 @@ with gr.Blocks(css=css) as demo:
                 container=False,)
           generate_button = gr.Button("Run", variant="primary", scale=0)
        img_output = gr.Image(label="Generated Image", interactive=False)
+       img_size= gr.Radio(choices=["square", "portrait", "landscape"], value="square", label="Image Size")
        with gr.Accordion("Advanced Settings", open=False):
             negative_prompt_input = gr.Textbox(label="Negative Prompt", value=default_negative_prompt, lines=3)
             seed_input = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=42, interactive=True)
@@ -176,7 +190,7 @@ with gr.Blocks(css=css) as demo:
         ],
         inputs=[prompt_input], outputs=[img_output, seed_input], fn=generate_image, cache_examples="lazy")            
     ui_inputs = [ 
-        prompt_input,
+        prompt_input,img_size,
         negative_prompt_input,
         guidance_scale_input, guidance_scale_2_input, steps_slider, seed_input, randomize_seed_checkbox
     ]