Upload README.md with huggingface_hub

README.md

@@ -1,3 +1,110 @@
 ---
 license: apache-2.0
 ---
+
+# 🚀 SenseFlow: Scaling Distribution Matching for Flow-based Text-to-Image Distillation
+
+[![arXiv](https://img.shields.io/badge/Arxiv-2506.00523-b31b1b)](https://arxiv.org/abs/2506.00523) 
+[![GitHub Repo stars](https://img.shields.io/github/stars/XingtongGe/SenseFlow.svg?style=social&label=Star&maxAge=60)](https://github.com/XingtongGe/SenseFlow)
+
+[![Model on HF](https://huggingface.co/datasets/huggingface/badges/resolve/main/model-on-hf-md-dark.svg)](https://huggingface.co/domiso/SenseFlow)
+
+<!-- [🤗 HuggingFace Model](https://huggingface.co/domiso/SenseFlow)  -->
+
+[Xingtong Ge](https://xingtongge.github.io/)<sup>1,2</sup>, Xin Zhang<sup>2</sup>, [Tongda Xu](https://tongdaxu.github.io/)<sup>3</sup>, [Yi Zhang](https://zhangyi-3.github.io/)<sup>4</sup>, [Xinjie Zhang](https://xinjie-q.github.io/)<sup>1</sup>, [Yan Wang](https://yanwang202199.github.io/)<sup>3</sup>, [Jun Zhang](https://eejzhang.people.ust.hk/)<sup>1</sup>
+
+<sup>1</sup>HKUST, <sup>2</sup>SenseTime Research, <sup>3</sup>Tsinghua University, <sup>4</sup>CUHK MMLab  
+
+## Abstract
+
+The Distribution Matching Distillation (DMD) has been successfully applied to text-to-image diffusion models such as Stable Diffusion (SD) 1.5. However, vanilla DMD suffers from convergence difficulties on large-scale flow-based text-to-image models, such as SD 3.5 and FLUX. In this paper, we first analyze the issues when applying vanilla DMD on large-scale models. Then, to overcome the scalability challenge, we propose implicit distribution alignment (IDA) to constrain the divergence between the generator and the fake distribution. Furthermore, we propose intra-segment guidance (ISG) to relocate the timestep denoising importance from the teacher model. With IDA alone, DMD converges for SD 3.5; employing both IDA and ISG, DMD converges for SD 3.5 and FLUX.1 dev. Together with a scaled VFM-based discriminator, our final model, dubbed **SenseFlow**, achieves superior performance in distillation for both diffusion based text-to-image models such as SDXL, and flow-matching models such as SD 3.5 Large and FLUX.1 dev.
+
+## SenseFlow-FLUX.1 dev (supports 4–8-step generation)
+* `SenseFlow-FLUX/diffusion_pytorch_model.safetensors`: the DiT checkpoint.
+* `SenseFlow-FLUX/config.json`: the config of DiT using in our model.
+
+
+### Usage
+
+1. prepare the base checkpoint of FLUX.1 dev to `Path/to/FLUX`
+2. Use `SenseFlow-FLUX` to replace the transformer folder `Path/to/FLUX/transformer`, obtaining the `Path/to/SenseFlow-FLUX`.
+
+#### Using the Euler sampler
+```python
+import torch
+from diffusers import FluxPipeline
+from diffusers import FlowMatchEulerDiscreteScheduler
+
+pipe = FluxPipeline.from_pretrained("Path/to/SenseFlow-FLUX", torch_dtype=torch.bfloat16).to("cuda")
+
+prompt="A cat sleeping on a windowsill with white curtains fluttering in the breeze"
+
+images = pipe(
+    prompt,
+    height=1024,
+    width=1024,
+    num_inference_steps=4,
+    max_sequence_length=512,
+).images[0]
+
+images.save("output.png")
+```
+#### Using the x0 sampler (similar to the LCMScheduler in diffusers)
+```python
+import torch
+from diffusers import FluxPipeline
+from diffusers import FlowMatchEulerDiscreteScheduler
+from typing import Union, Tuple, Optional
+
+class FlowMatchEulerX0Scheduler(FlowMatchEulerDiscreteScheduler):
+    def step(
+        self,
+        model_output: torch.FloatTensor,
+        timestep: Union[float, torch.FloatTensor],
+        sample: torch.FloatTensor,
+        generator: Optional[torch.Generator] = None,
+        return_dict: bool = True,
+    ) -> Union[FlowMatchEulerDiscreteSchedulerOutput, Tuple]:
+
+        if self.step_index is None:
+            self._init_step_index(timestep)
+
+        sample = sample.to(torch.float32)  # Ensure precision
+
+        sigma = self.sigmas[self.step_index]
+        sigma_next = self.sigmas[self.step_index + 1]
+
+        # 1. Compute x0 from model output (assuming model predicts noise)
+        x0 = sample - sigma * model_output
+
+        # 2. Add noise to x0 to get the sample for the next step
+        noise = torch.randn_like(sample)
+        prev_sample = (1 - sigma_next) * x0 + sigma_next * noise
+
+        prev_sample = prev_sample.to(model_output.dtype)  # Convert back to original dtype
+        self._step_index += 1  # Move to next step
+
+        if not return_dict:
+            return (prev_sample,)
+        
+        return FlowMatchEulerDiscreteSchedulerOutput(prev_sample=prev_sample)
+
+pipe = FluxPipeline.from_pretrained("Path/to/SenseFlow-FLUX", torch_dtype=torch.bfloat16).to("cuda")
+pipe.scheduler = FlowMatchEulerX0Scheduler.from_config(pipe.scheduler.config)
+
+prompt="A cat sleeping on a windowsill with white curtains fluttering in the breeze"
+
+images = pipe(
+    prompt,
+    height=1024,
+    width=1024,
+    num_inference_steps=4,
+    max_sequence_length=512,
+).images[0]
+
+images.save("output.png")
+```
+
+## DanceGRPO-SenseFlow (supports 4–8-step generation)
+
+comming soon!