Upload folder using huggingface_hub

.gitignore

@@ -1,5 +1,6 @@
 examples
 data
+checkpoints
 .cursorignore
 .specstory
 
@@ -69,6 +70,7 @@ ttt*
 
 *.pt
 *.pth
+*.ckpt
 
 
 *.pkl

README.md

@@ -1,8 +1,8 @@
 # UniBioTransfer
 
 
-### [Paper](https://arxiv.org/abs/2603.19637)
-### [Project page](https://scy639.github.io/UniBioTransfer.github.io/)
+[![Paper](https://img.shields.io/badge/cs.CV-Paper-b31b1b?logo=arxiv&logoColor=white)](https://arxiv.org/abs/2603.19637)
+[![Project Page](https://img.shields.io/badge/Project-Page-blue?logo=googlechrome&logoColor=white)](https://scy639.github.io/UniBioTransfer.github.io/)
 
 ## Abstract
 >Deepface generation has traditionally followed a task-driven paradigm, where distinct tasks (e.g., face transfer and hair transfer) are addressed by task-specific models. Nevertheless, this single-task setting severely limits model generalization and scalability. A unified model capable of solving multiple deepface generation tasks in a single pass represents a promising and practical direction, yet remains challenging due to data scarcity and cross-task conflicts arising from heterogeneous attribute transformations. To this end, we propose UniBioTransfer, the first unified framework capable of handling both conventional deepface tasks (e.g., face transfer and face reenactment) and shape-varying transformations (e.g., hair transfer and head transfer). Besides, UniBioTransfer naturally generalizes to unseen tasks, like lip, eye, and glasses transfer, with minimal fine-tuning. Generally, UniBioTransfer addresses data insufficiency in multi-task generation through a unified data construction strategy, including a swapping-based corruption mechanism designed for spatially dynamic attributes like hair. It further mitigates cross-task interference via an innovative BioMoE, a mixture-of-experts based model coupled with a novel two-stage training strategy that effectively disentangles task-specific knowledge. Extensive experiments demonstrate the effectiveness, generalization, and scalability of UniBioTransfer, outperforming both existing unified models and task-specific methods across a wide range of deepface generation tasks.
@@ -19,6 +19,8 @@
 
 
 ## Setup
+
+### 1. Environment setup
 ```
 conda create -n "unibio" python=3.10.16 -y
 conda activate unibio
@@ -26,11 +28,22 @@ sh setup.sh
 ```
 
 
+### 2. Download pre-trained weight
 Download pre-trained checkpoints via `python download_checkpoints.py`
 
 
 
 
+## Inference
+
+Minimum **11GB** VRAM required.
+
+```bash
+python infer.py --task-name=face   --tgt=examples/face/tgt.png   --ref=examples/face/ref.png   --out-dir=examples/face
+python infer.py --task-name=hair   --tgt=examples/hair/tgt.png   --ref=examples/hair/ref.png   --out-dir=examples/hair
+python infer.py --task-name=motion --tgt=examples/motion/tgt.png --ref=examples/motion/ref.png --out-dir=examples/motion
+python infer.py --task-name=head   --tgt=examples/head/tgt.png   --ref=examples/head/ref.png   --out-dir=examples/head
+```
 
 ## Citing Us
 If you find our work valuable, we kindly ask you to consider citing our paper.

download_checkpoints.py

@@ -5,8 +5,7 @@ from confs import *
 
 
 def _download(repo_id, filename, local_path: Path) -> Path:
-    if local_path.exists():
-        return local_path
+    local_path = Path(local_path)
     from huggingface_hub import hf_hub_download
     local_path.parent.mkdir(parents=True, exist_ok=True)
     token = os.getenv("HF_TOKEN") or os.getenv("HUGGINGFACE_HUB_TOKEN")
@@ -22,5 +21,5 @@ def _download(repo_id, filename, local_path: Path) -> Path:
 
 
 _download("CompVis/stable-diffusion-v-1-4-original",SD14_filename, SD14_localpath)
-_download("scy639/UniBioTransfer",PRETRAIN_CKPT_PATH, PRETRAIN_CKPT_PATH)
-_download("scy639/UniBioTransfer",PRETRAIN_JSON_PATH, PRETRAIN_JSON_PATH)
+_download("scy639/UniBioTransfer",PRETRAIN_CKPT_PATH, ".")
+_download("scy639/UniBioTransfer",PRETRAIN_JSON_PATH, ".")

util_and_constant.py

@@ -106,8 +106,8 @@ _gate_k2tu = { # id 2 (max_run,interval,prob)
     'vis-ffn_gate_input'            : ( 3, 3, None ),
     '[warning]: no param to sync'   : (10,1, None ),#-------------TP-----------
     '[TP] shared sync counts'       : (10,1, None ),
-    '[Conv2d param stats] count, name (sorted desc):': (1,1, None ),#-------------upcycle-----------
-    'avg full_name='                : ( 7, 1, None ),
+    '[Conv2d param stats] count, name (sorted desc):': (0,1, None ),#-------------upcycle-----------
+    'avg full_name='                : ( 0, 1, None ),
 }
 def gate_(id_, *args, **kw, ): # gate for some vis or print behaviour, just for vis/debug
     # return 0