fixed code readability

app.py

@@ -1,13 +1,11 @@
 import gradio as gr
 from PIL import Image
-import numpy as np
 from torchvision import transforms
 from load_model import sample
 import torch
-import glob
 import random
 import os
-import pathlib
+
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
 device = "mps" if torch.backends.mps.is_available() else device

load_model.py

@@ -1,12 +1,10 @@
-from models.structure.Unet_3 import Unet
+from models.structure.Advanced_Conditional_Unet import Unet
 from diffusers import DDPMScheduler
 import torch
 import os
 import glob
-from tqdm import tqdm
 from torchvision import transforms
 import pathlib
-from torchvision.utils import save_image
 from safetensors.torch import load_model, save_model
 import time as tm
 

models/structure/{Unet_3.py → Advanced_Conditional_Unet.py}

@@ -1,14 +1,8 @@
-import math
-from inspect import isfunction
 from functools import partial
-import matplotlib.pyplot as plt
-from tqdm.auto import tqdm
-from einops import rearrange
 import torch
-from torch import nn, einsum
+from torch import nn
 import torch.nn.functional as F
-from .Advanced_Network_Helpers_3 import *
-from transformers import PreTrainedModel
+from .Advanced_Network_Helpers import *
 
 
 class Unet(nn.Module):

models/structure/Advanced_Network_Helpers.py

@@ -143,23 +143,13 @@ class Attention(nn.Module):
         self.to_v = nn.Conv2d(dim, hidden_dim, 1, bias=False)
         self.to_out = nn.Conv2d(hidden_dim, dim, 1)
 
-    def forward(self, x, cross_attend=None):
+    def forward(self, x):
         b, c, h, w = x.shape
 
-        if cross_attend is not None:
-            assert cross_attend.shape == x.shape
-
-            q_att = self.to_q(x)
-            k_att = self.to_k(cross_attend)
-            v_att = self.to_v(cross_attend)
-            q = rearrange(q_att, "b (h c) x y -> b h c (x y)", h=self.heads)
-            k = rearrange(k_att, "b (h c) x y -> b h c (x y)", h=self.heads)
-            v = rearrange(v_att, "b (h c) x y -> b h c (x y)", h=self.heads)
-        else:
-            qkv = self.to_qkv(x).chunk(3, dim=1)
-            q, k, v = map(
-                lambda t: rearrange(t, "b (h c) x y -> b h c (x y)", h=self.heads), qkv
-            )
+        qkv = self.to_qkv(x).chunk(3, dim=1)
+        q, k, v = map(
+            lambda t: rearrange(t, "b (h c) x y -> b h c (x y)", h=self.heads), qkv
+        )
         q = q * self.scale
 
         sim = einsum("b h d i, b h d j -> b h i j", q, k)
@@ -173,7 +163,7 @@ class Attention(nn.Module):
 
 
 class LinearCrossAttention(nn.Module):
-    def __init__(self, dim, heads=12, dim_head=128) -> None:
+    def __init__(self, dim, heads=4, dim_head=32) -> None:
         super().__init__()
         self.scale = dim_head**-0.5
         self.heads = heads
@@ -210,25 +200,12 @@ class LinearAttention(nn.Module):
         self.to_v = nn.Conv2d(dim, hidden_dim, 1, bias=False)
         self.to_out = nn.Sequential(nn.Conv2d(hidden_dim, dim, 1), nn.GroupNorm(1, dim))
 
-    def forward(self, x, cross_attend=None):
+    def forward(self, x):
         b, c, h, w = x.shape
-        if cross_attend is not None:
-            assert (
-                cross_attend.shape == x.shape
-            ), f"cross_attend must be same shape as x is {cross_attend.shape} and x is {x.shape}"
-
-            q_att = self.to_q(x)
-            k_att = self.to_k(cross_attend)
-            v_att = self.to_v(cross_attend)
-            q = rearrange(q_att, "b (h c) x y -> b h c (x y)", h=self.heads)
-            k = rearrange(k_att, "b (h c) x y -> b h c (x y)", h=self.heads)
-            v = rearrange(v_att, "b (h c) x y -> b h c (x y)", h=self.heads)
-
-        else:
-            qkv = self.to_qkv(x).chunk(3, dim=1)
-            q, k, v = map(
-                lambda t: rearrange(t, "b (h c) x y -> b h c (x y)", h=self.heads), qkv
-            )
+        qkv = self.to_qkv(x).chunk(3, dim=1)
+        q, k, v = map(
+            lambda t: rearrange(t, "b (h c) x y -> b h c (x y)", h=self.heads), qkv
+        )
         # calculate the softmax with respect to columns softmax of equivalent to q^T with respect to last dim
         q = q.softmax(dim=-2)
         # calculate the softmax with respect to rows of k

models/structure/Advanced_Network_Helpers_2.py

