Delete src

src/__pycache__/model_converter.cpython-312.pyc

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:65ec381ffd1ecb7e843d62f4c98aab8630d90d1273f051e6761d1b9837281628
-size 170116

src/attention.py

@@ -1,69 +0,0 @@
-import torch
-from torch import nn
-from torch.nn import functional as F
-import math
-
-class SelfAttention(nn.Module):
-    def __init__(self, n_heads, d_embed, in_proj_bias=True, out_proj_bias=True):
-        super().__init__()
-        self.in_proj = nn.Linear(d_embed, 3 * d_embed, bias=in_proj_bias)
-        self.out_proj = nn.Linear(d_embed, d_embed, bias=out_proj_bias)
-        self.n_heads = n_heads
-        self.d_head = d_embed // n_heads
-
-    def forward(self, x, causal_mask=False):
-        input_shape = x.shape
-        batch_size, sequence_length, d_embed = input_shape
-        interim_shape = (batch_size, sequence_length, self.n_heads, self.d_head)
-
-        q, k, v = self.in_proj(x).chunk(3, dim=-1)
-        q = q.view(interim_shape).transpose(1, 2)
-        k = k.view(interim_shape).transpose(1, 2)
-        v = v.view(interim_shape).transpose(1, 2)
-
-        weight = q @ k.transpose(-1, -2)
-        
-        if causal_mask:
-            mask = torch.ones_like(weight, dtype=torch.bool).triu(1)
-            weight.masked_fill_(mask, -torch.inf)
-        
-        weight /= math.sqrt(self.d_head)
-        weight = F.softmax(weight, dim=-1)
-        output = weight @ v
-        output = output.transpose(1, 2).reshape(input_shape)
-        output = self.out_proj(output)
-        
-        return output
-
-class CrossAttention(nn.Module):
-    def __init__(self, n_heads, d_embed, d_cross, in_proj_bias=True, out_proj_bias=True):
-        super().__init__()
-        self.q_proj = nn.Linear(d_embed, d_embed, bias=in_proj_bias)
-        self.k_proj = nn.Linear(d_cross, d_embed, bias=in_proj_bias)
-        self.v_proj = nn.Linear(d_cross, d_embed, bias=in_proj_bias)
-        self.out_proj = nn.Linear(d_embed, d_embed, bias=out_proj_bias)
-        self.n_heads = n_heads
-        self.d_head = d_embed // n_heads
-    
-    def forward(self, x, y):
-        input_shape = x.shape
-        batch_size, sequence_length, d_embed = input_shape
-        interim_shape = (batch_size, -1, self.n_heads, self.d_head)
-        
-        q = self.q_proj(x)
-        k = self.k_proj(y)
-        v = self.v_proj(y)
-
-        q = q.view(interim_shape).transpose(1, 2)
-        k = k.view(interim_shape).transpose(1, 2)
-        v = v.view(interim_shape).transpose(1, 2)
-        
-        weight = q @ k.transpose(-1, -2)
-        weight /= math.sqrt(self.d_head)
-        weight = F.softmax(weight, dim=-1)
-        output = weight @ v
-        output = output.transpose(1, 2).contiguous().view(input_shape)
-        output = self.out_proj(output)
-
-        return output
-    

src/clip.py

@@ -1,54 +0,0 @@
-import torch
-from torch import nn
-import torch.nn.functional as F
-from attention import SelfAttention
-
-class CLIPEmbedding(nn.Module):
-    def __init__(self, n_vocab: int, n_embd: int, n_token: int):
-        super().__init__()
-        self.token_embedding = nn.Embedding(n_vocab, n_embd)
-        self.position_embedding = nn.Parameter(torch.zeros((n_token, n_embd)))
-    
-    def forward(self, tokens):
-        x = self.token_embedding(tokens)
-        x += self.position_embedding
-        return x
-
-class CLIPLayer(nn.Module):
-    def __init__(self, n_head: int, n_embd: int):
-        super().__init__()
-        self.layernorm_1 = nn.LayerNorm(n_embd)
-        self.attention = SelfAttention(n_head, n_embd)
-        self.layernorm_2 = nn.LayerNorm(n_embd)
-        self.linear_1 = nn.Linear(n_embd, 4 * n_embd)
-        self.linear_2 = nn.Linear(4 * n_embd, n_embd)
-        self.activation = nn.GELU()
-
-    def forward(self, x):
-        residue = x
-        x = self.layernorm_1(x)
-        x = self.attention(x, causal_mask=True)
-        x += residue
-
-        residue = x
-        x = self.layernorm_2(x)
-        x = self.linear_1(x)
-        x = self.activation(x)
-        x = self.linear_2(x)
-        x += residue
-
-        return x
-
-class CLIP(nn.Module):
-    def __init__(self):
-        super().__init__()
-        self.embedding = CLIPEmbedding(49408, 768, 77)
-        self.layers = nn.ModuleList([CLIPLayer(12, 768) for _ in range(12)])
-        self.layernorm = nn.LayerNorm(768)
-    
-    def forward(self, tokens: torch.LongTensor) -> torch.FloatTensor:
-        state = self.embedding(tokens)
-        for layer in self.layers:
-            state = layer(state)
-        output = self.layernorm(state)
-        return output

