---
base_model:
- black-forest-labs/FLUX.1-Fill-dev
- bytedance-research/OneReward
language:
- en
license: cc-by-nc-4.0
pipeline_tag: image-to-image
library_name: diffusers
---

# OneReward - ComfyUI

[![arXiv](https://img.shields.io/badge/arXiv-Paper-b31b1b.svg)](https://arxiv.org/abs/2508.21066) [![GitHub Repo](https://img.shields.io/badge/GitHub-Repo-green?logo=github)](https://github.com/bytedance/OneReward) [![GitHub Pages](https://img.shields.io/badge/GitHub-Project-blue?logo=github)](https://one-reward.github.io/)
<br>

This repo contains the checkpoint from [OneReward](https://huggingface.co/bytedance-research/OneReward) processed into a single model suitable for ComfyUI use.

**OneReward** is a novel RLHF methodology for the visual domain by employing Qwen2.5-VL as a generative reward model to enhance multitask reinforcement learning, significantly improving the policy model’s generation ability across multiple subtask. Building on OneReward, **FLUX.1-Fill-dev-OneReward** - based on FLUX Fill [dev], outperforms closed-source FLUX Fill [Pro] in inpainting and outpainting tasks, serving as a powerful new baseline for future research in unified image editing.

For more details and examples see original model repo: [**OneReward**](https://huggingface.co/bytedance-research/OneReward)

Partial redrawing workflow： [**OneReward redrawing**](https://www.runninghub.ai/post/1965826858283634690/?inviteCode=rh-v1091)  

Image expansion workflow： [**OneReward Image expansion**](https://www.runninghub.ai/post/1966338160613601281/?inviteCode=rh-v1091)
## Sample Usage

The following code snippet illustrates how to use the model with the `diffusers` library. Note that this requires the custom `FluxFillCFGPipeline` defined in the [official source code](https://github.com/bytedance/OneReward/blob/main/src/pipeline_flux_fill_with_cfg.py).

```python
import torch
from diffusers.utils import load_image
from diffusers import FluxTransformer2DModel

# Note: pipeline_flux_fill_with_cfg.py must be available in your local environment
from src.pipeline_flux_fill_with_cfg import FluxFillCFGPipeline

transformer_onereward = FluxTransformer2DModel.from_pretrained(
    "bytedance-research/OneReward",
    subfolder="flux.1-fill-dev-OneReward-transformer",
    torch_dtype=torch.bfloat16
)

pipe = FluxFillCFGPipeline.from_pretrained(
    "black-forest-labs/FLUX.1-Fill-dev", 
    transformer=transformer_onereward,
    torch_dtype=torch.bfloat16).to("cuda")

# Example: Image Fill
image = load_image('assets/image.png')
mask = load_image('assets/mask_fill.png')
image = pipe(
    prompt='the words "ByteDance", and in the next line "OneReward"',
    negative_prompt="nsfw",
    image=image,
    mask_image=mask,
    height=image.height,
    width=image.width,
    guidance_scale=1.0,
    true_cfg=4.0,
    num_inference_steps=50,
    generator=torch.Generator("cpu").manual_seed(0)
).images[0]
image.save(f"image_fill.jpg")
```

## Citation

```bibtex
@article{gong2025onereward,
  title={OneReward: Unified Mask-Guided Image Generation via Multi-Task Human Preference Learning},
  author={Gong, Yuan and Wang, Xionghui and Wu, Jie and Wang, Shiyin and Wang, Yitong and Wu, Xinglong},
  journal={arXiv preprint arXiv:2508.21066},
  year={2025}
}
```