import torch
import gradio as gr
from diffusers.pipelines.flux.pipeline_flux import FluxPipeline
from diffusers.models.controlnet_flux import FluxControlNetModel
from controlnet_aux import CannyDetector

base_model = 'black-forest-labs/FLUX.1-schnell'
controlnet_model = 'YishaoAI/flux-dev-controlnet-canny-kid-clothes'

controlnet = FluxControlNetModel.from_pretrained(controlnet_model, torch_dtype=torch.float16)
pipe = FluxPipeline.from_pretrained(base_model, controlnet=controlnet, torch_dtype=torch.float16)
pipe.enable_model_cpu_offload()
pipe.to("cuda")

canny = CannyDetector()

def inpaint(image, mask, prompt, strength, num_inference_steps, guidance_scale, controlnet_conditioning_scale):
    canny_image = canny(image)
    
    image_res = pipe(
        prompt,
        image=image,
        control_image=canny_image,
        controlnet_conditioning_scale=controlnet_conditioning_scale,
        mask_image=mask,
        strength=strength,
        num_inference_steps=num_inference_steps,
        guidance_scale=guidance_scale,
    ).images[0]
    
    return image_res

iface = gr.Interface(
    fn=inpaint,
    inputs=[
        gr.Image(type="pil", label="Input Image"),
        gr.Image(type="pil", label="Mask Image"),
        gr.Textbox(label="Prompt"),
        gr.Slider(0, 1, value=0.95, label="Strength"),
        gr.Slider(1, 100, value=50, step=1, label="Number of Inference Steps"),
        gr.Slider(0, 20, value=5, label="Guidance Scale"),
        gr.Slider(0, 1, value=0.5, label="ControlNet Conditioning Scale")
    ],
    outputs=gr.Image(type="pil", label="Output Image"),
    title="Flux Inpaint AI Model",
    description="Upload an image and a mask, then provide a prompt to generate an inpainted image."
)

iface.launch()