SD-online/Fooocus-Docker

1# https://github.com/city96/SD-Latent-Interposer/blob/main/interposer.py2 3import os4import torch5import safetensors.torch as sf6import torch.nn as nn7import ldm_patched.modules.model_management8 9from ldm_patched.modules.model_patcher import ModelPatcher10from modules.config import path_vae_approx11 12 13class Block(nn.Module):14    def __init__(self, size):15        super().__init__()16        self.join = nn.ReLU()17        self.long = nn.Sequential(18            nn.Conv2d(size, size, kernel_size=3, stride=1, padding=1),19            nn.LeakyReLU(0.1),20            nn.Conv2d(size, size, kernel_size=3, stride=1, padding=1),21            nn.LeakyReLU(0.1),22            nn.Conv2d(size, size, kernel_size=3, stride=1, padding=1),23        )24 25    def forward(self, x):26        y = self.long(x)27        z = self.join(y + x)28        return z29 30 31class Interposer(nn.Module):32    def __init__(self):33        super().__init__()34        self.chan = 435        self.hid = 12836 37        self.head_join = nn.ReLU()38        self.head_short = nn.Conv2d(self.chan, self.hid, kernel_size=3, stride=1, padding=1)39        self.head_long = nn.Sequential(40            nn.Conv2d(self.chan, self.hid, kernel_size=3, stride=1, padding=1),41            nn.LeakyReLU(0.1),42            nn.Conv2d(self.hid, self.hid, kernel_size=3, stride=1, padding=1),43            nn.LeakyReLU(0.1),44            nn.Conv2d(self.hid, self.hid, kernel_size=3, stride=1, padding=1),45        )46        self.core = nn.Sequential(47            Block(self.hid),48            Block(self.hid),49            Block(self.hid),50        )51        self.tail = nn.Sequential(52            nn.ReLU(),53            nn.Conv2d(self.hid, self.chan, kernel_size=3, stride=1, padding=1)54        )55 56    def forward(self, x):57        y = self.head_join(58            self.head_long(x) +59            self.head_short(x)60        )61        z = self.core(y)62        return self.tail(z)63 64 65vae_approx_model = None66vae_approx_filename = os.path.join(path_vae_approx, 'xl-to-v1_interposer-v3.1.safetensors')67 68 69def parse(x):70    global vae_approx_model71 72    x_origin = x.clone()73 74    if vae_approx_model is None:75        model = Interposer()76        model.eval()77        sd = sf.load_file(vae_approx_filename)78        model.load_state_dict(sd)79        fp16 = ldm_patched.modules.model_management.should_use_fp16()80        if fp16:81            model = model.half()82        vae_approx_model = ModelPatcher(83            model=model,84            load_device=ldm_patched.modules.model_management.get_torch_device(),85            offload_device=torch.device('cpu')86        )87        vae_approx_model.dtype = torch.float16 if fp16 else torch.float3288 89    ldm_patched.modules.model_management.load_model_gpu(vae_approx_model)90 91    x = x_origin.to(device=vae_approx_model.load_device, dtype=vae_approx_model.dtype)92    x = vae_approx_model.model(x).to(x_origin)93    return x94