app.py

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
from transformers import T5Tokenizer, T5TokenizerFast

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

# Try to load the fast tokenizer, fall back to slow if necessary
try:
    tokenizer = T5TokenizerFast.from_pretrained(base_model)
except ValueError:
    print("Fast tokenizer not available, falling back to slow tokenizer")
    tokenizer = T5Tokenizer.from_pretrained(base_model)

controlnet = FluxControlNetModel.from_pretrained(
    controlnet_model, torch_dtype=torch.float16
)
pipe = FluxPipeline.from_pretrained(
    base_model, controlnet=controlnet, torch_dtype=torch.float16, tokenizer=tokenizer
)
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()