harahan/meanflownft

MeanFlowNFT is the official code release for a research paper about training AI image and video generators to follow human preferences more closely, without slowing down how fast they can create pictures. It comes out of Tencent Hunyuan and the Hong Kong University of Science and Technology, and it is published alongside a paper on arXiv, a project webpage, and models hosted on Hugging Face. The starting point is a family of generators called MeanFlow, which can produce a finished image or video in very few steps by learning an average speed of change rather than simulating every tiny step of the generation process. The problem this project solves is how to further improve those fast generators using reinforcement learning, a training method that rewards outputs humans or scoring models rate highly. Normally, adjusting a fast average-speed generator this way is awkward, because the reward signal is easiest to compute using the model's underlying step by step behavior, not its shortcut average behavior. MeanFlowNFT works around this by optimizing rewards in that step by step space internally while keeping the actual model in its fast average-speed form, so generation stays just as quick as before. It builds on a method called DiffusionNFT and carries over a guarantee that, in an ideal case, each training update can only improve the model, not make it worse. According to the results shown, MeanFlowNFT beats other few-step image generation models on six out of eight measured qualities, and a four-step video version outperforms a video model that used fifty steps and its own reinforcement learning process. The current public code covers training and running the SD3.5-Medium image model, a version for the Wan2.1 video model is being built on a separate branch. Setup requires Python 3.10, PyTorch 2.6.0, and CUDA 12.4, installed through a provided requirements file, plus a separately downloaded copy of the SD3.5-Medium model after accepting its license on Hugging Face. Training also needs several reward scoring models such as PickScore, HPSv2, and CLIPScore, most of which download automatically, with instructions for using local copies instead. The training process itself is designed to run across multiple machines with GPUs and involves first generating training data from the base model before running the actual reinforcement learning step. This is research code aimed at people already working in image and video generation, not a beginner-friendly tool.