Spaces:
Sleeping
Sleeping
File size: 4,767 Bytes
4af0d61 | 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 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 | import torch
import numpy as np
from PIL import Image
import torch.nn.functional as F
class LouDiffusion:
"""Pipeline de diffusion DDPM optimisée pour Lou"""
def __init__(self, model, device='cuda', num_train_timesteps=1000):
self.model = model
self.device = device
self.num_train_timesteps = num_train_timesteps
# Beta schedule - cosine pour génération rapide
self.betas = self._cosine_beta_schedule(num_train_timesteps).to(device)
self.alphas = 1.0 - self.betas
self.alphas_cumprod = torch.cumprod(self.alphas, dim=0)
self.alphas_cumprod_prev = F.pad(self.alphas_cumprod[:-1], (1, 0), value=1.0)
self.sqrt_alphas_cumprod = torch.sqrt(self.alphas_cumprod)
self.sqrt_one_minus_alphas_cumprod = torch.sqrt(1.0 - self.alphas_cumprod)
def _cosine_beta_schedule(self, timesteps, s=0.008):
steps = timesteps + 1
x = torch.linspace(0, timesteps, steps)
alphas_cumprod = torch.cos(((x / timesteps) + s) / (1 + s) * torch.pi * 0.5) ** 2
alphas_cumprod = alphas_cumprod / alphas_cumprod[0]
betas = 1 - (alphas_cumprod[1:] / alphas_cumprod[:-1])
return torch.clip(betas, 0.0001, 0.9999)
def add_noise(self, x_start, t, noise=None):
if noise is None:
noise = torch.randn_like(x_start)
sqrt_alphas_cumprod_t = self.sqrt_alphas_cumprod[t][:, None, None, None]
sqrt_one_minus_alphas_cumprod_t = self.sqrt_one_minus_alphas_cumprod[t][:, None, None, None]
return sqrt_alphas_cumprod_t * x_start + sqrt_one_minus_alphas_cumprod_t * noise
@torch.no_grad()
def sample(self, shape, num_inference_steps=50, guidance_scale=7.5, prompt_embeds=None, generator=None):
"""Sampling avec classifier-free guidance"""
self.model.eval()
batch_size = shape[0]
# Scheduler pour sampling rapide
timesteps = torch.linspace(self.num_train_timesteps - 1, 0, num_inference_steps, dtype=torch.long, device=self.device)
# Bruit initial
x = torch.randn(shape, device=self.device, generator=generator)
for i, t in enumerate(timesteps):
t_batch = torch.full((batch_size,), t, device=self.device, dtype=torch.long)
# Prédiction conditionnée
noise_pred = self.model(x, t_batch)
# Guidance scale (simplifié)
if guidance_scale != 1.0:
noise_pred_uncond = self.model(x, t_batch)
noise_pred = noise_pred_uncond + guidance_scale * (noise_pred - noise_pred_uncond)
# DDPM step
alpha_t = self.alphas_cumprod[t]
alpha_t_prev = self.alphas_cumprod_prev[t] if t > 0 else torch.tensor(1.0, device=self.device)
pred_x0 = (x - torch.sqrt(1 - alpha_t) * noise_pred) / torch.sqrt(alpha_t)
pred_x0 = torch.clamp(pred_x0, -1, 1)
direction = torch.sqrt(1 - alpha_t_prev) * noise_pred
x = torch.sqrt(alpha_t_prev) * pred_x0 + direction
if i < len(timesteps) - 1:
noise = torch.randn_like(x) if i > 0 else torch.zeros_like(x)
x = x + torch.sqrt(self.betas[t]) * noise
return torch.clamp(x, -1, 1)
def training_step(self, images):
"""Un pas d'entraînement"""
self.model.train()
batch_size = images.shape[0]
device = images.device
# Timesteps aléatoires
t = torch.randint(0, self.num_train_timesteps, (batch_size,), device=device)
# Bruit
noise = torch.randn_like(images)
# Images bruitées
noisy_images = self.add_noise(images, t, noise)
# Prédiction
noise_pred = self.model(noisy_images, t)
# Loss MSE
loss = F.mse_loss(noise_pred, noise)
return loss
def tensor_to_pil(tensor):
"""Convertit un tensor [-1, 1] en PIL Image"""
tensor = (tensor + 1) / 2
tensor = torch.clamp(tensor, 0, 1)
tensor = tensor.cpu().permute(0, 2, 3, 1).numpy()
images = (tensor * 255).astype(np.uint8)
return [Image.fromarray(img) for img in images]
def pil_to_tensor(images, size=256):
"""Convertit PIL Images en tensor [-1, 1]"""
if not isinstance(images, list):
images = [images]
tensors = []
for img in images:
img = img.convert('RGB').resize((size, size))
arr = np.array(img).astype(np.float32) / 255.0
tensor = torch.from_numpy(arr).permute(2, 0, 1)
tensor = tensor * 2 - 1
tensors.append(tensor)
return torch.stack(tensors) |