@@ -1,232 +0,0 @@
-import math
-from inspect import isfunction
-from functools import partial
-import matplotlib.pyplot as plt
-from tqdm.auto import tqdm
-from einops import rearrange
-import torch
-from torch import nn, einsum
-import torch.nn.functional as F
-
-
-def exists(x):
-    return x is not None
-
-
-def default(val, d):
-    if exists(val):
-        return val
-    return d() if isfunction(d) else d
-
-
-class Residual(nn.Module):
-    def __init__(self, fn):
-        super().__init__()
-        self.fn = fn
-
-    def forward(self, x, *args, **kwargs):
-        return self.fn(x, *args, **kwargs) + x
-
-
-def Upsample(dim):
-    return nn.ConvTranspose2d(dim, dim, 4, 2, 1)
-
-
-def Downsample(dim):
-    return nn.Conv2d(dim, dim, 4, 2, 1)
-
-
-class SinusoidalPositionEmbeddings(nn.Module):
-    def __init__(self, dim):
-        super().__init__()
-        self.dim = dim
-
-    def forward(self, time):
-        device = time.device
-        half_dim = self.dim // 2
-        embeddings = math.log(10000) / (half_dim - 1)
-        embeddings = torch.exp(torch.arange(half_dim, device=device) * -embeddings)
-        embeddings = time[:, None] * embeddings[None, :]
-        embeddings = torch.cat((embeddings.sin(), embeddings.cos()), dim=-1)
-        return embeddings
-
-
-class Block(nn.Module):
-    def __init__(self, dim, dim_out, groups=8):
-        super().__init__()
-        self.proj = nn.Conv2d(dim, dim_out, 3, padding=1)
-        self.norm = nn.GroupNorm(groups, dim_out)
-        self.act = nn.SiLU()
-
-    def forward(self, x, scale_shift=None):
-        x = self.proj(x)
-        x = self.norm(x)
-
-        if exists(scale_shift):
-            scale, shift = scale_shift
-            x = x * (scale + 1) + shift
-
-        x = self.act(x)
-        return x
-
-
-class ResnetBlock(nn.Module):
-    """https://arxiv.org/abs/1512.03385"""
-
-    def __init__(self, dim, dim_out, *, time_emb_dim=None, groups=8):
-        super().__init__()
-        self.mlp = (
-            nn.Sequential(nn.SiLU(), nn.Linear(time_emb_dim, dim_out))
-            if exists(time_emb_dim)
-            else None
-        )
-
-        self.block1 = Block(dim, dim_out, groups=groups)
-        self.block2 = Block(dim_out, dim_out, groups=groups)
-        self.res_conv = nn.Conv2d(dim, dim_out, 1) if dim != dim_out else nn.Identity()
-
-    def forward(self, x, time_emb=None):
-        h = self.block1(x)
-
-        if exists(self.mlp) and exists(time_emb):
-            time_emb = self.mlp(time_emb)
-            h = rearrange(time_emb, "b c -> b c 1 1") + h
-
-        h = self.block2(h)
-        return h + self.res_conv(x)
-
-
-class ConvNextBlock(nn.Module):
-    """https://arxiv.org/abs/2201.03545"""
-
-    def __init__(self, dim, dim_out, *, time_emb_dim=None, mult=2, norm=True):
-        super().__init__()
-        self.mlp = (
-            nn.Sequential(nn.GELU(), nn.Linear(time_emb_dim, dim))
-            if exists(time_emb_dim)
-            else None
-        )
-
-        self.ds_conv = nn.Conv2d(dim, dim, 7, padding=3, groups=dim)
-
-        self.net = nn.Sequential(
-            nn.GroupNorm(1, dim) if norm else nn.Identity(),
-            nn.Conv2d(dim, dim_out * mult, 3, padding=1),
-            nn.GELU(),
-            nn.GroupNorm(1, dim_out * mult),
-            nn.Conv2d(dim_out * mult, dim_out, 3, padding=1),
-        )
-
-        self.res_conv = nn.Conv2d(dim, dim_out, 1) if dim != dim_out else nn.Identity()
-
-    def forward(self, x, time_emb=None):
-        h = self.ds_conv(x)
-
-        if exists(self.mlp) and exists(time_emb):
-            assert exists(time_emb), "time embedding must be passed in"
-            condition = self.mlp(time_emb)
-            h = h + rearrange(condition, "b c -> b c 1 1")
-
-        h = self.net(h)
-        return h + self.res_conv(x)
-
-
-class Attention(nn.Module):
-    def __init__(self, dim, heads=4, dim_head=32):
-        super().__init__()
-        self.scale = dim_head**-0.5
-        self.heads = heads
-        hidden_dim = dim_head * heads
-        self.to_qkv = nn.Conv2d(dim, hidden_dim * 3, 1, bias=False)
-        self.to_q = nn.Conv2d(dim, hidden_dim, 1, bias=False)
-        self.to_k = nn.Conv2d(dim, hidden_dim, 1, bias=False)
-        self.to_v = nn.Conv2d(dim, hidden_dim, 1, bias=False)
-        self.to_out = nn.Conv2d(hidden_dim, dim, 1)
-
-    def forward(self, x):
-        b, c, h, w = x.shape
-
-        qkv = self.to_qkv(x).chunk(3, dim=1)
-        q, k, v = map(
-            lambda t: rearrange(t, "b (h c) x y -> b h c (x y)", h=self.heads), qkv
-        )
-        q = q * self.scale
-
-        sim = einsum("b h d i, b h d j -> b h i j", q, k)
-        sim = sim - sim.amax(dim=-1, keepdim=True).detach()
-        attn = sim.softmax(dim=-1)
-
-        out = einsum("b h i j, b h d j -> b h i d", attn, v)
-        out = rearrange(out, "b h (x y) d -> b (h d) x y", x=h, y=w)
-
-        return self.to_out(out)
-
-
-class LinearCrossAttention(nn.Module):
-    def __init__(self, dim, heads=4, dim_head=32) -> None:
-        super().__init__()
-        self.scale = dim_head**-0.5
-        self.heads = heads
-        hidden_dim = dim_head * heads
-        self.to_kv = nn.Conv2d(dim, hidden_dim * 2, 1, bias=False)
-        self.to_q = nn.Conv2d(dim, hidden_dim, 1, bias=False)
-        self.out = nn.Conv2d(hidden_dim, dim, 1)
-
-    def forward(self, x, cross_attend):
-        b, c, h, w = x.shape
-        q = self.to_q(x)
-        k, v = self.to_kv(cross_attend).chunk(2, dim=1)
-        q = rearrange(q, "b (h c) x y -> b h c (x y)", h=self.heads)
-        k = rearrange(k, "b (h c) x y -> b h c (x y)", h=self.heads)
-        v = rearrange(v, "b (h c) x y -> b h c (x y)", h=self.heads)
-        q = q * self.scale
-        sim = einsum("b h d i, b h d j -> b h i j", q, k)
-        sim = sim - sim.amax(dim=-1, keepdim=True).detach()
-        attn = sim.softmax(dim=-1)
-        out = einsum("b h i j, b h d j -> b h i d", attn, v)
-        out = rearrange(out, "b h (x y) d -> b (h d) x y", x=h, y=w)
-        return self.out(out)
-
-
-class LinearAttention(nn.Module):
-    def __init__(self, dim, heads=4, dim_head=32):
-        super().__init__()
-        self.scale = dim_head**-0.5
-        self.heads = heads
-        hidden_dim = dim_head * heads
-        self.to_qkv = nn.Conv2d(dim, hidden_dim * 3, 1, bias=False)
-        self.to_q = nn.Conv2d(dim, hidden_dim, 1, bias=False)
-        self.to_k = nn.Conv2d(dim, hidden_dim, 1, bias=False)
-        self.to_v = nn.Conv2d(dim, hidden_dim, 1, bias=False)
-        self.to_out = nn.Sequential(nn.Conv2d(hidden_dim, dim, 1), nn.GroupNorm(1, dim))
-
-    def forward(self, x):
-        b, c, h, w = x.shape
-        qkv = self.to_qkv(x).chunk(3, dim=1)
-        q, k, v = map(
-            lambda t: rearrange(t, "b (h c) x y -> b h c (x y)", h=self.heads), qkv
-        )
-        # calculate the softmax with respect to columns softmax of equivalent to q^T with respect to last dim
-        q = q.softmax(dim=-2)
-        # calculate the softmax with respect to rows of k
-        k = k.softmax(dim=-1)
-        # normalize the values in the attention matrix
-        q = q * self.scale
-        # dot product of q and v matrices
-        context = torch.einsum("b h d n, b h e n -> b h d e", k, v)
-        # dot product of context and q
-        out = torch.einsum("b h d e, b h d n -> b h e n", context, q)
-        # rearrange the output to match the pytorch convention
-        out = rearrange(out, "b h c (x y) -> b (h c) x y", h=self.heads, x=h, y=w)
-        return self.to_out(out)
-
-
-class PreNorm(nn.Module):
-    def __init__(self, dim, fn):
-        super().__init__()
-        self.fn = fn
-        self.norm = nn.GroupNorm(1, dim)
-
-    def forward(self, x, *args, **kwargs):
-        x = self.norm(x)
-        return self.fn(x, *args, **kwargs)

