Upload RTDMD checkpoint and model card

README.md

@@ -0,0 +1,103 @@
+---
+license: apache-2.0
+library_name: diffusers
+pipeline_tag: text-to-image
+tags:
+  - text-to-image
+  - diffusion
+  - flow-matching
+  - lora
+  - rtdmd
+  - flux
+---
+
+<div align="center">
+
+<img width="70%" height="70%" alt="logo" src="https://github.com/user-attachments/assets/4d534e80-f8ec-4c0b-948f-730cc0311961" />
+
+<h2> Reinforcing Few-step Generators via Reward-Tilted Distribution Matching </h2>
+
+<p><b>Reward-Tilted DMD &nbsp;·&nbsp; Ambient-Consistent Distillation &nbsp;·&nbsp; Hybrid Policy Gradient</b></p>
+
+[![Paper](https://img.shields.io/badge/paper-arXiv-A42C25?style=for-the-badge&logo=arxiv&logoColor=white)](<TODO: arxiv link>)
+[![Github](https://img.shields.io/badge/Harahan%2FRTDMD-000000?style=for-the-badge&logo=github&logoColor=white)](https://github.com/Harahan/RTDMD)
+[![Hugging Face Collection](https://img.shields.io/badge/RTDMD_Collection-fcd022?style=for-the-badge&logo=huggingface&logoColor=000)](https://huggingface.co/collections/Harahan/rtdmd)
+
+[![License: Apache 2.0](https://img.shields.io/badge/License-Apache%202.0-blue.svg)](https://opensource.org/licenses/Apache-2.0)
+[![Python](https://img.shields.io/badge/Python-3.10%2B-blue.svg)](https://www.python.org/)
+
+</div>
+
+<div align="center">
+
+[Yushi Huang](https://harahan.github.io/)<sup>1, 2,</sup>\*<sup>†</sup>, [Xiangxin Zhou](https://zhouxiangxin1998.github.io/)<sup>2,</sup>\*, Ruoyu Wang<sup>2, 3,</sup>\*<sup>†</sup>, [Chi Zhang](https://icoz69.github.io/)<sup>3</sup>, [Jun Zhang](https://eejzhang.people.ust.hk/)<sup>1</sup>, [Tianyu Pang](https://p2333.github.io/)<sup>2,</sup>‡
+
+<sup>1</sup>The Hong Kong University of Science and Technology &nbsp;&nbsp;
+<sup>2</sup>Tencent Hunyuan &nbsp;&nbsp;
+<sup>3</sup>Westlake University
+
+\* Equal contribution &nbsp;·&nbsp; † Work done during internship at Tencent Hunyuan &nbsp;·&nbsp; ‡ Corresponding author
+
+</div>
+
+---
+
+## Abstract
+
+We propose **Reward-Tilted Distribution Matching Distillation (RTDMD)**, a two-stage framework that unifies distribution-matching distillation with reward-guided RL for few-step flow generators. Minimizing the KL divergence to a *reward-tilted teacher distribution* decomposes naturally into a **distribution-matching** term and a **reward-maximization** term, instantiated as **Ambient-Consistent DMD (AC-DMD)** for the cold start and a **hybrid policy gradient** for the RL stage.
+
+This repository hosts the **FLUX.2-klein 9B** RTDMD LoRA checkpoints for **4-NFE** text-to-image generation.
+
+<div align="center">
+  <img src="https://github.com/user-attachments/assets/cb1fb0da-d388-4846-9017-66bccebd0749" alt="RTDMD teaser" width="70%">
+  <br/>
+  <em>4-step samples from RTDMD-distilled generators (no classifier-free guidance).</em>
+</div>
+
+<div align="center">
+  <img src="https://github.com/user-attachments/assets/61a64fca-a143-40ae-9e36-79c6fcb5b696" alt="RTDMD method overview" width="70%">
+  <br/>
+  <em>RTDMD overview. <b>Det.</b> = deterministic final step, <b>Stoc.</b> = stochastic intermediate steps.</em>
+</div>
+
+## Contents
+
+```text
+.
+├── cold_start/
+│   └── generator_ema.pt    # Stage-1 AC-DMD LoRA
+└── rtdmd/
+    └── generator_ema.pt    # Stage-2 RTDMD LoRA
+```
+
+Each `generator_ema.pt` is a `torch.save`-d LoRA `state_dict` containing only adapter weights. The two adapters are intended to be stacked in order: first `cold_start`, then `rtdmd`.
+
+## Usage
+
+The recommended path is the RTDMD inference CLI, which handles LoRA stacking and the CPS scheduler used during training.
+
+```bash
+git clone https://github.com/Harahan/RTDMD.git
+cd RTDMD
+pip install -r requirements.txt
+pip install -e .
+
+huggingface-cli download Harahan/FLUX2-9B-RTDMD --local-dir ./ckpts/flux2_9b
+
+python inference.py configs/inference/flux2_9b.yaml \
+    --override lora_paths='["./ckpts/flux2_9b/cold_start/generator_ema.pt","./ckpts/flux2_9b/rtdmd/generator_ema.pt"]' \
+    --override eval_reward=false \
+    --prompt "a cute cat sitting on a windowsill"
+```
+
+RTDMD is trained for 4-step generation with CPS (`cps_eta=0.9`) and `guidance_scale=1.0`.
+
+## Citation
+
+```bibtex
+<TODO: add bibtex after arXiv release>
+```
+
+## License
+
+Apache 2.0, same as the upstream [RTDMD](https://github.com/Harahan/RTDMD) repository. The base model `black-forest-labs/FLUX.2-klein-9B` is governed by its own license; please review and comply with it separately.

cold_start/generator_ema.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:da6ee352e06684a8a472f204600a8505415a8d96c3b975c7a3e5f6dccb566098
+size 128008714

rtdmd/generator_ema.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:85ab08c37ad96d60476fffba7b42a9b5a05e16f205cbf30e4e59ca987d5fb39b
+size 128003914