import gradio as gr
import numpy as np
import spaces
import torch
import random
from PIL import Image

from diffusers import FluxKontextPipeline
from diffusers.utils import load_image

pipe = FluxKontextPipeline.from_pretrained("black-forest-labs/FLUX.1-Kontext-dev", torch_dtype=torch.bfloat16).to("cuda")

seed=42
guidance_scale=2.5
steps=28

@spaces.GPU
def infer(input_image, prompt, progress=gr.Progress(track_tqdm=True)):
    """
    Perform image editing using the FLUX.1 Kontext pipeline. This function takes an input image and a text 
    prompt to generate a modified version of the image based on the provided instructions. It uses the FLUX.1 
    Kontext model for contextual image editing tasks.
    
    Args:
        input_image (PIL.Image.Image): The input image to be edited. Will be converted
            to RGB format if not already in that format.
        prompt (str): Text description of the desired edit to apply to the image.
            Examples: "Remove glasses", "Add a hat", "Change background to beach".
    
    Returns:
        tuple: A 2-tuple containing:
            - tuple of (PIL.Image.Image, PIL.Image.Image): The original image, generated/edited image
            - int: The seed value used for generation (useful when randomize_seed=True)    
    """    
    if input_image:
        input_image = input_image.convert("RGB")
        image = pipe(
            image=input_image, 
            prompt=prompt,
            guidance_scale=guidance_scale,
            width = input_image.size[0],
            height = input_image.size[1],
            num_inference_steps=steps,
            generator=torch.Generator().manual_seed(seed),
        ).images[0]
    else:
        image = pipe(
            prompt=prompt,
            guidance_scale=guidance_scale,
            num_inference_steps=steps,
            generator=torch.Generator().manual_seed(seed),
        ).images[0]
    return input_image, image


demo = gr.Interface(
    fn=infer,
    inputs=[
        gr.Image(label="Upload the image for editing", type="pil"),
        gr.Text(
            label="Prompt",
            show_label=False,
            max_lines=1,
            placeholder="Enter your prompt for editing (e.g., 'Remove glasses', 'Add a hat')",
            container=False,
        )        
    ],
    outputs=[
        gr.ImageSlider(label="Result", show_label=False)
    ],
    examples=[
        ["flowers.png", "turn the flowers into sunflowers"],
        ["monster.png", "make this monster ride a skateboard on the beach"],
        ["cat.png", "make this cat happy"]
    ],
    cache_examples=True,
    cache_mode="lazy",
    title="FLUX.1 Kontext Image Editing"
)
    
demo.launch(mcp_server=True)