Create model.py

model.py

@@ -0,0 +1,227 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from transformers import CLIPTextModel, CLIPTokenizer
+
+class TimeEmbedding(nn.Module):
+    def __init__(self, dim):
+        super().__init__()
+        self.dim = dim
+        half_dim = dim // 2
+        emb = torch.log(torch.tensor(10000.0)) / (half_dim - 1)
+        emb = torch.exp(torch.arange(half_dim, dtype=torch.float) * -emb)
+        self.register_buffer('emb', emb)
+    
+    def forward(self, time):
+        emb = time[:, None] * self.emb[None, :]
+        emb = torch.cat((torch.sin(emb), torch.cos(emb)), dim=-1)
+        return emb
+
+class AttentionBlock(nn.Module):
+    def __init__(self, channels, num_heads=4):
+        super().__init__()
+        self.num_heads = num_heads
+        self.scale = (channels // num_heads) ** -0.5
+        
+        self.norm = nn.GroupNorm(32, channels)
+        self.qkv = nn.Conv2d(channels, channels * 3, 1)
+        self.proj = nn.Conv2d(channels, channels, 1)
+    
+    def forward(self, x):
+        b, c, h, w = x.shape
+        qkv = self.qkv(self.norm(x))
+        q, k, v = qkv.chunk(3, dim=1)
+        
+        q = q.view(b, self.num_heads, -1, h * w).permute(0, 1, 3, 2)
+        k = k.view(b, self.num_heads, -1, h * w)
+        v = v.view(b, self.num_heads, -1, h * w)
+        
+        attn = torch.softmax((q @ k) * self.scale, dim=-1)
+        x = (attn @ v).permute(0, 1, 3, 2).reshape(b, -1, h, w)
+        return self.proj(x) + x
+
+class ResBlock(nn.Module):
+    def __init__(self, in_ch, out_ch, time_emb_dim, text_emb_dim, dropout=0.1):
+        super().__init__()
+        self.mlp = nn.Sequential(
+            nn.SiLU(),
+            nn.Linear(time_emb_dim + text_emb_dim, out_ch * 2)
+        
+        self.block1 = nn.Sequential(
+            nn.GroupNorm(32, in_ch),
+            nn.SiLU(),
+            nn.Conv2d(in_ch, out_ch, 3, padding=1))
+        
+        self.block2 = nn.Sequential(
+            nn.GroupNorm(32, out_ch),
+            nn.SiLU(),
+            nn.Dropout(dropout),
+            nn.Conv2d(out_ch, out_ch, 3, padding=1))
+        
+        self.res_conv = nn.Conv2d(in_ch, out_ch, 1) if in_ch != out_ch else nn.Identity()
+    
+    def forward(self, x, time_emb, text_emb):
+        emb = self.mlp(torch.cat([time_emb, text_emb], dim=-1))
+        scale, shift = torch.chunk(emb, 2, dim=1)
+        
+        h = self.block1(x)
+        h = h * (1 + scale[:, :, None, None]) + shift[:, :, None, None]
+        h = self.block2(h)
+        
+        return h + self.res_conv(x)
+
+class UNet(nn.Module):
+    def __init__(self, in_channels=3, out_channels=3, dim=64, dim_mults=(1, 2, 4, 8)):
+        super().__init__()
+        dims = [dim * m for m in dim_mults]
+        in_out = list(zip(dims[:-1], dims[1:]))
+        
+        # Time and text embeddings
+        self.time_mlp = nn.Sequential(
+            TimeEmbedding(dim),
+            nn.Linear(dim, dim * 4),
+            nn.SiLU(),
+            nn.Linear(dim * 4, dim))
+        
+        # Text conditioning
+        self.text_proj = nn.Linear(768, dim * 4)
+        
+        # Initial convolution
+        self.init_conv = nn.Conv2d(in_channels, dim, 3, padding=1)
+        
+        # Downsample blocks
+        self.downs = nn.ModuleList()
+        for ind, (in_dim, out_dim) in enumerate(in_out):
+            is_last = ind >= (len(in_out) - 1)
+            self.downs.append(nn.ModuleList([
+                ResBlock(in_dim, in_dim, dim, dim * 4),
+                ResBlock(in_dim, in_dim, dim, dim * 4),
+                AttentionBlock(in_dim),
+                nn.Conv2d(in_dim, out_dim, 3, stride=2, padding=1) if not is_last else nn.Conv2d(in_dim, out_dim, 3, padding=1)
+            ]))
+        
+        # Middle blocks
+        self.mid_block1 = ResBlock(dims[-1], dims[-1], dim, dim * 4)
+        self.mid_attn = AttentionBlock(dims[-1])
+        self.mid_block2 = ResBlock(dims[-1], dims[-1], dim, dim * 4)
+        
+        # Upsample blocks
+        self.ups = nn.ModuleList()
+        for ind, (in_dim, out_dim) in enumerate(reversed(in_out)):
+            is_last = ind >= (len(in_out) - 1)
+            self.ups.append(nn.ModuleList([
+                ResBlock(out_dim + in_dim, out_dim, dim, dim * 4),
+                ResBlock(out_dim + in_dim, out_dim, dim, dim * 4),
+                AttentionBlock(out_dim),
+                nn.ConvTranspose2d(out_dim, out_dim, 4, 2, 1) if not is_last else nn.Identity()
+            ]))
+        
+        # Final blocks
+        self.final_block1 = ResBlock(dim * 2, dim, dim, dim * 4)
+        self.final_block2 = ResBlock(dim, dim, dim, dim * 4)
+        self.final_conv = nn.Conv2d(dim, out_channels, 3, padding=1)
+    
+    def forward(self, x, time, text_emb):
+        t = self.time_mlp(time)
+        text_emb = self.text_proj(text_emb)
+        
+        x = self.init_conv(x)
+        h = [x]
+        
+        # Downsample
+        for block1, block2, attn, downsample in self.downs:
+            x = block1(x, t, text_emb)
+            x = block2(x, t, text_emb)
+            x = attn(x)
+            h.append(x)
+            x = downsample(x)
+        
+        # Bottleneck
+        x = self.mid_block1(x, t, text_emb)
+        x = self.mid_attn(x)
+        x = self.mid_block2(x, t, text_emb)
+        
+        # Upsample
+        for block1, block2, attn, upsample in self.ups:
+            x = torch.cat([x, h.pop()], dim=1)
+            x = block1(x, t, text_emb)
+            x = block2(x, t, text_emb)
+            x = attn(x)
+            x = upsample(x)
+        
+        # Final
+        x = torch.cat([x, h.pop()], dim=1)
+        x = self.final_block1(x, t, text_emb)
+        x = self.final_block2(x, t, text_emb)
+        return self.final_conv(x)
+
+class DiffusionModel(nn.Module):
+    def __init__(self, model, betas, device):
+        super().__init__()
+        self.model = model
+        self.betas = betas
+        self.alphas = 1. - betas
+        self.alphas_cumprod = torch.cumprod(self.alphas, dim=0)
+        self.sqrt_alphas_cumprod = torch.sqrt(self.alphas_cumprod)
+        self.sqrt_one_minus_alphas_cumprod = torch.sqrt(1. - self.alphas_cumprod)
+        self.device = device
+        
+        # CLIP model for text conditioning
+        self.clip = CLIPTextModel.from_pretrained("openai/clip-vit-base-patch32")
+        self.tokenizer = CLIPTokenizer.from_pretrained("openai/clip-vit-base-patch32")
+        for param in self.clip.parameters():
+            param.requires_grad = False
+    
+    def get_text_emb(self, prompts):
+        inputs = self.tokenizer(prompts, padding=True, return_tensors="pt").to(self.device)
+        return self.clip(**inputs).last_hidden_state.mean(dim=1)
+    
+    def q_sample(self, x_start, t, noise=None):
+        if noise is None:
+            noise = torch.randn_like(x_start)
+        
+        sqrt_alpha_cumprod = self.sqrt_alphas_cumprod[t].view(-1, 1, 1, 1)
+        sqrt_one_minus_alpha_cumprod = self.sqrt_one_minus_alphas_cumprod[t].view(-1, 1, 1, 1)
+        
+        return sqrt_alpha_cumprod * x_start + sqrt_one_minus_alpha_cumprod * noise
+    
+    def p_losses(self, x_start, text, t, noise=None):
+        if noise is None:
+            noise = torch.randn_like(x_start)
+        
+        x_noisy = self.q_sample(x_start, t, noise)
+        text_emb = self.get_text_emb(text)
+        predicted_noise = self.model(x_noisy, t, text_emb)
+        
+        return F.mse_loss(noise, predicted_noise)
+    
+    @torch.no_grad()
+    def sample(self, prompts, image_size=256, batch_size=4, channels=3, cfg_scale=7.5):
+        shape = (batch_size, channels, image_size, image_size)
+        x = torch.randn(shape, device=self.device)
+        
+        text_emb = self.get_text_emb(prompts)
+        uncond_emb = self.get_text_emb([""] * batch_size)
+        
+        for i in reversed(range(0, len(self.betas))):
+            t = torch.full((batch_size,), i, device=self.device, dtype=torch.long)
+            
+            # Classifier-free guidance
+            noise_pred = self.model(x, t, text_emb)
+            noise_pred_uncond = self.model(x, t, uncond_emb)
+            noise_pred = noise_pred_uncond + cfg_scale * (noise_pred - noise_pred_uncond)
+            
+            alpha = self.alphas[t].view(-1, 1, 1, 1)
+            alpha_cumprod = self.alphas_cumprod[t].view(-1, 1, 1, 1)
+            beta = self.betas[t].view(-1, 1, 1, 1)
+            
+            if i > 0:
+                noise = torch.randn_like(x)
+            else:
+                noise = torch.zeros_like(x)
+            
+            x = (1 / torch.sqrt(alpha)) * (x - ((1 - alpha) / torch.sqrt(1 - alpha_cumprod)) * noise_pred) + torch.sqrt(beta) * noise
+        
+        x = (x.clamp(-1, 1) + 1) / 2
+        x = (x * 255).type(torch.uint8)
+        return x