import gradio as gr
import torch
from diffusers import FluxFillPipeline
from huggingface_hub import login
import os
import numpy as np

HF_TOKEN = os.environ["HF_TOKEN"]

# Authenticate with Hugging Face (if required)
login(token=HF_TOKEN)  # Use the actual token, not a string

# Load the inpainting model
pipe = FluxFillPipeline.from_pretrained(
    "black-forest-labs/FLUX.1-Fill-dev",
    torch_dtype=torch.float16
).to("cuda" if torch.cuda.is_available() else "cpu")

def inpaint_image(image, mask):
    # Convert PIL images to NumPy arrays
    image = np.array(image)
    mask = np.array(mask)

    result = pipe(image, mask).images[0]
    return result

# Gradio UI
with gr.Blocks() as demo:
    gr.Markdown("# FLUX.1 Fill - Inpainting Demo")
    with gr.Row():
        input_img = gr.Image(type="numpy", label="Upload Image")
        mask_img = gr.Image(type="numpy", label="Upload Mask (Black for removal)")
    output_img = gr.Image(label="Output Image")
    btn = gr.Button("Inpaint")
    btn.click(inpaint_image, inputs=[input_img, mask_img], outputs=output_img)

# Launch app
demo.launch()