expanded context window to 16384

Author: Plachtaa

modules/diffusion_transformer.py

@@ -106,7 +106,7 @@ def __init__(
         self.style_as_token = args.DiT.style_as_token if hasattr(args.DiT, 'style_as_token') else False
         self.uvit_skip_connection = args.DiT.uvit_skip_connection if hasattr(args.DiT, 'uvit_skip_connection') else False
         model_args = ModelArgs(
-            block_size=8192,#args.DiT.block_size,
+            block_size=16384,#args.DiT.block_size,
             n_layer=args.DiT.depth,
             n_head=args.DiT.num_heads,
             dim=args.DiT.hidden_dim,
@@ -139,7 +139,7 @@ def __init__(
         # self.style_embedder1 = weight_norm(nn.Linear(1024, args.DiT.hidden_dim, bias=True))
         # self.style_embedder2 = weight_norm(nn.Linear(1024, args.style_encoder.dim, bias=True))
 
-        input_pos = torch.arange(8192)
+        input_pos = torch.arange(16384)
         self.register_buffer("input_pos", input_pos)
 
         self.conv1 = nn.Linear(args.DiT.hidden_dim, args.wavenet.hidden_dim)

modules/flow_matching.py

@@ -6,6 +6,8 @@
 from modules.diffusion_transformer import DiT
 from modules.commons import sequence_mask
 
+from tqdm import tqdm
+
 class BASECFM(torch.nn.Module, ABC):
     def __init__(
         self,
@@ -76,7 +78,7 @@ def solve_euler(self, x, x_lens, prompt, mu, style, f0, t_span, inference_cfg_ra
         x[..., :prompt_len] = 0
         if self.zero_prompt_speech_token:
             mu[..., :prompt_len] = 0
-        for step in range(1, len(t_span)):
+        for step in tqdm(range(1, len(t_span))):
             dphi_dt = self.estimator(x, prompt_x, x_lens, t.unsqueeze(0), style, mu, f0)
             # Classifier-Free Guidance inference introduced in VoiceBox
             if inference_cfg_rate > 0: