Fix black formatting in irodori_tts model files

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

Author: yoshphys

mlx_audio/tts/models/irodori_tts/config.py

@@ -49,11 +49,19 @@ def speaker_patched_latent_dim(self) -> int:
 
     @property
     def text_mlp_ratio_resolved(self) -> float:
-        return self.mlp_ratio if self.text_mlp_ratio is None else float(self.text_mlp_ratio)
+        return (
+            self.mlp_ratio
+            if self.text_mlp_ratio is None
+            else float(self.text_mlp_ratio)
+        )
 
     @property
     def speaker_mlp_ratio_resolved(self) -> float:
-        return self.mlp_ratio if self.speaker_mlp_ratio is None else float(self.speaker_mlp_ratio)
+        return (
+            self.mlp_ratio
+            if self.speaker_mlp_ratio is None
+            else float(self.speaker_mlp_ratio)
+        )
 
 
 @dataclass

mlx_audio/tts/models/irodori_tts/irodori_tts.py

@@ -99,12 +99,14 @@ def post_load_hook(cls, model: "Model", model_path: Path) -> "Model":
                 dac = DACVAE(cfg)
                 dac.load_weights(str(local_dacvae / "model.safetensors"))
                 import mlx.core as _mx
+
                 _mx.eval(dac.parameters())
                 model.dacvae = dac
             else:
                 model.dacvae = DACVAE.from_pretrained(model.config.dacvae_repo)
         except Exception as e:
             import warnings
