Update app.py

app.py

@@ -11,6 +11,7 @@ from gradio_imageslider import ImageSlider
 from PIL import Image
 from huggingface_hub import snapshot_download
 import requests
+from transformers import T5TokenizerFast
 
 # For ESRGAN (requires pip install basicsr gfpgan)
 try:
@@ -151,6 +152,7 @@ def enhance_image(
     upscale_factor,
     denoising_strength,
     custom_prompt,
+    tile_size,
     progress=gr.Progress(track_tqdm=True),
 ):
     """Main enhancement function"""
@@ -162,11 +164,13 @@ def enhance_image(
             dtype = torch.bfloat16 if device == "cuda" else torch.float32
 
             print(f"📥 Loading FLUX Img2Img on {device}...")
+            tokenizer_2 = T5TokenizerFast.from_pretrained("black-forest-labs/FLUX.1-dev", subfolder="text_encoder_2")
             pipe = FluxImg2ImgPipeline.from_pretrained(
                 "black-forest-labs/FLUX.1-dev",
                 torch_dtype=dtype,
                 low_cpu_mem_usage=True,
-                device_map="balanced"
+                device_map="balanced",
+                tokenizer_2=tokenizer_2
             )
             pipe.enable_vae_tiling()
             pipe.enable_vae_slicing()
@@ -192,11 +196,13 @@ def enhance_image(
             device = "cpu"
             dtype = torch.float32
             # Reload on CPU if needed
+            tokenizer_2 = T5TokenizerFast.from_pretrained("black-forest-labs/FLUX.1-dev", subfolder="text_encoder_2")
             pipe = FluxImg2ImgPipeline.from_pretrained(
                 "black-forest-labs/FLUX.1-dev",
                 torch_dtype=dtype,
                 low_cpu_mem_usage=True,
-                device_map=None
+                device_map=None,
+                tokenizer_2=tokenizer_2
             )
             pipe.enable_vae_tiling()
             pipe.enable_vae_slicing()
@@ -241,7 +247,7 @@ def enhance_image(
         num_inference_steps,
         3.5,  # Updated guidance_scale to match workflow (3.5)
         generator,
-        tile_size=1024,
+        tile_size=tile_size,
         overlap=32
     )
 
@@ -326,6 +332,15 @@ with gr.Blocks(css=css, title="🎨 AI Image Upscaler - FLUX") as demo:
                 value=0.3,
                 info="Controls how much the image is transformed"
             )
+
+            tile_size = gr.Slider(
+                label="Tile Size",
+                minimum=256,
+                maximum=2048,
+                step=64,
+                value=1024,
+                info="Size of tiles for processing (larger = faster but more memory)"
+            )
             
             with gr.Row():
                 randomize_seed = gr.Checkbox(
@@ -370,6 +385,7 @@ with gr.Blocks(css=css, title="🎨 AI Image Upscaler - FLUX") as demo:
             upscale_factor,
             denoising_strength,
             custom_prompt,
+            tile_size
         ],
         outputs=[result_slider]
     )