Update app.py

app.py

@@ -9,8 +9,9 @@ from google.genai import types
 # Initialize Flask app
 app = Flask(__name__)
 
-# Set your Gemini API key (or set GEMINI_API_KEY in your environment)
-GEMINI_API_KEY = os.getenv("GEMINI_API_KEY")
+# Set your Gemini API key via Hugging Face Spaces environment variables.
+# Do not include a default fallback; the environment must supply GEMINI_API_KEY.
+GEMINI_API_KEY = os.environ["GEMINI_API_KEY"]
 client = genai.Client(api_key=GEMINI_API_KEY)
 
 # Create necessary directories
@@ -21,8 +22,8 @@ os.makedirs(RESULT_FOLDER, exist_ok=True)
 
 def generate_gemini_output(object_type, image_data_url):
     """
-    Generate output from Gemini by removing the object specified.
-    The image_data_url is expected to be a base64 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"
     files = []
@@ -35,10 +36,7 @@ def generate_gemini_output(object_type, image_data_url):
             raise ValueError("Invalid image data")
         binary_data = base64.b64decode(encoded)
         # Determine file extension from header
-        if "png" in header.lower():
-            ext = ".png"
-        else:
-            ext = ".jpg"
+        ext = ".png" if "png" in header.lower() else ".jpg"
         temp_filename = secure_filename("temp_image" + ext)
         temp_filepath = os.path.join(UPLOAD_FOLDER, temp_filename)
         with open(temp_filepath, "wb") as f:
@@ -52,7 +50,7 @@ def generate_gemini_output(object_type, image_data_url):
     if files:
         parts.append(types.Part.from_uri(file_uri=files[0].uri, mime_type=files[0].mime_type))
     if object_type:
-        # Create Gemini magic prompt for object removal
+        # Gemini magic prompt: instruct the model to remove the specified object
         magic_prompt = f"Remove {object_type} from the image"
         parts.append(types.Part.from_text(text=magic_prompt))
 
@@ -73,7 +71,7 @@ def generate_gemini_output(object_type, image_data_url):
     result_text = None
     result_image = None
 
-    # Stream the output from Gemini API
+    # Stream output from Gemini API
     for chunk in client.models.generate_content_stream(
         model=model,
         contents=contents,
@@ -96,26 +94,25 @@ def generate_gemini_output(object_type, image_data_url):
 
 @app.route("/")
 def index():
-    # Render the front-end HTML (which includes inline CSS and JS)
+    # Render the front-end HTML (which contains complete HTML/CSS/JS inline)
     return render_template("index.html")
 
 @app.route("/process", methods=["POST"])
 def process():
     try:
-        # Expecting JSON with keys "image" (base64 data URL) and "objectType"
+        # Expect JSON with keys "image" (base64 data URL) and "objectType"
         data = request.get_json(force=True)
         image_data = data.get("image")
         object_type = data.get("objectType", "").strip()
         if not image_data or not object_type:
             return jsonify({"success": False, "message": "Missing image data or object type."}), 400
 
-        # Generate Gemini output
+        # 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
 
         # Create a URL to serve the image from the static folder.
-        # Assuming your static folder is served at '/static'
         image_url = f"/static/{os.path.basename(result_image)}"
 
         return jsonify({"success": True, "resultPath": image_url, "resultText": result_text})
@@ -123,5 +120,5 @@ def process():
         return jsonify({"success": False, "message": f"Error: {str(e)}"}), 500
 
 if __name__ == "__main__":
-    # Run the app; debug can be set to False in production.
-    app.run(host="0.0.0.0", port=int(os.getenv("PORT", 5000)), debug=True)
+    # Run the app on port 5000 or the port provided by the environment (for Hugging Face Spaces)
+    app.run(host="0.0.0.0", port=int(os.environ.get("PORT", 5000)))