Trainer

    loss = loss * mask  # Zero out ignored positions
    loss = loss.sum() / mask.sum()  # Normalize by valid elements
    return loss

def masked_cross_entropy_loss(est_length_logits, tar_length_labels):
    first_zero_idx = (tar_length_labels == 0).int().argmax(dim=1)
    B, L = tar_length_labels.shape
    range_tensor = torch.arange(L, device=tar_length_labels.device).expand(B, L)
    mask = range_tensor <= first_zero_idx[:, None]  # Include the first 0
    loss = F.cross_entropy(
        est_length_logits.reshape(-1, est_length_logits.size(-1)),
        tar_length_labels.reshape(-1),
        reduction='none'
    ).reshape(B, L)
    loss = loss * mask
    loss = loss.sum() / mask.sum()
    return loss

class Trainer:
    def __init__(
        self,
        model,
        vocab_size,
        vocab_char_map,
        process_token_to_id=True,
        loss_fn='L1',
        lambda_L1=1,
        gumbel_tau=0.5,
        n_class=301,
        n_frame_per_class=10,
        epochs=15,
        learning_rate=1e-4,
        num_warmup_updates=20000,
        save_per_updates=1000,
        checkpoint_path=None,
        batch_size=32,
        batch_size_type: str = "sample",
        max_samples=32,
        grad_accumulation_steps=1,
        max_grad_norm=1.0,
        logger: str | None = "wandb",  # "wandb" | "tensorboard" | None
        wandb_project="test_e2-tts",
        wandb_run_name="test_run",
        wandb_resume_id: str = None,
        last_per_steps=None,
        accelerate_kwargs: dict = dict(),
        ema_kwargs: dict = dict(),
    ):
        ddp_kwargs = DistributedDataParallelKwargs(find_unused_parameters=False)

        if logger == "wandb" and not wandb.api.api_key:
            logger = None
        print(f"Using logger: {logger}")

        self.accelerator = Accelerator(
            log_with=logger if logger == "wandb" else None,
            kwargs_handlers=[ddp_kwargs],
            gradient_accumulation_steps=grad_accumulation_steps,
            **accelerate_kwargs,
        )

        self.logger = logger
        if self.logger == "wandb":
            if exists(wandb_resume_id):
                init_kwargs = {"wandb": {"resume": "allow", "name": wandb_run_name, "id": wandb_resume_id}}
            else:
                init_kwargs = {"wandb": {"resume": "allow", "name": wandb_run_name}}

            self.accelerator.init_trackers(
                project_name=wandb_project,
                init_kwargs=init_kwargs,
                config={
                    "epochs": epochs,
                    "learning_rate": learning_rate,
                    "num_warmup_updates": num_warmup_updates,
                    "batch_size": batch_size,
                    "batch_size_type": batch_size_type,
                    "max_samples": max_samples,
                    "grad_accumulation_steps": grad_accumulation_steps,
                    "max_grad_norm": max_grad_norm,
                    "gpus": self.accelerator.num_processes,
                },
            )

        elif self.logger == "tensorboard":
            from torch.utils.tensorboard import SummaryWriter

            self.writer = SummaryWriter(log_dir=f"runs/{wandb_run_name}")

        self.model = model
        self.vocab_size = vocab_size
        self.vocab_char_map = vocab_char_map
        self.process_token_to_id = process_token_to_id
        assert loss_fn in ['L1', 'CE', 'L1_and_CE']
        self.loss_fn = loss_fn
        self.lambda_L1 = lambda_L1
        self.n_class = n_class
        self.n_frame_per_class = n_frame_per_class
        self.gumbel_tau = gumbel_tau

        self.epochs = epochs
        self.num_warmup_updates = num_warmup_updates
        self.save_per_updates = save_per_updates
        self.last_per_steps = default(last_per_steps, save_per_updates * grad_accumulation_steps)
        self.checkpoint_path = default(checkpoint_path, "ckpts/test_e2-tts")

        self.batch_size = batch_size
        self.batch_size_type = batch_size_type
        self.max_samples = max_samples
        self.grad_accumulation_steps = grad_accumulation_steps
        self.max_grad_norm = max_grad_norm

        if bnb_optimizer:
            import bitsandbytes as bnb

            self.optimizer = bnb.optim.AdamW8bit(model.parameters(), lr=learning_rate)
        else:
            self.optimizer = AdamW(model.parameters(), lr=learning_rate)
        self.model, self.optimizer = self.accelerator.prepare(self.model, self.optimizer)

    @property
    def is_main(self):
        return self.accelerator.is_main_process

    def save_checkpoint(self, step, last=False):
        self.accelerator.wait_for_everyone()
        if self.is_main:
            checkpoint = dict(
                model_state_dict=self.accelerator.unwrap_model(self.model).state_dict(),
                optimizer_state_dict=self.accelerator.unwrap_model(self.optimizer).state_dict(),
                scheduler_state_dict=self.scheduler.state_dict(),
                step=step,
            )
            if not os.path.exists(self.checkpoint_path):
                os.makedirs(self.checkpoint_path)
            if last:
                self.accelerator.save(checkpoint, f"{self.checkpoint_path}/model_last.pt")
            else:
                self.accelerator.save(checkpoint, f"{self.checkpoint_path}/model_{step}.pt")

    def load_checkpoint(self):
        if (
            not exists(self.checkpoint_path)
            or not os.path.exists(self.checkpoint_path)
            or not any(filename.endswith(".pt") for filename in os.listdir(self.checkpoint_path))
        ):
            return 0

        self.accelerator.wait_for_everyone()
        if "model_last.pt" in os.listdir(self.checkpoint_path):
            latest_checkpoint = "model_last.pt"
        else:
            latest_checkpoint = sorted(
                [f for f in os.listdir(self.checkpoint_path) if f.endswith(".pt")],
                key=lambda x: int("".join(filter(str.isdigit, x))),
            )[-1]

        print(f'To load from {latest_checkpoint}.')

        # checkpoint = torch.load(f"{self.checkpoint_path}/{latest_checkpoint}", map_location=self.accelerator.device)  # rather use accelerator.load_state ಥ_ಥ
        checkpoint = torch.load(f"{self.checkpoint_path}/{latest_checkpoint}", weights_only=True, map_location="cpu")

        print(f'Loaded from {latest_checkpoint}.')

        if "step" in checkpoint:
            # patch for backward compatibility, 305e3ea
            for key in ["mel_spec.mel_stft.mel_scale.fb", "mel_spec.mel_stft.spectrogram.window"]:
                if key in checkpoint["model_state_dict"]:
                    del checkpoint["model_state_dict"][key]

            self.accelerator.unwrap_model(self.model).load_state_dict(checkpoint["model_state_dict"])
            self.accelerator.unwrap_model(self.optimizer).load_state_dict(checkpoint["optimizer_state_dict"])
            if self.scheduler:
                self.scheduler.load_state_dict(checkpoint["scheduler_state_dict"])
            step = checkpoint["step"]
        else:
            checkpoint["model_state_dict"] = {
                k.replace("ema_model.", ""): v
                for k, v in checkpoint["ema_model_state_dict"].items()
                if k not in ["initted", "step"]
            }
            self.accelerator.unwrap_model(self.model).load_state_dict(checkpoint["model_state_dict"])
            step = 0

        del checkpoint
        gc.collect()

        print(f'Exit load_checkpoint.')

        return step

    def validate(self, valid_dataloader, global_step):
        """
        Runs evaluation on the validation set, computes the average loss,
        and logs the average validation loss along with the CTC decoded strings.
        """
        self.model.eval()
        total_valid_loss = 0.0
        total_sec_error = 0.0
        count = 0
        # Iterate over the validation dataloader
        with torch.no_grad():
            for batch in valid_dataloader:
                # Inputs
                mel = batch['mel'].permute(0, 2, 1) # (B, L_mel, D)
                text = batch['text']

                if self.process_token_to_id:
                    text_ids = list_str_to_idx(text, self.vocab_char_map).to(mel.device)
                    text_ids = text_ids.masked_fill(text_ids==-1, self.vocab_size)
                else:
                    text_ids = text

                # Targets
                mel_lengths = batch['mel_lengths']
                tar_lengths = calculate_remaining_lengths(mel_lengths)
                predictions = self.model(text_ids=text_ids, mel=mel)

                if self.loss_fn == 'L1':
                    est_lengths = predictions
                    loss = masked_l1_loss(
                        est_lengths=est_lengths, tar_lengths=tar_lengths
                    )
                    frame_error = loss

                elif self.loss_fn == 'CE':
                    tar_length_labels = (tar_lengths // self.n_frame_per_class) \
                        .clamp(min=0, max=self.n_class-1) # [0, 1, ..., n_class-1]
                    est_length_logtis = predictions
                    est_length_labels = torch.argmax(est_length_logtis, dim=-1)
                    loss = masked_cross_entropy_loss(
                        est_length_logits=est_length_logtis, tar_length_labels=tar_length_labels
                    )
                    est_lengths = est_length_labels * self.n_frame_per_class
                    frame_error = masked_l1_loss(
                        est_lengths=est_lengths, tar_lengths=tar_lengths
                    )

                elif self.loss_fn == 'L1_and_CE':
                    tar_length_labels = (tar_lengths // self.n_frame_per_class) \
                        .clamp(min=0, max=self.n_class-1) # [0, 1, ..., n_class-1]
                    est_length_logtis = predictions
                    est_length_1hots = F.gumbel_softmax(
                        est_length_logtis, tau=self.gumbel_tau, hard=True, dim=-1
                    )
                    length_values = torch.arange(
                        self.n_class, device=est_length_1hots.device
                    ).float() * self.n_frame_per_class
                    est_lengths = (est_length_1hots * length_values).sum(-1)

                    loss_CE = masked_cross_entropy_loss(
                        est_length_logits=est_length_logtis, tar_length_labels=tar_length_labels
                    )

                    loss_L1 = masked_l1_loss(
                        est_lengths=est_lengths, tar_lengths=tar_lengths
                    )

                    loss = loss_CE + self.lambda_L1 * loss_L1

                    frame_error = loss_L1

                else:
                    raise NotImplementedError(self.loss_fn)

                sec_error = frame_error * 256 / 24000
                total_sec_error += sec_error.item()
                total_valid_loss += loss.item()
                count += 1

        avg_valid_loss = total_valid_loss / count if count > 0 else 0.0
        avg_valid_sec_error = total_sec_error / count if count > 0 else 0.0
        # Log validation metrics
        self.accelerator.log(
            {
                f"valid_loss": avg_valid_loss,
                f"valid_sec_error": avg_valid_sec_error
            },
            step=global_step
        )

        self.model.train()

    def train(self, train_dataset: Dataset, valid_dataset: Dataset,
        num_workers=64, resumable_with_seed: int = None):
        if exists(resumable_with_seed):
            generator = torch.Generator()
            generator.manual_seed(resumable_with_seed)
        else:
            generator = None

        # Create training dataloader using the appropriate batching strategy
        if self.batch_size_type == "sample":
            train_dataloader = DataLoader(
                train_dataset,
                collate_fn=collate_fn,
                num_workers=num_workers,
                pin_memory=True,
                persistent_workers=True,
                batch_size=self.batch_size,
                shuffle=True,
                generator=generator,
            )
            # Create validation dataloader (always sequential, no shuffling)
            valid_dataloader = DataLoader(
                valid_dataset,
                collate_fn=collate_fn,
                num_workers=num_workers,
                pin_memory=True,
                batch_size=self.batch_size,
                shuffle=False,
            )

        elif self.batch_size_type == "frame":
            self.accelerator.even_batches = False

            sampler = SequentialSampler(train_dataset)
            batch_sampler = DynamicBatchSampler(
                sampler, self.batch_size, max_samples=self.max_samples, random_seed=resumable_with_seed, drop_last=False
            )
            train_dataloader = DataLoader(
                train_dataset,
                collate_fn=collate_fn,
                num_workers=num_workers,
                pin_memory=True,
                persistent_workers=True,
                batch_sampler=batch_sampler,
            )

            sampler = SequentialSampler(valid_dataset)
            batch_sampler = DynamicBatchSampler(
                sampler, self.batch_size, max_samples=self.max_samples, random_seed=resumable_with_seed, drop_last=False
            )
            # Create validation dataloader (always sequential, no shuffling)
            valid_dataloader = DataLoader(
                valid_dataset,
                collate_fn=collate_fn,
                num_workers=num_workers,
                pin_memory=True,
                persistent_workers=True,
                batch_sampler=batch_sampler,
            )
        else:
            raise ValueError(f"batch_size_type must be either 'sample' or 'frame', but received {self.batch_size_type}")

        #  accelerator.prepare() dispatches batches to devices;
        #  which means the length of dataloader calculated before, should consider the number of devices
        warmup_steps = (
            self.num_warmup_updates * self.accelerator.num_processes
        )  # consider a fixed warmup steps while using accelerate multi-gpu ddp
        # otherwise by default with split_batches=False, warmup steps change with num_processes
        total_steps = len(train_dataloader) * self.epochs / self.grad_accumulation_steps
        decay_steps = total_steps - warmup_steps
        warmup_scheduler = LinearLR(self.optimizer, start_factor=1e-8, end_factor=1.0, total_iters=warmup_steps)
        decay_scheduler = LinearLR(self.optimizer, start_factor=1.0, end_factor=1e-8, total_iters=decay_steps)
        self.scheduler = SequentialLR(
            self.optimizer, schedulers=[warmup_scheduler, decay_scheduler], milestones=[warmup_steps]
        )
        train_dataloader, self.scheduler = self.accelerator.prepare(
            train_dataloader, self.scheduler
        )  # actual steps = 1 gpu steps / gpus
        start_step = self.load_checkpoint()
        global_step = start_step

        valid_dataloader = self.accelerator.prepare(valid_dataloader)

        if exists(resumable_with_seed):
            orig_epoch_step = len(train_dataloader)
            skipped_epoch = int(start_step // orig_epoch_step)
            skipped_batch = start_step % orig_epoch_step
            skipped_dataloader = self.accelerator.skip_first_batches(train_dataloader, num_batches=skipped_batch)
        else:
            skipped_epoch = 0

        for epoch in range(skipped_epoch, self.epochs):
            self.model.train()
            if exists(resumable_with_seed) and epoch == skipped_epoch:
                progress_bar = tqdm(
                    skipped_dataloader,
                    desc=f"Epoch {epoch+1}/{self.epochs}",
                    unit="step",
                    disable=not self.accelerator.is_local_main_process,
                    initial=skipped_batch,
                    total=orig_epoch_step,
                )
            else:
                progress_bar = tqdm(
                    train_dataloader,
                    desc=f"Epoch {epoch+1}/{self.epochs}",
                    unit="step",
                    disable=not self.accelerator.is_local_main_process,
                )

            for batch in progress_bar:
                with self.accelerator.accumulate(self.model):
                    # Inputs
                    mel = batch['mel'].permute(0, 2, 1) # (B, L_mel, D)
                    text = batch['text']

                    if self.process_token_to_id:
                        text_ids = list_str_to_idx(text, self.vocab_char_map).to(mel.device)
                        text_ids = text_ids.masked_fill(text_ids==-1, self.vocab_size)
                    else:
                        text_ids = text

                    # Targets
                    mel_lengths = batch['mel_lengths']
                    tar_lengths = calculate_remaining_lengths(mel_lengths)
                    predictions = self.model(text_ids=text_ids, mel=mel)

                    if self.loss_fn == 'L1':
                        est_lengths = predictions
                        loss = masked_l1_loss(
                            est_lengths=est_lengths, tar_lengths=tar_lengths
                        )

                        with torch.no_grad():
                            frame_error = loss
                            sec_error = frame_error * 256 / 24000

                        log_dict = {
                            'loss': loss.item(),
                            'loss_L1': loss.item(),
                            'sec_error': sec_error.item(),
                            'lr': self.scheduler.get_last_lr()[0]
                         }

                    elif self.loss_fn == 'CE':
                        tar_length_labels = (tar_lengths // self.n_frame_per_class) \
                            .clamp(min=0, max=self.n_class-1) # [0, 1, ..., n_class-1]
                        est_length_logtis = predictions
                        est_length_labels = torch.argmax(est_length_logtis, dim=-1)
                        loss = masked_cross_entropy_loss(
                            est_length_logits=est_length_logtis, tar_length_labels=tar_length_labels
                        )
                        with torch.no_grad():
                            est_lengths = est_length_labels * self.n_frame_per_class
                            frame_error = masked_l1_loss(
                                est_lengths=est_lengths, tar_lengths=tar_lengths
                            )
                            sec_error = frame_error * 256 / 24000

                        log_dict = {
                            'loss': loss.item(),
                            'loss_CE': loss.item(),
                            'sec_error': sec_error.item(),
                            'lr': self.scheduler.get_last_lr()[0]
                         }

                    elif self.loss_fn == 'L1_and_CE':
                        tar_length_labels = (tar_lengths // self.n_frame_per_class) \
                            .clamp(min=0, max=self.n_class-1) # [0, 1, ..., n_class-1]
                        est_length_logtis = predictions
                        est_length_1hots = F.gumbel_softmax(
                            est_length_logtis, tau=self.gumbel_tau, hard=True, dim=-1
                        )
                        length_values = torch.arange(
                            self.n_class, device=est_length_1hots.device
                        ).float() * self.n_frame_per_class
                        est_lengths = (est_length_1hots * length_values).sum(-1)

                        loss_CE = masked_cross_entropy_loss(
                            est_length_logits=est_length_logtis, tar_length_labels=tar_length_labels
                        )

                        loss_L1 =  masked_l1_loss(
                            est_lengths=e
# ... truncated ...


    loss = loss.sum() / mask.sum()  # Normalize by valid elements
    return loss

def masked_cross_entropy_loss(est_length_logits, tar_length_labels):
    first_zero_idx = (tar_length_labels == 0).int().argmax(dim=1)
    B, L = tar_length_labels.shape
    range_tensor = torch.arange(L, device=tar_length_labels.device).expand(B, L)
    mask = range_tensor <= first_zero_idx[:, None]  # Include the first 0
    loss = F.cross_entropy(
        est_length_logits.reshape(-1, est_length_logits.size(-1)),
        tar_length_labels.reshape(-1),
        reduction='none'
    ).reshape(B, L)
    loss = loss * mask
    loss = loss.sum() / mask.sum()
    return loss

class Trainer:
    def __init__(
        self,
        model,
        vocab_size,
        vocab_char_map,
        process_token_to_id=True,
        loss_fn='L1',
        lambda_L1=1,
        gumbel_tau=0.5,
        n_class=301,
        n_frame_per_class=10,
        epochs=15,
        learning_rate=1e-4,
        num_warmup_updates=20000,
        save_per_updates=1000,
        checkpoint_path=None,
        batch_size=32,
        batch_size_type: str = "sample",
        max_samples=32,
        grad_accumulation_steps=1,
        max_grad_norm=1.0,
        logger: str | None = "wandb",  # "wandb" | "tensorboard" | None
        wandb_project="test_e2-tts",
        wandb_run_name="test_run",
        wandb_resume_id: str = None,
        last_per_steps=None,
        accelerate_kwargs: dict = dict(),
        ema_kwargs: dict = dict(),
    ):
        ddp_kwargs = DistributedDataParallelKwargs(find_unused_parameters=False)

        if logger == "wandb" and not wandb.api.api_key:
            logger = None
        print(f"Using logger: {logger}")

        self.accelerator = Accelerator(
            log_with=logger if logger == "wandb" else None,
            kwargs_handlers=[ddp_kwargs],
            gradient_accumulation_steps=grad_accumulation_steps,
            **accelerate_kwargs,
        )

        self.logger = logger
        if self.logger == "wandb":
            if exists(wandb_resume_id):
                init_kwargs = {"wandb": {"resume": "allow", "name": wandb_run_name, "id": wandb_resume_id}}
            else:
                init_kwargs = {"wandb": {"resume": "allow", "name": wandb_run_name}}

            self.accelerator.init_trackers(
                project_name=wandb_project,
                init_kwargs=init_kwargs,
                config={
                    "epochs": epochs,
                    "learning_rate": learning_rate,
                    "num_warmup_updates": num_warmup_updates,
                    "batch_size": batch_size,
                    "batch_size_type": batch_size_type,
                    "max_samples": max_samples,
                    "grad_accumulation_steps": grad_accumulation_steps,
                    "max_grad_norm": max_grad_norm,
                    "gpus": self.accelerator.num_processes,
                },
            )

        elif self.logger == "tensorboard":
            from torch.utils.tensorboard import SummaryWriter

            self.writer = SummaryWriter(log_dir=f"runs/{wandb_run_name}")

        self.model = model
        self.vocab_size = vocab_size
        self.vocab_char_map = vocab_char_map
        self.process_token_to_id = process_token_to_id
        assert loss_fn in ['L1', 'CE', 'L1_and_CE']
        self.loss_fn = loss_fn
        self.lambda_L1 = lambda_L1
        self.n_class = n_class
        self.n_frame_per_class = n_frame_per_class
        self.gumbel_tau = gumbel_tau

        self.epochs = epochs
        self.num_warmup_updates = num_warmup_updates
        self.save_per_updates = save_per_updates
        self.last_per_steps = default(last_per_steps, save_per_updates * grad_accumulation_steps)
        self.checkpoint_path = default(checkpoint_path, "ckpts/test_e2-tts")

        self.batch_size = batch_size
        self.batch_size_type = batch_size_type
        self.max_samples = max_samples
        self.grad_accumulation_steps = grad_accumulation_steps
        self.max_grad_norm = max_grad_norm

        if bnb_optimizer:
            import bitsandbytes as bnb

            self.optimizer = bnb.optim.AdamW8bit(model.parameters(), lr=learning_rate)
        else:
            self.optimizer = AdamW(model.parameters(), lr=learning_rate)
        self.model, self.optimizer = self.accelerator.prepare(self.model, self.optimizer)

    @property
    def is_main(self):
        return self.accelerator.is_main_process

    def save_checkpoint(self, step, last=False):
        self.accelerator.wait_for_everyone()
        if self.is_main:
            checkpoint = dict(
                model_state_dict=self.accelerator.unwrap_model(self.model).state_dict(),
                optimizer_state_dict=self.accelerator.unwrap_model(self.optimizer).state_dict(),
                scheduler_state_dict=self.scheduler.state_dict(),
                step=step,
            )
            if not os.path.exists(self.checkpoint_path):
                os.makedirs(self.checkpoint_path)
            if last:
                self.accelerator.save(checkpoint, f"{self.checkpoint_path}/model_last.pt")
            else:
                self.accelerator.save(checkpoint, f"{self.checkpoint_path}/model_{step}.pt")


        self.num_warmup_updates = num_warmup_updates
        self.save_per_updates = save_per_updates
        self.last_per_steps = default(last_per_steps, save_per_updates * grad_accumulation_steps)
        self.checkpoint_path = default(checkpoint_path, "ckpts/test_e2-tts")

        self.batch_size = batch_size
        self.batch_size_type = batch_size_type
        self.max_samples = max_samples
        self.grad_accumulation_steps = grad_accumulation_steps
        self.max_grad_norm = max_grad_norm

        if bnb_optimizer:
            import bitsandbytes as bnb

            self.optimizer = bnb.optim.AdamW8bit(model.parameters(), lr=learning_rate)
        else:
            self.optimizer = AdamW(model.parameters(), lr=learning_rate)
        self.model, self.optimizer = self.accelerator.prepare(self.model, self.optimizer)

    @property
    def is_main(self):
        return self.accelerator.is_main_process

    def save_checkpoint(self, step, last=False):
        self.accelerator.wait_for_everyone()
        if self.is_main:
            checkpoint = dict(
                model_state_dict=self.accelerator.unwrap_model(self.model).state_dict(),
                optimizer_state_dict=self.accelerator.unwrap_model(self.optimizer).state_dict(),
                scheduler_state_dict=self.scheduler.state_dict(),
                step=step,
            )
            if not os.path.exists(self.checkpoint_path):
                os.makedirs(self.checkpoint_path)
            if last:
                self.accelerator.save(checkpoint, f"{self.checkpoint_path}/model_last.pt")
            else:
                self.accelerator.save(checkpoint, f"{self.checkpoint_path}/model_{step}.pt")

    def load_checkpoint(self):
        if (
            not exists(self.checkpoint_path)


        self.checkpoint_path = default(checkpoint_path, "ckpts/test_e2-tts")

        self.batch_size = batch_size
        self.batch_size_type = batch_size_type
        self.max_samples = max_samples
        self.grad_accumulation_steps = grad_accumulation_steps
        self.max_grad_norm = max_grad_norm

        if bnb_optimizer:
            import bitsandbytes as bnb

            self.optimizer = bnb.optim.AdamW8bit(model.parameters(), lr=learning_rate)
        else:
            self.optimizer = AdamW(model.parameters(), lr=learning_rate)
        self.model, self.optimizer = self.accelerator.prepare(self.model, self.optimizer)

    @property
    def is_main(self):
        return self.accelerator.is_main_process

    def save_checkpoint(self, step, last=False):
        self.accelerator.wait_for_everyone()
        if self.is_main:
            checkpoint = dict(
                model_state_dict=self.accelerator.unwrap_model(self.model).state_dict(),
                optimizer_state_dict=self.accelerator.unwrap_model(self.optimizer).state_dict(),
                scheduler_state_dict=self.scheduler.state_dict(),
                step=step,
            )
            if not os.path.exists(self.checkpoint_path):
                os.makedirs(self.checkpoint_path)
            if last:
                self.accelerator.save(checkpoint, f"{self.checkpoint_path}/model_last.pt")
            else:
                self.accelerator.save(checkpoint, f"{self.checkpoint_path}/model_{step}.pt")

    def load_checkpoint(self):
        if (
            not exists(self.checkpoint_path)
            or not os.path.exists(self.checkpoint_path)
            or not any(filename.endswith(".pt") for filename in os.listdir(self.checkpoint_path))
        ):
            return 0

        self.accelerator.wait_for_everyone()
        if "model_last.pt" in os.listdir(self.checkpoint_path):
            latest_checkpoint = "model_last.pt"
        else:
            latest_checkpoint = sorted(
                [f for f in os.listdir(self.checkpoint_path) if f.endswith(".pt")],
                key=lambda x: int("".join(filter(str.isdigit, x))),
            )[-1]

        print(f'To load from {latest_checkpoint}.')


    @property
    def is_main(self):
        return self.accelerator.is_main_process

    def save_checkpoint(self, step, last=False):
        self.accelerator.wait_for_everyone()
        if self.is_main:
            checkpoint = dict(
                model_state_dict=self.accelerator.unwrap_model(self.model).state_dict(),
                optimizer_state_dict=self.accelerator.unwrap_model(self.optimizer).state_dict(),
                scheduler_state_dict=self.scheduler.state_dict(),
                step=step,
            )
            if not os.path.exists(self.checkpoint_path):
                os.makedirs(self.checkpoint_path)
            if last:
                self.accelerator.save(checkpoint, f"{self.checkpoint_path}/model_last.pt")
            else:
                self.accelerator.save(checkpoint, f"{self.checkpoint_path}/model_{step}.pt")

    def load_checkpoint(self):
        if (
            not exists(self.checkpoint_path)
            or not os.path.exists(self.checkpoint_path)
            or not any(filename.endswith(".pt") for filename in os.listdir(self.checkpoint_path))
        ):
            return 0

        self.accelerator.wait_for_everyone()
        if "model_last.pt" in os.listdir(self.checkpoint_path):
            latest_checkpoint = "model_last.pt"
        else:
            latest_checkpoint = sorted(
                [f for f in os.listdir(self.checkpoint_path) if f.endswith(".pt")],
                key=lambda x: int("".join(filter(str.isdigit, x))),
            )[-1]

        print(f'To load from {latest_checkpoint}.')

        # checkpoint = torch.load(f"{self.checkpoint_path}/{latest_checkpoint}", map_location=self.accelerator.device)  # rather use accelerator.load_state ಥ_ಥ
        checkpoint = torch.load(f"{self.checkpoint_path}/{latest_checkpoint}", weights_only=True, map_location="cpu")

        print(f'Loaded from {latest_checkpoint}.')

        if "step" in checkpoint:
            # patch for backward compatibility, 305e3ea
            for key in ["mel_spec.mel_stft.mel_scale.fb", "mel_spec.mel_stft.spectrogram.window"]:
                if key in checkpoint["model_state_dict"]:
                    del checkpoint["model_state_dict"][key]

            self.accelerator.unwrap_model(self.model).load_state_dict(checkpoint["model_state_dict"])
            self.accelerator.unwrap_model(self.optimizer).load_state_dict(checkpoint["optimizer_state_dict"])
            if self.scheduler:
                self.scheduler.load_state_dict(checkpoint["scheduler_state_dict"])
            step = checkpoint["step"]
        else:
            checkpoint["model_state_dict"] = {
                k.replace("ema_model.", ""): v
                for k, v in checkpoint["ema_model_state_dict"].items()
                if k not in ["initted", "step"]
            }
            self.accelerator.unwrap_model(self.model).load_state_dict(checkpoint["model_state_dict"])
            step = 0

        del checkpoint
        gc.collect()

        print(f'Exit load_checkpoint.')

        return step

    def validate(self, valid_dataloader, global_step):
        """
        Runs evaluation on the validation set, computes the average loss,
        and logs the average validation loss along with the CTC decoded strings.
        """
        self.model.eval()
        total_valid_loss = 0.0
        total_sec_error = 0.0
        count = 0
        # Iterate over the validation dataloader
        with torch.no_grad():
            for batch in valid_dataloader:
                # Inputs
                mel = batch['mel'].permute(0, 2, 1) # (B, L_mel, D)
                text = batch['text']

                if self.process_token_to_id:
                    text_ids = list_str_to_idx(text, self.vocab_char_map).to(mel.device)


validate

            if self.scheduler:
                self.scheduler.load_state_dict(checkpoint["scheduler_state_dict"])
            step = checkpoint["step"]
        else:
            checkpoint["model_state_dict"] = {
                k.replace("ema_model.", ""): v
                for k, v in checkpoint["ema_model_state_dict"].items()
                if k not in ["initted", "step"]
            }
            self.accelerator.unwrap_model(self.model).load_state_dict(checkpoint["model_state_dict"])
            step = 0

        del checkpoint
        gc.collect()

        print(f'Exit load_checkpoint.')

        return step

    def validate(self, valid_dataloader, global_step):
        """
        Runs evaluation on the validation set, computes the average loss,
        and logs the average validation loss along with the CTC decoded strings.
        """
        self.model.eval()
        total_valid_loss = 0.0
        total_sec_error = 0.0
        count = 0
        # Iterate over the validation dataloader
        with torch.no_grad():
            for batch in valid_dataloader:
                # Inputs
                mel = batch['mel'].permute(0, 2, 1) # (B, L_mel, D)
                text = batch['text']

                if self.process_token_to_id:
                    text_ids = list_str_to_idx(text, self.vocab_char_map).to(mel.device)
                    text_ids = text_ids.masked_fill(text_ids==-1, self.vocab_size)
                else:
                    text_ids = text

                # Targets
                mel_lengths = batch['mel_lengths']
                tar_lengths = calculate_remaining_lengths(mel_lengths)
                predictions = self.model(text_ids=text_ids, mel=mel)

                if self.loss_fn == 'L1':
                    est_lengths = predictions
                    loss = masked_l1_loss(
                        est_lengths=est_lengths, tar_lengths=tar_lengths
                    )
                    frame_error = loss

                elif self.loss_fn == 'CE':
                    tar_length_labels = (tar_lengths // self.n_frame_per_class) \
                        .clamp(min=0, max=self.n_class-1) # [0, 1, ..., n_class-1]
                    est_length_logtis = predictions
                    est_length_labels = torch.argmax(est_length_logtis, dim=-1)
                    loss = masked_cross_entropy_loss(
                        est_length_logits=est_length_logtis, tar_length_labels=tar_length_labels
                    )
                    est_lengths = est_length_labels * self.n_frame_per_class
                    frame_error = masked_l1_loss(
                        est_lengths=est_lengths, tar_lengths=tar_lengths
                    )

                elif self.loss_fn == 'L1_and_CE':
                    tar_length_labels = (tar_lengths // self.n_frame_per_class) \
                        .clamp(min=0, max=self.n_class-1) # [0, 1, ..., n_class-1]
                    est_length_logtis = predictions
                    est_length_1hots = F.gumbel_softmax(
                        est_length_logtis, tau=self.gumbel_tau, hard=True, dim=-1
                    )
                    length_values = torch.arange(
                        self.n_class, device=est_length_1hots.device
                    ).float() * self.n_frame_per_class
                    est_lengths = (est_length_1hots * length_values).sum(-1)

                    loss_CE = masked_cross_entropy_loss(
                        est_length_logits=est_length_logtis, tar_length_labels=tar_length_labels
                    )

                    loss_L1 = masked_l1_loss(
                        est_lengths=est_lengths, tar_lengths=tar_lengths
                    )

                    loss = loss_CE + self.lambda_L1 * loss_L1

                    frame_error = loss_L1

                else:
                    raise NotImplementedError(self.loss_fn)

                sec_error = frame_error * 256 / 24000
                total_sec_error += sec_error.item()
                total_valid_loss += loss.item()
                count += 1

        avg_valid_loss = total_valid_loss / count if count > 0 else 0.0
        avg_valid_sec_error = total_sec_error / count if count > 0 else 0.0
        # Log validation metrics
        self.accelerator.log(
            {
                f"valid_loss": avg_valid_loss,
                f"valid_sec_error": avg_valid_sec_error
            },
            step=global_step
        )

        self.model.train()

    def train(self, train_dataset: Dataset, valid_dataset: Dataset,
        num_workers=64, resumable_with_seed: int = None):
        if exists(resumable_with_seed):
            generator = torch.Generator()
            generator.manual_seed(resumable_with_seed)
        else:
            generator = None

        # Create training dataloader using the appropriate batching strategy
        if self.batch_size_type == "sample":
            train_dataloader = DataLoader(
                train_dataset,
                collate_fn=collate_fn,
                num_workers=num_workers,
                pin_memory=True,
                persistent_workers=True,
                batch_size=self.batch_size,
                shuffle=True,


train

                sec_error = frame_error * 256 / 24000
                total_sec_error += sec_error.item()
                total_valid_loss += loss.item()
                count += 1

        avg_valid_loss = total_valid_loss / count if count > 0 else 0.0
        avg_valid_sec_error = total_sec_error / count if count > 0 else 0.0
        # Log validation metrics
        self.accelerator.log(
            {
                f"valid_loss": avg_valid_loss,
                f"valid_sec_error": avg_valid_sec_error
            },
            step=global_step
        )

        self.model.train()

    def train(self, train_dataset: Dataset, valid_dataset: Dataset,
        num_workers=64, resumable_with_seed: int = None):
        if exists(resumable_with_seed):
            generator = torch.Generator()
            generator.manual_seed(resumable_with_seed)
        else:
            generator = None

        # Create training dataloader using the appropriate batching strategy
        if self.batch_size_type == "sample":
            train_dataloader = DataLoader(
                train_dataset,
                collate_fn=collate_fn,
                num_workers=num_workers,
                pin_memory=True,
                persistent_workers=True,
                batch_size=self.batch_size,
                shuffle=True,
                generator=generator,
            )
            # Create validation dataloader (always sequential, no shuffling)
            valid_dataloader = DataLoader(
                valid_dataset,
                collate_fn=collate_fn,
                num_workers=num_workers,
                pin_memory=True,
                batch_size=self.batch_size,
                shuffle=False,
            )

        elif self.batch_size_type == "frame":
            self.accelerator.even_batches = False

            sampler = SequentialSampler(train_dataset)
            batch_sampler = DynamicBatchSampler(
                sampler, self.batch_size, max_samples=self.max_samples, random_seed=resumable_with_seed, drop_last=False
            )
            train_dataloader = DataLoader(
                train_dataset,
                collate_fn=collate_fn,
                num_workers=num_workers,
                pin_memory=True,
                persistent_workers=True,
                batch_sampler=batch_sampler,
            )

            sampler = SequentialSampler(valid_dataset)
            batch_sampler = DynamicBatchSampler(
                sampler, self.batch_size, max_samples=self.max_samples, random_seed=resumable_with_seed, drop_last=False
            )
            # Create validation dataloader (always sequential, no shuffling)
            valid_dataloader = DataLoader(
                valid_dataset,
                collate_fn=collate_fn,
                num_workers=num_workers,
                pin_memory=True,
                persistent_workers=True,
                batch_sampler=batch_sampler,
            )
        else:
            raise ValueError(f"batch_size_type must be either 'sample' or 'frame', but received {self.batch_size_type}")

        #  accelerator.prepare() dispatches batches to devices;
        #  which means the length of dataloader calculated before, should consider the number of devices
        warmup_steps = (
            self.num_warmup_updates * self.accelerator.num_processes
        )  # consider a fixed warmup steps while using accelerate multi-gpu ddp
        # otherwise by default with split_batches=False, warmup steps change with num_processes
        total_steps = len(train_dataloader) * self.epochs / self.grad_accumulation_steps
        decay_steps = total_steps - warmup_steps
        warmup_scheduler = LinearLR(self.optimizer, start_factor=1e-8, end_factor=1.0, total_iters=warmup_steps)
        decay_scheduler = LinearLR(self.optimizer, start_factor=1.0, end_factor=1e-8, total_iters=decay_steps)
        self.scheduler = SequentialLR(
            self.optimizer, schedulers=[warmup_scheduler, decay_scheduler], milestones=[warmup_steps]
        )
        train_dataloader, self.scheduler = self.accelerator.prepare(
            train_dataloader, self.scheduler
        )  # actual steps = 1 gpu steps / gpus
        start_step = self.load_checkpoint()
        global_step = start_step

        valid_dataloader = self.accelerator.prepare(valid_dataloader)

        if exists(resumable_with_seed):
            orig_epoch_step = len(train_dataloader)
            skipped_epoch = int(start_step // orig_epoch_step)
            skipped_batch = start_step % orig_epoch_step
            skipped_dataloader = self.accelerator.skip_first_batches(train_dataloader, num_batches=skipped_batch)
        else:
            skipped_epoch = 0

        for epoch in range(skipped_epoch, self.epochs):
            self.model.train()
            if exists(resumable_with_seed) and epoch == skipped_epoch:
                progress_bar = tqdm(
                    skipped_dataloader,
                    desc=f"Epoch {epoch+1}/{self.epochs}",
                    unit="step",
                    disable=not self.accelerator.is_local_main_process,
                    initial=skipped_batch,
                    total=orig_epoch_step,
                )
            else:
                progress_bar = tqdm(
                    train_dataloader,
                    desc=f"Epoch {epoch+1}/{self.epochs}",
                    unit="step",
                    disable=not self.accelerator.is_local_main_process,
                )

            for batch in progress_bar:
                with self.accelerator.accumulate(self.model):
                    # Inputs
                    mel = batch['mel'].permute(0, 2, 1) # (B, L_mel, D)
                    text = batch['text']

                    if self.process_token_to_id:
                        text_ids = list_str_to_idx(text, self.vocab_char_map).to(mel.device)
                        text_ids = text_ids.masked_fill(text_ids==-1, self.vocab_size)
                    else:
                        text_ids = text

                    # Targets
                    mel_lengths = batch['mel_lengths']
                    tar_lengths = calculate_remaining_lengths(mel_lengths)
                    predictions = self.model(text_ids=text_ids, mel=mel)

                    if self.loss_fn == 'L1':
                        est_lengths = predictions
                        loss = masked_l1_loss(
                            est_lengths=est_lengths, tar_lengths=tar_lengths
                        )

                        with torch.no_grad():
                            frame_error = loss
                            sec_error = frame_error * 256 / 24000

                        log_dict = {
                            'loss': loss.item(),
                            'loss_L1': loss.item(),
                            'sec_error': sec_error.item(),
                            'lr': self.scheduler.get_last_lr()[0]
                         }

                    elif self.loss_fn == 'CE':
                        tar_length_labels = (tar_lengths // self.n_frame_per_class) \
                            .clamp(min=0, max=self.n_class-1) # [0, 1, ..., n_class-1]
                        est_length_logtis = predictions
                        est_length_labels = torch.argmax(est_length_logtis, dim=-1)
                        loss = masked_cross_entropy_loss(
                            est_length_logits=est_length_logtis, tar_length_labels=tar_length_labels
                        )
                        with torch.no_grad():
                            est_lengths = est_length_labels * self.n_frame_per_class
                            frame_error = masked_l1_loss(
                                est_lengths=est_lengths, tar_lengths=tar_lengths
                            )
                            sec_error = frame_error * 256 / 24000

                        log_dict = {
                            'loss': loss.item(),
                            'loss_CE': loss.item(),
                            'sec_error': sec_error.item(),
                            'lr': self.scheduler.get_last_lr()[0]
                         }

                    elif self.loss_fn == 'L1_and_CE':
                        tar_length_labels = (tar_lengths // self.n_frame_per_class) \
                            .clamp(min=0, max=self.n_class-1) # [0, 1, ..., n_class-1]
                        est_length_logtis = predictions
                        est_length_1hots = F.gumbel_softmax(
                            est_length_logtis, tau=self.gumbel_tau, hard=True, dim=-1
                        )
                        length_values = torch.arange(
                            self.n_class, device=est_length_1hots.device
                        ).float() * self.n_frame_per_class
                        est_lengths = (est_length_1hots * length_values).sum(-1)

                        loss_CE = masked_cross_entropy_loss(
                            est_length_logits=est_length_logtis, tar_length_labels=tar_length_labels
                        )

                        loss_L1 =  masked_l1_loss(
                            est_lengths=est_lengths, tar_lengths=tar_lengths
                        )

                        loss = loss_CE + self.lambda_L1 * loss_L1

                        with torch.no_grad():
                            frame_error = loss_L1
                            sec_error = frame_error * 256 / 24000

                        log_dict = {
                            'loss': loss.item(),
                            'loss_L1': loss_L1.item(),
                            'loss_CE': loss_CE.item(),
                            'sec_error': sec_error.item(),
                            'lr': self.scheduler.get_last_lr()[0]
                        }

                    else:
                        raise NotImplementedError(self.loss_fn)

                    self.accelerator.backward(loss)

                    if self.max_grad_norm > 0 and self.accelerator.sync_gradients:
                        self.accelerator.clip_grad_norm_(self.model.parameters(), self.max_grad_norm)

                    self.optimizer.step()
                    self.scheduler.step()
                    self.optimizer.zero_grad()

                global_step += 1

                if self.accelerator.is_local_main_process:
                    self.accelerator.log(log_dict, step=global_step)
                progress_bar.set_postfix(step=str(global_step), loss=loss.item())

                if global_step % (self.save_per_updates * self.grad_accumulation_steps) == 0:
                    self.save_checkpoint(global_step)
                    # if self.log_samples and self.accelerator.is_local_main_process:
                    # Run validation at the end of each epoch (only on the main process)
                    if self.accelerator.is_local_main_process:
                        self.validate(valid_dataloader, global_step)
                # if global_step % self.last_per_steps == 0:
                #     self.save_checkpoint(global_step, last=True)

        self.save_checkpoint(global_step, last=True)
        self.accelerator.end_training()


# part of the code is borrowed from https://github.com/lawlict/ECAPA-TDNN

import torch
import torch.nn as nn
import torch.nn.functional as F
import torchaudio.transforms as trans
from ctcmodel import ConformerCTC
# from ctcmodel_nopool import ConformerCTC as ConformerCTCNoPool
from pathlib import Path

''' Res2Conv1d + BatchNorm1d + ReLU
'''

class Res2Conv1dReluBn(nn.Module):
    '''
    in_channels == out_channels == channels
    '''

    def __init__(self, channels, kernel_size=1, stride=1, padding=0, dilation=1, bias=True, scale=4):
        super().__init__()
        assert channels % scale == 0, "{} % {} != 0".format(channels, scale)
        self.scale = scale
        self.width = channels // scale
        self.nums = scale if scale == 1 else scale - 1

        self.convs = []
        self.bns = []
        for i in range(self.nums):
            self.convs.append(nn.Conv1d(self.width, self.width, kernel_size, stride, padding, dilation, bias=bias))
            self.bns.append(nn.BatchNorm1d(self.width))
        self.convs = nn.ModuleList(self.convs)
        self.bns = nn.ModuleList(self.bns)

    def forward(self, x):
        out = []
        spx = torch.split(x, self.width, 1)
        for i in range(self.nums):
            if i == 0:
                sp = spx[i]
            else:
                sp = sp + spx[i]
            # Order: conv -> relu -> bn
            sp = self.convs[i](sp)
            sp = self.bns[i](F.relu(sp))
            out.append(sp)
        if self.scale != 1:
            out.append(spx[self.nums])
        out = torch.cat(out, dim=1)

        return out

''' Conv1d + BatchNorm1d + ReLU
'''

class Conv1dReluBn(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size=1, stride=1, padding=0, dilation=1, bias=True):
        super().__init__()
        self.conv = nn.Conv1d(in_channels, out_channels, kernel_size, stride, padding, dilation, bias=bias)
        self.bn = nn.BatchNorm1d(out_channels)

    def forward(self, x):
        return self.bn(F.relu(self.conv(x)))

''' The SE connection of 1D case.
'''

class SE_Connect(nn.Module):
    def __init__(self, channels, se_bottleneck_dim=128):
        super().__init__()
        self.linear1 = nn.Linear(channels, se_bottleneck_dim)
        self.linear2 = nn.Linear(se_bottleneck_dim, channels)

    def forward(self, x):
        out = x.mean(dim=2)
        out = F.relu(self.linear1(out))
        out = torch.sigmoid(self.linear2(out))
        out = x * out.unsqueeze(2)

        return out

''' SE-Res2Block of the ECAPA-TDNN architecture.
'''

class SE_Res2Block(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size, stride, padding, dilation, scale, se_bottleneck_dim):
        super().__init__()
        self.Conv1dReluBn1 = Conv1dReluBn(in_channels, out_channels, kernel_size=1, stride=1, padding=0)
        self.Res2Conv1dReluBn = Res2Conv1dReluBn(out_channels, kernel_size, stride, padding, dilation, scale=scale)
        self.Conv1dReluBn2 = Conv1dReluBn(out_channels, out_channels, kernel_size=1, stride=1, padding=0)
        self.SE_Connect = SE_Connect(out_channels, se_bottleneck_dim)

        self.shortcut = None
        if in_channels != out_channels:
            self.shortcut = nn.Conv1d(
                in_channels=in_channels,
                out_channels=out_channels,
                kernel_size=1,
            )

    def forward(self, x):
        residual = x
        if self.shortcut:
            residual = self.shortcut(x)

        x = self.Conv1dReluBn1(x)
        x = self.Res2Conv1dReluBn(x)
        x = self.Conv1dReluBn2(x)
        x = self.SE_Connect(x)

        return x + residual

''' Attentive weighted mean and standard deviation pooling.
'''

class AttentiveStatsPool(nn.Module):
    def __init__(self, in_dim, attention_channels=128, global_context_att=False):
        super().__init__()
        self.global_context_att = global_context_att

        # Use Conv1d with stride == 1 rather than Linear, then we don't need to transpose inputs.
        if global_context_att:
            self.linear1 = nn.Conv1d(in_dim * 3, attention_channels, kernel_size=1)  # equals W and b in the paper
        else:
            self.linear1 = nn.Conv1d(in_dim, attention_channels, kernel_size=1)  # equals W and b in the paper
        self.linear2 = nn.Conv1d(attention_channels, in_dim, kernel_size=1)  # equals V and k in the paper

    def forward(self, x):

        if self.global_context_att:
            context_mean = torch.mean(x, dim=-1, keepdim=True).expand_as(x)
            context_std = torch.sqrt(torch.var(x, dim=-1, keepdim=True) + 1e-10).expand_as(x)
            x_in = torch.cat((x, context_mean, context_std), dim=1)
        else:
            x_in = x

        # DON'T use ReLU here! In experiments, I find ReLU hard to converge.
        alpha = torch.tanh(self.linear1(x_in))
        # alpha = F.relu(self.linear1(x_in))
        alpha = torch.softmax(self.linear2(alpha), dim=2)
        mean = torch.sum(alpha * x, dim=2)
        residuals = torch.sum(alpha * (x ** 2), dim=2) - mean ** 2
        std = torch.sqrt(residuals.clamp(min=1e-9))
        return torch.cat([mean, std], dim=1)

class ECAPA_TDNN(nn.Module):
    def __init__(self, channels=512, emb_dim=512,
        global_context_att=False, use_fp16=True,
        ctc_cls=ConformerCTC,
        ctc_path='/data4/F5TTS/ckpts/F5TTS_norm_ASR_vocos_pinyin_Emilia_ZH_EN/model_last.pt',
        ctc_args={'vocab_size': 2545, 'mel_dim': 100, 'num_heads': 8, 'd_hid': 512, 'nlayers': 6},
        ctc_no_grad=False
    ):
        super().__init__()
        if ctc_path != None:
            ctc_path = Path(ctc_path)
            model = ctc_cls(**ctc_args)
            state_dict = torch.load(ctc_path, map_location='cpu')
            model.load_state_dict(state_dict['model_state_dict'])
            print(f"Initialized pretrained ConformerCTC backbone from {ctc_path}.")
        else:
            raise ValueError(ctc_path)

        self.ctc_model = model
        self.ctc_model.out.requires_grad_(False)

        if ctc_cls == ConformerCTC:
            self.feat_num = ctc_args['nlayers'] + 2 + 1
        # elif ctc_cls == ConformerCTCNoPool:
        #     self.feat_num = ctc_args['nlayers'] + 1
        else:
            raise ValueError(ctc_cls)
        feat_dim = ctc_args['d_hid']

        self.emb_dim = emb_dim

        self.feature_weight = nn.Parameter(torch.zeros(self.feat_num))
        self.instance_norm = nn.InstanceNorm1d(feat_dim)

        # self.channels = [channels] * 4 + [channels * 3]
        self.channels = [channels] * 4 + [1536]

        self.layer1 = Conv1dReluBn(feat_dim, self.channels[0], kernel_size=5, padding=2)
        self.layer2 = SE_Res2Block(self.channels[0], self.channels[1], kernel_size=3, stride=1, padding=2, dilation=2, scale=8, se_bottleneck_dim=128)
        self.layer3 = SE_Res2Block(self.channels[1], self.channels[2], kernel_size=3, stride=1, padding=3, dilation=3, scale=8, se_bottleneck_dim=128)
        self.layer4 = SE_Res2Block(self.channels[2], self.channels[3], kernel_size=3, stride=1, padding=4, dilation=4, scale=8, se_bottleneck_dim=128)

        # self.conv = nn.Conv1d(self.channels[-1], self.channels[-1], kernel_size=1)
        cat_channels = channels * 3
        self.conv = nn.Conv1d(cat_channels, self.channels[-1], kernel_size=1)
        self.pooling = AttentiveStatsPool(self.channels[-1], attention_channels=128, global_context_att=global_context_att)
        self.bn = nn.BatchNorm1d(self.channels[-1] * 2)
        self.linear = nn.Linear(self.channels[-1] * 2, emb_dim)

        if ctc_no_grad:
            for param in self.ctc_model.parameters():
                param.requires_grad = False
            self.ctc_model = self.ctc_model.eval()
        else:
            self.ctc_model = self.ctc_model.train()
        self.ctc_no_grad = ctc_no_grad
        print('ctc_no_grad: ', self.ctc_no_grad)

    def forward(self, latent, input_lengths,  return_asr=False):
        if self.ctc_no_grad:
            with torch.no_grad():
                asr, h = self.ctc_model(latent, input_lengths)
        else:
            asr, h = self.ctc_model(latent, input_lengths)

        x = torch.stack(h, dim=0)
        norm_weights = F.softmax(self.feature_weight, dim=-1).unsqueeze(-1).unsqueeze(-1).unsqueeze(-1)
        x = (norm_weights * x).sum(dim=0)
        x = x + 1e-6
        # x = torch.transpose(x, 1, 2) + 1e-6

        x = self.instance_norm(x)
        # x = torch.transpose(x, 1, 2)

        out1 = self.layer1(x)
        out2 = self.layer2(out1)
        out3 = self.layer3(out2)
        out4 = self.layer4(out3)

        out = torch.cat([out2, out3, out4], dim=1)
        out = F.relu(self.conv(out))
        out = self.bn(self.pooling(out))
        out = self.linear(out)

        if return_asr:
            return out, asr
        return out

if __name__ == "__main__":
    from diffspeech.ldm.model import DiT
    from diffspeech.data.collate import get_mask_from_lengths
    from diffspeech.tools.text.vocab import IPA

    bsz = 3

    # Sample ipa
    ipa_lens = torch.randint(10, 50, (bsz,)).cuda()
    ipa_mask = get_mask_from_lengths(ipa_lens).cuda()
    ipa = torch.randint(0, len(IPA.vocab), (bsz, ipa_mask.size(-1))).cuda()

    # Sample latent
    latent_lens = torch.randint(50, 250, (bsz,)).cuda()
    latent_mask = get_mask_from_lengths(latent_lens).cuda()
    latent = torch.randn(bsz, latent_mask.size(-1), 64).cuda()

    # Sample prompt
    prompt_mask = get_mask_from_lengths(
        (latent_lens * 0.25).long(), max_len=latent_mask.size(-1)
    ).cuda()
    prompt_latent = latent * prompt_mask.unsqueeze(-1)

    model = ECAPA_TDNN(emb_dim=512).cuda()

    emb = model(latent, latent_mask.sum(axis=-1))

    print(emb.shape)


# part of the code is borrowed from https://github.com/lawlict/ECAPA-TDNN

import torch
import torch.nn as nn
import torch.nn.functional as F
import torchaudio.transforms as trans
from ctcmodel import ConformerCTC
# from ctcmodel_nopool import ConformerCTC as ConformerCTCNoPool
from pathlib import Path

''' Res2Conv1d + BatchNorm1d + ReLU
'''

class Res2Conv1dReluBn(nn.Module):
    '''
    in_channels == out_channels == channels
    '''

    def __init__(self, channels, kernel_size=1, stride=1, padding=0, dilation=1, bias=True, scale=4):
        super().__init__()
        assert channels % scale == 0, "{} % {} != 0".format(channels, scale)
        self.scale = scale
        self.width = channels // scale
        self.nums = scale if scale == 1 else scale - 1

        self.convs = []
        self.bns = []
        for i in range(self.nums):
            self.convs.append(nn.Conv1d(self.width, self.width, kernel_size, stride, padding, dilation, bias=bias))
            self.bns.append(nn.BatchNorm1d(self.width))
        self.convs = nn.ModuleList(self.convs)
        self.bns = nn.ModuleList(self.bns)

    def forward(self, x):
        out = []
        spx = torch.split(x, self.width, 1)
        for i in range(self.nums):
            if i == 0:
                sp = spx[i]
            else:
                sp = sp + spx[i]
            # Order: conv -> relu -> bn
            sp = self.convs[i](sp)
            sp = self.bns[i](F.relu(sp))
            out.append(sp)
        if self.scale != 1:
            out.append(spx[self.nums])
        out = torch.cat(out, dim=1)

        return out

''' Conv1d + BatchNorm1d + ReLU
'''

class Conv1dReluBn(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size=1, stride=1, padding=0, dilation=1, bias=True):
        super().__init__()
        self.conv = nn.Conv1d(in_channels, out_channels, kernel_size, stride, padding, dilation, bias=bias)
        self.bn = nn.BatchNorm1d(out_channels)

    def forward(self, x):
        return self.bn(F.relu(self.conv(x)))

''' The SE connection of 1D case.
'''


        for i in range(self.nums):
            if i == 0:
                sp = spx[i]
            else:
                sp = sp + spx[i]
            # Order: conv -> relu -> bn
            sp = self.convs[i](sp)
            sp = self.bns[i](F.relu(sp))
            out.append(sp)
        if self.scale != 1:
            out.append(spx[self.nums])
        out = torch.cat(out, dim=1)

        return out

''' Conv1d + BatchNorm1d + ReLU
'''

class Conv1dReluBn(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size=1, stride=1, padding=0, dilation=1, bias=True):
        super().__init__()
        self.conv = nn.Conv1d(in_channels, out_channels, kernel_size, stride, padding, dilation, bias=bias)
        self.bn = nn.BatchNorm1d(out_channels)

    def forward(self, x):
        return self.bn(F.relu(self.conv(x)))

''' The SE connection of 1D case.
'''

class SE_Connect(nn.Module):
    def __init__(self, channels, se_bottleneck_dim=128):
        super().__init__()
        self.linear1 = nn.Linear(channels, se_bottleneck_dim)
        self.linear2 = nn.Linear(se_bottleneck_dim, channels)

    def forward(self, x):
        out = x.mean(dim=2)
        out = F.relu(self.linear1(out))
        out = torch.sigmoid(self.linear2(out))
        out = x * out.unsqueeze(2)

        return out


''' Conv1d + BatchNorm1d + ReLU
'''

class Conv1dReluBn(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size=1, stride=1, padding=0, dilation=1, bias=True):
        super().__init__()
        self.conv = nn.Conv1d(in_channels, out_channels, kernel_size, stride, padding, dilation, bias=bias)
        self.bn = nn.BatchNorm1d(out_channels)

    def forward(self, x):
        return self.bn(F.relu(self.conv(x)))

''' The SE connection of 1D case.
'''

class SE_Connect(nn.Module):
    def __init__(self, channels, se_bottleneck_dim=128):
        super().__init__()
        self.linear1 = nn.Linear(channels, se_bottleneck_dim)
        self.linear2 = nn.Linear(se_bottleneck_dim, channels)

    def forward(self, x):
        out = x.mean(dim=2)
        out = F.relu(self.linear1(out))
        out = torch.sigmoid(self.linear2(out))
        out = x * out.unsqueeze(2)

        return out

''' SE-Res2Block of the ECAPA-TDNN architecture.
'''

class SE_Res2Block(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size, stride, padding, dilation, scale, se_bottleneck_dim):
        super().__init__()
        self.Conv1dReluBn1 = Conv1dReluBn(in_channels, out_channels, kernel_size=1, stride=1, padding=0)
        self.Res2Conv1dReluBn = Res2Conv1dReluBn(out_channels, kernel_size, stride, padding, dilation, scale=scale)
        self.Conv1dReluBn2 = Conv1dReluBn(out_channels, out_channels, kernel_size=1, stride=1, padding=0)
        self.SE_Connect = SE_Connect(out_channels, se_bottleneck_dim)

        self.shortcut = None
        if in_channels != out_channels:
            self.shortcut = nn.Conv1d(
                in_channels=in_channels,
                out_channels=out_channels,
                kernel_size=1,
            )


class SE_Connect(nn.Module):
    def __init__(self, channels, se_bottleneck_dim=128):
        super().__init__()
        self.linear1 = nn.Linear(channels, se_bottleneck_dim)
        self.linear2 = nn.Linear(se_bottleneck_dim, channels)

    def forward(self, x):
        out = x.mean(dim=2)
        out = F.relu(self.linear1(out))
        out = torch.sigmoid(self.linear2(out))
        out = x * out.unsqueeze(2)

        return out

''' SE-Res2Block of the ECAPA-TDNN architecture.
'''

class SE_Res2Block(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size, stride, padding, dilation, scale, se_bottleneck_dim):
        super().__init__()
        self.Conv1dReluBn1 = Conv1dReluBn(in_channels, out_channels, kernel_size=1, stride=1, padding=0)
        self.Res2Conv1dReluBn = Res2Conv1dReluBn(out_channels, kernel_size, stride, padding, dilation, scale=scale)
        self.Conv1dReluBn2 = Conv1dReluBn(out_channels, out_channels, kernel_size=1, stride=1, padding=0)
        self.SE_Connect = SE_Connect(out_channels, se_bottleneck_dim)

        self.shortcut = None
        if in_channels != out_channels:
            self.shortcut = nn.Conv1d(
                in_channels=in_channels,
                out_channels=out_channels,
                kernel_size=1,
            )

    def forward(self, x):
        residual = x
        if self.shortcut:
            residual = self.shortcut(x)

        x = self.Conv1dReluBn1(x)
        x = self.Res2Conv1dReluBn(x)
        x = self.Conv1dReluBn2(x)
        x = self.SE_Connect(x)

        return x + residual

''' Attentive weighted mean and standard deviation pooling.
'''

class AttentiveStatsPool(nn.Module):
    def __init__(self, in_dim, attention_channels=128, global_context_att=False):
        super().__init__()
        self.global_context_att = global_context_att

        # Use Conv1d with stride == 1 rather than Linear, then we don't need to transpose inputs.
        if global_context_att:
            self.linear1 = nn.Conv1d(in_dim * 3, attention_channels, kernel_size=1)  # equals W and b in the paper
        else:
            self.linear1 = nn.Conv1d(in_dim, attention_channels, kernel_size=1)  # equals W and b in the paper
        self.linear2 = nn.Conv1d(attention_channels, in_dim, kernel_size=1)  # equals V and k in the paper

    def forward(self, x):

        if self.global_context_att:


AttentiveStatsPool

                out_channels=out_channels,
                kernel_size=1,
            )

    def forward(self, x):
        residual = x
        if self.shortcut:
            residual = self.shortcut(x)

        x = self.Conv1dReluBn1(x)
        x = self.Res2Conv1dReluBn(x)
        x = self.Conv1dReluBn2(x)
        x = self.SE_Connect(x)

        return x + residual

''' Attentive weighted mean and standard deviation pooling.
'''

class AttentiveStatsPool(nn.Module):
    def __init__(self, in_dim, attention_channels=128, global_context_att=False):
        super().__init__()
        self.global_context_att = global_context_att

        # Use Conv1d with stride == 1 rather than Linear, then we don't need to transpose inputs.
        if global_context_att:
            self.linear1 = nn.Conv1d(in_dim * 3, attention_channels, kernel_size=1)  # equals W and b in the paper
        else:
            self.linear1 = nn.Conv1d(in_dim, attention_channels, kernel_size=1)  # equals W and b in the paper
        self.linear2 = nn.Conv1d(attention_channels, in_dim, kernel_size=1)  # equals V and k in the paper

    def forward(self, x):

        if self.global_context_att:
            context_mean = torch.mean(x, dim=-1, keepdim=True).expand_as(x)
            context_std = torch.sqrt(torch.var(x, dim=-1, keepdim=True) + 1e-10).expand_as(x)
            x_in = torch.cat((x, context_mean, context_std), dim=1)
        else:
            x_in = x

        # DON'T use ReLU here! In experiments, I find ReLU hard to converge.
        alpha = torch.tanh(self.linear1(x_in))
        # alpha = F.relu(self.linear1(x_in))
        alpha = torch.softmax(self.linear2(alpha), dim=2)
        mean = torch.sum(alpha * x, dim=2)
        residuals = torch.sum(alpha * (x ** 2), dim=2) - mean ** 2
        std = torch.sqrt(residuals.clamp(min=1e-9))
        return torch.cat([mean, std], dim=1)

class ECAPA_TDNN(nn.Module):
    def __init__(self, channels=512, emb_dim=512,
        global_context_att=False, use_fp16=True,
        ctc_cls=ConformerCTC,
        ctc_path='/data4/F5TTS/ckpts/F5TTS_norm_ASR_vocos_pinyin_Emilia_ZH_EN/model_last.pt',
        ctc_args={'vocab_size': 2545, 'mel_dim': 100, 'num_heads': 8, 'd_hid': 512, 'nlayers': 6},
        ctc_no_grad=False
    ):
        super().__init__()
        if ctc_path != None:
            ctc_path = Path(ctc_path)
            model = ctc_cls(**ctc_args)
            state_dict = torch.load(ctc_path, map_location='cpu')
            model.load_state_dict(state_dict['model_state_dict'])
            print(f"Initialized pretrained ConformerCTC backbone from {ctc_path}.")
        else:
            raise ValueError(ctc_path)


    def forward(self, x):

        if self.global_context_att:
            context_mean = torch.mean(x, dim=-1, keepdim=True).expand_as(x)
            context_std = torch.sqrt(torch.var(x, dim=-1, keepdim=True) + 1e-10).expand_as(x)
            x_in = torch.cat((x, context_mean, context_std), dim=1)
        else:
            x_in = x

        # DON'T use ReLU here! In experiments, I find ReLU hard to converge.
        alpha = torch.tanh(self.linear1(x_in))
        # alpha = F.relu(self.linear1(x_in))
        alpha = torch.softmax(self.linear2(alpha), dim=2)
        mean = torch.sum(alpha * x, dim=2)
        residuals = torch.sum(alpha * (x ** 2), dim=2) - mean ** 2
        std = torch.sqrt(residuals.clamp(min=1e-9))
        return torch.cat([mean, std], dim=1)

class ECAPA_TDNN(nn.Module):
    def __init__(self, channels=512, emb_dim=512,
        global_context_att=False, use_fp16=True,
        ctc_cls=ConformerCTC,
        ctc_path='/data4/F5TTS/ckpts/F5TTS_norm_ASR_vocos_pinyin_Emilia_ZH_EN/model_last.pt',
        ctc_args={'vocab_size': 2545, 'mel_dim': 100, 'num_heads': 8, 'd_hid': 512, 'nlayers': 6},
        ctc_no_grad=False
    ):
        super().__init__()
        if ctc_path != None:
            ctc_path = Path(ctc_path)
            model = ctc_cls(**ctc_args)
            state_dict = torch.load(ctc_path, map_location='cpu')
            model.load_state_dict(state_dict['model_state_dict'])
            print(f"Initialized pretrained ConformerCTC backbone from {ctc_path}.")
        else:
            raise ValueError(ctc_path)

        self.ctc_model = model
        self.ctc_model.out.requires_grad_(False)

        if ctc_cls == ConformerCTC:
            self.feat_num = ctc_args['nlayers'] + 2 + 1
        # elif ctc_cls == ConformerCTCNoPool:
        #     self.feat_num = ctc_args['nlayers'] + 1
        else:
            raise ValueError(ctc_cls)
        feat_dim = ctc_args['d_hid']

        self.emb_dim = emb_dim

        self.feature_weight = nn.Parameter(torch.zeros(self.feat_num))
        self.instance_norm = nn.InstanceNorm1d(feat_dim)

        # self.channels = [channels] * 4 + [channels * 3]
        self.channels = [channels] * 4 + [1536]

        self.layer1 = Conv1dReluBn(feat_dim, self.channels[0], kernel_size=5, padding=2)
        self.layer2 = SE_Res2Block(self.channels[0], self.channels[1], kernel_size=3, stride=1, padding=2, dilation=2, scale=8, se_bottleneck_dim=128)
        self.layer3 = SE_Res2Block(self.channels[1], self.channels[2], kernel_size=3, stride=1, padding=3, dilation=3, scale=8, se_bottleneck_dim=128)
        self.layer4 = SE_Res2Block(self.channels[2], self.channels[3], kernel_size=3, stride=1, padding=4, dilation=4, scale=8, se_bottleneck_dim=128)

        # self.conv = nn.Conv1d(self.channels[-1], self.channels[-1], kernel_size=1)
        cat_channels = channels * 3
        self.conv = nn.Conv1d(cat_channels, self.channels[-1], kernel_size=1)
        self.pooling = AttentiveStatsPool(self.channels[-1], attention_channels=128, global_context_att=global_context_att)
        self.bn = nn.BatchNorm1d(self.channels[-1] * 2)
        self.linear = nn.Linear(self.channels[-1] * 2, emb_dim)

        if ctc_no_grad:
            for param in self.ctc_model.parameters():
                param.requires_grad = False
            self.ctc_model = self.ctc_model.eval()
        else:
            self.ctc_model = self.ctc_model.train()
        self.ctc_no_grad = ctc_no_grad
        print('ctc_no_grad: ', self.ctc_no_grad)

    def forward(self, latent, input_lengths,  return_asr=False):
        if self.ctc_no_grad:
            with torch.no_grad():
                asr, h = self.ctc_model(latent, input_lengths)
        else:
            asr, h = self.ctc_model(latent, input_lengths)

        x = torch.stack(h, dim=0)
        norm_weights = F.softmax(self.feature_weight, dim=-1).unsqueeze(-1).unsqueeze(-1).unsqueeze(-1)
        x = (norm_weights * x).sum(dim=0)
        x = x + 1e-6
        # x = torch.transpose(x, 1, 2) + 1e-6

        x = self.instance_norm(x)
        # x = torch.transpose(x, 1, 2)

        out1 = self.layer1(x)
        out2 = self.layer2(out1)
        out3 = self.layer3(out2)
        out4 = self.layer4(out3)

        out = torch.cat([out2, out3, out4], dim=1)
        out = F.relu(self.conv(out))
        out = self.bn(self.pooling(out))
        out = self.linear(out)

        if return_asr:
            return out, asr
        return out

if __name__ == "__main__":
    from diffspeech.ldm.model import DiT
    from diffspeech.data.collate import get_mask_from_lengths
    from diffspeech.tools.text.vocab import IPA

    bsz = 3

    # Sample ipa
    ipa_lens = torch.randint(10, 50, (bsz,)).cuda()
    ipa_mask = get_mask_from_lengths(ipa_lens).cuda()
    ipa = torch.randint(0, len(IPA.vocab), (bsz, ipa_mask.size(-1))).cuda()

    # Sample latent
    latent_lens = torch.randint(50, 250, (bsz,)).cuda()
    latent_mask = get_mask_from_lengths(latent_lens).cuda()
    latent = torch.randn(bsz, latent_mask.size(-1), 64).cuda()

    # Sample prompt
    prompt_mask = get_mask_from_lengths(


    def forward(self, x):

        if self.global_context_att:
            context_mean = torch.mean(x, dim=-1, keepdim=True).expand_as(x)
            context_std = torch.sqrt(torch.var(x, dim=-1, keepdim=True) + 1e-10).expand_as(x)
            x_in = torch.cat((x, context_mean, context_std), dim=1)
        else:
            x_in = x

        # DON'T use ReLU here! In experiments, I find ReLU hard to converge.
        alpha = torch.tanh(self.linear1(x_in))
        # alpha = F.relu(self.linear1(x_in))
        alpha = torch.softmax(self.linear2(alpha), dim=2)
        mean = torch.sum(alpha * x, dim=2)
        residuals = torch.sum(alpha * (x ** 2), dim=2) - mean ** 2
        std = torch.sqrt(residuals.clamp(min=1e-9))
        return torch.cat([mean, std], dim=1)

class ECAPA_TDNN(nn.Module):
    def __init__(self, channels=512, emb_dim=512,
        global_context_att=False, use_fp16=True,
        ctc_cls=ConformerCTC,
        ctc_path='/data4/F5TTS/ckpts/F5TTS_norm_ASR_vocos_pinyin_Emilia_ZH_EN/model_last.pt',
        ctc_args={'vocab_size': 2545, 'mel_dim': 100, 'num_heads': 8, 'd_hid': 512, 'nlayers': 6},
        ctc_no_grad=False
    ):
        super().__init__()
        if ctc_path != None:
            ctc_path = Path(ctc_path)
            model = ctc_cls(**ctc_args)
            state_dict = torch.load(ctc_path, map_location='cpu')
            model.load_state_dict(state_dict['model_state_dict'])
            print(f"Initialized pretrained ConformerCTC backbone from {ctc_path}.")
        else:
            raise ValueError(ctc_path)

        self.ctc_model = model
        self.ctc_model.out.requires_grad_(False)

        if ctc_cls == ConformerCTC:
            self.feat_num = ctc_args['nlayers'] + 2 + 1
        # elif ctc_cls == ConformerCTCNoPool:
        #     self.feat_num = ctc_args['nlayers'] + 1
        else:
            raise ValueError(ctc_cls)
        feat_dim = ctc_args['d_hid']

        self.emb_dim = emb_dim

        self.feature_weight = nn.Parameter(torch.zeros(self.feat_num))
        self.instance_norm = nn.InstanceNorm1d(feat_dim)

        # self.channels = [channels] * 4 + [channels * 3]
        self.channels = [channels] * 4 + [1536]

        self.layer1 = Conv1dReluBn(feat_dim, self.channels[0], kernel_size=5, padding=2)
        self.layer2 = SE_Res2Block(self.channels[0], self.channels[1], kernel_size=3, stride=1, padding=2, dilation=2, scale=8, se_bottleneck_dim=128)
        self.layer3 = SE_Res2Block(self.channels[1], self.channels[2], kernel_size=3, stride=1, padding=3, dilation=3, scale=8, se_bottleneck_dim=128)
        self.layer4 = SE_Res2Block(self.channels[2], self.channels[3], kernel_size=3, stride=1, padding=4, dilation=4, scale=8, se_bottleneck_dim=128)

        # self.conv = nn.Conv1d(self.channels[-1], self.channels[-1], kernel_size=1)
        cat_channels = channels * 3
        self.conv = nn.Conv1d(cat_channels, self.channels[-1], kernel_size=1)
        self.pooling = AttentiveStatsPool(self.channels[-1], attention_channels=128, global_context_att=global_context_att)
        self.bn = nn.BatchNorm1d(self.channels[-1] * 2)
        self.linear = nn.Linear(self.channels[-1] * 2, emb_dim)

        if ctc_no_grad:
            for param in self.ctc_model.parameters():
                param.requires_grad = False
            self.ctc_model = self.ctc_model.eval()
        else:
            self.ctc_model = self.ctc_model.train()
        self.ctc_no_grad = ctc_no_grad
        print('ctc_no_grad: ', self.ctc_no_grad)

    def forward(self, latent, input_lengths,  return_asr=False):
        if self.ctc_no_grad:
            with torch.no_grad():
                asr, h = self.ctc_model(latent, input_lengths)
        else:
            asr, h = self.ctc_model(latent, input_lengths)

        x = torch.stack(h, dim=0)
        norm_weights = F.softmax(self.feature_weight, dim=-1).unsqueeze(-1).unsqueeze(-1).unsqueeze(-1)
        x = (norm_weights * x).sum(dim=0)
        x = x + 1e-6
        # x = torch.transpose(x, 1, 2) + 1e-6

        x = self.instance_norm(x)
        # x = torch.transpose(x, 1, 2)

        out1 = self.layer1(x)
        out2 = self.layer2(out1)
        out3 = self.layer3(out2)


forward

        self.layer2 = SE_Res2Block(self.channels[0], self.channels[1], kernel_size=3, stride=1, padding=2, dilation=2, scale=8, se_bottleneck_dim=128)
        self.layer3 = SE_Res2Block(self.channels[1], self.channels[2], kernel_size=3, stride=1, padding=3, dilation=3, scale=8, se_bottleneck_dim=128)
        self.layer4 = SE_Res2Block(self.channels[2], self.channels[3], kernel_size=3, stride=1, padding=4, dilation=4, scale=8, se_bottleneck_dim=128)

        # self.conv = nn.Conv1d(self.channels[-1], self.channels[-1], kernel_size=1)
        cat_channels = channels * 3
        self.conv = nn.Conv1d(cat_channels, self.channels[-1], kernel_size=1)
        self.pooling = AttentiveStatsPool(self.channels[-1], attention_channels=128, global_context_att=global_context_att)
        self.bn = nn.BatchNorm1d(self.channels[-1] * 2)
        self.linear = nn.Linear(self.channels[-1] * 2, emb_dim)

        if ctc_no_grad:
            for param in self.ctc_model.parameters():
                param.requires_grad = False
            self.ctc_model = self.ctc_model.eval()
        else:
            self.ctc_model = self.ctc_model.train()
        self.ctc_no_grad = ctc_no_grad
        print('ctc_no_grad: ', self.ctc_no_grad)

    def forward(self, latent, input_lengths,  return_asr=False):
        if self.ctc_no_grad:
            with torch.no_grad():
                asr, h = self.ctc_model(latent, input_lengths)
        else:
            asr, h = self.ctc_model(latent, input_lengths)

        x = torch.stack(h, dim=0)
        norm_weights = F.softmax(self.feature_weight, dim=-1).unsqueeze(-1).unsqueeze(-1).unsqueeze(-1)
        x = (norm_weights * x).sum(dim=0)
        x = x + 1e-6
        # x = torch.transpose(x, 1, 2) + 1e-6

        x = self.instance_norm(x)
        # x = torch.transpose(x, 1, 2)

        out1 = self.layer1(x)
        out2 = self.layer2(out1)
        out3 = self.layer3(out2)
        out4 = self.layer4(out3)

        out = torch.cat([out2, out3, out4], dim=1)
        out = F.relu(self.conv(out))
        out = self.bn(self.pooling(out))
        out = self.linear(out)

        if return_asr:
            return out, asr
        return out

if __name__ == "__main__":
    from diffspeech.ldm.model import DiT
    from diffspeech.data.collate import get_mask_from_lengths
    from diffspeech.tools.text.vocab import IPA

    bsz = 3

    # Sample ipa
    ipa_lens = torch.randint(10, 50, (bsz,)).cuda()
    ipa_mask = get_mask_from_lengths(ipa_lens).cuda()
    ipa = torch.randint(0, len(IPA.vocab), (bsz, ipa_mask.size(-1))).cuda()

    # Sample latent
    latent_lens = torch.randint(50, 250, (bsz,)).cuda()
    latent_mask = get_mask_from_lengths(latent_lens).cuda()
    latent = torch.randn(bsz, latent_mask.size(-1), 64).cuda()

    # Sample prompt
    prompt_mask = get_mask_from_lengths(


src.unimodel

from __future__ import annotations
from typing import Callable
from random import random

import contextlib

from torch import nn
import torch
import copy
import os

from f5_tts.model import DiT, UNetT
from pathlib import Path
from guidance_model import Guidance
from f5_tts.model.utils import (
    default,
    exists,
    list_str_to_idx,
    list_str_to_tensor,
    lens_to_mask,
    mask_from_frac_lengths,
    sample_consecutive_steps,
    sample_from_list,
)

class UniModel(nn.Module):
    def __init__(self,
                 model: DiT, # teacher model (dit model)
                 checkpoint_path: str = "",
                 second_time: bool = True,
                 use_fp16: bool = True,
                 real_guidance_scale: float = 2.0,
                 fake_guidance_scale: float = 0.0,
                 gen_cls_loss: bool = False,
                 sway_coeff: float = -1.0,
                 vocab_char_map: dict[str, int] | None = None,
                 frac_lengths_mask: tuple[float, float] = (0.7, 1.0)):

        super().__init__()

        if checkpoint_path != "":
            if "model_last.pt" in os.listdir(checkpoint_path):
                latest_checkpoint = "model_last.pt"
            else:
                latest_checkpoint = sorted(
                    [f for f in os.listdir(checkpoint_path) if f.endswith(".pt")],
                    key=lambda x: int("".join(filter(str.isdigit, x))),
                )[-1]
            checkpoint = torch.load(f"{checkpoint_path}/{latest_checkpoint}", weights_only=True, map_location="cpu")

            if "scale" in checkpoint:
                self.scale = checkpoint["scale"]
            else:
                self.scale = 1.0
            print(f"Loaded teacher model with scale: {self.scale}")

            if "step" in checkpoint:
                state = checkpoint["model_state_dict"]
            else:
                checkpoint["model_state_dict"] = {
                    k.replace("ema_model.", ""): v
                    for k, v in checkpoint["ema_model_state_dict"].items()
                    if k not in ["initted", "step"]
                }
                state = checkpoint["model_state_dict"]

            # only load the DiT module
            filtered_state_dict = {
                k.replace("transformer.", ""): v
                for k, v in state.items()
                if k.startswith("transformer.")
            }

            model.load_state_dict(filtered_state_dict, strict=False)
        else:
            self.scale = 1.0

        real_unet = copy.deepcopy(model)
        real_unet.time_embed2 = None

        fake_unet = copy.deepcopy(model)

        # Instantiate Guidance, which internally uses real_unet and fake_unet initialized from the teacher
        self.guidance_model = Guidance(
            real_unet=real_unet,
            fake_unet=fake_unet,
            use_fp16=use_fp16,
            real_guidance_scale=real_guidance_scale,
            fake_guidance_scale=fake_guidance_scale,
            gen_cls_loss=gen_cls_loss,
            sway_coeff=sway_coeff,
        )

        self.feedforward_model = copy.deepcopy(model) # initialize the student model
        self.feedforward_model.requires_grad_(True)
        self.feedforward_model.time_embed2 = None

        self.vocab_char_map = vocab_char_map
        self.frac_lengths_mask = frac_lengths_mask

        self.second_time = second_time # fake_unet.time_embed2 is not None

    def forward(self,
                inp: float["b n d"],  # mel
                text: int["b nt"] | list[str],
                *,
                lens: int["b"] | None = None,
                student_steps: list[int] = [0, 0.25, 0.5, 0.75],
                update_generator: bool = False,
    ):
        """
        Forward pass that routes to either generator_forward or guidance_forward
        in the Guidance class, depending on the arguments.

        Parameters:
        -----------
        generator_turn: bool
            If True, run the generator forward pass (distribution matching loss, etc.)
        guidance_turn: bool
            If True, run the guidance forward pass (fake loss, cls loss, etc.)
        data_dict: dict
            Input dictionary containing the necessary keys for the forward passes.
            Expected keys may include:
                "inp": Tensor (B, N, D) - input mel or latent
                "text": Tensor or list[str] - text input
                "rand_span_mask": Tensor (B, N) - boolean mask
                "real_data": dict with keys like:
                     "inp", "text", "rand_span_mask"

        Returns:
        --------
        loss_dict: dict[str, Tensor]
            Dictionary of losses.
        log_dict: dict[str, Tensor or float]
            Dictionary of logging tensors or values.
        """

        batch, seq_len, dtype, device = *inp.shape[:2], inp.dtype, inp.device

        # handle text as string
        if isinstance(text, list):
            if exists(self.vocab_char_map):
                text = list_str_to_idx(text, self.vocab_char_map).to(device)
            else:
                text = list_str_to_tensor(text).to(device)
            assert text.shape[0] == batch

        # lens and mask
        if not exists(lens):
            lens = torch.full((batch,), seq_len, device=device)

        mask = lens_to_mask(lens, length=seq_len)  # useless here, as collate_fn will pad to max length in batch

        # sample from the list of student steps
        time = sample_from_list(student_steps, batch).to(device)
        c_time, p_time = sample_consecutive_steps(student_steps)
        time = torch.ones_like(time) * c_time
        p_time = torch.ones_like(time) * p_time

        frac_lengths = torch.zeros((batch,), device=device).float().uniform_(*self.frac_lengths_mask)
        rand_span_mask = mask_from_frac_lengths(lens, frac_lengths)

        if exists(mask):
            rand_span_mask &= mask

        # # use generated output from previous step as input
        with torch.no_grad():
            x1 = inp
            x0 = torch.randn_like(x1)
            t = p_time.unsqueeze(-1).unsqueeze(-1)
            phi = (1 - t) * x0 + t * x1
            cond = torch.where(rand_span_mask[..., None], torch.zeros_like(x1), x1)

            pred = self.feedforward_model(
                x=phi,
                cond=cond,
                text=text,
                time=p_time,
                drop_audio_cond=False,
                drop_text=False # make sure the cfg=1
            ) # flow prediction

            # predicted mel spectrogram
            output = phi + (1 - t) * pred
            output[~rand_span_mask] = inp[~rand_span_mask]

        # forward diffusion
        x1 = output
        x0 = torch.randn_like(x1)
        t = time.unsqueeze(-1).unsqueeze(-1)
        phi = (1 - t) * x0 + t * x1
        cond = torch.where(rand_span_mask[..., None], torch.zeros_like(x1), x1)

        with torch.no_grad() if not update_generator else contextlib.nullcontext():
            pred = self.feedforward_model(
                x=phi,
                cond=cond,
                text=text,
                time=time,
                drop_audio_cond=False,
                drop_text=False # make sure no cfg is used
            )

            # predicted mel spectrogram
            output = phi + (1 - t) * pred
            output[~rand_span_mask] = inp[~rand_span_mask]

        if update_generator:
            generator_data_dict = {
                "inp": output,
                "text": text,
                "rand_span_mask": rand_span_mask,
                "second_time": time if self.second_time else None,
                "mse_loss": time.mean() == student_steps[-1].mean(),
                "real_data": {
                    "inp": inp,
                    "text": text,
                    "rand_span_mask": rand_span_mask
                }
            }

            # avoid any side effects of gradient accumulation
            # self.guidance_model.requires_grad_(False)
            # self.feedforward_model.requires_grad_(True)
            generator_loss_dict, generator_log_dict = self.guidance_model(
                generator_turn=True,
                guidance_turn=False,
                generator_data_dict=generator_data_dict,
                guidance_data_dict=None
            )

            generator_log_dict['ground_truth'] = x1
            generator_log_dict['generator_input'] = phi
            generator_log_dict['generator_output'] = output
            generator_log_dict['generator_cond'] = cond
            generator_log_dict['time'] = time

            return generator_loss_dict, generator_log_dict
        else:
            guidance_data_dict = {
                "inp": output.detach(),
                "text": text,
                "rand_span_mask": rand_span_mask,
                "second_time": time if self.second_time else None,
                "real_data": {
                    "inp": inp,
                    "text": text,
                    "rand_span_mask": rand_span_mask
                }
            }

            # avoid any side effects of gradient accumulation
            # self.feedforward_model.requires_grad_(False)
            # self.guidance_model.requires_grad_(True)
            guidance_loss_dict, guidance_log_dict = self.guidance_model(
                generator_turn=False,
                guidance_turn=True,
                generator_data_dict=None,
                guidance_data_dict=guidance_data_dict
            )
            # self.feedforward_model.requires_grad_(True)

            return guidance_loss_dict, guidance_log_dict

        # return guidance_loss_dict, guidance_log_dict, generator_loss_dict, generator_log_dict

if __name__ == "__main__":

    from f5_tts.model.utils import get_tokenizer
    from torch.utils.data import DataLoader, Dataset, SequentialSampler
    from f5_tts.model.dataset import load_dataset
    from f5_tts.model.dataset import DynamicBatchSampler, collate_fn

    bsz = 16

    tokenizer = "pinyin"  # 'pinyin', 'char', or 'custom'
    tokenizer_path = None  # if tokenizer = 'custom', define the path to the tokenizer you want to use (should be vocab.txt)
    dataset_name = "Emilia_ZH_EN"
    if tokenizer == "custom":
        tokenizer_path = tokenizer_path
    else:
        tokenizer_path = dataset_name
    vocab_char_map, vocab_size = get_tokenizer(tokenizer_path, tokenizer)

    dit = DiT(dim=1024, depth=22, heads=16, ff_mult=2, text_dim=512, conv_layers=4, text_num_embeds=vocab_size, mel_dim=100)

    model = UniModel(dit,
                     checkpoint_path="/data4/F5TTS/ckpts/F5TTS_Base_norm_flow_8GPU_vocos_pinyin_Emilia_ZH_EN",
                     gen_cls_loss=True,
                     vocab_char_map=vocab_char_map,
                     frac_lengths_mask=(0.7, 1.0)
                     ).cuda()

    # batch = next(iter(train_dataloader))
    # torch.save(batch, "batch.pt")
    batch = torch.load("batch.pt")
    inp, text, lens = batch["mel"].permute(0, 2, 1).cuda(), batch["text"], batch["mel_lengths"].cuda()

    # text = ["hello world"] * bsz
    # lens = torch.randint(1, 1000, (bsz,)).cuda()
    # inp = torch.randn(bsz, lens.max(), 100).cuda()
    with torch.autocast(device_type="cuda", dtype=torch.float16):
        num_student_step = 4

        guidance_loss_dict, guidance_log_dict = model(inp, text, lens=lens, update_generator=False, student_steps=(torch.linspace(0.0, 1.0, num_student_step + 1)[:-1]))

        generator_loss_dict, generator_log_dict = model(inp, text, lens=lens, update_generator=True, student_steps=(torch.linspace(0.0, 1.0, num_student_step + 1)[:-1]))

        print(guidance_loss_dict)
        print(generator_loss_dict)

        guidance_loss = 0
        guidance_loss += guidance_loss_dict["loss_fake_mean"]
        guidance_loss += guidance_loss_dict["guidance_cls_loss"]

        generator_loss = 0
        generator_loss += generator_loss_dict["loss_dm"]
        generator_loss += generator_loss_dict["loss_ctc"]
        generator_loss += generator_loss_dict["loss_sim"]
        generator_loss += generator_loss_dict["gen_cls_loss"]
        generator_loss += generator_loss_dict["loss_mse"]

        guidance_loss.backward()
        generator_loss.backward()