src/config.py

@@ -1,72 +0,0 @@
-from dataclasses import dataclass, field
-from typing import Optional, Dict, Any
-import torch
-
-@dataclass
-class ModelConfig:
-    # Image dimensions
-    width: int = 512
-    height: int = 512
-    latents_width: int = 64  # width // 8
-    latents_height: int = 64  # height // 8
-    
-    # Model architecture parameters
-    n_embd: int = 1280
-    n_head: int = 8
-    d_context: int = 768
-    
-    # UNet parameters
-    n_time: int = 1280
-    n_channels: int = 4
-    n_residual_blocks: int = 2
-    
-    # Attention parameters
-    attention_heads: int = 8
-    attention_dim: int = 1280
-
-@dataclass
-class DiffusionConfig:
-    # Sampling parameters
-    n_inference_steps: int = 50
-    guidance_scale: float = 7.5
-    strength: float = 0.8
-    
-    # Sampler configuration
-    sampler_name: str = "ddpm"
-    beta_start: float = 0.00085
-    beta_end: float = 0.0120
-    beta_schedule: str = "linear"
-    
-    # Conditioning parameters
-    do_cfg: bool = True
-    cfg_scale: float = 7.5
-
-@dataclass
-class DeviceConfig:
-    device: Optional[str] = None
-    idle_device: Optional[str] = None
-    
-    def __post_init__(self):
-        if self.device is None:
-            self.device = "cuda" if torch.cuda.is_available() else "cpu"
-        if self.idle_device is None:
-            self.idle_device = "cpu"
-
-@dataclass
-class Config:
-    model: ModelConfig = field(default_factory=ModelConfig)
-    diffusion: DiffusionConfig = field(default_factory=DiffusionConfig)
-    device: DeviceConfig = field(default_factory=DeviceConfig)
-    
-    # Additional settings
-    seed: Optional[int] = None
-    tokenizer: Optional[Any] = None
-    models: Dict[str, Any] = field(default_factory=dict)
-    
-    def __post_init__(self):
-        # Update latent dimensions based on image dimensions
-        self.model.latents_width = self.model.width // 8
-        self.model.latents_height = self.model.height // 8
-
-# Default configuration instance
-default_config = Config() 

src/ddpm.py

