show parameter memory usage, upgrad diffusers to 0.31

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"{parallel_info}: 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/sd3_example.py

@@ -24,6 +24,9 @@ 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()
@@ -63,7 +66,7 @@ def main():
 
     if get_world_group().rank == get_world_group().world_size - 1:
         print(
-            f"{parallel_info} epoch time: {elapsed_time:.2f} sec, memory: {peak_memory/1e9} GB"
+            f"{parallel_info} epoch time: {elapsed_time:.2f} sec, overall memory: {peak_memory/1e9:.2f} GB, model memory: {model_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.dev is necessary for CogVideoX and Flux
             "transformers>=4.39.1",
             "sentencepiece>=0.1.99",
             "beautifulsoup4>=4.12.3",