Add UniTranslator model card

README.md

@@ -0,0 +1,273 @@
+---
+license: other
+license_name: flux-1-dev-non-commercial-license
+license_link: https://huggingface.co/black-forest-labs/FLUX.1-Kontext-dev/blob/main/LICENSE.md
+library_name: transformers
+pipeline_tag: image-to-image
+inference: false
+base_model:
+  - Qwen/Qwen2.5-VL-3B-Instruct
+  - black-forest-labs/FLUX.1-Kontext-dev
+datasets:
+  - yztian/IIMT30k
+  - yztian/MTedIIMT
+  - yztian/PRIM
+language:
+  - de
+  - en
+  - fr
+  - ro
+  - cs
+  - ru
+tags:
+  - in-image-machine-translation
+  - image-translation
+  - image-editing
+  - multimodal
+  - qwen2.5-vl
+  - flux
+  - arxiv:2606.24333
+---
+
+# UniTranslator
+
+## A Unified Multimodal Framework for End-to-End In-Image Machine Translation
+
+<p align="center">
+  📄 <a href="https://arxiv.org/abs/2606.24333">Paper</a> &nbsp;·&nbsp;
+  💻 <a href="https://github.com/SeerRay-Lab/Unitranslator">Code</a> &nbsp;·&nbsp;
+  🤗 <a href="https://huggingface.co/SeerRay-Lab/Unitranslator">Model</a>
+</p>
+
+**UniTranslator**, accepted at **ECCV 2026**, is an end-to-end framework for **in-image machine translation (IIMT)**. Given an image and a translation instruction, it predicts the translated text and renders that translation back into the source text regions while preserving the surrounding scene, layout, and typography as closely as possible.
+
+![UniTranslator overview](https://raw.githubusercontent.com/SeerRay-Lab/Unitranslator/main/assets/uni-framework.png)
+
+UniTranslator introduces two components:
+
+- **Understand-Generation Alignment Module (UGAM):** aligns translation-understanding representations with the image-generation condition, reducing semantic inconsistency between predicted and rendered text.
+- **Spatial Mask Decoder (SMD):** adds pixel-level supervision over text regions to improve localization, geometric alignment, and layout-preserving text replacement.
+
+> This repository contains a research checkpoint that uses custom model classes from the GitHub repository. It is not directly compatible with `AutoPipeline.from_pretrained("SeerRay-Lab/Unitranslator")` or the hosted Hugging Face Inference API.
+
+## Model details
+
+| Property | Description |
+|---|---|
+| Task | End-to-end in-image machine translation |
+| Input | Source image, source language, and target language |
+| Output | Predicted translation text and an edited image containing the translated text |
+| Multimodal backbone | Qwen2.5-VL-3B-based checkpoint |
+| Image generator | FLUX.1-Kontext-dev-based denoiser |
+| Main components | UGAM and SMD |
+| Training | Two-stage warm-up and joint fine-tuning |
+| Recommended precision | BF16 |
+
+The released checkpoint combines a multimodal understanding branch, a FLUX-based generation branch, and task-specific alignment and spatial-supervision modules. Inference first autoregressively predicts the translation and then uses the resulting representation to condition image generation.
+
+## Supported and evaluated translation directions
+
+The paper evaluates the following directions:
+
+- German → English (`De → En`)
+- English → German (`En → De`)
+- French → English (`Fr → En`)
+- Romanian → English (`Ro → En`)
+- English → French (`En → Fr`)
+- English → Czech (`En → Cs`)
+- English → Russian (`En → Ru`)
+- English → Romanian (`En → Ro`)
+
+Other language directions are not guaranteed to provide comparable quality.
+
+## Checkpoint contents
+
+| Path | Purpose |
+|---|---|
+| `univa/` | Final combined UniTranslator task checkpoint, including the multimodal model, denoiser, UGAM, and SMD weights |
+| `lora/` | Rank-64 LoRA adapter used in stage-two Qwen2.5-VL fine-tuning |
+| `denoise_projector.bin` | Standalone denoise-projector/UGAM checkpoint artifact |
+| `pytorch_model/` | DeepSpeed training state for resuming training |
+| `random_states_*.pkl`, `scheduler.bin`, `latest` | Training-resume metadata |
+
+For inference, download `univa/` and `lora/`. The large DeepSpeed optimizer and random-state files are not required.
+
+## Installation
+
+```bash
+git clone https://github.com/SeerRay-Lab/Unitranslator.git
+cd Unitranslator
+
+conda create -n univa python=3.10 -y
+conda activate univa
+pip install -r requirements.txt
+pip install flash_attn --no-build-isolation
+```
+
+The reference environment uses PyTorch 2.7.1, Transformers 4.57.0, Diffusers 0.32.2, Accelerate 1.5.2, and PEFT 0.10.0.
+
+## Download
+
+Download only the files needed for inference:
+
+```bash
+hf download SeerRay-Lab/Unitranslator \
+  --include "univa/*" \
+  --include "lora/*" \
+  --local-dir checkpoints/Unitranslator
+```
+
+UniTranslator also depends on Qwen2.5-VL-3B-Instruct and FLUX.1-Kontext-dev:
+
+```bash
+hf download Qwen/Qwen2.5-VL-3B-Instruct \
+  --local-dir checkpoints/Qwen2.5-VL-3B-Instruct
+
+hf download black-forest-labs/FLUX.1-Kontext-dev \
+  --local-dir checkpoints/FLUX.1-Kontext-dev
+```
+
+FLUX.1-Kontext-dev is gated. You must first accept its license on the [model page](https://huggingface.co/black-forest-labs/FLUX.1-Kontext-dev) and authenticate with Hugging Face.
+
+## Inference
+
+### 1. Construct the base hybrid checkpoint
+
+The released inference code reconstructs the base Qwen2.5-VL + FLUX hybrid model before loading the UniTranslator task weights and LoRA adapter:
+
+```bash
+python scripts/make_univa_qwen2p5vl_tf.py \
+  --origin_qwenvl_ckpt_path checkpoints/Qwen2.5-VL-3B-Instruct \
+  --origin_flux_ckpt_path checkpoints/FLUX.1-Kontext-dev \
+  --save_path checkpoints/UniWorld_Kontext_3b_TF
+```
+
+### 2. Translate a directory of images
+
+```bash
+python infer_dir_tf.py \
+  --base_model_path checkpoints/UniWorld_Kontext_3b_TF \
+  --lora_adapter_path checkpoints/Unitranslator/lora \
+  --flux_finetune_path checkpoints/Unitranslator/univa \
+  --flux_base_path checkpoints/FLUX.1-Kontext-dev \
+  --input_dir path/to/input_images \
+  --output_dir results/de_to_en \
+  --gpu_id 0 \
+  --total_gpus 1 \
+  --source_language German \
+  --target_language English \
+  --dtype bf16 \
+  --height 1024 \
+  --width 1024 \
+  --num_inference_steps 50 \
+  --guidance_scale 5.0
+```
+
+The generated images and a JSONL file containing the predicted translations are written to `--output_dir`.
+
+For multi-GPU directory inference, launch one process per GPU with different `--gpu_id` values and the same `--total_gpus`. Input images are assigned to processes by round-robin sharding.
+
+The prompt format used by the inference script is:
+
+```text
+Translate all {source_language} texts into {target_language}.
+```
+
+### Hardware note
+
+The paper reports approximately **50 GB peak GPU memory** and **9.51 seconds per image** under its evaluation setting. Actual memory use and latency depend on resolution, precision, hardware, and inference steps. A high-memory CUDA GPU is recommended.
+
+## Training
+
+UniTranslator uses a two-stage training strategy:
+
+1. **Module warm-up:** freeze the pretrained Qwen2.5-VL and diffusion backbones and optimize the task-specific alignment and spatial modules.
+2. **Joint fine-tuning:** jointly train the understanding and generation paths, including rank-64 Qwen2.5-VL LoRA adapters, UGAM, SMD, and MMDiT attention projections.
+
+The paper reports BF16 mixed precision, AdamW, gradient checkpointing, gradient accumulation of 8, and NVIDIA H800 GPUs. See the configuration files under [`scripts/denoiser/`](https://github.com/SeerRay-Lab/Unitranslator/tree/main/scripts/denoiser) for the released training setup.
+
+The repository provides `train_stage1.sh` and `train_stage2.sh` as reference launchers. Update their configuration paths for your environment before running them; the checked-in scripts and YAML files contain project-local paths and may require adaptation.
+
+### Data preparation
+
+```bash
+# Stage-one supervision
+python convert_en_de.py
+
+# Stage-two mask supervision
+python convert_transv_mask.py
+```
+
+Datasets used by the project include:
+
+- [Translatotron-V](https://drive.google.com/drive/folders/12r54tAQ98Oxtp6Eb3dvaoiOKzu4lD_7h?usp=sharing)
+- [IIMT30k](https://huggingface.co/datasets/yztian/IIMT30k)
+- [MTedIIMT](https://huggingface.co/datasets/yztian/MTedIIMT)
+- [PRIM](https://huggingface.co/datasets/yztian/PRIM)
+
+Users are responsible for complying with the licenses and terms of each dataset.
+
+## Evaluation results
+
+All numbers below are reported in the UniTranslator paper.
+
+### Translatotron-V
+
+| Direction | BLEU ↑ | Structure-BLEU ↑ | SSIM ↑ |
+|---|---:|---:|---:|
+| De → En | **25.03** | **24.86** | **0.8184** |
+| En → De | **13.41** | **13.36** | **0.7887** |
+| Fr → En | **27.77** | **27.14** | **0.8060** |
+| Ro → En | **18.45** | **18.29** | **0.8045** |
+
+### IIMT30k test set
+
+| Direction | BLEU ↑ | COMET ↑ | FID ↓ |
+|---|---:|---:|---:|
+| De → En | **14.7** | **59.8** | **8.9** |
+| En → De | **13.0** | **45.5** | 12.5 |
+
+### PRIM
+
+| System | Average BLEU ↑ | Average COMET ↑ | Average FID ↓ |
+|---|---:|---:|---:|
+| Translatotron-V | 1.4 | 32.2 | 69.1 |
+| VisTrans | 11.3 | 47.0 | 28.8 |
+| **UniTranslator** | **12.8** | **50.7** | **22.9** |
+
+BLEU and COMET evaluate translation quality, Structure-BLEU additionally considers text-region alignment, SSIM measures source/target structural similarity, and FID evaluates generated-image distribution quality.
+
+![UniTranslator results](https://raw.githubusercontent.com/SeerRay-Lab/Unitranslator/main/assets/uni-vis.png)
+
+## Limitations and risks
+
+- Low-resource language settings may produce missing words or incorrect character rendering.
+- Highly stylized typography can lead to imperfect preservation of strokes, glow, erosion, or cursive deformation.
+- Complex backgrounds may be altered outside the intended text region, including texture, color, or local appearance changes.
+- Small, dense, curved, occluded, or low-resolution text remains challenging.
+- Results may vary for language directions and domains not represented in the evaluated datasets.
+- Generated translations and images should be verified before use in safety-critical, legal, medical, financial, or public-facing contexts.
+- Input images may contain personal or copyrighted material; users are responsible for lawful processing and distribution.
+
+## License
+
+The source code repository is released under the [Apache 2.0 License](https://github.com/SeerRay-Lab/Unitranslator/blob/main/LICENSE).
+
+The released model is based in part on [FLUX.1-Kontext-dev](https://huggingface.co/black-forest-labs/FLUX.1-Kontext-dev). Use of the model weights is therefore also subject to the [FLUX.1 Dev Non-Commercial License](https://huggingface.co/black-forest-labs/FLUX.1-Kontext-dev/blob/main/LICENSE.md) and its Acceptable Use Policy. Users must comply with all applicable upstream model and dataset licenses; the more restrictive terms apply where relevant.
+
+## Citation
+
+If you find this work useful, please cite:
+
+```bibtex
+@article{lyu2026unitranslator,
+  title={UniTranslator: A Unified Multi-modal Framework for End-to-end In-Image Machine Translation},
+  author={Lyu, Jiahao and Fu, Pei and Li, Zhenhang and Zhang, Shaojie and Yang, Jiahui and Ma, Can and Zhou, Yu and Luo, Zhenbo and Luan, Jian},
+  journal={arXiv preprint arXiv:2606.24333},
+  year={2026}
+}
+```
+
+## Acknowledgements
+
+This project builds on [UniWorld](https://github.com/PKU-YuanGroup/UniWorld-V1), [Qwen2.5-VL](https://huggingface.co/Qwen/Qwen2.5-VL-3B-Instruct), and [FLUX.1-Kontext-dev](https://huggingface.co/black-forest-labs/FLUX.1-Kontext-dev). See the paper and GitHub repository for the complete acknowledgements.