Update app.py

app.py

@@ -2,8 +2,8 @@ from flask import Flask, render_template, request, jsonify
 import google.generativeai as genai
 import os
 import tempfile
+import base64
 from dotenv import load_dotenv
-from google.generativeai.types import Content, Part
 
 # Load environment variables
 load_dotenv()
@@ -14,12 +14,10 @@ app = Flask(__name__)
 # Configure Gemini API
 genai.configure(api_key=os.getenv("GEMINI_API_KEY"))
 
-# Route for home page
 @app.route("/")
 def home():
     return render_template("index.html")
 
-# Route for processing images
 @app.route("/process", methods=["POST"])
 def process_image():
     try:
@@ -32,33 +30,26 @@ def process_image():
             return jsonify({"success": False, "message": "Invalid input data"})
 
         # Decode base64 image data
-        import base64
         image_bytes = base64.b64decode(image_data.split(",")[1])
 
         # Create temporary directory
         temp_dir = tempfile.mkdtemp()
-        upload_path = os.path.join(temp_dir, "input_image.png")
-        with open(upload_path, "wb") as f:
+        input_path = os.path.join(temp_dir, "input.png")
+        with open(input_path, "wb") as f:
             f.write(image_bytes)
 
-        # Construct Gemini prompt
-        gemini_prompt = f"Remove the {object_type} from the image and fill the area naturally."
+        # Create the model
+        model = genai.GenerativeModel('gemini-2.0-flash-exp-image-generation')
 
-        # Create request content
-        contents = [
-            Content(
-                role="user",
-                parts=[
-                    Part.from_text(gemini_prompt),
-                    Part.from_data(mime_type="image/png", data=image_bytes)
-                ]
-            )
-        ]
+        # Build the prompt
+        prompt = f"Remove the {object_type} from the image and fill the area naturally."
 
-        # Generate content with Gemini
-        model = genai.GenerativeModel('gemini-2.0-flash-exp-image-generation')
+        # Generate content
         response = model.generate_content(
-            contents=contents,
+            [
+                prompt,
+                genai.upload_file(input_path)
+            ],
             generation_config={
                 "temperature": 1,
                 "top_p": 0.95,
@@ -71,16 +62,16 @@ def process_image():
         )
 
         # Process response
-        if response.candidates:
-            for candidate in response.candidates:
-                for part in candidate.content.parts:
+        output_path = os.path.join(temp_dir, "result.png")
+        for chunk in response:
+            if chunk.candidates:
+                for part in chunk.candidates[0].content.parts:
                     if hasattr(part, 'inline_data'):
-                        output_path = os.path.join(temp_dir, "result.png")
                         with open(output_path, "wb") as f:
                             f.write(part.inline_data.data)
                         return jsonify({
                             "success": True,
-                            "resultPath": f"/static/results/result.png"
+                            "resultPath": output_path
                         })
 
         return jsonify({"success": False, "message": "No valid image data found in response"})
@@ -88,6 +79,5 @@ def process_image():
     except Exception as e:
         return jsonify({"success": False, "message": str(e)})
 
-# Run the app
 if __name__ == "__main__":
-    app.run(debug=True)
+    app.run(host="0.0.0.0", port=7860)