import os
import requests
import base64
import io
import gradio as gr
from PIL import Image

API_KEY = os.getenv("NVIDIA_API_KEY")
INVOKE_URL = "https://ai.api.nvidia.com/v1/genai/black-forest-labs/flux.2-klein-4b"

HEADERS = {
    "Authorization": f"Bearer {API_KEY}",
    "Accept": "application/json",
    "Content-Type": "application/json",
}

def generate_or_edit(prompt, input_image=None):
    payload = {
        "prompt": prompt,
        "width": 1024,
        "height": 1024,
        "seed": 0,
        "steps": 4
    }

    if input_image is not None:
        buffered = io.BytesIO()
        input_image.save(buffered, format="PNG")
        img_str = base64.b64encode(buffered.getvalue()).decode()
        payload["image"] = img_str

    try:
        response = requests.post(INVOKE_URL, headers=HEADERS, json=payload)
        response.raise_for_status()
        result = response.json()
        
        # अब हम सीधे उस स्ट्रक्चर का उपयोग कर रहे हैं जो आपने दिखाया है
        image_b64 = result["artifacts"][0]["base64"]
        image_data = base64.b64decode(image_b64)
            
        return Image.open(io.BytesIO(image_data))
        
    except Exception as e:
        print(f"Error: {e}")
        return None

# Gradio इंटरफ़ेस
with gr.Blocks() as demo:
    gr.Markdown("# Vedika AI - Image Studio")
    prompt_input = gr.Textbox(label="Prompt")
    image_input = gr.Image(label="Input Image (Optional)", type="pil")
    submit_btn = gr.Button("Generate")
    result_output = gr.Image(label="Result")

    submit_btn.click(fn=generate_or_edit, inputs=[prompt_input, image_input], outputs=result_output)

demo.launch(server_name="0.0.0.0", server_port=7860)