@@ -1,76 +0,0 @@
-import torch
-import numpy as np
-
-class DDPMSampler:
-
-    def __init__(self, generator: torch.Generator, num_training_steps=1000, beta_start: float = 0.00085, beta_end: float = 0.0120):
-        self.betas = torch.linspace(beta_start ** 0.5, beta_end ** 0.5, num_training_steps, dtype=torch.float32) ** 2
-        self.alphas = 1.0 - self.betas
-        self.alphas_cumprod = torch.cumprod(self.alphas, dim=0)
-        self.one = torch.tensor(1.0)
-
-        self.generator = generator
-        self.num_train_timesteps = num_training_steps
-        self.timesteps = torch.from_numpy(np.arange(0, num_training_steps)[::-1].copy())
-
-    def set_inference_timesteps(self, num_inference_steps=50):
-        self.num_inference_steps = num_inference_steps
-        step_ratio = self.num_train_timesteps // self.num_inference_steps
-        inference_timesteps = (np.arange(0, num_inference_steps) * step_ratio).round()[::-1].copy().astype(np.int64)
-        self.timesteps = torch.from_numpy(inference_timesteps)
-
-    def _get_previous_timestep(self, timestep: int) -> int:
-        return timestep - self.num_train_timesteps // self.num_inference_steps
-    
-    def _get_variance(self, timestep: int) -> torch.Tensor:
-        prev_timestep = self._get_previous_timestep(timestep)
-        alpha_prod_t = self.alphas_cumprod[timestep]
-        alpha_prod_t_prev = self.alphas_cumprod[prev_timestep] if prev_timestep >= 0 else self.one
-        current_beta_t = 1 - alpha_prod_t / alpha_prod_t_prev
-        variance = (1 - alpha_prod_t_prev) / (1 - alpha_prod_t) * current_beta_t
-        return torch.clamp(variance, min=1e-20)
-    
-    def set_strength(self, strength=1):
-        start_step = self.num_inference_steps - int(self.num_inference_steps * strength)
-        self.timesteps = self.timesteps[start_step:]
-        self.start_step = start_step
-
-    def step(self, timestep: int, latents: torch.Tensor, model_output: torch.Tensor):
-        prev_timestep = self._get_previous_timestep(timestep)
-        alpha_prod_t = self.alphas_cumprod[timestep]
-        alpha_prod_t_prev = self.alphas_cumprod[prev_timestep] if prev_timestep >= 0 else self.one
-        beta_prod_t = 1 - alpha_prod_t
-        beta_prod_t_prev = 1 - alpha_prod_t_prev
-        current_alpha_t = alpha_prod_t / alpha_prod_t_prev
-        current_beta_t = 1 - current_alpha_t
-
-        pred_original_sample = (latents - beta_prod_t ** 0.5 * model_output) / alpha_prod_t ** 0.5
-        pred_original_sample_coeff = (alpha_prod_t_prev ** 0.5 * current_beta_t) / beta_prod_t
-        current_sample_coeff = current_alpha_t ** 0.5 * beta_prod_t_prev / beta_prod_t
-        pred_prev_sample = pred_original_sample_coeff * pred_original_sample + current_sample_coeff * latents
-
-        variance = 0
-        if timestep > 0:
-            device = model_output.device
-            noise = torch.randn(model_output.shape, generator=self.generator, device=device, dtype=model_output.dtype)
-            variance = (self._get_variance(timestep) ** 0.5) * noise
-        
-        return pred_prev_sample + variance
-    
-    def add_noise(self, original_samples: torch.FloatTensor, timesteps: torch.IntTensor) -> torch.FloatTensor:
-        alphas_cumprod = self.alphas_cumprod.to(device=original_samples.device, dtype=original_samples.dtype)
-        timesteps = timesteps.to(original_samples.device)
-
-        sqrt_alpha_prod = alphas_cumprod[timesteps] ** 0.5
-        sqrt_alpha_prod = sqrt_alpha_prod.flatten()
-        while len(sqrt_alpha_prod.shape) < len(original_samples.shape):
-            sqrt_alpha_prod = sqrt_alpha_prod.unsqueeze(-1)
-
-        sqrt_one_minus_alpha_prod = (1 - alphas_cumprod[timesteps]) ** 0.5
-        sqrt_one_minus_alpha_prod = sqrt_one_minus_alpha_prod.flatten()
-        while len(sqrt_one_minus_alpha_prod.shape) < len(original_samples.shape):
-            sqrt_one_minus_alpha_prod = sqrt_one_minus_alpha_prod.unsqueeze(-1)
-
-        noise = torch.randn(original_samples.shape, generator=self.generator, device=original_samples.device, dtype=original_samples.dtype)
-        return sqrt_alpha_prod * original_samples + sqrt_one_minus_alpha_prod * noise
-    

src/decoder.py

