Update app.py

app.py

@@ -9,8 +9,7 @@ from google.genai import types
 # Initialize Flask app
 app = Flask(__name__)
 
-# Set your Gemini API key via Hugging Face Spaces environment variables.
-# Do not include a default fallback; the environment must supply GEMINI_API_KEY.
+# Read the Gemini API key from environment variables (set in Hugging Face Spaces)
 GEMINI_API_KEY = os.environ["GEMINI_API_KEY"]
 client = genai.Client(api_key=GEMINI_API_KEY)
 
@@ -25,7 +24,7 @@ def generate_gemini_output(object_type, image_data_url):
     Generate output from Gemini by removing the specified object.
     Expects the image_data_url to be a base64 data URL.
     """
-    model = "gemini-2.0-flash-exp-image-generation"
+    model = "gemini-2.0-flash-lite"  # Use the lite model for text-based responses
     files = []
 
     # Decode the image data from the data URL
@@ -61,15 +60,14 @@ def generate_gemini_output(object_type, image_data_url):
         top_p=0.95,
         top_k=40,
         max_output_tokens=8192,
-        response_modalities=["image", "text"],
-        safety_settings=[
-            types.SafetySetting(category="HARM_CATEGORY_CIVIC_INTEGRITY", threshold="OFF"),
-        ],
         response_mime_type="text/plain",
+        system_instruction=[
+            types.Part.from_text(text="""Your AI finds user requests about removing objects from images.
+If the user asks to remove a person or animal, respond with 'No'."""),
+        ],
     )
 
     result_text = None
-    result_image = None
 
     # Stream output from Gemini API
     for chunk in client.models.generate_content_stream(
@@ -80,17 +78,31 @@ def generate_gemini_output(object_type, image_data_url):
         if not chunk.candidates or not chunk.candidates[0].content or not chunk.candidates[0].content.parts:
             continue
         part = chunk.candidates[0].content.parts[0]
-        if part.inline_data:
-            file_extension = mimetypes.guess_extension(part.inline_data.mime_type) or ".png"
-            output_filename = secure_filename("generated_output" + file_extension)
-            result_image_path = os.path.join(RESULT_FOLDER, output_filename)
-            with open(result_image_path, "wb") as f:
-                f.write(part.inline_data.data)
-            result_image = result_image_path  # Path relative to static folder
-        else:
+        if part.text:
             result_text = part.text
 
-    return result_text, result_image
+    # If the response is "No", switch to the image generation model
+    if result_text and "no" in result_text.lower():
+        model = "gemini-2.0-flash-exp-image-generation"
+        generate_content_config.response_modalities = ["image", "text"]
+        for chunk in client.models.generate_content_stream(
+            model=model,
+            contents=contents,
+            config=generate_content_config,
+        ):
+            if not chunk.candidates or not chunk.candidates[0].content or not chunk.candidates[0].content.parts:
+                continue
+            part = chunk.candidates[0].content.parts[0]
+            if part.inline_data:
+                file_extension = mimetypes.guess_extension(part.inline_data.mime_type) or ".png"
+                output_filename = secure_filename("generated_output" + file_extension)
+                result_image_path = os.path.join(RESULT_FOLDER, output_filename)
+                with open(result_image_path, "wb") as f:
+                    f.write(part.inline_data.data)
+                result_image = result_image_path  # Path relative to static folder
+                return result_text, result_image
+
+    return result_text, None
 
 @app.route("/")
 def index():
@@ -110,7 +122,7 @@ def process():
         # Generate output using Gemini
         result_text, result_image = generate_gemini_output(object_type, image_data)
         if not result_image:
-            return jsonify({"success": False, "message": "Failed to generate image."}), 500
+            return jsonify({"success": False, "message": result_text or "Failed to generate image."}), 500
 
         # Create a URL to serve the image from the static folder.
         image_url = f"/static/{os.path.basename(result_image)}"