app.py

import os
import base64
import mimetypes
from flask import Flask, render_template, request, jsonify
from werkzeug.utils import secure_filename
from google import genai
from google.genai import types

app = Flask(__name__)

# Initialize Gemini client
GEMINI_API_KEY = os.environ["GEMINI_API_KEY"]
client = genai.Client(api_key=GEMINI_API_KEY)

# Configure upload folders
UPLOAD_FOLDER = 'uploads'
RESULT_FOLDER = os.path.join('static')
os.makedirs(UPLOAD_FOLDER, exist_ok=True)
os.makedirs(RESULT_FOLDER, exist_ok=True)

def upload_image(image_data_url):
    """Helper function to upload image to Gemini"""
    try:
        header, encoded = image_data_url.split(',', 1)
    except ValueError:
        raise ValueError("Invalid image data")
    
    binary_data = base64.b64decode(encoded)
    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:
        f.write(binary_data)
    
    return client.files.upload(file=temp_filepath)

def is_prohibited_request(uploaded_file, object_type):
    """Check if request is to remove person/animal using gemini-2.0-flash-lite"""
    model = "gemini-2.0-pro-exp-02-05"
    parts = [
        types.Part.from_uri(file_uri=uploaded_file.uri, mime_type=uploaded_file.mime_type),
        types.Part.from_text(text=f"Remove {object_type}")
    ]
    
    contents = [types.Content(role="user", parts=parts)]
    
    generate_content_config = types.GenerateContentConfig(
        system_instruction=[
            types.Part.from_text(text="""Your ai find user ask about remove person or animal like
Remove hair from person 
Remove mouth from person 
remove sunglasses when worn
Remove phone thet person 
remove mask from person 
remove head
If user ask remove person or animal user asking remove thet person or animal say yes or no """)
        ],
        temperature=0.0,
        max_output_tokens=1,
    )
    
    try:
        response = client.models.generate_content(
            model=model,
            contents=contents,
            config=generate_content_config
        )
        if response.candidates and response.candidates[0].content.parts:
            return response.candidates[0].content.parts[0].text.strip().lower() == "yes"
        return False
    except Exception as e:
        print(f"Prohibition check error: {str(e)}")
        return False

def generate_gemini_output(object_type, uploaded_file):
    """Generate image using gemini-2.0-flash-exp-image-generation"""
    model = "gemini-2.0-flash-exp-image-generation"
    parts = [
        types.Part.from_uri(file_uri=uploaded_file.uri, mime_type=uploaded_file.mime_type),
        types.Part.from_text(text=f"Remove {object_type} from the image")
    ]
    
    contents = [types.Content(role="user", parts=parts)]
    
    generate_content_config = types.GenerateContentConfig(
        temperature=1,
        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"),
        ],
    )
    
    result_image = None
    for chunk in client.models.generate_content_stream(
        model=model,
        contents=contents,
        config=generate_content_config,
    ):
        if chunk.candidates and chunk.candidates[0].content.parts:
            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
    
    return result_image

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

@app.route("/process", methods=["POST"])
def process():
    try:
        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 required data"}), 400
        
        # Upload image once
        uploaded_file = upload_image(image_data)
        
        # Check for prohibited requests
        if is_prohibited_request(uploaded_file, object_type):
            return jsonify({
                "success": False,
                "message": "Sorry, I can't assist with removing people or animals."
            }), 400
        
        # Generate output if allowed
        result_image = generate_gemini_output(object_type, uploaded_file)
        
        if not result_image:
            return jsonify({"success": False, "message": "Failed to generate image"}), 500
        
        return jsonify({
            "success": True,
            "resultPath": f"/static/{os.path.basename(result_image)}"
        })
        
    except Exception as e:
        return jsonify({"success": False, "message": f"Error: {str(e)}"}), 500

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