♥️

README.md

@@ -1,12 +1,12 @@
 ---
-title: Remove Watermark
-emoji: 📚
-colorFrom: gray
-colorTo: green
+title: Remove-WM
+emoji: 🔍🗑️
+colorFrom: pink
+colorTo: purple
 sdk: gradio
 sdk_version: 4.38.1
 app_file: app.py
 pinned: false
+license: apache-2.0
 ---
-
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+https://github.com/sponsors/Damarcreative

app.py

@@ -0,0 +1,102 @@
+import gradio as gr
+import torch
+from lama_cleaner.model_manager import ModelManager
+from lama_cleaner.schema import Config, HDStrategy, LDMSampler
+from transformers import AutoProcessor, AutoModelForCausalLM
+import cv2
+import numpy as np
+from PIL import Image, ImageDraw
+import spaces
+import subprocess
+
+# Install necessary packages
+subprocess.run('pip install flash-attn --no-build-isolation', env={'FLASH_ATTENTION_SKIP_CUDA_BUILD': "TRUE"}, shell=True)
+
+# Initialize Florence model
+model_id = 'microsoft/Florence-2-large'
+florence_model = AutoModelForCausalLM.from_pretrained(model_id, trust_remote_code=True).to("cuda").eval()
+florence_processor = AutoProcessor.from_pretrained(model_id, trust_remote_code=True)
+
+# Initialize Llama Cleaner model
+device = 'cuda' if torch.cuda.is_available() else 'cpu'
+
+@spaces.GPU()
+def process_image(image, mask, strategy, sampler, fx=1, fy=1):
+    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+    mask = cv2.cvtColor(mask, cv2.COLOR_BGR2GRAY)
+
+    if fx != 1 or fy != 1:
+        image = cv2.resize(image, None, fx=fx, fy=fy, interpolation=cv2.INTER_AREA)
+        mask = cv2.resize(mask, None, fx=fx, fy=fy, interpolation=cv2.INTER_NEAREST)
+    
+    config = Config(
+        ldm_steps=1,
+        ldm_sampler=sampler,
+        hd_strategy=strategy,
+        hd_strategy_crop_margin=32,
+        hd_strategy_crop_trigger_size=200,
+        hd_strategy_resize_limit=200,
+    )
+
+    model = ModelManager(name="lama", device=device)
+    result = model(image, mask, config)
+    return result
+
+def create_mask(image, prediction):
+    mask = Image.new("RGBA", image.size, (0, 0, 0, 255))  # Black background
+    draw = ImageDraw.Draw(mask)
+    scale = 1
+    for polygons in prediction['polygons']:
+        for _polygon in polygons:
+            _polygon = np.array(_polygon).reshape(-1, 2)
+            if len(_polygon) < 3:
+                continue
+            _polygon = (_polygon * scale).reshape(-1).tolist()
+            draw.polygon(_polygon, fill=(255, 255, 255, 255))  # Make selected area white
+    return mask
+
+@spaces.GPU()
+def process_images_florence_lama(image):
+    # Convert image to OpenCV format
+    image_cv = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)
+    
+    # Run Florence to get mask
+    text_input = 'watermark'  # Teks untuk Florence agar mengenali watermark
+    task_prompt = '<REGION_TO_SEGMENTATION>'
+    image_pil = Image.fromarray(image_cv)  # Convert array to PIL Image
+    inputs = florence_processor(text=task_prompt + text_input, images=image_pil, return_tensors="pt").to("cuda")
+    generated_ids = florence_model.generate(
+        input_ids=inputs["input_ids"],
+        pixel_values=inputs["pixel_values"],
+        max_new_tokens=1024,
+        early_stopping=False,
+        do_sample=False,
+        num_beams=3,
+    )
+    generated_text = florence_processor.batch_decode(generated_ids, skip_special_tokens=False)[0]
+    parsed_answer = florence_processor.post_process_generation(
+        generated_text,
+        task=task_prompt,
+        image_size=(image_pil.width, image_pil.height)
+    )
+    
+    # Create mask and process image with Llama Cleaner
+    mask_image = create_mask(image_pil, parsed_answer['<REGION_TO_SEGMENTATION>'])
+    result_image = process_image(image_cv, np.array(mask_image), HDStrategy.RESIZE, LDMSampler.ddim)
+    
+    # Convert result back to PIL Image
+    result_image_pil = Image.fromarray(cv2.cvtColor(result_image, cv2.COLOR_BGR2RGB))
+    
+    return result_image_pil
+
+# Define Gradio interface
+demo = gr.Interface(
+    fn=process_images_florence_lama,
+    inputs=gr.Image(type="pil", label="Input Image"),
+    outputs=gr.Image(type="pil", label="Output Image"),
+    title="Watermark Remover.",
+    description="Upload images and remove selected watermarks using Florence and Llama Cleaner."
+)
+# Launch Gradio interface with example images
+if __name__ == "__main__":
+    demo.launch()

lama_cleaner/asgi.py

@@ -0,0 +1,16 @@
+"""
+ASGI config for lama_cleaner project.
+
+It exposes the ASGI callable as a module-level variable named ``application``.
+
+For more information on this file, see
+https://docs.djangoproject.com/en/4.1/howto/deployment/asgi/
+"""
+
+import os
+
+from django.core.asgi import get_asgi_application
+
+os.environ.setdefault('DJANGO_SETTINGS_MODULE', 'lama_cleaner.settings')
+
+application = get_asgi_application()

lama_cleaner/helper.py