@@ -1,76 +0,0 @@
-import torch
-from torch import nn
-from torch.nn import functional as F
-from attention import SelfAttention
-
-class VAE_AttentionBlock(nn.Module):
-    def __init__(self, channels):
-        super().__init__()
-        self.groupnorm = nn.GroupNorm(32, channels)
-        self.attention = SelfAttention(1, channels)
-    
-    def forward(self, x):
-        residue = x
-        x = self.groupnorm(x)
-        n, c, h, w = x.shape
-        x = x.view((n, c, h * w)).transpose(-1, -2)
-        x = self.attention(x)
-        x = x.transpose(-1, -2).view((n, c, h, w))
-        return x + residue
-
-class VAE_ResidualBlock(nn.Module):
-    def __init__(self, in_channels, out_channels):
-        super().__init__()
-        self.groupnorm_1 = nn.GroupNorm(32, in_channels)
-        self.conv_1 = nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1)
-        self.groupnorm_2 = nn.GroupNorm(32, out_channels)
-        self.conv_2 = nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1)
-        self.residual_layer = nn.Identity() if in_channels == out_channels else nn.Conv2d(in_channels, out_channels, kernel_size=1, padding=0)
-    
-    def forward(self, x):
-        residue = x
-        x = self.groupnorm_1(x)
-        x = F.silu(x)
-        x = self.conv_1(x)
-        x = self.groupnorm_2(x)
-        x = F.silu(x)
-        x = self.conv_2(x)
-        return x + self.residual_layer(residue)
-
-class VAE_Decoder(nn.Sequential):
-    def __init__(self):
-        super().__init__(
-            nn.Conv2d(4, 4, kernel_size=1, padding=0),
-            nn.Conv2d(4, 512, kernel_size=3, padding=1),
-            VAE_ResidualBlock(512, 512),
-            VAE_AttentionBlock(512),
-            VAE_ResidualBlock(512, 512),
-            VAE_ResidualBlock(512, 512),
-            VAE_ResidualBlock(512, 512),
-            VAE_ResidualBlock(512, 512),
-            nn.Upsample(scale_factor=2),
-            nn.Conv2d(512, 512, kernel_size=3, padding=1),
-            VAE_ResidualBlock(512, 512),
-            VAE_ResidualBlock(512, 512),
-            VAE_ResidualBlock(512, 512),
-            nn.Upsample(scale_factor=2),
-            nn.Conv2d(512, 512, kernel_size=3, padding=1),
-            VAE_ResidualBlock(512, 256),
-            VAE_ResidualBlock(256, 256),
-            VAE_ResidualBlock(256, 256),
-            nn.Upsample(scale_factor=2),
-            nn.Conv2d(256, 256, kernel_size=3, padding=1),
-            VAE_ResidualBlock(256, 128),
-            VAE_ResidualBlock(128, 128),
-            VAE_ResidualBlock(128, 128),
-            nn.GroupNorm(32, 128),
-            nn.SiLU(),
-            nn.Conv2d(128, 3, kernel_size=3, padding=1),
-        )
-
-    def forward(self, x):
-        x /= 0.18215
-        for module in self:
-            x = module(x)
-        return x
-    

src/demo.py

@@ -1,48 +0,0 @@
-import model_loader
-import pipeline
-from PIL import Image
-from pathlib import Path
-from transformers import CLIPTokenizer
-import torch
-from config import Config, default_config, DeviceConfig
-
-# Device configuration
-ALLOW_CUDA = False
-ALLOW_MPS = False
-
-device = "cpu"
-if torch.cuda.is_available() and ALLOW_CUDA:
-    device = "cuda"
-elif (torch.backends.mps.is_built() or torch.backends.mps.is_available()) and ALLOW_MPS:
-    device = "mps"
-print(f"Using device: {device}")
-
-# Initialize configuration
-config = Config(
-    device=DeviceConfig(device=device),
-    seed=42,
-    tokenizer=CLIPTokenizer.from_pretrained("openai/clip-vit-base-patch32")
-)
-
-# Update diffusion parameters
-config.diffusion.strength = 0.75
-config.diffusion.cfg_scale = 8.0
-config.diffusion.n_inference_steps = 50
-
-# Load models with SE blocks enabled
-model_file = "data/v1-5-pruned-emaonly.ckpt"
-config.models = model_loader.load_models(model_file, device, use_se=True)
-
-# Generate image
-prompt = "A ultra sharp photorealtici painting of a futuristic cityscape at night with neon lights and flying cars"
-uncond_prompt = ""
-
-output_image = pipeline.generate(
-    prompt=prompt,
-    uncond_prompt=uncond_prompt,
-    config=config
-)
-
-# Save output
-output_image = Image.fromarray(output_image)
-output_image.save("output.png")

src/diffusion.py

