Update app.py

app.py

@@ -73,7 +73,8 @@ else:
         INITIALIZATION_SUCCESS = False
 
 # --- 2. The Core Image Generation Function ---
-def generate_image(prompt: str, negative_prompt: str, seed: int):
+# UPDATE: Added 'aspect_ratio' parameter to the function signature
+def generate_image(prompt: str, negative_prompt: str, seed: int, aspect_ratio: str):
     """Generates an image using the Vertex AI Imagen model."""
     if not INITIALIZATION_SUCCESS:
         # Provide the clear, generic error message in the UI if initialization failed.
@@ -81,14 +82,19 @@ def generate_image(prompt: str, negative_prompt: str, seed: int):
 
     # <<< START: PRIVACY UPDATE >>>
     # Do not log the user's prompt content.
-    print("Received new image generation request.")
+    print(f"Received new image generation request with aspect ratio: {aspect_ratio}.")
     # <<< END: PRIVACY UPDATE >>>
     
     try:
         images = generation_model.generate_images(
-            prompt=prompt, number_of_images=1, aspect_ratio="1:1",
-            negative_prompt=negative_prompt, add_watermark=False,
-            safety_filter_level="block_few", person_generation="allow_adult",
+            prompt=prompt,
+            number_of_images=1,
+            # UPDATE: Use the aspect_ratio from the UI instead of a hardcoded value
+            aspect_ratio=aspect_ratio,
+            negative_prompt=negative_prompt,
+            add_watermark=False,
+            safety_filter_level="block_few",
+            person_generation="allow_adult",
             seed=seed
         )
         print("Image generation successful for the request.")
@@ -102,6 +108,9 @@ def generate_image(prompt: str, negative_prompt: str, seed: int):
         # <<< END: SECURITY UPDATE >>>
 
 # --- 3. Gradio User Interface Definition ---
+# UPDATE: Added a list of all supported aspect ratios for Imagen 4
+IMAGEN4_ASPECT_RATIOS = ["1:1", "16:9", "9:16", "4:3", "3:4"]
+
 with gr.Blocks(theme=gr.themes.Soft(), css=".gradio-container {max-width: 960px !important}") as demo:
     gr.Markdown("# 🎨 Vertex AI Imagen 4.0 Image Generator")
     gr.Markdown("Generate high-quality images with Google's latest Imagen 4.0 model.")
@@ -110,6 +119,15 @@ with gr.Blocks(theme=gr.themes.Soft(), css=".gradio-container {max-width: 960px
         with gr.Column(scale=2):
             prompt_input = gr.Textbox(label="Prompt", placeholder="A photorealistic image of a futuristic city...", lines=3)
             negative_prompt_input = gr.Textbox(label="Negative Prompt", placeholder="text, watermark, blurry...", lines=2)
+            
+            # UPDATE: Added a radio button selector for the aspect ratio
+            aspect_ratio_input = gr.Radio(
+                choices=IMAGEN4_ASPECT_RATIOS,
+                label="Aspect Ratio",
+                value="1:1",
+                info="Select the output image dimensions (Square, Landscape, Portrait)."
+            )
+            
             seed_input = gr.Slider(label="Seed", minimum=0, maximum=99999, step=1, randomize=True, info="Controls randomness.")
             submit_button = gr.Button("Generate Image", variant="primary")
         with gr.Column(scale=1):
@@ -125,7 +143,8 @@ with gr.Blocks(theme=gr.themes.Soft(), css=".gradio-container {max-width: 960px
 
     submit_button.click(
         fn=generate_image,
-        inputs=[prompt_input, negative_prompt_input, seed_input],
+        # UPDATE: Added aspect_ratio_input to the list of inputs
+        inputs=[prompt_input, negative_prompt_input, seed_input, aspect_ratio_input],
         outputs=image_output
     )