fix dataloader deadlock bug (#225)

* update README.md;

Signed-off-by: lawrence-cj <[email protected]>

* fix dataloader deadlock bug;

Signed-off-by: lawrence-cj <[email protected]>

* tmp change;

Signed-off-by: lawrence-cj <[email protected]>

* tmp change;

Signed-off-by: lawrence-cj <[email protected]>

* tmp change;

Signed-off-by: lawrence-cj <[email protected]>

* pre-commit;

Signed-off-by: lawrence-cj <[email protected]>

* Revert "tmp change;"

This reverts commit 2da18bc.

---------

Signed-off-by: lawrence-cj <[email protected]>

Author: lawrence-cj

README.md

@@ -10,8 +10,7 @@
   <a href="https://nvlabs.github.io/Sana/"><img src="https://img.shields.io/static/v1?label=Project&message=Github&color=blue&logo=github-pages"></a> &ensp;
   <a href="https://hanlab.mit.edu/projects/sana/"><img src="https://img.shields.io/static/v1?label=Page&message=MIT&color=darkred&logo=github-pages"></a> &ensp;
   <a href="https://arxiv.org/abs/2410.10629"><img src="https://img.shields.io/static/v1?label=Arxiv&message=Sana&color=red&logo=arxiv"></a> &ensp;
-  <a href="https://nv-sana.mit.edu/"><img src="https://img.shields.io/static/v1?label=Demo:5x3090&message=SANA&color=yellow"></a> &ensp;
-  <a href="https://nv-sana.mit.edu/sprint/"><img src="https://img.shields.io/static/v1?label=Demo:1x3090&message=SANA-Sprint&color=yellow"></a> &ensp;
+  <a href="https://nv-sana.mit.edu/"><img src="https://img.shields.io/static/v1?label=Demo:6x3090&message=MIT&color=yellow"></a> &ensp;
   <a href="https://nv-sana.mit.edu/4bit/"><img src="https://img.shields.io/static/v1?label=Demo:1x3090&message=4bit&color=yellow"></a> &ensp;
   <a href="https://nv-sana.mit.edu/ctrlnet/"><img src="https://img.shields.io/static/v1?label=Demo:1x3090&message=ControlNet&color=yellow"></a> &ensp;
   <a href="https://replicate.com/chenxwh/sana"><img src="https://img.shields.io/static/v1?label=API:H100&message=Replicate&color=pink"></a> &ensp;
@@ -26,7 +25,7 @@
 
 ### 🚶 Basic:
 
-**Demo**: [SANA-1.5](https://nv-sana.mit.edu/) | [SANA-ControlNet](https://nv-sana.mit.edu/ctrlnet/) | [SANA-4bit](https://nv-sana.mit.edu/4bit/) | [SANA-Sprint](https://nv-sana.mit.edu/sprint/) <br>
+**Demo**: [SANA-1.5](https://nv-sana.mit.edu/) | [SANA-ControlNet](https://nv-sana.mit.edu/ctrlnet/) | [SANA-4bit](https://nv-sana.mit.edu/4bit/) | [SANA-Sprint (Coming)](<>) <br>
 **ComfyUI**: [ComfyUI Guidance](asset/docs/ComfyUI/comfyui.md) <br>
 **Model Zoo:** [Model Card Collects All Models](asset/docs/model_zoo.md) <br>
 **Env Preparation:** [One-Click Env Install](#-1-dependencies-and-installation) <br>
@@ -183,8 +182,8 @@ cd Sana
 DEMO_PORT=15432 \
 python app/app_sana.py \
     --share \
-    --config=configs/sana_config/1024ms/Sana_1600M_img1024.yaml \
-    --model_path=hf://Efficient-Large-Model/Sana_1600M_1024px/checkpoints/Sana_1600M_1024px.pth \
+    --config=hf://Efficient-Large-Model/SANA1.5_1.6B_1024px/checkpoints/SANA1.5_1.6B_1024px.pth \
+    --model_path=hf://Efficient-Large-Model/SANA1.5_1.6B_1024px/checkpoints/SANA1.5_1.6B_1024px.pth \
     --image_size=1024
 ```
 

tests/bash/entry.sh

@@ -1,7 +1,9 @@
 #/bin/bash
 set -e
 
-for t in tests/bash/test_*.sh; do
-    echo "========================== Testing $t =================================="
-    bash $t;
-done
+
+#echo "Testing inference"
+#bash tests/bash/test_inference.sh
+
+echo "Testing training"
+bash tests/bash/test_training_1epoch.sh

train_scripts/train.py

@@ -269,13 +269,7 @@ def train(
     skip_step = max(config.train.skip_step, global_step) % train_dataloader_len
     skip_step = skip_step if skip_step < (train_dataloader_len - 20) else 0
     loss_nan_timer = 0
-
-    if config.train.use_fsdp:
-        model_instance = model
-    elif model_ema is not None:
-        model_instance = model_ema
-    else:
-        model_instance = model
+    model_instance.to(accelerator.device)
 
     # Cache Dataset for BatchSampler
     if args.caching and config.model.multi_scale:
@@ -542,9 +536,11 @@ def train(
                     merged_state_dict = accelerator.get_state_dict(model)
 
                 accelerator.wait_for_everyone()
+                print(rank, 111111)
                 if accelerator.is_main_process:
                     if config.train.use_fsdp:
                         model_instance.load_state_dict(merged_state_dict)
+                    print(rank, 222222)
                     if validation_noise is not None:
                         log_validation(
                             accelerator=accelerator,
@@ -567,6 +563,7 @@ def train(
                             vae=vae,
                         )
 
+                print(rank, 333333)
             # avoid dead-lock of multiscale data batch sampler
             if (
                 config.model.multi_scale
@@ -629,7 +626,7 @@ def main(cfg: SanaConfig) -> None:
     global train_dataloader_len, start_epoch, start_step, vae, generator, num_replicas, rank, training_start_time
     global load_vae_feat, load_text_feat, validation_noise, text_encoder, tokenizer
     global max_length, validation_prompts, latent_size, valid_prompt_embed_suffix, null_embed_path
-    global image_size, cache_file, total_steps, vae_dtype
+    global image_size, cache_file, total_steps, vae_dtype, model_instance
 
     config = cfg
     args = cfg
@@ -870,6 +867,13 @@ def main(cfg: SanaConfig) -> None:
         )
     )
 
+    if config.train.use_fsdp:
+        model_instance = deepcopy(model)
+    elif model_ema is not None:
+        model_instance = deepcopy(model_ema)
+    else:
+        model_instance = model
+
     # 4-1. load model
     if args.load_from is not None:
         config.model.load_from = args.load_from