Spaces:
Runtime error
Runtime error
File size: 2,147 Bytes
c744305 5745ae4 0c68dbd c744305 47a3e09 c744305 47a3e09 c744305 47a3e09 c744305 47a3e09 c744305 47a3e09 c744305 47a3e09 c744305 47a3e09 c744305 47a3e09 c744305 47a3e09 c744305 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 |
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()
|