feat:support t2i

app.py

@@ -72,9 +72,9 @@ badges_text = r"""
 <div style="text-align: center; display: flex; justify-content: center; gap: 5px;">
 <a href="https://github.com/bytedance/UMO"><img alt="Build" src="https://img.shields.io/github/stars/bytedance/UMO"></a> 
 <a href="https://bytedance.github.io/UMO/"><img alt="Build" src="https://img.shields.io/badge/Project%20Page-UMO-yellow"></a> 
-<a href="https://arxiv.org/abs/25xx.xxxxx"><img alt="Build" src="https://img.shields.io/badge/arXiv%20paper-UMO-b31b1b.svg"></a>
+<a href="https://arxiv.org/abs/2509.06818"><img alt="Build" src="https://img.shields.io/badge/arXiv%20paper-UMO-b31b1b.svg"></a>
 <a href="https://huggingface.co/bytedance-research/UMO"><img src="https://img.shields.io/static/v1?label=%F0%9F%A4%97%20Hugging%20Face&message=Model&color=orange"></a>
-<a href="https://huggingface.co/spaces/bytedance-research/UMO-FLUX"><img src="https://img.shields.io/static/v1?label=%F0%9F%A4%97%20Hugging%20Face&message=demo&color=orange"></a>
+# <a href="https://huggingface.co/spaces/bytedance-research/UMO-FLUX"><img src="https://img.shields.io/static/v1?label=%F0%9F%A4%97%20Hugging%20Face&message=demo&color=orange"></a>
 </div>
 """.strip()
 
@@ -95,11 +95,14 @@ tips = """
 
 article = """
 ```bibtex
-@article{cheng2025umo,
-  title={UMO: Scaling Multi-Identity Consistency for Image Customization via Matching Reward},
-  author={Cheng, Yufeng and Wu, Wenxu and Wu, Shaojin and Huang, Mengqi and Ding, Fei and He, Qian},
-  journal={arXiv preprint arXiv:25xx.xxxxx},
-  year={2025}
+@misc{cheng2025umoscalingmultiidentityconsistency,
+      title={UMO: Scaling Multi-Identity Consistency for Image Customization via Matching Reward}, 
+      author={Yufeng Cheng and Wenxu Wu and Shaojin Wu and Mengqi Huang and Fei Ding and Qian He},
+      year={2025},
+      eprint={2509.06818},
+      archivePrefix={arXiv},
+      primaryClass={cs.CV},
+      url={https://arxiv.org/abs/2509.06818}, 
 }
 ```
 """.strip()

uno/flux/model.py

@@ -177,8 +177,9 @@ class Flux(nn.Module):
         img_end = img.shape[1]
         if ref_img is not None:
             if isinstance(ref_img, tuple) or isinstance(ref_img, list):
-                img = torch.cat([img, self.img_in(torch.cat(ref_img, dim=1))], dim=1)
+                img_in = [img] + [self.img_in(ref) for ref in ref_img]
                 img_ids = [ids] + [ref_ids for ref_ids in ref_img_ids]
+                img = torch.cat(img_in, dim=1)
                 ids = torch.cat(img_ids, dim=1)
             else:
                 img = torch.cat((img, self.img_in(ref_img)), dim=1)