| 476 | |
| 477 | |
| 478 | class WanVaceMP(WanVace): |
| 479 | |
| 480 | def __init__(self, |
| 481 | config, |
| 482 | checkpoint_dir, |
| 483 | use_usp=False, |
| 484 | ulysses_size=None, |
| 485 | ring_size=None): |
| 486 | self.config = config |
| 487 | self.checkpoint_dir = checkpoint_dir |
| 488 | self.use_usp = use_usp |
| 489 | os.environ['MASTER_ADDR'] = 'localhost' |
| 490 | os.environ['MASTER_PORT'] = '12345' |
| 491 | os.environ['RANK'] = '0' |
| 492 | os.environ['WORLD_SIZE'] = '1' |
| 493 | self.in_q_list = None |
| 494 | self.out_q = None |
| 495 | self.inference_pids = None |
| 496 | self.ulysses_size = ulysses_size |
| 497 | self.ring_size = ring_size |
| 498 | self.dynamic_load() |
| 499 | |
| 500 | self.device = 'cpu' if torch.cuda.is_available() else 'cpu' |
| 501 | self.vid_proc = VaceVideoProcessor( |
| 502 | downsample=tuple( |
| 503 | [x * y for x, y in zip(config.vae_stride, config.patch_size)]), |
| 504 | min_area=480 * 832, |
| 505 | max_area=480 * 832, |
| 506 | min_fps=self.config.sample_fps, |
| 507 | max_fps=self.config.sample_fps, |
| 508 | zero_start=True, |
| 509 | seq_len=32760, |
| 510 | keep_last=True) |
| 511 | |
| 512 | def dynamic_load(self): |
| 513 | if hasattr(self, 'inference_pids') and self.inference_pids is not None: |
| 514 | return |
| 515 | gpu_infer = os.environ.get( |
| 516 | 'LOCAL_WORLD_SIZE') or torch.cuda.device_count() |
| 517 | pmi_rank = int(os.environ['RANK']) |
| 518 | pmi_world_size = int(os.environ['WORLD_SIZE']) |
| 519 | in_q_list = [ |
| 520 | torch.multiprocessing.Manager().Queue() for _ in range(gpu_infer) |
| 521 | ] |
| 522 | out_q = torch.multiprocessing.Manager().Queue() |
| 523 | initialized_events = [ |
| 524 | torch.multiprocessing.Manager().Event() for _ in range(gpu_infer) |
| 525 | ] |
| 526 | context = mp.spawn( |
| 527 | self.mp_worker, |
| 528 | nprocs=gpu_infer, |
| 529 | args=(gpu_infer, pmi_rank, pmi_world_size, in_q_list, out_q, |
| 530 | initialized_events, self), |
| 531 | join=False) |
| 532 | all_initialized = False |
| 533 | while not all_initialized: |
| 534 | all_initialized = all( |
| 535 | event.is_set() for event in initialized_events) |
nothing calls this directly
no outgoing calls
no test coverage detected