models/structure/Advanced_Network_Helpers_3.py

@@ -1,232 +0,0 @@
-import math
-from inspect import isfunction
-from functools import partial
-import matplotlib.pyplot as plt
-from tqdm.auto import tqdm
-from einops import rearrange
-import torch
-from torch import nn, einsum
-import torch.nn.functional as F
-
-
-def exists(x):
-    return x is not None
-
-
-def default(val, d):
-    if exists(val):
-        return val
-    return d() if isfunction(d) else d
-
-
-class Residual(nn.Module):
-    def __init__(self, fn):
-        super().__init__()
-        self.fn = fn
-
-    def forward(self, x, *args, **kwargs):
-        return self.fn(x, *args, **kwargs) + x
-
-
-def Upsample(dim):
-    return nn.ConvTranspose2d(dim, dim, 4, 2, 1)
-
-
-def Downsample(dim):
-    return nn.Conv2d(dim, dim, 4, 2, 1)
-
-
-class SinusoidalPositionEmbeddings(nn.Module):
-    def __init__(self, dim):
-        super().__init__()
-        self.dim = dim
-
-    def forward(self, time):
-        device = time.device
-        half_dim = self.dim // 2
-        embeddings = math.log(10000) / (half_dim - 1)
-        embeddings = torch.exp(torch.arange(half_dim, device=device) * -embeddings)
-        embeddings = time[:, None] * embeddings[None, :]
-        embeddings = torch.cat((embeddings.sin(), embeddings.cos()), dim=-1)
-        return embeddings
-
-
-class Block(nn.Module):
-    def __init__(self, dim, dim_out, groups=8):
-        super().__init__()
-        self.proj = nn.Conv2d(dim, dim_out, 3, padding=1)
-        self.norm = nn.GroupNorm(groups, dim_out)
-        self.act = nn.SiLU()
-
-    def forward(self, x, scale_shift=None):
-        x = self.proj(x)
-        x = self.norm(x)
-
-        if exists(scale_shift):
-            scale, shift = scale_shift
-            x = x * (scale + 1) + shift
-
-        x = self.act(x)
-        return x
-
-
-class ResnetBlock(nn.Module):
-    """https://arxiv.org/abs/1512.03385"""
-
-    def __init__(self, dim, dim_out, *, time_emb_dim=None, groups=8):
-        super().__init__()
-        self.mlp = (
-            nn.Sequential(nn.SiLU(), nn.Linear(time_emb_dim, dim_out))
-            if exists(time_emb_dim)
-            else None
-        )
-
-        self.block1 = Block(dim, dim_out, groups=groups)
-        self.block2 = Block(dim_out, dim_out, groups=groups)
-        self.res_conv = nn.Conv2d(dim, dim_out, 1) if dim != dim_out else nn.Identity()
-
-    def forward(self, x, time_emb=None):
-        h = self.block1(x)
-
-        if exists(self.mlp) and exists(time_emb):
-            time_emb = self.mlp(time_emb)
-            h = rearrange(time_emb, "b c -> b c 1 1") + h
-
-        h = self.block2(h)
-        return h + self.res_conv(x)
-
-
-class ConvNextBlock(nn.Module):
-    """https://arxiv.org/abs/2201.03545"""
-
-    def __init__(self, dim, dim_out, *, time_emb_dim=None, mult=2, norm=True):
-        super().__init__()
-        self.mlp = (
-            nn.Sequential(nn.GELU(), nn.Linear(time_emb_dim, dim))
-            if exists(time_emb_dim)
-            else None
-        )
-
-        self.ds_conv = nn.Conv2d(dim, dim, 7, padding=3, groups=dim)
-
-        self.net = nn.Sequential(
-            nn.GroupNorm(1, dim) if norm else nn.Identity(),
-            nn.Conv2d(dim, dim_out * mult, 3, padding=1),
-            nn.GELU(),
-            nn.GroupNorm(1, dim_out * mult),
-            nn.Conv2d(dim_out * mult, dim_out, 3, padding=1),
-        )
-
-        self.res_conv = nn.Conv2d(dim, dim_out, 1) if dim != dim_out else nn.Identity()
-
-    def forward(self, x, time_emb=None):
-        h = self.ds_conv(x)
-
-        if exists(self.mlp) and exists(time_emb):
-            assert exists(time_emb), "time embedding must be passed in"
-            condition = self.mlp(time_emb)
-            h = h + rearrange(condition, "b c -> b c 1 1")
-
-        h = self.net(h)
-        return h + self.res_conv(x)
-
-
-class Attention(nn.Module):
-    def __init__(self, dim, heads=4, dim_head=32):
-        super().__init__()
-        self.scale = dim_head**-0.5
-        self.heads = heads
-        hidden_dim = dim_head * heads
-        self.to_qkv = nn.Conv2d(dim, hidden_dim * 3, 1, bias=False)
-        self.to_q = nn.Conv2d(dim, hidden_dim, 1, bias=False)
-        self.to_k = nn.Conv2d(dim, hidden_dim, 1, bias=False)
-        self.to_v = nn.Conv2d(dim, hidden_dim, 1, bias=False)
-        self.to_out = nn.Conv2d(hidden_dim, dim, 1)
-
-    def forward(self, x):
-        b, c, h, w = x.shape
-
-        qkv = self.to_qkv(x).chunk(3, dim=1)
-        q, k, v = map(
-            lambda t: rearrange(t, "b (h c) x y -> b h c (x y)", h=self.heads), qkv
-        )
-        q = q * self.scale
-
-        sim = einsum("b h d i, b h d j -> b h i j", q, k)
-        sim = sim - sim.amax(dim=-1, keepdim=True).detach()
-        attn = sim.softmax(dim=-1)
-
-        out = einsum("b h i j, b h d j -> b h i d", attn, v)
-        out = rearrange(out, "b h (x y) d -> b (h d) x y", x=h, y=w)
-
-        return self.to_out(out)
-
-
-class LinearCrossAttention(nn.Module):
-    def __init__(self, dim, heads=4, dim_head=32) -> None:
-        super().__init__()
-        self.scale = dim_head**-0.5
-        self.heads = heads
-        hidden_dim = dim_head * heads
-        self.to_kv = nn.Conv2d(dim, hidden_dim * 2, 1, bias=False)
-        self.to_q = nn.Conv2d(dim, hidden_dim, 1, bias=False)
-        self.out = nn.Conv2d(hidden_dim, dim, 1)
-
-    def forward(self, x, cross_attend):
-        b, c, h, w = x.shape
-        q = self.to_q(x)
-        k, v = self.to_kv(cross_attend).chunk(2, dim=1)
-        q = rearrange(q, "b (h c) x y -> b h c (x y)", h=self.heads)
-        k = rearrange(k, "b (h c) x y -> b h c (x y)", h=self.heads)
-        v = rearrange(v, "b (h c) x y -> b h c (x y)", h=self.heads)
-        q = q * self.scale
-        sim = einsum("b h d i, b h d j -> b h i j", q, k)
-        sim = sim - sim.amax(dim=-1, keepdim=True).detach()
-        attn = sim.softmax(dim=-1)
-        out = einsum("b h i j, b h d j -> b h i d", attn, v)
-        out = rearrange(out, "b h (x y) d -> b (h d) x y", x=h, y=w)
-        return self.out(out)
-
-
-class LinearAttention(nn.Module):
-    def __init__(self, dim, heads=4, dim_head=32):
-        super().__init__()
-        self.scale = dim_head**-0.5
-        self.heads = heads
-        hidden_dim = dim_head * heads
-        self.to_qkv = nn.Conv2d(dim, hidden_dim * 3, 1, bias=False)
-        self.to_q = nn.Conv2d(dim, hidden_dim, 1, bias=False)
-        self.to_k = nn.Conv2d(dim, hidden_dim, 1, bias=False)
-        self.to_v = nn.Conv2d(dim, hidden_dim, 1, bias=False)
-        self.to_out = nn.Sequential(nn.Conv2d(hidden_dim, dim, 1), nn.GroupNorm(1, dim))
-
-    def forward(self, x):
-        b, c, h, w = x.shape
-        qkv = self.to_qkv(x).chunk(3, dim=1)
-        q, k, v = map(
-            lambda t: rearrange(t, "b (h c) x y -> b h c (x y)", h=self.heads), qkv
-        )
-        # calculate the softmax with respect to columns softmax of equivalent to q^T with respect to last dim
-        q = q.softmax(dim=-2)
-        # calculate the softmax with respect to rows of k
-        k = k.softmax(dim=-1)
-        # normalize the values in the attention matrix
-        q = q * self.scale
-        # dot product of q and v matrices
-        context = torch.einsum("b h d n, b h e n -> b h d e", k, v)
-        # dot product of context and q
-        out = torch.einsum("b h d e, b h d n -> b h e n", context, q)
-        # rearrange the output to match the pytorch convention
-        out = rearrange(out, "b h c (x y) -> b (h c) x y", h=self.heads, x=h, y=w)
-        return self.to_out(out)
-
-
-class PreNorm(nn.Module):
-    def __init__(self, dim, fn):
-        super().__init__()
-        self.fn = fn
-        self.norm = nn.GroupNorm(1, dim)
-
-    def forward(self, x, *args, **kwargs):
-        x = self.norm(x)
-        return self.fn(x, *args, **kwargs)

