Zengtudor/Wan2.1
1
0
Fork 0
mirror of https://github.com/Wan-Video/Wan2.1.git synced 2025-11-04 06:15:17 +00:00
Code Issues Actions Packages Projects Releases Wiki Activity

Remove torch cuda memory profiler

Browse Source
This commit is contained in:
wangang.wa 2025-05-16 11:54:34 +08:00
parent b8c5404b65
commit 98ccd16d3d
1 changed files with 0 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
generate.py
View File

@ -497,7 +497,6 @@ def generate(args):
seed=args.base_seed, seed=args.base_seed,
offload_model=args.offload_model) offload_model=args.offload_model)
elif "vace" in args.task: elif "vace" in args.task:
torch.cuda.memory._record_memory_history(max_entries=1000000)
if args.prompt is None: if args.prompt is None:
args.prompt = EXAMPLE_PROMPT[args.task]["prompt"] args.prompt = EXAMPLE_PROMPT[args.task]["prompt"]
args.src_video = EXAMPLE_PROMPT[args.task].get("src_video", None) args.src_video = EXAMPLE_PROMPT[args.task].get("src_video", None)
@ -580,8 +579,6 @@ def generate(args):
nrow=1, nrow=1,
normalize=True, normalize=True,
value_range=(-1, 1)) value_range=(-1, 1))
torch.cuda.memory._record_memory_history(enabled=None)
torch.cuda.memory._dump_snapshot(f"memory.pickle")
logging.info("Finished.") logging.info("Finished.")