thu-ml/DiT-Extrapolation

Diffusion-Transformer Extrapolation for Long Video Generation

This repository provides the official implementation of RIFLEx (ICML 2025), UltraViCo (ICLR 2026) and UltraImage , which achieve diffusion-transformer extrapolation for long video generation and high resolution image generation in a plug-and-play way.

This repository hosts RIFLEx and UltraViCo on separate branches, and the code is fully open source.

• RIFLEx:

  • main: HunyuanVideo-diffusers and CogVideoX-diffusers
  • multi-gpu: multi-GPU inference for HunyuanVideo

• UltraViCo:

  • ultra-wan: UltraViCo for Wan2.1
  • ultra-hunyuan:UltraViCo for HunyuanVideo

• UltraImage:

  • ultra-flux: UltraImage for Flux
  • ultra-qwenimage: UltraImage for Qwen-Image

🎉 Supported Models

Here, we list the SOTA video diffusion transformers that RIFLEx has been applied to. We are continuously working to support more models. Feel free to suggest additional models you would like us to support!

To be continued……

🔥🔥 News

• 2026.1.26 : UltraViCo is accepted by ICLR 2026!
• 2025.5.1 : RIFLEx is accepted by ICML 2025!
• 2025.3.17 : RIFLEx is integrated to the official HunyuanVideo-I2V repository!
• 2025.3.1 : The code for CogVideoX-5B and fine-tuned CogVideoX-RIFLEx are released.
• 2025.2.26 RIFLEx is supported in HunyuanVideoGP, where a 10.5 s video at 1280x720 can be generated on an RTX 4090.
• 2025.2.26 RIFLEx is supported in ComfyUI-HunyuanVideoWrapper by KiJai.
• 2025.2.25 : The paper, project page, code and fine-tuned HunyuanVideo-RIFLEx are released.

RIFLEx Code

RIFLEx only adds a single line of code on the original 1D RoPE.

