fix: format

app.py

@@ -31,6 +31,7 @@ import torch
 from torchvision.transforms.functional import to_pil_image, to_tensor
 
 from accelerate import Accelerator
+from huggingface_hub import hf_hub_download
 from peft import LoraConfig
 from safetensors.torch import load_file
 
@@ -41,7 +42,6 @@ from omnigen2.schedulers.scheduling_dpmsolver_multistep import DPMSolverMultiste
 from omnigen2.utils.img_util import create_collage
 
 NEGATIVE_PROMPT = "(((deformed))), blurry, over saturation, bad anatomy, disfigured, poorly drawn face, mutation, mutated, (extra_limb), (ugly), (poorly drawn hands), fused fingers, messy drawing, broken legs censor, censored, censor_bar"
-ROOT_DIR = "projects/OmniGen2"
 SAVE_DIR = "output/gradio"
 
 pipeline = None
@@ -59,20 +59,22 @@ def load_pipeline(accelerator, weight_dtype, args):
         subfolder="transformer",
         torch_dtype=weight_dtype,
     )
-    if args.lora_path is not None:
-        target_modules = ["to_k", "to_q", "to_v", "to_out.0"]
-        lora_config = LoraConfig(
-            r=512,
-            lora_alpha=512,
-            lora_dropout=0,
-            init_lora_weights="gaussian",
-            target_modules=target_modules,
-        )
-        pipeline.transformer.add_adapter(lora_config)
-        lora_state_dict = load_file(args.lora_path, device=accelerator.device.__str__())
-        pipeline.transformer.load_state_dict(lora_state_dict, strict=False)
-        pipeline.transformer.fuse_lora(lora_scale=1, safe_fusing=False, adapter_names=["default"])
-        pipeline.transformer.unload_lora()
+    
+    lora_path = hf_hub_download("bytedance-research/UMO", "UMO_OmniGen2.safetensors") if args.lora_path is None else args.lora_path
+    target_modules = ["to_k", "to_q", "to_v", "to_out.0"]
+    lora_config = LoraConfig(
+        r=512,
+        lora_alpha=512,
+        lora_dropout=0,
+        init_lora_weights="gaussian",
+        target_modules=target_modules,
+    )
+    pipeline.transformer.add_adapter(lora_config)
+    lora_state_dict = load_file(lora_path, device=accelerator.device.__str__())
+    pipeline.transformer.load_state_dict(lora_state_dict, strict=False)
+    pipeline.transformer.fuse_lora(lora_scale=1, safe_fusing=False, adapter_names=["default"])
+    pipeline.transformer.unload_lora()
+    
     if args.enable_sequential_cpu_offload:
         pipeline.enable_sequential_cpu_offload()
     elif args.enable_model_cpu_offload:
@@ -153,6 +155,7 @@ def run(
     vis_images = [to_tensor(image) * 2 - 1 for image in results.images]
     output_image = create_collage(vis_images)
 
+    output_path = ""
     if save_images:
         # Create outputs directory if it doesn't exist
         output_dir = SAVE_DIR
@@ -171,7 +174,7 @@ def run(
             for i, image in enumerate(results.images):
                 image_name, ext = os.path.splitext(output_path)
                 image.save(f"{image_name}_{i}{ext}")
-    return output_image
+    return output_image, output_path
 
 
 def get_examples(base_dir="assets/examples/OmniGen2"):
@@ -199,9 +202,8 @@ 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>
 </div>
 """.strip()
 
@@ -227,11 +229,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()
@@ -384,9 +389,11 @@ def main(args):
                     # output image
                     output_image = gr.Image(label="Output Image")
                     global save_images
-                    save_images = gr.Checkbox(label="Save generated images", value=True)
+                    # save_images = gr.Checkbox(label="Save generated images", value=True)
+                    save_images = True
                     with gr.Accordion("Examples Comparison with OmniGen2", open=False):
                         output_image_omnigen2 = gr.Image(label="Generated Image (OmniGen2)")
+                    download_btn = gr.File(label="Download full-resolution", type="filepath", interactive=False)
 
         gr.Markdown(star)
         
@@ -422,7 +429,7 @@ def main(args):
                 seed_input,
                 align_res,
             ],
-            outputs=output_image,
+            outputs=[output_image, download_btn],
         )
 
         gr.Examples(
@@ -444,7 +451,7 @@ def main(args):
         )
 
     # launch
-    demo.launch(share=args.share, server_port=args.port, allowed_paths=[ROOT_DIR], server_name=args.server_name)
+    demo.launch(share=args.share, server_port=args.port, server_name=args.server_name, ssr_mode=False)
 
 def parse_args():
     parser = argparse.ArgumentParser()