models/structure/Unet.py

@@ -1,152 +0,0 @@
-import math
-from inspect import isfunction
-from functools import partial
-import matplotlib.pyplot as plt
-from tqdm.auto import tqdm
-from einops import rearrange
-import torch
-from torch import nn, einsum
-import torch.nn.functional as F
-from .Advanced_Network_Helpers import *
-
-
-class Unet(nn.Module):
-    def __init__(
-        self,
-        dim,
-        init_dim=None,
-        out_dim=None,
-        dim_mults=(1, 2, 4, 8),
-        channels=3,
-        with_time_emb=True,
-        resnet_block_groups=8,
-        use_convnext=True,
-        convnext_mult=2,
-    ):
-        super().__init__()
-
-        # determine dimensions
-        self.channels = channels  # since we are concatenating the images and the conditionings along the channel dimension
-
-        init_dim = default(init_dim, dim // 3 * 2)
-        self.init_conv = nn.Conv2d(self.channels * 2, init_dim, 7, padding=3)
-        self.conditioning_init = nn.Conv2d(self.channels * 2, init_dim, 7, padding=3)
-        dims = [init_dim, *map(lambda m: dim * m, dim_mults)]
-        in_out = list(zip(dims[:-1], dims[1:]))
-        self.in_out = in_out
-
-        if use_convnext:
-            block_klass = partial(ConvNextBlock, mult=convnext_mult)
-        else:
-            block_klass = partial(ResnetBlock, groups=resnet_block_groups)
-
-        # time embeddings
-        if with_time_emb:
-            time_dim = dim * 4
-            self.time_mlp = nn.Sequential(
-                SinusoidalPositionEmbeddings(dim),
-                nn.Linear(dim, time_dim),
-                nn.GELU(),
-                nn.Linear(time_dim, time_dim),
-            )
-        else:
-            time_dim = None
-            self.time_mlp = None
-
-        # layers
-        self.downs = nn.ModuleList([])
-        self.ups = nn.ModuleList([])
-        self.conditioning_encoder = nn.ModuleList([])
-        num_resolutions = len(in_out)
-        self.num_resolutions = num_resolutions
-
-        # conditioning encoder
-        for ind, (dim_in, dim_out) in enumerate(in_out):
-            is_last = ind >= (num_resolutions - 1)
-
-            self.conditioning_encoder.append(
-                nn.ModuleList(
-                    [
-                        block_klass(dim_in, dim_out),
-                        Residual(PreNorm(dim_out, LinearAttention(dim_out))),
-                        Downsample(dim_out) if not is_last else nn.Identity(),
-                    ]
-                )
-            )
-
-        for ind, (dim_in, dim_out) in enumerate(in_out):
-            is_last = ind >= (num_resolutions - 1)
-
-            self.downs.append(
-                nn.ModuleList(
-                    [
-                        block_klass(dim_in, dim_out, time_emb_dim=time_dim),
-                        block_klass(dim_out, dim_out, time_emb_dim=time_dim),
-                        Residual(PreNorm(dim_out, LinearAttention(dim_out))),
-                        Downsample(dim_out) if not is_last else nn.Identity(),
-                    ]
-                )
-            )
-
-        mid_dim = dims[-1]
-
-        self.mid_block1 = block_klass(mid_dim, mid_dim, time_emb_dim=time_dim)
-        self.cross_attention = Residual(PreNorm(mid_dim, LinearCrossAttention(mid_dim)))
-        self.mid_block2 = block_klass(mid_dim, mid_dim, time_emb_dim=time_dim)
-
-        for ind, (dim_in, dim_out) in enumerate(reversed(in_out[1:])):
-            is_last = ind >= (num_resolutions - 1)
-            self.ups.append(
-                nn.ModuleList(
-                    [
-                        block_klass(dim_out * 2, dim_in, time_emb_dim=time_dim),
-                        block_klass(dim_in, dim_in, time_emb_dim=time_dim),
-                        Residual(PreNorm(dim_in, LinearAttention(dim_in))),
-                        Upsample(dim_in) if not is_last else nn.Identity(),
-                    ]
-                )
-            )
-
-        out_dim = default(out_dim, channels)
-        self.final_conv = nn.Sequential(
-            block_klass(dim, dim), nn.Conv2d(dim, out_dim, 1)
-        )
-
-    def forward(self, x, time, implicit_conditioning, explicit_conditioning):
-        x = torch.cat((x, explicit_conditioning), dim=1)
-        conditioning = torch.cat((implicit_conditioning, explicit_conditioning), dim=1)
-        x = self.init_conv(x)
-
-        conditioning = self.conditioning_init(conditioning)
-
-        t = self.time_mlp(time) if exists(self.time_mlp) else None
-
-        h = []
-
-        # conditioning encoder
-
-        for block1, attn, downsample in self.conditioning_encoder:
-            conditioning = block1(conditioning)
-            conditioning = attn(conditioning)
-            conditioning = downsample(conditioning)
-
-        for block1, block2, attn, downsample in self.downs:
-            x = block1(x, t)
-            x = block2(x, t)
-            x = attn(x)
-            h.append(x)
-            x = downsample(x)
-
-        # bottleneck
-        x = self.mid_block1(x, t)
-        x = self.cross_attention(x, conditioning)
-        x = self.mid_block2(x, t)
-
-        for block1, block2, attn, upsample in self.ups:
-            x = torch.cat((x, h.pop()), dim=1)
-            x = block1(x, t)
-            x = block2(x, t)
-            x = attn(x)
-            x = upsample(x)
-
-        return self.final_conv(x)

models/structure/Unet_2.py

@@ -1,152 +0,0 @@
-import math
-from inspect import isfunction
-from functools import partial
-import matplotlib.pyplot as plt
-from tqdm.auto import tqdm
-from einops import rearrange
-import torch
-from torch import nn, einsum
-import torch.nn.functional as F
-from .Advanced_Network_Helpers_2 import *
-
-
-class Unet(nn.Module):
-    def __init__(
-        self,
-        dim,
-        init_dim=None,
-        out_dim=None,
-        dim_mults=(1, 2, 4, 8),
-        channels=3,
-        with_time_emb=True,
-        resnet_block_groups=8,
-        use_convnext=True,
-        convnext_mult=2,
-    ):
-        super().__init__()
-
-        # determine dimensions
-        self.channels = channels  # since we are concatenating the images and the conditionings along the channel dimension
-
-        init_dim = default(init_dim, dim // 3 * 2)
-        self.init_conv = nn.Conv2d(self.channels * 2, init_dim, 7, padding=3)
-        self.conditioning_init = nn.Conv2d(self.channels * 2, init_dim, 7, padding=3)
-        dims = [init_dim, *map(lambda m: dim * m, dim_mults)]
-        in_out = list(zip(dims[:-1], dims[1:]))
-        self.in_out = in_out
-
-        if use_convnext:
-            block_klass = partial(ConvNextBlock, mult=convnext_mult)
-        else:
-            block_klass = partial(ResnetBlock, groups=resnet_block_groups)
-
-        # time embeddings
-        if with_time_emb:
-            time_dim = dim * 4
-            self.time_mlp = nn.Sequential(
-                SinusoidalPositionEmbeddings(dim),
-                nn.Linear(dim, time_dim),
-                nn.GELU(),
-                nn.Linear(time_dim, time_dim),
-            )
-        else:
-            time_dim = None
-            self.time_mlp = None
-
-        # layers
-        self.downs = nn.ModuleList([])
-        self.ups = nn.ModuleList([])
-        self.conditioning_encoder = nn.ModuleList([])
-        num_resolutions = len(in_out)
-        self.num_resolutions = num_resolutions
-
-        # conditioning encoder
-        for ind, (dim_in, dim_out) in enumerate(in_out):
-            is_last = ind >= (num_resolutions - 1)
-
-            self.conditioning_encoder.append(
-                nn.ModuleList(
-                    [
-                        block_klass(dim_in, dim_out),
-                        Residual(PreNorm(dim_out, LinearAttention(dim_out))),
-                        Downsample(dim_out) if not is_last else nn.Identity(),
-                    ]
-                )
-            )
-
-        for ind, (dim_in, dim_out) in enumerate(in_out):
-            is_last = ind >= (num_resolutions - 1)
-
-            self.downs.append(
-                nn.ModuleList(
-                    [
-                        block_klass(dim_in, dim_out, time_emb_dim=time_dim),
-                        block_klass(dim_out, dim_out, time_emb_dim=time_dim),
-                        Residual(PreNorm(dim_out, LinearAttention(dim_out))),
-                        Downsample(dim_out) if not is_last else nn.Identity(),
-                    ]
-                )
-            )
-
-        mid_dim = dims[-1]
-
-        self.mid_block1 = block_klass(mid_dim, mid_dim, time_emb_dim=time_dim)
-        self.cross_attention = Residual(PreNorm(mid_dim, LinearCrossAttention(mid_dim)))
-        self.mid_block2 = block_klass(mid_dim, mid_dim, time_emb_dim=time_dim)
-
-        for ind, (dim_in, dim_out) in enumerate(reversed(in_out[1:])):
-            is_last = ind >= (num_resolutions - 1)
-            self.ups.append(
-                nn.ModuleList(
-                    [
-                        block_klass(dim_out * 2, dim_in, time_emb_dim=time_dim),
-                        block_klass(dim_in, dim_in, time_emb_dim=time_dim),
-                        Residual(PreNorm(dim_in, LinearAttention(dim_in))),
-                        Upsample(dim_in) if not is_last else nn.Identity(),
-                    ]
-                )
-            )
-
-        out_dim = default(out_dim, channels)
-        self.final_conv = nn.Sequential(
-            block_klass(dim, dim), nn.Conv2d(dim, out_dim, 1)
-        )
-
-    def forward(self, x, time, implicit_conditioning, explicit_conditioning):
-        x = torch.cat((x, explicit_conditioning), dim=1)
-        conditioning = torch.cat((implicit_conditioning, explicit_conditioning), dim=1)
-        x = self.init_conv(x)
-
-        conditioning = self.conditioning_init(conditioning)
-
-        t = self.time_mlp(time) if exists(self.time_mlp) else None
-
-        h = []
-
-        # conditioning encoder
-
-        for block1, attn, downsample in self.conditioning_encoder:
-            conditioning = block1(conditioning)
-            conditioning = attn(conditioning)
-            conditioning = downsample(conditioning)
-
-        for block1, block2, attn, downsample in self.downs:
-            x = block1(x, t)
-            x = block2(x, t)
-            x = attn(x)
-            h.append(x)
-            x = downsample(x)
-
-        # bottleneck
-        x = self.mid_block1(x, t)
-        x = self.cross_attention(x, conditioning)
-        x = self.mid_block2(x, t)
-
-        for block1, block2, attn, upsample in self.ups:
-            x = torch.cat((x, h.pop()), dim=1)
-            x = block1(x, t)
-            x = block2(x, t)
-            x = attn(x)
-            x = upsample(x)
-
-        return self.final_conv(x)

models/structure/hf_compatible_model.py

@@ -1,192 +0,0 @@
-from transformers import PretrainedConfig, PreTrainedModel
-import math
-from inspect import isfunction
-from functools import partial
-import matplotlib.pyplot as plt
-from tqdm.auto import tqdm
-from einops import rearrange
-import torch
-from torch import nn, einsum
-import torch.nn.functional as F
-from transformers import PreTrainedModel
-from .Advanced_Network_Helpers_3 import *
-import os
-
-
-class UnetConfig(PretrainedConfig):
-    model_type = "unet"
-
-    def __init__(
-        self,
-        dim=64,
-        init_dim=None,
-        out_dim=None,
-        dim_mults=(1, 2, 4, 8),
-        channels=3,
-        with_time_emb=True,
-        resnet_block_groups=8,
-        use_convnext=True,
-        convnext_mult=2,
-        **kwargs
-    ):
-        super().__init__(**kwargs)
-        self.dim = dim
-        self.init_dim = init_dim
-        self.out_dim = out_dim
-        self.dim_mults = dim_mults
-        self.channels = channels
-        self.with_time_emb = with_time_emb
-        self.resnet_block_groups = resnet_block_groups
-        self.use_convnext = use_convnext
-        self.convnext_mult = convnext_mult
-
-
-class Unet(PreTrainedModel):
-    config_class = UnetConfig
-
-    def __init__(
-        self,
-        config,
-    ):
-        super().__init__(config)
-
-        # determine dimensions
-        self.channels = (
-            config.channels
-        )  # since we are concatenating the images and the conditionings along the channel dimension
-
-        init_dim = default(config.init_dim, config.dim // 3 * 2)
-        self.init_conv = nn.Conv2d(self.channels * 2, init_dim, 7, padding=3)
-        self.conditioning_init = nn.Conv2d(self.channels, init_dim, 7, padding=3)
-        dims = [init_dim, *map(lambda m: config.dim * m, config.dim_mults)]
-        in_out = list(zip(dims[:-1], dims[1:]))
-        self.in_out = in_out
-
-        if config.use_convnext:
-            block_klass = partial(ConvNextBlock, mult=config.convnext_mult)
-        else:
-            block_klass = partial(ResnetBlock, groups=config.resnet_block_groups)
-
-        # time embeddings
-        if config.with_time_emb:
-            time_dim = config.dim * 4
-            self.time_mlp = nn.Sequential(
-                SinusoidalPositionEmbeddings(config.dim),
-                nn.Linear(config.dim, time_dim),
-                nn.GELU(),
-                nn.Linear(time_dim, time_dim),
-            )
-        else:
-            time_dim = None
-            self.time_mlp = None
-
-        # layers
-        self.downs = nn.ModuleList([])
-        self.ups = nn.ModuleList([])
-        self.conditioning_encoder = nn.ModuleList([])
-        num_resolutions = len(in_out)
-        self.num_resolutions = num_resolutions
-
-        # conditioning encoder
-        for ind, (dim_in, dim_out) in enumerate(in_out):
-            is_last = ind >= (num_resolutions - 1)
-
-            self.conditioning_encoder.append(
-                nn.ModuleList(
-                    [
-                        block_klass(dim_in, dim_out),
-                        Residual(PreNorm(dim_out, LinearAttention(dim_out))),
-                        Downsample(dim_out) if not is_last else nn.Identity(),
-                    ]
-                )
-            )
-
-        for ind, (dim_in, dim_out) in enumerate(in_out):
-            is_last = ind >= (num_resolutions - 1)
-
-            self.downs.append(
-                nn.ModuleList(
-                    [
-                        block_klass(dim_in, dim_out, time_emb_dim=time_dim),
-                        block_klass(dim_out, dim_out, time_emb_dim=time_dim),
-                        Residual(PreNorm(dim_out, LinearAttention(dim_out))),
-                        Downsample(dim_out) if not is_last else nn.Identity(),
-                    ]
-                )
-            )
-
-        mid_dim = dims[-1]
-
-        self.mid_block1 = block_klass(mid_dim, mid_dim, time_emb_dim=time_dim)
-        self.cross_attention_1 = Residual(
-            PreNorm(mid_dim, LinearCrossAttention(mid_dim))
-        )
-        self.cross_attention_2 = Residual(
-            PreNorm(mid_dim, LinearCrossAttention(mid_dim))
-        )
-        self.cross_attention_3 = Residual(
-            PreNorm(mid_dim, LinearCrossAttention(mid_dim))
-        )
-        self.mid_block2 = block_klass(mid_dim, mid_dim, time_emb_dim=time_dim)
-
-        for ind, (dim_in, dim_out) in enumerate(reversed(in_out[1:])):
-            is_last = ind >= (num_resolutions - 1)
-            self.ups.append(
-                nn.ModuleList(
-                    [
-                        block_klass(dim_out * 2, dim_in, time_emb_dim=time_dim),
-                        block_klass(dim_in, dim_in, time_emb_dim=time_dim),
-                        Residual(PreNorm(dim_in, LinearAttention(dim_in))),
-                        Upsample(dim_in) if not is_last else nn.Identity(),
-                    ]
-                )
-            )
-
-        out_dim = default(config.out_dim, config.channels)
-        self.final_conv = nn.Sequential(
-            block_klass(config.dim, config.dim), nn.Conv2d(config.dim, out_dim, 1)
-        )
-
-    def forward(self, x, time, implicit_conditioning, explicit_conditioning):
-        x = torch.cat((x, explicit_conditioning), dim=1)
-
-        x = self.init_conv(x)
-
-        conditioning = self.conditioning_init(implicit_conditioning)
-
-        t = self.time_mlp(time) if exists(self.time_mlp) else None
-
-        h = []
-
-        # conditioning encoder
-
-        for block1, attn, downsample in self.conditioning_encoder:
-            conditioning = block1(conditioning)
-            conditioning = attn(conditioning)
-            conditioning = downsample(conditioning)
-
-        for block1, block2, attn, downsample in self.downs:
-            x = block1(x, t)
-            x = block2(x, t)
-            x = attn(x)
-            h.append(x)
-            x = downsample(x)
-
-        # reverse the c list
-
-        # bottleneck
-
-        x = self.cross_attention_1(x, conditioning)
-        x = self.mid_block1(x, t)
-        x = self.cross_attention_2(x, conditioning)
-        x = self.mid_block2(x, t)
-        x = self.cross_attention_3(x, conditioning)
-
-        for block1, block2, attn, upsample in self.ups:
-            x = torch.cat((x, h.pop()), dim=1)
-            x = block1(x, t)
-            x = block2(x, t)
-            x = attn(x)
-            x = upsample(x)
-
-        return self.final_conv(x)

requirements.txt

@@ -1,14 +1,10 @@
 einops 
 datasets 
-matplotlib 
 tqdm 
 accelerate 
-jax[cpu] 
 torchinfo 
-wandb
-ema_pytorch
-lpips
-pyyaml
 diffusers
 transformers
-torch-ema
+pathlib
+safetensors
+