from typing import Tuple
import uuid
import random
import os
import numpy as np
import gradio as gr
import spaces
import torch
from PIL import Image
from diffusers import FluxInpaintPipeline
from gradio_client import Client, handle_file

# Set an environment variable
HF_TOKEN = os.environ.get("HF_TOKEN", None)

MARKDOWN = """
# FLUX.1 Inpainting with Text guided Mask🔥
Shoutout to [Black Forest Labs](https://huggingface.co/black-forest-labs) team for FLUX!
Special thanks to [Piotr Skalski](https://huggingface.co/SkalskiP) and [Gothos](https://github.com/Gothos) 
for their work on enabling and [showcasing inpainting](https://huggingface.co/spaces/SkalskiP/FLUX.1-inpaint) with the FLUX.
<br>We have used Gradio clients to access [EVF-SAM Spaces demo](https://huggingface.co/spaces/wondervictor/evf-sam) for text-guided segmentation. 
Visit [Gradio docs](https://www.gradio.app) to start building! 
"""

MAX_SEED = np.iinfo(np.int32).max
MAX_IMAGE_SIZE = 2048
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"

# Using Gradio Python Client to query EVF-SAM demo, hosted on SPaces, as an endpoint   
client = Client("ysharma/evf-sam", hf_token=HF_TOKEN)

pipe = FluxInpaintPipeline.from_pretrained(
    "black-forest-labs/FLUX.1-schnell", torch_dtype=torch.bfloat16).to(DEVICE)


def resize_image_dimensions(
    original_resolution_wh: Tuple[int, int],
    maximum_dimension: int = 2048
) -> Tuple[int, int]:
    width, height = original_resolution_wh

    if width <= maximum_dimension and height <= maximum_dimension:
        width = width - (width % 32)
        height = height - (height % 32)
        return width, height

    if width > height:
        scaling_factor = maximum_dimension / width
    else:
        scaling_factor = maximum_dimension / height

    new_width = int(width * scaling_factor)
    new_height = int(height * scaling_factor)

    new_width = new_width - (new_width % 32)
    new_height = new_height - (new_height % 32)

    return new_width, new_height


def evf_sam_mask(image, prompt):
    images = client.predict(
      image_np=handle_file(image),
      prompt=prompt, 
      api_name="/predict")
    # Open the mask image
    pil_image = Image.open(images[1])
    return pil_image

@spaces.GPU(duration=150)
def process(
    input_image: dict,
    input_text: str,
    inpaint_text: str,
    seed_slicer: int,
    randomize_seed_checkbox: bool,
    strength_slider: float,
    num_inference_steps_slider: int,
    progress=gr.Progress(track_tqdm=True)
):
    if not input_text:
        gr.Info("Please enter a text prompt.")
        return None

    mask = evf_sam_mask(input_image, input_text)
    
    if not input_image:
        gr.Info("Please upload an image.")
        return None
    else:
        input_image = Image.open(input_image)

    if not mask:
        gr.Info("Please draw a mask on the image.")
        return None

    width, height = resize_image_dimensions(original_resolution_wh=input_image.size)
    resized_image = input_image.resize((width, height), Image.LANCZOS)
    resized_mask = mask.resize((width, height), Image.NEAREST)

    if randomize_seed_checkbox:
        seed_slicer = random.randint(0, MAX_SEED)
    generator = torch.Generator().manual_seed(seed_slicer)
    result = pipe(
        prompt=inpaint_text,
        image=resized_image,
        mask_image=resized_mask,
        width=width,
        height=height,
        strength=strength_slider,
        generator=generator,
        num_inference_steps=num_inference_steps_slider
    ).images[0]
    print('INFERENCE DONE')
    return result, resized_mask


with gr.Blocks() as demo:
    gr.Markdown(MARKDOWN)
    with gr.Row():
        with gr.Column():
            input_image = gr.Image( 
                label='Image',
                type='filepath',
                sources=["upload", "webcam", "clipboard"],
                image_mode='RGB',
            )
            with gr.Row():
                with gr.Column():
                    input_text_component = gr.Text(
                        label="Text-guided segmentation",
                        show_label=True,
                        max_lines=1,
                        placeholder="Enter text for generating the segmentation mask",
                        container=False,
                    )
                    inpaint_text_component = gr.Text(
                        label="Text-guided Inpainting",
                        show_label=True,
                        max_lines=1,
                        placeholder="Enter text to generate Inpainting",
                        container=False,
                    )
                submit_button_component = gr.Button(value='Submit', variant='primary', scale=0)

            with gr.Accordion("Advanced Settings", open=False):
                seed_slicer_component = gr.Slider(
                    label="Seed",
                    minimum=0,
                    maximum=MAX_SEED,
                    step=1,
                    value=42,
                )

                randomize_seed_checkbox_component = gr.Checkbox(
                    label="Randomize seed", value=False)

                with gr.Row():
                    strength_slider_component = gr.Slider(
                        label="Strength",
                        minimum=0,
                        maximum=1,
                        step=0.01,
                        value=0.75,
                    )

                    num_inference_steps_slider_component = gr.Slider(
                        label="Number of inference steps",
                        minimum=1,
                        maximum=50,
                        step=1,
                        value=20,
                    )
        with gr.Column():
            output_image_component = gr.Image(
                type='pil', image_mode='RGB', label='Generated image')
            with gr.Accordion("Generated Mask", open=False):
                output_mask_component = gr.Image(
                    type='pil', image_mode='RGB', label='Input mask')

    submit_button_component.click(
        fn=process,
        inputs=[
            input_image, #input_image_editor_component,
            input_text_component,
            inpaint_text_component,
            seed_slicer_component,
            randomize_seed_checkbox_component,
            strength_slider_component,
            num_inference_steps_slider_component
        ],
        outputs=[
            output_image_component,
            output_mask_component,
        ]
    )

demo.launch(debug=True)