import gradio as gr
import torch
import os
import spaces
from diffusers import AutoPipelineForImage2Image
from PIL import Image

# Gated मॉडल के लिए Hugging Face Token (Space Settings -> Secrets में HF_TOKEN सेट करें)
hf_token = os.environ.get("HF_TOKEN")

# मॉडल लोड करना (Flux Image-to-Image)
model_id = "black-forest-labs/FLUX.1-dev"
pipe = AutoPipelineForImage2Image.from_pretrained(
    model_id,
    token=hf_token,
    torch_dtype=torch.bfloat16
)
pipe.to("cuda")

@spaces.GPU
def edit_image(input_image, prompt, strength):
    if input_image is None or not prompt.strip():
        raise gr.Error("कृपया एक फोटो अपलोड करें और प्रॉम्प्ट लिखें।")
    
    # इमेज को सही फॉर्मेट में लाना
    input_image = input_image.convert("RGB")
    
    # 48GB VRAM का पूरा फायदा उठाते हुए हाई-क्वालिटी इमेज जनरेट करना
    edited_image = pipe(
        prompt=prompt,
        image=input_image,
        strength=strength, # 0.0 से 1.0 के बीच
        guidance_scale=3.5,
        num_inference_steps=28
    ).images[0]
    
    return edited_image

# Gradio इंटरफेस
with gr.Blocks(theme='soft') as demo:
    gr.Markdown("# 📸 FLUX.1 Ultra-Realistic Image Editor")
    gr.Markdown("अपनी फोटो अपलोड करें और बताएं कि आप **फाइनल फोटो कैसी चाहते हैं**। यह एकदम असली (Realistic) रिजल्ट देगा।")
    
    with gr.Row():
        with gr.Column():
            input_image = gr.Image(type="pil", label="ओरिजिनल फोटो")
            prompt = gr.Textbox(
                label="प्रॉम्प्ट (जैसे: 'A photorealistic image of a man in a snowy forest')", 
                placeholder="यहाँ फाइनल रिजल्ट का वर्णन करें..."
            )
            strength = gr.Slider(
                label="Strength (बदलाव की तीव्रता)",
                minimum=0.1,
                maximum=1.0,
                step=0.05,
                value=0.6,
                info="कम वैल्यू = ओरिजिनल फोटो जैसी, ज्यादा वैल्यू = नए प्रॉम्प्ट जैसी।"
            )
            btn = gr.Button("जनरेट करें", variant="primary")
            
        with gr.Column():
            output_image = gr.Image(label="एडिट की गई फोटो (Ultra-Realistic)")

    btn.click(
        fn=edit_image,
        inputs=[input_image, prompt, strength],
        outputs=output_image
    )

if __name__ == "__main__":
    demo.launch()