def get_1d_rotary_pos_embed_riflex(
    dim: int,
    pos: Union[np.ndarray, int],
    theta: float = 10000.0,
    k: Optional[int] = None,
    L_test: Optional[int] = None,
):
    '''
        k: the index for the intrinsic frequency in RoPE
        L_test: the number of frames for inference
    '''
    
    assert dim % 2 == 0
    if isinstance(pos, int):
        pos = torch.arange(pos)
    if isinstance(pos, np.ndarray):
        pos = torch.from_numpy(pos)
    freqs = 1.0 / (theta ** (torch.arange(0, dim, 2, device=pos.device)[: (dim // 2)].float() / dim)) 

    # === RIFLEx modification start ===
    # Reduce intrinsic frequency to stay within a single period after extrapolation (Eq.(8)).
    # Empirical observations show that a few videos may exhibit repetition in the tail frames.
    # To be conservative, we multiply 0.9 to keep extrapolated length below 90% of a period. 
    freqs[k-1] = 0.9 * 2 * torch.pi / L_test
    # === RIFLEx modification end ===

    freqs = torch.outer(pos, freqs)  
    freqs_cos = freqs.cos().repeat_interleave(2, dim=1).float()  
    freqs_sin = freqs.sin().repeat_interleave(2, dim=1).float()  
    return freqs_cos, freqs_sin

In riflex_utils.py, we show how to identify the intrinsic frequency in a RoPE-based pre-trained diffusion transformer.

Single GPU Inference with Diffusers for Quick Start

Installation

conda create -n riflex python=3.10
pip install -r requirements.txt
pip install -U bitsandbytes

Prompts

The example prompts for all models are listed in assets/prompts. The prompts on the project page can be found in assets/prompts/free_hunyuan.txt and assets/prompts/finetune_hunyuan.txt.

Please note that for single GPU inference with HunyuanVideo, Diffusers use DiffusersBitsAndBytesConfig to save memory, which may affect performance. To produce the demo on the project page, please refer to the Multi-GPU Inference section.

Inference for HunyuanVideo

python hunyuanvideo.py --k 4 --N_k 50 --num_frames 261 --prompt "A white and orange tabby cat is seen happily darting through a dense garden, as if chasing something. Its eyes are wide and happy as it jogs forward, scanning the branches, flowers, and leaves as it walks. The path is narrow as it makes its way between all the plants. the scene is captured from a ground-level angle, following the cat closely, giving a low and intimate perspective. The image is cinematic with warm tones and a grainy texture. The scattered daylight between the leaves and plants above creates a warm contrast, accentuating the cat’s orange fur. The shot is clear and sharp, with a shallow depth of field."

python hunyuanvideo.py --k 4 --N_k 66 --num_frames 261 --finetune --model_id "thu-ml/Hunyuan-RIFLEx-diffusers" --prompt "3D animation of a small, round, fluffy creature with big, expressive eyes explores a vibrant, enchanted forest. The creature, a whimsical blend of a rabbit and a squirrel, has soft blue fur and a bushy, striped tail. It hops along a sparkling stream, its eyes wide with wonder. The forest is alive with magical elements: flowers that glow and change colors, trees with leaves in shades of purple and silver, and small floating lights that resemble fireflies. The creature stops to interact playfully with a group of tiny, fairy-like beings dancing around a mushroom ring. The creature looks up in awe at a large, glowing tree that seems to be the heart of the forest."

Inference for CogVideoX

python cogvideox.py --k 2 --N_k 20 --num_frames 97 --prompt "3D animation of a small, round, fluffy creature with big, expressive eyes explores a vibrant, enchanted forest. The creature, a whimsical blend of a rabbit and a squirrel, has soft blue fur and a bushy, striped tail. It hops along a sparkling stream, its eyes wide with wonder. The forest is alive with magical elements: flowers that glow and change colors, trees with leaves in shades of purple and silver, and small floating lights that resemble fireflies. The creature stops to interact playfully with a group of tiny, fairy-like beings dancing around a mushroom ring. The creature looks up in awe at a large, glowing tree that seems to be the heart of the forest."

python cogvideox.py --k 2 --N_k 25 --num_frames 97 --finetune --model_id "thu-ml/CogVideoX-RIFLEx-diffusers" --prompt "A drone camera circles around a beautiful historic church built on a rocky outcropping along the Amalfi Coast, the view showcases historic and magnificent architectural details and tiered pathways and patios, waves are seen crashing against the rocks below as the view overlooks the horizon of the coastal waters and hilly landscapes of the Amalfi Coast Italy, several distant people are seen walking and enjoying vistas on patios of the dramatic ocean views, the warm glow of the afternoon sun creates a magical and romantic feeling to the scene, the view is stunning captured with beautiful photography."

Multi-GPU Inference ( Recommended )

To enhance inference speed and reproduce the demos in our project page, please use the multi-gpu inference. Details can be found in the multi-gpu branch.

References

If you find the code useful, please cite

@article{zhao2025riflex,
  title={Riflex: A free lunch for length extrapolation in video diffusion transformers},
  author={Zhao, Min and He, Guande and Chen, Yixiao and Zhu, Hongzhou and Li, Chongxuan and Zhu, Jun},
  journal={arXiv preprint arXiv:2502.15894},
  year={2025}
}

@article{zhao2025ultravico,
  title={UltraViCo: Breaking Extrapolation Limits in Video Diffusion Transformers},
  author={Zhao, Min and Zhu, Hongzhou and Wang, Yingze and Yan, Bokai and Zhang, Jintao and He, Guande and Yang, Ling and Li, Chongxuan and Zhu, Jun},
  journal={arXiv preprint arXiv:2511.20123},
  year={2025}
}
@article{zhao2025ultraimage,
  title={UltraImage: Rethinking Resolution Extrapolation in Image Diffusion Transformers},
  author={Zhao, Min and Yan, Bokai and Yang, Xue and Zhu, Hongzhou and Zhang, Jintao and Liu, Shilong and Li, Chongxuan and Zhu, Jun},
  journal={arXiv preprint arXiv:2512.04504},
  year={2025}
}