Method diffusion_process

trainer-api.py:76–153

(
        self,
        duration,
        encoder_text_hidden_states,
        text_attention_mask,
        speaker_embds,
        lyric_token_ids,
        lyric_mask,
        random_generator=None,
        infer_steps=60,
        guidance_scale=15.0,
        omega_scale=10.0,
    )

        return last_hidden_states, attention_mask

    def diffusion_process(
        self,
        duration,
        encoder_text_hidden_states,
        text_attention_mask,
        speaker_embds,
        lyric_token_ids,
        lyric_mask,
        random_generator=None,
        infer_steps=60,
        guidance_scale=15.0,
        omega_scale=10.0,
    ):
        do_classifier_free_guidance = guidance_scale > 1.0
        device = encoder_text_hidden_states.device
        dtype = encoder_text_hidden_states.dtype
        bsz = encoder_text_hidden_states.shape[0]

        timesteps, num_inference_steps = retrieve_timesteps(
            self.scheduler, num_inference_steps=infer_steps, device=device
        )

        frame_length = int(duration * 44100 / 512 / 8)
        target_latents = randn_tensor(
            shape=(bsz, 8, 16, frame_length),
            generator=random_generator,
            device=device,
            dtype=dtype,
        )
        attention_mask = torch.ones(bsz, frame_length, device=device, dtype=dtype)

        if do_classifier_free_guidance:
            attention_mask = torch.cat([attention_mask] * 2, dim=0)
            encoder_text_hidden_states = torch.cat(
                [encoder_text_hidden_states, torch.zeros_like(encoder_text_hidden_states)],
                0,
            )
            text_attention_mask = torch.cat([text_attention_mask] * 2, dim=0)
            speaker_embds = torch.cat([speaker_embds, torch.zeros_like(speaker_embds)], 0)
            lyric_token_ids = torch.cat([lyric_token_ids, torch.zeros_like(lyric_token_ids)], 0)
            lyric_mask = torch.cat([lyric_mask, torch.zeros_like(lyric_mask)], 0)

        momentum_buffer = MomentumBuffer()

        for t in timesteps:
            latent_model_input = (
                torch.cat([target_latents] * 2) if do_classifier_free_guidance else target_latents
            )
            timestep = t.expand(latent_model_input.shape[0])
            with torch.no_grad():
                noise_pred = self.transformers(
                    hidden_states=latent_model_input,
                    attention_mask=attention_mask,
                    encoder_text_hidden_states=encoder_text_hidden_states,
                    text_attention_mask=text_attention_mask,
                    speaker_embeds=speaker_embds,
                    lyric_token_idx=lyric_token_ids,
                    lyric_mask=lyric_mask,