@@ -0,0 +1,182 @@
+import os
+import sys
+from typing import List, Optional
+
+from urllib.parse import urlparse
+import cv2
+import numpy as np
+import torch
+from loguru import logger
+from torch.hub import download_url_to_file, get_dir
+
+
+def get_cache_path_by_url(url):
+    parts = urlparse(url)
+    hub_dir = get_dir()
+    model_dir = os.path.join(hub_dir, "checkpoints")
+    if not os.path.isdir(model_dir):
+        os.makedirs(os.path.join(model_dir, "hub", "checkpoints"))
+    filename = os.path.basename(parts.path)
+    cached_file = os.path.join(model_dir, filename)
+    return cached_file
+
+
+def download_model(url):
+    cached_file = get_cache_path_by_url(url)
+    if not os.path.exists(cached_file):
+        sys.stderr.write('Downloading: "{}" to {}\n'.format(url, cached_file))
+        hash_prefix = None
+        download_url_to_file(url, cached_file, hash_prefix, progress=True)
+    return cached_file
+
+
+def ceil_modulo(x, mod):
+    if x % mod == 0:
+        return x
+    return (x // mod + 1) * mod
+
+
+def load_jit_model(url_or_path, device):
+    if os.path.exists(url_or_path):
+        model_path = url_or_path
+    else:
+        model_path = download_model(url_or_path)
+    logger.info(f"Load model from: {model_path}")
+    try:
+        model = torch.jit.load(model_path).to(device)
+    except:
+        logger.error(
+            f"Failed to load {model_path}, delete model and restart lama-cleaner"
+        )
+        exit(-1)
+    model.eval()
+    return model
+
+
+def load_model(model: torch.nn.Module, url_or_path, device):
+    if os.path.exists(url_or_path):
+        model_path = url_or_path
+    else:
+        model_path = download_model(url_or_path)
+
+    try:
+        state_dict = torch.load(model_path, map_location='cpu')
+        model.load_state_dict(state_dict, strict=True)
+        model.to(device)
+        logger.info(f"Load model from: {model_path}")
+    except:
+        logger.error(
+            f"Failed to load {model_path}, delete model and restart lama-cleaner"
+        )
+        exit(-1)
+    model.eval()
+    return model
+
+
+def numpy_to_bytes(image_numpy: np.ndarray, ext: str) -> bytes:
+    data = cv2.imencode(
+        f".{ext}",
+        image_numpy,
+        [int(cv2.IMWRITE_JPEG_QUALITY), 100, int(cv2.IMWRITE_PNG_COMPRESSION), 0],
+    )[1]
+    image_bytes = data.tobytes()
+    return image_bytes
+
+
+def load_img(img_bytes, gray: bool = False):
+    alpha_channel = None
+    nparr = np.frombuffer(img_bytes, np.uint8)
+    if gray:
+        np_img = cv2.imdecode(nparr, cv2.IMREAD_GRAYSCALE)
+    else:
+        np_img = cv2.imdecode(nparr, cv2.IMREAD_UNCHANGED)
+        if len(np_img.shape) == 3 and np_img.shape[2] == 4:
+            alpha_channel = np_img[:, :, -1]
+            np_img = cv2.cvtColor(np_img, cv2.COLOR_BGRA2RGB)
+        else:
+            np_img = cv2.cvtColor(np_img, cv2.COLOR_BGR2RGB)
+
+    return np_img, alpha_channel
+
+
+def norm_img(np_img):
+    if len(np_img.shape) == 2:
+        np_img = np_img[:, :, np.newaxis]
+    np_img = np.transpose(np_img, (2, 0, 1))
+    np_img = np_img.astype("float32") / 255
+    return np_img
+
+
+def resize_max_size(
+    np_img, size_limit: int, interpolation=cv2.INTER_CUBIC
+) -> np.ndarray:
+    # Resize image's longer size to size_limit if longer size larger than size_limit
+    h, w = np_img.shape[:2]
+    if max(h, w) > size_limit:
+        ratio = size_limit / max(h, w)
+        new_w = int(w * ratio + 0.5)
+        new_h = int(h * ratio + 0.5)
+        return cv2.resize(np_img, dsize=(new_w, new_h), interpolation=interpolation)
+    else:
+        return np_img
+
+
+def pad_img_to_modulo(
+    img: np.ndarray, mod: int, square: bool = False, min_size: Optional[int] = None
+):
+    """
+
+    Args:
+        img: [H, W, C]
+        mod:
+        square: 是否为正方形
+        min_size:
+
+    Returns:
+
+    """
+    if len(img.shape) == 2:
+        img = img[:, :, np.newaxis]
+    height, width = img.shape[:2]
+    out_height = ceil_modulo(height, mod)
+    out_width = ceil_modulo(width, mod)
+
+    if min_size is not None:
+        assert min_size % mod == 0
+        out_width = max(min_size, out_width)
+        out_height = max(min_size, out_height)
+
+    if square:
+        max_size = max(out_height, out_width)
+        out_height = max_size
+        out_width = max_size
+
+    return np.pad(
+        img,
+        ((0, out_height - height), (0, out_width - width), (0, 0)),
+        mode="symmetric",
+    )
+
+
+def boxes_from_mask(mask: np.ndarray) -> List[np.ndarray]:
+    """
+    Args:
+        mask: (h, w, 1)  0~255
+
+    Returns:
+
+    """
+    height, width = mask.shape[:2]
+    _, thresh = cv2.threshold(mask, 127, 255, 0)
+    contours, _ = cv2.findContours(thresh, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
+
+    boxes = []
+    for cnt in contours:
+        x, y, w, h = cv2.boundingRect(cnt)
+        box = np.array([x, y, x + w, y + h]).astype(int)
+
+        box[::2] = np.clip(box[::2], 0, width)
+        box[1::2] = np.clip(box[1::2], 0, height)
+        boxes.append(box)
+
+    return boxes

lama_cleaner/model/base.py

@@ -0,0 +1,183 @@
+import abc
+from typing import Optional
+
+import cv2
+import torch
+from loguru import logger
+
+from lama_cleaner.helper import boxes_from_mask, resize_max_size, pad_img_to_modulo
+from lama_cleaner.schema import Config, HDStrategy
+
+
+class InpaintModel:
+    min_size: Optional[int] = None
+    pad_mod = 8
+    pad_to_square = False
+
+    def __init__(self, device, **kwargs):
+        """
+
+        Args:
+            device:
+        """
+        self.device = device
+        self.init_model(device, **kwargs)
+
+    @abc.abstractmethod
+    def init_model(self, device, **kwargs):
+        ...
+
+    @staticmethod
+    @abc.abstractmethod
+    def is_downloaded() -> bool:
+        ...
+
+    @abc.abstractmethod
+    def forward(self, image, mask, config: Config):
+        """Input images and output images have same size
+        images: [H, W, C] RGB
+        masks: [H, W, 1] 255 为 masks 区域
+        return: BGR IMAGE
+        """
+        ...
+
+    def _pad_forward(self, image, mask, config: Config):
+        origin_height, origin_width = image.shape[:2]
+        pad_image = pad_img_to_modulo(
+            image, mod=self.pad_mod, square=self.pad_to_square, min_size=self.min_size
+        )
+        pad_mask = pad_img_to_modulo(
+            mask, mod=self.pad_mod, square=self.pad_to_square, min_size=self.min_size
+        )
+
+        logger.info(f"final forward pad size: {pad_image.shape}")
+
+        result = self.forward(pad_image, pad_mask, config)
+        result = result[0:origin_height, 0:origin_width, :]
+
+        original_pixel_indices = mask < 127
+        result[original_pixel_indices] = image[:, :, ::-1][original_pixel_indices]
+        return result
+
+    @torch.no_grad()
+    def __call__(self, image, mask, config: Config):
+        """
+        images: [H, W, C] RGB, not normalized
+        masks: [H, W]
+        return: BGR IMAGE
+        """
+        inpaint_result = None
+        logger.info(f"hd_strategy: {config.hd_strategy}")
+        if config.hd_strategy == HDStrategy.CROP:
+            if max(image.shape) > config.hd_strategy_crop_trigger_size:
+                logger.info(f"Run crop strategy")
+                boxes = boxes_from_mask(mask)
+                crop_result = []
+                for box in boxes:
+                    crop_image, crop_box = self._run_box(image, mask, box, config)
+                    crop_result.append((crop_image, crop_box))
+
+                inpaint_result = image[:, :, ::-1]
+                for crop_image, crop_box in crop_result:
+                    x1, y1, x2, y2 = crop_box
+                    inpaint_result[y1:y2, x1:x2, :] = crop_image
+
+        elif config.hd_strategy == HDStrategy.RESIZE:
+            if max(image.shape) > config.hd_strategy_resize_limit:
+                origin_size = image.shape[:2]
+                downsize_image = resize_max_size(
+                    image, size_limit=config.hd_strategy_resize_limit
+                )
+                downsize_mask = resize_max_size(
+                    mask, size_limit=config.hd_strategy_resize_limit
+                )
+
+                logger.info(
+                    f"Run resize strategy, origin size: {image.shape} forward size: {downsize_image.shape}"
+                )
+                inpaint_result = self._pad_forward(
+                    downsize_image, downsize_mask, config
+                )
+
+                # only paste masked area result
+                inpaint_result = cv2.resize(
+                    inpaint_result,
+                    (origin_size[1], origin_size[0]),
+                    interpolation=cv2.INTER_CUBIC,
+                )
+                original_pixel_indices = mask < 127
+                inpaint_result[original_pixel_indices] = image[:, :, ::-1][
+                    original_pixel_indices
+                ]
+
+        if inpaint_result is None:
+            inpaint_result = self._pad_forward(image, mask, config)
+
+        return inpaint_result
+
+    def _crop_box(self, image, mask, box, config: Config):
+        """
+
+        Args:
+            image: [H, W, C] RGB
+            mask: [H, W, 1]
+            box: [left,top,right,bottom]
+
+        Returns:
+            BGR IMAGE, (l, r, r, b)
+        """
+        box_h = box[3] - box[1]
+        box_w = box[2] - box[0]
+        cx = (box[0] + box[2]) // 2
+        cy = (box[1] + box[3]) // 2
+        img_h, img_w = image.shape[:2]
+
+        w = box_w + config.hd_strategy_crop_margin * 2
+        h = box_h + config.hd_strategy_crop_margin * 2
+
+        _l = cx - w // 2
+        _r = cx + w // 2
+        _t = cy - h // 2
+        _b = cy + h // 2
+
+        l = max(_l, 0)
+        r = min(_r, img_w)
+        t = max(_t, 0)
+        b = min(_b, img_h)
+
+        # try to get more context when crop around image edge
+        if _l < 0:
+            r += abs(_l)
+        if _r > img_w:
+            l -= _r - img_w
+        if _t < 0:
+            b += abs(_t)
+        if _b > img_h:
+            t -= _b - img_h
+
+        l = max(l, 0)
+        r = min(r, img_w)
+        t = max(t, 0)
+        b = min(b, img_h)
+
+        crop_img = image[t:b, l:r, :]
+        crop_mask = mask[t:b, l:r]
+
+        logger.info(f"box size: ({box_h},{box_w}) crop size: {crop_img.shape}")
+
+        return crop_img, crop_mask, [l, t, r, b]
+
+    def _run_box(self, image, mask, box, config: Config):
+        """
+
+        Args:
+            image: [H, W, C] RGB
+            mask: [H, W, 1]
+            box: [left,top,right,bottom]
+
+        Returns:
+            BGR IMAGE
+        """
+        crop_img, crop_mask, [l, t, r, b] = self._crop_box(image, mask, box, config)
+
+        return self._pad_forward(crop_img, crop_mask, config), [l, t, r, b]

lama_cleaner/model/ddim_sampler.py

@@ -0,0 +1,193 @@
+import torch
+import numpy as np
+from tqdm import tqdm
+
+from lama_cleaner.model.utils import make_ddim_timesteps, make_ddim_sampling_parameters, noise_like
+
+from loguru import logger
+
+
+class DDIMSampler(object):
+    def __init__(self, model, schedule="linear"):
+        super().__init__()
+        self.model = model
+        self.ddpm_num_timesteps = model.num_timesteps
+        self.schedule = schedule
+
+    def register_buffer(self, name, attr):
+        setattr(self, name, attr)
+
+    def make_schedule(
+        self, ddim_num_steps, ddim_discretize="uniform", ddim_eta=0.0, verbose=True
+    ):
+        self.ddim_timesteps = make_ddim_timesteps(
+            ddim_discr_method=ddim_discretize,
+            num_ddim_timesteps=ddim_num_steps,
+            # array([1])
+            num_ddpm_timesteps=self.ddpm_num_timesteps,
+            verbose=verbose,
+        )
+        alphas_cumprod = self.model.alphas_cumprod  # torch.Size([1000])
+        assert (
+                alphas_cumprod.shape[0] == self.ddpm_num_timesteps
+        ), "alphas have to be defined for each timestep"
+        to_torch = lambda x: x.clone().detach().to(torch.float32).to(self.model.device)
+
+        self.register_buffer("betas", to_torch(self.model.betas))
+        self.register_buffer("alphas_cumprod", to_torch(alphas_cumprod))
+        self.register_buffer(
+            "alphas_cumprod_prev", to_torch(self.model.alphas_cumprod_prev)
+        )
+
+        # calculations for diffusion q(x_t | x_{t-1}) and others
+        self.register_buffer(
+            "sqrt_alphas_cumprod", to_torch(np.sqrt(alphas_cumprod.cpu()))
+        )
+        self.register_buffer(
+            "sqrt_one_minus_alphas_cumprod",
+            to_torch(np.sqrt(1.0 - alphas_cumprod.cpu())),
+        )
+        self.register_buffer(
+            "log_one_minus_alphas_cumprod", to_torch(np.log(1.0 - alphas_cumprod.cpu()))
+        )
+        self.register_buffer(
+            "sqrt_recip_alphas_cumprod", to_torch(np.sqrt(1.0 / alphas_cumprod.cpu()))
+        )
+        self.register_buffer(
+            "sqrt_recipm1_alphas_cumprod",
+            to_torch(np.sqrt(1.0 / alphas_cumprod.cpu() - 1)),
+        )
+
+        # ddim sampling parameters
+        ddim_sigmas, ddim_alphas, ddim_alphas_prev = make_ddim_sampling_parameters(
+            alphacums=alphas_cumprod.cpu(),
+            ddim_timesteps=self.ddim_timesteps,
+            eta=ddim_eta,
+            verbose=verbose,
+        )
+        self.register_buffer("ddim_sigmas", ddim_sigmas)
+        self.register_buffer("ddim_alphas", ddim_alphas)
+        self.register_buffer("ddim_alphas_prev", ddim_alphas_prev)
+        self.register_buffer("ddim_sqrt_one_minus_alphas", np.sqrt(1.0 - ddim_alphas))
+        sigmas_for_original_sampling_steps = ddim_eta * torch.sqrt(
+            (1 - self.alphas_cumprod_prev)
+            / (1 - self.alphas_cumprod)
+            * (1 - self.alphas_cumprod / self.alphas_cumprod_prev)
+        )
+        self.register_buffer(
+            "ddim_sigmas_for_original_num_steps", sigmas_for_original_sampling_steps
+        )
+
+    @torch.no_grad()
+    def sample(self, steps, conditioning, batch_size, shape):
+        self.make_schedule(ddim_num_steps=steps, ddim_eta=0, verbose=False)
+        # sampling
+        C, H, W = shape
+        size = (batch_size, C, H, W)
+
+        # samples: 1,3,128,128
+        return self.ddim_sampling(
+            conditioning,
+            size,
+            quantize_denoised=False,
+            ddim_use_original_steps=False,
+            noise_dropout=0,
+            temperature=1.0,
+        )
+
+    @torch.no_grad()
+    def ddim_sampling(
+        self,
+        cond,
+        shape,
+        ddim_use_original_steps=False,
+        quantize_denoised=False,
+        temperature=1.0,
+        noise_dropout=0.0,
+    ):
+        device = self.model.betas.device
+        b = shape[0]
+        img = torch.randn(shape, device=device, dtype=cond.dtype)
+        timesteps = (
+            self.ddpm_num_timesteps if ddim_use_original_steps else self.ddim_timesteps
+        )
+
+        time_range = (
+            reversed(range(0, timesteps))
+            if ddim_use_original_steps
+            else np.flip(timesteps)
+        )
+        total_steps = timesteps if ddim_use_original_steps else timesteps.shape[0]
+        logger.info(f"Running DDIM Sampling with {total_steps} timesteps")
+
+        iterator = tqdm(time_range, desc="DDIM Sampler", total=total_steps)
+
+        for i, step in enumerate(iterator):
+            index = total_steps - i - 1
+            ts = torch.full((b,), step, device=device, dtype=torch.long)
+
+            outs = self.p_sample_ddim(
+                img,
+                cond,
+                ts,
+                index=index,
+                use_original_steps=ddim_use_original_steps,
+                quantize_denoised=quantize_denoised,
+                temperature=temperature,
+                noise_dropout=noise_dropout,
+            )
+            img, _ = outs
+
+        return img
+
+    @torch.no_grad()
+    def p_sample_ddim(
+        self,
+        x,
+        c,
+        t,
+        index,
+        repeat_noise=False,
+        use_original_steps=False,
+        quantize_denoised=False,
+        temperature=1.0,
+        noise_dropout=0.0,
+    ):
+        b, *_, device = *x.shape, x.device
+        e_t = self.model.apply_model(x, t, c)
+
+        alphas = self.model.alphas_cumprod if use_original_steps else self.ddim_alphas
+        alphas_prev = (
+            self.model.alphas_cumprod_prev
+            if use_original_steps
+            else self.ddim_alphas_prev
+        )
+        sqrt_one_minus_alphas = (
+            self.model.sqrt_one_minus_alphas_cumprod
+            if use_original_steps
+            else self.ddim_sqrt_one_minus_alphas
+        )
+        sigmas = (
+            self.model.ddim_sigmas_for_original_num_steps
+            if use_original_steps
+            else self.ddim_sigmas
+        )
+        # select parameters corresponding to the currently considered timestep
+        a_t = torch.full((b, 1, 1, 1), alphas[index], device=device)
+        a_prev = torch.full((b, 1, 1, 1), alphas_prev[index], device=device)
+        sigma_t = torch.full((b, 1, 1, 1), sigmas[index], device=device)
+        sqrt_one_minus_at = torch.full(
+            (b, 1, 1, 1), sqrt_one_minus_alphas[index], device=device
+        )
+
+        # current prediction for x_0
+        pred_x0 = (x - sqrt_one_minus_at * e_t) / a_t.sqrt()
+        if quantize_denoised:  # 没用
+            pred_x0, _, *_ = self.model.first_stage_model.quantize(pred_x0)
+        # direction pointing to x_t
+        dir_xt = (1.0 - a_prev - sigma_t ** 2).sqrt() * e_t
+        noise = sigma_t * noise_like(x.shape, device, repeat_noise) * temperature
+        if noise_dropout > 0.0:  # 没用
+            noise = torch.nn.functional.dropout(noise, p=noise_dropout)
+        x_prev = a_prev.sqrt() * pred_x0 + dir_xt + noise
+        return x_prev, pred_x0

lama_cleaner/model/fcf.py

@@ -0,0 +1,1214 @@
+import os
+import random
+
+import cv2
+import torch
+import numpy as np
+import torch.fft as fft
+
+from lama_cleaner.schema import Config
+
+from lama_cleaner.helper import load_model, get_cache_path_by_url, norm_img, boxes_from_mask, resize_max_size
+from lama_cleaner.model.base import InpaintModel
+from torch import conv2d, nn
+import torch.nn.functional as F
+
+from lama_cleaner.model.utils import setup_filter, _parse_scaling, _parse_padding, Conv2dLayer, FullyConnectedLayer, \
+    MinibatchStdLayer, activation_funcs, conv2d_resample, bias_act, upsample2d, normalize_2nd_moment, downsample2d
+
+
+def upfirdn2d(x, f, up=1, down=1, padding=0, flip_filter=False, gain=1, impl='cuda'):
+    assert isinstance(x, torch.Tensor)
+    return _upfirdn2d_ref(x, f, up=up, down=down, padding=padding, flip_filter=flip_filter, gain=gain)
+
+
+def _upfirdn2d_ref(x, f, up=1, down=1, padding=0, flip_filter=False, gain=1):
+    """Slow reference implementation of `upfirdn2d()` using standard PyTorch ops.
+    """
+    # Validate arguments.
+    assert isinstance(x, torch.Tensor) and x.ndim == 4
+    if f is None:
+        f = torch.ones([1, 1], dtype=torch.float32, device=x.device)
+    assert isinstance(f, torch.Tensor) and f.ndim in [1, 2]
+    assert f.dtype == torch.float32 and not f.requires_grad
+    batch_size, num_channels, in_height, in_width = x.shape
+    upx, upy = _parse_scaling(up)
+    downx, downy = _parse_scaling(down)
+    padx0, padx1, pady0, pady1 = _parse_padding(padding)
+
+    # Upsample by inserting zeros.
+    x = x.reshape([batch_size, num_channels, in_height, 1, in_width, 1])
+    x = torch.nn.functional.pad(x, [0, upx - 1, 0, 0, 0, upy - 1])
+    x = x.reshape([batch_size, num_channels, in_height * upy, in_width * upx])
+
+    # Pad or crop.
+    x = torch.nn.functional.pad(x, [max(padx0, 0), max(padx1, 0), max(pady0, 0), max(pady1, 0)])
+    x = x[:, :, max(-pady0, 0): x.shape[2] - max(-pady1, 0), max(-padx0, 0): x.shape[3] - max(-padx1, 0)]
+
+    # Setup filter.
+    f = f * (gain ** (f.ndim / 2))
+    f = f.to(x.dtype)
+    if not flip_filter:
+        f = f.flip(list(range(f.ndim)))
+
+    # Convolve with the filter.
+    f = f[np.newaxis, np.newaxis].repeat([num_channels, 1] + [1] * f.ndim)
+    if f.ndim == 4:
+        x = conv2d(input=x, weight=f, groups=num_channels)
+    else:
+        x = conv2d(input=x, weight=f.unsqueeze(2), groups=num_channels)
+        x = conv2d(input=x, weight=f.unsqueeze(3), groups=num_channels)
+
+    # Downsample by throwing away pixels.
+    x = x[:, :, ::downy, ::downx]
+    return x
+
+
+class EncoderEpilogue(torch.nn.Module):
+    def __init__(self,
+                 in_channels,  # Number of input channels.
+                 cmap_dim,  # Dimensionality of mapped conditioning label, 0 = no label.
+                 z_dim,  # Output Latent (Z) dimensionality.
+                 resolution,  # Resolution of this block.
+                 img_channels,  # Number of input color channels.
+                 architecture='resnet',  # Architecture: 'orig', 'skip', 'resnet'.
+                 mbstd_group_size=4,  # Group size for the minibatch standard deviation layer, None = entire minibatch.
+                 mbstd_num_channels=1,  # Number of features for the minibatch standard deviation layer, 0 = disable.
+                 activation='lrelu',  # Activation function: 'relu', 'lrelu', etc.
+                 conv_clamp=None,  # Clamp the output of convolution layers to +-X, None = disable clamping.
+                 ):
+        assert architecture in ['orig', 'skip', 'resnet']
+        super().__init__()
+        self.in_channels = in_channels
+        self.cmap_dim = cmap_dim
+        self.resolution = resolution
+        self.img_channels = img_channels
+        self.architecture = architecture
+
+        if architecture == 'skip':
+            self.fromrgb = Conv2dLayer(self.img_channels, in_channels, kernel_size=1, activation=activation)
+        self.mbstd = MinibatchStdLayer(group_size=mbstd_group_size,
+                                       num_channels=mbstd_num_channels) if mbstd_num_channels > 0 else None
+        self.conv = Conv2dLayer(in_channels + mbstd_num_channels, in_channels, kernel_size=3, activation=activation,
+                                conv_clamp=conv_clamp)
+        self.fc = FullyConnectedLayer(in_channels * (resolution ** 2), z_dim, activation=activation)
+        self.dropout = torch.nn.Dropout(p=0.5)
+
+    def forward(self, x, cmap, force_fp32=False):
+        _ = force_fp32  # unused
+        dtype = torch.float32
+        memory_format = torch.contiguous_format
+
+        # FromRGB.
+        x = x.to(dtype=dtype, memory_format=memory_format)
+
+        # Main layers.
+        if self.mbstd is not None:
+            x = self.mbstd(x)
+        const_e = self.conv(x)
+        x = self.fc(const_e.flatten(1))
+        x = self.dropout(x)
+
+        # Conditioning.
+        if self.cmap_dim > 0:
+            x = (x * cmap).sum(dim=1, keepdim=True) * (1 / np.sqrt(self.cmap_dim))
+
+        assert x.dtype == dtype
+        return x, const_e
+
+
+class EncoderBlock(torch.nn.Module):
+    def __init__(self,
+                 in_channels,  # Number of input channels, 0 = first block.
+                 tmp_channels,  # Number of intermediate channels.
+                 out_channels,  # Number of output channels.
+                 resolution,  # Resolution of this block.
+                 img_channels,  # Number of input color channels.
+                 first_layer_idx,  # Index of the first layer.
+                 architecture='skip',  # Architecture: 'orig', 'skip', 'resnet'.
+                 activation='lrelu',  # Activation function: 'relu', 'lrelu', etc.
+                 resample_filter=[1, 3, 3, 1],  # Low-pass filter to apply when resampling activations.
+                 conv_clamp=None,  # Clamp the output of convolution layers to +-X, None = disable clamping.
+                 use_fp16=False,  # Use FP16 for this block?
+                 fp16_channels_last=False,  # Use channels-last memory format with FP16?
+                 freeze_layers=0,  # Freeze-D: Number of layers to freeze.
+                 ):
+        assert in_channels in [0, tmp_channels]
+        assert architecture in ['orig', 'skip', 'resnet']
+        super().__init__()
+        self.in_channels = in_channels
+        self.resolution = resolution
+        self.img_channels = img_channels + 1
+        self.first_layer_idx = first_layer_idx
+        self.architecture = architecture
+        self.use_fp16 = use_fp16
+        self.channels_last = (use_fp16 and fp16_channels_last)
+        self.register_buffer('resample_filter', setup_filter(resample_filter))
+
+        self.num_layers = 0
+
+        def trainable_gen():
+            while True:
+                layer_idx = self.first_layer_idx + self.num_layers
+                trainable = (layer_idx >= freeze_layers)
+                self.num_layers += 1
+                yield trainable
+
+        trainable_iter = trainable_gen()
+
+        if in_channels == 0:
+            self.fromrgb = Conv2dLayer(self.img_channels, tmp_channels, kernel_size=1, activation=activation,
+                                       trainable=next(trainable_iter), conv_clamp=conv_clamp,
+                                       channels_last=self.channels_last)
+
+        self.conv0 = Conv2dLayer(tmp_channels, tmp_channels, kernel_size=3, activation=activation,
+                                 trainable=next(trainable_iter), conv_clamp=conv_clamp,
+                                 channels_last=self.channels_last)
+
+        self.conv1 = Conv2dLayer(tmp_channels, out_channels, kernel_size=3, activation=activation, down=2,
+                                 trainable=next(trainable_iter), resample_filter=resample_filter, conv_clamp=conv_clamp,
+                                 channels_last=self.channels_last)
+
+        if architecture == 'resnet':
+            self.skip = Conv2dLayer(tmp_channels, out_channels, kernel_size=1, bias=False, down=2,
+                                    trainable=next(trainable_iter), resample_filter=resample_filter,
+                                    channels_last=self.channels_last)
+
+    def forward(self, x, img, force_fp32=False):
+        # dtype = torch.float16 if self.use_fp16 and not force_fp32 else torch.float32
+        dtype = torch.float32
+        memory_format = torch.channels_last if self.channels_last and not force_fp32 else torch.contiguous_format
+
+        # Input.
+        if x is not None:
+            x = x.to(dtype=dtype, memory_format=memory_format)
+
+        # FromRGB.
+        if self.in_channels == 0:
+            img = img.to(dtype=dtype, memory_format=memory_format)
+            y = self.fromrgb(img)
+            x = x + y if x is not None else y
+            img = downsample2d(img, self.resample_filter) if self.architecture == 'skip' else None
+
+        # Main layers.
+        if self.architecture == 'resnet':
+            y = self.skip(x, gain=np.sqrt(0.5))
+            x = self.conv0(x)
+            feat = x.clone()
+            x = self.conv1(x, gain=np.sqrt(0.5))
+            x = y.add_(x)
+        else:
+            x = self.conv0(x)
+            feat = x.clone()
+            x = self.conv1(x)
+
+        assert x.dtype == dtype
+        return x, img, feat
+
+
+class EncoderNetwork(torch.nn.Module):
+    def __init__(self,
+                 c_dim,  # Conditioning label (C) dimensionality.
+                 z_dim,  # Input latent (Z) dimensionality.
+                 img_resolution,  # Input resolution.
+                 img_channels,  # Number of input color channels.
+                 architecture='orig',  # Architecture: 'orig', 'skip', 'resnet'.
+                 channel_base=16384,  # Overall multiplier for the number of channels.
+                 channel_max=512,  # Maximum number of channels in any layer.
+                 num_fp16_res=0,  # Use FP16 for the N highest resolutions.
+                 conv_clamp=None,  # Clamp the output of convolution layers to +-X, None = disable clamping.
+                 cmap_dim=None,  # Dimensionality of mapped conditioning label, None = default.
+                 block_kwargs={},  # Arguments for DiscriminatorBlock.
+                 mapping_kwargs={},  # Arguments for MappingNetwork.
+                 epilogue_kwargs={},  # Arguments for EncoderEpilogue.
+                 ):
+        super().__init__()
+        self.c_dim = c_dim
+        self.z_dim = z_dim
+        self.img_resolution = img_resolution
+        self.img_resolution_log2 = int(np.log2(img_resolution))
+        self.img_channels = img_channels
+        self.block_resolutions = [2 ** i for i in range(self.img_resolution_log2, 2, -1)]
+        channels_dict = {res: min(channel_base // res, channel_max) for res in self.block_resolutions + [4]}
+        fp16_resolution = max(2 ** (self.img_resolution_log2 + 1 - num_fp16_res), 8)
+
+        if cmap_dim is None:
+            cmap_dim = channels_dict[4]
+        if c_dim == 0:
+            cmap_dim = 0
+
+        common_kwargs = dict(img_channels=img_channels, architecture=architecture, conv_clamp=conv_clamp)
+        cur_layer_idx = 0
+        for res in self.block_resolutions:
+            in_channels = channels_dict[res] if res < img_resolution else 0
+            tmp_channels = channels_dict[res]
+            out_channels = channels_dict[res // 2]
+            use_fp16 = (res >= fp16_resolution)
+            use_fp16 = False
+            block = EncoderBlock(in_channels, tmp_channels, out_channels, resolution=res,
+                                 first_layer_idx=cur_layer_idx, use_fp16=use_fp16, **block_kwargs, **common_kwargs)
+            setattr(self, f'b{res}', block)
+            cur_layer_idx += block.num_layers
+        if c_dim > 0:
+            self.mapping = MappingNetwork(z_dim=0, c_dim=c_dim, w_dim=cmap_dim, num_ws=None, w_avg_beta=None,
+                                          **mapping_kwargs)
+        self.b4 = EncoderEpilogue(channels_dict[4], cmap_dim=cmap_dim, z_dim=z_dim * 2, resolution=4, **epilogue_kwargs,
+                                  **common_kwargs)
+
+    def forward(self, img, c, **block_kwargs):
+        x = None
+        feats = {}
+        for res in self.block_resolutions:
+            block = getattr(self, f'b{res}')
+            x, img, feat = block(x, img, **block_kwargs)
+            feats[res] = feat
+
+        cmap = None
+        if self.c_dim > 0:
+            cmap = self.mapping(None, c)
+        x, const_e = self.b4(x, cmap)
+        feats[4] = const_e
+
+        B, _ = x.shape
+        z = torch.zeros((B, self.z_dim), requires_grad=False, dtype=x.dtype,
+                        device=x.device)  ## Noise for Co-Modulation
+        return x, z, feats
+
+
+def fma(a, b, c):  # => a * b + c
+    return _FusedMultiplyAdd.apply(a, b, c)
+
+
+class _FusedMultiplyAdd(torch.autograd.Function):  # a * b + c
+    @staticmethod
+    def forward(ctx, a, b, c):  # pylint: disable=arguments-differ
+        out = torch.addcmul(c, a, b)
+        ctx.save_for_backward(a, b)
+        ctx.c_shape = c.shape
+        return out
+
+    @staticmethod
+    def backward(ctx, dout):  # pylint: disable=arguments-differ
+        a, b = ctx.saved_tensors
+        c_shape = ctx.c_shape
+        da = None
+        db = None
+        dc = None
+
+        if ctx.needs_input_grad[0]:
+            da = _unbroadcast(dout * b, a.shape)
+
+        if ctx.needs_input_grad[1]:
+            db = _unbroadcast(dout * a, b.shape)
+
+        if ctx.needs_input_grad[2]:
+            dc = _unbroadcast(dout, c_shape)
+
+        return da, db, dc
+
+
+def _unbroadcast(x, shape):
+    extra_dims = x.ndim - len(shape)
+    assert extra_dims >= 0
+    dim = [i for i in range(x.ndim) if x.shape[i] > 1 and (i < extra_dims or shape[i - extra_dims] == 1)]
+    if len(dim):
+        x = x.sum(dim=dim, keepdim=True)
+    if extra_dims:
+        x = x.reshape(-1, *x.shape[extra_dims + 1:])
+    assert x.shape == shape
+    return x
+
+
+def modulated_conv2d(
+    x,  # Input tensor of shape [batch_size, in_channels, in_height, in_width].
+    weight,  # Weight tensor of shape [out_channels, in_channels, kernel_height, kernel_width].
+    styles,  # Modulation coefficients of shape [batch_size, in_channels].
+    noise=None,  # Optional noise tensor to add to the output activations.
+    up=1,  # Integer upsampling factor.
+    down=1,  # Integer downsampling factor.
+    padding=0,  # Padding with respect to the upsampled image.
+    resample_filter=None,
+    # Low-pass filter to apply when resampling activations. Must be prepared beforehand by calling upfirdn2d.setup_filter().
+    demodulate=True,  # Apply weight demodulation?
+    flip_weight=True,  # False = convolution, True = correlation (matches torch.nn.functional.conv2d).
+    fused_modconv=True,  # Perform modulation, convolution, and demodulation as a single fused operation?
+):
+    batch_size = x.shape[0]
+    out_channels, in_channels, kh, kw = weight.shape
+
+    # Pre-normalize inputs to avoid FP16 overflow.
+    if x.dtype == torch.float16 and demodulate:
+        weight = weight * (1 / np.sqrt(in_channels * kh * kw) / weight.norm(float('inf'), dim=[1, 2, 3],
+                                                                            keepdim=True))  # max_Ikk
+        styles = styles / styles.norm(float('inf'), dim=1, keepdim=True)  # max_I
+
+    # Calculate per-sample weights and demodulation coefficients.
+    w = None
+    dcoefs = None
+    if demodulate or fused_modconv:
+        w = weight.unsqueeze(0)  # [NOIkk]
+        w = w * styles.reshape(batch_size, 1, -1, 1, 1)  # [NOIkk]
+    if demodulate:
+        dcoefs = (w.square().sum(dim=[2, 3, 4]) + 1e-8).rsqrt()  # [NO]
+    if demodulate and fused_modconv:
+        w = w * dcoefs.reshape(batch_size, -1, 1, 1, 1)  # [NOIkk]
+    # Execute by scaling the activations before and after the convolution.
+    if not fused_modconv:
+        x = x * styles.to(x.dtype).reshape(batch_size, -1, 1, 1)
+        x = conv2d_resample.conv2d_resample(x=x, w=weight.to(x.dtype), f=resample_filter, up=up, down=down,
+                                            padding=padding, flip_weight=flip_weight)
+        if demodulate and noise is not None:
+            x = fma(x, dcoefs.to(x.dtype).reshape(batch_size, -1, 1, 1), noise.to(x.dtype))
+        elif demodulate:
+            x = x * dcoefs.to(x.dtype).reshape(batch_size, -1, 1, 1)
+        elif noise is not None:
+            x = x.add_(noise.to(x.dtype))
+        return x
+
+    # Execute as one fused op using grouped convolution.
+    batch_size = int(batch_size)
+    x = x.reshape(1, -1, *x.shape[2:])
+    w = w.reshape(-1, in_channels, kh, kw)
+    x = conv2d_resample(x=x, w=w.to(x.dtype), f=resample_filter, up=up, down=down, padding=padding,
+                        groups=batch_size, flip_weight=flip_weight)
+    x = x.reshape(batch_size, -1, *x.shape[2:])
+    if noise is not None:
+        x = x.add_(noise)
+    return x
+
+
+class SynthesisLayer(torch.nn.Module):
+    def __init__(self,
+                 in_channels,  # Number of input channels.
+                 out_channels,  # Number of output channels.
+                 w_dim,  # Intermediate latent (W) dimensionality.
+                 resolution,  # Resolution of this layer.
+                 kernel_size=3,  # Convolution kernel size.
+                 up=1,  # Integer upsampling factor.
+                 use_noise=True,  # Enable noise input?
+                 activation='lrelu',  # Activation function: 'relu', 'lrelu', etc.
+                 resample_filter=[1, 3, 3, 1],  # Low-pass filter to apply when resampling activations.
+                 conv_clamp=None,  # Clamp the output of convolution layers to +-X, None = disable clamping.
+                 channels_last=False,  # Use channels_last format for the weights?
+                 ):
+        super().__init__()
+        self.resolution = resolution
+        self.up = up
+        self.use_noise = use_noise
+        self.activation = activation
+        self.conv_clamp = conv_clamp
+        self.register_buffer('resample_filter', setup_filter(resample_filter))
+        self.padding = kernel_size // 2
+        self.act_gain = activation_funcs[activation].def_gain
+
+        self.affine = FullyConnectedLayer(w_dim, in_channels, bias_init=1)
+        memory_format = torch.channels_last if channels_last else torch.contiguous_format
+        self.weight = torch.nn.Parameter(
+            torch.randn([out_channels, in_channels, kernel_size, kernel_size]).to(memory_format=memory_format))
+        if use_noise:
+            self.register_buffer('noise_const', torch.randn([resolution, resolution]))
+            self.noise_strength = torch.nn.Parameter(torch.zeros([]))
+        self.bias = torch.nn.Parameter(torch.zeros([out_channels]))
+
+    def forward(self, x, w, noise_mode='none', fused_modconv=True, gain=1):
+        assert noise_mode in ['random', 'const', 'none']
+        in_resolution = self.resolution // self.up
+        styles = self.affine(w)
+
+        noise = None
+        if self.use_noise and noise_mode == 'random':
+            noise = torch.randn([x.shape[0], 1, self.resolution, self.resolution],
+                                device=x.device) * self.noise_strength
+        if self.use_noise and noise_mode == 'const':
+            noise = self.noise_const * self.noise_strength
+
+        flip_weight = (self.up == 1)  # slightly faster
+        x = modulated_conv2d(x=x, weight=self.weight, styles=styles, noise=noise, up=self.up,
+                             padding=self.padding, resample_filter=self.resample_filter, flip_weight=flip_weight,
+                             fused_modconv=fused_modconv)
+
+        act_gain = self.act_gain * gain
+        act_clamp = self.conv_clamp * gain if self.conv_clamp is not None else None
+        x = F.leaky_relu(x, negative_slope=0.2, inplace=False)
+        if act_gain != 1:
+            x = x * act_gain
+        if act_clamp is not None:
+            x = x.clamp(-act_clamp, act_clamp)
+        return x
+
+
+class ToRGBLayer(torch.nn.Module):
+    def __init__(self, in_channels, out_channels, w_dim, kernel_size=1, conv_clamp=None, channels_last=False):
+        super().__init__()
+        self.conv_clamp = conv_clamp
+        self.affine = FullyConnectedLayer(w_dim, in_channels, bias_init=1)
+        memory_format = torch.channels_last if channels_last else torch.contiguous_format
+        self.weight = torch.nn.Parameter(
+            torch.randn([out_channels, in_channels, kernel_size, kernel_size]).to(memory_format=memory_format))
+        self.bias = torch.nn.Parameter(torch.zeros([out_channels]))
+        self.weight_gain = 1 / np.sqrt(in_channels * (kernel_size ** 2))
+
+    def forward(self, x, w, fused_modconv=True):
+        styles = self.affine(w) * self.weight_gain
+        x = modulated_conv2d(x=x, weight=self.weight, styles=styles, demodulate=False, fused_modconv=fused_modconv)
+        x = bias_act(x, self.bias.to(x.dtype), clamp=self.conv_clamp)
+        return x
+
+
+class SynthesisForeword(torch.nn.Module):
+    def __init__(self,
+                 z_dim,  # Output Latent (Z) dimensionality.
+                 resolution,  # Resolution of this block.
+                 in_channels,
+                 img_channels,  # Number of input color channels.
+                 architecture='skip',  # Architecture: 'orig', 'skip', 'resnet'.
+                 activation='lrelu',  # Activation function: 'relu', 'lrelu', etc.
+
+                 ):
+        super().__init__()
+        self.in_channels = in_channels
+        self.z_dim = z_dim
+        self.resolution = resolution
+        self.img_channels = img_channels
+        self.architecture = architecture
+
+        self.fc = FullyConnectedLayer(self.z_dim, (self.z_dim // 2) * 4 * 4, activation=activation)
+        self.conv = SynthesisLayer(self.in_channels, self.in_channels, w_dim=(z_dim // 2) * 3, resolution=4)
+
+        if architecture == 'skip':
+            self.torgb = ToRGBLayer(self.in_channels, self.img_channels, kernel_size=1, w_dim=(z_dim // 2) * 3)
+
+    def forward(self, x, ws, feats, img, force_fp32=False):
+        _ = force_fp32  # unused
+        dtype = torch.float32
+        memory_format = torch.contiguous_format
+
+        x_global = x.clone()
+        # ToRGB.
+        x = self.fc(x)
+        x = x.view(-1, self.z_dim // 2, 4, 4)
+        x = x.to(dtype=dtype, memory_format=memory_format)
+
+        # Main layers.
+        x_skip = feats[4].clone()
+        x = x + x_skip
+
+        mod_vector = []
+        mod_vector.append(ws[:, 0])
+        mod_vector.append(x_global.clone())
+        mod_vector = torch.cat(mod_vector, dim=1)
+
+        x = self.conv(x, mod_vector)
+
+        mod_vector = []
+        mod_vector.append(ws[:, 2 * 2 - 3])
+        mod_vector.append(x_global.clone())
+        mod_vector = torch.cat(mod_vector, dim=1)
+
+        if self.architecture == 'skip':
+            img = self.torgb(x, mod_vector)
+            img = img.to(dtype=torch.float32, memory_format=torch.contiguous_format)
+
+        assert x.dtype == dtype
+        return x, img
+
+
+class SELayer(nn.Module):
+    def __init__(self, channel, reduction=16):
+        super(SELayer, self).__init__()
+        self.avg_pool = nn.AdaptiveAvgPool2d(1)
+        self.fc = nn.Sequential(
+            nn.Linear(channel, channel // reduction, bias=False),
+            nn.ReLU(inplace=False),
+            nn.Linear(channel // reduction, channel, 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)
+        res = x * y.expand_as(x)
+        return res
+
+
+class FourierUnit(nn.Module):
+
+    def __init__(self, in_channels, out_channels, groups=1, spatial_scale_factor=None, spatial_scale_mode='bilinear',
+                 spectral_pos_encoding=False, use_se=False, se_kwargs=None, ffc3d=False, fft_norm='ortho'):
+        # bn_layer not used
+        super(FourierUnit, self).__init__()
+        self.groups = groups
+
+        self.conv_layer = torch.nn.Conv2d(in_channels=in_channels * 2 + (2 if spectral_pos_encoding else 0),
+                                          out_channels=out_channels * 2,
+                                          kernel_size=1, stride=1, padding=0, groups=self.groups, bias=False)
+        self.relu = torch.nn.ReLU(inplace=False)
+
+        # squeeze and excitation block
+        self.use_se = use_se
+        if use_se:
+            if se_kwargs is None:
+                se_kwargs = {}
+            self.se = SELayer(self.conv_layer.in_channels, **se_kwargs)
+
+        self.spatial_scale_factor = spatial_scale_factor
+        self.spatial_scale_mode = spatial_scale_mode
+        self.spectral_pos_encoding = spectral_pos_encoding
+        self.ffc3d = ffc3d
+        self.fft_norm = fft_norm
+
+    def forward(self, x):
+        batch = x.shape[0]
+
+        if self.spatial_scale_factor is not None:
+            orig_size = x.shape[-2:]
+            x = F.interpolate(x, scale_factor=self.spatial_scale_factor, mode=self.spatial_scale_mode,
+                              align_corners=False)
+
+        r_size = x.size()
+        # (batch, c, h, w/2+1, 2)
+        fft_dim = (-3, -2, -1) if self.ffc3d else (-2, -1)
+        ffted = fft.rfftn(x, dim=fft_dim, norm=self.fft_norm)
+        ffted = torch.stack((ffted.real, ffted.imag), dim=-1)
+        ffted = ffted.permute(0, 1, 4, 2, 3).contiguous()  # (batch, c, 2, h, w/2+1)
+        ffted = ffted.view((batch, -1,) + ffted.size()[3:])
+
+        if self.spectral_pos_encoding:
+            height, width = ffted.shape[-2:]
+            coords_vert = torch.linspace(0, 1, height)[None, None, :, None].expand(batch, 1, height, width).to(ffted)
+            coords_hor = torch.linspace(0, 1, width)[None, None, None, :].expand(batch, 1, height, width).to(ffted)
+            ffted = torch.cat((coords_vert, coords_hor, ffted), dim=1)
+
+        if self.use_se:
+            ffted = self.se(ffted)
+
+        ffted = self.conv_layer(ffted)  # (batch, c*2, h, w/2+1)
+        ffted = self.relu(ffted)
+
+        ffted = ffted.view((batch, -1, 2,) + ffted.size()[2:]).permute(
+            0, 1, 3, 4, 2).contiguous()  # (batch,c, t, h, w/2+1, 2)
+        ffted = torch.complex(ffted[..., 0], ffted[..., 1])
+
+        ifft_shape_slice = x.shape[-3:] if self.ffc3d else x.shape[-2:]
+        output = torch.fft.irfftn(ffted, s=ifft_shape_slice, dim=fft_dim, norm=self.fft_norm)
+
+        if self.spatial_scale_factor is not None:
+            output = F.interpolate(output, size=orig_size, mode=self.spatial_scale_mode, align_corners=False)
+
+        return output
+
+
+class SpectralTransform(nn.Module):
+
+    def __init__(self, in_channels, out_channels, stride=1, groups=1, enable_lfu=True, **fu_kwargs):
+        # bn_layer not used
+        super(SpectralTransform, self).__init__()
+        self.enable_lfu = enable_lfu
+        if stride == 2:
+            self.downsample = nn.AvgPool2d(kernel_size=(2, 2), stride=2)
+        else:
+            self.downsample = nn.Identity()
+
+        self.stride = stride
+        self.conv1 = nn.Sequential(
+            nn.Conv2d(in_channels, out_channels //
+                      2, kernel_size=1, groups=groups, bias=False),
+            # nn.BatchNorm2d(out_channels // 2),
+            nn.ReLU(inplace=True)
+        )
+        self.fu = FourierUnit(
+            out_channels // 2, out_channels // 2, groups, **fu_kwargs)
+        if self.enable_lfu:
+            self.lfu = FourierUnit(
+                out_channels // 2, out_channels // 2, groups)
+        self.conv2 = torch.nn.Conv2d(
+            out_channels // 2, out_channels, kernel_size=1, groups=groups, bias=False)
+
+    def forward(self, x):
+
+        x = self.downsample(x)
+        x = self.conv1(x)
+        output = self.fu(x)
+
+        if self.enable_lfu:
+            n, c, h, w = x.shape
+            split_no = 2
+            split_s = h // split_no
+            xs = torch.cat(torch.split(
+                x[:, :c // 4], split_s, dim=-2), dim=1).contiguous()
+            xs = torch.cat(torch.split(xs, split_s, dim=-1),
+                           dim=1).contiguous()
+            xs = self.lfu(xs)
+            xs = xs.repeat(1, 1, split_no, split_no).contiguous()
+        else:
+            xs = 0
+
+        output = self.conv2(x + output + xs)
+
+        return output
+
+
+class FFC(nn.Module):
+
+    def __init__(self, in_channels, out_channels, kernel_size,
+                 ratio_gin, ratio_gout, stride=1, padding=0,
+                 dilation=1, groups=1, bias=False, enable_lfu=True,
+                 padding_type='reflect', gated=False, **spectral_kwargs):
+        super(FFC, self).__init__()
+
+        assert stride == 1 or stride == 2, "Stride should be 1 or 2."
+        self.stride = stride
+
+        in_cg = int(in_channels * ratio_gin)
+        in_cl = in_channels - in_cg
+        out_cg = int(out_channels * ratio_gout)
+        out_cl = out_channels - out_cg
+        # groups_g = 1 if groups == 1 else int(groups * ratio_gout)
+        # groups_l = 1 if groups == 1 else groups - groups_g
+
+        self.ratio_gin = ratio_gin
+        self.ratio_gout = ratio_gout
+        self.global_in_num = in_cg
+
+        module = nn.Identity if in_cl == 0 or out_cl == 0 else nn.Conv2d
+        self.convl2l = module(in_cl, out_cl, kernel_size,
+                              stride, padding, dilation, groups, bias, padding_mode=padding_type)
+        module = nn.Identity if in_cl == 0 or out_cg == 0 else nn.Conv2d
+        self.convl2g = module(in_cl, out_cg, kernel_size,
+                              stride, padding, dilation, groups, bias, padding_mode=padding_type)
+        module = nn.Identity if in_cg == 0 or out_cl == 0 else nn.Conv2d
+        self.convg2l = module(in_cg, out_cl, kernel_size,
+                              stride, padding, dilation, groups, bias, padding_mode=padding_type)
+        module = nn.Identity if in_cg == 0 or out_cg == 0 else SpectralTransform
+        self.convg2g = module(
+            in_cg, out_cg, stride, 1 if groups == 1 else groups // 2, enable_lfu, **spectral_kwargs)
+
+        self.gated = gated
+        module = nn.Identity if in_cg == 0 or out_cl == 0 or not self.gated else nn.Conv2d
+        self.gate = module(in_channels, 2, 1)
+
+    def forward(self, x, fname=None):
+        x_l, x_g = x if type(x) is tuple else (x, 0)
+        out_xl, out_xg = 0, 0
+
+        if self.gated:
+            total_input_parts = [x_l]
+            if torch.is_tensor(x_g):
+                total_input_parts.append(x_g)
+            total_input = torch.cat(total_input_parts, dim=1)
+
+            gates = torch.sigmoid(self.gate(total_input))
+            g2l_gate, l2g_gate = gates.chunk(2, dim=1)
+        else:
+            g2l_gate, l2g_gate = 1, 1
+
+        spec_x = self.convg2g(x_g)
+
+        if self.ratio_gout != 1:
+            out_xl = self.convl2l(x_l) + self.convg2l(x_g) * g2l_gate
+        if self.ratio_gout != 0:
+            out_xg = self.convl2g(x_l) * l2g_gate + spec_x
+
+        return out_xl, out_xg
+
+
+class FFC_BN_ACT(nn.Module):
+
+    def __init__(self, in_channels, out_channels,
+                 kernel_size, ratio_gin, ratio_gout,
+                 stride=1, padding=0, dilation=1, groups=1, bias=False,
+                 norm_layer=nn.SyncBatchNorm, activation_layer=nn.Identity,
+                 padding_type='reflect',
+                 enable_lfu=True, **kwargs):
+        super(FFC_BN_ACT, self).__init__()
+        self.ffc = FFC(in_channels, out_channels, kernel_size,
+                       ratio_gin, ratio_gout, stride, padding, dilation,
+                       groups, bias, enable_lfu, padding_type=padding_type, **kwargs)
+        lnorm = nn.Identity if ratio_gout == 1 else norm_layer
+        gnorm = nn.Identity if ratio_gout == 0 else norm_layer
+        global_channels = int(out_channels * ratio_gout)
+        # self.bn_l = lnorm(out_channels - global_channels)
+        # self.bn_g = gnorm(global_channels)
+
+        lact = nn.Identity if ratio_gout == 1 else activation_layer
+        gact = nn.Identity if ratio_gout == 0 else activation_layer
+        self.act_l = lact(inplace=True)
+        self.act_g = gact(inplace=True)
+
+    def forward(self, x, fname=None):
+        x_l, x_g = self.ffc(x, fname=fname, )
+        x_l = self.act_l(x_l)
+        x_g = self.act_g(x_g)
+        return x_l, x_g
+
+
+class FFCResnetBlock(nn.Module):
+    def __init__(self, dim, padding_type, norm_layer, activation_layer=nn.ReLU, dilation=1,
+                 spatial_transform_kwargs=None, inline=False, ratio_gin=0.75, ratio_gout=0.75):
+        super().__init__()
+        self.conv1 = FFC_BN_ACT(dim, dim, kernel_size=3, padding=dilation, dilation=dilation,
+                                norm_layer=norm_layer,
+                                activation_layer=activation_layer,
+                                padding_type=padding_type,
+                                ratio_gin=ratio_gin, ratio_gout=ratio_gout)
+        self.conv2 = FFC_BN_ACT(dim, dim, kernel_size=3, padding=dilation, dilation=dilation,
+                                norm_layer=norm_layer,
+                                activation_layer=activation_layer,
+                                padding_type=padding_type,
+                                ratio_gin=ratio_gin, ratio_gout=ratio_gout)
+        self.inline = inline
+
+    def forward(self, x, fname=None):
+        if self.inline:
+            x_l, x_g = x[:, :-self.conv1.ffc.global_in_num], x[:, -self.conv1.ffc.global_in_num:]
+        else:
+            x_l, x_g = x if type(x) is tuple else (x, 0)
+
+        id_l, id_g = x_l, x_g
+
+        x_l, x_g = self.conv1((x_l, x_g), fname=fname)
+        x_l, x_g = self.conv2((x_l, x_g), fname=fname)
+
+        x_l, x_g = id_l + x_l, id_g + x_g
+        out = x_l, x_g
+        if self.inline:
+            out = torch.cat(out, dim=1)
+        return out
+
+
+class ConcatTupleLayer(nn.Module):
+    def forward(self, x):
+        assert isinstance(x, tuple)
+        x_l, x_g = x
+        assert torch.is_tensor(x_l) or torch.is_tensor(x_g)
+        if not torch.is_tensor(x_g):
+            return x_l
+        return torch.cat(x, dim=1)
+
+
+class FFCBlock(torch.nn.Module):
+    def __init__(self,
+                 dim,  # Number of output/input channels.
+                 kernel_size,  # Width and height of the convolution kernel.
+                 padding,
+                 ratio_gin=0.75,
+                 ratio_gout=0.75,
+                 activation='linear',  # Activation function: 'relu', 'lrelu', etc.
+                 ):
+        super().__init__()
+        if activation == 'linear':
+            self.activation = nn.Identity
+        else:
+            self.activation = nn.ReLU
+        self.padding = padding
+        self.kernel_size = kernel_size
+        self.ffc_block = FFCResnetBlock(dim=dim,
+                                        padding_type='reflect',
+                                        norm_layer=nn.SyncBatchNorm,
+                                        activation_layer=self.activation,
+                                        dilation=1,
+                                        ratio_gin=ratio_gin,
+                                        ratio_gout=ratio_gout)
+
+        self.concat_layer = ConcatTupleLayer()
+
+    def forward(self, gen_ft, mask, fname=None):
+        x = gen_ft.float()
+
+        x_l, x_g = x[:, :-self.ffc_block.conv1.ffc.global_in_num], x[:, -self.ffc_block.conv1.ffc.global_in_num:]
+        id_l, id_g = x_l, x_g
+
+        x_l, x_g = self.ffc_block((x_l, x_g), fname=fname)
+        x_l, x_g = id_l + x_l, id_g + x_g
+        x = self.concat_layer((x_l, x_g))
+
+        return x + gen_ft.float()
+
+
+class FFCSkipLayer(torch.nn.Module):
+    def __init__(self,
+                 dim,  # Number of input/output channels.
+                 kernel_size=3,  # Convolution kernel size.
+                 ratio_gin=0.75,
+                 ratio_gout=0.75,
+                 ):
+        super().__init__()
+        self.padding = kernel_size // 2
+
+        self.ffc_act = FFCBlock(dim=dim, kernel_size=kernel_size, activation=nn.ReLU,
+                                padding=self.padding, ratio_gin=ratio_gin, ratio_gout=ratio_gout)
+
+    def forward(self, gen_ft, mask, fname=None):
+        x = self.ffc_act(gen_ft, mask, fname=fname)
+        return x
+
+
+class SynthesisBlock(torch.nn.Module):
+    def __init__(self,
+                 in_channels,  # Number of input channels, 0 = first block.
+                 out_channels,  # Number of output channels.
+                 w_dim,  # Intermediate latent (W) dimensionality.
+                 resolution,  # Resolution of this block.
+                 img_channels,  # Number of output color channels.
+                 is_last,  # Is this the last block?
+                 architecture='skip',  # Architecture: 'orig', 'skip', 'resnet'.
+                 resample_filter=[1, 3, 3, 1],  # Low-pass filter to apply when resampling activations.
+                 conv_clamp=None,  # Clamp the output of convolution layers to +-X, None = disable clamping.
+                 use_fp16=False,  # Use FP16 for this block?
+                 fp16_channels_last=False,  # Use channels-last memory format with FP16?
+                 **layer_kwargs,  # Arguments for SynthesisLayer.
+                 ):
+        assert architecture in ['orig', 'skip', 'resnet']
+        super().__init__()
+        self.in_channels = in_channels
+        self.w_dim = w_dim
+        self.resolution = resolution
+        self.img_channels = img_channels
+        self.is_last = is_last
+        self.architecture = architecture
+        self.use_fp16 = use_fp16
+        self.channels_last = (use_fp16 and fp16_channels_last)
+        self.register_buffer('resample_filter', setup_filter(resample_filter))
+        self.num_conv = 0
+        self.num_torgb = 0
+        self.res_ffc = {4: 0, 8: 0, 16: 0, 32: 1, 64: 1, 128: 1, 256: 1, 512: 1}
+
+        if in_channels != 0 and resolution >= 8:
+            self.ffc_skip = nn.ModuleList()
+            for _ in range(self.res_ffc[resolution]):
+                self.ffc_skip.append(FFCSkipLayer(dim=out_channels))
+
+        if in_channels == 0:
+            self.const = torch.nn.Parameter(torch.randn([out_channels, resolution, resolution]))
+
+        if in_channels != 0:
+            self.conv0 = SynthesisLayer(in_channels, out_channels, w_dim=w_dim * 3, resolution=resolution, up=2,
+                                        resample_filter=resample_filter, conv_clamp=conv_clamp,
+                                        channels_last=self.channels_last, **layer_kwargs)
+            self.num_conv += 1
+
+        self.conv1 = SynthesisLayer(out_channels, out_channels, w_dim=w_dim * 3, resolution=resolution,
+                                    conv_clamp=conv_clamp, channels_last=self.channels_last, **layer_kwargs)
+        self.num_conv += 1
+
+        if is_last or architecture == 'skip':
+            self.torgb = ToRGBLayer(out_channels, img_channels, w_dim=w_dim * 3,
+                                    conv_clamp=conv_clamp, channels_last=self.channels_last)
+            self.num_torgb += 1
+
+        if in_channels != 0 and architecture == 'resnet':
+            self.skip = Conv2dLayer(in_channels, out_channels, kernel_size=1, bias=False, up=2,
+                                    resample_filter=resample_filter, channels_last=self.channels_last)
+
+    def forward(self, x, mask, feats, img, ws, fname=None, force_fp32=False, fused_modconv=None, **layer_kwargs):
+        dtype = torch.float16 if self.use_fp16 and not force_fp32 else torch.float32
+        dtype = torch.float32
+        memory_format = torch.channels_last if self.channels_last and not force_fp32 else torch.contiguous_format
+        if fused_modconv is None:
+            fused_modconv = (not self.training) and (dtype == torch.float32 or int(x.shape[0]) == 1)
+
+        x = x.to(dtype=dtype, memory_format=memory_format)
+        x_skip = feats[self.resolution].clone().to(dtype=dtype, memory_format=memory_format)
+
+        # Main layers.
+        if self.in_channels == 0:
+            x = self.conv1(x, ws[1], fused_modconv=fused_modconv, **layer_kwargs)
+        elif self.architecture == 'resnet':
+            y = self.skip(x, gain=np.sqrt(0.5))
+            x = self.conv0(x, ws[0].clone(), fused_modconv=fused_modconv, **layer_kwargs)
+            if len(self.ffc_skip) > 0:
+                mask = F.interpolate(mask, size=x_skip.shape[2:], )
+                z = x + x_skip
+                for fres in self.ffc_skip:
+                    z = fres(z, mask)
+                x = x + z
+            else:
+                x = x + x_skip
+            x = self.conv1(x, ws[1].clone(), fused_modconv=fused_modconv, gain=np.sqrt(0.5), **layer_kwargs)
+            x = y.add_(x)
+        else:
+            x = self.conv0(x, ws[0].clone(), fused_modconv=fused_modconv, **layer_kwargs)
+            if len(self.ffc_skip) > 0:
+                mask = F.interpolate(mask, size=x_skip.shape[2:], )
+                z = x + x_skip
+                for fres in self.ffc_skip:
+                    z = fres(z, mask)
+                x = x + z
+            else:
+                x = x + x_skip
+            x = self.conv1(x, ws[1].clone(), fused_modconv=fused_modconv, **layer_kwargs)
+        # ToRGB.
+        if img is not None:
+            img = upsample2d(img, self.resample_filter)
+        if self.is_last or self.architecture == 'skip':
+            y = self.torgb(x, ws[2].clone(), fused_modconv=fused_modconv)
+            y = y.to(dtype=torch.float32, memory_format=torch.contiguous_format)
+            img = img.add_(y) if img is not None else y
+
+        x = x.to(dtype=dtype)
+        assert x.dtype == dtype
+        assert img is None or img.dtype == torch.float32
+        return x, img
+
+
+class SynthesisNetwork(torch.nn.Module):
+    def __init__(self,
+                 w_dim,  # Intermediate latent (W) dimensionality.
+                 z_dim,  # Output Latent (Z) dimensionality.
+                 img_resolution,  # Output image resolution.
+                 img_channels,  # Number of color channels.
+                 channel_base=16384,  # Overall multiplier for the number of channels.
+                 channel_max=512,  # Maximum number of channels in any layer.
+                 num_fp16_res=0,  # Use FP16 for the N highest resolutions.
+                 **block_kwargs,  # Arguments for SynthesisBlock.
+                 ):
+        assert img_resolution >= 4 and img_resolution & (img_resolution - 1) == 0
+        super().__init__()
+        self.w_dim = w_dim
+        self.img_resolution = img_resolution
+        self.img_resolution_log2 = int(np.log2(img_resolution))
+        self.img_channels = img_channels
+        self.block_resolutions = [2 ** i for i in range(3, self.img_resolution_log2 + 1)]
+        channels_dict = {res: min(channel_base // res, channel_max) for res in self.block_resolutions}
+        fp16_resolution = max(2 ** (self.img_resolution_log2 + 1 - num_fp16_res), 8)
+
+        self.foreword = SynthesisForeword(img_channels=img_channels, in_channels=min(channel_base // 4, channel_max),
+                                          z_dim=z_dim * 2, resolution=4)
+
+        self.num_ws = self.img_resolution_log2 * 2 - 2
+        for res in self.block_resolutions:
+            if res // 2 in channels_dict.keys():
+                in_channels = channels_dict[res // 2] if res > 4 else 0
+            else:
+                in_channels = min(channel_base // (res // 2), channel_max)
+            out_channels = channels_dict[res]
+            use_fp16 = (res >= fp16_resolution)
+            use_fp16 = False
+            is_last = (res == self.img_resolution)
+            block = SynthesisBlock(in_channels, out_channels, w_dim=w_dim, resolution=res,
+                                   img_channels=img_channels, is_last=is_last, use_fp16=use_fp16, **block_kwargs)
+            setattr(self, f'b{res}', block)
+
+    def forward(self, x_global, mask, feats, ws, fname=None, **block_kwargs):
+
+        img = None
+
+        x, img = self.foreword(x_global, ws, feats, img)
+
+        for res in self.block_resolutions:
+            block = getattr(self, f'b{res}')
+            mod_vector0 = []
+            mod_vector0.append(ws[:, int(np.log2(res)) * 2 - 5])
+            mod_vector0.append(x_global.clone())
+            mod_vector0 = torch.cat(mod_vector0, dim=1)
+
+            mod_vector1 = []
+            mod_vector1.append(ws[:, int(np.log2(res)) * 2 - 4])
+            mod_vector1.append(x_global.clone())
+            mod_vector1 = torch.cat(mod_vector1, dim=1)
+
+            mod_vector_rgb = []
+            mod_vector_rgb.append(ws[:, int(np.log2(res)) * 2 - 3])
+            mod_vector_rgb.append(x_global.clone())
+            mod_vector_rgb = torch.cat(mod_vector_rgb, dim=1)
+            x, img = block(x, mask, feats, img, (mod_vector0, mod_vector1, mod_vector_rgb), fname=fname, **block_kwargs)
+        return img
+
+
+class MappingNetwork(torch.nn.Module):
+    def __init__(self,
+                 z_dim,  # Input latent (Z) dimensionality, 0 = no latent.
+                 c_dim,  # Conditioning label (C) dimensionality, 0 = no label.
+                 w_dim,  # Intermediate latent (W) dimensionality.
+                 num_ws,  # Number of intermediate latents to output, None = do not broadcast.
+                 num_layers=8,  # Number of mapping layers.
+                 embed_features=None,  # Label embedding dimensionality, None = same as w_dim.
+                 layer_features=None,  # Number of intermediate features in the mapping layers, None = same as w_dim.
+                 activation='lrelu',  # Activation function: 'relu', 'lrelu', etc.
+                 lr_multiplier=0.01,  # Learning rate multiplier for the mapping layers.
+                 w_avg_beta=0.995,  # Decay for tracking the moving average of W during training, None = do not track.
+                 ):
+        super().__init__()
+        self.z_dim = z_dim
+        self.c_dim = c_dim
+        self.w_dim = w_dim
+        self.num_ws = num_ws
+        self.num_layers = num_layers
+        self.w_avg_beta = w_avg_beta
+
+        if embed_features is None:
+            embed_features = w_dim
+        if c_dim == 0:
+            embed_features = 0
+        if layer_features is None:
+            layer_features = w_dim
+        features_list = [z_dim + embed_features] + [layer_features] * (num_layers - 1) + [w_dim]
+
+        if c_dim > 0:
+            self.embed = FullyConnectedLayer(c_dim, embed_features)
+        for idx in range(num_layers):
+            in_features = features_list[idx]
+            out_features = features_list[idx + 1]
+            layer = FullyConnectedLayer(in_features, out_features, activation=activation, lr_multiplier=lr_multiplier)
+            setattr(self, f'fc{idx}', layer)
+
+        if num_ws is not None and w_avg_beta is not None:
+            self.register_buffer('w_avg', torch.zeros([w_dim]))
+
+    def forward(self, z, c, truncation_psi=1, truncation_cutoff=None, skip_w_avg_update=False):
+        # Embed, normalize, and concat inputs.
+        x = None
+        with torch.autograd.profiler.record_function('input'):
+            if self.z_dim > 0:
+                x = normalize_2nd_moment(z.to(torch.float32))
+            if self.c_dim > 0:
+                y = normalize_2nd_moment(self.embed(c.to(torch.float32)))
+                x = torch.cat([x, y], dim=1) if x is not None else y
+
+        # Main layers.
+        for idx in range(self.num_layers):
+            layer = getattr(self, f'fc{idx}')
+            x = layer(x)
+
+        # Update moving average of W.
+        if self.w_avg_beta is not None and self.training and not skip_w_avg_update:
+            with torch.autograd.profiler.record_function('update_w_avg'):
+                self.w_avg.copy_(x.detach().mean(dim=0).lerp(self.w_avg, self.w_avg_beta))
+
+        # Broadcast.
+        if self.num_ws is not None:
+            with torch.autograd.profiler.record_function('broadcast'):
+                x = x.unsqueeze(1).repeat([1, self.num_ws, 1])
+
+        # Apply truncation.
+        if truncation_psi != 1:
+            with torch.autograd.profiler.record_function('truncate'):
+                assert self.w_avg_beta is not None
+                if self.num_ws is None or truncation_cutoff is None:
+                    x = self.w_avg.lerp(x, truncation_psi)
+                else:
+                    x[:, :truncation_cutoff] = self.w_avg.lerp(x[:, :truncation_cutoff], truncation_psi)
+        return x
+
+
+class Generator(torch.nn.Module):
+    def __init__(self,
+                 z_dim,  # Input latent (Z) dimensionality.
+                 c_dim,  # Conditioning label (C) dimensionality.
+                 w_dim,  # Intermediate latent (W) dimensionality.
+                 img_resolution,  # Output resolution.
+                 img_channels,  # Number of output color channels.
+                 encoder_kwargs={},  # Arguments for EncoderNetwork.
+                 mapping_kwargs={},  # Arguments for MappingNetwork.
+                 synthesis_kwargs={},  # Arguments for SynthesisNetwork.
+                 ):
+        super().__init__()
+        self.z_dim = z_dim
+        self.c_dim = c_dim
+        self.w_dim = w_dim
+        self.img_resolution = img_resolution
+        self.img_channels = img_channels
+        self.encoder = EncoderNetwork(c_dim=c_dim, z_dim=z_dim, img_resolution=img_resolution,
+                                      img_channels=img_channels, **encoder_kwargs)
+        self.synthesis = SynthesisNetwork(z_dim=z_dim, w_dim=w_dim, img_resolution=img_resolution,
+                                          img_channels=img_channels, **synthesis_kwargs)
+        self.num_ws = self.synthesis.num_ws
+        self.mapping = MappingNetwork(z_dim=z_dim, c_dim=c_dim, w_dim=w_dim, num_ws=self.num_ws, **mapping_kwargs)
+
+    def forward(self, img, c, fname=None, truncation_psi=1, truncation_cutoff=None, **synthesis_kwargs):
+        mask = img[:, -1].unsqueeze(1)
+        x_global, z, feats = self.encoder(img, c)
+        ws = self.mapping(z, c, truncation_psi=truncation_psi, truncation_cutoff=truncation_cutoff)
+        img = self.synthesis(x_global, mask, feats, ws, fname=fname, **synthesis_kwargs)
+        return img
+
+
+FCF_MODEL_URL = os.environ.get(
+    "FCF_MODEL_URL",
+    "https://github.com/Sanster/models/releases/download/add_fcf/places_512_G.pth",
+)
+
+
+class FcF(InpaintModel):
+    min_size = 512
+    pad_mod = 512
+    pad_to_square = True
+
+    def init_model(self, device, **kwargs):
+        seed = 0
+        random.seed(seed)
+        np.random.seed(seed)
+        torch.manual_seed(seed)
+        torch.cuda.manual_seed_all(seed)
+        torch.backends.cudnn.deterministic = True
+        torch.backends.cudnn.benchmark = False
+
+        kwargs = {'channel_base': 1 * 32768, 'channel_max': 512, 'num_fp16_res': 4, 'conv_clamp': 256}
+        G = Generator(z_dim=512, c_dim=0, w_dim=512, img_resolution=512, img_channels=3,
+                      synthesis_kwargs=kwargs, encoder_kwargs=kwargs, mapping_kwargs={'num_layers': 2})
+        self.model = load_model(G, FCF_MODEL_URL, device)
+        self.label = torch.zeros([1, self.model.c_dim], device=device)
+
+    @staticmethod
+    def is_downloaded() -> bool:
+        return os.path.exists(get_cache_path_by_url(FCF_MODEL_URL))
+
+    @torch.no_grad()
+    def __call__(self, image, mask, config: Config):
+        """
+        images: [H, W, C] RGB, not normalized
+        masks: [H, W]
+        return: BGR IMAGE
+        """
+        if image.shape[0] == 512 and image.shape[1] == 512:
+            return self._pad_forward(image, mask, config)
+
+        boxes = boxes_from_mask(mask)
+        crop_result = []
+        config.hd_strategy_crop_margin = 128
+        for box in boxes:
+            crop_image, crop_mask, crop_box = self._crop_box(image, mask, box, config)
+            origin_size = crop_image.shape[:2]
+            resize_image = resize_max_size(crop_image, size_limit=512)
+            resize_mask = resize_max_size(crop_mask, size_limit=512)
+            inpaint_result = self._pad_forward(resize_image, resize_mask, config)
+
+            # only paste masked area result
+            inpaint_result = cv2.resize(inpaint_result, (origin_size[1], origin_size[0]), interpolation=cv2.INTER_CUBIC)
+
+            original_pixel_indices = crop_mask < 127
+            inpaint_result[original_pixel_indices] = crop_image[:, :, ::-1][original_pixel_indices]
+
+            crop_result.append((inpaint_result, crop_box))
+
+        inpaint_result = image[:, :, ::-1]
+        for crop_image, crop_box in crop_result:
+            x1, y1, x2, y2 = crop_box
+            inpaint_result[y1:y2, x1:x2, :] = crop_image
+
+        return inpaint_result
+
+    def forward(self, image, mask, config: Config):
+        """Input images and output images have same size
+        images: [H, W, C] RGB
+        masks: [H, W] mask area == 255
+        return: BGR IMAGE
+        """
+
+        image = norm_img(image)  # [0, 1]
+        image = image * 2 - 1  # [0, 1] -> [-1, 1]
+        mask = (mask > 120) * 255
+        mask = norm_img(mask)
+
+        image = torch.from_numpy(image).unsqueeze(0).to(self.device)
+        mask = torch.from_numpy(mask).unsqueeze(0).to(self.device)
+
+        erased_img = image * (1 - mask)
+        input_image = torch.cat([0.5 - mask, erased_img], dim=1)
+
+        output = self.model(input_image, self.label, truncation_psi=0.1, noise_mode='none')
+        output = (output.permute(0, 2, 3, 1) * 127.5 + 127.5).round().clamp(0, 255).to(torch.uint8)
+        output = output[0].cpu().numpy()
+        cur_res = cv2.cvtColor(output, cv2.COLOR_RGB2BGR)
+        return cur_res

lama_cleaner/model/lama.py

@@ -0,0 +1,61 @@
+import os
+
+import cv2
+import numpy as np
+import torch
+from loguru import logger
+
+from lama_cleaner.helper import pad_img_to_modulo, download_model, norm_img, get_cache_path_by_url
+from lama_cleaner.model.base import InpaintModel
+from lama_cleaner.schema import Config
+
+LAMA_MODEL_URL = os.environ.get(
+    "LAMA_MODEL_URL",
+    "https://github.com/Sanster/models/releases/download/add_big_lama/big-lama.pt",
+)
+
+#"https://drive.google.com/file/d/1bMD06F9hkkS1oi8cEmb4cSjXz54Pxs6A/view?usp=sharing"   #big-lama.pt file
+
+
+class LaMa(InpaintModel):
+    pad_mod = 8
+
+    def init_model(self, device, **kwargs):
+        if os.environ.get("LAMA_MODEL"):
+            model_path = os.environ.get("LAMA_MODEL")
+            if not os.path.exists(model_path):
+                raise FileNotFoundError(
+                    f"lama torchscript model not found: {model_path}"
+                )
+        else:
+            model_path = download_model(LAMA_MODEL_URL)
+        logger.info(f"Load LaMa model from: {model_path}")
+        model = torch.jit.load(model_path, map_location="cpu")
+        model = model.to(device)
+        model.eval()
+        self.model = model
+        self.model_path = model_path
+
+    @staticmethod
+    def is_downloaded() -> bool:
+        return os.path.exists(get_cache_path_by_url(LAMA_MODEL_URL))
+
+    def forward(self, image, mask, config: Config):
+        """Input image and output image have same size
+        image: [H, W, C] RGB
+        mask: [H, W]
+        return: BGR IMAGE
+        """
+        image = norm_img(image)
+        mask = norm_img(mask)
+
+        mask = (mask > 0) * 1
+        image = torch.from_numpy(image).unsqueeze(0).to(self.device)
+        mask = torch.from_numpy(mask).unsqueeze(0).to(self.device)
+
+        inpainted_image = self.model(image, mask)
+
+        cur_res = inpainted_image[0].permute(1, 2, 0).detach().cpu().numpy()
+        cur_res = np.clip(cur_res * 255, 0, 255).astype("uint8")
+        cur_res = cv2.cvtColor(cur_res, cv2.COLOR_RGB2BGR)
+        return cur_res

lama_cleaner/model/ldm.py

@@ -0,0 +1,312 @@
+import os
+
+import numpy as np
+import torch
+from loguru import logger
+
+from lama_cleaner.model.base import InpaintModel
+from lama_cleaner.model.ddim_sampler import DDIMSampler
+from lama_cleaner.model.plms_sampler import PLMSSampler
+from lama_cleaner.schema import Config, LDMSampler
+
+torch.manual_seed(42)
+import torch.nn as nn
+from lama_cleaner.helper import (
+    download_model,
+    norm_img,
+    get_cache_path_by_url,
+    load_jit_model,
+)
+from lama_cleaner.model.utils import (
+    make_beta_schedule,
+    timestep_embedding,
+)
+
+LDM_ENCODE_MODEL_URL = os.environ.get(
+    "LDM_ENCODE_MODEL_URL",
+    "https://github.com/Sanster/models/releases/download/add_ldm/cond_stage_model_encode.pt",
+)
+
+LDM_DECODE_MODEL_URL = os.environ.get(
+    "LDM_DECODE_MODEL_URL",
+    "https://github.com/Sanster/models/releases/download/add_ldm/cond_stage_model_decode.pt",
+)
+
+LDM_DIFFUSION_MODEL_URL = os.environ.get(
+    "LDM_DIFFUSION_MODEL_URL",
+    "https://github.com/Sanster/models/releases/download/add_ldm/diffusion.pt",
+)
+
+
+class DDPM(nn.Module):
+    # classic DDPM with Gaussian diffusion, in image space
+    def __init__(
+        self,
+        device,
+        timesteps=1000,
+        beta_schedule="linear",
+        linear_start=0.0015,
+        linear_end=0.0205,
+        cosine_s=0.008,
+        original_elbo_weight=0.0,
+        v_posterior=0.0,  # weight for choosing posterior variance as sigma = (1-v) * beta_tilde + v * beta
+        l_simple_weight=1.0,
+        parameterization="eps",  # all assuming fixed variance schedules
+        use_positional_encodings=False,
+    ):
+        super().__init__()
+        self.device = device
+        self.parameterization = parameterization
+        self.use_positional_encodings = use_positional_encodings
+
+        self.v_posterior = v_posterior
+        self.original_elbo_weight = original_elbo_weight
+        self.l_simple_weight = l_simple_weight
+
+        self.register_schedule(
+            beta_schedule=beta_schedule,
+            timesteps=timesteps,
+            linear_start=linear_start,
+            linear_end=linear_end,
+            cosine_s=cosine_s,
+        )
+
+    def register_schedule(
+        self,
+        given_betas=None,
+        beta_schedule="linear",
+        timesteps=1000,
+        linear_start=1e-4,
+        linear_end=2e-2,
+        cosine_s=8e-3,
+    ):
+        betas = make_beta_schedule(
+            self.device,
+            beta_schedule,
+            timesteps,
+            linear_start=linear_start,
+            linear_end=linear_end,
+            cosine_s=cosine_s,
+        )
+        alphas = 1.0 - betas
+        alphas_cumprod = np.cumprod(alphas, axis=0)
+        alphas_cumprod_prev = np.append(1.0, alphas_cumprod[:-1])
+
+        (timesteps,) = betas.shape
+        self.num_timesteps = int(timesteps)
+        self.linear_start = linear_start
+        self.linear_end = linear_end
+        assert (
+            alphas_cumprod.shape[0] == self.num_timesteps
+        ), "alphas have to be defined for each timestep"
+
+        to_torch = lambda x: torch.tensor(x, dtype=torch.float32).to(self.device)
+
+        self.register_buffer("betas", to_torch(betas))
+        self.register_buffer("alphas_cumprod", to_torch(alphas_cumprod))
+        self.register_buffer("alphas_cumprod_prev", to_torch(alphas_cumprod_prev))
+
+        # calculations for diffusion q(x_t | x_{t-1}) and others
+        self.register_buffer("sqrt_alphas_cumprod", to_torch(np.sqrt(alphas_cumprod)))
+        self.register_buffer(
+            "sqrt_one_minus_alphas_cumprod", to_torch(np.sqrt(1.0 - alphas_cumprod))
+        )
+        self.register_buffer(
+            "log_one_minus_alphas_cumprod", to_torch(np.log(1.0 - alphas_cumprod))
+        )
+        self.register_buffer(
+            "sqrt_recip_alphas_cumprod", to_torch(np.sqrt(1.0 / alphas_cumprod))
+        )
+        self.register_buffer(
+            "sqrt_recipm1_alphas_cumprod", to_torch(np.sqrt(1.0 / alphas_cumprod - 1))
+        )
+
+        # calculations for posterior q(x_{t-1} | x_t, x_0)
+        posterior_variance = (1 - self.v_posterior) * betas * (
+            1.0 - alphas_cumprod_prev
+        ) / (1.0 - alphas_cumprod) + self.v_posterior * betas
+        # above: equal to 1. / (1. / (1. - alpha_cumprod_tm1) + alpha_t / beta_t)
+        self.register_buffer("posterior_variance", to_torch(posterior_variance))
+        # below: log calculation clipped because the posterior variance is 0 at the beginning of the diffusion chain
+        self.register_buffer(
+            "posterior_log_variance_clipped",
+            to_torch(np.log(np.maximum(posterior_variance, 1e-20))),
+        )
+        self.register_buffer(
+            "posterior_mean_coef1",
+            to_torch(betas * np.sqrt(alphas_cumprod_prev) / (1.0 - alphas_cumprod)),
+        )
+        self.register_buffer(
+            "posterior_mean_coef2",
+            to_torch(
+                (1.0 - alphas_cumprod_prev) * np.sqrt(alphas) / (1.0 - alphas_cumprod)
+            ),
+        )
+
+        if self.parameterization == "eps":
+            lvlb_weights = self.betas**2 / (
+                2
+                * self.posterior_variance
+                * to_torch(alphas)
+                * (1 - self.alphas_cumprod)
+            )
+        elif self.parameterization == "x0":
+            lvlb_weights = (
+                0.5
+                * np.sqrt(torch.Tensor(alphas_cumprod))
+                / (2.0 * 1 - torch.Tensor(alphas_cumprod))
+            )
+        else:
+            raise NotImplementedError("mu not supported")
+        # TODO how to choose this term
+        lvlb_weights[0] = lvlb_weights[1]
+        self.register_buffer("lvlb_weights", lvlb_weights, persistent=False)
+        assert not torch.isnan(self.lvlb_weights).all()
+
+
+class LatentDiffusion(DDPM):
+    def __init__(
+        self,
+        diffusion_model,
+        device,
+        cond_stage_key="image",
+        cond_stage_trainable=False,
+        concat_mode=True,
+        scale_factor=1.0,
+        scale_by_std=False,
+        *args,
+        **kwargs,
+    ):
+        self.num_timesteps_cond = 1
+        self.scale_by_std = scale_by_std
+        super().__init__(device, *args, **kwargs)
+        self.diffusion_model = diffusion_model
+        self.concat_mode = concat_mode
+        self.cond_stage_trainable = cond_stage_trainable
+        self.cond_stage_key = cond_stage_key
+        self.num_downs = 2
+        self.scale_factor = scale_factor
+
+    def make_cond_schedule(
+        self,
+    ):
+        self.cond_ids = torch.full(
+            size=(self.num_timesteps,),
+            fill_value=self.num_timesteps - 1,
+            dtype=torch.long,
+        )
+        ids = torch.round(
+            torch.linspace(0, self.num_timesteps - 1, self.num_timesteps_cond)
+        ).long()
+        self.cond_ids[: self.num_timesteps_cond] = ids
+
+    def register_schedule(
+        self,
+        given_betas=None,
+        beta_schedule="linear",
+        timesteps=1000,
+        linear_start=1e-4,
+        linear_end=2e-2,
+        cosine_s=8e-3,
+    ):
+        super().register_schedule(
+            given_betas, beta_schedule, timesteps, linear_start, linear_end, cosine_s
+        )
+
+        self.shorten_cond_schedule = self.num_timesteps_cond > 1
+        if self.shorten_cond_schedule:
+            self.make_cond_schedule()
+
+    def apply_model(self, x_noisy, t, cond):
+        # x_recon = self.model(x_noisy, t, cond['c_concat'][0])  # cond['c_concat'][0].shape 1,4,128,128
+        t_emb = timestep_embedding(x_noisy.device, t, 256, repeat_only=False)
+        x_recon = self.diffusion_model(x_noisy, t_emb, cond)
+        return x_recon
+
+
+class LDM(InpaintModel):
+    pad_mod = 32
+
+    def __init__(self, device, fp16: bool = True, **kwargs):
+        self.fp16 = fp16
+        super().__init__(device)
+        self.device = device
+
+    def init_model(self, device, **kwargs):
+        self.diffusion_model = load_jit_model(LDM_DIFFUSION_MODEL_URL, device)
+        self.cond_stage_model_decode = load_jit_model(LDM_DECODE_MODEL_URL, device)
+        self.cond_stage_model_encode = load_jit_model(LDM_ENCODE_MODEL_URL, device)
+        if self.fp16 and "cuda" in str(device):
+            self.diffusion_model = self.diffusion_model.half()
+            self.cond_stage_model_decode = self.cond_stage_model_decode.half()
+            self.cond_stage_model_encode = self.cond_stage_model_encode.half()
+
+        self.model = LatentDiffusion(self.diffusion_model, device)
+
+    @staticmethod
+    def is_downloaded() -> bool:
+        model_paths = [
+            get_cache_path_by_url(LDM_DIFFUSION_MODEL_URL),
+            get_cache_path_by_url(LDM_DECODE_MODEL_URL),
+            get_cache_path_by_url(LDM_ENCODE_MODEL_URL),
+        ]
+        return all([os.path.exists(it) for it in model_paths])
+
+    @torch.cuda.amp.autocast()
+    def forward(self, image, mask, config: Config):
+        """
+        image: [H, W, C] RGB
+        mask: [H, W, 1]
+        return: BGR IMAGE
+        """
+        # image [1,3,512,512] float32
+        # mask: [1,1,512,512] float32
+        # masked_image: [1,3,512,512] float32
+        if config.ldm_sampler == LDMSampler.ddim:
+            sampler = DDIMSampler(self.model)
+        elif config.ldm_sampler == LDMSampler.plms:
+            sampler = PLMSSampler(self.model)
+        else:
+            raise ValueError()
+
+        steps = config.ldm_steps
+        image = norm_img(image)
+        mask = norm_img(mask)
+
+        mask[mask < 0.5] = 0
+        mask[mask >= 0.5] = 1
+
+        image = torch.from_numpy(image).unsqueeze(0).to(self.device)
+        mask = torch.from_numpy(mask).unsqueeze(0).to(self.device)
+        masked_image = (1 - mask) * image
+
+        mask = self._norm(mask)
+        masked_image = self._norm(masked_image)
+
+        c = self.cond_stage_model_encode(masked_image)
+        torch.cuda.empty_cache()
+
+        cc = torch.nn.functional.interpolate(mask, size=c.shape[-2:])  # 1,1,128,128
+        c = torch.cat((c, cc), dim=1)  # 1,4,128,128
+
+        shape = (c.shape[1] - 1,) + c.shape[2:]
+        samples_ddim = sampler.sample(
+            steps=steps, conditioning=c, batch_size=c.shape[0], shape=shape
+        )
+        torch.cuda.empty_cache()
+        x_samples_ddim = self.cond_stage_model_decode(
+            samples_ddim
+        )  # samples_ddim: 1, 3, 128, 128 float32
+        torch.cuda.empty_cache()
+
+        # image = torch.clamp((image + 1.0) / 2.0, min=0.0, max=1.0)
+        # mask = torch.clamp((mask + 1.0) / 2.0, min=0.0, max=1.0)
+        inpainted_image = torch.clamp((x_samples_ddim + 1.0) / 2.0, min=0.0, max=1.0)
+
+        # inpainted = (1 - mask) * image + mask * predicted_image
+        inpainted_image = inpainted_image.cpu().numpy().transpose(0, 2, 3, 1)[0] * 255
+        inpainted_image = inpainted_image.astype(np.uint8)[:, :, ::-1]
+        return inpainted_image
+
+    def _norm(self, tensor):
+        return tensor * 2.0 - 1.0

lama_cleaner/model/mat.py

@@ -0,0 +1,1444 @@
+import os
+import random
+
+import cv2
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.utils.checkpoint as checkpoint
+
+from lama_cleaner.helper import load_model, get_cache_path_by_url, norm_img
+from lama_cleaner.model.base import InpaintModel
+from lama_cleaner.model.utils import setup_filter, Conv2dLayer, FullyConnectedLayer, conv2d_resample, bias_act, \
+    upsample2d, activation_funcs, MinibatchStdLayer, to_2tuple, normalize_2nd_moment
+from lama_cleaner.schema import Config
+
+
+class ModulatedConv2d(nn.Module):
+    def __init__(self,
+                 in_channels,  # Number of input channels.
+                 out_channels,  # Number of output channels.
+                 kernel_size,  # Width and height of the convolution kernel.
+                 style_dim,  # dimension of the style code
+                 demodulate=True,  # perfrom demodulation
+                 up=1,  # Integer upsampling factor.
+                 down=1,  # Integer downsampling factor.
+                 resample_filter=[1, 3, 3, 1],  # Low-pass filter to apply when resampling activations.
+                 conv_clamp=None,  # Clamp the output to +-X, None = disable clamping.
+                 ):
+        super().__init__()
+        self.demodulate = demodulate
+
+        self.weight = torch.nn.Parameter(torch.randn([1, out_channels, in_channels, kernel_size, kernel_size]))
+        self.out_channels = out_channels
+        self.kernel_size = kernel_size
+        self.weight_gain = 1 / np.sqrt(in_channels * (kernel_size ** 2))
+        self.padding = self.kernel_size // 2
+        self.up = up
+        self.down = down
+        self.register_buffer('resample_filter', setup_filter(resample_filter))
+        self.conv_clamp = conv_clamp
+
+        self.affine = FullyConnectedLayer(style_dim, in_channels, bias_init=1)
+
+    def forward(self, x, style):
+        batch, in_channels, height, width = x.shape
+        style = self.affine(style).view(batch, 1, in_channels, 1, 1)
+        weight = self.weight * self.weight_gain * style
+
+        if self.demodulate:
+            decoefs = (weight.pow(2).sum(dim=[2, 3, 4]) + 1e-8).rsqrt()
+            weight = weight * decoefs.view(batch, self.out_channels, 1, 1, 1)
+
+        weight = weight.view(batch * self.out_channels, in_channels, self.kernel_size, self.kernel_size)
+        x = x.view(1, batch * in_channels, height, width)
+        x = conv2d_resample(x=x, w=weight, f=self.resample_filter, up=self.up, down=self.down,
+                            padding=self.padding, groups=batch)
+        out = x.view(batch, self.out_channels, *x.shape[2:])
+
+        return out
+
+
+class StyleConv(torch.nn.Module):
+    def __init__(self,
+                 in_channels,  # Number of input channels.
+                 out_channels,  # Number of output channels.
+                 style_dim,  # Intermediate latent (W) dimensionality.
+                 resolution,  # Resolution of this layer.
+                 kernel_size=3,  # Convolution kernel size.
+                 up=1,  # Integer upsampling factor.
+                 use_noise=False,  # Enable noise input?
+                 activation='lrelu',  # Activation function: 'relu', 'lrelu', etc.
+                 resample_filter=[1, 3, 3, 1],  # Low-pass filter to apply when resampling activations.
+                 conv_clamp=None,  # Clamp the output of convolution layers to +-X, None = disable clamping.
+                 demodulate=True,  # perform demodulation
+                 ):
+        super().__init__()
+
+        self.conv = ModulatedConv2d(in_channels=in_channels,
+                                    out_channels=out_channels,
+                                    kernel_size=kernel_size,
+                                    style_dim=style_dim,
+                                    demodulate=demodulate,
+                                    up=up,
+                                    resample_filter=resample_filter,
+                                    conv_clamp=conv_clamp)
+
+        self.use_noise = use_noise
+        self.resolution = resolution
+        if use_noise:
+            self.register_buffer('noise_const', torch.randn([resolution, resolution]))
+            self.noise_strength = torch.nn.Parameter(torch.zeros([]))
+
+        self.bias = torch.nn.Parameter(torch.zeros([out_channels]))
+        self.activation = activation
+        self.act_gain = activation_funcs[activation].def_gain
+        self.conv_clamp = conv_clamp
+
+    def forward(self, x, style, noise_mode='random', gain=1):
+        x = self.conv(x, style)
+
+        assert noise_mode in ['random', 'const', 'none']
+
+        if self.use_noise:
+            if noise_mode == 'random':
+                xh, xw = x.size()[-2:]
+                noise = torch.randn([x.shape[0], 1, xh, xw], device=x.device) \
+                        * self.noise_strength
+            if noise_mode == 'const':
+                noise = self.noise_const * self.noise_strength
+            x = x + noise
+
+        act_gain = self.act_gain * gain
+        act_clamp = self.conv_clamp * gain if self.conv_clamp is not None else None
+        out = bias_act(x, self.bias, act=self.activation, gain=act_gain, clamp=act_clamp)
+
+        return out
+
+
+class ToRGB(torch.nn.Module):
+    def __init__(self,
+                 in_channels,
+                 out_channels,
+                 style_dim,
+                 kernel_size=1,
+                 resample_filter=[1, 3, 3, 1],
+                 conv_clamp=None,
+                 demodulate=False):
+        super().__init__()
+
+        self.conv = ModulatedConv2d(in_channels=in_channels,
+                                    out_channels=out_channels,
+                                    kernel_size=kernel_size,
+                                    style_dim=style_dim,
+                                    demodulate=demodulate,
+                                    resample_filter=resample_filter,
+                                    conv_clamp=conv_clamp)
+        self.bias = torch.nn.Parameter(torch.zeros([out_channels]))
+        self.register_buffer('resample_filter', setup_filter(resample_filter))
+        self.conv_clamp = conv_clamp
+
+    def forward(self, x, style, skip=None):
+        x = self.conv(x, style)
+        out = bias_act(x, self.bias, clamp=self.conv_clamp)
+
+        if skip is not None:
+            if skip.shape != out.shape:
+                skip = upsample2d(skip, self.resample_filter)
+            out = out + skip
+
+        return out
+
+
+def get_style_code(a, b):
+    return torch.cat([a, b], dim=1)
+
+
+class DecBlockFirst(nn.Module):
+    def __init__(self, in_channels, out_channels, activation, style_dim, use_noise, demodulate, img_channels):
+        super().__init__()
+        self.fc = FullyConnectedLayer(in_features=in_channels * 2,
+                                      out_features=in_channels * 4 ** 2,
+                                      activation=activation)
+        self.conv = StyleConv(in_channels=in_channels,
+                              out_channels=out_channels,
+                              style_dim=style_dim,
+                              resolution=4,
+                              kernel_size=3,
+                              use_noise=use_noise,
+                              activation=activation,
+                              demodulate=demodulate,
+                              )
+        self.toRGB = ToRGB(in_channels=out_channels,
+                           out_channels=img_channels,
+                           style_dim=style_dim,
+                           kernel_size=1,
+                           demodulate=False,
+                           )
+
+    def forward(self, x, ws, gs, E_features, noise_mode='random'):
+        x = self.fc(x).view(x.shape[0], -1, 4, 4)
+        x = x + E_features[2]
+        style = get_style_code(ws[:, 0], gs)
+        x = self.conv(x, style, noise_mode=noise_mode)
+        style = get_style_code(ws[:, 1], gs)
+        img = self.toRGB(x, style, skip=None)
+
+        return x, img
+
+
+class DecBlockFirstV2(nn.Module):
+    def __init__(self, in_channels, out_channels, activation, style_dim, use_noise, demodulate, img_channels):
+        super().__init__()
+        self.conv0 = Conv2dLayer(in_channels=in_channels,
+                                 out_channels=in_channels,
+                                 kernel_size=3,
+                                 activation=activation,
+                                 )
+        self.conv1 = StyleConv(in_channels=in_channels,
+                               out_channels=out_channels,
+                               style_dim=style_dim,
+                               resolution=4,
+                               kernel_size=3,
+                               use_noise=use_noise,
+                               activation=activation,
+                               demodulate=demodulate,
+                               )
+        self.toRGB = ToRGB(in_channels=out_channels,
+                           out_channels=img_channels,
+                           style_dim=style_dim,
+                           kernel_size=1,
+                           demodulate=False,
+                           )
+
+    def forward(self, x, ws, gs, E_features, noise_mode='random'):
+        # x = self.fc(x).view(x.shape[0], -1, 4, 4)
+        x = self.conv0(x)
+        x = x + E_features[2]
+        style = get_style_code(ws[:, 0], gs)
+        x = self.conv1(x, style, noise_mode=noise_mode)
+        style = get_style_code(ws[:, 1], gs)
+        img = self.toRGB(x, style, skip=None)
+
+        return x, img
+
+
+class DecBlock(nn.Module):
+    def __init__(self, res, in_channels, out_channels, activation, style_dim, use_noise, demodulate,
+                 img_channels):  # res = 2, ..., resolution_log2
+        super().__init__()
+        self.res = res
+
+        self.conv0 = StyleConv(in_channels=in_channels,
+                               out_channels=out_channels,
+                               style_dim=style_dim,
+                               resolution=2 ** res,
+                               kernel_size=3,
+                               up=2,
+                               use_noise=use_noise,
+                               activation=activation,
+                               demodulate=demodulate,
+                               )
+        self.conv1 = StyleConv(in_channels=out_channels,
+                               out_channels=out_channels,
+                               style_dim=style_dim,
+                               resolution=2 ** res,
+                               kernel_size=3,
+                               use_noise=use_noise,
+                               activation=activation,
+                               demodulate=demodulate,
+                               )
+        self.toRGB = ToRGB(in_channels=out_channels,
+                           out_channels=img_channels,
+                           style_dim=style_dim,
+                           kernel_size=1,
+                           demodulate=False,
+                           )
+
+    def forward(self, x, img, ws, gs, E_features, noise_mode='random'):
+        style = get_style_code(ws[:, self.res * 2 - 5], gs)
+        x = self.conv0(x, style, noise_mode=noise_mode)
+        x = x + E_features[self.res]
+        style = get_style_code(ws[:, self.res * 2 - 4], gs)
+        x = self.conv1(x, style, noise_mode=noise_mode)
+        style = get_style_code(ws[:, self.res * 2 - 3], gs)
+        img = self.toRGB(x, style, skip=img)
+
+        return x, img
+
+
+class MappingNet(torch.nn.Module):
+    def __init__(self,
+                 z_dim,  # Input latent (Z) dimensionality, 0 = no latent.
+                 c_dim,  # Conditioning label (C) dimensionality, 0 = no label.
+                 w_dim,  # Intermediate latent (W) dimensionality.
+                 num_ws,  # Number of intermediate latents to output, None = do not broadcast.
+                 num_layers=8,  # Number of mapping layers.
+                 embed_features=None,  # Label embedding dimensionality, None = same as w_dim.
+                 layer_features=None,  # Number of intermediate features in the mapping layers, None = same as w_dim.
+                 activation='lrelu',  # Activation function: 'relu', 'lrelu', etc.
+                 lr_multiplier=0.01,  # Learning rate multiplier for the mapping layers.
+                 w_avg_beta=0.995,  # Decay for tracking the moving average of W during training, None = do not track.
+                 ):
+        super().__init__()
+        self.z_dim = z_dim
+        self.c_dim = c_dim
+        self.w_dim = w_dim
+        self.num_ws = num_ws
+        self.num_layers = num_layers
+        self.w_avg_beta = w_avg_beta
+
+        if embed_features is None:
+            embed_features = w_dim
+        if c_dim == 0:
+            embed_features = 0
+        if layer_features is None:
+            layer_features = w_dim
+        features_list = [z_dim + embed_features] + [layer_features] * (num_layers - 1) + [w_dim]
+
+        if c_dim > 0:
+            self.embed = FullyConnectedLayer(c_dim, embed_features)
+        for idx in range(num_layers):
+            in_features = features_list[idx]
+            out_features = features_list[idx + 1]
+            layer = FullyConnectedLayer(in_features, out_features, activation=activation, lr_multiplier=lr_multiplier)
+            setattr(self, f'fc{idx}', layer)
+
+        if num_ws is not None and w_avg_beta is not None:
+            self.register_buffer('w_avg', torch.zeros([w_dim]))
+
+    def forward(self, z, c, truncation_psi=1, truncation_cutoff=None, skip_w_avg_update=False):
+        # Embed, normalize, and concat inputs.
+        x = None
+        with torch.autograd.profiler.record_function('input'):
+            if self.z_dim > 0:
+                x = normalize_2nd_moment(z.to(torch.float32))
+            if self.c_dim > 0:
+                y = normalize_2nd_moment(self.embed(c.to(torch.float32)))
+                x = torch.cat([x, y], dim=1) if x is not None else y
+
+        # Main layers.
+        for idx in range(self.num_layers):
+            layer = getattr(self, f'fc{idx}')
+            x = layer(x)
+
+        # Update moving average of W.
+        if self.w_avg_beta is not None and self.training and not skip_w_avg_update:
+            with torch.autograd.profiler.record_function('update_w_avg'):
+                self.w_avg.copy_(x.detach().mean(dim=0).lerp(self.w_avg, self.w_avg_beta))
+
+        # Broadcast.
+        if self.num_ws is not None:
+            with torch.autograd.profiler.record_function('broadcast'):
+                x = x.unsqueeze(1).repeat([1, self.num_ws, 1])
+
+        # Apply truncation.
+        if truncation_psi != 1:
+            with torch.autograd.profiler.record_function('truncate'):
+                assert self.w_avg_beta is not None
+                if self.num_ws is None or truncation_cutoff is None:
+                    x = self.w_avg.lerp(x, truncation_psi)
+                else:
+                    x[:, :truncation_cutoff] = self.w_avg.lerp(x[:, :truncation_cutoff], truncation_psi)
+
+        return x
+
+
+class DisFromRGB(nn.Module):
+    def __init__(self, in_channels, out_channels, activation):  # res = 2, ..., resolution_log2
+        super().__init__()
+        self.conv = Conv2dLayer(in_channels=in_channels,
+                                out_channels=out_channels,
+                                kernel_size=1,
+                                activation=activation,
+                                )
+
+    def forward(self, x):
+        return self.conv(x)
+
+
+class DisBlock(nn.Module):
+    def __init__(self, in_channels, out_channels, activation):  # res = 2, ..., resolution_log2
+        super().__init__()
+        self.conv0 = Conv2dLayer(in_channels=in_channels,
+                                 out_channels=in_channels,
+                                 kernel_size=3,
+                                 activation=activation,
+                                 )
+        self.conv1 = Conv2dLayer(in_channels=in_channels,
+                                 out_channels=out_channels,
+                                 kernel_size=3,
+                                 down=2,
+                                 activation=activation,
+                                 )
+        self.skip = Conv2dLayer(in_channels=in_channels,
+                                out_channels=out_channels,
+                                kernel_size=1,
+                                down=2,
+                                bias=False,
+                                )
+
+    def forward(self, x):
+        skip = self.skip(x, gain=np.sqrt(0.5))
+        x = self.conv0(x)
+        x = self.conv1(x, gain=np.sqrt(0.5))
+        out = skip + x
+
+        return out
+
+
+class Discriminator(torch.nn.Module):
+    def __init__(self,
+                 c_dim,  # Conditioning label (C) dimensionality.
+                 img_resolution,  # Input resolution.
+                 img_channels,  # Number of input color channels.
+                 channel_base=32768,  # Overall multiplier for the number of channels.
+                 channel_max=512,  # Maximum number of channels in any layer.
+                 channel_decay=1,
+                 cmap_dim=None,  # Dimensionality of mapped conditioning label, None = default.
+                 activation='lrelu',
+                 mbstd_group_size=4,  # Group size for the minibatch standard deviation layer, None = entire minibatch.
+                 mbstd_num_channels=1,  # Number of features for the minibatch standard deviation layer, 0 = disable.
+                 ):
+        super().__init__()
+        self.c_dim = c_dim
+        self.img_resolution = img_resolution
+        self.img_channels = img_channels
+
+        resolution_log2 = int(np.log2(img_resolution))
+        assert img_resolution == 2 ** resolution_log2 and img_resolution >= 4
+        self.resolution_log2 = resolution_log2
+
+        def nf(stage):
+            return np.clip(int(channel_base / 2 ** (stage * channel_decay)), 1, channel_max)
+
+        if cmap_dim == None:
+            cmap_dim = nf(2)
+        if c_dim == 0:
+            cmap_dim = 0
+        self.cmap_dim = cmap_dim
+
+        if c_dim > 0:
+            self.mapping = MappingNet(z_dim=0, c_dim=c_dim, w_dim=cmap_dim, num_ws=None, w_avg_beta=None)
+
+        Dis = [DisFromRGB(img_channels + 1, nf(resolution_log2), activation)]
+        for res in range(resolution_log2, 2, -1):
+            Dis.append(DisBlock(nf(res), nf(res - 1), activation))
+
+        if mbstd_num_channels > 0:
+            Dis.append(MinibatchStdLayer(group_size=mbstd_group_size, num_channels=mbstd_num_channels))
+        Dis.append(Conv2dLayer(nf(2) + mbstd_num_channels, nf(2), kernel_size=3, activation=activation))
+        self.Dis = nn.Sequential(*Dis)
+
+        self.fc0 = FullyConnectedLayer(nf(2) * 4 ** 2, nf(2), activation=activation)
+        self.fc1 = FullyConnectedLayer(nf(2), 1 if cmap_dim == 0 else cmap_dim)
+
+    def forward(self, images_in, masks_in, c):
+        x = torch.cat([masks_in - 0.5, images_in], dim=1)
+        x = self.Dis(x)
+        x = self.fc1(self.fc0(x.flatten(start_dim=1)))
+
+        if self.c_dim > 0:
+            cmap = self.mapping(None, c)
+
+        if self.cmap_dim > 0:
+            x = (x * cmap).sum(dim=1, keepdim=True) * (1 / np.sqrt(self.cmap_dim))
+
+        return x
+
+
+def nf(stage, channel_base=32768, channel_decay=1.0, channel_max=512):
+    NF = {512: 64, 256: 128, 128: 256, 64: 512, 32: 512, 16: 512, 8: 512, 4: 512}
+    return NF[2 ** stage]
+
+
+class Mlp(nn.Module):
+    def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.GELU, drop=0.):
+        super().__init__()
+        out_features = out_features or in_features
+        hidden_features = hidden_features or in_features
+        self.fc1 = FullyConnectedLayer(in_features=in_features, out_features=hidden_features, activation='lrelu')
+        self.fc2 = FullyConnectedLayer(in_features=hidden_features, out_features=out_features)
+
+    def forward(self, x):
+        x = self.fc1(x)
+        x = self.fc2(x)
+        return x
+
+
+def window_partition(x, window_size):
+    """
+    Args:
+        x: (B, H, W, C)
+        window_size (int): window size
+    Returns:
+        windows: (num_windows*B, window_size, window_size, C)
+    """
+    B, H, W, C = x.shape
+    x = x.view(B, H // window_size, window_size, W // window_size, window_size, C)
+    windows = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(-1, window_size, window_size, C)
+    return windows
+
+
+def window_reverse(windows, window_size: int, H: int, W: int):
+    """
+    Args:
+        windows: (num_windows*B, window_size, window_size, C)
+        window_size (int): Window size
+        H (int): Height of image
+        W (int): Width of image
+    Returns:
+        x: (B, H, W, C)
+    """
+    B = int(windows.shape[0] / (H * W / window_size / window_size))
+    # B = windows.shape[0] / (H * W / window_size / window_size)
+    x = windows.view(B, H // window_size, W // window_size, window_size, window_size, -1)
+    x = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(B, H, W, -1)
+    return x
+
+
+class Conv2dLayerPartial(nn.Module):
+    def __init__(self,
+                 in_channels,  # Number of input channels.
+                 out_channels,  # Number of output channels.
+                 kernel_size,  # Width and height of the convolution kernel.
+                 bias=True,  # Apply additive bias before the activation function?
+                 activation='linear',  # Activation function: 'relu', 'lrelu', etc.
+                 up=1,  # Integer upsampling factor.
+                 down=1,  # Integer downsampling factor.
+                 resample_filter=[1, 3, 3, 1],  # Low-pass filter to apply when resampling activations.
+                 conv_clamp=None,  # Clamp the output to +-X, None = disable clamping.
+                 trainable=True,  # Update the weights of this layer during training?
+                 ):
+        super().__init__()
+        self.conv = Conv2dLayer(in_channels, out_channels, kernel_size, bias, activation, up, down, resample_filter,
+                                conv_clamp, trainable)
+
+        self.weight_maskUpdater = torch.ones(1, 1, kernel_size, kernel_size)
+        self.slide_winsize = kernel_size ** 2
+        self.stride = down
+        self.padding = kernel_size // 2 if kernel_size % 2 == 1 else 0
+
+    def forward(self, x, mask=None):
+        if mask is not None:
+            with torch.no_grad():
+                if self.weight_maskUpdater.type() != x.type():
+                    self.weight_maskUpdater = self.weight_maskUpdater.to(x)
+                update_mask = F.conv2d(mask, self.weight_maskUpdater, bias=None, stride=self.stride,
+                                       padding=self.padding)
+                mask_ratio = self.slide_winsize / (update_mask + 1e-8)
+                update_mask = torch.clamp(update_mask, 0, 1)  # 0 or 1
+                mask_ratio = torch.mul(mask_ratio, update_mask)
+            x = self.conv(x)
+            x = torch.mul(x, mask_ratio)
+            return x, update_mask
+        else:
+            x = self.conv(x)
+            return x, None
+
+
+class WindowAttention(nn.Module):
+    r""" Window based multi-head self attention (W-MSA) module with relative position bias.
+    It supports both of shifted and non-shifted window.
+    Args:
+        dim (int): Number of input channels.
+        window_size (tuple[int]): The height and width of the window.
+        num_heads (int): Number of attention heads.
+        qkv_bias (bool, optional):  If True, add a learnable bias to query, key, value. Default: True
+        qk_scale (float | None, optional): Override default qk scale of head_dim ** -0.5 if set
+        attn_drop (float, optional): Dropout ratio of attention weight. Default: 0.0
+        proj_drop (float, optional): Dropout ratio of output. Default: 0.0
+    """
+
+    def __init__(self, dim, window_size, num_heads, down_ratio=1, qkv_bias=True, qk_scale=None, attn_drop=0.,
+                 proj_drop=0.):
+
+        super().__init__()
+        self.dim = dim
+        self.window_size = window_size  # Wh, Ww
+        self.num_heads = num_heads
+        head_dim = dim // num_heads
+        self.scale = qk_scale or head_dim ** -0.5
+
+        self.q = FullyConnectedLayer(in_features=dim, out_features=dim)
+        self.k = FullyConnectedLayer(in_features=dim, out_features=dim)
+        self.v = FullyConnectedLayer(in_features=dim, out_features=dim)
+        self.proj = FullyConnectedLayer(in_features=dim, out_features=dim)
+
+        self.softmax = nn.Softmax(dim=-1)
+
+    def forward(self, x, mask_windows=None, mask=None):
+        """
+        Args:
+            x: input features with shape of (num_windows*B, N, C)
+            mask: (0/-inf) mask with shape of (num_windows, Wh*Ww, Wh*Ww) or None
+        """
+        B_, N, C = x.shape
+        norm_x = F.normalize(x, p=2.0, dim=-1)
+        q = self.q(norm_x).reshape(B_, N, self.num_heads, C // self.num_heads).permute(0, 2, 1, 3)
+        k = self.k(norm_x).view(B_, -1, self.num_heads, C // self.num_heads).permute(0, 2, 3, 1)
+        v = self.v(x).view(B_, -1, self.num_heads, C // self.num_heads).permute(0, 2, 1, 3)
+
+        attn = (q @ k) * self.scale
+
+        if mask is not None:
+            nW = mask.shape[0]
+            attn = attn.view(B_ // nW, nW, self.num_heads, N, N) + mask.unsqueeze(1).unsqueeze(0)
+            attn = attn.view(-1, self.num_heads, N, N)
+
+        if mask_windows is not None:
+            attn_mask_windows = mask_windows.squeeze(-1).unsqueeze(1).unsqueeze(1)
+            attn = attn + attn_mask_windows.masked_fill(attn_mask_windows == 0, float(-100.0)).masked_fill(
+                attn_mask_windows == 1, float(0.0))
+            with torch.no_grad():
+                mask_windows = torch.clamp(torch.sum(mask_windows, dim=1, keepdim=True), 0, 1).repeat(1, N, 1)
+
+        attn = self.softmax(attn)
+
+        x = (attn @ v).transpose(1, 2).reshape(B_, N, C)
+        x = self.proj(x)
+        return x, mask_windows
+
+
+class SwinTransformerBlock(nn.Module):
+    r""" Swin Transformer Block.
+    Args:
+        dim (int): Number of input channels.
+        input_resolution (tuple[int]): Input resulotion.
+        num_heads (int): Number of attention heads.
+        window_size (int): Window size.
+        shift_size (int): Shift size for SW-MSA.
+        mlp_ratio (float): Ratio of mlp hidden dim to embedding dim.
+        qkv_bias (bool, optional): If True, add a learnable bias to query, key, value. Default: True
+        qk_scale (float | None, optional): Override default qk scale of head_dim ** -0.5 if set.
+        drop (float, optional): Dropout rate. Default: 0.0
+        attn_drop (float, optional): Attention dropout rate. Default: 0.0
+        drop_path (float, optional): Stochastic depth rate. Default: 0.0
+        act_layer (nn.Module, optional): Activation layer. Default: nn.GELU
+        norm_layer (nn.Module, optional): Normalization layer.  Default: nn.LayerNorm
+    """
+
+    def __init__(self, dim, input_resolution, num_heads, down_ratio=1, window_size=7, shift_size=0,
+                 mlp_ratio=4., qkv_bias=True, qk_scale=None, drop=0., attn_drop=0., drop_path=0.,
+                 act_layer=nn.GELU, norm_layer=nn.LayerNorm):
+        super().__init__()
+        self.dim = dim
+        self.input_resolution = input_resolution
+        self.num_heads = num_heads
+        self.window_size = window_size
+        self.shift_size = shift_size
+        self.mlp_ratio = mlp_ratio
+        if min(self.input_resolution) <= self.window_size:
+            # if window size is larger than input resolution, we don't partition windows
+            self.shift_size = 0
+            self.window_size = min(self.input_resolution)
+        assert 0 <= self.shift_size < self.window_size, "shift_size must in 0-window_size"
+
+        if self.shift_size > 0:
+            down_ratio = 1
+        self.attn = WindowAttention(dim, window_size=to_2tuple(self.window_size), num_heads=num_heads,
+                                    down_ratio=down_ratio, qkv_bias=qkv_bias, qk_scale=qk_scale, attn_drop=attn_drop,
+                                    proj_drop=drop)
+
+        self.fuse = FullyConnectedLayer(in_features=dim * 2, out_features=dim, activation='lrelu')
+
+        mlp_hidden_dim = int(dim * mlp_ratio)
+        self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop)
+
+        if self.shift_size > 0:
+            attn_mask = self.calculate_mask(self.input_resolution)
+        else:
+            attn_mask = None
+
+        self.register_buffer("attn_mask", attn_mask)
+
+    def calculate_mask(self, x_size):
+        # calculate attention mask for SW-MSA
+        H, W = x_size
+        img_mask = torch.zeros((1, H, W, 1))  # 1 H W 1
+        h_slices = (slice(0, -self.window_size),
+                    slice(-self.window_size, -self.shift_size),
+                    slice(-self.shift_size, None))
+        w_slices = (slice(0, -self.window_size),
+                    slice(-self.window_size, -self.shift_size),
+                    slice(-self.shift_size, None))
+        cnt = 0
+        for h in h_slices:
+            for w in w_slices:
+                img_mask[:, h, w, :] = cnt
+                cnt += 1
+
+        mask_windows = window_partition(img_mask, self.window_size)  # nW, window_size, window_size, 1
+        mask_windows = mask_windows.view(-1, self.window_size * self.window_size)
+        attn_mask = mask_windows.unsqueeze(1) - mask_windows.unsqueeze(2)
+        attn_mask = attn_mask.masked_fill(attn_mask != 0, float(-100.0)).masked_fill(attn_mask == 0, float(0.0))
+
+        return attn_mask
+
+    def forward(self, x, x_size, mask=None):
+        # H, W = self.input_resolution
+        H, W = x_size
+        B, L, C = x.shape
+        # assert L == H * W, "input feature has wrong size"
+
+        shortcut = x
+        x = x.view(B, H, W, C)
+        if mask is not None:
+            mask = mask.view(B, H, W, 1)
+
+        # cyclic shift
+        if self.shift_size > 0:
+            shifted_x = torch.roll(x, shifts=(-self.shift_size, -self.shift_size), dims=(1, 2))
+            if mask is not None:
+                shifted_mask = torch.roll(mask, shifts=(-self.shift_size, -self.shift_size), dims=(1, 2))
+        else:
+            shifted_x = x
+            if mask is not None:
+                shifted_mask = mask
+
+        # partition windows
+        x_windows = window_partition(shifted_x, self.window_size)  # nW*B, window_size, window_size, C
+        x_windows = x_windows.view(-1, self.window_size * self.window_size, C)  # nW*B, window_size*window_size, C
+        if mask is not None:
+            mask_windows = window_partition(shifted_mask, self.window_size)
+            mask_windows = mask_windows.view(-1, self.window_size * self.window_size, 1)
+        else:
+            mask_windows = None
+
+        # W-MSA/SW-MSA (to be compatible for testing on images whose shapes are the multiple of window size
+        if self.input_resolution == x_size:
+            attn_windows, mask_windows = self.attn(x_windows, mask_windows,
+                                                   mask=self.attn_mask)  # nW*B, window_size*window_size, C
+        else:
+            attn_windows, mask_windows = self.attn(x_windows, mask_windows, mask=self.calculate_mask(x_size).to(
+                x.device))  # nW*B, window_size*window_size, C
+
+        # merge windows
+        attn_windows = attn_windows.view(-1, self.window_size, self.window_size, C)
+        shifted_x = window_reverse(attn_windows, self.window_size, H, W)  # B H' W' C
+        if mask is not None:
+            mask_windows = mask_windows.view(-1, self.window_size, self.window_size, 1)
+            shifted_mask = window_reverse(mask_windows, self.window_size, H, W)
+
+        # reverse cyclic shift
+        if self.shift_size > 0:
+            x = torch.roll(shifted_x, shifts=(self.shift_size, self.shift_size), dims=(1, 2))
+            if mask is not None:
+                mask = torch.roll(shifted_mask, shifts=(self.shift_size, self.shift_size), dims=(1, 2))
+        else:
+            x = shifted_x
+            if mask is not None:
+                mask = shifted_mask
+        x = x.view(B, H * W, C)
+        if mask is not None:
+            mask = mask.view(B, H * W, 1)
+
+        # FFN
+        x = self.fuse(torch.cat([shortcut, x], dim=-1))
+        x = self.mlp(x)
+
+        return x, mask
+
+
+class PatchMerging(nn.Module):
+    def __init__(self, in_channels, out_channels, down=2):
+        super().__init__()
+        self.conv = Conv2dLayerPartial(in_channels=in_channels,
+                                       out_channels=out_channels,
+                                       kernel_size=3,
+                                       activation='lrelu',
+                                       down=down,
+                                       )
+        self.down = down
+
+    def forward(self, x, x_size, mask=None):
+        x = token2feature(x, x_size)
+        if mask is not None:
+            mask = token2feature(mask, x_size)
+        x, mask = self.conv(x, mask)
+        if self.down != 1:
+            ratio = 1 / self.down
+            x_size = (int(x_size[0] * ratio), int(x_size[1] * ratio))
+        x = feature2token(x)
+        if mask is not None:
+            mask = feature2token(mask)
+        return x, x_size, mask
+
+
+class PatchUpsampling(nn.Module):
+    def __init__(self, in_channels, out_channels, up=2):
+        super().__init__()
+        self.conv = Conv2dLayerPartial(in_channels=in_channels,
+                                       out_channels=out_channels,
+                                       kernel_size=3,
+                                       activation='lrelu',
+                                       up=up,
+                                       )
+        self.up = up
+
+    def forward(self, x, x_size, mask=None):
+        x = token2feature(x, x_size)
+        if mask is not None:
+            mask = token2feature(mask, x_size)
+        x, mask = self.conv(x, mask)
+        if self.up != 1:
+            x_size = (int(x_size[0] * self.up), int(x_size[1] * self.up))
+        x = feature2token(x)
+        if mask is not None:
+            mask = feature2token(mask)
+        return x, x_size, mask
+
+
+class BasicLayer(nn.Module):
+    """ A basic Swin Transformer layer for one stage.
+    Args:
+        dim (int): Number of input channels.
+        input_resolution (tuple[int]): Input resolution.
+        depth (int): Number of blocks.
+        num_heads (int): Number of attention heads.
+        window_size (int): Local window size.
+        mlp_ratio (float): Ratio of mlp hidden dim to embedding dim.
+        qkv_bias (bool, optional): If True, add a learnable bias to query, key, value. Default: True
+        qk_scale (float | None, optional): Override default qk scale of head_dim ** -0.5 if set.
+        drop (float, optional): Dropout rate. Default: 0.0
+        attn_drop (float, optional): Attention dropout rate. Default: 0.0
+        drop_path (float | tuple[float], optional): Stochastic depth rate. Default: 0.0
+        norm_layer (nn.Module, optional): Normalization layer. Default: nn.LayerNorm
+        downsample (nn.Module | None, optional): Downsample layer at the end of the layer. Default: None
+        use_checkpoint (bool): Whether to use checkpointing to save memory. Default: False.
+    """
+
+    def __init__(self, dim, input_resolution, depth, num_heads, window_size, down_ratio=1,
+                 mlp_ratio=2., qkv_bias=True, qk_scale=None, drop=0., attn_drop=0.,
+                 drop_path=0., norm_layer=nn.LayerNorm, downsample=None, use_checkpoint=False):
+
+        super().__init__()
+        self.dim = dim
+        self.input_resolution = input_resolution
+        self.depth = depth
+        self.use_checkpoint = use_checkpoint
+
+        # patch merging layer
+        if downsample is not None:
+            # self.downsample = downsample(input_resolution, dim=dim, norm_layer=norm_layer)
+            self.downsample = downsample
+        else:
+            self.downsample = None
+
+        # build blocks
+        self.blocks = nn.ModuleList([
+            SwinTransformerBlock(dim=dim, input_resolution=input_resolution,
+                                 num_heads=num_heads, down_ratio=down_ratio, window_size=window_size,
+                                 shift_size=0 if (i % 2 == 0) else window_size // 2,
+                                 mlp_ratio=mlp_ratio,
+                                 qkv_bias=qkv_bias, qk_scale=qk_scale,
+                                 drop=drop, attn_drop=attn_drop,
+                                 drop_path=drop_path[i] if isinstance(drop_path, list) else drop_path,
+                                 norm_layer=norm_layer)
+            for i in range(depth)])
+
+        self.conv = Conv2dLayerPartial(in_channels=dim, out_channels=dim, kernel_size=3, activation='lrelu')
+
+    def forward(self, x, x_size, mask=None):
+        if self.downsample is not None:
+            x, x_size, mask = self.downsample(x, x_size, mask)
+        identity = x
+        for blk in self.blocks:
+            if self.use_checkpoint:
+                x, mask = checkpoint.checkpoint(blk, x, x_size, mask)
+            else:
+                x, mask = blk(x, x_size, mask)
+        if mask is not None:
+            mask = token2feature(mask, x_size)
+        x, mask = self.conv(token2feature(x, x_size), mask)
+        x = feature2token(x) + identity
+        if mask is not None:
+            mask = feature2token(mask)
+        return x, x_size, mask
+
+
+class ToToken(nn.Module):
+    def __init__(self, in_channels=3, dim=128, kernel_size=5, stride=1):
+        super().__init__()
+
+        self.proj = Conv2dLayerPartial(in_channels=in_channels, out_channels=dim, kernel_size=kernel_size,
+                                       activation='lrelu')
+
+    def forward(self, x, mask):
+        x, mask = self.proj(x, mask)
+
+        return x, mask
+
+
+class EncFromRGB(nn.Module):
+    def __init__(self, in_channels, out_channels, activation):  # res = 2, ..., resolution_log2
+        super().__init__()
+        self.conv0 = Conv2dLayer(in_channels=in_channels,
+                                 out_channels=out_channels,
+                                 kernel_size=1,
+                                 activation=activation,
+                                 )
+        self.conv1 = Conv2dLayer(in_channels=out_channels,
+                                 out_channels=out_channels,
+                                 kernel_size=3,
+                                 activation=activation,
+                                 )
+
+    def forward(self, x):
+        x = self.conv0(x)
+        x = self.conv1(x)
+
+        return x
+
+
+class ConvBlockDown(nn.Module):
+    def __init__(self, in_channels, out_channels, activation):  # res = 2, ..., resolution_log
+        super().__init__()
+
+        self.conv0 = Conv2dLayer(in_channels=in_channels,
+                                 out_channels=out_channels,
+                                 kernel_size=3,
+                                 activation=activation,
+                                 down=2,
+                                 )
+        self.conv1 = Conv2dLayer(in_channels=out_channels,
+                                 out_channels=out_channels,
+                                 kernel_size=3,
+                                 activation=activation,
+                                 )
+
+    def forward(self, x):
+        x = self.conv0(x)
+        x = self.conv1(x)
+
+        return x
+
+
+def token2feature(x, x_size):
+    B, N, C = x.shape
+    h, w = x_size
+    x = x.permute(0, 2, 1).reshape(B, C, h, w)
+    return x
+
+
+def feature2token(x):
+    B, C, H, W = x.shape
+    x = x.view(B, C, -1).transpose(1, 2)
+    return x
+
+
+class Encoder(nn.Module):
+    def __init__(self, res_log2, img_channels, activation, patch_size=5, channels=16, drop_path_rate=0.1):
+        super().__init__()
+
+        self.resolution = []
+
+        for idx, i in enumerate(range(res_log2, 3, -1)):  # from input size to 16x16
+            res = 2 ** i
+            self.resolution.append(res)
+            if i == res_log2:
+                block = EncFromRGB(img_channels * 2 + 1, nf(i), activation)
+            else:
+                block = ConvBlockDown(nf(i + 1), nf(i), activation)
+            setattr(self, 'EncConv_Block_%dx%d' % (res, res), block)
+
+    def forward(self, x):
+        out = {}
+        for res in self.resolution:
+            res_log2 = int(np.log2(res))
+            x = getattr(self, 'EncConv_Block_%dx%d' % (res, res))(x)
+            out[res_log2] = x
+
+        return out
+
+
+class ToStyle(nn.Module):
+    def __init__(self, in_channels, out_channels, activation, drop_rate):
+        super().__init__()
+        self.conv = nn.Sequential(
+            Conv2dLayer(in_channels=in_channels, out_channels=in_channels, kernel_size=3, activation=activation,
+                        down=2),
+            Conv2dLayer(in_channels=in_channels, out_channels=in_channels, kernel_size=3, activation=activation,
+                        down=2),
+            Conv2dLayer(in_channels=in_channels, out_channels=in_channels, kernel_size=3, activation=activation,
+                        down=2),
+        )
+
+        self.pool = nn.AdaptiveAvgPool2d(1)
+        self.fc = FullyConnectedLayer(in_features=in_channels,
+                                      out_features=out_channels,
+                                      activation=activation)
+        # self.dropout = nn.Dropout(drop_rate)
+
+    def forward(self, x):
+        x = self.conv(x)
+        x = self.pool(x)
+        x = self.fc(x.flatten(start_dim=1))
+        # x = self.dropout(x)
+
+        return x
+
+
+class DecBlockFirstV2(nn.Module):
+    def __init__(self, res, in_channels, out_channels, activation, style_dim, use_noise, demodulate, img_channels):
+        super().__init__()
+        self.res = res
+
+        self.conv0 = Conv2dLayer(in_channels=in_channels,
+                                 out_channels=in_channels,
+                                 kernel_size=3,
+                                 activation=activation,
+                                 )
+        self.conv1 = StyleConv(in_channels=in_channels,
+                               out_channels=out_channels,
+                               style_dim=style_dim,
+                               resolution=2 ** res,
+                               kernel_size=3,
+                               use_noise=use_noise,
+                               activation=activation,
+                               demodulate=demodulate,
+                               )
+        self.toRGB = ToRGB(in_channels=out_channels,
+                           out_channels=img_channels,
+                           style_dim=style_dim,
+                           kernel_size=1,
+                           demodulate=False,
+                           )
+
+    def forward(self, x, ws, gs, E_features, noise_mode='random'):
+        # x = self.fc(x).view(x.shape[0], -1, 4, 4)
+        x = self.conv0(x)
+        x = x + E_features[self.res]
+        style = get_style_code(ws[:, 0], gs)
+        x = self.conv1(x, style, noise_mode=noise_mode)
+        style = get_style_code(ws[:, 1], gs)
+        img = self.toRGB(x, style, skip=None)
+
+        return x, img
+
+
+class DecBlock(nn.Module):
+    def __init__(self, res, in_channels, out_channels, activation, style_dim, use_noise, demodulate,
+                 img_channels):  # res = 4, ..., resolution_log2
+        super().__init__()
+        self.res = res
+
+        self.conv0 = StyleConv(in_channels=in_channels,
+                               out_channels=out_channels,
+                               style_dim=style_dim,
+                               resolution=2 ** res,
+                               kernel_size=3,
+                               up=2,
+                               use_noise=use_noise,
+                               activation=activation,
+                               demodulate=demodulate,
+                               )
+        self.conv1 = StyleConv(in_channels=out_channels,
+                               out_channels=out_channels,
+                               style_dim=style_dim,
+                               resolution=2 ** res,
+                               kernel_size=3,
+                               use_noise=use_noise,
+                               activation=activation,
+                               demodulate=demodulate,
+                               )
+        self.toRGB = ToRGB(in_channels=out_channels,
+                           out_channels=img_channels,
+                           style_dim=style_dim,
+                           kernel_size=1,
+                           demodulate=False,
+                           )
+
+    def forward(self, x, img, ws, gs, E_features, noise_mode='random'):
+        style = get_style_code(ws[:, self.res * 2 - 9], gs)
+        x = self.conv0(x, style, noise_mode=noise_mode)
+        x = x + E_features[self.res]
+        style = get_style_code(ws[:, self.res * 2 - 8], gs)
+        x = self.conv1(x, style, noise_mode=noise_mode)
+        style = get_style_code(ws[:, self.res * 2 - 7], gs)
+        img = self.toRGB(x, style, skip=img)
+
+        return x, img
+
+
+class Decoder(nn.Module):
+    def __init__(self, res_log2, activation, style_dim, use_noise, demodulate, img_channels):
+        super().__init__()
+        self.Dec_16x16 = DecBlockFirstV2(4, nf(4), nf(4), activation, style_dim, use_noise, demodulate, img_channels)
+        for res in range(5, res_log2 + 1):
+            setattr(self, 'Dec_%dx%d' % (2 ** res, 2 ** res),
+                    DecBlock(res, nf(res - 1), nf(res), activation, style_dim, use_noise, demodulate, img_channels))
+        self.res_log2 = res_log2
+
+    def forward(self, x, ws, gs, E_features, noise_mode='random'):
+        x, img = self.Dec_16x16(x, ws, gs, E_features, noise_mode=noise_mode)
+        for res in range(5, self.res_log2 + 1):
+            block = getattr(self, 'Dec_%dx%d' % (2 ** res, 2 ** res))
+            x, img = block(x, img, ws, gs, E_features, noise_mode=noise_mode)
+
+        return img
+
+
+class DecStyleBlock(nn.Module):
+    def __init__(self, res, in_channels, out_channels, activation, style_dim, use_noise, demodulate, img_channels):
+        super().__init__()
+        self.res = res
+
+        self.conv0 = StyleConv(in_channels=in_channels,
+                               out_channels=out_channels,
+                               style_dim=style_dim,
+                               resolution=2 ** res,
+                               kernel_size=3,
+                               up=2,
+                               use_noise=use_noise,
+                               activation=activation,
+                               demodulate=demodulate,
+                               )
+        self.conv1 = StyleConv(in_channels=out_channels,
+                               out_channels=out_channels,
+                               style_dim=style_dim,
+                               resolution=2 ** res,
+                               kernel_size=3,
+                               use_noise=use_noise,
+                               activation=activation,
+                               demodulate=demodulate,
+                               )
+        self.toRGB = ToRGB(in_channels=out_channels,
+                           out_channels=img_channels,
+                           style_dim=style_dim,
+                           kernel_size=1,
+                           demodulate=False,
+                           )
+
+    def forward(self, x, img, style, skip, noise_mode='random'):
+        x = self.conv0(x, style, noise_mode=noise_mode)
+        x = x + skip
+        x = self.conv1(x, style, noise_mode=noise_mode)
+        img = self.toRGB(x, style, skip=img)
+
+        return x, img
+
+
+class FirstStage(nn.Module):
+    def __init__(self, img_channels, img_resolution=256, dim=180, w_dim=512, use_noise=False, demodulate=True,
+                 activation='lrelu'):
+        super().__init__()
+        res = 64
+
+        self.conv_first = Conv2dLayerPartial(in_channels=img_channels + 1, out_channels=dim, kernel_size=3,
+                                             activation=activation)
+        self.enc_conv = nn.ModuleList()
+        down_time = int(np.log2(img_resolution // res))
+        # 根据图片尺寸构建 swim transformer 的层数
+        for i in range(down_time):  # from input size to 64
+            self.enc_conv.append(
+                Conv2dLayerPartial(in_channels=dim, out_channels=dim, kernel_size=3, down=2, activation=activation)
+            )
+
+        # from 64 -> 16 -> 64
+        depths = [2, 3, 4, 3, 2]
+        ratios = [1, 1 / 2, 1 / 2, 2, 2]
+        num_heads = 6
+        window_sizes = [8, 16, 16, 16, 8]
+        drop_path_rate = 0.1
+        dpr = [x.item() for x in torch.linspace(0, drop_path_rate, sum(depths))]
+
+        self.tran = nn.ModuleList()
+        for i, depth in enumerate(depths):
+            res = int(res * ratios[i])
+            if ratios[i] < 1:
+                merge = PatchMerging(dim, dim, down=int(1 / ratios[i]))
+            elif ratios[i] > 1:
+                merge = PatchUpsampling(dim, dim, up=ratios[i])
+            else:
+                merge = None
+            self.tran.append(
+                BasicLayer(dim=dim, input_resolution=[res, res], depth=depth, num_heads=num_heads,
+                           window_size=window_sizes[i], drop_path=dpr[sum(depths[:i]):sum(depths[:i + 1])],
+                           downsample=merge)
+            )
+
+        # global style
+        down_conv = []
+        for i in range(int(np.log2(16))):
+            down_conv.append(
+                Conv2dLayer(in_channels=dim, out_channels=dim, kernel_size=3, down=2, activation=activation))
+        down_conv.append(nn.AdaptiveAvgPool2d((1, 1)))
+        self.down_conv = nn.Sequential(*down_conv)
+        self.to_style = FullyConnectedLayer(in_features=dim, out_features=dim * 2, activation=activation)
+        self.ws_style = FullyConnectedLayer(in_features=w_dim, out_features=dim, activation=activation)
+        self.to_square = FullyConnectedLayer(in_features=dim, out_features=16 * 16, activation=activation)
+
+        style_dim = dim * 3
+        self.dec_conv = nn.ModuleList()
+        for i in range(down_time):  # from 64 to input size
+            res = res * 2
+            self.dec_conv.append(
+                DecStyleBlock(res, dim, dim, activation, style_dim, use_noise, demodulate, img_channels))
+
+    def forward(self, images_in, masks_in, ws, noise_mode='random'):
+        x = torch.cat([masks_in - 0.5, images_in * masks_in], dim=1)
+
+        skips = []
+        x, mask = self.conv_first(x, masks_in)  # input size
+        skips.append(x)
+        for i, block in enumerate(self.enc_conv):  # input size to 64
+            x, mask = block(x, mask)
+            if i != len(self.enc_conv) - 1:
+                skips.append(x)
+
+        x_size = x.size()[-2:]
+        x = feature2token(x)
+        mask = feature2token(mask)
+        mid = len(self.tran) // 2
+        for i, block in enumerate(self.tran):  # 64 to 16
+            if i < mid:
+                x, x_size, mask = block(x, x_size, mask)
+                skips.append(x)
+            elif i > mid:
+                x, x_size, mask = block(x, x_size, None)
+                x = x + skips[mid - i]
+            else:
+                x, x_size, mask = block(x, x_size, None)
+
+                mul_map = torch.ones_like(x) * 0.5
+                mul_map = F.dropout(mul_map, training=True)
+                ws = self.ws_style(ws[:, -1])
+                add_n = self.to_square(ws).unsqueeze(1)
+                add_n = F.interpolate(add_n, size=x.size(1), mode='linear', align_corners=False).squeeze(1).unsqueeze(
+                    -1)
+                x = x * mul_map + add_n * (1 - mul_map)
+                gs = self.to_style(self.down_conv(token2feature(x, x_size)).flatten(start_dim=1))
+                style = torch.cat([gs, ws], dim=1)
+
+        x = token2feature(x, x_size).contiguous()
+        img = None
+        for i, block in enumerate(self.dec_conv):
+            x, img = block(x, img, style, skips[len(self.dec_conv) - i - 1], noise_mode=noise_mode)
+
+        # ensemble
+        img = img * (1 - masks_in) + images_in * masks_in
+
+        return img
+
+
+class SynthesisNet(nn.Module):
+    def __init__(self,
+                 w_dim,  # Intermediate latent (W) dimensionality.
+                 img_resolution,  # Output image resolution.
+                 img_channels=3,  # Number of color channels.
+                 channel_base=32768,  # Overall multiplier for the number of channels.
+                 channel_decay=1.0,
+                 channel_max=512,  # Maximum number of channels in any layer.
+                 activation='lrelu',  # Activation function: 'relu', 'lrelu', etc.
+                 drop_rate=0.5,
+                 use_noise=False,
+                 demodulate=True,
+                 ):
+        super().__init__()
+        resolution_log2 = int(np.log2(img_resolution))
+        assert img_resolution == 2 ** resolution_log2 and img_resolution >= 4
+
+        self.num_layers = resolution_log2 * 2 - 3 * 2
+        self.img_resolution = img_resolution
+        self.resolution_log2 = resolution_log2
+
+        # first stage
+        self.first_stage = FirstStage(img_channels, img_resolution=img_resolution, w_dim=w_dim, use_noise=False,
+                                      demodulate=demodulate)
+
+        # second stage
+        self.enc = Encoder(resolution_log2, img_channels, activation, patch_size=5, channels=16)
+        self.to_square = FullyConnectedLayer(in_features=w_dim, out_features=16 * 16, activation=activation)
+        self.to_style = ToStyle(in_channels=nf(4), out_channels=nf(2) * 2, activation=activation, drop_rate=drop_rate)
+        style_dim = w_dim + nf(2) * 2
+        self.dec = Decoder(resolution_log2, activation, style_dim, use_noise, demodulate, img_channels)
+
+    def forward(self, images_in, masks_in, ws, noise_mode='random', return_stg1=False):
+        out_stg1 = self.first_stage(images_in, masks_in, ws, noise_mode=noise_mode)
+
+        # encoder
+        x = images_in * masks_in + out_stg1 * (1 - masks_in)
+        x = torch.cat([masks_in - 0.5, x, images_in * masks_in], dim=1)
+        E_features = self.enc(x)
+
+        fea_16 = E_features[4]
+        mul_map = torch.ones_like(fea_16) * 0.5
+        mul_map = F.dropout(mul_map, training=True)
+        add_n = self.to_square(ws[:, 0]).view(-1, 16, 16).unsqueeze(1)
+        add_n = F.interpolate(add_n, size=fea_16.size()[-2:], mode='bilinear', align_corners=False)
+        fea_16 = fea_16 * mul_map + add_n * (1 - mul_map)
+        E_features[4] = fea_16
+
+        # style
+        gs = self.to_style(fea_16)
+
+        # decoder
+        img = self.dec(fea_16, ws, gs, E_features, noise_mode=noise_mode)
+
+        # ensemble
+        img = img * (1 - masks_in) + images_in * masks_in
+
+        if not return_stg1:
+            return img
+        else:
+            return img, out_stg1
+
+
+class Generator(nn.Module):
+    def __init__(self,
+                 z_dim,  # Input latent (Z) dimensionality, 0 = no latent.
+                 c_dim,  # Conditioning label (C) dimensionality, 0 = no label.
+                 w_dim,  # Intermediate latent (W) dimensionality.
+                 img_resolution,  # resolution of generated image
+                 img_channels,  # Number of input color channels.
+                 synthesis_kwargs={},  # Arguments for SynthesisNetwork.
+                 mapping_kwargs={},  # Arguments for MappingNetwork.
+                 ):
+        super().__init__()
+        self.z_dim = z_dim
+        self.c_dim = c_dim
+        self.w_dim = w_dim
+        self.img_resolution = img_resolution
+        self.img_channels = img_channels
+
+        self.synthesis = SynthesisNet(w_dim=w_dim,
+                                      img_resolution=img_resolution,
+                                      img_channels=img_channels,
+                                      **synthesis_kwargs)
+        self.mapping = MappingNet(z_dim=z_dim,
+                                  c_dim=c_dim,
+                                  w_dim=w_dim,
+                                  num_ws=self.synthesis.num_layers,
+                                  **mapping_kwargs)
+
+    def forward(self, images_in, masks_in, z, c, truncation_psi=1, truncation_cutoff=None, skip_w_avg_update=False,
+                noise_mode='none', return_stg1=False):
+        ws = self.mapping(z, c, truncation_psi=truncation_psi, truncation_cutoff=truncation_cutoff,
+                          skip_w_avg_update=skip_w_avg_update)
+        img = self.synthesis(images_in, masks_in, ws, noise_mode=noise_mode)
+        return img
+
+
+class Discriminator(torch.nn.Module):
+    def __init__(self,
+                 c_dim,  # Conditioning label (C) dimensionality.
+                 img_resolution,  # Input resolution.
+                 img_channels,  # Number of input color channels.
+                 channel_base=32768,  # Overall multiplier for the number of channels.
+                 channel_max=512,  # Maximum number of channels in any layer.
+                 channel_decay=1,
+                 cmap_dim=None,  # Dimensionality of mapped conditioning label, None = default.
+                 activation='lrelu',
+                 mbstd_group_size=4,  # Group size for the minibatch standard deviation layer, None = entire minibatch.
+                 mbstd_num_channels=1,  # Number of features for the minibatch standard deviation layer, 0 = disable.
+                 ):
+        super().__init__()
+        self.c_dim = c_dim
+        self.img_resolution = img_resolution
+        self.img_channels = img_channels
+
+        resolution_log2 = int(np.log2(img_resolution))
+        assert img_resolution == 2 ** resolution_log2 and img_resolution >= 4
+        self.resolution_log2 = resolution_log2
+
+        if cmap_dim == None:
+            cmap_dim = nf(2)
+        if c_dim == 0:
+            cmap_dim = 0
+        self.cmap_dim = cmap_dim
+
+        if c_dim > 0:
+            self.mapping = MappingNet(z_dim=0, c_dim=c_dim, w_dim=cmap_dim, num_ws=None, w_avg_beta=None)
+
+        Dis = [DisFromRGB(img_channels + 1, nf(resolution_log2), activation)]
+        for res in range(resolution_log2, 2, -1):
+            Dis.append(DisBlock(nf(res), nf(res - 1), activation))
+
+        if mbstd_num_channels > 0:
+            Dis.append(MinibatchStdLayer(group_size=mbstd_group_size, num_channels=mbstd_num_channels))
+        Dis.append(Conv2dLayer(nf(2) + mbstd_num_channels, nf(2), kernel_size=3, activation=activation))
+        self.Dis = nn.Sequential(*Dis)
+
+        self.fc0 = FullyConnectedLayer(nf(2) * 4 ** 2, nf(2), activation=activation)
+        self.fc1 = FullyConnectedLayer(nf(2), 1 if cmap_dim == 0 else cmap_dim)
+
+        # for 64x64
+        Dis_stg1 = [DisFromRGB(img_channels + 1, nf(resolution_log2) // 2, activation)]
+        for res in range(resolution_log2, 2, -1):
+            Dis_stg1.append(DisBlock(nf(res) // 2, nf(res - 1) // 2, activation))
+
+        if mbstd_num_channels > 0:
+            Dis_stg1.append(MinibatchStdLayer(group_size=mbstd_group_size, num_channels=mbstd_num_channels))
+        Dis_stg1.append(Conv2dLayer(nf(2) // 2 + mbstd_num_channels, nf(2) // 2, kernel_size=3, activation=activation))
+        self.Dis_stg1 = nn.Sequential(*Dis_stg1)
+
+        self.fc0_stg1 = FullyConnectedLayer(nf(2) // 2 * 4 ** 2, nf(2) // 2, activation=activation)
+        self.fc1_stg1 = FullyConnectedLayer(nf(2) // 2, 1 if cmap_dim == 0 else cmap_dim)
+
+    def forward(self, images_in, masks_in, images_stg1, c):
+        x = self.Dis(torch.cat([masks_in - 0.5, images_in], dim=1))
+        x = self.fc1(self.fc0(x.flatten(start_dim=1)))
+
+        x_stg1 = self.Dis_stg1(torch.cat([masks_in - 0.5, images_stg1], dim=1))
+        x_stg1 = self.fc1_stg1(self.fc0_stg1(x_stg1.flatten(start_dim=1)))
+
+        if self.c_dim > 0:
+            cmap = self.mapping(None, c)
+
+        if self.cmap_dim > 0:
+            x = (x * cmap).sum(dim=1, keepdim=True) * (1 / np.sqrt(self.cmap_dim))
+            x_stg1 = (x_stg1 * cmap).sum(dim=1, keepdim=True) * (1 / np.sqrt(self.cmap_dim))
+
+        return x, x_stg1
+
+
+MAT_MODEL_URL = os.environ.get(
+    "MAT_MODEL_URL",
+    "https://github.com/Sanster/models/releases/download/add_mat/Places_512_FullData_G.pth",
+)
+
+
+class MAT(InpaintModel):
+    min_size = 512
+    pad_mod = 512
+    pad_to_square = True
+
+    def init_model(self, device, **kwargs):
+        seed = 240  # pick up a random number
+        random.seed(seed)
+        np.random.seed(seed)
+        torch.manual_seed(seed)
+
+        G = Generator(z_dim=512, c_dim=0, w_dim=512, img_resolution=512, img_channels=3)
+        self.model = load_model(G, MAT_MODEL_URL, device)
+        self.z = torch.from_numpy(np.random.randn(1, G.z_dim)).to(device)  # [1., 512]
+        self.label = torch.zeros([1, self.model.c_dim], device=device)
+
+    @staticmethod
+    def is_downloaded() -> bool:
+        return os.path.exists(get_cache_path_by_url(MAT_MODEL_URL))
+
+    def forward(self, image, mask, config: Config):
+        """Input images and output images have same size
+        images: [H, W, C] RGB
+        masks: [H, W] mask area == 255
+        return: BGR IMAGE
+        """
+
+        image = norm_img(image)  # [0, 1]
+        image = image * 2 - 1  # [0, 1] -> [-1, 1]
+
+        mask = (mask > 127) * 255
+        mask = 255 - mask
+        mask = norm_img(mask)
+
+        image = torch.from_numpy(image).unsqueeze(0).to(self.device)
+        mask = torch.from_numpy(mask).unsqueeze(0).to(self.device)
+
+        output = self.model(image, mask, self.z, self.label, truncation_psi=1, noise_mode='none')
+        output = (output.permute(0, 2, 3, 1) * 127.5 + 127.5).round().clamp(0, 255).to(torch.uint8)
+        output = output[0].cpu().numpy()
+        cur_res = cv2.cvtColor(output, cv2.COLOR_RGB2BGR)
+        return cur_res

lama_cleaner/model/opencv2.py

@@ -0,0 +1,24 @@
+import cv2
+from lama_cleaner.model.base import InpaintModel
+from lama_cleaner.schema import Config
+
+flag_map = {
+    "INPAINT_NS": cv2.INPAINT_NS,
+    "INPAINT_TELEA": cv2.INPAINT_TELEA
+}
+
+class OpenCV2(InpaintModel):
+    pad_mod = 1
+
+    @staticmethod
+    def is_downloaded() -> bool:
+        return True
+
+    def forward(self, image, mask, config: Config):
+        """Input image and output image have same size
+        image: [H, W, C] RGB
+        mask: [H, W, 1]
+        return: BGR IMAGE
+        """
+        cur_res = cv2.inpaint(image[:,:,::-1], mask, inpaintRadius=config.cv2_radius, flags=flag_map[config.cv2_flag])
+        return cur_res

lama_cleaner/model/plms_sampler.py

@@ -0,0 +1,225 @@
+# From: https://github.com/CompVis/latent-diffusion/blob/main/ldm/models/diffusion/plms.py
+import torch
+import numpy as np
+from lama_cleaner.model.utils import make_ddim_timesteps, make_ddim_sampling_parameters, noise_like
+from tqdm import tqdm
+
+
+class PLMSSampler(object):
+    def __init__(self, model, schedule="linear", **kwargs):
+        super().__init__()
+        self.model = model
+        self.ddpm_num_timesteps = model.num_timesteps
+        self.schedule = schedule
+
+    def register_buffer(self, name, attr):
+        setattr(self, name, attr)
+
+    def make_schedule(self, ddim_num_steps, ddim_discretize="uniform", ddim_eta=0., verbose=True):
+        if ddim_eta != 0:
+            raise ValueError('ddim_eta must be 0 for PLMS')
+        self.ddim_timesteps = make_ddim_timesteps(ddim_discr_method=ddim_discretize, num_ddim_timesteps=ddim_num_steps,
+                                                  num_ddpm_timesteps=self.ddpm_num_timesteps, verbose=verbose)
+        alphas_cumprod = self.model.alphas_cumprod
+        assert alphas_cumprod.shape[0] == self.ddpm_num_timesteps, 'alphas have to be defined for each timestep'
+        to_torch = lambda x: x.clone().detach().to(torch.float32).to(self.model.device)
+
+        self.register_buffer('betas', to_torch(self.model.betas))
+        self.register_buffer('alphas_cumprod', to_torch(alphas_cumprod))
+        self.register_buffer('alphas_cumprod_prev', to_torch(self.model.alphas_cumprod_prev))
+
+        # calculations for diffusion q(x_t | x_{t-1}) and others
+        self.register_buffer('sqrt_alphas_cumprod', to_torch(np.sqrt(alphas_cumprod.cpu())))
+        self.register_buffer('sqrt_one_minus_alphas_cumprod', to_torch(np.sqrt(1. - alphas_cumprod.cpu())))
+        self.register_buffer('log_one_minus_alphas_cumprod', to_torch(np.log(1. - alphas_cumprod.cpu())))
+        self.register_buffer('sqrt_recip_alphas_cumprod', to_torch(np.sqrt(1. / alphas_cumprod.cpu())))
+        self.register_buffer('sqrt_recipm1_alphas_cumprod', to_torch(np.sqrt(1. / alphas_cumprod.cpu() - 1)))
+
+        # ddim sampling parameters
+        ddim_sigmas, ddim_alphas, ddim_alphas_prev = make_ddim_sampling_parameters(alphacums=alphas_cumprod.cpu(),
+                                                                                   ddim_timesteps=self.ddim_timesteps,
+                                                                                   eta=ddim_eta, verbose=verbose)
+        self.register_buffer('ddim_sigmas', ddim_sigmas)
+        self.register_buffer('ddim_alphas', ddim_alphas)
+        self.register_buffer('ddim_alphas_prev', ddim_alphas_prev)
+        self.register_buffer('ddim_sqrt_one_minus_alphas', np.sqrt(1. - ddim_alphas))
+        sigmas_for_original_sampling_steps = ddim_eta * torch.sqrt(
+            (1 - self.alphas_cumprod_prev) / (1 - self.alphas_cumprod) * (
+                    1 - self.alphas_cumprod / self.alphas_cumprod_prev))
+        self.register_buffer('ddim_sigmas_for_original_num_steps', sigmas_for_original_sampling_steps)
+
+    @torch.no_grad()
+    def sample(self,
+               steps,
+               batch_size,
+               shape,
+               conditioning=None,
+               callback=None,
+               normals_sequence=None,
+               img_callback=None,
+               quantize_x0=False,
+               eta=0.,
+               mask=None,
+               x0=None,
+               temperature=1.,
+               noise_dropout=0.,
+               score_corrector=None,
+               corrector_kwargs=None,
+               verbose=False,
+               x_T=None,
+               log_every_t=100,
+               unconditional_guidance_scale=1.,
+               unconditional_conditioning=None,
+               # this has to come in the same format as the conditioning, # e.g. as encoded tokens, ...
+               **kwargs
+               ):
+        if conditioning is not None:
+            if isinstance(conditioning, dict):
+                cbs = conditioning[list(conditioning.keys())[0]].shape[0]
+                if cbs != batch_size:
+                    print(f"Warning: Got {cbs} conditionings but batch-size is {batch_size}")
+            else:
+                if conditioning.shape[0] != batch_size:
+                    print(f"Warning: Got {conditioning.shape[0]} conditionings but batch-size is {batch_size}")
+
+        self.make_schedule(ddim_num_steps=steps, ddim_eta=eta, verbose=verbose)
+        # sampling
+        C, H, W = shape
+        size = (batch_size, C, H, W)
+        print(f'Data shape for PLMS sampling is {size}')
+
+        samples = self.plms_sampling(conditioning, size,
+                                     callback=callback,
+                                     img_callback=img_callback,
+                                     quantize_denoised=quantize_x0,
+                                     mask=mask, x0=x0,
+                                     ddim_use_original_steps=False,
+                                     noise_dropout=noise_dropout,
+                                     temperature=temperature,
+                                     score_corrector=score_corrector,
+                                     corrector_kwargs=corrector_kwargs,
+                                     x_T=x_T,
+                                     log_every_t=log_every_t,
+                                     unconditional_guidance_scale=unconditional_guidance_scale,
+                                     unconditional_conditioning=unconditional_conditioning,
+                                     )
+        return samples
+
+    @torch.no_grad()
+    def plms_sampling(self, cond, shape,
+                      x_T=None, ddim_use_original_steps=False,
+                      callback=None, timesteps=None, quantize_denoised=False,
+                      mask=None, x0=None, img_callback=None, log_every_t=100,
+                      temperature=1., noise_dropout=0., score_corrector=None, corrector_kwargs=None,
+                      unconditional_guidance_scale=1., unconditional_conditioning=None, ):
+        device = self.model.betas.device
+        b = shape[0]
+        if x_T is None:
+            img = torch.randn(shape, device=device)
+        else:
+            img = x_T
+
+        if timesteps is None:
+            timesteps = self.ddpm_num_timesteps if ddim_use_original_steps else self.ddim_timesteps
+        elif timesteps is not None and not ddim_use_original_steps:
+            subset_end = int(min(timesteps / self.ddim_timesteps.shape[0], 1) * self.ddim_timesteps.shape[0]) - 1
+            timesteps = self.ddim_timesteps[:subset_end]
+
+        time_range = list(reversed(range(0, timesteps))) if ddim_use_original_steps else np.flip(timesteps)
+        total_steps = timesteps if ddim_use_original_steps else timesteps.shape[0]
+        print(f"Running PLMS Sampling with {total_steps} timesteps")
+
+        iterator = tqdm(time_range, desc='PLMS Sampler', total=total_steps)
+        old_eps = []
+
+        for i, step in enumerate(iterator):
+            index = total_steps - i - 1
+            ts = torch.full((b,), step, device=device, dtype=torch.long)
+            ts_next = torch.full((b,), time_range[min(i + 1, len(time_range) - 1)], device=device, dtype=torch.long)
+
+            if mask is not None:
+                assert x0 is not None
+                img_orig = self.model.q_sample(x0, ts)  # TODO: deterministic forward pass?
+                img = img_orig * mask + (1. - mask) * img
+
+            outs = self.p_sample_plms(img, cond, ts, index=index, use_original_steps=ddim_use_original_steps,
+                                      quantize_denoised=quantize_denoised, temperature=temperature,
+                                      noise_dropout=noise_dropout, score_corrector=score_corrector,
+                                      corrector_kwargs=corrector_kwargs,
+                                      unconditional_guidance_scale=unconditional_guidance_scale,
+                                      unconditional_conditioning=unconditional_conditioning,
+                                      old_eps=old_eps, t_next=ts_next)
+            img, pred_x0, e_t = outs
+            old_eps.append(e_t)
+            if len(old_eps) >= 4:
+                old_eps.pop(0)
+            if callback: callback(i)
+            if img_callback: img_callback(pred_x0, i)
+
+        return img
+
+    @torch.no_grad()
+    def p_sample_plms(self, x, c, t, index, repeat_noise=False, use_original_steps=False, quantize_denoised=False,
+                      temperature=1., noise_dropout=0., score_corrector=None, corrector_kwargs=None,
+                      unconditional_guidance_scale=1., unconditional_conditioning=None, old_eps=None, t_next=None):
+        b, *_, device = *x.shape, x.device
+
+        def get_model_output(x, t):
+            if unconditional_conditioning is None or unconditional_guidance_scale == 1.:
+                e_t = self.model.apply_model(x, t, c)
+            else:
+                x_in = torch.cat([x] * 2)
+                t_in = torch.cat([t] * 2)
+                c_in = torch.cat([unconditional_conditioning, c])
+                e_t_uncond, e_t = self.model.apply_model(x_in, t_in, c_in).chunk(2)
+                e_t = e_t_uncond + unconditional_guidance_scale * (e_t - e_t_uncond)
+
+            if score_corrector is not None:
+                assert self.model.parameterization == "eps"
+                e_t = score_corrector.modify_score(self.model, e_t, x, t, c, **corrector_kwargs)
+
+            return e_t
+
+        alphas = self.model.alphas_cumprod if use_original_steps else self.ddim_alphas
+        alphas_prev = self.model.alphas_cumprod_prev if use_original_steps else self.ddim_alphas_prev
+        sqrt_one_minus_alphas = self.model.sqrt_one_minus_alphas_cumprod if use_original_steps else self.ddim_sqrt_one_minus_alphas
+        sigmas = self.model.ddim_sigmas_for_original_num_steps if use_original_steps else self.ddim_sigmas
+
+        def get_x_prev_and_pred_x0(e_t, index):
+            # select parameters corresponding to the currently considered timestep
+            a_t = torch.full((b, 1, 1, 1), alphas[index], device=device)
+            a_prev = torch.full((b, 1, 1, 1), alphas_prev[index], device=device)
+            sigma_t = torch.full((b, 1, 1, 1), sigmas[index], device=device)
+            sqrt_one_minus_at = torch.full((b, 1, 1, 1), sqrt_one_minus_alphas[index], device=device)
+
+            # current prediction for x_0
+            pred_x0 = (x - sqrt_one_minus_at * e_t) / a_t.sqrt()
+            if quantize_denoised:
+                pred_x0, _, *_ = self.model.first_stage_model.quantize(pred_x0)
+            # direction pointing to x_t
+            dir_xt = (1. - a_prev - sigma_t ** 2).sqrt() * e_t
+            noise = sigma_t * noise_like(x.shape, device, repeat_noise) * temperature
+            if noise_dropout > 0.:
+                noise = torch.nn.functional.dropout(noise, p=noise_dropout)
+            x_prev = a_prev.sqrt() * pred_x0 + dir_xt + noise
+            return x_prev, pred_x0
+
+        e_t = get_model_output(x, t)
+        if len(old_eps) == 0:
+            # Pseudo Improved Euler (2nd order)
+            x_prev, pred_x0 = get_x_prev_and_pred_x0(e_t, index)
+            e_t_next = get_model_output(x_prev, t_next)
+            e_t_prime = (e_t + e_t_next) / 2
+        elif len(old_eps) == 1:
+            # 2nd order Pseudo Linear Multistep (Adams-Bashforth)
+            e_t_prime = (3 * e_t - old_eps[-1]) / 2
+        elif len(old_eps) == 2:
+            # 3nd order Pseudo Linear Multistep (Adams-Bashforth)
+            e_t_prime = (23 * e_t - 16 * old_eps[-1] + 5 * old_eps[-2]) / 12
+        elif len(old_eps) >= 3:
+            # 4nd order Pseudo Linear Multistep (Adams-Bashforth)
+            e_t_prime = (55 * e_t - 59 * old_eps[-1] + 37 * old_eps[-2] - 9 * old_eps[-3]) / 24
+
+        x_prev, pred_x0 = get_x_prev_and_pred_x0(e_t_prime, index)
+
+        return x_prev, pred_x0, e_t

lama_cleaner/model/sd.py

@@ -0,0 +1,215 @@
+import random
+
+import PIL.Image
+import cv2
+import numpy as np
+import torch
+from diffusers import PNDMScheduler, DDIMScheduler
+from loguru import logger
+from transformers import FeatureExtractionMixin, ImageFeatureExtractionMixin
+
+from lama_cleaner.helper import norm_img
+
+from lama_cleaner.model.base import InpaintModel
+from lama_cleaner.schema import Config, SDSampler
+
+
+#
+#
+# def preprocess_image(image):
+#     w, h = image.size
+#     w, h = map(lambda x: x - x % 32, (w, h))  # resize to integer multiple of 32
+#     image = image.resize((w, h), resample=PIL.Image.LANCZOS)
+#     image = np.array(image).astype(np.float32) / 255.0
+#     image = image[None].transpose(0, 3, 1, 2)
+#     image = torch.from_numpy(image)
+#     # [-1, 1]
+#     return 2.0 * image - 1.0
+#
+#
+# def preprocess_mask(mask):
+#     mask = mask.convert("L")
+#     w, h = mask.size
+#     w, h = map(lambda x: x - x % 32, (w, h))  # resize to integer multiple of 32
+#     mask = mask.resize((w // 8, h // 8), resample=PIL.Image.NEAREST)
+#     mask = np.array(mask).astype(np.float32) / 255.0
+#     mask = np.tile(mask, (4, 1, 1))
+#     mask = mask[None].transpose(0, 1, 2, 3)  # what does this step do?
+#     mask = 1 - mask  # repaint white, keep black
+#     mask = torch.from_numpy(mask)
+#     return mask
+
+class DummyFeatureExtractorOutput:
+    def __init__(self, pixel_values):
+        self.pixel_values = pixel_values
+
+    def to(self, device):
+        return self
+
+
+class DummyFeatureExtractor(FeatureExtractionMixin, ImageFeatureExtractionMixin):
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+
+    def __call__(self, *args, **kwargs):
+        return DummyFeatureExtractorOutput(torch.empty(0, 3))
+
+
+class DummySafetyChecker:
+    def __init__(self, *args, **kwargs):
+        pass
+
+    def __call__(self, clip_input, images):
+        return images, False
+
+
+class SD(InpaintModel):
+    pad_mod = 64  # current diffusers only support 64 https://github.com/huggingface/diffusers/pull/505
+    min_size = 512
+
+    def init_model(self, device: torch.device, **kwargs):
+        from .sd_pipeline import StableDiffusionInpaintPipeline
+
+        model_kwargs = {"local_files_only": kwargs['sd_run_local']}
+        if kwargs['sd_disable_nsfw']:
+            logger.info("Disable Stable Diffusion Model NSFW checker")
+            model_kwargs.update(dict(
+                feature_extractor=DummyFeatureExtractor(),
+                safety_checker=DummySafetyChecker(),
+            ))
+
+        self.model = StableDiffusionInpaintPipeline.from_pretrained(
+            self.model_id_or_path,
+            revision="fp16" if torch.cuda.is_available() else "main",
+            torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
+            use_auth_token=kwargs["hf_access_token"],
+            **model_kwargs
+        )
+        # https://huggingface.co/docs/diffusers/v0.3.0/en/api/pipelines/stable_diffusion#diffusers.StableDiffusionInpaintPipeline.enable_attention_slicing
+        self.model.enable_attention_slicing()
+        self.model = self.model.to(device)
+
+        if kwargs['sd_cpu_textencoder']:
+            logger.info("Run Stable Diffusion TextEncoder on CPU")
+            self.model.text_encoder = self.model.text_encoder.to(torch.device('cpu'), non_blocking=True)
+            self.model.text_encoder = self.model.text_encoder.to(torch.float32, non_blocking=True )
+
+        self.callbacks = kwargs.pop("callbacks", None)
+
+    @torch.cuda.amp.autocast()
+    def forward(self, image, mask, config: Config):
+        """Input image and output image have same size
+        image: [H, W, C] RGB
+        mask: [H, W, 1] 255 means area to repaint
+        return: BGR IMAGE
+        """
+
+        # image = norm_img(image)  # [0, 1]
+        # image = image * 2 - 1  # [0, 1] -> [-1, 1]
+
+        # resize to latent feature map size
+        # h, w = mask.shape[:2]
+        # mask = cv2.resize(mask, (h // 8, w // 8), interpolation=cv2.INTER_AREA)
+        # mask = norm_img(mask)
+        #
+        # image = torch.from_numpy(image).unsqueeze(0).to(self.device)
+        # mask = torch.from_numpy(mask).unsqueeze(0).to(self.device)
+
+        if config.sd_sampler == SDSampler.ddim:
+            scheduler = DDIMScheduler(
+                beta_start=0.00085,
+                beta_end=0.012,
+                beta_schedule="scaled_linear",
+                clip_sample=False,
+                set_alpha_to_one=False,
+            )
+        elif config.sd_sampler == SDSampler.pndm:
+            PNDM_kwargs = {
+                "tensor_format": "pt",
+                "beta_schedule": "scaled_linear",
+                "beta_start": 0.00085,
+                "beta_end": 0.012,
+                "num_train_timesteps": 1000,
+                "skip_prk_steps": True,
+            }
+            scheduler = PNDMScheduler(**PNDM_kwargs)
+        else:
+            raise ValueError(config.sd_sampler)
+
+        self.model.scheduler = scheduler
+
+        seed = config.sd_seed
+        random.seed(seed)
+        np.random.seed(seed)
+        torch.manual_seed(seed)
+        torch.cuda.manual_seed_all(seed)
+
+        if config.sd_mask_blur != 0:
+            k = 2 * config.sd_mask_blur + 1
+            mask = cv2.GaussianBlur(mask, (k, k), 0)[:, :, np.newaxis]
+
+        output = self.model(
+            prompt=config.prompt,
+            init_image=PIL.Image.fromarray(image),
+            mask_image=PIL.Image.fromarray(mask[:, :, -1], mode="L"),
+            strength=config.sd_strength,
+            num_inference_steps=config.sd_steps,
+            guidance_scale=config.sd_guidance_scale,
+            output_type="np.array",
+            callbacks=self.callbacks,
+        ).images[0]
+
+        output = (output * 255).round().astype("uint8")
+        output = cv2.cvtColor(output, cv2.COLOR_RGB2BGR)
+        return output
+
+    @torch.no_grad()
+    def __call__(self, image, mask, config: Config):
+        """
+        images: [H, W, C] RGB, not normalized
+        masks: [H, W]
+        return: BGR IMAGE
+        """
+        img_h, img_w = image.shape[:2]
+
+        # boxes = boxes_from_mask(mask)
+        if config.use_croper:
+            logger.info("use croper")
+            l, t, w, h = (
+                config.croper_x,
+                config.croper_y,
+                config.croper_width,
+                config.croper_height,
+            )
+            r = l + w
+            b = t + h
+
+            l = max(l, 0)
+            r = min(r, img_w)
+            t = max(t, 0)
+            b = min(b, img_h)
+
+            crop_img = image[t:b, l:r, :]
+            crop_mask = mask[t:b, l:r]
+
+            crop_image = self._pad_forward(crop_img, crop_mask, config)
+
+            inpaint_result = image[:, :, ::-1]
+            inpaint_result[t:b, l:r, :] = crop_image
+        else:
+            inpaint_result = self._pad_forward(image, mask, config)
+
+        return inpaint_result
+
+    @staticmethod
+    def is_downloaded() -> bool:
+        # model will be downloaded when app start, and can't switch in frontend settings
+        return True
+
+
+class SD14(SD):
+    model_id_or_path = "CompVis/stable-diffusion-v1-4"
+
+
+class SD15(SD):
+    model_id_or_path = "CompVis/stable-diffusion-v1-5"

lama_cleaner/model/sd_pipeline.py

@@ -0,0 +1,310 @@
+import inspect
+from typing import List, Optional, Union, Callable
+
+import numpy as np
+import torch
+
+import PIL
+from diffusers import DiffusionPipeline, AutoencoderKL, UNet2DConditionModel, DDIMScheduler, PNDMScheduler
+from diffusers.pipelines.stable_diffusion import StableDiffusionSafetyChecker, StableDiffusionPipelineOutput
+from diffusers.utils import logging
+from tqdm.auto import tqdm
+from transformers import CLIPFeatureExtractor, CLIPTextModel, CLIPTokenizer
+
+logger = logging.get_logger(__name__)
+
+
+def preprocess_image(image):
+    w, h = image.size
+    w, h = map(lambda x: x - x % 32, (w, h))  # resize to integer multiple of 32
+    image = image.resize((w, h), resample=PIL.Image.LANCZOS)
+    image = np.array(image).astype(np.float32) / 255.0
+    image = image[None].transpose(0, 3, 1, 2)
+    image = torch.from_numpy(image)
+    return 2.0 * image - 1.0
+
+
+def preprocess_mask(mask):
+    mask = mask.convert("L")
+    w, h = mask.size
+    w, h = map(lambda x: x - x % 32, (w, h))  # resize to integer multiple of 32
+    mask = mask.resize((w // 8, h // 8), resample=PIL.Image.NEAREST)
+    mask = np.array(mask).astype(np.float32) / 255.0
+    mask = np.tile(mask, (4, 1, 1))
+    mask = mask[None].transpose(0, 1, 2, 3)  # what does this step do?
+    mask = 1 - mask  # repaint white, keep black
+    mask = torch.from_numpy(mask)
+    return mask
+
+
+class StableDiffusionInpaintPipeline(DiffusionPipeline):
+    r"""
+    Pipeline for text-guided image inpainting using Stable Diffusion. *This is an experimental feature*.
+
+    This model inherits from [`DiffusionPipeline`]. Check the superclass documentation for the generic methods the
+    library implements for all the pipelines (such as downloading or saving, running on a particular device, etc.)
+
+    Args:
+        vae ([`AutoencoderKL`]):
+            Variational Auto-Encoder (VAE) Model to encode and decode images to and from latent representations.
+        text_encoder ([`CLIPTextModel`]):
+            Frozen text-encoder. Stable Diffusion uses the text portion of
+            [CLIP](https://huggingface.co/docs/transformers/model_doc/clip#transformers.CLIPTextModel), specifically
+            the [clip-vit-large-patch14](https://huggingface.co/openai/clip-vit-large-patch14) variant.
+        tokenizer (`CLIPTokenizer`):
+            Tokenizer of class
+            [CLIPTokenizer](https://huggingface.co/docs/transformers/v4.21.0/en/model_doc/clip#transformers.CLIPTokenizer).
+        unet ([`UNet2DConditionModel`]): Conditional U-Net architecture to denoise the encoded image latents.
+        scheduler ([`SchedulerMixin`]):
+            A scheduler to be used in combination with `unet` to denoise the encoded image latens. Can be one of
+            [`DDIMScheduler`], [`LMSDiscreteScheduler`], or [`PNDMScheduler`].
+        safety_checker ([`StableDiffusionSafetyChecker`]):
+            Classification module that estimates whether generated images could be considered offsensive or harmful.
+            Please, refer to the [model card](https://huggingface.co/CompVis/stable-diffusion-v1-4) for details.
+        feature_extractor ([`CLIPFeatureExtractor`]):
+            Model that extracts features from generated images to be used as inputs for the `safety_checker`.
+    """
+
+    def __init__(
+        self,
+        vae: AutoencoderKL,
+        text_encoder: CLIPTextModel,
+        tokenizer: CLIPTokenizer,
+        unet: UNet2DConditionModel,
+        scheduler: Union[DDIMScheduler, PNDMScheduler],
+        safety_checker: StableDiffusionSafetyChecker,
+        feature_extractor: CLIPFeatureExtractor,
+    ):
+        super().__init__()
+        scheduler = scheduler.set_format("pt")
+        logger.info("`StableDiffusionInpaintPipeline` is experimental and will very likely change in the future.")
+        self.register_modules(
+            vae=vae,
+            text_encoder=text_encoder,
+            tokenizer=tokenizer,
+            unet=unet,
+            scheduler=scheduler,
+            safety_checker=safety_checker,
+            feature_extractor=feature_extractor,
+        )
+
+    def enable_attention_slicing(self, slice_size: Optional[Union[str, int]] = "auto"):
+        r"""
+        Enable sliced attention computation.
+
+        When this option is enabled, the attention module will split the input tensor in slices, to compute attention
+        in several steps. This is useful to save some memory in exchange for a small speed decrease.
+
+        Args:
+            slice_size (`str` or `int`, *optional*, defaults to `"auto"`):
+                When `"auto"`, halves the input to the attention heads, so attention will be computed in two steps. If
+                a number is provided, uses as many slices as `attention_head_dim // slice_size`. In this case,
+                `attention_head_dim` must be a multiple of `slice_size`.
+        """
+        if slice_size == "auto":
+            # half the attention head size is usually a good trade-off between
+            # speed and memory
+            slice_size = self.unet.config.attention_head_dim // 2
+        self.unet.set_attention_slice(slice_size)
+
+    def disable_attention_slicing(self):
+        r"""
+        Disable sliced attention computation. If `enable_attention_slicing` was previously invoked, this method will go
+        back to computing attention in one step.
+        """
+        # set slice_size = `None` to disable `set_attention_slice`
+        self.enable_attention_slice(None)
+
+    @torch.no_grad()
+    def __call__(
+        self,
+        prompt: Union[str, List[str]],
+        init_image: Union[torch.FloatTensor, PIL.Image.Image],
+        mask_image: Union[torch.FloatTensor, PIL.Image.Image],
+        strength: float = 0.8,
+        num_inference_steps: Optional[int] = 50,
+        guidance_scale: Optional[float] = 7.5,
+        eta: Optional[float] = 0.0,
+        generator: Optional[torch.Generator] = None,
+        output_type: Optional[str] = "pil",
+        return_dict: bool = True,
+        callbacks: List[Callable[[int], None]] = None
+    ):
+        r"""
+        Function invoked when calling the pipeline for generation.
+
+        Args:
+            prompt (`str` or `List[str]`):
+                The prompt or prompts to guide the image generation.
+            init_image (`torch.FloatTensor` or `PIL.Image.Image`):
+                `Image`, or tensor representing an image batch, that will be used as the starting point for the
+                process. This is the image whose masked region will be inpainted.
+            mask_image (`torch.FloatTensor` or `PIL.Image.Image`):
+                `Image`, or tensor representing an image batch, to mask `init_image`. White pixels in the mask will be
+                replaced by noise and therefore repainted, while black pixels will be preserved. The mask image will be
+                converted to a single channel (luminance) before use.
+            strength (`float`, *optional*, defaults to 0.8):
+                Conceptually, indicates how much to inpaint the masked area. Must be between 0 and 1. When `strength`
+                is 1, the denoising process will be run on the masked area for the full number of iterations specified
+                in `num_inference_steps`. `init_image` will be used as a reference for the masked area, adding more
+                noise to that region the larger the `strength`. If `strength` is 0, no inpainting will occur.
+            num_inference_steps (`int`, *optional*, defaults to 50):
+                The reference number of denoising steps. More denoising steps usually lead to a higher quality image at
+                the expense of slower inference. This parameter will be modulated by `strength`, as explained above.
+            guidance_scale (`float`, *optional*, defaults to 7.5):
+                Guidance scale as defined in [Classifier-Free Diffusion Guidance](https://arxiv.org/abs/2207.12598).
+                `guidance_scale` is defined as `w` of equation 2. of [Imagen
+                Paper](https://arxiv.org/pdf/2205.11487.pdf). Guidance scale is enabled by setting `guidance_scale >
+                1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`,
+                usually at the expense of lower image quality.
+            eta (`float`, *optional*, defaults to 0.0):
+                Corresponds to parameter eta (η) in the DDIM paper: https://arxiv.org/abs/2010.02502. Only applies to
+                [`schedulers.DDIMScheduler`], will be ignored for others.
+            generator (`torch.Generator`, *optional*):
+                A [torch generator](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation
+                deterministic.
+            output_type (`str`, *optional*, defaults to `"pil"`):
+                The output format of the generate image. Choose between
+                [PIL](https://pillow.readthedocs.io/en/stable/): `PIL.Image.Image` or `nd.array`.
+            return_dict (`bool`, *optional*, defaults to `True`):
+                Whether or not to return a [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] instead of a
+                plain tuple.
+
+        Returns:
+            [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] or `tuple`:
+            [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] if `return_dict` is True, otherwise a `tuple.
+            When returning a tuple, the first element is a list with the generated images, and the second element is a
+            list of `bool`s denoting whether the corresponding generated image likely represents "not-safe-for-work"
+            (nsfw) content, according to the `safety_checker`.
+        """
+        if isinstance(prompt, str):
+            batch_size = 1
+        elif isinstance(prompt, list):
+            batch_size = len(prompt)
+        else:
+            raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}")
+
+        if strength < 0 or strength > 1:
+            raise ValueError(f"The value of strength should in [0.0, 1.0] but is {strength}")
+
+        # set timesteps
+        accepts_offset = "offset" in set(inspect.signature(self.scheduler.set_timesteps).parameters.keys())
+        extra_set_kwargs = {}
+        offset = 0
+        if accepts_offset:
+            offset = 1
+            extra_set_kwargs["offset"] = 1
+
+        self.scheduler.set_timesteps(num_inference_steps, **extra_set_kwargs)
+
+        # preprocess image
+        init_image = preprocess_image(init_image).to(self.device)
+
+        # encode the init image into latents and scale the latents
+        init_latent_dist = self.vae.encode(init_image.to(self.device)).latent_dist
+        init_latents = init_latent_dist.sample(generator=generator)
+
+        init_latents = 0.18215 * init_latents
+
+        # Expand init_latents for batch_size
+        init_latents = torch.cat([init_latents] * batch_size)
+        init_latents_orig = init_latents
+
+        # preprocess mask
+        mask = preprocess_mask(mask_image).to(self.device)
+        mask = torch.cat([mask] * batch_size)
+
+        # check sizes
+        if not mask.shape == init_latents.shape:
+            raise ValueError("The mask and init_image should be the same size!")
+
+        # get the original timestep using init_timestep
+        init_timestep = int(num_inference_steps * strength) + offset
+        init_timestep = min(init_timestep, num_inference_steps)
+        timesteps = self.scheduler.timesteps[-init_timestep]
+        timesteps = torch.tensor([timesteps] * batch_size, dtype=torch.long, device=self.device)
+
+        # add noise to latents using the timesteps
+        noise = torch.randn(init_latents.shape, generator=generator, device=self.device)
+        init_latents = self.scheduler.add_noise(init_latents, noise, timesteps)
+
+        # get prompt text embeddings
+        text_input = self.tokenizer(
+            prompt,
+            padding="max_length",
+            max_length=self.tokenizer.model_max_length,
+            truncation=True,
+            return_tensors="pt",
+        )
+        text_encoder_device = self.text_encoder.device
+
+        text_embeddings = self.text_encoder(text_input.input_ids.to(text_encoder_device, non_blocking=True))[0].to(self.device, non_blocking=True)
+
+        # here `guidance_scale` is defined analog to the guidance weight `w` of equation (2)
+        # of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1`
+        # corresponds to doing no classifier free guidance.
+        do_classifier_free_guidance = guidance_scale > 1.0
+        # get unconditional embeddings for classifier free guidance
+        if do_classifier_free_guidance:
+            max_length = text_input.input_ids.shape[-1]
+            uncond_input = self.tokenizer(
+                [""] * batch_size, padding="max_length", max_length=max_length, return_tensors="pt"
+            )
+            uncond_embeddings = self.text_encoder(uncond_input.input_ids.to(text_encoder_device, non_blocking=True))[0].to(self.device, non_blocking=True)
+
+            # For classifier free guidance, we need to do two forward passes.
+            # Here we concatenate the unconditional and text embeddings into a single batch
+            # to avoid doing two forward passes
+            text_embeddings = torch.cat([uncond_embeddings, text_embeddings])
+
+        # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature
+        # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers.
+        # eta corresponds to η in DDIM paper: https://arxiv.org/abs/2010.02502
+        # and should be between [0, 1]
+        accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys())
+        extra_step_kwargs = {}
+        if accepts_eta:
+            extra_step_kwargs["eta"] = eta
+
+        latents = init_latents
+        t_start = max(num_inference_steps - init_timestep + offset, 0)
+        for i, t in tqdm(enumerate(self.scheduler.timesteps[t_start:])):
+            # expand the latents if we are doing classifier free guidance
+            latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents
+            # predict the noise residual
+            noise_pred = self.unet(latent_model_input, t, encoder_hidden_states=text_embeddings).sample
+
+            # perform guidance
+            if do_classifier_free_guidance:
+                noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
+                noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)
+
+            # compute the previous noisy sample x_t -> x_t-1
+            latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs).prev_sample
+
+            # masking
+            init_latents_proper = self.scheduler.add_noise(init_latents_orig, noise, t)
+            latents = (init_latents_proper * mask) + (latents * (1 - mask))
+
+            if callbacks is not None:
+                for callback in callbacks:
+                    callback(i)
+
+        # scale and decode the image latents with vae
+        latents = 1 / 0.18215 * latents
+        image = self.vae.decode(latents).sample
+
+        image = (image / 2 + 0.5).clamp(0, 1)
+        image = image.cpu().permute(0, 2, 3, 1).numpy()
+
+        # run safety checker
+        safety_cheker_input = self.feature_extractor(self.numpy_to_pil(image), return_tensors="pt").to(self.device)
+        image, has_nsfw_concept = self.safety_checker(images=image, clip_input=safety_cheker_input.pixel_values)
+
+        if output_type == "pil":
+            image = self.numpy_to_pil(image)
+
+        if not return_dict:
+            return (image, has_nsfw_concept)
+
+        return StableDiffusionPipelineOutput(images=image, nsfw_content_detected=has_nsfw_concept)

lama_cleaner/model/utils.py

@@ -0,0 +1,709 @@
+import math
+from typing import Any
+
+import torch
+import numpy as np
+import collections
+from itertools import repeat
+
+from torch import conv2d, conv_transpose2d
+
+
+def make_beta_schedule(device, schedule, n_timestep, linear_start=1e-4, linear_end=2e-2, cosine_s=8e-3):
+    if schedule == "linear":
+        betas = (
+                torch.linspace(linear_start ** 0.5, linear_end ** 0.5, n_timestep, dtype=torch.float64) ** 2
+        )
+
+    elif schedule == "cosine":
+        timesteps = (torch.arange(n_timestep + 1, dtype=torch.float64) / n_timestep + cosine_s).to(device)
+        alphas = timesteps / (1 + cosine_s) * np.pi / 2
+        alphas = torch.cos(alphas).pow(2).to(device)
+        alphas = alphas / alphas[0]
+        betas = 1 - alphas[1:] / alphas[:-1]
+        betas = np.clip(betas, a_min=0, a_max=0.999)
+
+    elif schedule == "sqrt_linear":
+        betas = torch.linspace(linear_start, linear_end, n_timestep, dtype=torch.float64)
+    elif schedule == "sqrt":
+        betas = torch.linspace(linear_start, linear_end, n_timestep, dtype=torch.float64) ** 0.5
+    else:
+        raise ValueError(f"schedule '{schedule}' unknown.")
+    return betas.numpy()
+
+
+def make_ddim_sampling_parameters(alphacums, ddim_timesteps, eta, verbose=True):
+    # select alphas for computing the variance schedule
+    alphas = alphacums[ddim_timesteps]
+    alphas_prev = np.asarray([alphacums[0]] + alphacums[ddim_timesteps[:-1]].tolist())
+
+    # according the the formula provided in https://arxiv.org/abs/2010.02502
+    sigmas = eta * np.sqrt((1 - alphas_prev) / (1 - alphas) * (1 - alphas / alphas_prev))
+    if verbose:
+        print(f'Selected alphas for ddim sampler: a_t: {alphas}; a_(t-1): {alphas_prev}')
+        print(f'For the chosen value of eta, which is {eta}, '
+              f'this results in the following sigma_t schedule for ddim sampler {sigmas}')
+    return sigmas, alphas, alphas_prev
+
+
+def make_ddim_timesteps(ddim_discr_method, num_ddim_timesteps, num_ddpm_timesteps, verbose=True):
+    if ddim_discr_method == 'uniform':
+        c = num_ddpm_timesteps // num_ddim_timesteps
+        ddim_timesteps = np.asarray(list(range(0, num_ddpm_timesteps, c)))
+    elif ddim_discr_method == 'quad':
+        ddim_timesteps = ((np.linspace(0, np.sqrt(num_ddpm_timesteps * .8), num_ddim_timesteps)) ** 2).astype(int)
+    else:
+        raise NotImplementedError(f'There is no ddim discretization method called "{ddim_discr_method}"')
+
+    # assert ddim_timesteps.shape[0] == num_ddim_timesteps
+    # add one to get the final alpha values right (the ones from first scale to data during sampling)
+    steps_out = ddim_timesteps + 1
+    if verbose:
+        print(f'Selected timesteps for ddim sampler: {steps_out}')
+    return steps_out
+
+
+def noise_like(shape, device, repeat=False):
+    repeat_noise = lambda: torch.randn((1, *shape[1:]), device=device).repeat(shape[0], *((1,) * (len(shape) - 1)))
+    noise = lambda: torch.randn(shape, device=device)
+    return repeat_noise() if repeat else noise()
+
+
+def timestep_embedding(device, timesteps, dim, max_period=10000, repeat_only=False):
+    """
+    Create sinusoidal timestep embeddings.
+    :param timesteps: a 1-D Tensor of N indices, one per batch element.
+                      These may be fractional.
+    :param dim: the dimension of the output.
+    :param max_period: controls the minimum frequency of the embeddings.
+    :return: an [N x dim] Tensor of positional embeddings.
+    """
+    half = dim // 2
+    freqs = torch.exp(
+        -math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32) / half
+    ).to(device=device)
+
+    args = timesteps[:, None].float() * freqs[None]
+
+    embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
+    if dim % 2:
+        embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1)
+    return embedding
+
+
+###### MAT and FcF #######
+
+
+def normalize_2nd_moment(x, dim=1, eps=1e-8):
+    return x * (x.square().mean(dim=dim, keepdim=True) + eps).rsqrt()
+
+
+class EasyDict(dict):
+    """Convenience class that behaves like a dict but allows access with the attribute syntax."""
+
+    def __getattr__(self, name: str) -> Any:
+        try:
+            return self[name]
+        except KeyError:
+            raise AttributeError(name)
+
+    def __setattr__(self, name: str, value: Any) -> None:
+        self[name] = value
+
+    def __delattr__(self, name: str) -> None:
+        del self[name]
+
+
+def _bias_act_ref(x, b=None, dim=1, act='linear', alpha=None, gain=None, clamp=None):
+    """Slow reference implementation of `bias_act()` using standard TensorFlow ops.
+    """
+    assert isinstance(x, torch.Tensor)
+    assert clamp is None or clamp >= 0
+    spec = activation_funcs[act]
+    alpha = float(alpha if alpha is not None else spec.def_alpha)
+    gain = float(gain if gain is not None else spec.def_gain)
+    clamp = float(clamp if clamp is not None else -1)
+
+    # Add bias.
+    if b is not None:
+        assert isinstance(b, torch.Tensor) and b.ndim == 1
+        assert 0 <= dim < x.ndim
+        assert b.shape[0] == x.shape[dim]
+        x = x + b.reshape([-1 if i == dim else 1 for i in range(x.ndim)])
+
+    # Evaluate activation function.
+    alpha = float(alpha)
+    x = spec.func(x, alpha=alpha)
+
+    # Scale by gain.
+    gain = float(gain)
+    if gain != 1:
+        x = x * gain
+
+    # Clamp.
+    if clamp >= 0:
+        x = x.clamp(-clamp, clamp)  # pylint: disable=invalid-unary-operand-type
+    return x
+
+
+def bias_act(x, b=None, dim=1, act='linear', alpha=None, gain=None, clamp=None, impl='ref'):
+    r"""Fused bias and activation function.
+
+    Adds bias `b` to activation tensor `x`, evaluates activation function `act`,
+    and scales the result by `gain`. Each of the steps is optional. In most cases,
+    the fused op is considerably more efficient than performing the same calculation
+    using standard PyTorch ops. It supports first and second order gradients,
+    but not third order gradients.
+
+    Args:
+        x:      Input activation tensor. Can be of any shape.
+        b:      Bias vector, or `None` to disable. Must be a 1D tensor of the same type
+                as `x`. The shape must be known, and it must match the dimension of `x`
+                corresponding to `dim`.
+        dim:    The dimension in `x` corresponding to the elements of `b`.
+                The value of `dim` is ignored if `b` is not specified.
+        act:    Name of the activation function to evaluate, or `"linear"` to disable.
+                Can be e.g. `"relu"`, `"lrelu"`, `"tanh"`, `"sigmoid"`, `"swish"`, etc.
+                See `activation_funcs` for a full list. `None` is not allowed.
+        alpha:  Shape parameter for the activation function, or `None` to use the default.
+        gain:   Scaling factor for the output tensor, or `None` to use default.
+                See `activation_funcs` for the default scaling of each activation function.
+                If unsure, consider specifying 1.
+        clamp:  Clamp the output values to `[-clamp, +clamp]`, or `None` to disable
+                the clamping (default).
+        impl:   Name of the implementation to use. Can be `"ref"` or `"cuda"` (default).
+
+    Returns:
+        Tensor of the same shape and datatype as `x`.
+    """
+    assert isinstance(x, torch.Tensor)
+    assert impl in ['ref', 'cuda']
+    return _bias_act_ref(x=x, b=b, dim=dim, act=act, alpha=alpha, gain=gain, clamp=clamp)
+
+
+def _get_filter_size(f):
+    if f is None:
+        return 1, 1
+
+    assert isinstance(f, torch.Tensor) and f.ndim in [1, 2]
+    fw = f.shape[-1]
+    fh = f.shape[0]
+
+    fw = int(fw)
+    fh = int(fh)
+    assert fw >= 1 and fh >= 1
+    return fw, fh
+
+
+def _get_weight_shape(w):
+    shape = [int(sz) for sz in w.shape]
+    return shape
+
+
+def _parse_scaling(scaling):
+    if isinstance(scaling, int):
+        scaling = [scaling, scaling]
+    assert isinstance(scaling, (list, tuple))
+    assert all(isinstance(x, int) for x in scaling)
+    sx, sy = scaling
+    assert sx >= 1 and sy >= 1
+    return sx, sy
+
+
+def _parse_padding(padding):
+    if isinstance(padding, int):
+        padding = [padding, padding]
+    assert isinstance(padding, (list, tuple))
+    assert all(isinstance(x, int) for x in padding)
+    if len(padding) == 2:
+        padx, pady = padding
+        padding = [padx, padx, pady, pady]
+    padx0, padx1, pady0, pady1 = padding
+    return padx0, padx1, pady0, pady1
+
+
+def setup_filter(f, device=torch.device('cpu'), normalize=True, flip_filter=False, gain=1, separable=None):
+    r"""Convenience function to setup 2D FIR filter for `upfirdn2d()`.
+
+    Args:
+        f:           Torch tensor, numpy array, or python list of the shape
+                     `[filter_height, filter_width]` (non-separable),
+                     `[filter_taps]` (separable),
+                     `[]` (impulse), or
+                     `None` (identity).
+        device:      Result device (default: cpu).
+        normalize:   Normalize the filter so that it retains the magnitude
+                     for constant input signal (DC)? (default: True).
+        flip_filter: Flip the filter? (default: False).
+        gain:        Overall scaling factor for signal magnitude (default: 1).
+        separable:   Return a separable filter? (default: select automatically).
+
+    Returns:
+        Float32 tensor of the shape
+        `[filter_height, filter_width]` (non-separable) or
+        `[filter_taps]` (separable).
+    """
+    # Validate.
+    if f is None:
+        f = 1
+    f = torch.as_tensor(f, dtype=torch.float32)
+    assert f.ndim in [0, 1, 2]
+    assert f.numel() > 0
+    if f.ndim == 0:
+        f = f[np.newaxis]
+
+    # Separable?
+    if separable is None:
+        separable = (f.ndim == 1 and f.numel() >= 8)
+    if f.ndim == 1 and not separable:
+        f = f.ger(f)
+    assert f.ndim == (1 if separable else 2)
+
+    # Apply normalize, flip, gain, and device.
+    if normalize:
+        f /= f.sum()
+    if flip_filter:
+        f = f.flip(list(range(f.ndim)))
+    f = f * (gain ** (f.ndim / 2))
+    f = f.to(device=device)
+    return f
+
+
+def _ntuple(n):
+    def parse(x):
+        if isinstance(x, collections.abc.Iterable):
+            return x
+        return tuple(repeat(x, n))
+
+    return parse
+
+
+to_2tuple = _ntuple(2)
+
+activation_funcs = {
+    'linear': EasyDict(func=lambda x, **_: x, def_alpha=0, def_gain=1, cuda_idx=1, ref='', has_2nd_grad=False),
+    'relu': EasyDict(func=lambda x, **_: torch.nn.functional.relu(x), def_alpha=0, def_gain=np.sqrt(2), cuda_idx=2,
+                     ref='y', has_2nd_grad=False),
+    'lrelu': EasyDict(func=lambda x, alpha, **_: torch.nn.functional.leaky_relu(x, alpha), def_alpha=0.2,
+                      def_gain=np.sqrt(2), cuda_idx=3, ref='y', has_2nd_grad=False),
+    'tanh': EasyDict(func=lambda x, **_: torch.tanh(x), def_alpha=0, def_gain=1, cuda_idx=4, ref='y',
+                     has_2nd_grad=True),
+    'sigmoid': EasyDict(func=lambda x, **_: torch.sigmoid(x), def_alpha=0, def_gain=1, cuda_idx=5, ref='y',
+                        has_2nd_grad=True),
+    'elu': EasyDict(func=lambda x, **_: torch.nn.functional.elu(x), def_alpha=0, def_gain=1, cuda_idx=6, ref='y',
+                    has_2nd_grad=True),
+    'selu': EasyDict(func=lambda x, **_: torch.nn.functional.selu(x), def_alpha=0, def_gain=1, cuda_idx=7, ref='y',
+                     has_2nd_grad=True),
+    'softplus': EasyDict(func=lambda x, **_: torch.nn.functional.softplus(x), def_alpha=0, def_gain=1, cuda_idx=8,
+                         ref='y', has_2nd_grad=True),
+    'swish': EasyDict(func=lambda x, **_: torch.sigmoid(x) * x, def_alpha=0, def_gain=np.sqrt(2), cuda_idx=9, ref='x',
+                      has_2nd_grad=True),
+}
+
+
+def upfirdn2d(x, f, up=1, down=1, padding=0, flip_filter=False, gain=1, impl='cuda'):
+    r"""Pad, upsample, filter, and downsample a batch of 2D images.
+
+    Performs the following sequence of operations for each channel:
+
+    1. Upsample the image by inserting N-1 zeros after each pixel (`up`).
+
+    2. Pad the image with the specified number of zeros on each side (`padding`).
+       Negative padding corresponds to cropping the image.
+
+    3. Convolve the image with the specified 2D FIR filter (`f`), shrinking it
+       so that the footprint of all output pixels lies within the input image.
+
+    4. Downsample the image by keeping every Nth pixel (`down`).
+
+    This sequence of operations bears close resemblance to scipy.signal.upfirdn().
+    The fused op is considerably more efficient than performing the same calculation
+    using standard PyTorch ops. It supports gradients of arbitrary order.
+
+    Args:
+        x:           Float32/float64/float16 input tensor of the shape
+                     `[batch_size, num_channels, in_height, in_width]`.
+        f:           Float32 FIR filter of the shape
+                     `[filter_height, filter_width]` (non-separable),
+                     `[filter_taps]` (separable), or
+                     `None` (identity).
+        up:          Integer upsampling factor. Can be a single int or a list/tuple
+                     `[x, y]` (default: 1).
+        down:        Integer downsampling factor. Can be a single int or a list/tuple
+                     `[x, y]` (default: 1).
+        padding:     Padding with respect to the upsampled image. Can be a single number
+                     or a list/tuple `[x, y]` or `[x_before, x_after, y_before, y_after]`
+                     (default: 0).
+        flip_filter: False = convolution, True = correlation (default: False).
+        gain:        Overall scaling factor for signal magnitude (default: 1).
+        impl:        Implementation to use. Can be `'ref'` or `'cuda'` (default: `'cuda'`).
+
+    Returns:
+        Tensor of the shape `[batch_size, num_channels, out_height, out_width]`.
+    """
+    # assert isinstance(x, torch.Tensor)
+    # assert impl in ['ref', 'cuda']
+    return _upfirdn2d_ref(x, f, up=up, down=down, padding=padding, flip_filter=flip_filter, gain=gain)
+
+
+def _upfirdn2d_ref(x, f, up=1, down=1, padding=0, flip_filter=False, gain=1):
+    """Slow reference implementation of `upfirdn2d()` using standard PyTorch ops.
+    """
+    # Validate arguments.
+    assert isinstance(x, torch.Tensor) and x.ndim == 4
+    if f is None:
+        f = torch.ones([1, 1], dtype=torch.float32, device=x.device)
+    assert isinstance(f, torch.Tensor) and f.ndim in [1, 2]
+    assert f.dtype == torch.float32 and not f.requires_grad
+    batch_size, num_channels, in_height, in_width = x.shape
+    # upx, upy = _parse_scaling(up)
+    # downx, downy = _parse_scaling(down)
+
+    upx, upy = up, up
+    downx, downy = down, down
+
+    # padx0, padx1, pady0, pady1 = _parse_padding(padding)
+    padx0, padx1, pady0, pady1 = padding[0], padding[1], padding[2], padding[3]
+
+    # Upsample by inserting zeros.
+    x = x.reshape([batch_size, num_channels, in_height, 1, in_width, 1])
+    x = torch.nn.functional.pad(x, [0, upx - 1, 0, 0, 0, upy - 1])
+    x = x.reshape([batch_size, num_channels, in_height * upy, in_width * upx])
+
+    # Pad or crop.
+    x = torch.nn.functional.pad(x, [max(padx0, 0), max(padx1, 0), max(pady0, 0), max(pady1, 0)])
+    x = x[:, :, max(-pady0, 0): x.shape[2] - max(-pady1, 0), max(-padx0, 0): x.shape[3] - max(-padx1, 0)]
+
+    # Setup filter.
+    f = f * (gain ** (f.ndim / 2))
+    f = f.to(x.dtype)
+    if not flip_filter:
+        f = f.flip(list(range(f.ndim)))
+
+    # Convolve with the filter.
+    f = f[np.newaxis, np.newaxis].repeat([num_channels, 1] + [1] * f.ndim)
+    if f.ndim == 4:
+        x = conv2d(input=x, weight=f, groups=num_channels)
+    else:
+        x = conv2d(input=x, weight=f.unsqueeze(2), groups=num_channels)
+        x = conv2d(input=x, weight=f.unsqueeze(3), groups=num_channels)
+
+    # Downsample by throwing away pixels.
+    x = x[:, :, ::downy, ::downx]
+    return x
+
+
+def downsample2d(x, f, down=2, padding=0, flip_filter=False, gain=1, impl='cuda'):
+    r"""Downsample a batch of 2D images using the given 2D FIR filter.
+
+    By default, the result is padded so that its shape is a fraction of the input.
+    User-specified padding is applied on top of that, with negative values
+    indicating cropping. Pixels outside the image are assumed to be zero.
+
+    Args:
+        x:           Float32/float64/float16 input tensor of the shape
+                     `[batch_size, num_channels, in_height, in_width]`.
+        f:           Float32 FIR filter of the shape
+                     `[filter_height, filter_width]` (non-separable),
+                     `[filter_taps]` (separable), or
+                     `None` (identity).
+        down:        Integer downsampling factor. Can be a single int or a list/tuple
+                     `[x, y]` (default: 1).
+        padding:     Padding with respect to the input. Can be a single number or a
+                     list/tuple `[x, y]` or `[x_before, x_after, y_before, y_after]`
+                     (default: 0).
+        flip_filter: False = convolution, True = correlation (default: False).
+        gain:        Overall scaling factor for signal magnitude (default: 1).
+        impl:        Implementation to use. Can be `'ref'` or `'cuda'` (default: `'cuda'`).
+
+    Returns:
+        Tensor of the shape `[batch_size, num_channels, out_height, out_width]`.
+    """
+    downx, downy = _parse_scaling(down)
+    # padx0, padx1, pady0, pady1 = _parse_padding(padding)
+    padx0, padx1, pady0, pady1 = padding, padding, padding, padding
+
+    fw, fh = _get_filter_size(f)
+    p = [
+        padx0 + (fw - downx + 1) // 2,
+        padx1 + (fw - downx) // 2,
+        pady0 + (fh - downy + 1) // 2,
+        pady1 + (fh - downy) // 2,
+    ]
+    return upfirdn2d(x, f, down=down, padding=p, flip_filter=flip_filter, gain=gain, impl=impl)
+
+
+def upsample2d(x, f, up=2, padding=0, flip_filter=False, gain=1, impl='cuda'):
+    r"""Upsample a batch of 2D images using the given 2D FIR filter.
+
+    By default, the result is padded so that its shape is a multiple of the input.
+    User-specified padding is applied on top of that, with negative values
+    indicating cropping. Pixels outside the image are assumed to be zero.
+
+    Args:
+        x:           Float32/float64/float16 input tensor of the shape
+                     `[batch_size, num_channels, in_height, in_width]`.
+        f:           Float32 FIR filter of the shape
+                     `[filter_height, filter_width]` (non-separable),
+                     `[filter_taps]` (separable), or
+                     `None` (identity).
+        up:          Integer upsampling factor. Can be a single int or a list/tuple
+                     `[x, y]` (default: 1).
+        padding:     Padding with respect to the output. Can be a single number or a
+                     list/tuple `[x, y]` or `[x_before, x_after, y_before, y_after]`
+                     (default: 0).
+        flip_filter: False = convolution, True = correlation (default: False).
+        gain:        Overall scaling factor for signal magnitude (default: 1).
+        impl:        Implementation to use. Can be `'ref'` or `'cuda'` (default: `'cuda'`).
+
+    Returns:
+        Tensor of the shape `[batch_size, num_channels, out_height, out_width]`.
+    """
+    upx, upy = _parse_scaling(up)
+    # upx, upy = up, up
+    padx0, padx1, pady0, pady1 = _parse_padding(padding)
+    # padx0, padx1, pady0, pady1 = padding, padding, padding, padding
+    fw, fh = _get_filter_size(f)
+    p = [
+        padx0 + (fw + upx - 1) // 2,
+        padx1 + (fw - upx) // 2,
+        pady0 + (fh + upy - 1) // 2,
+        pady1 + (fh - upy) // 2,
+    ]
+    return upfirdn2d(x, f, up=up, padding=p, flip_filter=flip_filter, gain=gain * upx * upy, impl=impl)
+
+
+class MinibatchStdLayer(torch.nn.Module):
+    def __init__(self, group_size, num_channels=1):
+        super().__init__()
+        self.group_size = group_size
+        self.num_channels = num_channels
+
+    def forward(self, x):
+        N, C, H, W = x.shape
+        G = torch.min(torch.as_tensor(self.group_size),
+                      torch.as_tensor(N)) if self.group_size is not None else N
+        F = self.num_channels
+        c = C // F
+
+        y = x.reshape(G, -1, F, c, H,
+                      W)  # [GnFcHW] Split minibatch N into n groups of size G, and channels C into F groups of size c.
+        y = y - y.mean(dim=0)  # [GnFcHW] Subtract mean over group.
+        y = y.square().mean(dim=0)  # [nFcHW]  Calc variance over group.
+        y = (y + 1e-8).sqrt()  # [nFcHW]  Calc stddev over group.
+        y = y.mean(dim=[2, 3, 4])  # [nF]     Take average over channels and pixels.
+        y = y.reshape(-1, F, 1, 1)  # [nF11]   Add missing dimensions.
+        y = y.repeat(G, 1, H, W)  # [NFHW]   Replicate over group and pixels.
+        x = torch.cat([x, y], dim=1)  # [NCHW]   Append to input as new channels.
+        return x
+
+
+class FullyConnectedLayer(torch.nn.Module):
+    def __init__(self,
+                 in_features,  # Number of input features.
+                 out_features,  # Number of output features.
+                 bias=True,  # Apply additive bias before the activation function?
+                 activation='linear',  # Activation function: 'relu', 'lrelu', etc.
+                 lr_multiplier=1,  # Learning rate multiplier.
+                 bias_init=0,  # Initial value for the additive bias.
+                 ):
+        super().__init__()
+        self.weight = torch.nn.Parameter(torch.randn([out_features, in_features]) / lr_multiplier)
+        self.bias = torch.nn.Parameter(torch.full([out_features], np.float32(bias_init))) if bias else None
+        self.activation = activation
+
+        self.weight_gain = lr_multiplier / np.sqrt(in_features)
+        self.bias_gain = lr_multiplier
+
+    def forward(self, x):
+        w = self.weight * self.weight_gain
+        b = self.bias
+        if b is not None and self.bias_gain != 1:
+            b = b * self.bias_gain
+
+        if self.activation == 'linear' and b is not None:
+            # out = torch.addmm(b.unsqueeze(0), x, w.t())
+            x = x.matmul(w.t())
+            out = x + b.reshape([-1 if i == x.ndim - 1 else 1 for i in range(x.ndim)])
+        else:
+            x = x.matmul(w.t())
+            out = bias_act(x, b, act=self.activation, dim=x.ndim - 1)
+        return out
+
+
+def _conv2d_wrapper(x, w, stride=1, padding=0, groups=1, transpose=False, flip_weight=True):
+    """Wrapper for the underlying `conv2d()` and `conv_transpose2d()` implementations.
+    """
+    out_channels, in_channels_per_group, kh, kw = _get_weight_shape(w)
+
+    # Flip weight if requested.
+    if not flip_weight:  # conv2d() actually performs correlation (flip_weight=True) not convolution (flip_weight=False).
+        w = w.flip([2, 3])
+
+    # Workaround performance pitfall in cuDNN 8.0.5, triggered when using
+    # 1x1 kernel + memory_format=channels_last + less than 64 channels.
+    if kw == 1 and kh == 1 and stride == 1 and padding in [0, [0, 0], (0, 0)] and not transpose:
+        if x.stride()[1] == 1 and min(out_channels, in_channels_per_group) < 64:
+            if out_channels <= 4 and groups == 1:
+                in_shape = x.shape
+                x = w.squeeze(3).squeeze(2) @ x.reshape([in_shape[0], in_channels_per_group, -1])
+                x = x.reshape([in_shape[0], out_channels, in_shape[2], in_shape[3]])
+            else:
+                x = x.to(memory_format=torch.contiguous_format)
+                w = w.to(memory_format=torch.contiguous_format)
+                x = conv2d(x, w, groups=groups)
+            return x.to(memory_format=torch.channels_last)
+
+    # Otherwise => execute using conv2d_gradfix.
+    op = conv_transpose2d if transpose else conv2d
+    return op(x, w, stride=stride, padding=padding, groups=groups)
+
+
+def conv2d_resample(x, w, f=None, up=1, down=1, padding=0, groups=1, flip_weight=True, flip_filter=False):
+    r"""2D convolution with optional up/downsampling.
+
+    Padding is performed only once at the beginning, not between the operations.
+
+    Args:
+        x:              Input tensor of shape
+                        `[batch_size, in_channels, in_height, in_width]`.
+        w:              Weight tensor of shape
+                        `[out_channels, in_channels//groups, kernel_height, kernel_width]`.
+        f:              Low-pass filter for up/downsampling. Must be prepared beforehand by
+                        calling setup_filter(). None = identity (default).
+        up:             Integer upsampling factor (default: 1).
+        down:           Integer downsampling factor (default: 1).
+        padding:        Padding with respect to the upsampled image. Can be a single number
+                        or a list/tuple `[x, y]` or `[x_before, x_after, y_before, y_after]`
+                        (default: 0).
+        groups:         Split input channels into N groups (default: 1).
+        flip_weight:    False = convolution, True = correlation (default: True).
+        flip_filter:    False = convolution, True = correlation (default: False).
+
+    Returns:
+        Tensor of the shape `[batch_size, num_channels, out_height, out_width]`.
+    """
+    # Validate arguments.
+    assert isinstance(x, torch.Tensor) and (x.ndim == 4)
+    assert isinstance(w, torch.Tensor) and (w.ndim == 4) and (w.dtype == x.dtype)
+    assert f is None or (isinstance(f, torch.Tensor) and f.ndim in [1, 2] and f.dtype == torch.float32)
+    assert isinstance(up, int) and (up >= 1)
+    assert isinstance(down, int) and (down >= 1)
+    # assert isinstance(groups, int) and (groups >= 1), f"!!!!!! groups: {groups} isinstance(groups, int)  {isinstance(groups, int)} {type(groups)}"
+    out_channels, in_channels_per_group, kh, kw = _get_weight_shape(w)
+    fw, fh = _get_filter_size(f)
+    # px0, px1, py0, py1 = _parse_padding(padding)
+    px0, px1, py0, py1 = padding, padding, padding, padding
+
+    # Adjust padding to account for up/downsampling.
+    if up > 1:
+        px0 += (fw + up - 1) // 2
+        px1 += (fw - up) // 2
+        py0 += (fh + up - 1) // 2
+        py1 += (fh - up) // 2
+    if down > 1:
+        px0 += (fw - down + 1) // 2
+        px1 += (fw - down) // 2
+        py0 += (fh - down + 1) // 2
+        py1 += (fh - down) // 2
+
+    # Fast path: 1x1 convolution with downsampling only => downsample first, then convolve.
+    if kw == 1 and kh == 1 and (down > 1 and up == 1):
+        x = upfirdn2d(x=x, f=f, down=down, padding=[px0, px1, py0, py1], flip_filter=flip_filter)
+        x = _conv2d_wrapper(x=x, w=w, groups=groups, flip_weight=flip_weight)
+        return x
+
+    # Fast path: 1x1 convolution with upsampling only => convolve first, then upsample.
+    if kw == 1 and kh == 1 and (up > 1 and down == 1):
+        x = _conv2d_wrapper(x=x, w=w, groups=groups, flip_weight=flip_weight)
+        x = upfirdn2d(x=x, f=f, up=up, padding=[px0, px1, py0, py1], gain=up ** 2, flip_filter=flip_filter)
+        return x
+
+    # Fast path: downsampling only => use strided convolution.
+    if down > 1 and up == 1:
+        x = upfirdn2d(x=x, f=f, padding=[px0, px1, py0, py1], flip_filter=flip_filter)
+        x = _conv2d_wrapper(x=x, w=w, stride=down, groups=groups, flip_weight=flip_weight)
+        return x
+
+    # Fast path: upsampling with optional downsampling => use transpose strided convolution.
+    if up > 1:
+        if groups == 1:
+            w = w.transpose(0, 1)
+        else:
+            w = w.reshape(groups, out_channels // groups, in_channels_per_group, kh, kw)
+            w = w.transpose(1, 2)
+            w = w.reshape(groups * in_channels_per_group, out_channels // groups, kh, kw)
+        px0 -= kw - 1
+        px1 -= kw - up
+        py0 -= kh - 1
+        py1 -= kh - up
+        pxt = max(min(-px0, -px1), 0)
+        pyt = max(min(-py0, -py1), 0)
+        x = _conv2d_wrapper(x=x, w=w, stride=up, padding=[pyt, pxt], groups=groups, transpose=True,
+                            flip_weight=(not flip_weight))
+        x = upfirdn2d(x=x, f=f, padding=[px0 + pxt, px1 + pxt, py0 + pyt, py1 + pyt], gain=up ** 2,
+                      flip_filter=flip_filter)
+        if down > 1:
+            x = upfirdn2d(x=x, f=f, down=down, flip_filter=flip_filter)
+        return x
+
+    # Fast path: no up/downsampling, padding supported by the underlying implementation => use plain conv2d.
+    if up == 1 and down == 1:
+        if px0 == px1 and py0 == py1 and px0 >= 0 and py0 >= 0:
+            return _conv2d_wrapper(x=x, w=w, padding=[py0, px0], groups=groups, flip_weight=flip_weight)
+
+    # Fallback: Generic reference implementation.
+    x = upfirdn2d(x=x, f=(f if up > 1 else None), up=up, padding=[px0, px1, py0, py1], gain=up ** 2,
+                  flip_filter=flip_filter)
+    x = _conv2d_wrapper(x=x, w=w, groups=groups, flip_weight=flip_weight)
+    if down > 1:
+        x = upfirdn2d(x=x, f=f, down=down, flip_filter=flip_filter)
+    return x
+
+
+class Conv2dLayer(torch.nn.Module):
+    def __init__(self,
+                 in_channels,  # Number of input channels.
+                 out_channels,  # Number of output channels.
+                 kernel_size,  # Width and height of the convolution kernel.
+                 bias=True,  # Apply additive bias before the activation function?
+                 activation='linear',  # Activation function: 'relu', 'lrelu', etc.
+                 up=1,  # Integer upsampling factor.
+                 down=1,  # Integer downsampling factor.
+                 resample_filter=[1, 3, 3, 1],  # Low-pass filter to apply when resampling activations.
+                 conv_clamp=None,  # Clamp the output to +-X, None = disable clamping.
+                 channels_last=False,  # Expect the input to have memory_format=channels_last?
+                 trainable=True,  # Update the weights of this layer during training?
+                 ):
+        super().__init__()
+        self.activation = activation
+        self.up = up
+        self.down = down
+        self.register_buffer('resample_filter', setup_filter(resample_filter))
+        self.conv_clamp = conv_clamp
+        self.padding = kernel_size // 2
+        self.weight_gain = 1 / np.sqrt(in_channels * (kernel_size ** 2))
+        self.act_gain = activation_funcs[activation].def_gain
+
+        memory_format = torch.channels_last if channels_last else torch.contiguous_format
+        weight = torch.randn([out_channels, in_channels, kernel_size, kernel_size]).to(memory_format=memory_format)
+        bias = torch.zeros([out_channels]) if bias else None
+        if trainable:
+            self.weight = torch.nn.Parameter(weight)
+            self.bias = torch.nn.Parameter(bias) if bias is not None else None
+        else:
+            self.register_buffer('weight', weight)
+            if bias is not None:
+                self.register_buffer('bias', bias)
+            else:
+                self.bias = None
+
+    def forward(self, x, gain=1):
+        w = self.weight * self.weight_gain
+        x = conv2d_resample(x=x, w=w, f=self.resample_filter, up=self.up, down=self.down,
+                            padding=self.padding)
+
+        act_gain = self.act_gain * gain
+        act_clamp = self.conv_clamp * gain if self.conv_clamp is not None else None
+        out = bias_act(x, self.bias, act=self.activation, gain=act_gain, clamp=act_clamp)
+        return out

lama_cleaner/model/zits.py

@@ -0,0 +1,427 @@
+import os
+import time
+
+import cv2
+import skimage
+import torch
+import torch.nn.functional as F
+
+from lama_cleaner.helper import get_cache_path_by_url, load_jit_model
+from lama_cleaner.schema import Config
+import numpy as np
+
+from lama_cleaner.model.base import InpaintModel
+
+ZITS_INPAINT_MODEL_URL = os.environ.get(
+    "ZITS_INPAINT_MODEL_URL",
+    "https://github.com/Sanster/models/releases/download/add_zits/zits-inpaint-0717.pt",
+)
+
+ZITS_EDGE_LINE_MODEL_URL = os.environ.get(
+    "ZITS_EDGE_LINE_MODEL_URL",
+    "https://github.com/Sanster/models/releases/download/add_zits/zits-edge-line-0717.pt",
+)
+
+ZITS_STRUCTURE_UPSAMPLE_MODEL_URL = os.environ.get(
+    "ZITS_STRUCTURE_UPSAMPLE_MODEL_URL",
+    "https://github.com/Sanster/models/releases/download/add_zits/zits-structure-upsample-0717.pt",
+)
+
+ZITS_WIRE_FRAME_MODEL_URL = os.environ.get(
+    "ZITS_WIRE_FRAME_MODEL_URL",
+    "https://github.com/Sanster/models/releases/download/add_zits/zits-wireframe-0717.pt",
+)
+
+
+def resize(img, height, width, center_crop=False):
+    imgh, imgw = img.shape[0:2]
+
+    if center_crop and imgh != imgw:
+        # center crop
+        side = np.minimum(imgh, imgw)
+        j = (imgh - side) // 2
+        i = (imgw - side) // 2
+        img = img[j : j + side, i : i + side, ...]
+
+    if imgh > height and imgw > width:
+        inter = cv2.INTER_AREA
+    else:
+        inter = cv2.INTER_LINEAR
+    img = cv2.resize(img, (height, width), interpolation=inter)
+
+    return img
+
+
+def to_tensor(img, scale=True, norm=False):
+    if img.ndim == 2:
+        img = img[:, :, np.newaxis]
+    c = img.shape[-1]
+
+    if scale:
+        img_t = torch.from_numpy(img).permute(2, 0, 1).float().div(255)
+    else:
+        img_t = torch.from_numpy(img).permute(2, 0, 1).float()
+
+    if norm:
+        mean = torch.tensor([0.5, 0.5, 0.5]).reshape(c, 1, 1)
+        std = torch.tensor([0.5, 0.5, 0.5]).reshape(c, 1, 1)
+        img_t = (img_t - mean) / std
+    return img_t
+
+
+def load_masked_position_encoding(mask):
+    ones_filter = np.ones((3, 3), dtype=np.float32)
+    d_filter1 = np.array([[1, 1, 0], [1, 1, 0], [0, 0, 0]], dtype=np.float32)
+    d_filter2 = np.array([[0, 0, 0], [1, 1, 0], [1, 1, 0]], dtype=np.float32)
+    d_filter3 = np.array([[0, 1, 1], [0, 1, 1], [0, 0, 0]], dtype=np.float32)
+    d_filter4 = np.array([[0, 0, 0], [0, 1, 1], [0, 1, 1]], dtype=np.float32)
+    str_size = 256
+    pos_num = 128
+
+    ori_mask = mask.copy()
+    ori_h, ori_w = ori_mask.shape[0:2]
+    ori_mask = ori_mask / 255
+    mask = cv2.resize(mask, (str_size, str_size), interpolation=cv2.INTER_AREA)
+    mask[mask > 0] = 255
+    h, w = mask.shape[0:2]
+    mask3 = mask.copy()
+    mask3 = 1.0 - (mask3 / 255.0)
+    pos = np.zeros((h, w), dtype=np.int32)
+    direct = np.zeros((h, w, 4), dtype=np.int32)
+    i = 0
+    while np.sum(1 - mask3) > 0:
+        i += 1
+        mask3_ = cv2.filter2D(mask3, -1, ones_filter)
+        mask3_[mask3_ > 0] = 1
+        sub_mask = mask3_ - mask3
+        pos[sub_mask == 1] = i
+
+        m = cv2.filter2D(mask3, -1, d_filter1)
+        m[m > 0] = 1
+        m = m - mask3
+        direct[m == 1, 0] = 1
+
+        m = cv2.filter2D(mask3, -1, d_filter2)
+        m[m > 0] = 1
+        m = m - mask3
+        direct[m == 1, 1] = 1
+
+        m = cv2.filter2D(mask3, -1, d_filter3)
+        m[m > 0] = 1
+        m = m - mask3
+        direct[m == 1, 2] = 1
+
+        m = cv2.filter2D(mask3, -1, d_filter4)
+        m[m > 0] = 1
+        m = m - mask3
+        direct[m == 1, 3] = 1
+
+        mask3 = mask3_
+
+    abs_pos = pos.copy()
+    rel_pos = pos / (str_size / 2)  # to 0~1 maybe larger than 1
+    rel_pos = (rel_pos * pos_num).astype(np.int32)
+    rel_pos = np.clip(rel_pos, 0, pos_num - 1)
+
+    if ori_w != w or ori_h != h:
+        rel_pos = cv2.resize(rel_pos, (ori_w, ori_h), interpolation=cv2.INTER_NEAREST)
+        rel_pos[ori_mask == 0] = 0
+        direct = cv2.resize(direct, (ori_w, ori_h), interpolation=cv2.INTER_NEAREST)
+        direct[ori_mask == 0, :] = 0
+
+    return rel_pos, abs_pos, direct
+
+
+def load_image(img, mask, device, sigma256=3.0):
+    """
+    Args:
+        img: [H, W, C] RGB
+        mask: [H, W] 255 为 masks 区域
+        sigma256:
+
+    Returns:
+
+    """
+    h, w, _ = img.shape
+    imgh, imgw = img.shape[0:2]
+    img_256 = resize(img, 256, 256)
+
+    mask = (mask > 127).astype(np.uint8) * 255
+    mask_256 = cv2.resize(mask, (256, 256), interpolation=cv2.INTER_AREA)
+    mask_256[mask_256 > 0] = 255
+
+    mask_512 = cv2.resize(mask, (512, 512), interpolation=cv2.INTER_AREA)
+    mask_512[mask_512 > 0] = 255
+
+    # original skimage implemention
+    # https://scikit-image.org/docs/stable/api/skimage.feature.html#skimage.feature.canny
+    # low_threshold: Lower bound for hysteresis thresholding (linking edges). If None, low_threshold is set to 10% of dtype’s max.
+    # high_threshold: Upper bound for hysteresis thresholding (linking edges). If None, high_threshold is set to 20% of dtype’s max.
+    gray_256 = skimage.color.rgb2gray(img_256)
+    edge_256 = skimage.feature.canny(gray_256, sigma=sigma256, mask=None).astype(float)
+    # cv2.imwrite("skimage_gray.jpg", (_gray_256*255).astype(np.uint8))
+    # cv2.imwrite("skimage_edge.jpg", (_edge_256*255).astype(np.uint8))
+
+    # gray_256 = cv2.cvtColor(img_256, cv2.COLOR_RGB2GRAY)
+    # gray_256_blured = cv2.GaussianBlur(gray_256, ksize=(3,3), sigmaX=sigma256, sigmaY=sigma256)
+    # edge_256 = cv2.Canny(gray_256_blured, threshold1=int(255*0.1), threshold2=int(255*0.2))
+    # cv2.imwrite("edge.jpg", edge_256)
+
+    # line
+    img_512 = resize(img, 512, 512)
+
+    rel_pos, abs_pos, direct = load_masked_position_encoding(mask)
+
+    batch = dict()
+    batch["images"] = to_tensor(img.copy()).unsqueeze(0).to(device)
+    batch["img_256"] = to_tensor(img_256, norm=True).unsqueeze(0).to(device)
+    batch["masks"] = to_tensor(mask).unsqueeze(0).to(device)
+    batch["mask_256"] = to_tensor(mask_256).unsqueeze(0).to(device)
+    batch["mask_512"] = to_tensor(mask_512).unsqueeze(0).to(device)
+    batch["edge_256"] = to_tensor(edge_256, scale=False).unsqueeze(0).to(device)
+    batch["img_512"] = to_tensor(img_512).unsqueeze(0).to(device)
+    batch["rel_pos"] = torch.LongTensor(rel_pos).unsqueeze(0).to(device)
+    batch["abs_pos"] = torch.LongTensor(abs_pos).unsqueeze(0).to(device)
+    batch["direct"] = torch.LongTensor(direct).unsqueeze(0).to(device)
+    batch["h"] = imgh
+    batch["w"] = imgw
+
+    return batch
+
+
+def to_device(data, device):
+    if isinstance(data, torch.Tensor):
+        return data.to(device)
+    if isinstance(data, dict):
+        for key in data:
+            if isinstance(data[key], torch.Tensor):
+                data[key] = data[key].to(device)
+        return data
+    if isinstance(data, list):
+        return [to_device(d, device) for d in data]
+
+
+class ZITS(InpaintModel):
+    min_size = 256
+    pad_mod = 32
+    pad_to_square = True
+
+    def __init__(self, device, **kwargs):
+        """
+
+        Args:
+            device:
+        """
+        super().__init__(device)
+        self.device = device
+        self.sample_edge_line_iterations = 1
+
+    def init_model(self, device, **kwargs):
+        self.wireframe = load_jit_model(ZITS_WIRE_FRAME_MODEL_URL, device)
+        self.edge_line = load_jit_model(ZITS_EDGE_LINE_MODEL_URL, device)
+        self.structure_upsample = load_jit_model(
+            ZITS_STRUCTURE_UPSAMPLE_MODEL_URL, device
+        )
+        self.inpaint = load_jit_model(ZITS_INPAINT_MODEL_URL, device)
+
+    @staticmethod
+    def is_downloaded() -> bool:
+        model_paths = [
+            get_cache_path_by_url(ZITS_WIRE_FRAME_MODEL_URL),
+            get_cache_path_by_url(ZITS_EDGE_LINE_MODEL_URL),
+            get_cache_path_by_url(ZITS_STRUCTURE_UPSAMPLE_MODEL_URL),
+            get_cache_path_by_url(ZITS_INPAINT_MODEL_URL),
+        ]
+        return all([os.path.exists(it) for it in model_paths])
+
+    def wireframe_edge_and_line(self, items, enable: bool):
+        # 最终向 items 中添加 edge 和 line key
+        if not enable:
+            items["edge"] = torch.zeros_like(items["masks"])
+            items["line"] = torch.zeros_like(items["masks"])
+            return
+
+        start = time.time()
+        try:
+            line_256 = self.wireframe_forward(
+                items["img_512"],
+                h=256,
+                w=256,
+                masks=items["mask_512"],
+                mask_th=0.85,
+            )
+        except:
+            line_256 = torch.zeros_like(items["mask_256"])
+
+        print(f"wireframe_forward time: {(time.time() - start) * 1000:.2f}ms")
+
+        # np_line = (line[0][0].numpy() * 255).astype(np.uint8)
+        # cv2.imwrite("line.jpg", np_line)
+
+        start = time.time()
+        edge_pred, line_pred = self.sample_edge_line_logits(
+            context=[items["img_256"], items["edge_256"], line_256],
+            mask=items["mask_256"].clone(),
+            iterations=self.sample_edge_line_iterations,
+            add_v=0.05,
+            mul_v=4,
+        )
+        print(f"sample_edge_line_logits time: {(time.time() - start) * 1000:.2f}ms")
+
+        # np_edge_pred = (edge_pred[0][0].numpy() * 255).astype(np.uint8)
+        # cv2.imwrite("edge_pred.jpg", np_edge_pred)
+        # np_line_pred = (line_pred[0][0].numpy() * 255).astype(np.uint8)
+        # cv2.imwrite("line_pred.jpg", np_line_pred)
+        # exit()
+
+        input_size = min(items["h"], items["w"])
+        if input_size != 256 and input_size > 256:
+            while edge_pred.shape[2] < input_size:
+                edge_pred = self.structure_upsample(edge_pred)
+                edge_pred = torch.sigmoid((edge_pred + 2) * 2)
+
+                line_pred = self.structure_upsample(line_pred)
+                line_pred = torch.sigmoid((line_pred + 2) * 2)
+
+            edge_pred = F.interpolate(
+                edge_pred,
+                size=(input_size, input_size),
+                mode="bilinear",
+                align_corners=False,
+            )
+            line_pred = F.interpolate(
+                line_pred,
+                size=(input_size, input_size),
+                mode="bilinear",
+                align_corners=False,
+            )
+
+        # np_edge_pred = (edge_pred[0][0].numpy() * 255).astype(np.uint8)
+        # cv2.imwrite("edge_pred_upsample.jpg", np_edge_pred)
+        # np_line_pred = (line_pred[0][0].numpy() * 255).astype(np.uint8)
+        # cv2.imwrite("line_pred_upsample.jpg", np_line_pred)
+        # exit()
+
+        items["edge"] = edge_pred.detach()
+        items["line"] = line_pred.detach()
+
+    @torch.no_grad()
+    def forward(self, image, mask, config: Config):
+        """Input images and output images have same size
+        images: [H, W, C] RGB
+        masks: [H, W]
+        return: BGR IMAGE
+        """
+        mask = mask[:, :, 0]
+        items = load_image(image, mask, device=self.device)
+
+        self.wireframe_edge_and_line(items, config.zits_wireframe)
+
+        inpainted_image = self.inpaint(
+            items["images"],
+            items["masks"],
+            items["edge"],
+            items["line"],
+            items["rel_pos"],
+            items["direct"],
+        )
+
+        inpainted_image = inpainted_image * 255.0
+        inpainted_image = (
+            inpainted_image.cpu().permute(0, 2, 3, 1)[0].numpy().astype(np.uint8)
+        )
+        inpainted_image = inpainted_image[:, :, ::-1]
+
+        # cv2.imwrite("inpainted.jpg", inpainted_image)
+        # exit()
+
+        return inpainted_image
+
+    def wireframe_forward(self, images, h, w, masks, mask_th=0.925):
+        lcnn_mean = torch.tensor([109.730, 103.832, 98.681]).reshape(1, 3, 1, 1)
+        lcnn_std = torch.tensor([22.275, 22.124, 23.229]).reshape(1, 3, 1, 1)
+        images = images * 255.0
+        # the masks value of lcnn is 127.5
+        masked_images = images * (1 - masks) + torch.ones_like(images) * masks * 127.5
+        masked_images = (masked_images - lcnn_mean) / lcnn_std
+
+        def to_int(x):
+            return tuple(map(int, x))
+
+        lines_tensor = []
+        lmap = np.zeros((h, w))
+
+        output_masked = self.wireframe(masked_images)
+
+        output_masked = to_device(output_masked, "cpu")
+        if output_masked["num_proposals"] == 0:
+            lines_masked = []
+            scores_masked = []
+        else:
+            lines_masked = output_masked["lines_pred"].numpy()
+            lines_masked = [
+                [line[1] * h, line[0] * w, line[3] * h, line[2] * w]
+                for line in lines_masked
+            ]
+            scores_masked = output_masked["lines_score"].numpy()
+
+        for line, score in zip(lines_masked, scores_masked):
+            if score > mask_th:
+                rr, cc, value = skimage.draw.line_aa(
+                    *to_int(line[0:2]), *to_int(line[2:4])
+                )
+                lmap[rr, cc] = np.maximum(lmap[rr, cc], value)
+
+        lmap = np.clip(lmap * 255, 0, 255).astype(np.uint8)
+        lines_tensor.append(to_tensor(lmap).unsqueeze(0))
+
+        lines_tensor = torch.cat(lines_tensor, dim=0)
+        return lines_tensor.detach().to(self.device)
+
+    def sample_edge_line_logits(
+        self, context, mask=None, iterations=1, add_v=0, mul_v=4
+    ):
+        [img, edge, line] = context
+
+        img = img * (1 - mask)
+        edge = edge * (1 - mask)
+        line = line * (1 - mask)
+
+        for i in range(iterations):
+            edge_logits, line_logits = self.edge_line(img, edge, line, masks=mask)
+
+            edge_pred = torch.sigmoid(edge_logits)
+            line_pred = torch.sigmoid((line_logits + add_v) * mul_v)
+            edge = edge + edge_pred * mask
+            edge[edge >= 0.25] = 1
+            edge[edge < 0.25] = 0
+            line = line + line_pred * mask
+
+            b, _, h, w = edge_pred.shape
+            edge_pred = edge_pred.reshape(b, -1, 1)
+            line_pred = line_pred.reshape(b, -1, 1)
+            mask = mask.reshape(b, -1)
+
+            edge_probs = torch.cat([1 - edge_pred, edge_pred], dim=-1)
+            line_probs = torch.cat([1 - line_pred, line_pred], dim=-1)
+            edge_probs[:, :, 1] += 0.5
+            line_probs[:, :, 1] += 0.5
+            edge_max_probs = edge_probs.max(dim=-1)[0] + (1 - mask) * (-100)
+            line_max_probs = line_probs.max(dim=-1)[0] + (1 - mask) * (-100)
+
+            indices = torch.sort(
+                edge_max_probs + line_max_probs, dim=-1, descending=True
+            )[1]
+
+            for ii in range(b):
+                keep = int((i + 1) / iterations * torch.sum(mask[ii, ...]))
+
+                assert torch.sum(mask[ii][indices[ii, :keep]]) == keep, "Error!!!"
+                mask[ii][indices[ii, :keep]] = 0
+
+            mask = mask.reshape(b, 1, h, w)
+            edge = edge * (1 - mask)
+            line = line * (1 - mask)
+
+        edge, line = edge.to(torch.float32), line.to(torch.float32)
+        return edge, line

lama_cleaner/model_manager.py

@@ -0,0 +1,43 @@
+from lama_cleaner.model.fcf import FcF
+from lama_cleaner.model.lama import LaMa
+from lama_cleaner.model.ldm import LDM
+from lama_cleaner.model.mat import MAT
+from lama_cleaner.model.sd import SD14
+from lama_cleaner.model.zits import ZITS
+from lama_cleaner.model.opencv2 import OpenCV2
+from lama_cleaner.schema import Config
+
+models = {"lama": LaMa, "ldm": LDM, "zits": ZITS, "mat": MAT, "fcf": FcF, "sd1.4": SD14, "cv2": OpenCV2}
+
+
+class ModelManager:
+    def __init__(self, name: str, device, **kwargs):
+        self.name = name
+        self.device = device
+        self.kwargs = kwargs
+        self.model = self.init_model(name, device, **kwargs)
+
+    def init_model(self, name: str, device, **kwargs):
+        if name in models:
+            model = models[name](device, **kwargs)
+        else:
+            raise NotImplementedError(f"Not supported model: {name}")
+        return model
+
+    def is_downloaded(self, name: str) -> bool:
+        if name in models:
+            return models[name].is_downloaded()
+        else:
+            raise NotImplementedError(f"Not supported model: {name}")
+
+    def __call__(self, image, mask, config: Config):
+        return self.model(image, mask, config)
+
+    def switch(self, new_name: str):
+        if new_name == self.name:
+            return
+        try:
+            self.model = self.init_model(new_name, self.device, **self.kwargs)
+            self.name = new_name
+        except NotImplementedError as e:
+            raise e

lama_cleaner/schema.py

@@ -0,0 +1,50 @@
+from enum import Enum
+
+from pydantic import BaseModel
+
+
+class HDStrategy(str, Enum):
+    ORIGINAL = "Original"
+    RESIZE = "Resize"
+    CROP = "Crop"
+
+
+class LDMSampler(str, Enum):
+    ddim = "ddim"
+    plms = "plms"
+
+
+class SDSampler(str, Enum):
+    ddim = "ddim"
+    pndm = "pndm"
+
+
+class Config(BaseModel):
+    ldm_steps: int
+    ldm_sampler: str = LDMSampler.plms
+    zits_wireframe: bool = True
+    hd_strategy: str
+    hd_strategy_crop_margin: int
+    hd_strategy_crop_trigger_size: int
+    hd_strategy_resize_limit: int
+
+    prompt: str = ""
+    # 始终是在原图尺度上的值
+    use_croper: bool = False
+    croper_x: int = None
+    croper_y: int = None
+    croper_height: int = None
+    croper_width: int = None
+
+    # sd
+    sd_mask_blur: int = 0
+    sd_strength: float = 0.75
+    sd_steps: int = 50
+    sd_guidance_scale: float = 7.5
+    sd_sampler: str = SDSampler.ddim
+    # -1 mean random seed
+    sd_seed: int = 42
+
+    # cv2
+    cv2_flag: str = 'INPAINT_NS'
+    cv2_radius: int = 4

lama_cleaner/settings.py

@@ -0,0 +1,124 @@
+"""
+Django settings for lama_cleaner project.
+
+Generated by 'django-admin startproject' using Django 4.1.2.
+
+For more information on this file, see
+https://docs.djangoproject.com/en/4.1/topics/settings/
+
+For the full list of settings and their values, see
+https://docs.djangoproject.com/en/4.1/ref/settings/
+"""
+
+from pathlib import Path
+
+# Build paths inside the project like this: BASE_DIR / 'subdir'.
+BASE_DIR = Path(__file__).resolve().parent.parent
+
+
+# Quick-start development settings - unsuitable for production
+# See https://docs.djangoproject.com/en/4.1/howto/deployment/checklist/
+
+# SECURITY WARNING: keep the secret key used in production secret!
+SECRET_KEY = 'django-insecure-=x2n@zasb2nkq$)frp(&h*tsozyka+jb5(&3^7@u5@ven@-sdu'
+
+# SECURITY WARNING: don't run with debug turned on in production!
+DEBUG = True
+
+ALLOWED_HOSTS = []
+
+
+# Application definition
+
+INSTALLED_APPS = [
+    'django.contrib.admin',
+    'django.contrib.auth',
+    'django.contrib.contenttypes',
+    'django.contrib.sessions',
+    'django.contrib.messages',
+    'django.contrib.staticfiles',
+	'inpainting',
+]
+
+MIDDLEWARE = [
+    'django.middleware.security.SecurityMiddleware',
+    'django.contrib.sessions.middleware.SessionMiddleware',
+    'django.middleware.common.CommonMiddleware',
+    'django.middleware.csrf.CsrfViewMiddleware',
+    'django.contrib.auth.middleware.AuthenticationMiddleware',
+    'django.contrib.messages.middleware.MessageMiddleware',
+    'django.middleware.clickjacking.XFrameOptionsMiddleware',
+]
+
+ROOT_URLCONF = 'lama_cleaner.urls'
+
+TEMPLATES = [
+    {
+        'BACKEND': 'django.template.backends.django.DjangoTemplates',
+        'DIRS': [],
+        'APP_DIRS': True,
+        'OPTIONS': {
+            'context_processors': [
+                'django.template.context_processors.debug',
+                'django.template.context_processors.request',
+                'django.contrib.auth.context_processors.auth',
+                'django.contrib.messages.context_processors.messages',
+            ],
+        },
+    },
+]
+
+WSGI_APPLICATION = 'lama_cleaner.wsgi.application'
+
+
+# Database
+# https://docs.djangoproject.com/en/4.1/ref/settings/#databases
+
+DATABASES = {
+    'default': {
+        'ENGINE': 'django.db.backends.sqlite3',
+        'NAME': BASE_DIR / 'db.sqlite3',
+    }
+}
+
+
+# Password validation
+# https://docs.djangoproject.com/en/4.1/ref/settings/#auth-password-validators
+
+AUTH_PASSWORD_VALIDATORS = [
+    {
+        'NAME': 'django.contrib.auth.password_validation.UserAttributeSimilarityValidator',
+    },
+    {
+        'NAME': 'django.contrib.auth.password_validation.MinimumLengthValidator',
+    },
+    {
+        'NAME': 'django.contrib.auth.password_validation.CommonPasswordValidator',
+    },
+    {
+        'NAME': 'django.contrib.auth.password_validation.NumericPasswordValidator',
+    },
+]
+
+
+# Internationalization
+# https://docs.djangoproject.com/en/4.1/topics/i18n/
+
+LANGUAGE_CODE = 'en-us'
+
+TIME_ZONE = 'UTC'
+
+USE_I18N = True
+
+USE_TZ = True
+
+
+# Static files (CSS, JavaScript, Images)
+# https://docs.djangoproject.com/en/4.1/howto/static-files/
+
+STATIC_URL = 'static/'
+
+# Default primary key field type
+# https://docs.djangoproject.com/en/4.1/ref/settings/#default-auto-field
+
+DEFAULT_AUTO_FIELD = 'django.db.models.BigAutoField'

lama_cleaner/urls.py

@@ -0,0 +1,22 @@
+"""lama_cleaner URL Configuration
+
+The `urlpatterns` list routes URLs to views. For more information please see:
+    https://docs.djangoproject.com/en/4.1/topics/http/urls/
+Examples:
+Function views
+    1. Add an import:  from my_app import views
+    2. Add a URL to urlpatterns:  path('', views.home, name='home')
+Class-based views
+    1. Add an import:  from other_app.views import Home
+    2. Add a URL to urlpatterns:  path('', Home.as_view(), name='home')
+Including another URLconf
+    1. Import the include() function: from django.urls import include, path
+    2. Add a URL to urlpatterns:  path('blog/', include('blog.urls'))
+"""
+from django.contrib import admin
+from django.urls import path,include
+
+urlpatterns = [
+    path('admin/', admin.site.urls),
+    path('inpainting/',include('inpainting.urls')),
+]

lama_cleaner/wsgi.py

@@ -0,0 +1,16 @@
+"""
+WSGI config for lama_cleaner project.
+
+It exposes the WSGI callable as a module-level variable named ``application``.
+
+For more information on this file, see
+https://docs.djangoproject.com/en/4.1/howto/deployment/wsgi/
+"""
+
+import os
+
+from django.core.wsgi import get_wsgi_application
+
+os.environ.setdefault('DJANGO_SETTINGS_MODULE', 'lama_cleaner.settings')
+
+application = get_wsgi_application()

requirements.txt

@@ -0,0 +1,12 @@
+opencv-python==4.6.0.66
+pytest==7.1.3
+torch==2.2.0
+pydantic==1.10.2
+loguru==0.6.0
+tqdm==4.64.1
+Pillow==9.2.0
+diffusers==0.4.2
+transformers
+scikit-image==0.19.3
+gradio
+timm