Avoid YuEGP pinned-memory exhaustion on 32GB Windows hosts
This commit is contained in:
+21
-2
@@ -61,6 +61,16 @@ def load_codec_on_cpu(ns, torch):
|
||||
return codec
|
||||
|
||||
|
||||
def memory_profile_options(profile, compile_enabled, platform, total_ram):
|
||||
options = dict(profile_no=profile, quantizeTransformer=profile == 3,
|
||||
compile=compile_enabled, verboseLevel=1)
|
||||
# Windows CUDA pinned allocations can fail before inference on 32GB hosts.
|
||||
# This affects transfer staging only, not weight precision or VRAM budgets.
|
||||
if platform == 'win32' and total_ram < 48 * 1024 ** 3:
|
||||
options['pinnedMemory'] = False
|
||||
return options
|
||||
|
||||
|
||||
def main():
|
||||
import psutil
|
||||
parent = psutil.Process(os.getppid())
|
||||
@@ -126,8 +136,16 @@ def main():
|
||||
model2.generation_config.cache_implementation = None
|
||||
model._validate_model_kwargs = lambda _: None
|
||||
model2._validate_model_kwargs = lambda _: None
|
||||
offloader = ns['offload'].profile({'transformer': model, 'stage2': model2}, profile_no=cli.profile,
|
||||
quantizeTransformer=cli.profile == 3, compile=compile_enabled, verboseLevel=1)
|
||||
profile_options = memory_profile_options(cli.profile, compile_enabled, sys.platform, psutil.virtual_memory().total)
|
||||
emit(stage='loading', message='Preparing YuEGP memory', profileOptions=profile_options)
|
||||
offloader = ns['offload'].profile({'transformer': model, 'stage2': model2}, **profile_options)
|
||||
def memory_checkpoint(stage):
|
||||
torch.cuda.synchronize()
|
||||
free, total = torch.cuda.mem_get_info()
|
||||
emit(stage=stage, message=f'YuEGP {stage}', gpuFreeMiB=round(free / 1024 ** 2),
|
||||
gpuTotalMiB=round(total / 1024 ** 2), gpuAllocatedMiB=round(torch.cuda.memory_allocated() / 1024 ** 2),
|
||||
ramAvailableMiB=round(psutil.virtual_memory().available / 1024 ** 2))
|
||||
memory_checkpoint('memory-ready')
|
||||
args = SimpleNamespace(use_audio_prompt=False, use_dual_tracks_prompt=False, rescale=True,
|
||||
output_dir=str(output), cuda_idx=0)
|
||||
ns.update(model=model, model_stage2=model2, device=device, codec_model=None, args=args,
|
||||
@@ -150,6 +168,7 @@ def main():
|
||||
progress=round(100 * done / max(1, total), 1))
|
||||
emit(stage='stage1', message='Generating song tokens', progress=0)
|
||||
stems = ns['stage1_inference'](tags, lyrics, 1, max_tokens, seed, state, callback)
|
||||
memory_checkpoint('stage1-finished')
|
||||
emit(stage='stage2', message='Generating audio detail', progress=0)
|
||||
results = ns['stage2_inference'](model2, stems, str(output / 'stage2'),
|
||||
batch_size=20 if cli.profile == 1 else 4, state=state, callback=callback)
|
||||
|
||||
Reference in New Issue
Block a user