+
             warnings.warn(
                 f"Could not load DACVAE: {e}\n"
                 "Set model.dacvae manually before calling generate()."
@@ -119,6 +121,7 @@ def post_load_hook(cls, model: "Model", model_path: Path) -> "Model":
     def _get_tokenizer(self):
         if self._tokenizer is None:
             from transformers import AutoTokenizer
+
             self._tokenizer = AutoTokenizer.from_pretrained(
                 self.config.dit.text_tokenizer_repo
             )
@@ -147,23 +150,23 @@ def _prepare_text(
     # Reference audio encoding
     # ------------------------------------------------------------------
 
-    def _encode_ref_audio(
-        self, audio: mx.array
-    ) -> tuple[mx.array, mx.array]:
+    def _encode_ref_audio(self, audio: mx.array) -> tuple[mx.array, mx.array]:
         """
         Encode reference waveform with DACVAE.
         audio: (1, samples) at config.sample_rate
         Returns (latent, mask): latent (1, T, 128), mask (1, T) bool
         """
         assert self.dacvae is not None, "DACVAE not loaded"
 
-        max_samples = self.config.max_speaker_latent_length * self.config.audio_downsample_factor
+        max_samples = (
+            self.config.max_speaker_latent_length * self.config.audio_downsample_factor
+        )
         audio = audio[:, :max_samples]
 
         # DACVAE encode expects (B, L, 1)
-        audio_in = audio[:, :, None]                   # (1, L, 1)
-        latent = self.dacvae.encode(audio_in)          # (1, 128, T) channels-first
-        latent = mx.transpose(latent, (0, 2, 1))       # (1, T, 128) sequence-first
+        audio_in = audio[:, :, None]  # (1, L, 1)
+        latent = self.dacvae.encode(audio_in)  # (1, 128, T) channels-first
+        latent = mx.transpose(latent, (0, 2, 1))  # (1, T, 128) sequence-first
 
         actual_t = int(audio.shape[1]) // self.config.audio_downsample_factor
         actual_t = min(actual_t, latent.shape[1])
@@ -266,8 +269,8 @@ def generate(
         # Decode latent → waveform
         # latent_out: (1, T, 128)
         latent_for_decode = mx.transpose(latent_out, (0, 2, 1))  # (1, 128, T)
-        audio_out = self.dacvae.decode(latent_for_decode)        # (1, L, 1)
-        audio_out = audio_out[:, :, 0]                           # (1, L)
+        audio_out = self.dacvae.decode(latent_for_decode)  # (1, L, 1)
+        audio_out = audio_out[:, :, 0]  # (1, L)
 
         # Trim trailing silence
         silence_t = _find_silence_point(latent_out[0])
@@ -277,7 +280,9 @@ def generate(
         audio = audio_out[0]  # (L,)
         samples = int(audio.shape[0])
         elapsed = max(time.perf_counter() - start_time, 1e-6)
-        audio_duration_seconds = samples / self.sample_rate if self.sample_rate > 0 else 0.0
+        audio_duration_seconds = (
+            samples / self.sample_rate if self.sample_rate > 0 else 0.0
+        )
 
         h = int(audio_duration_seconds // 3600)
         m = int((audio_duration_seconds % 3600) // 60)

mlx_audio/tts/models/irodori_tts/sampling.py

@@ -58,8 +58,8 @@ def _temporal_score_rescale(
     if t >= 1.0:
         return v_pred
     one_minus_t = 1.0 - t
-    snr = (one_minus_t ** 2) / (t ** 2)
-    sigma_sq = rescale_sigma ** 2
+    snr = (one_minus_t**2) / (t**2)
+    sigma_sq = rescale_sigma**2
     ratio = (snr * sigma_sq + 1.0) / (snr * sigma_sq / rescale_k + 1.0)
     return (ratio * (one_minus_t * v_pred + x_t) - x_t) / one_minus_t
 
@@ -237,13 +237,19 @@ def sample_euler_cfg(
                         axis=0,
                     )
                     v_out = model.forward_with_conditions(
-                        x_t=x_cfg, t=t_cfg,
+                        x_t=x_cfg,
+                        t=t_cfg,
                         text_state=mx.concatenate(
-                            [text_state_cond, text_state_uncond, text_state_cond], axis=0
+                            [text_state_cond, text_state_uncond, text_state_cond],
+                            axis=0,
                         ),
                         text_mask=text_mask_cfg,
                         speaker_state=mx.concatenate(
-                            [speaker_state_cond, speaker_state_cond, speaker_state_uncond],
+                            [
+                                speaker_state_cond,
+                                speaker_state_cond,
+                                speaker_state_uncond,
+                            ],
                             axis=0,
                         ),
                         speaker_mask=speaker_mask_cfg,
@@ -261,7 +267,8 @@ def sample_euler_cfg(
                     x_cfg = mx.concatenate([x_t, x_t], axis=0)
                     t_cfg = mx.full((batch_size * 2,), t, dtype=mx.float32)
                     v_out = model.forward_with_conditions(
-                        x_t=x_cfg, t=t_cfg,
+                        x_t=x_cfg,
+                        t=t_cfg,
                         text_state=mx.concatenate(
                             [text_state_cond, text_state_uncond], axis=0
                         ),
@@ -284,7 +291,8 @@ def sample_euler_cfg(
                     x_cfg = mx.concatenate([x_t, x_t], axis=0)
                     t_cfg = mx.full((batch_size * 2,), t, dtype=mx.float32)
                     v_out = model.forward_with_conditions(
-                        x_t=x_cfg, t=t_cfg,
+                        x_t=x_cfg,
+                        t=t_cfg,
                         text_state=mx.concatenate(
                             [text_state_cond, text_state_cond], axis=0
                         ),
@@ -315,53 +323,77 @@ def sample_euler_cfg(
                     joint_scale = cfg_scale_text if has_text_cfg else cfg_scale_speaker
 
                 v_cond = model.forward_with_conditions(
-                    x_t=x_t, t=t_arr,
-                    text_state=text_state_cond, text_mask=text_mask_cond,
-                    speaker_state=speaker_state_cond, speaker_mask=speaker_mask_cond,
-                    kv_text=kv_text_cond, kv_speaker=kv_speaker_cond,
+                    x_t=x_t,
+                    t=t_arr,
+                    text_state=text_state_cond,
+                    text_mask=text_mask_cond,
+                    speaker_state=speaker_state_cond,
+                    speaker_mask=speaker_mask_cond,
+                    kv_text=kv_text_cond,
+                    kv_speaker=kv_speaker_cond,
                 )
                 v_uncond = model.forward_with_conditions(
-                    x_t=x_t, t=t_arr,
-                    text_state=text_state_uncond, text_mask=text_mask_uncond,
-                    speaker_state=speaker_state_uncond, speaker_mask=speaker_mask_uncond,
-                    kv_text=kv_text_uncond_joint, kv_speaker=kv_speaker_uncond_joint,
+                    x_t=x_t,
+                    t=t_arr,
+                    text_state=text_state_uncond,
+                    text_mask=text_mask_uncond,
+                    speaker_state=speaker_state_uncond,
+                    speaker_mask=speaker_mask_uncond,
+                    kv_text=kv_text_uncond_joint,
+                    kv_speaker=kv_speaker_uncond_joint,
                 )
                 v_pred = v_cond + joint_scale * (v_cond - v_uncond)
 
             else:  # alternating
                 v_cond = model.forward_with_conditions(
-                    x_t=x_t, t=t_arr,
-                    text_state=text_state_cond, text_mask=text_mask_cond,
-                    speaker_state=speaker_state_cond, speaker_mask=speaker_mask_cond,
-                    kv_text=kv_text_cond, kv_speaker=kv_speaker_cond,
+                    x_t=x_t,
+                    t=t_arr,
+                    text_state=text_state_cond,
+                    text_mask=text_mask_cond,
+                    speaker_state=speaker_state_cond,
+                    speaker_mask=speaker_mask_cond,
+                    kv_text=kv_text_cond,
+                    kv_speaker=kv_speaker_cond,
                 )
                 use_text_uncond = (has_text_cfg and has_speaker_cfg and i % 2 == 0) or (
                     has_text_cfg and not has_speaker_cfg
                 )
                 if use_text_uncond:
                     v_uncond = model.forward_with_conditions(
-                        x_t=x_t, t=t_arr,
-                        text_state=text_state_uncond, text_mask=text_mask_uncond,
-                        speaker_state=speaker_state_cond, speaker_mask=speaker_mask_cond,
-                        kv_text=kv_text_uncond_alt, kv_speaker=kv_speaker_cond,
+                        x_t=x_t,
+                        t=t_arr,
+                        text_state=text_state_uncond,
+                        text_mask=text_mask_uncond,
+                        speaker_state=speaker_state_cond,
+                        speaker_mask=speaker_mask_cond,
+                        kv_text=kv_text_uncond_alt,
+                        kv_speaker=kv_speaker_cond,
                     )
                     v_pred = v_cond + cfg_scale_text * (v_cond - v_uncond)
                 else:
                     v_uncond = model.forward_with_conditions(
-                        x_t=x_t, t=t_arr,
-                        text_state=text_state_cond, text_mask=text_mask_cond,
-                        speaker_state=speaker_state_uncond, speaker_mask=speaker_mask_uncond,
-                        kv_text=kv_text_cond, kv_speaker=kv_speaker_uncond_alt,
+                        x_t=x_t,
+                        t=t_arr,
+                        text_state=text_state_cond,
+                        text_mask=text_mask_cond,
+                        speaker_state=speaker_state_uncond,
+                        speaker_mask=speaker_mask_uncond,
+                        kv_text=kv_text_cond,
+                        kv_speaker=kv_speaker_uncond_alt,
                     )
                     v_pred = v_cond + cfg_scale_speaker * (v_cond - v_uncond)
 
         else:
             # no CFG this step
             v_pred = model.forward_with_conditions(
-                x_t=x_t, t=t_arr,
-                text_state=text_state_cond, text_mask=text_mask_cond,
-                speaker_state=speaker_state_cond, speaker_mask=speaker_mask_cond,
-                kv_text=kv_text_cond, kv_speaker=kv_speaker_cond,
+                x_t=x_t,
+                t=t_arr,
+                text_state=text_state_cond,
+                text_mask=text_mask_cond,
+                speaker_state=speaker_state_cond,
+                speaker_mask=speaker_mask_cond,
+                kv_text=kv_text_cond,
+                kv_speaker=kv_speaker_cond,
             )
 
         # optional temporal score rescaling