app.py

from flask import Flask, render_template, request, jsonify
import google.generativeai as genai
import os
import tempfile
import base64
import logging
from dotenv import load_dotenv

# Configure logging
logging.basicConfig(level=logging.INFO)

# Load environment variables
load_dotenv()

app = Flask(__name__)
genai.configure(api_key=os.getenv("GEMINI_API_KEY"))

@app.route("/")
def home():
    return render_template("index.html")

@app.route("/process", methods=["POST"])
def process_image():
    try:
        data = request.json
        image_data = data.get("image")
        object_type = data.get("objectType", "").strip()

        # Validate inputs
        if not image_data or not object_type:
            return jsonify({"success": False, "message": "Missing required parameters"})

        # Decode image
        try:
            header, encoded = image_data.split(",", 1)
            image_bytes = base64.b64decode(encoded)
        except Exception as e:
            logging.error(f"Image decoding failed: {str(e)}")
            return jsonify({"success": False, "message": "Invalid image data"})

        # Temporary files
        with tempfile.TemporaryDirectory() as temp_dir:
            input_path = os.path.join(temp_dir, "input.png")
            with open(input_path, "wb") as f:
                f.write(image_bytes)

            # Configure model with safety settings
            model = genai.GenerativeModel('gemini-2.0-flash-exp-image-generation')
            prompt = f"Remove {object_type} from image naturally without text or artifacts"
            
            response = model.generate_content(
                [prompt, genai.upload_file(input_path)],
                generation_config={
                    "temperature": 0.9,
                    "top_p": 0.95,
                    "top_k": 32,
                    "max_output_tokens": 4096,
                },
                safety_settings={
                    "HARM_CATEGORY_CIVIC_INTEGRITY": "BLOCK_ONLY_HIGH",
                    "HARM_CATEGORY_HARASSMENT": "BLOCK_NONE",
                    "HARM_CATEGORY_HATE_SPEECH": "BLOCK_NONE",
                    "HARM_CATEGORY_SEXUALLY_EXPLICIT": "BLOCK_NONE",
                    "HARM_CATEGORY_DANGEROUS_CONTENT": "BLOCK_NONE"
                }
            )

            # Process response
            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'):
                            with open(output_path, "wb") as f:
                                f.write(part.inline_data.data)
                            return jsonify({
                                "success": True,
                                "resultPath": output_path
                            })
                        elif hasattr(part, 'text'):
                            logging.info(f"Text response: {part.text}")

            return jsonify({"success": False, "message": "No valid output generated"})

    except Exception as e:
        logging.error(f"Processing error: {str(e)}")
        return jsonify({"success": False, "message": f"Processing error: {str(e)}"})

if __name__ == "__main__":
    app.run(host="0.0.0.0", port=7860)