| 27 | |
| 28 | |
| 29 | class WanT2V: |
| 30 | |
| 31 | def __init__( |
| 32 | self, |
| 33 | config, |
| 34 | checkpoint_dir, |
| 35 | device_id=0, |
| 36 | rank=0, |
| 37 | t5_fsdp=False, |
| 38 | dit_fsdp=False, |
| 39 | use_usp=False, |
| 40 | t5_cpu=False, |
| 41 | ): |
| 42 | r""" |
| 43 | Initializes the Wan text-to-video generation model components. |
| 44 | |
| 45 | Args: |
| 46 | config (EasyDict): |
| 47 | Object containing model parameters initialized from config.py |
| 48 | checkpoint_dir (`str`): |
| 49 | Path to directory containing model checkpoints |
| 50 | device_id (`int`, *optional*, defaults to 0): |
| 51 | Id of target GPU device |
| 52 | rank (`int`, *optional*, defaults to 0): |
| 53 | Process rank for distributed training |
| 54 | t5_fsdp (`bool`, *optional*, defaults to False): |
| 55 | Enable FSDP sharding for T5 model |
| 56 | dit_fsdp (`bool`, *optional*, defaults to False): |
| 57 | Enable FSDP sharding for DiT model |
| 58 | use_usp (`bool`, *optional*, defaults to False): |
| 59 | Enable distribution strategy of USP. |
| 60 | t5_cpu (`bool`, *optional*, defaults to False): |
| 61 | Whether to place T5 model on CPU. Only works without t5_fsdp. |
| 62 | """ |
| 63 | self.device = torch.device(f"cuda:{device_id}") |
| 64 | self.config = config |
| 65 | self.rank = rank |
| 66 | self.t5_cpu = t5_cpu |
| 67 | |
| 68 | self.num_train_timesteps = config.num_train_timesteps |
| 69 | self.param_dtype = config.param_dtype |
| 70 | |
| 71 | shard_fn = partial(shard_model, device_id=device_id) |
| 72 | self.text_encoder = T5EncoderModel( |
| 73 | text_len=config.text_len, |
| 74 | dtype=config.t5_dtype, |
| 75 | device=torch.device('cpu'), |
| 76 | checkpoint_path=os.path.join(checkpoint_dir, config.t5_checkpoint), |
| 77 | tokenizer_path=os.path.join(checkpoint_dir, config.t5_tokenizer), |
| 78 | shard_fn=shard_fn if t5_fsdp else None) |
| 79 | |
| 80 | self.vae_stride = config.vae_stride |
| 81 | self.patch_size = config.patch_size |
| 82 | self.vae = WanVAE( |
| 83 | vae_pth=os.path.join(checkpoint_dir, config.vae_checkpoint), |
| 84 | device=self.device) |
| 85 | |
| 86 | logging.info(f"Creating WanModel from {checkpoint_dir}") |
nothing calls this directly
no outgoing calls
no test coverage detected