Update app.py

app.py

@@ -1,8 +1,9 @@
 
 import gradio as gr
 import os
-import google.generativeai as genai
-from google.generativeai.types import HarmCategory, HarmBlockThreshold
+# Updated import to use the new 'google.genai' SDK
+from google import genai
+from google.genai import types
 import json
 from PIL import Image
 from io import BytesIO
@@ -13,12 +14,21 @@ import re
 
 # --- Configuration and Initialization ---
 
-# Configure the Gemini API key from Hugging Face secrets
+# Configure the Gemini API client. It will automatically pick up the
+# GEMINI_API_KEY from the environment variables.
+# Make sure to set this secret in your Hugging Face Space.
 try:
-    API_KEY = os.environ["API_KEY"]
-    genai.configure(api_key=API_KEY)
-except KeyError:
-    raise ValueError("API_KEY environment variable not set. Please add it to your Hugging Face Space secrets.")
+    client = genai.Client()
+    # A quick check to see if the client is configured.
+    # This will throw an exception if the key is not found.
+    client.get_api_key()
+except Exception as e:
+    raise ValueError(
+        "Could not configure Gemini API client. "
+        "Please make sure the GEMINI_API_KEY secret is set in your Hugging Face Space. "
+        f"Original error: {e}"
+    )
+
 
 # --- Default Data Structures (Ported from TypeScript) ---
 
@@ -86,12 +96,13 @@ def generate_image(prompt: str, aspect_ratio: str):
     """Generates an image using Gemini and returns a base64 data URL."""
     print(f"Generating image for prompt: {prompt}")
     try:
-        model = genai.GenerativeModel("gemini-2.0-flash-preview-image-generation")
+        # Refactored to use the new client.models.generate_content method
         full_prompt = f'Generate a single, high-quality image. Description: "{prompt}". The image should have a {aspect_ratio} aspect ratio.'
         
-        response = model.generate_content(
+        response = client.models.generate_content(
+            model="gemini-2.0-flash-preview-image-generation",
             contents=full_prompt,
-            generation_config=genai.types.GenerateContentConfig(response_modalities=['TEXT', 'IMAGE'])
+            config=types.GenerateContentConfig(response_modalities=['TEXT', 'IMAGE'])
         )
         
         image_part = next((part for part in response.candidates[0].content.parts if part.inline_data), None)
@@ -295,7 +306,7 @@ def create_gradio_app():
         
         # --- Step 1: Generate Layout ---
         try:
-            model = genai.GenerativeModel('gemini-2.5-flash')
+            # Refactored to use the new client.models.generate_content method
             prompt = create_prompt_from_answers(answers)
             schema = {
                 "type": "object",
@@ -319,17 +330,19 @@ def create_gradio_app():
                 }
             }
             
-            response = model.generate_content(
-                prompt,
-                generation_config=genai.GenerationConfig(
+            response = client.models.generate_content(
+                model='gemini-2.5-flash',
+                contents=prompt,
+                config=types.GenerateContentConfig(
                     response_mime_type="application/json",
                     response_schema=schema
                 ),
+                # Updated to use types.HarmCategory and types.HarmBlockThreshold
                 safety_settings={
-                    HarmCategory.HARM_CATEGORY_HARASSMENT: HarmBlockThreshold.BLOCK_NONE,
-                    HarmCategory.HARM_CATEGORY_HATE_SPEECH: HarmBlockThreshold.BLOCK_NONE,
-                    HarmCategory.HARM_CATEGORY_SEXUALLY_EXPLICIT: HarmBlockThreshold.BLOCK_NONE,
-                    HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT: HarmBlockThreshold.BLOCK_NONE,
+                    types.HarmCategory.HARM_CATEGORY_HARASSMENT: types.HarmBlockThreshold.BLOCK_NONE,
+                    types.HarmCategory.HARM_CATEGORY_HATE_SPEECH: types.HarmBlockThreshold.BLOCK_NONE,
+                    types.HarmCategory.HARM_CATEGORY_SEXUALLY_EXPLICIT: types.HarmBlockThreshold.BLOCK_NONE,
+                    types.HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT: types.HarmBlockThreshold.BLOCK_NONE,
                 }
             )
             design_plan = json.loads(response.text)