Coverage for wrapper/vibevoice/modular_vibevoice_diffusion_head.py: 33%

1# mypy: ignore-errors 

2# Copy from https://github.com/microsoft/VibeVoice/blob/main/vibevoice/modular/modular_vibevoice_diffusion_head.py 

3 

4import math 

5from typing import cast 

6 

7import torch 

8import torch.nn as nn 

9 

10from transformers import AutoModel 

11from transformers.modeling_utils import PreTrainedModel 

12from transformers.activations import ACT2FN 

13from transformers.utils import logging 

14 

15from configuration_vibevoice import VibeVoiceDiffusionHeadConfig 

16 

17 

18logger = logging.get_logger(__name__) 

19 

20 

21class RMSNorm(nn.Module): 

22 def __init__( 

23 self, 

24 dim: int, 

25 eps: float = 1e-6, 

26 elementwise_affine: bool = True, 

27 memory_efficient: bool = False 

28 ) -> None: 

29 super().__init__() 

30 self.dim = dim 

31 self.eps = eps 

32 self.elementwise_affine = elementwise_affine 

33 if self.elementwise_affine: 

34 self.weight = nn.Parameter(torch.ones(dim)) 

35 else: 

36 self.register_parameter('weight', None) 

37 

38 def _norm( 

39 self, 

40 x: torch.Tensor 

41 ) -> torch.Tensor: 

42 return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps) 

43 

44 def forward( 

45 self, 

46 x: torch.Tensor 

47 ) -> torch.Tensor: 

48 output = self._norm(x.float()).type_as(x) 

49 if self.weight is not None: 

50 output = output * self.weight 

51 return output 

52 

53 def extra_repr(self) -> str: 

54 return f'dim={self.dim}, eps={self.eps}, elementwise_affine={self.elementwise_affine}' 

55 

56 

57def modulate( 

58 x: torch.Tensor, 

59 shift: torch.Tensor, 

60 scale: torch.Tensor, 

61) -> torch.Tensor: 

62 """Apply modulation to input tensor.""" 

63 return x * (1 + scale) + shift 

64 

65 

66class TimestepEmbedder(nn.Module): 

67 """ 

68 Embeds scalar timesteps into vector representations. 

69 Args: 

70 hidden_size (`int`): Size of the output embedding 

71 frequency_embedding_size (`int`, optional): Size of the intermediate frequency embedding 

72 """ 

73 def __init__(self, hidden_size: int, frequency_embedding_size: int = 256) -> None: 

74 super().__init__() 

75 self.mlp = nn.Sequential( 

76 nn.Linear(frequency_embedding_size, hidden_size, bias=False), 

77 # nn.SiLU(), 

78 ACT2FN['silu'], 

79 nn.Linear(hidden_size, hidden_size, bias=False), 

80 ) 

81 self.frequency_embedding_size = frequency_embedding_size 

82 

83 @staticmethod 

84 def timestep_embedding(t: torch.Tensor, dim: int, max_period: int = 10000) -> torch.Tensor: 

85 """ 

86 Create sinusoidal timestep embeddings. 

87 Args: 

88 t (`torch.Tensor`): A 1-D Tensor of N indices, one per batch element. 

89 These may be fractional. 

90 dim (`int`): The dimension of the output. 

91 max_period (`int`, optional): Controls the minimum frequency of the embeddings. 

92 Returns: 

93 `torch.Tensor`: An [N, D] Tensor of positional embeddings. 

94 """ 

95 half = dim // 2 

96 freqs = torch.exp( 

97 -math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32) / half 

98 ).to(t.device) 

99 args = t[:, None].float() * freqs[None] 

100 embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1) 

101 if dim % 2: 

102 embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1) 

103 return embedding.to(t.dtype) 

104 

105 def forward(self, t: torch.Tensor) -> torch.Tensor: 

106 t_freq = self.timestep_embedding(t, self.frequency_embedding_size) 

107 t_emb = self.mlp(t_freq) 