@@ -1,187 +0,0 @@
-import torch
-from torch import nn
-from torch.nn import functional as F
-from attention import SelfAttention, CrossAttention
-
-class TimeEmbedding(nn.Module):
-    def __init__(self, n_embd):
-        super().__init__()
-        self.linear_1 = nn.Linear(n_embd, 4 * n_embd)
-        self.linear_2 = nn.Linear(4 * n_embd, 4 * n_embd)
-
-    def forward(self, x):
-        x = F.silu(self.linear_1(x))
-        return self.linear_2(x)
-
-class SqueezeExcitation(nn.Module):
-    def __init__(self, channels, reduction=16):
-        super().__init__()
-        self.avg_pool = nn.AdaptiveAvgPool2d(1)
-        self.fc = nn.Sequential(
-            nn.Linear(channels, channels // reduction, bias=False),
-            nn.ReLU(inplace=True),
-            nn.Linear(channels // reduction, channels, bias=False),
-            nn.Sigmoid()
-        )
-
-    def forward(self, x):
-        b, c, _, _ = x.size()
-        y = self.avg_pool(x).view(b, c)
-        y = self.fc(y).view(b, c, 1, 1)
-        return x * y.expand_as(x)
-
-class UNET_ResidualBlock(nn.Module):
-    def __init__(self, in_channels, out_channels, n_time=1280, use_se=False):
-        super().__init__()
-        self.groupnorm_feature = nn.GroupNorm(32, in_channels)
-        self.conv_feature = nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1)
-        self.linear_time = nn.Linear(n_time, out_channels)
-        self.groupnorm_merged = nn.GroupNorm(32, out_channels)
-        self.conv_merged = nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1)
-        self.residual_layer = nn.Identity() if in_channels == out_channels else nn.Conv2d(in_channels, out_channels, kernel_size=1, padding=0)
-        
-        # Add Squeeze-Excitation blocks only if use_se is True
-        self.use_se = use_se
-        if use_se:
-            self.se1 = SqueezeExcitation(out_channels)
-            self.se2 = SqueezeExcitation(out_channels)
-    
-    def forward(self, feature, time):
-        residue = feature
-        feature = F.silu(self.groupnorm_feature(feature))
-        feature = self.conv_feature(feature)
-        if self.use_se:
-            feature = self.se1(feature)  # Apply SE after first conv
-        
-        time = self.linear_time(F.silu(time))
-        merged = feature + time.unsqueeze(-1).unsqueeze(-1)
-        merged = F.silu(self.groupnorm_merged(merged))
-        merged = self.conv_merged(merged)
-        if self.use_se:
-            merged = self.se2(merged)  # Apply SE after second conv
-        
-        return merged + self.residual_layer(residue)
-
-class UNET_AttentionBlock(nn.Module):
-    def __init__(self, n_head: int, n_embd: int, d_context=768):
-        super().__init__()
-        channels = n_head * n_embd
-        self.groupnorm = nn.GroupNorm(32, channels, eps=1e-6)
-        self.conv_input = nn.Conv2d(channels, channels, kernel_size=1, padding=0)
-        self.layernorm_1 = nn.LayerNorm(channels)
-        self.attention_1 = SelfAttention(n_head, channels, in_proj_bias=False)
-        self.layernorm_2 = nn.LayerNorm(channels)
-        self.attention_2 = CrossAttention(n_head, channels, d_context, in_proj_bias=False)
-        self.layernorm_3 = nn.LayerNorm(channels)
-        self.linear_geglu_1 = nn.Linear(channels, 4 * channels * 2)
-        self.linear_geglu_2 = nn.Linear(4 * channels, channels)
-        self.conv_output = nn.Conv2d(channels, channels, kernel_size=1, padding=0)
-    
-    def forward(self, x, context):
-        residue_long = x
-        x = self.conv_input(self.groupnorm(x))
-        n, c, h, w = x.shape
-        x = x.view((n, c, h * w)).transpose(-1, -2)
-        residue_short = x
-        x = self.attention_1(self.layernorm_1(x)) + residue_short
-        residue_short = x
-        x = self.attention_2(self.layernorm_2(x), context) + residue_short
-        residue_short = x
-        x, gate = self.linear_geglu_1(self.layernorm_3(x)).chunk(2, dim=-1)
-        x = self.linear_geglu_2(x * F.gelu(gate)) + residue_short
-        x = x.transpose(-1, -2).view((n, c, h, w))
-        return self.conv_output(x) + residue_long
-
-class Upsample(nn.Module):
-    def __init__(self, channels):
-        super().__init__()
-        self.conv = nn.Conv2d(channels, channels, kernel_size=3, padding=1)
-    
-    def forward(self, x):
-        return self.conv(F.interpolate(x, scale_factor=2, mode='nearest'))
-
-class SwitchSequential(nn.Sequential):
-    def forward(self, x, context, time):
-        for layer in self:
-            if isinstance(layer, UNET_AttentionBlock):
-                x = layer(x, context)
-            elif isinstance(layer, UNET_ResidualBlock):
-                x = layer(x, time)
-            else:
-                x = layer(x)
-        return x
-
-class UNET(nn.Module):
-    def __init__(self, use_se=False):
-        super().__init__()
-        self.encoders = nn.ModuleList([
-            SwitchSequential(nn.Conv2d(4, 320, kernel_size=3, padding=1)),
-            SwitchSequential(UNET_ResidualBlock(320, 320, use_se=use_se), UNET_AttentionBlock(8, 40)),
-            SwitchSequential(UNET_ResidualBlock(320, 320, use_se=use_se), UNET_AttentionBlock(8, 40)),
-            SwitchSequential(nn.Conv2d(320, 320, kernel_size=3, stride=2, padding=1)),
-            SwitchSequential(UNET_ResidualBlock(320, 640, use_se=use_se), UNET_AttentionBlock(8, 80)),
-            SwitchSequential(UNET_ResidualBlock(640, 640, use_se=use_se), UNET_AttentionBlock(8, 80)),
-            SwitchSequential(nn.Conv2d(640, 640, kernel_size=3, stride=2, padding=1)),
-            SwitchSequential(UNET_ResidualBlock(640, 1280, use_se=use_se), UNET_AttentionBlock(8, 160)),
-            SwitchSequential(UNET_ResidualBlock(1280, 1280, use_se=use_se), UNET_AttentionBlock(8, 160)),
-            SwitchSequential(nn.Conv2d(1280, 1280, kernel_size=3, stride=2, padding=1)),
-            SwitchSequential(UNET_ResidualBlock(1280, 1280, use_se=use_se)),
-            SwitchSequential(UNET_ResidualBlock(1280, 1280, use_se=use_se)),
-        ])
-
-        self.bottleneck = SwitchSequential(
-            UNET_ResidualBlock(1280, 1280, use_se=use_se),
-            UNET_AttentionBlock(8, 160),
-            UNET_ResidualBlock(1280, 1280, use_se=use_se),
-        )
-        
-        self.decoders = nn.ModuleList([
-            SwitchSequential(UNET_ResidualBlock(2560, 1280, use_se=use_se)),
-            SwitchSequential(UNET_ResidualBlock(2560, 1280, use_se=use_se)),
-            SwitchSequential(UNET_ResidualBlock(2560, 1280, use_se=use_se), Upsample(1280)),
-            SwitchSequential(UNET_ResidualBlock(2560, 1280, use_se=use_se), UNET_AttentionBlock(8, 160)),
-            SwitchSequential(UNET_ResidualBlock(2560, 1280, use_se=use_se), UNET_AttentionBlock(8, 160)),
-            SwitchSequential(UNET_ResidualBlock(1920, 1280, use_se=use_se), UNET_AttentionBlock(8, 160), Upsample(1280)),
-            SwitchSequential(UNET_ResidualBlock(1920, 640, use_se=use_se), UNET_AttentionBlock(8, 80)),
-            SwitchSequential(UNET_ResidualBlock(1280, 640, use_se=use_se), UNET_AttentionBlock(8, 80)),
-            SwitchSequential(UNET_ResidualBlock(960, 640, use_se=use_se), UNET_AttentionBlock(8, 80), Upsample(640)),
-            SwitchSequential(UNET_ResidualBlock(960, 320, use_se=use_se), UNET_AttentionBlock(8, 40)),
-            SwitchSequential(UNET_ResidualBlock(640, 320, use_se=use_se), UNET_AttentionBlock(8, 40)),
-            SwitchSequential(UNET_ResidualBlock(640, 320, use_se=use_se), UNET_AttentionBlock(8, 40)),
-        ])
-
-    def forward(self, x, context, time):
-        skip_connections = []
-        for layers in self.encoders:
-            x = layers(x, context, time)
-            skip_connections.append(x)
-
-        x = self.bottleneck(x, context, time)
-
-        for layers in self.decoders:
-            x = torch.cat((x, skip_connections.pop()), dim=1)
-            x = layers(x, context, time)
-        
-        return x
-
-class UNET_OutputLayer(nn.Module):
-    def __init__(self, in_channels, out_channels):
-        super().__init__()
-        self.groupnorm = nn.GroupNorm(32, in_channels)
-        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1)
-    
-    def forward(self, x):
-        x = F.silu(self.groupnorm(x))
-        return self.conv(x)
-
-class Diffusion(nn.Module):
-    def __init__(self, use_se=False):
-        super().__init__()
-        self.time_embedding = TimeEmbedding(320)
-        self.unet = UNET(use_se=use_se)
-        self.final = UNET_OutputLayer(320, 4)
-    
-    def forward(self, latent, context, time):
-        time = self.time_embedding(time)
-        output = self.unet(latent, context, time)
-        return self.final(output)

