Update app.py

app.py

@@ -1,166 +1,83 @@
-import json
+from flask import Flask, render_template, request, jsonify
+import google.generativeai as genai
 import os
-import time
-import uuid
 import tempfile
-from PIL import Image
-import gradio as gr
 import base64
-import mimetypes
+from dotenv import load_dotenv
 
-from google import genai
-from google.genai import types
+# Load environment variables
+load_dotenv()
 
-def save_binary_file(file_name, data):
-    with open(file_name, "wb") as f:
-        f.write(data)
+# Configure Flask app
+app = Flask(__name__)
 
-def generate(text, file_name, api_key, model="gemini-2.0-flash-exp"):
-    # Initialize client using provided api_key (or fallback to env variable)
-    client = genai.Client(api_key=(api_key.strip() if api_key and api_key.strip() != ""
-                                     else os.environ.get("GEMINI_API_KEY")))
-    
-    files = [
-        client.files.upload(file=file_name),
-    ]
+# Configure Gemini API
+genai.configure(api_key=os.getenv("GEMINI_API_KEY"))
 
-    contents = [
-        types.Content(
-            role="user",
-            parts=[
-                types.Part.from_uri(
-                    file_uri=files[0].uri,
-                    mime_type=files[0].mime_type,
-                ),
-                types.Part.from_text(text=text),
-            ],
-        ),
-    ]
-    generate_content_config = types.GenerateContentConfig(
-        temperature=1,
-        top_p=0.95,
-        top_k=40,
-        max_output_tokens=8192,
-        response_modalities=[
-            "image",
-            "text",
-        ],
-        response_mime_type="text/plain",
-    )
-
-    with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmp:
-        temp_path = tmp.name
-        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
-            inline_data = chunk.candidates[0].content.parts[0].inline_data
-            if inline_data:
-                save_binary_file(temp_path, inline_data.data)
-                print(
-                    "File of mime type "
-                    f"{inline_data.mime_type} saved to: {temp_path} and prompt input :{text}"
-                )
-            else:
-                print(chunk.text)
-
-    del files
-    return temp_path
- 
-
-def process_image_and_prompt(composite_pil, prompt, gemini_api_key):
-    # Save the composite image to a temporary file.
-    with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmp:
-        composite_path = tmp.name
-        composite_pil.save(composite_path)
+@app.route("/")
+def home():
+    return render_template("index.html")
 
-    file_name = composite_path  
-    input_text = prompt 
-    model = "gemini-2.0-flash-exp" 
+@app.route("/process", methods=["POST"])
+def process_image():
+    try:
+        # Get data from request
+        data = request.json
+        image_data = data.get("image")
+        object_type = data.get("objectType")
 
-    gemma_edited_image_path = generate(text=input_text, file_name=file_name, api_key=gemini_api_key, model=model)
-    print("image_path ", gemma_edited_image_path)
-    result_img = Image.open(gemma_edited_image_path)
-    if result_img.mode == "RGBA":
-        result_img = result_img.convert("RGB")
-    return [result_img]
+        if not image_data or not object_type:
+            return jsonify({"success": False, "message": "Invalid input data"})
 
-# Build a Blocks-based interface to include the custom HTML header.
-with gr.Blocks() as demo:
-    # HTML Header for the application.
-    gr.HTML(
-    """
-    <div style='display: flex; align-items: center; justify-content: center; gap: 20px'>
-    <div style="background-color: var(--block-background-fill); border-radius: 8px">
-        <img src="https://www.gstatic.com/lamda/images/gemini_favicon_f069958c85030456e93de685481c559f160ea06b.png" style="width: 100px; height: 100px;">
-    </div>
-    <div>
-        <h1>Gen AI Image Editing</h1>
-        <p>Gemini using for Image Editing</p>
-        <p>Powered by <a href="https://gradio.app/">Gradio</a> ⚡️</p>
-        <p>Get an API Key <a href="https://aistudio.google.com/apikey">here</a></p>
-        <p>Follow me on Twitter: <a href="https://x.com/Ameerazam18">Ameerazam18</a></p>
-    </div>
-    </div>
-    """
-    )
+        # Decode base64 image data
+        image_bytes = base64.b64decode(image_data.split(",")[1])
 
-    # Title and description.
+        # Create temporary directory
+        temp_dir = tempfile.mkdtemp()
+        input_path = os.path.join(temp_dir, "input.png")
+        with open(input_path, "wb") as f:
+            f.write(image_bytes)
 
-    # Define examples to be shown within the Gradio interface
-    examples = [
-        # Each example is a list corresponding to the inputs:
-        # [Input Image, Prompt, Guidance Scale, Number of Steps, LoRA Name]
-        ["data/1.webp", 'change text to "AMEER"'],
-        ["data/2.webp", "remove the spoon from  hand only"],
-        ["data/3.webp", 'change text to "Make it "'],
-        ["data/1.jpg", "add  joker style only on face"],
-         ["data/1777043.jpg", "add  joker style only on face"],
-         ["data/2807615.jpg","add lipstick on lip only "],
-         
-         ["data/76860.jpg", "add lipstick on lip only "],
-         ["data/2807615.jpg", "make it happy looking face only"],
-        
-    
-    ]
+        # Create the model
+        model = genai.GenerativeModel('gemini-2.0-flash-exp-image-generation')
 
-    gr.Markdown("Upload an image and enter a prompt to generate outputs in the gallery. Do not Use NFSW Images")
+        # Build the prompt
+        prompt = f"Remove the {object_type} from the image and fill the area naturally."
 
-    with gr.Row():
-        with gr.Column():
-            image_input = gr.Image(
-                type="pil",
-                label="Upload Image",
-                image_mode="RGBA"
-            )
-            gemini_api_key = gr.Textbox(
-                lines=1,
-                placeholder="Enter Gemini API Key (optional)",
-                label="Gemini API Key (optional) Generate and fill here"
-            )
-            prompt_input = gr.Textbox(
-                lines=2,
-                placeholder="Enter prompt here...",
-                label="Prompt"
-            )
-            submit_btn = gr.Button("Generate")
-        with gr.Column():
-            output_gallery = gr.Gallery(label="Generated Outputs")
-
-    # Set up the interaction.
-    submit_btn.click(
-        fn=process_image_and_prompt,
-        inputs=[image_input, prompt_input, gemini_api_key],
-        outputs=output_gallery,
- 
-    )
-    gr.Examples(
-        examples=examples,
-        inputs=[image_input, prompt_input, gemini_api_key],
-        label="Try these examples"
-    )
-
-demo.launch(share=True)
+        # Generate content
+        response = model.generate_content(
+            [
+                prompt,
+                genai.upload_file(input_path)
+            ],
+            generation_config={
+                "temperature": 1,
+                "top_p": 0.95,
+                "top_k": 40,
+                "max_output_tokens": 8192,
+            },
+            safety_settings={
+                "HARM_CATEGORY_CIVIC_INTEGRITY": "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
+                        })
+
+        return jsonify({"success": False, "message": "No valid image data found in response"})
+
+    except Exception as e:
+        return jsonify({"success": False, "message": str(e)})
+
+if __name__ == "__main__":
+    app.run(host="0.0.0.0", port=5000)