108 return t_emb 

109 

110 

111class FeedForwardNetwork(nn.Module): 

112 """ 

113 Standard feed-forward network with SwiGLU activation. 

114 Args: 

115 embed_dim (`int`): Input dimension 

116 ffn_dim (`int`): Hidden dimension 

117 """ 

118 def __init__( 

119 self, 

120 embed_dim: int, 

121 ffn_dim: int, 

122 ) -> None: 

123 super().__init__() 

124 self.embed_dim = embed_dim 

125 self.gate_proj = nn.Linear(self.embed_dim, ffn_dim, bias=False) 

126 self.up_proj = nn.Linear(self.embed_dim, ffn_dim, bias=False) 

127 self.down_proj = nn.Linear(ffn_dim, self.embed_dim, bias=False) 

128 self.act_fn = ACT2FN['silu'] # Using SiLU as the activation function 

129 

130 def forward(self, x: torch.Tensor) -> torch.Tensor: 

131 gate = self.gate_proj(x) 

132 up = self.up_proj(x) 

133 

134 # SwiGLU activation 

135 # gate = F.silu(gate) 

136 gate = self.act_fn(gate) 

137 return self.down_proj(gate * up) 

138 

139 

140class HeadLayer(nn.Module): 

141 """ 

142 A layer in the diffusion head. 

143 Args: 

144 embed_dim (`int`): Input dimension 

145 ffn_dim (`int`): Hidden dimension 

146 cond_dim (`int`): Condition embedding dimension 

147 norm_eps (`float`, optional): Epsilon for normalization 

148 """ 

149 def __init__( 

150 self, 

151 embed_dim: int, 

152 ffn_dim: int, 

153 cond_dim: int, 

154 norm_eps: float = 1e-5, 

155 ) -> None: 

156 super().__init__() 

157 self.embed_dim = embed_dim 

158 self.cond_dim = cond_dim 

159 self.ffn_dim = ffn_dim 

160 self.ffn = FeedForwardNetwork( 

161 self.embed_dim, 

162 self.ffn_dim, 

163 ) 

164 self.norm = RMSNorm(self.embed_dim, eps=norm_eps) 

165 self.adaLN_modulation = nn.Sequential( 

166 # nn.SiLU(), 

167 ACT2FN['silu'], 

168 nn.Linear(cond_dim, 3 * self.embed_dim, bias=False) 

169 ) 

170 

171 def forward(self, x: torch.Tensor, c: torch.Tensor) -> torch.Tensor: 

172 shift_ffn, scale_ffn, gate_ffn = self.adaLN_modulation(c).chunk(3, dim=-1) 

173 x = x + gate_ffn * self.ffn(modulate(self.norm(x), shift_ffn, scale_ffn)) 

174 return x 

175 

176 

177class FinalLayer(nn.Module): 

178 """ 

179 Final layer in the diffusion head. 

180 Args: 

181 hidden_size (`int`): Input dimension 

182 output_size (`int`): Output dimension 

183 cond_size (`int`): Condition embedding dimension 

184 norm_eps (`float`, optional): Epsilon for normalization 

185 """ 

186 def __init__(self, hidden_size: int, output_size: int, cond_size: int, norm_eps: float = 1e-5) -> None: 

187 super().__init__() 

188 self.norm_final = RMSNorm(hidden_size, eps=norm_eps, elementwise_affine=False) 

189 self.linear = nn.Linear(hidden_size, output_size, bias=False) 

190 self.adaLN_modulation = nn.Sequential( 

191 # nn.SiLU(), 

192 ACT2FN['silu'], 

193 nn.Linear(cond_size, 2 * hidden_size, bias=False) 

194 ) 

195 

196 def forward(self, x: torch.Tensor, c: torch.Tensor) -> torch.Tensor: 

197 shift, scale = self.adaLN_modulation(c).chunk(2, dim=-1) 

198 x = modulate(self.norm_final(x), shift, scale) 

