| import numpy as np |
| import torch |
| from torchvision import transforms |
| from tqdm import tqdm |
| from PIL import Image |
| import soundfile as sf |
| from mel_module import Mel |
|
|
| class Generator: |
| def __init__(self, config, unet, scheduler, vae, embedding, progress_callback=None): |
| self.config = config |
| self.unet = unet |
| self.scheduler = scheduler |
| self.vae = vae |
| self.embedding = embedding |
| self.progress_callback = progress_callback |
|
|
| def tensor_to_mel(self, tensor): |
| denormalize = transforms.Normalize( |
| mean=[-m/s for m, s in zip([0.5], [0.5])], |
| std=[1/s for s in [0.5]] |
| ) |
| dn_tensor= denormalize(tensor.detach().cpu()) |
| s = np.array(dn_tensor.squeeze())*255 |
| return Mel(spectrogram=s) |
|
|
| def generate(self): |
| with torch.no_grad(): |
| uncond_image = torch.zeros((1, 1, self.config.image_size, self.config.image_size), device=self.config.device) |
| mu, log_var = self.vae.encode(uncond_image) |
| uncond_latent = torch.cat((mu, log_var), dim=1) |
| uncond_latent = uncond_latent.unsqueeze(0) |
| print("uncond", uncond_latent.shape) |
|
|
| embeddings = torch.cat([uncond_latent, self.embedding]) |
|
|
| generator = torch.Generator(device=self.config.device) |
|
|
| noise = torch.randn( |
| (1, 1, self.config.image_size, self.config.image_size), |
| generator=generator, |
| device=self.config.device, |
| ) |
|
|
| total_steps = len(self.scheduler.timesteps) |
|
|
| for i, t in enumerate(self.progress_callback.tqdm(self.scheduler.timesteps)): |
| image_model_input = torch.cat([noise] * 2) |
| image_model_input = self.scheduler.scale_model_input(image_model_input, timestep=t) |
|
|
| with torch.no_grad(): |
| noise_pred = self.unet(image_model_input, t, encoder_hidden_states=embeddings).sample |
| noise_pred_uncond, noise_pred_img = noise_pred.chunk(2) |
| noise_pred = noise_pred_uncond + self.config.guidance_scale * (noise_pred_img - noise_pred_uncond) |
| noise = self.scheduler.step(noise_pred, t, noise).prev_sample |
|
|
| image_tensor = noise.squeeze(1) |
| mel = self.tensor_to_mel(image_tensor) |
| mel.save_audio() |
|
|
