File size: 2,305 Bytes
e599c74 6f849d1 e599c74 6f849d1 e599c74 6f849d1 e599c74 6f849d1 e599c74 6f849d1 e599c74 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 |
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) # [1, 512, 512]
mel = self.tensor_to_mel(image_tensor)
mel.save_audio()
|