Update README.md

README.md

@@ -28,7 +28,7 @@ Distilled 4-step and 8-step FLUX.1 models proposed in the paper:
 [Sai Bi](https://sai-bi.github.io/)<sup>2</sup><br>
 <sup>1</sup>Stanford University, <sup>2</sup>Adobe Research
 <br>
-[[arXiv](https://arxiv.org/abs/2510.14974)] [[Code](https://github.com/Lakonik/piFlow)] [[pi-Qwen Demo🤗](https://huggingface.co/spaces/Lakonik/pi-Qwen)] [[pi-FLUX Demo🤗](https://huggingface.co/spaces/Lakonik/pi-FLUX.1)] 
+[[arXiv](https://arxiv.org/abs/2510.14974)] [[Code](https://github.com/Lakonik/piFlow)] [[pi-Qwen Demo🤗](https://huggingface.co/spaces/Lakonik/pi-Qwen)] [[pi-FLUX Demo🤗](https://huggingface.co/spaces/Lakonik/pi-FLUX.1)] [[pi-FLUX.2 Demo🤗](https://huggingface.co/spaces/Lakonik/pi-FLUX.2)]
 
 ![teaser](https://cdn-uploads.huggingface.co/production/uploads/638067fcb334960c987fbeda/H0J1LYUcSS5YqOwZqQ0Jb.jpeg)
 
@@ -42,8 +42,8 @@ We provide diffusers pipelines for easy inference. The following code demonstrat
 Note: For the 8-NFE version, replace `gmflux_k8_piid_4step` with `gmflux_k8_piid_8step` and set `num_inference_steps=8`.
 ```python
 import torch
-from diffusers import FlowMatchEulerDiscreteScheduler
-from lakonlab.pipelines.piflux_pipeline import PiFluxPipeline
+from lakonlab.models.diffusions.schedulers import FlowMapSDEScheduler
+from lakonlab.pipelines.pipeline_piflux import PiFluxPipeline
 
 pipe = PiFluxPipeline.from_pretrained(
     'black-forest-labs/FLUX.1-dev',
@@ -52,8 +52,8 @@ adapter_name = pipe.load_piflow_adapter(  # you may later call `pipe.set_adapter
     'Lakonik/pi-FLUX.1',
     subfolder='gmflux_k8_piid_4step',
     target_module_name='transformer')
-pipe.scheduler = FlowMatchEulerDiscreteScheduler.from_config(  # use fixed shift=3.2
-    pipe.scheduler.config, shift=3.2, use_dynamic_shifting=False)
+pipe.scheduler = FlowMapSDEScheduler.from_config(  # use fixed shift=3.2
+    pipe.scheduler.config, shift=3.2, use_dynamic_shifting=False, final_step_size_scale=0.5)
 pipe = pipe.to('cuda')
 
 out = pipe(
@@ -71,8 +71,8 @@ out.save('gmflux_4nfe.png')
 ### 4-NFE DX-FLUX (DX Policy)
 ```
 import torch
-from diffusers import FlowMatchEulerDiscreteScheduler
-from lakonlab.pipelines.piflux_pipeline import PiFluxPipeline
+from lakonlab.models.diffusions.schedulers import FlowMapSDEScheduler
+from lakonlab.pipelines.pipeline_piflux import PiFluxPipeline
 
 pipe = PiFluxPipeline.from_pretrained(
     'black-forest-labs/FLUX.1-dev',
@@ -85,8 +85,8 @@ adapter_name = pipe.load_piflow_adapter(  # you may later call `pipe.set_adapter
     'Lakonik/pi-FLUX.1',
     subfolder='dxflux_n10_piid_4step',
     target_module_name='transformer')
-pipe.scheduler = FlowMatchEulerDiscreteScheduler.from_config(  # use fixed shift=3.2
-    pipe.scheduler.config, shift=3.2, use_dynamic_shifting=False)
+pipe.scheduler = FlowMapSDEScheduler.from_config(  # use fixed shift=3.2
+    pipe.scheduler.config, shift=3.2, use_dynamic_shifting=False, final_step_size_scale=0.5)
 pipe = pipe.to('cuda')
 
 out = pipe(