upgrade diffusers to 0.31 (#316)

benchmark/run.sh

@@ -6,13 +6,19 @@ set -x
 # MODEL="/mnt/models/SD/stable-diffusion-3-medium-diffusers"
 # SCRIPT="./examples/sd3_example.py"
 
-MODEL="/mnt/models/SD/HunyuanDiT-v1.2-Diffusers"
-SCRIPT="./examples/hunyuandit_example.py"
+# MODEL="/mnt/models/SD/HunyuanDiT-v1.2-Diffusers"
+# SCRIPT="./examples/hunyuandit_example.py"
+
+MODEL="/cfs/dit/FLUX.1-dev/"
+SCRIPT="./examples/flux_example.py"
 
 export PYTHONPATH=$PWD:$PYTHONPATH
 
 python benchmark/single_node_latency_test.py \
 --model_id $MODEL \
 --script $SCRIPT \
 --sizes 1024 \
---no_use_resolution_binning
+--no_use_resolution_binning \
+--num_inference_steps 28 \
+--no_use_cfg_parallel \
+--n_gpus 4

examples/pixartalpha_example.py

@@ -24,6 +24,7 @@ def main():
         engine_config=engine_config,
         torch_dtype=torch.float16,
     ).to(f"cuda:{local_rank}")
+    model_memory = torch.cuda.max_memory_allocated(device=f"cuda:{local_rank}")
     pipe.prepare_run(input_config)
 
     torch.cuda.reset_peak_memory_stats()
@@ -62,7 +63,9 @@ def main():
                 print(img_file)
 
     if get_world_group().rank == get_world_group().world_size - 1:
-        print(f"epoch time: {elapsed_time:.2f} sec, memory: {peak_memory/1e9} GB")
+        print(
+            f"epoch time: {elapsed_time:.2f} sec, model memory: {model_memory/1e9:.2f} GB, overall memory: {peak_memory/1e9:.2f} GB"
+        )
     get_runtime_state().destory_distributed_env()
 
 

examples/run.sh

@@ -27,8 +27,10 @@ mkdir -p ./results
 # task args
 TASK_ARGS="--height 1024 --width 1024 --no_use_resolution_binning"
 
+
+# On 8 gpus, pp=2, ulysses=2, ring=1, cfg_parallel=2 (split batch)
 N_GPUS=8
-PARALLEL_ARGS="--ulysses_degree 1 --ring_degree 1 --pipefusion_parallel_degree 8"
+PARALLEL_ARGS="--pipefusion_parallel_degree 2 --ulysses_degree 2 --ring_degree 2"
 
 # CFG_ARGS="--use_cfg_parallel"
 
@@ -43,14 +45,17 @@ PARALLEL_ARGS="--ulysses_degree 1 --ring_degree 1 --pipefusion_parallel_degree 8
 # Another compile option is `--use_onediff` which will use onediff's compiler.
 # COMPILE_FLAG="--use_torch_compile"
 
+
+# export CUDA_VISIBLE_DEVICES=4,5,6,7
+
 torchrun --nproc_per_node=$N_GPUS ./examples/$SCRIPT \
 --model $MODEL_ID \
 $PARALLEL_ARGS \
 $TASK_ARGS \
 $PIPEFUSION_ARGS \
 $OUTPUT_ARGS \
 --num_inference_steps $INFERENCE_STEP \
---warmup_steps 0 \
+--warmup_steps 1 \
 --prompt "brown dog laying on the ground with a metal bowl in front of him." \
 $CFG_ARGS \
 $PARALLLEL_VAE \

examples/sd3_example.py

@@ -66,8 +66,9 @@ def main():
 
     if get_world_group().rank == get_world_group().world_size - 1:
         print(
-            f"{parallel_info} epoch time: {elapsed_time:.2f} sec, parameter memory: {parameter_peak_memory/1e9:.2f} GB, peak memory: {peak_memory/1e9:.2f} GB"
+            f"epoch time: {elapsed_time:.2f} sec, parameter memory: {parameter_peak_memory/1e9:.2f} GB, peak memory: {peak_memory/1e9:.2f} GB"
         )
+
     get_runtime_state().destory_distributed_env()
 
 

setup.py

@@ -28,7 +28,7 @@ def get_cuda_version():
         install_requires=[
             "torch>=2.1.0",
             "accelerate>=0.33.0",
-            "diffusers@git+https://github.com/huggingface/diffusers",  # NOTE: diffusers>=0.31.0.dev is necessary for CogVideoX and Flux
+            "diffusers>=0.31",  # NOTE: diffusers>=0.31.0 is necessary for CogVideoX and Flux
             "transformers>=4.39.1",
             "sentencepiece>=0.1.99",
             "beautifulsoup4>=4.12.3",

xfuser/__version__.py

@@ -1 +1 @@
-__version__ = "0.3.2"
+__version__ = "0.3.3"