src/encoder.py

@@ -1,42 +0,0 @@
-import torch
-from torch import nn
-from torch.nn import functional as F
-from decoder import VAE_AttentionBlock, VAE_ResidualBlock
-
-class VAE_Encoder(nn.Sequential):
-    def __init__(self):
-        super().__init__(
-            nn.Conv2d(3, 128, kernel_size=3, padding=1),
-            VAE_ResidualBlock(128, 128),
-            VAE_ResidualBlock(128, 128),
-            nn.Conv2d(128, 128, kernel_size=3, stride=2, padding=0),
-            VAE_ResidualBlock(128, 256),
-            VAE_ResidualBlock(256, 256),
-            nn.Conv2d(256, 256, kernel_size=3, stride=2, padding=0),
-            VAE_ResidualBlock(256, 512),
-            VAE_ResidualBlock(512, 512),
-            nn.Conv2d(512, 512, kernel_size=3, stride=2, padding=0),
-            VAE_ResidualBlock(512, 512),
-            VAE_ResidualBlock(512, 512),
-            VAE_ResidualBlock(512, 512),
-            VAE_AttentionBlock(512),
-            VAE_ResidualBlock(512, 512),
-            nn.GroupNorm(32, 512),
-            nn.SiLU(),
-            nn.Conv2d(512, 8, kernel_size=3, padding=1),
-            nn.Conv2d(8, 8, kernel_size=1, padding=0),
-        )
-
-    def forward(self, x, noise):
-        for module in self:
-            if getattr(module, 'stride', None) == (2, 2):
-                x = F.pad(x, (0, 1, 0, 1))
-            x = module(x)
-        mean, log_variance = torch.chunk(x, 2, dim=1)
-        log_variance = torch.clamp(log_variance, -30, 20)
-        variance = log_variance.exp()
-        stdev = variance.sqrt()
-        x = mean + stdev * noise
-        x *= 0.18215
-        return x
-    