199 x = self.linear(x) 

200 return x 

201 

202 

203class VibeVoiceDiffusionHead(PreTrainedModel): 

204 """ 

205 Diffusion head model for vibevoice. 

206 Args: 

207 config (`VibeVoiceDiffusionHeadConfig`): Model configuration 

208 latent_size (`int`, optional): Size of the latent space. If not provided, uses `config.latent_size`. 

209 """ 

210 config_class = VibeVoiceDiffusionHeadConfig 

211 supports_gradient_checkpointing = True 

212 _supports_flash_attn_2 = True 

213 _supports_sdpa = True 

214 

215 def __init__( 

216 self, 

217 config: VibeVoiceDiffusionHeadConfig, 

218 ) -> None: 

219 super().__init__(config) 

220 self.config = config 

221 self.cond_dim = config.hidden_size 

222 latent_size = config.latent_size 

223 

224 self.noisy_images_proj = nn.Linear(latent_size, config.hidden_size, bias=False) 

225 self.cond_proj = nn.Linear(config.hidden_size, self.cond_dim, bias=False) 

226 self.t_embedder = TimestepEmbedder(self.cond_dim) 

227 

228 ffn_dim = int(config.hidden_size * config.head_ffn_ratio) 

229 

230 # Create the intermediate layers 

231 self.layers = nn.ModuleList([ 

232 HeadLayer( 

233 embed_dim=config.hidden_size, 

234 ffn_dim=ffn_dim, 

235 cond_dim=self.cond_dim, 

236 norm_eps=config.rms_norm_eps 

237 ) 

238 for _ in range(config.head_layers) 

239 ]) 

240 

241 # Final layer for output 

242 self.final_layer = FinalLayer( 

243 hidden_size=config.hidden_size, 

244 output_size=latent_size, 

245 cond_size=self.cond_dim, 

246 norm_eps=config.rms_norm_eps 

247 ) 

248 

249 self.initialize_weights() 

250 

251 def initialize_weights(self) -> None: 

252 """Initialize the weights of the model.""" 

253 # Initialize timestep embedder 

254 nn.init.normal_(cast(nn.Linear, self.t_embedder.mlp[0]).weight, std=0.02) 

255 nn.init.normal_(cast(nn.Linear, self.t_embedder.mlp[2]).weight, std=0.02) 

256 

257 # Zero-out adaLN modulation layers 

258 for layer in self.layers: 

259 head_layer = cast(HeadLayer, layer) 

260 nn.init.constant_(cast(nn.Linear, head_layer.adaLN_modulation[-1]).weight, 0) 

261 

262 # Zero-out output layers 

263 nn.init.constant_(cast(nn.Linear, self.final_layer.adaLN_modulation[-1]).weight, 0) 

264 nn.init.constant_(self.final_layer.linear.weight, 0) 

265 

266 def forward( 

267 self, 

268 noisy_images: torch.Tensor, 

269 timesteps: torch.Tensor, 

270 condition: torch.Tensor, 

271 ) -> torch.Tensor: 

272 """ 

273 Forward pass of the prediction head. 

274 Args: 

275 noisy_images (`torch.Tensor`): Noisy images/latents to denoise 

276 timesteps (`torch.Tensor`): Timesteps for diffusion 

277 condition (`torch.Tensor`): Conditioning information 

278 Returns: 

279 `torch.Tensor`: The predicted noise/velocity 

280 """ 

281 x = self.noisy_images_proj(noisy_images) 

282 t = self.t_embedder(timesteps) 

283 condition = self.cond_proj(condition) 

284 c = condition + t 

285 

286 for layer in self.layers: 

287 x = layer(x, c) 

288 

289 x = self.final_layer(x, c) 

290 return x 

291 

292 

293AutoModel.register(VibeVoiceDiffusionHeadConfig, VibeVoiceDiffusionHead) 

294 

295__all__ = [ 

296 "VibeVoiceDiffusionHead", 

297]