Update app.py

app.py

@@ -164,9 +164,9 @@ 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-Krea-dev", subfolder="tokenizer_2", token=huggingface_token)
+            tokenizer_2 = T5TokenizerFast.from_pretrained("black-forest-labs/FLUX.1-schnell", subfolder="tokenizer_2", token=huggingface_token)
             pipe = FluxImg2ImgPipeline.from_pretrained(
-                "black-forest-labs/FLUX.1-Krea-dev",
+                "black-forest-labs/FLUX.1-schnell",
                 torch_dtype=dtype,
                 low_cpu_mem_usage=True,
                 device_map="balanced",
@@ -197,9 +197,9 @@ def enhance_image(
             device = "cpu"
             dtype = torch.float32
             # Reload on CPU if needed
-            tokenizer_2 = T5TokenizerFast.from_pretrained("black-forest-labs/FLUX.1-Krea-dev", subfolder="tokenizer_2", token=huggingface_token)
+            tokenizer_2 = T5TokenizerFast.from_pretrained("black-forest-labs/FLUX.1-schnell", subfolder="tokenizer_2", token=huggingface_token)
             pipe = FluxImg2ImgPipeline.from_pretrained(
-                "black-forest-labs/FLUX.1-Krea-dev",
+                "black-forest-labs/FLUX.1-schnell",
                 torch_dtype=dtype,
                 low_cpu_mem_usage=True,
                 device_map=None,
@@ -319,11 +319,11 @@ with gr.Blocks(css=css, title="🎨 AI Image Upscaler - FLUX") as demo:
             
             num_inference_steps = gr.Slider(
                 label="Number of Inference Steps",
-                minimum=8,
+                minimum=1,
                 maximum=50,
                 step=1,
-                value=25,
-                info="More steps = better quality but slower"
+                value=4,
+                info="More steps = better quality but slower (default 4 for schnell)"
             )
             
             denoising_strength = gr.Slider(
@@ -394,7 +394,7 @@ with gr.Blocks(css=css, title="🎨 AI Image Upscaler - FLUX") as demo:
     
     gr.HTML("""
     <div style="margin-top: 2rem; padding: 1rem; background: #f0f0f0; border-radius: 8px;">
-        <p><strong>Note:</strong> This upscaler uses the Flux dev model. Users are responsible for obtaining commercial rights if used commercially under their license.</p>
+        <p><strong>Note:</strong> This upscaler uses the Flux.1-schnell model. Users are responsible for obtaining commercial rights if used commercially under their license.</p>
     </div>
     """)