tokenizer

app.py

@@ -3,20 +3,41 @@ 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'
+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)
+# 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):
+
+def inpaint(
+    image,
+    mask,
+    prompt,
+    strength,
+    num_inference_steps,
+    guidance_scale,
+    controlnet_conditioning_scale,
+):
     canny_image = canny(image)
-    
+
     image_res = pipe(
         prompt,
         image=image,
@@ -27,9 +48,10 @@ def inpaint(image, mask, prompt, strength, num_inference_steps, guidance_scale,
         num_inference_steps=num_inference_steps,
         guidance_scale=guidance_scale,
     ).images[0]
-    
+
     return image_res
 
+
 iface = gr.Interface(
     fn=inpaint,
     inputs=[
@@ -39,11 +61,11 @@ iface = gr.Interface(
         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")
+        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."
+    description="Upload an image and a mask, then provide a prompt to generate an inpainted image.",
 )
 
 iface.launch()

requirements.txt

@@ -4,3 +4,5 @@ transformers
 accelerate
 controlnet_aux
 gradio
+sentencepiece
+tokenizers