app.py

import json
import os
import time
import uuid
import tempfile
from PIL import Image
import gradio as gr
import base64
import mimetypes

from google import genai
from google.genai import types

# Aapki Gemini API Key yahan daal di gayi hai.
GEMINI_API_KEY = "AIzaSyCrhWiAEQmCidtE2QZw3CTiLt7F8yv5M7A"

## <<< CHANGE START: API Client ko globally initialize karein taaki baar-baar na bane.
# API key configure karein
genai.configure(api_key=GEMINI_API_KEY)
# File service ke liye client banayein
file_client = genai.FilesClient(api_key=GEMINI_API_KEY)
## <<< CHANGE END

def save_binary_file(file_name, data):
    with open(file_name, "wb") as f:
        f.write(data)

def generate(text, file_name, model="gemini-1.5-flash-latest"): # <<< CHANGE: Model ko stable flash version se update kiya gaya
    
    # Client ko global scope se istemal karein
    print(f"Uploading file: {file_name}")
    uploaded_file = file_client.upload_file(path=file_name)
    print(f"Completed uploading file: {uploaded_file.name}")

    # Model ko initialize karein
    model_instance = genai.GenerativeModel(model_name=model)

    contents = [
        types.Content(
            role="user",
            parts=[
                types.Part.from_uri(
                    file_uri=uploaded_file.uri,
                    mime_type=uploaded_file.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
        print("Generating content...")
        # Stream ki jagah seedha generate karein, image generation ke liye yeh aam taur par behtar hai
        response = model_instance.generate_content(
            contents=contents,
            generation_config=generate_content_config,
        )
        print("Content generated.")
        
        # Response se image data extract karein
        if response.candidates and response.candidates[0].content and response.candidates[0].content.parts:
            inline_data = response.candidates[0].content.parts[0].inline_data
            if inline_data:
                save_binary_file(temp_path, inline_data.data)
                print(
                    f"File of mime type {inline_data.mime_type} saved to: {temp_path}"
                )
            else:
                 # Agar image nahi hai, to text response print karein
                if response.text:
                    print(f"Model Response (Text): {response.text}")
                else:
                    print("No image data or text found in the response.")
        else:
            print("Could not get a valid response from the model.")


    # Upload ki gayi file ko delete karein taaki storage na bhare
    print(f"Deleting uploaded file: {uploaded_file.name}")
    file_client.delete_file(name=uploaded_file.name)
    
    return temp_path
 

def process_image_and_prompt(composite_pil, prompt):
    ## <<< CHANGE START: Image ko API par bhejne se pehle resize karein
    MAX_SIZE = (1024, 1024)  # Maximum dimensions
    composite_pil.thumbnail(MAX_SIZE, Image.Resampling.LANCZOS)
    print(f"Image resized to fit within {MAX_SIZE} dimensions.")
    ## <<< CHANGE END

    # Resized image ko temporary file mein save karein
    with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmp:
        composite_path = tmp.name
        # Agar image mein transparency (alpha channel) hai to use RGB mein convert karein
        if composite_pil.mode == "RGBA":
            composite_pil = composite_pil.convert("RGB")
        composite_pil.save(composite_path, "PNG")

    file_name = composite_path  
    input_text = prompt 
    
    gemma_edited_image_path = generate(text=input_text, file_name=file_name)
    
    if os.path.exists(gemma_edited_image_path) and os.path.getsize(gemma_edited_image_path) > 0:
        print("image_path ", gemma_edited_image_path)
        result_img = Image.open(gemma_edited_image_path)
        # Safai ke liye temporary files delete karein
        os.remove(composite_path)
        os.remove(gemma_edited_image_path)
        return [result_img]
    else:
        # Agar koi image generate nahi hui to error handle karein
        print("Failed to generate image or the generated file is empty.")
        os.remove(composite_path) # Input temp file ko phir bhi delete karein
        # Aap yahan user ko ek error message dikha sakte hain
        # For now, we return nothing, which will clear the gallery
        return []


# Build a Blocks-based interface.
with gr.Blocks() as demo:
    gr.Markdown("## Gen AI Image Editing\nUpload an image and enter a prompt to generate outputs in the gallery. Do not Use NFSW Images")
    
    # Define examples to be shown within the Gradio interface
    examples = [
        ["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"],
    ]

    with gr.Row():
        with gr.Column():
            image_input = gr.Image(
                type="pil",
                label="Upload Image",
                image_mode="RGBA"
            )
            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],
        outputs=output_gallery,
        show_progress="full"  # <<< CHANGE: UI mein loading progress dikhayein
    )
    gr.Examples(
        examples=examples,
        inputs=[image_input, prompt_input],
        label="Try these examples"
    )

demo.launch(share=True)