import os
import torch
import spaces
import gradio as gr
from diffusers import FluxFillPipeline
import random
import numpy as np
from huggingface_hub import hf_hub_download
from PIL import Image, ImageOps


CSS = """
h1 {
    margin-top: 10px
}
"""


os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "1"
MAX_SEED = np.iinfo(np.int32).max

repo_id = "black-forest-labs/FLUX.1-Fill-dev"

if torch.cuda.is_available():
    pipe = FluxFillPipeline.from_pretrained(repo_id, torch_dtype=torch.bfloat16).to("cuda")


def create_mask_image(mask_array):
    # Convert the mask to a numpy array if it's not already
    if not isinstance(mask_array, np.ndarray):
        mask_array = np.array(mask_array)
    
    # Create a new array with the same shape as the mask, but only for RGB channels
    processed_mask = np.zeros((mask_array.shape[0], mask_array.shape[1], 3), dtype=np.uint8)
    
    # Set transparent parts (alpha=0) to black (0, 0, 0)
    transparent_mask = mask_array[:, :, 3] == 0
    processed_mask[transparent_mask] = [0, 0, 0]
    
    # Set black parts (RGB=0, 0, 0 and alpha=255) to white (255, 255, 255)
    black_mask = (mask_array[:, :, :3] == [0, 0, 0]).all(axis=2) & (mask_array[:, :, 3] == 255)
    processed_mask[black_mask] = [255, 255, 255]
    
    return Image.fromarray(processed_mask)

@spaces.GPU()
def inpaintGen(
        imgMask,
        inpaint_prompt: str,
        guidance: float,
        num_steps: int,
        seed: int,
        randomize_seed: bool,
        progress=gr.Progress(track_tqdm=True)):

    source_path = imgMask["background"]
    mask_path = imgMask["layers"][0]

    print(f'source_path: {source_path}')
    print(f'mask_path: {mask_path}')  
    
    if not source_path:
        raise gr.Error("Please upload an image.")

    if not mask_path:
        raise gr.Error("Please draw a mask on the image.")

    source_img = Image.open(source_path).convert("RGB")
    mask_img = Image.open(mask_path)
    if mask_img.mode != 'L':
        mask_img = mask_img.convert('L')
    mask_img = ImageOps.invert(mask_img)

    #mask_img = create_mask_image(mask_img)
    
    width, height = source_img.size

    if randomize_seed:
        seed = random.randint(0, MAX_SEED)
    generator = torch.Generator("cpu").manual_seed(seed)

    result = pipe(
        prompt=inpaint_prompt,
        image=source_img,
        mask_image=mask_img,
        width=width,
        height=height,
        num_inference_steps=num_steps,
        generator=generator,
        guidance_scale=guidance,
        max_sequence_length=512,
    ).images[0]
    
    return result, seed
 

with gr.Blocks(theme="ocean", title="Flux.1 Fill dev", css=CSS) as demo:
    gr.HTML("<h1><center>Flux.1 Fill dev</center></h1>")
    gr.HTML("""
        <p>
            <center>
                A partial redraw of the image based on your prompt words and occluded parts.
            </center>
        </p>
    """)
    with gr.Row():
        with gr.Column():
            imgMask = gr.ImageMask(type="filepath", label="Image", layers=False, height=800)
            inpaint_prompt = gr.Textbox(label='Prompts ✏️', placeholder="A hat...")
            with gr.Row():
                Inpaint_sendBtn = gr.Button(value="Submit", variant='primary')
                Inpaint_clearBtn = gr.ClearButton([imgMask, inpaint_prompt], value="Clear")
        image_out = gr.Image(type="pil", label="Output", height=960)
    with gr.Accordion("Advanced ⚙️", open=False):
        guidance = gr.Slider(label="Guidance scale", minimum=1, maximum=50, value=30.0, step=0.1)
        num_steps = gr.Slider(label="Steps", minimum=1, maximum=50, value=20, step=1)
        seed = gr.Number(label="Seed", value=42, precision=0)
        randomize_seed = gr.Checkbox(label="Randomize seed", value=True)

    gr.on(
        triggers = [
            inpaint_prompt.submit,
            Inpaint_sendBtn.click,
        ],
        fn = inpaintGen,
        inputs = [
            imgMask,
            inpaint_prompt,
            guidance,
            num_steps,
            seed,
            randomize_seed
        ],
        outputs = [image_out, seed]
    )

if __name__ == "__main__":
    demo.queue(api_open=False).launch(show_api=False, share=False)