src/model_loader.py

@@ -1,40 +0,0 @@
-from clip import CLIP
-from encoder import VAE_Encoder
-from decoder import VAE_Decoder
-from diffusion import Diffusion
-
-import model_converter
-import torch
-
-def load_models(ckpt_path, device, use_se=False):
-    state_dict = model_converter.load_from_standard_weights(ckpt_path, device)
-
-    encoder = VAE_Encoder().to(device)
-    encoder.load_state_dict(state_dict['encoder'], strict=True)
-
-    decoder = VAE_Decoder().to(device)
-    decoder.load_state_dict(state_dict['decoder'], strict=True)
-
-    # Initialize diffusion model with SE blocks disabled for loading pre-trained weights
-    diffusion = Diffusion(use_se=False).to(device)
-    diffusion.load_state_dict(state_dict['diffusion'], strict=True)
-    
-    # If SE blocks are requested, reinitialize the model with them
-    if use_se:
-        diffusion = Diffusion(use_se=True).to(device)
-        # Copy the weights from the loaded model
-        with torch.no_grad():
-            for name, param in diffusion.named_parameters():
-                if 'se' not in name:  # Skip SE block parameters
-                    if name in state_dict['diffusion']:
-                        param.copy_(state_dict['diffusion'][name])
-
-    clip = CLIP().to(device)
-    clip.load_state_dict(state_dict['clip'], strict=True)
-
-    return {
-        'clip': clip,
-        'encoder': encoder,
-        'decoder': decoder,
-        'diffusion': diffusion,
-    }

src/pipeline.py

