Upload 26 files

app.py

@@ -0,0 +1,189 @@
+import gradio as gr
+import numpy as np
+import random
+import torch
+from PIL import Image
+import os
+from huggingface_hub import hf_hub_download
+from pathlib import Path
+import sys
+
+# Add src directory to Python path
+sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+
+from src import model_loader
+from src import pipeline
+from src.config import Config, DeviceConfig
+from transformers import CLIPTokenizer
+
+# Create data directory if it doesn't exist
+data_dir = Path("data")
+data_dir.mkdir(exist_ok=True)
+
+# Model configuration
+MODEL_REPO = "stable-diffusion-v1-5/stable-diffusion-v1-5"
+MODEL_FILENAME = "v1-5-pruned-emaonly.ckpt"
+model_file = data_dir / MODEL_FILENAME
+
+# Download model if it doesn't exist
+if not model_file.exists():
+    print(f"Downloading model from {MODEL_REPO}...")
+    model_file = hf_hub_download(
+        repo_id=MODEL_REPO,
+        filename=MODEL_FILENAME,
+        local_dir=data_dir,
+        local_dir_use_symlinks=False
+    )
+    print("Model downloaded successfully!")
+
+# Device configuration
+device = "cuda" if torch.cuda.is_available() else "cpu"
+print(f"Using device: {device}")
+
+# Initialize configuration
+config = Config(
+    device=DeviceConfig(device=device),
+    tokenizer=CLIPTokenizer.from_pretrained("openai/clip-vit-base-patch32")
+)
+
+# Load models with SE blocks enabled
+config.models = model_loader.load_models(str(model_file), device, use_se=True)
+
+MAX_SEED = np.iinfo(np.int32).max
+MAX_IMAGE_SIZE = 1024
+
+def infer(
+    prompt,
+    negative_prompt,
+    seed,
+    randomize_seed,
+    width,
+    height,
+    guidance_scale,
+    num_inference_steps,
+    progress=gr.Progress(track_tqdm=True),
+):
+    if randomize_seed:
+        seed = random.randint(0, MAX_SEED)
+    
+    # Update config with user settings
+    config.seed = seed
+    config.diffusion.cfg_scale = guidance_scale
+    config.diffusion.n_inference_steps = num_inference_steps
+    config.model.width = width
+    config.model.height = height
+    
+    # Generate image
+    output_image = pipeline.generate(
+        prompt=prompt,
+        uncond_prompt=negative_prompt,
+        config=config
+    )
+    
+    # Convert numpy array to PIL Image
+    image = Image.fromarray(output_image)
+    
+    return image, seed
+
+examples = [
+    "A ultra sharp photorealtici painting of a futuristic cityscape at night with neon lights and flying cars",
+    "A serene mountain landscape at sunset with snow-capped peaks and a clear lake reflection",
+    "A detailed portrait of a cyberpunk character with glowing neon implants and holographic tattoos",
+]
+
+css = """
+#col-container {
+    margin: 0 auto;
+    max-width: 640px;
+}
+"""
+
+with gr.Blocks(css=css) as demo:
+    with gr.Column(elem_id="col-container"):
+        gr.Markdown(" # Custom Diffusion Model Text-to-Image Generator")
+
+        with gr.Row():
+            prompt = gr.Text(
+                label="Prompt",
+                show_label=False,
+                max_lines=1,
+                placeholder="Enter your prompt",
+                container=False,
+            )
+
+            run_button = gr.Button("Run", scale=0, variant="primary")
+
+        result = gr.Image(label="Result", show_label=False)
+
+        with gr.Accordion("Advanced Settings", open=False):
+            negative_prompt = gr.Text(
+                label="Negative prompt",
+                max_lines=1,
+                placeholder="Enter a negative prompt",
+                visible=False,
+            )
+
+            seed = gr.Slider(
+                label="Seed",
+                minimum=0,
+                maximum=MAX_SEED,
+                step=1,
+                value=42,
+            )
+
+            randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
+
+            with gr.Row():
+                width = gr.Slider(
+                    label="Width",
+                    minimum=256,
+                    maximum=MAX_IMAGE_SIZE,
+                    step=32,
+                    value=512,
+                )
+
+                height = gr.Slider(
+                    label="Height",
+                    minimum=256,
+                    maximum=MAX_IMAGE_SIZE,
+                    step=32,
+                    value=512,
+                )
+
+            with gr.Row():
+                guidance_scale = gr.Slider(
+                    label="Guidance scale",
+                    minimum=0.0,
+                    maximum=10.0,
+                    step=0.1,
+                    value=7.5,
+                )
+
+                num_inference_steps = gr.Slider(
+                    label="Number of inference steps",
+                    minimum=1,
+                    maximum=50,
+                    step=1,
+                    value=50,
+                )
+
+        gr.Examples(examples=examples, inputs=[prompt])
+        
+    gr.on(
+        triggers=[run_button.click, prompt.submit],
+        fn=infer,
+        inputs=[
+            prompt,
+            negative_prompt,
+            seed,
+            randomize_seed,
+            width,
+            height,
+            guidance_scale,
+            num_inference_steps,
+        ],
+        outputs=[result, seed],
+    )
+
+if __name__ == "__main__":
+    demo.launch() 

requirements.txt

@@ -0,0 +1,9 @@
+torch>=2.0.0
+gradio>=4.0.0
+transformers>=4.30.0
+numpy>=1.24.0
+Pillow>=10.0.0
+huggingface_hub>=0.19.0
+accelerate>=0.25.0
+safetensors>=0.4.0
+setuptools>=65.5.1

setup.py

@@ -0,0 +1,17 @@
+from setuptools import setup, find_packages
+
+setup(
+    name="custom-diffusion",
+    version="0.1.0",
+    packages=find_packages(),
+    install_requires=[
+        "torch>=2.0.0",
+        "gradio>=4.0.0",
+        "transformers>=4.30.0",
+        "numpy>=1.24.0",
+        "Pillow>=10.0.0",
+        "huggingface_hub>=0.19.0",
+        "accelerate>=0.25.0",
+        "safetensors>=0.4.0",
+    ],
+) 

src/__init__.py

@@ -0,0 +1 @@
+# This file makes the src directory a Python package 

src/__pycache__/model_converter.cpython-312.pyc

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

src/attention.py

@@ -0,0 +1,69 @@
+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

@@ -0,0 +1,54 @@
+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

@@ -0,0 +1,72 @@
+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

@@ -0,0 +1,76 @@
+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

@@ -0,0 +1,76 @@
+import torch
+import torch.nn as nn
+import torch.nn.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

@@ -0,0 +1,48 @@
+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

@@ -0,0 +1,187 @@
+import torch
+import torch.nn as nn
+import torch.nn.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

@@ -0,0 +1,42 @@
+import torch
+import torch.nn as nn
+import torch.nn.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

@@ -0,0 +1,40 @@
+from .clip import CLIP
+from .encoder import VAE_Encoder
+from .decoder import VAE_Decoder
+from .diffusion import Diffusion
+
+from . 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

@@ -0,0 +1,124 @@
+import torch
+import torch.nn.functional as F
+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)