Update app.py

app.py

@@ -1,6 +1,7 @@
 import gradio as gr
 from gradio_client import Client, handle_file
 from google import genai
+from google.genai import types
 import os
 from typing import Optional, List, Tuple, Union
 from huggingface_hub import whoami
@@ -126,7 +127,7 @@ GOOGLE_API_KEY = os.getenv("GOOGLE_API_KEY", "")
 if not GOOGLE_API_KEY:
     raise ValueError("GOOGLE_API_KEY environment variable not set.")
 client = genai.Client(api_key=os.environ.get("GOOGLE_API_KEY"))
-GEMINI_MODEL_NAME = 'gemini-2.5-flash-image-preview'
+GEMINI_MODEL_NAME = 'gemini-2.5-flash-image'
 
 def verify_pro_status(token: Optional[Union[gr.OAuthToken, str]]) -> bool:
     """Verifies if the user is a Hugging Face PRO user or part of an enterprise org."""
@@ -250,7 +251,7 @@ def _generate_video_segment(input_image_path: str, output_image_path: str, promp
     )
     return result[0]["video"]
 
-def unified_image_generator(prompt: str, images: Optional[List[str]], previous_video_path: Optional[str], last_frame_path: Optional[str], manual_token: str, oauth_token: Optional[gr.OAuthToken]) -> tuple:
+def unified_image_generator(prompt: str, images: Optional[List[str]], previous_video_path: Optional[str], last_frame_path: Optional[str], aspect_ratio: str, manual_token: str, oauth_token: Optional[gr.OAuthToken]) -> tuple:
     if not (verify_pro_status(oauth_token) or verify_pro_status(manual_token)): 
         raise gr.Error("Access Denied.")
     
@@ -265,7 +266,25 @@ def unified_image_generator(prompt: str, images: Optional[List[str]], previous_v
     try:
         contents = [Image.open(image_path[0]) for image_path in images] if images else []
         contents.append(prompt)
-        response = client.models.generate_content(model=GEMINI_MODEL_NAME, contents=contents)
+        
+        # Create config with aspect ratio (omit image_config if Auto is selected)
+        if aspect_ratio == "Auto":
+            generate_content_config = types.GenerateContentConfig(
+                response_modalities=["IMAGE", "TEXT"],
+            )
+        else:
+            generate_content_config = types.GenerateContentConfig(
+                response_modalities=["IMAGE", "TEXT"],
+                image_config=types.ImageConfig(
+                    aspect_ratio=aspect_ratio,
+                ),
+            )
+        
+        response = client.models.generate_content(
+            model=GEMINI_MODEL_NAME, 
+            contents=contents,
+            config=generate_content_config
+        )
         image_data = _extract_image_data_from_response(response)
         if not image_data: raise gr.Error("No image data in response")
         with tempfile.NamedTemporaryFile(delete=False, suffix=".png") as tmp:
@@ -339,6 +358,12 @@ with gr.Blocks(theme=gr.themes.Citrus(), css=css) as demo:
             with gr.Column(scale=1):
                 image_input_gallery = gr.Gallery(label="Upload one or more images here. Leave empty for text-to-image", file_types=["image"], height="auto")
                 prompt_input = gr.Textbox(label="Prompt", placeholder="Turns this photo into a masterpiece")
+                aspect_ratio_dropdown = gr.Dropdown(
+                    label="Aspect Ratio",
+                    choices=["Auto", "1:1", "9:16", "16:9", "3:4", "4:3", "3:2", "2:3", "5:4", "4:5", "21:9"],
+                    value="Auto",
+                    interactive=True
+                )
                 generate_button = gr.Button("Generate", variant="primary")
             with gr.Column(scale=1):
                 output_image = gr.Image(label="Output", interactive=False, elem_id="output", type="filepath")
@@ -357,7 +382,7 @@ with gr.Blocks(theme=gr.themes.Citrus(), css=css) as demo:
     gr.on(
         triggers=[generate_button.click, prompt_input.submit],
         fn=unified_image_generator,
-        inputs=[prompt_input, image_input_gallery, previous_video_state, last_frame_of_video_state, manual_token],
+        inputs=[prompt_input, image_input_gallery, previous_video_state, last_frame_of_video_state, aspect_ratio_dropdown, manual_token],
         outputs=[output_image, create_video_button, extend_video_button, video_group],
         api_name=False
     )