@@ -1,123 +0,0 @@
-import torch
-import numpy as np
-from tqdm import tqdm
-from ddpm import DDPMSampler
-import logging
-from config import Config, default_config
-
-WIDTH = 512
-HEIGHT = 512
-LATENTS_WIDTH = WIDTH // 8
-LATENTS_HEIGHT = HEIGHT // 8
-
-logging.basicConfig(level=logging.INFO)
-
-def generate(
-    prompt,
-    uncond_prompt=None,
-    input_image=None,
-    config: Config = default_config,
-):
-    with torch.no_grad():
-        validate_strength(config.diffusion.strength)
-        generator = initialize_generator(config.seed, config.device.device)
-        context = encode_prompt(prompt, uncond_prompt, config.diffusion.do_cfg, config.tokenizer, config.models["clip"], config.device.device)
-        latents = initialize_latents(input_image, config.diffusion.strength, generator, config.models, config.device.device, config.diffusion.sampler_name, config.diffusion.n_inference_steps)
-        images = run_diffusion(latents, context, config.diffusion.do_cfg, config.diffusion.cfg_scale, config.models, config.device.device, config.diffusion.sampler_name, config.diffusion.n_inference_steps, generator)
-        return postprocess_images(images)
-
-def validate_strength(strength):
-    if not 0 < strength <= 1:
-        raise ValueError("Strength must be between 0 and 1")
-
-def initialize_generator(seed, device):
-    generator = torch.Generator(device=device)
-    if seed is None:
-        generator.seed()
-    else:
-        generator.manual_seed(seed)
-    return generator
-
-def encode_prompt(prompt, uncond_prompt, do_cfg, tokenizer, clip, device):
-    clip.to(device)
-    if do_cfg:
-        cond_tokens = tokenizer.batch_encode_plus([prompt], padding="max_length", max_length=77).input_ids
-        cond_tokens = torch.tensor(cond_tokens, dtype=torch.long, device=device)
-        cond_context = clip(cond_tokens)
-        uncond_tokens = tokenizer.batch_encode_plus([uncond_prompt], padding="max_length", max_length=77).input_ids
-        uncond_tokens = torch.tensor(uncond_tokens, dtype=torch.long, device=device)
-        uncond_context = clip(uncond_tokens)
-        context = torch.cat([cond_context, uncond_context])
-    else:
-        tokens = tokenizer.batch_encode_plus([prompt], padding="max_length", max_length=77).input_ids
-        tokens = torch.tensor(tokens, dtype=torch.long, device=device)
-        context = clip(tokens)
-    return context
-
-def initialize_latents(input_image, strength, generator, models, device, sampler_name, n_inference_steps):
-    if input_image is None:
-        # Initialize with random noise
-        latents = torch.randn((1, 4, 64, 64), generator=generator, device=device)
-    else:
-        # Initialize with encoded input image
-        latents = encode_image(input_image, models, device)
-        # Add noise based on strength
-        noise = torch.randn_like(latents, generator=generator)
-        latents = (1 - strength) * latents + strength * noise
-    return latents
-
-def preprocess_image(input_image):
-    input_image_tensor = input_image.resize((WIDTH, HEIGHT))
-    input_image_tensor = np.array(input_image_tensor)
-    input_image_tensor = torch.tensor(input_image_tensor, dtype=torch.float32)
-    input_image_tensor = rescale(input_image_tensor, (0, 255), (-1, 1))
-    input_image_tensor = input_image_tensor.unsqueeze(0)
-    input_image_tensor = input_image_tensor.permute(0, 3, 1, 2)
-    return input_image_tensor
-
-def get_sampler(sampler_name, generator, n_inference_steps):
-    if sampler_name == "ddpm":
-        sampler = DDPMSampler(generator)
-        sampler.set_inference_timesteps(n_inference_steps)
-    else:
-        raise ValueError(f"Unknown sampler value {sampler_name}.")
-    return sampler
-
-def run_diffusion(latents, context, do_cfg, cfg_scale, models, device, sampler_name, n_inference_steps, generator):
-    diffusion = models["diffusion"]
-    diffusion.to(device)
-    sampler = get_sampler(sampler_name, generator, n_inference_steps)
-    timesteps = tqdm(sampler.timesteps)
-    for timestep in timesteps:
-        time_embedding = get_time_embedding(timestep).to(device)
-        model_input = latents.repeat(2, 1, 1, 1) if do_cfg else latents
-        model_output = diffusion(model_input, context, time_embedding)
-        if do_cfg:
-            output_cond, output_uncond = model_output.chunk(2)
-            model_output = cfg_scale * (output_cond - output_uncond) + output_uncond
-        latents = sampler.step(timestep, latents, model_output)
-    decoder = models["decoder"]
-    decoder.to(device)
-    images = decoder(latents)
-    return images
-
-def postprocess_images(images):
-    images = rescale(images, (-1, 1), (0, 255), clamp=True)
-    images = images.permute(0, 2, 3, 1)
-    images = images.to("cpu", torch.uint8).numpy()
-    return images[0]
-
-def rescale(x, old_range, new_range, clamp=False):
-    old_min, old_max = old_range
-    new_min, new_max = new_range
-    x -= old_min
-    x *= (new_max - new_min) / (old_max - old_min)
-    x += new_min
-    if clamp:
-        x = x.clamp(new_min, new_max)
-    return x
-
-def get_time_embedding(timestep):
-    freqs = torch.pow(10000, -torch.arange(start=0, end=160, dtype=torch.float32) / 160)
-    x = torch.tensor([timestep], dtype=torch.float32)[:, None] * freqs[None]
-    return torch.cat([torch.cos(x), torch.sin(x)], dim=-1)