Update app.py

app.py

@@ -23,6 +23,9 @@ device = "cuda" if torch.cuda.is_available() else "cpu"
 MAX_SEED = np.iinfo(np.int32).max
 MAX_IMAGE_SIZE = 1024
 
+# Pre-shifted custom sigmas for 8-step turbo inference
+TURBO_SIGMAS = [1.0, 0.6509, 0.4374, 0.2932, 0.1893, 0.1108, 0.0495, 0.00031]
+
 hf_client = InferenceClient(
     api_key=os.environ.get("HF_TOKEN"),
 )
@@ -78,6 +81,13 @@ pipe = Flux2Pipeline.from_pretrained(
     transformer=dit,
     torch_dtype=torch.bfloat16
 )
+
+# Load the Turbo LoRA
+pipe.load_lora_weights(
+    "fal/FLUX.2-Turbo", 
+    weight_name="flux.2-turbo-lora.safetensors"
+)
+
 pipe.to(device)
 
 # Pull pre-compiled Flux2 Transformer blocks from HF hub
@@ -157,14 +167,17 @@ def update_dimensions_from_image(image_list):
     
     return new_width, new_height
 
-# Updated duration function to match generate_image arguments (including progress)
-def get_duration(prompt_embeds, image_list, width, height, num_inference_steps, guidance_scale, seed, progress=gr.Progress(track_tqdm=True)):
+# Updated duration function for Turbo (much faster with fewer steps)
+def get_duration(prompt_embeds, image_list, width, height, num_inference_steps, guidance_scale, seed, use_turbo, progress=gr.Progress(track_tqdm=True)):
     num_images = 0 if image_list is None else len(image_list)
     step_duration = 1 + 0.8 * num_images
+    # Turbo mode uses fewer steps, so shorter duration
+    if use_turbo:
+        return max(30, 8 * step_duration + 10)  # Fixed 8 steps for turbo
     return max(65, num_inference_steps * step_duration + 10)
 
 @spaces.GPU(duration=get_duration)
-def generate_image(prompt_embeds, image_list, width, height, num_inference_steps, guidance_scale, seed, progress=gr.Progress(track_tqdm=True)):
+def generate_image(prompt_embeds, image_list, width, height, num_inference_steps, guidance_scale, seed, use_turbo, progress=gr.Progress(track_tqdm=True)):
     # Move embeddings to GPU only when inside the GPU decorated function
     prompt_embeds = prompt_embeds.to(device)
     
@@ -173,13 +186,19 @@ def generate_image(prompt_embeds, image_list, width, height, num_inference_steps
     pipe_kwargs = {
         "prompt_embeds": prompt_embeds,
         "image": image_list,
-        "num_inference_steps": num_inference_steps,
         "guidance_scale": guidance_scale,
         "generator": generator,
         "width": width,
         "height": height,
     }
     
+    # Use Turbo sigmas or regular inference steps
+    if use_turbo:
+        pipe_kwargs["sigmas"] = TURBO_SIGMAS
+        pipe_kwargs["num_inference_steps"] = 8  # Turbo always uses 8 steps
+    else:
+        pipe_kwargs["num_inference_steps"] = num_inference_steps
+    
     # Progress bar for the actual generation steps
     if progress:
         progress(0, desc="Starting generation...")
@@ -187,7 +206,7 @@ def generate_image(prompt_embeds, image_list, width, height, num_inference_steps
     image = pipe(**pipe_kwargs).images[0]
     return image
 
-def infer(prompt, input_images=None, seed=42, randomize_seed=False, width=1024, height=1024, num_inference_steps=50, guidance_scale=2.5, prompt_upsampling=False, progress=gr.Progress(track_tqdm=True)):
+def infer(prompt, input_images=None, seed=42, randomize_seed=False, width=1024, height=1024, num_inference_steps=8, guidance_scale=2.5, prompt_upsampling=False, use_turbo=True, progress=gr.Progress(track_tqdm=True)):
     
     if randomize_seed:
         seed = random.randint(0, MAX_SEED)
@@ -221,7 +240,8 @@ def infer(prompt, input_images=None, seed=42, randomize_seed=False, width=1024,
         height, 
         num_inference_steps, 
         guidance_scale, 
-        seed, 
+        seed,
+        use_turbo,
         progress
     )
     
@@ -252,8 +272,8 @@ css="""
 with gr.Blocks() as demo:
     
     with gr.Column(elem_id="col-container"):
-        gr.Markdown(f"""# FLUX.2 [dev]
-FLUX.2 [dev] is a 32B model rectified flow capable of generating, editing and combining images based on text instructions model [[model](https://huggingface.co/black-forest-labs/FLUX.2-dev)], [[blog](https://bfl.ai/blog/flux-2)]
+        gr.Markdown(f"""# FLUX.2 [dev] Turbo
+FLUX.2 [dev] with [Turbo LoRA by fal](https://huggingface.co/fal/FLUX.2-Turbo) - a 32B rectified flow model capable of generating, editing and combining images based on text instructions in just 8 steps [[model](https://huggingface.co/black-forest-labs/FLUX.2-dev)], [[blog](https://bfl.ai/blog/flux-2)]
         """)
         with gr.Row():
             with gr.Column():
@@ -278,6 +298,12 @@ FLUX.2 [dev] is a 32B model rectified flow capable of generating, editing and co
                     )
                 
                 with gr.Accordion("Advanced Settings", open=False):
+                    use_turbo = gr.Checkbox(
+                        label="Use Turbo Mode (8 steps)",
+                        value=True,
+                        info="Enable Turbo LoRA for fast 8-step generation"
+                    )
+                    
                     prompt_upsampling = gr.Checkbox(
                         label="Prompt Upsampling",
                         value=True,
@@ -315,7 +341,7 @@ FLUX.2 [dev] is a 32B model rectified flow capable of generating, editing and co
                     with gr.Row():
                         
                         num_inference_steps = gr.Slider(
-                            label="Number of inference steps",
+                            label="Number of inference steps (ignored in Turbo mode)",
                             minimum=1,
                             maximum=100,
                             step=1,
@@ -327,7 +353,7 @@ FLUX.2 [dev] is a 32B model rectified flow capable of generating, editing and co
                             minimum=0.0,
                             maximum=10.0,
                             step=0.1,
-                            value=4,
+                            value=2.5,
                         )
                 
                 
@@ -363,7 +389,7 @@ FLUX.2 [dev] is a 32B model rectified flow capable of generating, editing and co
     gr.on(
         triggers=[run_button.click, prompt.submit],
         fn=infer,
-        inputs=[prompt, input_images, seed, randomize_seed, width, height, num_inference_steps, guidance_scale, prompt_upsampling],
+        inputs=[prompt, input_images, seed, randomize_seed, width, height, num_inference_steps, guidance_scale, prompt_upsampling, use_turbo],
         outputs=[result, seed]
     )