diff --git a/LICENSE b/LICENSE
new file mode 100644
index 0000000000000000000000000000000000000000..46f3b49d54a23bc2474776988c3596595c1e5f3e
--- /dev/null
+++ b/LICENSE
@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2021 Du Ang
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.
diff --git a/data/dataset.py b/data/dataset.py
new file mode 100644
index 0000000000000000000000000000000000000000..c9e43a9cd2d0ac3aa23f442b0800454f5c2cdae8
--- /dev/null
+++ b/data/dataset.py
@@ -0,0 +1,75 @@
+import sys
+import torch.utils.data as data
+from os import listdir
+from utils.tools import default_loader, is_image_file, normalize
+import os
+
+import torchvision.transforms as transforms
+
+
+class Dataset(data.Dataset):
+ def __init__(self, data_path, image_shape, with_subfolder=False, random_crop=True, return_name=False):
+ super(Dataset, self).__init__()
+ if with_subfolder:
+ self.samples = self._find_samples_in_subfolders(data_path)
+ else:
+ self.samples = [x for x in listdir(data_path) if is_image_file(x)]
+ self.data_path = data_path
+ self.image_shape = image_shape[:-1]
+ self.random_crop = random_crop
+ self.return_name = return_name
+
+ def __getitem__(self, index):
+ path = os.path.join(self.data_path, self.samples[index])
+ img = default_loader(path)
+
+ if self.random_crop:
+ imgw, imgh = img.size
+ if imgh < self.image_shape[0] or imgw < self.image_shape[1]:
+ img = transforms.Resize(min(self.image_shape))(img)
+ img = transforms.RandomCrop(self.image_shape)(img)
+ else:
+ img = transforms.Resize(self.image_shape)(img)
+ img = transforms.RandomCrop(self.image_shape)(img)
+
+ img = transforms.ToTensor()(img) # turn the image to a tensor
+ img = normalize(img)
+
+ if self.return_name:
+ return self.samples[index], img
+ else:
+ return img
+
+ def _find_samples_in_subfolders(self, dir):
+ """
+ Finds the class folders in a dataset.
+ Args:
+ dir (string): Root directory path.
+ Returns:
+ tuple: (classes, class_to_idx) where classes are relative to (dir), and class_to_idx is a dictionary.
+ Ensures:
+ No class is a subdirectory of another.
+ """
+ if sys.version_info >= (3, 5):
+ # Faster and available in Python 3.5 and above
+ classes = [d.name for d in os.scandir(dir) if d.is_dir()]
+ else:
+ classes = [d for d in os.listdir(dir) if os.path.isdir(os.path.join(dir, d))]
+ classes.sort()
+ class_to_idx = {classes[i]: i for i in range(len(classes))}
+ samples = []
+ for target in sorted(class_to_idx.keys()):
+ d = os.path.join(dir, target)
+ if not os.path.isdir(d):
+ continue
+ for root, _, fnames in sorted(os.walk(d)):
+ for fname in sorted(fnames):
+ if is_image_file(fname):
+ path = os.path.join(root, fname)
+ # item = (path, class_to_idx[target])
+ # samples.append(item)
+ samples.append(path)
+ return samples
+
+ def __len__(self):
+ return len(self.samples)
diff --git a/full-stack-server.py b/full-stack-server.py
new file mode 100644
index 0000000000000000000000000000000000000000..95754a6744815152718e37d1a379921ef410e118
--- /dev/null
+++ b/full-stack-server.py
@@ -0,0 +1,231 @@
+import os
+import base64
+import io
+import uuid
+from ultralytics import YOLO
+import cv2
+import torch
+import numpy as np
+from PIL import Image
+from torchvision import transforms
+import imageio.v2 as imageio
+from trainer import Trainer
+from utils.tools import get_config
+import torch.nn.functional as F
+from iopaint.single_processing import batch_inpaint
+from pathlib import Path
+from flask import Flask, request, jsonify,render_template
+from flask_cors import CORS
+
+app = Flask(__name__)
+CORS(app)
+
+# set current working directory cache instead of default
+os.environ["TORCH_HOME"] = "./pretrained-model"
+os.environ["HUGGINGFACE_HUB_CACHE"] = "./pretrained-model"
+
+
+def resize_image(input_image_base64, width=640, height=640):
+ """Resizes an image from base64 data and returns the resized image as bytes."""
+ try:
+ # Decode base64 string to bytes
+ input_image_data = base64.b64decode(input_image_base64)
+ # Convert bytes to NumPy array
+ img = np.frombuffer(input_image_data, dtype=np.uint8)
+ # Decode NumPy array as an image
+ img = cv2.imdecode(img, cv2.IMREAD_COLOR)
+
+ # Resize while maintaining the aspect ratio
+ shape = img.shape[:2] # current shape [height, width]
+ new_shape = (width, height) # the shape to resize to
+
+ # Scale ratio (new / old)
+ r = min(new_shape[0] / shape[0], new_shape[1] / shape[1])
+ ratio = r, r # width, height ratios
+ new_unpad = int(round(shape[1] * r)), int(round(shape[0] * r))
+
+ # Resize the image
+ im = cv2.resize(img, new_unpad, interpolation=cv2.INTER_LINEAR)
+
+ # Pad the image
+ color = (114, 114, 114) # color used for padding
+ dw, dh = new_shape[1] - new_unpad[0], new_shape[0] - new_unpad[1] # wh padding
+ # divide padding into 2 sides
+ dw /= 2
+ dh /= 2
+ # compute padding on all corners
+ top, bottom = int(round(dh - 0.1)), int(round(dh + 0.1))
+ left, right = int(round(dw - 0.1)), int(round(dw + 0.1))
+ im = cv2.copyMakeBorder(im, top, bottom, left, right, cv2.BORDER_CONSTANT, value=color) # add border
+
+ # Convert the resized and padded image to bytes
+ resized_image_bytes = cv2.imencode('.png', im)[1].tobytes()
+ return resized_image_bytes
+
+ except Exception as e:
+ print(f"Error resizing image: {e}")
+ return None # Or handle differently as needed
+
+
+def load_weights(path, device):
+ model_weights = torch.load(path)
+ return {
+ k: v.to(device)
+ for k, v in model_weights.items()
+ }
+
+
+# Function to convert image to base64
+def convert_image_to_base64(image):
+ # Convert image to bytes
+ _, buffer = cv2.imencode('.png', image)
+ # Convert bytes to base64
+ image_base64 = base64.b64encode(buffer).decode('utf-8')
+ return image_base64
+
+
+def convert_to_base64(image):
+ # Read the image file as binary data
+ image_data = image.read()
+ # Encode the binary data as base64
+ base64_encoded = base64.b64encode(image_data).decode('utf-8')
+ return base64_encoded
+
+
+@app.route('/')
+def index():
+ return render_template('index.html')
+
+
+@app.route('/process_images', methods=['POST'])
+def process_images():
+ # Static paths
+ config_path = Path('configs/config.yaml')
+ model_path = Path('pretrained-model/torch_model.p')
+
+ # Check if the request contains files
+ if 'input_image' not in request.files or 'append_image' not in request.files:
+ return jsonify({'error': 'No files found'}), 419
+
+ # Get the objectName from the request or use default "chair" if not provided
+ default_class = request.form.get('objectName', 'chair')
+
+ # Convert the images to base64
+ try:
+ input_base64 = convert_to_base64(request.files['input_image'])
+ append_base64 = convert_to_base64(request.files['append_image'])
+ except Exception as e:
+ return jsonify({'error': 'Failed to read files'}), 419
+
+ # Resize input image and get base64 data of resized image
+ input_resized_image_bytes = resize_image(input_base64)
+
+ # Convert resized image bytes to base64
+ input_resized_base64 = base64.b64encode(input_resized_image_bytes).decode('utf-8')
+
+ # Decode the resized image from base64 data directly
+ img = cv2.imdecode(np.frombuffer(input_resized_image_bytes, np.uint8), cv2.IMREAD_COLOR)
+
+ if img is None:
+ return jsonify({'error': 'Failed to decode resized image'}), 419
+
+ H, W, _ = img.shape
+ x_point = 0
+ y_point = 0
+ width = 1
+ height = 1
+
+ # Load a model
+ model = YOLO('pretrained-model/yolov8m-seg.pt') # pretrained YOLOv8m-seg model
+
+ # Run batched inference on a list of images
+ results = model(img, imgsz=(W,H), conf=0.5) # chair class 56 with confidence >= 0.5
+ names = model.names
+ # print(names)
+
+ class_found = False
+ for result in results:
+ for i, label in enumerate(result.boxes.cls):
+ # Check if the label matches the chair label
+ if names[int(label)] == default_class:
+ class_found = True
+ # Convert the tensor to a numpy array
+ chair_mask_np = result.masks.data[i].numpy()
+
+ kernel = np.ones((5, 5), np.uint8) # Create a 5x5 kernel for dilation
+ chair_mask_np = cv2.dilate(chair_mask_np, kernel, iterations=2) # Apply dilation
+
+ # Find contours to get bounding box
+ contours, _ = cv2.findContours((chair_mask_np == 1).astype(np.uint8), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
+
+ # Iterate over contours to find the bounding box of each object
+ for contour in contours:
+ x, y, w, h = cv2.boundingRect(contour)
+ x_point = x
+ y_point = y
+ width = w
+ height = h
+
+ # Get the corresponding mask
+ mask = result.masks.data[i].numpy() * 255
+ dilated_mask = cv2.dilate(mask, kernel, iterations=2) # Apply dilation
+ # Resize the mask to match the dimensions of the original image
+ resized_mask = cv2.resize(dilated_mask, (img.shape[1], img.shape[0]))
+ # Convert mask to base64
+ mask_base64 = convert_image_to_base64(resized_mask)
+
+ # call repainting and merge function
+ output_base64 = repaitingAndMerge(append_base64,str(model_path), str(config_path),width, height, x_point, y_point, input_resized_base64, mask_base64)
+ # Return the output base64 image in the API response
+ return jsonify({'output_base64': output_base64}), 200
+
+ # return class not found in prediction
+ if not class_found:
+ return jsonify({'message': f'{default_class} object not found in the image'}), 200
+
+def repaitingAndMerge(append_image_base64_image, model_path, config_path, width, height, xposition, yposition, input_base64, mask_base64):
+ config = get_config(config_path)
+ device = torch.device("cpu")
+ trainer = Trainer(config)
+ trainer.load_state_dict(load_weights(model_path, device), strict=False)
+ trainer.eval()
+
+ # lama inpainting start
+ print("lama inpainting start")
+ inpaint_result_base64 = batch_inpaint('lama', 'cpu', input_base64, mask_base64)
+ print("lama inpainting end")
+
+ # Decode base64 to bytes
+ inpaint_result_bytes = base64.b64decode(inpaint_result_base64)
+
+ # Convert bytes to NumPy array
+ inpaint_result_np = np.array(Image.open(io.BytesIO(inpaint_result_bytes)))
+
+ # Create PIL Image from NumPy array
+ final_image = Image.fromarray(inpaint_result_np)
+
+ print("merge start")
+ # Decode base64 to binary data
+ decoded_image_data = base64.b64decode(append_image_base64_image)
+ # Convert binary data to a NumPy array
+ append_image = cv2.imdecode(np.frombuffer(decoded_image_data, np.uint8), cv2.IMREAD_UNCHANGED)
+ # Resize the append image while preserving transparency
+ resized_image = cv2.resize(append_image, (width, height), interpolation=cv2.INTER_AREA)
+ # Convert the resized image to RGBA format (assuming it's in BGRA format)
+ resized_image = cv2.cvtColor(resized_image, cv2.COLOR_BGRA2RGBA)
+ # Create a PIL Image from the resized image with transparent background
+ append_image_pil = Image.fromarray(resized_image)
+ # Paste the append image onto the final image
+ final_image.paste(append_image_pil, (xposition, yposition), append_image_pil)
+ # Save the resulting image
+ print("merge end")
+ # Convert the final image to base64
+ with io.BytesIO() as output_buffer:
+ final_image.save(output_buffer, format='PNG')
+ output_base64 = base64.b64encode(output_buffer.getvalue()).decode('utf-8')
+
+ return output_base64
+
+
+if __name__ == '__main__':
+ app.run(host='0.0.0.0',debug=True)
diff --git a/iopaint/cli.py b/iopaint/cli.py
new file mode 100644
index 0000000000000000000000000000000000000000..951fbb4e6ae97dbf66175aec8736b95d5ef4c76c
--- /dev/null
+++ b/iopaint/cli.py
@@ -0,0 +1,223 @@
+import webbrowser
+from contextlib import asynccontextmanager
+from pathlib import Path
+from typing import Dict, Optional
+
+import typer
+from fastapi import FastAPI
+from loguru import logger
+from typer import Option
+from typer_config import use_json_config
+
+from iopaint.const import *
+from iopaint.runtime import setup_model_dir, dump_environment_info, check_device
+from iopaint.schema import InteractiveSegModel, Device, RealESRGANModel, RemoveBGModel
+
+typer_app = typer.Typer(pretty_exceptions_show_locals=False, add_completion=False)
+
+
+@typer_app.command(help="Install all plugins dependencies")
+def install_plugins_packages():
+ from iopaint.installer import install_plugins_package
+
+ install_plugins_package()
+
+
+@typer_app.command(help="Download SD/SDXL normal/inpainting model from HuggingFace")
+def download(
+ model: str = Option(
+ ..., help="Model id on HuggingFace e.g: runwayml/stable-diffusion-inpainting"
+ ),
+ model_dir: Path = Option(
+ DEFAULT_MODEL_DIR,
+ help=MODEL_DIR_HELP,
+ file_okay=False,
+ callback=setup_model_dir,
+ ),
+):
+ from iopaint.download import cli_download_model
+
+ cli_download_model(model)
+
+
+@typer_app.command(name="list", help="List downloaded models")
+def list_model(
+ model_dir: Path = Option(
+ DEFAULT_MODEL_DIR,
+ help=MODEL_DIR_HELP,
+ file_okay=False,
+ callback=setup_model_dir,
+ ),
+):
+ from iopaint.download import scan_models
+
+ scanned_models = scan_models()
+ for it in scanned_models:
+ print(it.name)
+
+
+@typer_app.command(help="Batch processing images")
+def run(
+ model: str = Option("lama"),
+ device: Device = Option(Device.cpu),
+ image: Path = Option(..., help="Image folders or file path"),
+ mask: Path = Option(
+ ...,
+ help="Mask folders or file path. "
+ "If it is a directory, the mask images in the directory should have the same name as the original image."
+ "If it is a file, all images will use this mask."
+ "Mask will automatically resize to the same size as the original image.",
+ ),
+ output: Path = Option(..., help="Output directory or file path"),
+ config: Path = Option(
+ None, help="Config file path. You can use dump command to create a base config."
+ ),
+ concat: bool = Option(
+ False, help="Concat original image, mask and output images into one image"
+ ),
+ model_dir: Path = Option(
+ DEFAULT_MODEL_DIR,
+ help=MODEL_DIR_HELP,
+ file_okay=False,
+ callback=setup_model_dir,
+ ),
+):
+ from iopaint.download import cli_download_model, scan_models
+
+ scanned_models = scan_models()
+ if model not in [it.name for it in scanned_models]:
+ logger.info(f"{model} not found in {model_dir}, try to downloading")
+ cli_download_model(model)
+
+ from iopaint.batch_processing import batch_inpaint
+
+ batch_inpaint(model, device, image, mask, output, config, concat)
+
+
+@typer_app.command(help="Start IOPaint server")
+@use_json_config()
+def start(
+ host: str = Option("127.0.0.1"),
+ port: int = Option(8080),
+ inbrowser: bool = Option(False, help=INBROWSER_HELP),
+ model: str = Option(
+ DEFAULT_MODEL,
+ help=f"Erase models: [{', '.join(AVAILABLE_MODELS)}].\n"
+ f"Diffusion models: [{', '.join(DIFFUSION_MODELS)}] or any SD/SDXL normal/inpainting models on HuggingFace.",
+ ),
+ model_dir: Path = Option(
+ DEFAULT_MODEL_DIR,
+ help=MODEL_DIR_HELP,
+ dir_okay=True,
+ file_okay=False,
+ callback=setup_model_dir,
+ ),
+ low_mem: bool = Option(False, help=LOW_MEM_HELP),
+ no_half: bool = Option(False, help=NO_HALF_HELP),
+ cpu_offload: bool = Option(False, help=CPU_OFFLOAD_HELP),
+ disable_nsfw_checker: bool = Option(False, help=DISABLE_NSFW_HELP),
+ cpu_textencoder: bool = Option(False, help=CPU_TEXTENCODER_HELP),
+ local_files_only: bool = Option(False, help=LOCAL_FILES_ONLY_HELP),
+ device: Device = Option(Device.cpu),
+ input: Optional[Path] = Option(None, help=INPUT_HELP),
+ output_dir: Optional[Path] = Option(
+ None, help=OUTPUT_DIR_HELP, dir_okay=True, file_okay=False
+ ),
+ quality: int = Option(95, help=QUALITY_HELP),
+ enable_interactive_seg: bool = Option(False, help=INTERACTIVE_SEG_HELP),
+ interactive_seg_model: InteractiveSegModel = Option(
+ InteractiveSegModel.vit_b, help=INTERACTIVE_SEG_MODEL_HELP
+ ),
+ interactive_seg_device: Device = Option(Device.cpu),
+ enable_remove_bg: bool = Option(False, help=REMOVE_BG_HELP),
+ remove_bg_model: RemoveBGModel = Option(RemoveBGModel.briaai_rmbg_1_4),
+ enable_anime_seg: bool = Option(False, help=ANIMESEG_HELP),
+ enable_realesrgan: bool = Option(False),
+ realesrgan_device: Device = Option(Device.cpu),
+ realesrgan_model: RealESRGANModel = Option(RealESRGANModel.realesr_general_x4v3),
+ enable_gfpgan: bool = Option(False),
+ gfpgan_device: Device = Option(Device.cpu),
+ enable_restoreformer: bool = Option(False),
+ restoreformer_device: Device = Option(Device.cpu),
+):
+ dump_environment_info()
+ device = check_device(device)
+ if input and not input.exists():
+ logger.error(f"invalid --input: {input} not exists")
+ exit(-1)
+ if input and input.is_dir() and not output_dir:
+ logger.error(f"invalid --output-dir: must be set when --input is a directory")
+ exit(-1)
+ if output_dir:
+ output_dir = output_dir.expanduser().absolute()
+ logger.info(f"Image will be saved to {output_dir}")
+ if not output_dir.exists():
+ logger.info(f"Create output directory {output_dir}")
+ output_dir.mkdir(parents=True)
+
+ model_dir = model_dir.expanduser().absolute()
+
+ if local_files_only:
+ os.environ["TRANSFORMERS_OFFLINE"] = "1"
+ os.environ["HF_HUB_OFFLINE"] = "1"
+
+ from iopaint.download import cli_download_model, scan_models
+
+ scanned_models = scan_models()
+ if model not in [it.name for it in scanned_models]:
+ logger.info(f"{model} not found in {model_dir}, try to downloading")
+ cli_download_model(model)
+
+ from iopaint.api import Api
+ from iopaint.schema import ApiConfig
+
+ @asynccontextmanager
+ async def lifespan(app: FastAPI):
+ if inbrowser:
+ webbrowser.open(f"http://localhost:{port}", new=0, autoraise=True)
+ yield
+
+ app = FastAPI(lifespan=lifespan)
+
+ api_config = ApiConfig(
+ host=host,
+ port=port,
+ inbrowser=inbrowser,
+ model=model,
+ no_half=no_half,
+ low_mem=low_mem,
+ cpu_offload=cpu_offload,
+ disable_nsfw_checker=disable_nsfw_checker,
+ local_files_only=local_files_only,
+ cpu_textencoder=cpu_textencoder if device == Device.cuda else False,
+ device=device,
+ input=input,
+ output_dir=output_dir,
+ quality=quality,
+ enable_interactive_seg=enable_interactive_seg,
+ interactive_seg_model=interactive_seg_model,
+ interactive_seg_device=interactive_seg_device,
+ enable_remove_bg=enable_remove_bg,
+ remove_bg_model=remove_bg_model,
+ enable_anime_seg=enable_anime_seg,
+ enable_realesrgan=enable_realesrgan,
+ realesrgan_device=realesrgan_device,
+ realesrgan_model=realesrgan_model,
+ enable_gfpgan=enable_gfpgan,
+ gfpgan_device=gfpgan_device,
+ enable_restoreformer=enable_restoreformer,
+ restoreformer_device=restoreformer_device,
+ )
+ print(api_config.model_dump_json(indent=4))
+ api = Api(app, api_config)
+ api.launch()
+
+
+@typer_app.command(help="Start IOPaint web config page")
+def start_web_config(
+ config_file: Path = Option("config.json"),
+):
+ dump_environment_info()
+ from iopaint.web_config import main
+
+ main(config_file)
diff --git a/iopaint/const.py b/iopaint/const.py
new file mode 100644
index 0000000000000000000000000000000000000000..f7eb1077cc3b6b552e2a1b136b17ea4dbfa08709
--- /dev/null
+++ b/iopaint/const.py
@@ -0,0 +1,121 @@
+import os
+from typing import List
+
+INSTRUCT_PIX2PIX_NAME = "timbrooks/instruct-pix2pix"
+KANDINSKY22_NAME = "kandinsky-community/kandinsky-2-2-decoder-inpaint"
+POWERPAINT_NAME = "Sanster/PowerPaint-V1-stable-diffusion-inpainting"
+ANYTEXT_NAME = "Sanster/AnyText"
+
+
+DIFFUSERS_SD_CLASS_NAME = "StableDiffusionPipeline"
+DIFFUSERS_SD_INPAINT_CLASS_NAME = "StableDiffusionInpaintPipeline"
+DIFFUSERS_SDXL_CLASS_NAME = "StableDiffusionXLPipeline"
+DIFFUSERS_SDXL_INPAINT_CLASS_NAME = "StableDiffusionXLInpaintPipeline"
+
+MPS_UNSUPPORT_MODELS = [
+ "lama",
+ "ldm",
+ "zits",
+ "mat",
+ "fcf",
+ "cv2",
+ "manga",
+]
+
+DEFAULT_MODEL = "lama"
+AVAILABLE_MODELS = ["lama", "ldm", "zits", "mat", "fcf", "manga", "cv2", "migan"]
+DIFFUSION_MODELS = [
+ "runwayml/stable-diffusion-inpainting",
+ "Uminosachi/realisticVisionV51_v51VAE-inpainting",
+ "redstonehero/dreamshaper-inpainting",
+ "Sanster/anything-4.0-inpainting",
+ "diffusers/stable-diffusion-xl-1.0-inpainting-0.1",
+ "Fantasy-Studio/Paint-by-Example",
+ POWERPAINT_NAME,
+ ANYTEXT_NAME,
+]
+
+NO_HALF_HELP = """
+Using full precision(fp32) model.
+If your diffusion model generate result is always black or green, use this argument.
+"""
+
+CPU_OFFLOAD_HELP = """
+Offloads diffusion model's weight to CPU RAM, significantly reducing vRAM usage.
+"""
+
+LOW_MEM_HELP = "Enable attention slicing and vae tiling to save memory."
+
+DISABLE_NSFW_HELP = """
+Disable NSFW checker for diffusion model.
+"""
+
+CPU_TEXTENCODER_HELP = """
+Run diffusion models text encoder on CPU to reduce vRAM usage.
+"""
+
+SD_CONTROLNET_CHOICES: List[str] = [
+ "lllyasviel/control_v11p_sd15_canny",
+ # "lllyasviel/control_v11p_sd15_seg",
+ "lllyasviel/control_v11p_sd15_openpose",
+ "lllyasviel/control_v11p_sd15_inpaint",
+ "lllyasviel/control_v11f1p_sd15_depth",
+]
+
+SD2_CONTROLNET_CHOICES = [
+ "thibaud/controlnet-sd21-canny-diffusers",
+ "thibaud/controlnet-sd21-depth-diffusers",
+ "thibaud/controlnet-sd21-openpose-diffusers",
+]
+
+SDXL_CONTROLNET_CHOICES = [
+ "thibaud/controlnet-openpose-sdxl-1.0",
+ "destitech/controlnet-inpaint-dreamer-sdxl",
+ "diffusers/controlnet-canny-sdxl-1.0",
+ "diffusers/controlnet-canny-sdxl-1.0-mid",
+ "diffusers/controlnet-canny-sdxl-1.0-small",
+ "diffusers/controlnet-depth-sdxl-1.0",
+ "diffusers/controlnet-depth-sdxl-1.0-mid",
+ "diffusers/controlnet-depth-sdxl-1.0-small",
+]
+
+LOCAL_FILES_ONLY_HELP = """
+When loading diffusion models, using local files only, not connect to HuggingFace server.
+"""
+
+DEFAULT_MODEL_DIR = os.path.abspath(
+ os.getenv("XDG_CACHE_HOME", os.path.join(os.path.expanduser("~"), ".cache"))
+)
+#DEFAULT_MODEL_DIR = os.path.abspath("pretrained-models")
+
+MODEL_DIR_HELP = f"""
+Model download directory (by setting XDG_CACHE_HOME environment variable), by default model download to {DEFAULT_MODEL_DIR}
+"""
+
+OUTPUT_DIR_HELP = """
+Result images will be saved to output directory automatically.
+"""
+
+INPUT_HELP = """
+If input is image, it will be loaded by default.
+If input is directory, you can browse and select image in file manager.
+"""
+
+GUI_HELP = """
+Launch Lama Cleaner as desktop app
+"""
+
+QUALITY_HELP = """
+Quality of image encoding, 0-100. Default is 95, higher quality will generate larger file size.
+"""
+
+INTERACTIVE_SEG_HELP = "Enable interactive segmentation using Segment Anything."
+INTERACTIVE_SEG_MODEL_HELP = "Model size: mobile_sam < vit_b < vit_l < vit_h. Bigger model size means better segmentation but slower speed."
+REMOVE_BG_HELP = "Enable remove background plugin. Always run on CPU"
+ANIMESEG_HELP = "Enable anime segmentation plugin. Always run on CPU"
+REALESRGAN_HELP = "Enable realesrgan super resolution"
+GFPGAN_HELP = "Enable GFPGAN face restore. To also enhance background, use with --enable-realesrgan"
+RESTOREFORMER_HELP = "Enable RestoreFormer face restore. To also enhance background, use with --enable-realesrgan"
+GIF_HELP = "Enable GIF plugin. Make GIF to compare original and cleaned image"
+
+INBROWSER_HELP = "Automatically launch IOPaint in a new tab on the default browser"
diff --git a/iopaint/download.py b/iopaint/download.py
new file mode 100644
index 0000000000000000000000000000000000000000..2ebd7fcbf907465708a0a4de0d67c0609d069cb3
--- /dev/null
+++ b/iopaint/download.py
@@ -0,0 +1,294 @@
+import json
+import os
+from functools import lru_cache
+from typing import List
+
+from iopaint.schema import ModelType, ModelInfo
+from loguru import logger
+from pathlib import Path
+
+from iopaint.const import (
+ DEFAULT_MODEL_DIR,
+ DIFFUSERS_SD_CLASS_NAME,
+ DIFFUSERS_SD_INPAINT_CLASS_NAME,
+ DIFFUSERS_SDXL_CLASS_NAME,
+ DIFFUSERS_SDXL_INPAINT_CLASS_NAME,
+ ANYTEXT_NAME,
+)
+from iopaint.model.original_sd_configs import get_config_files
+
+
+def cli_download_model(model: str):
+ from iopaint.model import models
+ from iopaint.model.utils import handle_from_pretrained_exceptions
+
+ if model in models and models[model].is_erase_model:
+ logger.info(f"Downloading {model}...")
+ models[model].download()
+ logger.info(f"Done.")
+ elif model == ANYTEXT_NAME:
+ logger.info(f"Downloading {model}...")
+ models[model].download()
+ logger.info(f"Done.")
+ else:
+ logger.info(f"Downloading model from Huggingface: {model}")
+ from diffusers import DiffusionPipeline
+
+ downloaded_path = handle_from_pretrained_exceptions(
+ DiffusionPipeline.download,
+ pretrained_model_name=model,
+ variant="fp16",
+ resume_download=True,
+ )
+ logger.info(f"Done. Downloaded to {downloaded_path}")
+
+
+def folder_name_to_show_name(name: str) -> str:
+ return name.replace("models--", "").replace("--", "/")
+
+
+@lru_cache(maxsize=512)
+def get_sd_model_type(model_abs_path: str) -> ModelType:
+ if "inpaint" in Path(model_abs_path).name.lower():
+ model_type = ModelType.DIFFUSERS_SD_INPAINT
+ else:
+ # load once to check num_in_channels
+ from diffusers import StableDiffusionInpaintPipeline
+
+ try:
+ StableDiffusionInpaintPipeline.from_single_file(
+ model_abs_path,
+ load_safety_checker=False,
+ num_in_channels=9,
+ config_files=get_config_files(),
+ )
+ model_type = ModelType.DIFFUSERS_SD_INPAINT
+ except ValueError as e:
+ if "Trying to set a tensor of shape torch.Size([320, 4, 3, 3])" in str(e):
+ model_type = ModelType.DIFFUSERS_SD
+ else:
+ raise e
+ return model_type
+
+
+@lru_cache()
+def get_sdxl_model_type(model_abs_path: str) -> ModelType:
+ if "inpaint" in model_abs_path:
+ model_type = ModelType.DIFFUSERS_SDXL_INPAINT
+ else:
+ # load once to check num_in_channels
+ from diffusers import StableDiffusionXLInpaintPipeline
+
+ try:
+ model = StableDiffusionXLInpaintPipeline.from_single_file(
+ model_abs_path,
+ load_safety_checker=False,
+ num_in_channels=9,
+ config_files=get_config_files(),
+ )
+ if model.unet.config.in_channels == 9:
+ # https://github.com/huggingface/diffusers/issues/6610
+ model_type = ModelType.DIFFUSERS_SDXL_INPAINT
+ else:
+ model_type = ModelType.DIFFUSERS_SDXL
+ except ValueError as e:
+ if "Trying to set a tensor of shape torch.Size([320, 4, 3, 3])" in str(e):
+ model_type = ModelType.DIFFUSERS_SDXL
+ else:
+ raise e
+ return model_type
+
+
+def scan_single_file_diffusion_models(cache_dir) -> List[ModelInfo]:
+ cache_dir = Path(cache_dir)
+ stable_diffusion_dir = cache_dir / "stable_diffusion"
+ cache_file = stable_diffusion_dir / "iopaint_cache.json"
+ model_type_cache = {}
+ if cache_file.exists():
+ try:
+ with open(cache_file, "r", encoding="utf-8") as f:
+ model_type_cache = json.load(f)
+ assert isinstance(model_type_cache, dict)
+ except:
+ pass
+
+ res = []
+ for it in stable_diffusion_dir.glob(f"*.*"):
+ if it.suffix not in [".safetensors", ".ckpt"]:
+ continue
+ model_abs_path = str(it.absolute())
+ model_type = model_type_cache.get(it.name)
+ if model_type is None:
+ model_type = get_sd_model_type(model_abs_path)
+ model_type_cache[it.name] = model_type
+ res.append(
+ ModelInfo(
+ name=it.name,
+ path=model_abs_path,
+ model_type=model_type,
+ is_single_file_diffusers=True,
+ )
+ )
+ if stable_diffusion_dir.exists():
+ with open(cache_file, "w", encoding="utf-8") as fw:
+ json.dump(model_type_cache, fw, indent=2, ensure_ascii=False)
+
+ stable_diffusion_xl_dir = cache_dir / "stable_diffusion_xl"
+ sdxl_cache_file = stable_diffusion_xl_dir / "iopaint_cache.json"
+ sdxl_model_type_cache = {}
+ if sdxl_cache_file.exists():
+ try:
+ with open(sdxl_cache_file, "r", encoding="utf-8") as f:
+ sdxl_model_type_cache = json.load(f)
+ assert isinstance(sdxl_model_type_cache, dict)
+ except:
+ pass
+
+ for it in stable_diffusion_xl_dir.glob(f"*.*"):
+ if it.suffix not in [".safetensors", ".ckpt"]:
+ continue
+ model_abs_path = str(it.absolute())
+ model_type = sdxl_model_type_cache.get(it.name)
+ if model_type is None:
+ model_type = get_sdxl_model_type(model_abs_path)
+ sdxl_model_type_cache[it.name] = model_type
+ if stable_diffusion_xl_dir.exists():
+ with open(sdxl_cache_file, "w", encoding="utf-8") as fw:
+ json.dump(sdxl_model_type_cache, fw, indent=2, ensure_ascii=False)
+
+ res.append(
+ ModelInfo(
+ name=it.name,
+ path=model_abs_path,
+ model_type=model_type,
+ is_single_file_diffusers=True,
+ )
+ )
+ return res
+
+
+def scan_inpaint_models(model_dir: Path) -> List[ModelInfo]:
+ res = []
+ from iopaint.model import models
+
+ # logger.info(f"Scanning inpaint models in {model_dir}")
+
+ for name, m in models.items():
+ if m.is_erase_model and m.is_downloaded():
+ res.append(
+ ModelInfo(
+ name=name,
+ path=name,
+ model_type=ModelType.INPAINT,
+ )
+ )
+ return res
+
+
+def scan_diffusers_models() -> List[ModelInfo]:
+ from huggingface_hub.constants import HF_HUB_CACHE
+
+ available_models = []
+ cache_dir = Path(HF_HUB_CACHE)
+ # logger.info(f"Scanning diffusers models in {cache_dir}")
+ diffusers_model_names = []
+ for it in cache_dir.glob("**/*/model_index.json"):
+ with open(it, "r", encoding="utf-8") as f:
+ try:
+ data = json.load(f)
+ except:
+ continue
+
+ _class_name = data["_class_name"]
+ name = folder_name_to_show_name(it.parent.parent.parent.name)
+ if name in diffusers_model_names:
+ continue
+ if "PowerPaint" in name:
+ model_type = ModelType.DIFFUSERS_OTHER
+ elif _class_name == DIFFUSERS_SD_CLASS_NAME:
+ model_type = ModelType.DIFFUSERS_SD
+ elif _class_name == DIFFUSERS_SD_INPAINT_CLASS_NAME:
+ model_type = ModelType.DIFFUSERS_SD_INPAINT
+ elif _class_name == DIFFUSERS_SDXL_CLASS_NAME:
+ model_type = ModelType.DIFFUSERS_SDXL
+ elif _class_name == DIFFUSERS_SDXL_INPAINT_CLASS_NAME:
+ model_type = ModelType.DIFFUSERS_SDXL_INPAINT
+ elif _class_name in [
+ "StableDiffusionInstructPix2PixPipeline",
+ "PaintByExamplePipeline",
+ "KandinskyV22InpaintPipeline",
+ "AnyText",
+ ]:
+ model_type = ModelType.DIFFUSERS_OTHER
+ else:
+ continue
+
+ diffusers_model_names.append(name)
+ available_models.append(
+ ModelInfo(
+ name=name,
+ path=name,
+ model_type=model_type,
+ )
+ )
+ return available_models
+
+
+def _scan_converted_diffusers_models(cache_dir) -> List[ModelInfo]:
+ cache_dir = Path(cache_dir)
+ available_models = []
+ diffusers_model_names = []
+ for it in cache_dir.glob("**/*/model_index.json"):
+ with open(it, "r", encoding="utf-8") as f:
+ try:
+ data = json.load(f)
+ except:
+ logger.error(
+ f"Failed to load {it}, please try revert from original model or fix model_index.json by hand."
+ )
+ continue
+
+ _class_name = data["_class_name"]
+ name = folder_name_to_show_name(it.parent.name)
+ if name in diffusers_model_names:
+ continue
+ elif _class_name == DIFFUSERS_SD_CLASS_NAME:
+ model_type = ModelType.DIFFUSERS_SD
+ elif _class_name == DIFFUSERS_SD_INPAINT_CLASS_NAME:
+ model_type = ModelType.DIFFUSERS_SD_INPAINT
+ elif _class_name == DIFFUSERS_SDXL_CLASS_NAME:
+ model_type = ModelType.DIFFUSERS_SDXL
+ elif _class_name == DIFFUSERS_SDXL_INPAINT_CLASS_NAME:
+ model_type = ModelType.DIFFUSERS_SDXL_INPAINT
+ else:
+ continue
+
+ diffusers_model_names.append(name)
+ available_models.append(
+ ModelInfo(
+ name=name,
+ path=str(it.parent.absolute()),
+ model_type=model_type,
+ )
+ )
+ return available_models
+
+
+def scan_converted_diffusers_models(cache_dir) -> List[ModelInfo]:
+ cache_dir = Path(cache_dir)
+ available_models = []
+ stable_diffusion_dir = cache_dir / "stable_diffusion"
+ stable_diffusion_xl_dir = cache_dir / "stable_diffusion_xl"
+ available_models.extend(_scan_converted_diffusers_models(stable_diffusion_dir))
+ available_models.extend(_scan_converted_diffusers_models(stable_diffusion_xl_dir))
+ return available_models
+
+
+def scan_models() -> List[ModelInfo]:
+ model_dir = os.getenv("XDG_CACHE_HOME", DEFAULT_MODEL_DIR)
+ available_models = []
+ available_models.extend(scan_inpaint_models(model_dir))
+ available_models.extend(scan_single_file_diffusion_models(model_dir))
+ available_models.extend(scan_diffusers_models())
+ available_models.extend(scan_converted_diffusers_models(model_dir))
+ return available_models
diff --git a/iopaint/file_manager/file_manager.py b/iopaint/file_manager/file_manager.py
new file mode 100644
index 0000000000000000000000000000000000000000..413162cfedb88e9b7dfeed872be45049eaf1ba0f
--- /dev/null
+++ b/iopaint/file_manager/file_manager.py
@@ -0,0 +1,215 @@
+import os
+from io import BytesIO
+from pathlib import Path
+from typing import List
+
+from PIL import Image, ImageOps, PngImagePlugin
+from fastapi import FastAPI, UploadFile, HTTPException
+from starlette.responses import FileResponse
+
+from ..schema import MediasResponse, MediaTab
+
+LARGE_ENOUGH_NUMBER = 100
+PngImagePlugin.MAX_TEXT_CHUNK = LARGE_ENOUGH_NUMBER * (1024**2)
+from .storage_backends import FilesystemStorageBackend
+from .utils import aspect_to_string, generate_filename, glob_img
+
+
+class FileManager:
+ def __init__(self, app: FastAPI, input_dir: Path, output_dir: Path):
+ self.app = app
+ self.input_dir: Path = input_dir
+ self.output_dir: Path = output_dir
+
+ self.image_dir_filenames = []
+ self.output_dir_filenames = []
+ if not self.thumbnail_directory.exists():
+ self.thumbnail_directory.mkdir(parents=True)
+
+ # fmt: off
+ self.app.add_api_route("/api/v1/medias", self.api_medias, methods=["GET"], response_model=List[MediasResponse])
+ self.app.add_api_route("/api/v1/media_file", self.api_media_file, methods=["GET"])
+ self.app.add_api_route("/api/v1/media_thumbnail_file", self.api_media_thumbnail_file, methods=["GET"])
+ # fmt: on
+
+ def api_medias(self, tab: MediaTab) -> List[MediasResponse]:
+ img_dir = self._get_dir(tab)
+ return self._media_names(img_dir)
+
+ def api_media_file(self, tab: MediaTab, filename: str) -> FileResponse:
+ file_path = self._get_file(tab, filename)
+ return FileResponse(file_path, media_type="image/png")
+
+ # tab=${tab}?filename=${filename.name}?width=${width}&height=${height}
+ def api_media_thumbnail_file(
+ self, tab: MediaTab, filename: str, width: int, height: int
+ ) -> FileResponse:
+ img_dir = self._get_dir(tab)
+ thumb_filename, (width, height) = self.get_thumbnail(
+ img_dir, filename, width=width, height=height
+ )
+ thumbnail_filepath = self.thumbnail_directory / thumb_filename
+ return FileResponse(
+ thumbnail_filepath,
+ headers={
+ "X-Width": str(width),
+ "X-Height": str(height),
+ },
+ media_type="image/jpeg",
+ )
+
+ def _get_dir(self, tab: MediaTab) -> Path:
+ if tab == "input":
+ return self.input_dir
+ elif tab == "output":
+ return self.output_dir
+ else:
+ raise HTTPException(status_code=422, detail=f"tab not found: {tab}")
+
+ def _get_file(self, tab: MediaTab, filename: str) -> Path:
+ file_path = self._get_dir(tab) / filename
+ if not file_path.exists():
+ raise HTTPException(status_code=422, detail=f"file not found: {file_path}")
+ return file_path
+
+ @property
+ def thumbnail_directory(self) -> Path:
+ return self.output_dir / "thumbnails"
+
+ @staticmethod
+ def _media_names(directory: Path) -> List[MediasResponse]:
+ names = sorted([it.name for it in glob_img(directory)])
+ res = []
+ for name in names:
+ path = os.path.join(directory, name)
+ img = Image.open(path)
+ res.append(
+ MediasResponse(
+ name=name,
+ height=img.height,
+ width=img.width,
+ ctime=os.path.getctime(path),
+ mtime=os.path.getmtime(path),
+ )
+ )
+ return res
+
+ def get_thumbnail(
+ self, directory: Path, original_filename: str, width, height, **options
+ ):
+ directory = Path(directory)
+ storage = FilesystemStorageBackend(self.app)
+ crop = options.get("crop", "fit")
+ background = options.get("background")
+ quality = options.get("quality", 90)
+
+ original_path, original_filename = os.path.split(original_filename)
+ original_filepath = os.path.join(directory, original_path, original_filename)
+ image = Image.open(BytesIO(storage.read(original_filepath)))
+
+ # keep ratio resize
+ if not width and not height:
+ width = 256
+
+ if width != 0:
+ height = int(image.height * width / image.width)
+ else:
+ width = int(image.width * height / image.height)
+
+ thumbnail_size = (width, height)
+
+ thumbnail_filename = generate_filename(
+ directory,
+ original_filename,
+ aspect_to_string(thumbnail_size),
+ crop,
+ background,
+ quality,
+ )
+
+ thumbnail_filepath = os.path.join(
+ self.thumbnail_directory, original_path, thumbnail_filename
+ )
+
+ if storage.exists(thumbnail_filepath):
+ return thumbnail_filepath, (width, height)
+
+ try:
+ image.load()
+ except (IOError, OSError):
+ self.app.logger.warning("Thumbnail not load image: %s", original_filepath)
+ return thumbnail_filepath, (width, height)
+
+ # get original image format
+ options["format"] = options.get("format", image.format)
+
+ image = self._create_thumbnail(
+ image, thumbnail_size, crop, background=background
+ )
+
+ raw_data = self.get_raw_data(image, **options)
+ storage.save(thumbnail_filepath, raw_data)
+
+ return thumbnail_filepath, (width, height)
+
+ def get_raw_data(self, image, **options):
+ data = {
+ "format": self._get_format(image, **options),
+ "quality": options.get("quality", 90),
+ }
+
+ _file = BytesIO()
+ image.save(_file, **data)
+ return _file.getvalue()
+
+ @staticmethod
+ def colormode(image, colormode="RGB"):
+ if colormode == "RGB" or colormode == "RGBA":
+ if image.mode == "RGBA":
+ return image
+ if image.mode == "LA":
+ return image.convert("RGBA")
+ return image.convert(colormode)
+
+ if colormode == "GRAY":
+ return image.convert("L")
+
+ return image.convert(colormode)
+
+ @staticmethod
+ def background(original_image, color=0xFF):
+ size = (max(original_image.size),) * 2
+ image = Image.new("L", size, color)
+ image.paste(
+ original_image,
+ tuple(map(lambda x: (x[0] - x[1]) / 2, zip(size, original_image.size))),
+ )
+
+ return image
+
+ def _get_format(self, image, **options):
+ if options.get("format"):
+ return options.get("format")
+ if image.format:
+ return image.format
+
+ return "JPEG"
+
+ def _create_thumbnail(self, image, size, crop="fit", background=None):
+ try:
+ resample = Image.Resampling.LANCZOS
+ except AttributeError: # pylint: disable=raise-missing-from
+ resample = Image.ANTIALIAS
+
+ if crop == "fit":
+ image = ImageOps.fit(image, size, resample)
+ else:
+ image = image.copy()
+ image.thumbnail(size, resample=resample)
+
+ if background is not None:
+ image = self.background(image)
+
+ image = self.colormode(image)
+
+ return image
diff --git a/iopaint/helper.py b/iopaint/helper.py
new file mode 100644
index 0000000000000000000000000000000000000000..80dafb416b158f585fbc4f9e5b12fd65fb1a941a
--- /dev/null
+++ b/iopaint/helper.py
@@ -0,0 +1,425 @@
+import base64
+import imghdr
+import io
+import os
+import sys
+from typing import List, Optional, Dict, Tuple
+
+from urllib.parse import urlparse
+import cv2
+from PIL import Image, ImageOps, PngImagePlugin
+import numpy as np
+import torch
+from iopaint.const import MPS_UNSUPPORT_MODELS
+from loguru import logger
+from torch.hub import download_url_to_file, get_dir
+import hashlib
+
+
+def md5sum(filename):
+ md5 = hashlib.md5()
+ with open(filename, "rb") as f:
+ for chunk in iter(lambda: f.read(128 * md5.block_size), b""):
+ md5.update(chunk)
+ return md5.hexdigest()
+
+
+def switch_mps_device(model_name, device):
+ if model_name in MPS_UNSUPPORT_MODELS and str(device) == "mps":
+ logger.info(f"{model_name} not support mps, switch to cpu")
+ return torch.device("cpu")
+ return device
+
+
+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(model_dir)
+ filename = os.path.basename(parts.path)
+ cached_file = os.path.join(model_dir, filename)
+ return cached_file
+
+def get_cache_path_by_local(url):
+ root_path = os.getcwd()
+ model_path = os.path.join(root_path, 'pretrained-model', 'big-lama.pt')
+ return model_path
+
+def download_model(url, model_md5: str = None):
+ cached_file = get_cache_path_by_url(url)
+ # cached_file = get_cache_path_by_local(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)
+ if model_md5:
+ _md5 = md5sum(cached_file)
+ if model_md5 == _md5:
+ logger.info(f"Download model success, md5: {_md5}")
+ else:
+ try:
+ os.remove(cached_file)
+ logger.error(
+ f"Model md5: {_md5}, expected md5: {model_md5}, wrong model deleted. Please restart iopaint."
+ f"If you still have errors, please try download model manually first https://lama-cleaner-docs.vercel.app/install/download_model_manually.\n"
+ )
+ except:
+ logger.error(
+ f"Model md5: {_md5}, expected md5: {model_md5}, please delete {cached_file} and restart iopaint."
+ )
+ exit(-1)
+
+ return cached_file
+
+
+def ceil_modulo(x, mod):
+ if x % mod == 0:
+ return x
+ return (x // mod + 1) * mod
+
+
+def handle_error(model_path, model_md5, e):
+ _md5 = md5sum(model_path)
+ if _md5 != model_md5:
+ try:
+ os.remove(model_path)
+ logger.error(
+ f"Model md5: {_md5}, expected md5: {model_md5}, wrong model deleted. Please restart iopaint."
+ f"If you still have errors, please try download model manually first https://lama-cleaner-docs.vercel.app/install/download_model_manually.\n"
+ )
+ except:
+ logger.error(
+ f"Model md5: {_md5}, expected md5: {model_md5}, please delete {model_path} and restart iopaint."
+ )
+ else:
+ logger.error(
+ f"Failed to load model {model_path},"
+ f"please submit an issue at https://github.com/Sanster/lama-cleaner/issues and include a screenshot of the error:\n{e}"
+ )
+ exit(-1)
+
+
+def load_jit_model(url_or_path, device, model_md5: str):
+ if os.path.exists(url_or_path):
+ model_path = url_or_path
+ else:
+ model_path = download_model(url_or_path, model_md5)
+
+ logger.info(f"Loading model from: {model_path}")
+ try:
+ model = torch.jit.load(model_path, map_location="cpu").to(device)
+ except Exception as e:
+ handle_error(model_path, model_md5, e)
+ model.eval()
+ return model
+
+
+def load_model(model: torch.nn.Module, url_or_path, device, model_md5):
+ if os.path.exists(url_or_path):
+ model_path = url_or_path
+ else:
+ model_path = download_model(url_or_path, model_md5)
+
+ try:
+ logger.info(f"Loading model from: {model_path}")
+ state_dict = torch.load(model_path, map_location="cpu")
+ model.load_state_dict(state_dict, strict=True)
+ model.to(device)
+ except Exception as e:
+ handle_error(model_path, model_md5, e)
+ 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 pil_to_bytes(pil_img, ext: str, quality: int = 95, infos={}) -> bytes:
+ with io.BytesIO() as output:
+ kwargs = {k: v for k, v in infos.items() if v is not None}
+ if ext == "jpg":
+ ext = "jpeg"
+ if "png" == ext.lower() and "parameters" in kwargs:
+ pnginfo_data = PngImagePlugin.PngInfo()
+ pnginfo_data.add_text("parameters", kwargs["parameters"])
+ kwargs["pnginfo"] = pnginfo_data
+
+ pil_img.save(output, format=ext, quality=quality, **kwargs)
+ image_bytes = output.getvalue()
+ return image_bytes
+
+def pil_to_bytes_single(pil_img, ext: str, quality: int = 95, infos=None) -> bytes:
+ infos = infos or {} # Use an empty dictionary if infos is None
+ with io.BytesIO() as output:
+ kwargs = {k: v for k, v in infos.items() if v is not None}
+ if ext == "jpg":
+ ext = "jpeg"
+ if "png" == ext.lower() and "parameters" in kwargs:
+ pnginfo_data = PngImagePlugin.PngInfo()
+ pnginfo_data.add_text("parameters", kwargs["parameters"])
+ kwargs["pnginfo"] = pnginfo_data
+
+ pil_img.save(output, format=ext, quality=quality, **kwargs)
+ image_bytes = output.getvalue()
+ return image_bytes
+
+
+def load_img(img_bytes, gray: bool = False, return_info: bool = False):
+ alpha_channel = None
+ image = Image.open(io.BytesIO(img_bytes))
+
+ if return_info:
+ infos = image.info
+
+ try:
+ image = ImageOps.exif_transpose(image)
+ except:
+ pass
+
+ if gray:
+ image = image.convert("L")
+ np_img = np.array(image)
+ else:
+ if image.mode == "RGBA":
+ np_img = np.array(image)
+ alpha_channel = np_img[:, :, -1]
+ np_img = cv2.cvtColor(np_img, cv2.COLOR_RGBA2RGB)
+ else:
+ image = image.convert("RGB")
+ np_img = np.array(image)
+
+ if return_info:
+ return np_img, alpha_channel, infos
+ 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
+
+
+def only_keep_largest_contour(mask: np.ndarray) -> List[np.ndarray]:
+ """
+ Args:
+ mask: (h, w) 0~255
+
+ Returns:
+
+ """
+ _, thresh = cv2.threshold(mask, 127, 255, 0)
+ contours, _ = cv2.findContours(thresh, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
+
+ max_area = 0
+ max_index = -1
+ for i, cnt in enumerate(contours):
+ area = cv2.contourArea(cnt)
+ if area > max_area:
+ max_area = area
+ max_index = i
+
+ if max_index != -1:
+ new_mask = np.zeros_like(mask)
+ return cv2.drawContours(new_mask, contours, max_index, 255, -1)
+ else:
+ return mask
+
+
+def is_mac():
+ return sys.platform == "darwin"
+
+
+def get_image_ext(img_bytes):
+ w = imghdr.what("", img_bytes)
+ if w is None:
+ w = "jpeg"
+ return w
+
+
+def decode_base64_to_image(
+ encoding: str, gray=False
+) -> Tuple[np.array, Optional[np.array], Dict]:
+ if encoding.startswith("data:image/") or encoding.startswith(
+ "data:application/octet-stream;base64,"
+ ):
+ encoding = encoding.split(";")[1].split(",")[1]
+ image = Image.open(io.BytesIO(base64.b64decode(encoding)))
+
+ alpha_channel = None
+ try:
+ image = ImageOps.exif_transpose(image)
+ except:
+ pass
+ # exif_transpose will remove exif rotate info，we must call image.info after exif_transpose
+ infos = image.info
+
+ if gray:
+ image = image.convert("L")
+ np_img = np.array(image)
+ else:
+ if image.mode == "RGBA":
+ np_img = np.array(image)
+ alpha_channel = np_img[:, :, -1]
+ np_img = cv2.cvtColor(np_img, cv2.COLOR_RGBA2RGB)
+ else:
+ image = image.convert("RGB")
+ np_img = np.array(image)
+
+ return np_img, alpha_channel, infos
+
+
+def encode_pil_to_base64(image: Image, quality: int, infos: Dict) -> bytes:
+ img_bytes = pil_to_bytes(
+ image,
+ "png",
+ quality=quality,
+ infos=infos,
+ )
+ return base64.b64encode(img_bytes)
+
+
+def concat_alpha_channel(rgb_np_img, alpha_channel) -> np.ndarray:
+ if alpha_channel is not None:
+ if alpha_channel.shape[:2] != rgb_np_img.shape[:2]:
+ alpha_channel = cv2.resize(
+ alpha_channel, dsize=(rgb_np_img.shape[1], rgb_np_img.shape[0])
+ )
+ rgb_np_img = np.concatenate(
+ (rgb_np_img, alpha_channel[:, :, np.newaxis]), axis=-1
+ )
+ return rgb_np_img
+
+
+def adjust_mask(mask: np.ndarray, kernel_size: int, operate):
+ # fronted brush color "ffcc00bb"
+ # kernel_size = kernel_size*2+1
+ mask[mask >= 127] = 255
+ mask[mask < 127] = 0
+
+ if operate == "reverse":
+ mask = 255 - mask
+ else:
+ kernel = cv2.getStructuringElement(
+ cv2.MORPH_ELLIPSE, (2 * kernel_size + 1, 2 * kernel_size + 1)
+ )
+ if operate == "expand":
+ mask = cv2.dilate(
+ mask,
+ kernel,
+ iterations=1,
+ )
+ else:
+ mask = cv2.erode(
+ mask,
+ kernel,
+ iterations=1,
+ )
+ res_mask = np.zeros((mask.shape[0], mask.shape[1], 4), dtype=np.uint8)
+ res_mask[mask > 128] = [255, 203, 0, int(255 * 0.73)]
+ res_mask = cv2.cvtColor(res_mask, cv2.COLOR_BGRA2RGBA)
+ return res_mask
+
+
+def gen_frontend_mask(bgr_or_gray_mask):
+ if len(bgr_or_gray_mask.shape) == 3 and bgr_or_gray_mask.shape[2] != 1:
+ bgr_or_gray_mask = cv2.cvtColor(bgr_or_gray_mask, cv2.COLOR_BGR2GRAY)
+
+ # fronted brush color "ffcc00bb"
+ # TODO: how to set kernel size?
+ kernel_size = 9
+ bgr_or_gray_mask = cv2.dilate(
+ bgr_or_gray_mask,
+ np.ones((kernel_size, kernel_size), np.uint8),
+ iterations=1,
+ )
+ res_mask = np.zeros(
+ (bgr_or_gray_mask.shape[0], bgr_or_gray_mask.shape[1], 4), dtype=np.uint8
+ )
+ res_mask[bgr_or_gray_mask > 128] = [255, 203, 0, int(255 * 0.73)]
+ res_mask = cv2.cvtColor(res_mask, cv2.COLOR_BGRA2RGBA)
+ return res_mask
diff --git a/iopaint/installer.py b/iopaint/installer.py
new file mode 100644
index 0000000000000000000000000000000000000000..f255e33548d4fb561e753fb96961d6e2b2a9e446
--- /dev/null
+++ b/iopaint/installer.py
@@ -0,0 +1,12 @@
+import subprocess
+import sys
+
+
+def install(package):
+ subprocess.check_call([sys.executable, "-m", "pip", "install", package])
+
+
+def install_plugins_package():
+ install("rembg")
+ install("realesrgan")
+ install("gfpgan")
diff --git a/iopaint/model/anytext/cldm/cldm.py b/iopaint/model/anytext/cldm/cldm.py
new file mode 100644
index 0000000000000000000000000000000000000000..ad9692a458e027c2ae77d7ca7334b8deaa438393
--- /dev/null
+++ b/iopaint/model/anytext/cldm/cldm.py
@@ -0,0 +1,630 @@
+import os
+from pathlib import Path
+
+import einops
+import torch
+import torch as th
+import torch.nn as nn
+import copy
+from easydict import EasyDict as edict
+
+from iopaint.model.anytext.ldm.modules.diffusionmodules.util import (
+ conv_nd,
+ linear,
+ zero_module,
+ timestep_embedding,
+)
+
+from einops import rearrange, repeat
+from iopaint.model.anytext.ldm.modules.attention import SpatialTransformer
+from iopaint.model.anytext.ldm.modules.diffusionmodules.openaimodel import UNetModel, TimestepEmbedSequential, ResBlock, Downsample, AttentionBlock
+from iopaint.model.anytext.ldm.models.diffusion.ddpm import LatentDiffusion
+from iopaint.model.anytext.ldm.util import log_txt_as_img, exists, instantiate_from_config
+from iopaint.model.anytext.ldm.models.diffusion.ddim import DDIMSampler
+from iopaint.model.anytext.ldm.modules.distributions.distributions import DiagonalGaussianDistribution
+from .recognizer import TextRecognizer, create_predictor
+
+CURRENT_DIR = Path(os.path.dirname(os.path.abspath(__file__)))
+
+
+def count_parameters(model):
+ return sum(p.numel() for p in model.parameters() if p.requires_grad)
+
+
+class ControlledUnetModel(UNetModel):
+ def forward(self, x, timesteps=None, context=None, control=None, only_mid_control=False, **kwargs):
+ hs = []
+ with torch.no_grad():
+ t_emb = timestep_embedding(timesteps, self.model_channels, repeat_only=False)
+ if self.use_fp16:
+ t_emb = t_emb.half()
+ emb = self.time_embed(t_emb)
+ h = x.type(self.dtype)
+ for module in self.input_blocks:
+ h = module(h, emb, context)
+ hs.append(h)
+ h = self.middle_block(h, emb, context)
+
+ if control is not None:
+ h += control.pop()
+
+ for i, module in enumerate(self.output_blocks):
+ if only_mid_control or control is None:
+ h = torch.cat([h, hs.pop()], dim=1)
+ else:
+ h = torch.cat([h, hs.pop() + control.pop()], dim=1)
+ h = module(h, emb, context)
+
+ h = h.type(x.dtype)
+ return self.out(h)
+
+
+class ControlNet(nn.Module):
+ def __init__(
+ self,
+ image_size,
+ in_channels,
+ model_channels,
+ glyph_channels,
+ position_channels,
+ num_res_blocks,
+ attention_resolutions,
+ dropout=0,
+ channel_mult=(1, 2, 4, 8),
+ conv_resample=True,
+ dims=2,
+ use_checkpoint=False,
+ use_fp16=False,
+ num_heads=-1,
+ num_head_channels=-1,
+ num_heads_upsample=-1,
+ use_scale_shift_norm=False,
+ resblock_updown=False,
+ use_new_attention_order=False,
+ use_spatial_transformer=False, # custom transformer support
+ transformer_depth=1, # custom transformer support
+ context_dim=None, # custom transformer support
+ n_embed=None, # custom support for prediction of discrete ids into codebook of first stage vq model
+ legacy=True,
+ disable_self_attentions=None,
+ num_attention_blocks=None,
+ disable_middle_self_attn=False,
+ use_linear_in_transformer=False,
+ ):
+ super().__init__()
+ if use_spatial_transformer:
+ assert context_dim is not None, 'Fool!! You forgot to include the dimension of your cross-attention conditioning...'
+
+ if context_dim is not None:
+ assert use_spatial_transformer, 'Fool!! You forgot to use the spatial transformer for your cross-attention conditioning...'
+ from omegaconf.listconfig import ListConfig
+ if type(context_dim) == ListConfig:
+ context_dim = list(context_dim)
+
+ if num_heads_upsample == -1:
+ num_heads_upsample = num_heads
+
+ if num_heads == -1:
+ assert num_head_channels != -1, 'Either num_heads or num_head_channels has to be set'
+
+ if num_head_channels == -1:
+ assert num_heads != -1, 'Either num_heads or num_head_channels has to be set'
+ self.dims = dims
+ self.image_size = image_size
+ self.in_channels = in_channels
+ self.model_channels = model_channels
+ if isinstance(num_res_blocks, int):
+ self.num_res_blocks = len(channel_mult) * [num_res_blocks]
+ else:
+ if len(num_res_blocks) != len(channel_mult):
+ raise ValueError("provide num_res_blocks either as an int (globally constant) or "
+ "as a list/tuple (per-level) with the same length as channel_mult")
+ self.num_res_blocks = num_res_blocks
+ if disable_self_attentions is not None:
+ # should be a list of booleans, indicating whether to disable self-attention in TransformerBlocks or not
+ assert len(disable_self_attentions) == len(channel_mult)
+ if num_attention_blocks is not None:
+ assert len(num_attention_blocks) == len(self.num_res_blocks)
+ assert all(map(lambda i: self.num_res_blocks[i] >= num_attention_blocks[i], range(len(num_attention_blocks))))
+ print(f"Constructor of UNetModel received num_attention_blocks={num_attention_blocks}. "
+ f"This option has LESS priority than attention_resolutions {attention_resolutions}, "
+ f"i.e., in cases where num_attention_blocks[i] > 0 but 2**i not in attention_resolutions, "
+ f"attention will still not be set.")
+ self.attention_resolutions = attention_resolutions
+ self.dropout = dropout
+ self.channel_mult = channel_mult
+ self.conv_resample = conv_resample
+ self.use_checkpoint = use_checkpoint
+ self.use_fp16 = use_fp16
+ self.dtype = th.float16 if use_fp16 else th.float32
+ self.num_heads = num_heads
+ self.num_head_channels = num_head_channels
+ self.num_heads_upsample = num_heads_upsample
+ self.predict_codebook_ids = n_embed is not None
+
+ time_embed_dim = model_channels * 4
+ self.time_embed = nn.Sequential(
+ linear(model_channels, time_embed_dim),
+ nn.SiLU(),
+ linear(time_embed_dim, time_embed_dim),
+ )
+
+ self.input_blocks = nn.ModuleList(
+ [
+ TimestepEmbedSequential(
+ conv_nd(dims, in_channels, model_channels, 3, padding=1)
+ )
+ ]
+ )
+ self.zero_convs = nn.ModuleList([self.make_zero_conv(model_channels)])
+
+ self.glyph_block = TimestepEmbedSequential(
+ conv_nd(dims, glyph_channels, 8, 3, padding=1),
+ nn.SiLU(),
+ conv_nd(dims, 8, 8, 3, padding=1),
+ nn.SiLU(),
+ conv_nd(dims, 8, 16, 3, padding=1, stride=2),
+ nn.SiLU(),
+ conv_nd(dims, 16, 16, 3, padding=1),
+ nn.SiLU(),
+ conv_nd(dims, 16, 32, 3, padding=1, stride=2),
+ nn.SiLU(),
+ conv_nd(dims, 32, 32, 3, padding=1),
+ nn.SiLU(),
+ conv_nd(dims, 32, 96, 3, padding=1, stride=2),
+ nn.SiLU(),
+ conv_nd(dims, 96, 96, 3, padding=1),
+ nn.SiLU(),
+ conv_nd(dims, 96, 256, 3, padding=1, stride=2),
+ nn.SiLU(),
+ )
+
+ self.position_block = TimestepEmbedSequential(
+ conv_nd(dims, position_channels, 8, 3, padding=1),
+ nn.SiLU(),
+ conv_nd(dims, 8, 8, 3, padding=1),
+ nn.SiLU(),
+ conv_nd(dims, 8, 16, 3, padding=1, stride=2),
+ nn.SiLU(),
+ conv_nd(dims, 16, 16, 3, padding=1),
+ nn.SiLU(),
+ conv_nd(dims, 16, 32, 3, padding=1, stride=2),
+ nn.SiLU(),
+ conv_nd(dims, 32, 32, 3, padding=1),
+ nn.SiLU(),
+ conv_nd(dims, 32, 64, 3, padding=1, stride=2),
+ nn.SiLU(),
+ )
+
+ self.fuse_block = zero_module(conv_nd(dims, 256+64+4, model_channels, 3, padding=1))
+
+ self._feature_size = model_channels
+ input_block_chans = [model_channels]
+ ch = model_channels
+ ds = 1
+ for level, mult in enumerate(channel_mult):
+ for nr in range(self.num_res_blocks[level]):
+ layers = [
+ ResBlock(
+ ch,
+ time_embed_dim,
+ dropout,
+ out_channels=mult * model_channels,
+ dims=dims,
+ use_checkpoint=use_checkpoint,
+ use_scale_shift_norm=use_scale_shift_norm,
+ )
+ ]
+ ch = mult * model_channels
+ if ds in attention_resolutions:
+ if num_head_channels == -1:
+ dim_head = ch // num_heads
+ else:
+ num_heads = ch // num_head_channels
+ dim_head = num_head_channels
+ if legacy:
+ # num_heads = 1
+ dim_head = ch // num_heads if use_spatial_transformer else num_head_channels
+ if exists(disable_self_attentions):
+ disabled_sa = disable_self_attentions[level]
+ else:
+ disabled_sa = False
+
+ if not exists(num_attention_blocks) or nr < num_attention_blocks[level]:
+ layers.append(
+ AttentionBlock(
+ ch,
+ use_checkpoint=use_checkpoint,
+ num_heads=num_heads,
+ num_head_channels=dim_head,
+ use_new_attention_order=use_new_attention_order,
+ ) if not use_spatial_transformer else SpatialTransformer(
+ ch, num_heads, dim_head, depth=transformer_depth, context_dim=context_dim,
+ disable_self_attn=disabled_sa, use_linear=use_linear_in_transformer,
+ use_checkpoint=use_checkpoint
+ )
+ )
+ self.input_blocks.append(TimestepEmbedSequential(*layers))
+ self.zero_convs.append(self.make_zero_conv(ch))
+ self._feature_size += ch
+ input_block_chans.append(ch)
+ if level != len(channel_mult) - 1:
+ out_ch = ch
+ self.input_blocks.append(
+ TimestepEmbedSequential(
+ ResBlock(
+ ch,
+ time_embed_dim,
+ dropout,
+ out_channels=out_ch,
+ dims=dims,
+ use_checkpoint=use_checkpoint,
+ use_scale_shift_norm=use_scale_shift_norm,
+ down=True,
+ )
+ if resblock_updown
+ else Downsample(
+ ch, conv_resample, dims=dims, out_channels=out_ch
+ )
+ )
+ )
+ ch = out_ch
+ input_block_chans.append(ch)
+ self.zero_convs.append(self.make_zero_conv(ch))
+ ds *= 2
+ self._feature_size += ch
+
+ if num_head_channels == -1:
+ dim_head = ch // num_heads
+ else:
+ num_heads = ch // num_head_channels
+ dim_head = num_head_channels
+ if legacy:
+ # num_heads = 1
+ dim_head = ch // num_heads if use_spatial_transformer else num_head_channels
+ self.middle_block = TimestepEmbedSequential(
+ ResBlock(
+ ch,
+ time_embed_dim,
+ dropout,
+ dims=dims,
+ use_checkpoint=use_checkpoint,
+ use_scale_shift_norm=use_scale_shift_norm,
+ ),
+ AttentionBlock(
+ ch,
+ use_checkpoint=use_checkpoint,
+ num_heads=num_heads,
+ num_head_channels=dim_head,
+ use_new_attention_order=use_new_attention_order,
+ ) if not use_spatial_transformer else SpatialTransformer( # always uses a self-attn
+ ch, num_heads, dim_head, depth=transformer_depth, context_dim=context_dim,
+ disable_self_attn=disable_middle_self_attn, use_linear=use_linear_in_transformer,
+ use_checkpoint=use_checkpoint
+ ),
+ ResBlock(
+ ch,
+ time_embed_dim,
+ dropout,
+ dims=dims,
+ use_checkpoint=use_checkpoint,
+ use_scale_shift_norm=use_scale_shift_norm,
+ ),
+ )
+ self.middle_block_out = self.make_zero_conv(ch)
+ self._feature_size += ch
+
+ def make_zero_conv(self, channels):
+ return TimestepEmbedSequential(zero_module(conv_nd(self.dims, channels, channels, 1, padding=0)))
+
+ def forward(self, x, hint, text_info, timesteps, context, **kwargs):
+ t_emb = timestep_embedding(timesteps, self.model_channels, repeat_only=False)
+ if self.use_fp16:
+ t_emb = t_emb.half()
+ emb = self.time_embed(t_emb)
+
+ # guided_hint from text_info
+ B, C, H, W = x.shape
+ glyphs = torch.cat(text_info['glyphs'], dim=1).sum(dim=1, keepdim=True)
+ positions = torch.cat(text_info['positions'], dim=1).sum(dim=1, keepdim=True)
+ enc_glyph = self.glyph_block(glyphs, emb, context)
+ enc_pos = self.position_block(positions, emb, context)
+ guided_hint = self.fuse_block(torch.cat([enc_glyph, enc_pos, text_info['masked_x']], dim=1))
+
+ outs = []
+
+ h = x.type(self.dtype)
+ for module, zero_conv in zip(self.input_blocks, self.zero_convs):
+ if guided_hint is not None:
+ h = module(h, emb, context)
+ h += guided_hint
+ guided_hint = None
+ else:
+ h = module(h, emb, context)
+ outs.append(zero_conv(h, emb, context))
+
+ h = self.middle_block(h, emb, context)
+ outs.append(self.middle_block_out(h, emb, context))
+
+ return outs
+
+
+class ControlLDM(LatentDiffusion):
+
+ def __init__(self, control_stage_config, control_key, glyph_key, position_key, only_mid_control, loss_alpha=0, loss_beta=0, with_step_weight=False, use_vae_upsample=False, latin_weight=1.0, embedding_manager_config=None, *args, **kwargs):
+ self.use_fp16 = kwargs.pop('use_fp16', False)
+ super().__init__(*args, **kwargs)
+ self.control_model = instantiate_from_config(control_stage_config)
+ self.control_key = control_key
+ self.glyph_key = glyph_key
+ self.position_key = position_key
+ self.only_mid_control = only_mid_control
+ self.control_scales = [1.0] * 13
+ self.loss_alpha = loss_alpha
+ self.loss_beta = loss_beta
+ self.with_step_weight = with_step_weight
+ self.use_vae_upsample = use_vae_upsample
+ self.latin_weight = latin_weight
+
+ if embedding_manager_config is not None and embedding_manager_config.params.valid:
+ self.embedding_manager = self.instantiate_embedding_manager(embedding_manager_config, self.cond_stage_model)
+ for param in self.embedding_manager.embedding_parameters():
+ param.requires_grad = True
+ else:
+ self.embedding_manager = None
+ if self.loss_alpha > 0 or self.loss_beta > 0 or self.embedding_manager:
+ if embedding_manager_config.params.emb_type == 'ocr':
+ self.text_predictor = create_predictor().eval()
+ args = edict()
+ args.rec_image_shape = "3, 48, 320"
+ args.rec_batch_num = 6
+ args.rec_char_dict_path = str(CURRENT_DIR.parent / "ocr_recog" / "ppocr_keys_v1.txt")
+ args.use_fp16 = self.use_fp16
+ self.cn_recognizer = TextRecognizer(args, self.text_predictor)
+ for param in self.text_predictor.parameters():
+ param.requires_grad = False
+ if self.embedding_manager:
+ self.embedding_manager.recog = self.cn_recognizer
+
+ @torch.no_grad()
+ def get_input(self, batch, k, bs=None, *args, **kwargs):
+ if self.embedding_manager is None: # fill in full caption
+ self.fill_caption(batch)
+ x, c, mx = super().get_input(batch, self.first_stage_key, mask_k='masked_img', *args, **kwargs)
+ control = batch[self.control_key] # for log_images and loss_alpha, not real control
+ if bs is not None:
+ control = control[:bs]
+ control = control.to(self.device)
+ control = einops.rearrange(control, 'b h w c -> b c h w')
+ control = control.to(memory_format=torch.contiguous_format).float()
+
+ inv_mask = batch['inv_mask']
+ if bs is not None:
+ inv_mask = inv_mask[:bs]
+ inv_mask = inv_mask.to(self.device)
+ inv_mask = einops.rearrange(inv_mask, 'b h w c -> b c h w')
+ inv_mask = inv_mask.to(memory_format=torch.contiguous_format).float()
+
+ glyphs = batch[self.glyph_key]
+ gly_line = batch['gly_line']
+ positions = batch[self.position_key]
+ n_lines = batch['n_lines']
+ language = batch['language']
+ texts = batch['texts']
+ assert len(glyphs) == len(positions)
+ for i in range(len(glyphs)):
+ if bs is not None:
+ glyphs[i] = glyphs[i][:bs]
+ gly_line[i] = gly_line[i][:bs]
+ positions[i] = positions[i][:bs]
+ n_lines = n_lines[:bs]
+ glyphs[i] = glyphs[i].to(self.device)
+ gly_line[i] = gly_line[i].to(self.device)
+ positions[i] = positions[i].to(self.device)
+ glyphs[i] = einops.rearrange(glyphs[i], 'b h w c -> b c h w')
+ gly_line[i] = einops.rearrange(gly_line[i], 'b h w c -> b c h w')
+ positions[i] = einops.rearrange(positions[i], 'b h w c -> b c h w')
+ glyphs[i] = glyphs[i].to(memory_format=torch.contiguous_format).float()
+ gly_line[i] = gly_line[i].to(memory_format=torch.contiguous_format).float()
+ positions[i] = positions[i].to(memory_format=torch.contiguous_format).float()
+ info = {}
+ info['glyphs'] = glyphs
+ info['positions'] = positions
+ info['n_lines'] = n_lines
+ info['language'] = language
+ info['texts'] = texts
+ info['img'] = batch['img'] # nhwc, (-1,1)
+ info['masked_x'] = mx
+ info['gly_line'] = gly_line
+ info['inv_mask'] = inv_mask
+ return x, dict(c_crossattn=[c], c_concat=[control], text_info=info)
+
+ def apply_model(self, x_noisy, t, cond, *args, **kwargs):
+ assert isinstance(cond, dict)
+ diffusion_model = self.model.diffusion_model
+ _cond = torch.cat(cond['c_crossattn'], 1)
+ _hint = torch.cat(cond['c_concat'], 1)
+ if self.use_fp16:
+ x_noisy = x_noisy.half()
+ control = self.control_model(x=x_noisy, timesteps=t, context=_cond, hint=_hint, text_info=cond['text_info'])
+ control = [c * scale for c, scale in zip(control, self.control_scales)]
+ eps = diffusion_model(x=x_noisy, timesteps=t, context=_cond, control=control, only_mid_control=self.only_mid_control)
+
+ return eps
+
+ def instantiate_embedding_manager(self, config, embedder):
+ model = instantiate_from_config(config, embedder=embedder)
+ return model
+
+ @torch.no_grad()
+ def get_unconditional_conditioning(self, N):
+ return self.get_learned_conditioning(dict(c_crossattn=[[""] * N], text_info=None))
+
+ def get_learned_conditioning(self, c):
+ if self.cond_stage_forward is None:
+ if hasattr(self.cond_stage_model, 'encode') and callable(self.cond_stage_model.encode):
+ if self.embedding_manager is not None and c['text_info'] is not None:
+ self.embedding_manager.encode_text(c['text_info'])
+ if isinstance(c, dict):
+ cond_txt = c['c_crossattn'][0]
+ else:
+ cond_txt = c
+ if self.embedding_manager is not None:
+ cond_txt = self.cond_stage_model.encode(cond_txt, embedding_manager=self.embedding_manager)
+ else:
+ cond_txt = self.cond_stage_model.encode(cond_txt)
+ if isinstance(c, dict):
+ c['c_crossattn'][0] = cond_txt
+ else:
+ c = cond_txt
+ if isinstance(c, DiagonalGaussianDistribution):
+ c = c.mode()
+ else:
+ c = self.cond_stage_model(c)
+ else:
+ assert hasattr(self.cond_stage_model, self.cond_stage_forward)
+ c = getattr(self.cond_stage_model, self.cond_stage_forward)(c)
+ return c
+
+ def fill_caption(self, batch, place_holder='*'):
+ bs = len(batch['n_lines'])
+ cond_list = copy.deepcopy(batch[self.cond_stage_key])
+ for i in range(bs):
+ n_lines = batch['n_lines'][i]
+ if n_lines == 0:
+ continue
+ cur_cap = cond_list[i]
+ for j in range(n_lines):
+ r_txt = batch['texts'][j][i]
+ cur_cap = cur_cap.replace(place_holder, f'"{r_txt}"', 1)
+ cond_list[i] = cur_cap
+ batch[self.cond_stage_key] = cond_list
+
+ @torch.no_grad()
+ def log_images(self, batch, N=4, n_row=2, sample=False, ddim_steps=50, ddim_eta=0.0, return_keys=None,
+ quantize_denoised=True, inpaint=True, plot_denoise_rows=False, plot_progressive_rows=True,
+ plot_diffusion_rows=False, unconditional_guidance_scale=9.0, unconditional_guidance_label=None,
+ use_ema_scope=True,
+ **kwargs):
+ use_ddim = ddim_steps is not None
+
+ log = dict()
+ z, c = self.get_input(batch, self.first_stage_key, bs=N)
+ if self.cond_stage_trainable:
+ with torch.no_grad():
+ c = self.get_learned_conditioning(c)
+ c_crossattn = c["c_crossattn"][0][:N]
+ c_cat = c["c_concat"][0][:N]
+ text_info = c["text_info"]
+ text_info['glyphs'] = [i[:N] for i in text_info['glyphs']]
+ text_info['gly_line'] = [i[:N] for i in text_info['gly_line']]
+ text_info['positions'] = [i[:N] for i in text_info['positions']]
+ text_info['n_lines'] = text_info['n_lines'][:N]
+ text_info['masked_x'] = text_info['masked_x'][:N]
+ text_info['img'] = text_info['img'][:N]
+
+ N = min(z.shape[0], N)
+ n_row = min(z.shape[0], n_row)
+ log["reconstruction"] = self.decode_first_stage(z)
+ log["masked_image"] = self.decode_first_stage(text_info['masked_x'])
+ log["control"] = c_cat * 2.0 - 1.0
+ log["img"] = text_info['img'].permute(0, 3, 1, 2) # log source image if needed
+ # get glyph
+ glyph_bs = torch.stack(text_info['glyphs'])
+ glyph_bs = torch.sum(glyph_bs, dim=0) * 2.0 - 1.0
+ log["glyph"] = torch.nn.functional.interpolate(glyph_bs, size=(512, 512), mode='bilinear', align_corners=True,)
+ # fill caption
+ if not self.embedding_manager:
+ self.fill_caption(batch)
+ captions = batch[self.cond_stage_key]
+ log["conditioning"] = log_txt_as_img((512, 512), captions, size=16)
+
+ if plot_diffusion_rows:
+ # get diffusion row
+ diffusion_row = list()
+ z_start = z[:n_row]
+ for t in range(self.num_timesteps):
+ if t % self.log_every_t == 0 or t == self.num_timesteps - 1:
+ t = repeat(torch.tensor([t]), '1 -> b', b=n_row)
+ t = t.to(self.device).long()
+ noise = torch.randn_like(z_start)
+ z_noisy = self.q_sample(x_start=z_start, t=t, noise=noise)
+ diffusion_row.append(self.decode_first_stage(z_noisy))
+
+ diffusion_row = torch.stack(diffusion_row) # n_log_step, n_row, C, H, W
+ diffusion_grid = rearrange(diffusion_row, 'n b c h w -> b n c h w')
+ diffusion_grid = rearrange(diffusion_grid, 'b n c h w -> (b n) c h w')
+ diffusion_grid = make_grid(diffusion_grid, nrow=diffusion_row.shape[0])
+ log["diffusion_row"] = diffusion_grid
+
+ if sample:
+ # get denoise row
+ samples, z_denoise_row = self.sample_log(cond={"c_concat": [c_cat], "c_crossattn": [c], "text_info": text_info},
+ batch_size=N, ddim=use_ddim,
+ ddim_steps=ddim_steps, eta=ddim_eta)
+ x_samples = self.decode_first_stage(samples)
+ log["samples"] = x_samples
+ if plot_denoise_rows:
+ denoise_grid = self._get_denoise_row_from_list(z_denoise_row)
+ log["denoise_row"] = denoise_grid
+
+ if unconditional_guidance_scale > 1.0:
+ uc_cross = self.get_unconditional_conditioning(N)
+ uc_cat = c_cat # torch.zeros_like(c_cat)
+ uc_full = {"c_concat": [uc_cat], "c_crossattn": [uc_cross['c_crossattn'][0]], "text_info": text_info}
+ samples_cfg, tmps = self.sample_log(cond={"c_concat": [c_cat], "c_crossattn": [c_crossattn], "text_info": text_info},
+ batch_size=N, ddim=use_ddim,
+ ddim_steps=ddim_steps, eta=ddim_eta,
+ unconditional_guidance_scale=unconditional_guidance_scale,
+ unconditional_conditioning=uc_full,
+ )
+ x_samples_cfg = self.decode_first_stage(samples_cfg)
+ log[f"samples_cfg_scale_{unconditional_guidance_scale:.2f}"] = x_samples_cfg
+ pred_x0 = False # wether log pred_x0
+ if pred_x0:
+ for idx in range(len(tmps['pred_x0'])):
+ pred_x0 = self.decode_first_stage(tmps['pred_x0'][idx])
+ log[f"pred_x0_{tmps['index'][idx]}"] = pred_x0
+
+ return log
+
+ @torch.no_grad()
+ def sample_log(self, cond, batch_size, ddim, ddim_steps, **kwargs):
+ ddim_sampler = DDIMSampler(self)
+ b, c, h, w = cond["c_concat"][0].shape
+ shape = (self.channels, h // 8, w // 8)
+ samples, intermediates = ddim_sampler.sample(ddim_steps, batch_size, shape, cond, verbose=False, log_every_t=5, **kwargs)
+ return samples, intermediates
+
+ def configure_optimizers(self):
+ lr = self.learning_rate
+ params = list(self.control_model.parameters())
+ if self.embedding_manager:
+ params += list(self.embedding_manager.embedding_parameters())
+ if not self.sd_locked:
+ # params += list(self.model.diffusion_model.input_blocks.parameters())
+ # params += list(self.model.diffusion_model.middle_block.parameters())
+ params += list(self.model.diffusion_model.output_blocks.parameters())
+ params += list(self.model.diffusion_model.out.parameters())
+ if self.unlockKV:
+ nCount = 0
+ for name, param in self.model.diffusion_model.named_parameters():
+ if 'attn2.to_k' in name or 'attn2.to_v' in name:
+ params += [param]
+ nCount += 1
+ print(f'Cross attention is unlocked, and {nCount} Wk or Wv are added to potimizers!!!')
+
+ opt = torch.optim.AdamW(params, lr=lr)
+ return opt
+
+ def low_vram_shift(self, is_diffusing):
+ if is_diffusing:
+ self.model = self.model.cuda()
+ self.control_model = self.control_model.cuda()
+ self.first_stage_model = self.first_stage_model.cpu()
+ self.cond_stage_model = self.cond_stage_model.cpu()
+ else:
+ self.model = self.model.cpu()
+ self.control_model = self.control_model.cpu()
+ self.first_stage_model = self.first_stage_model.cuda()
+ self.cond_stage_model = self.cond_stage_model.cuda()
diff --git a/iopaint/model/anytext/cldm/ddim_hacked.py b/iopaint/model/anytext/cldm/ddim_hacked.py
new file mode 100644
index 0000000000000000000000000000000000000000..b23a883e7e4353dc378469de1398aac2f612045d
--- /dev/null
+++ b/iopaint/model/anytext/cldm/ddim_hacked.py
@@ -0,0 +1,486 @@
+"""SAMPLING ONLY."""
+
+import torch
+import numpy as np
+from tqdm import tqdm
+
+from iopaint.model.anytext.ldm.modules.diffusionmodules.util import (
+ make_ddim_sampling_parameters,
+ make_ddim_timesteps,
+ noise_like,
+ extract_into_tensor,
+)
+
+
+class DDIMSampler(object):
+ def __init__(self, model, device, schedule="linear", **kwargs):
+ super().__init__()
+ self.device = device
+ self.model = model
+ self.ddpm_num_timesteps = model.num_timesteps
+ self.schedule = schedule
+
+ def register_buffer(self, name, attr):
+ if type(attr) == torch.Tensor:
+ if attr.device != torch.device(self.device):
+ attr = attr.to(torch.device(self.device))
+ 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,
+ 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.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,
+ S,
+ batch_size,
+ shape,
+ conditioning=None,
+ callback=None,
+ normals_sequence=None,
+ img_callback=None,
+ quantize_x0=False,
+ eta=0.0,
+ mask=None,
+ x0=None,
+ temperature=1.0,
+ noise_dropout=0.0,
+ score_corrector=None,
+ corrector_kwargs=None,
+ verbose=True,
+ x_T=None,
+ log_every_t=100,
+ unconditional_guidance_scale=1.0,
+ unconditional_conditioning=None, # this has to come in the same format as the conditioning, # e.g. as encoded tokens, ...
+ dynamic_threshold=None,
+ ucg_schedule=None,
+ **kwargs,
+ ):
+ if conditioning is not None:
+ if isinstance(conditioning, dict):
+ ctmp = conditioning[list(conditioning.keys())[0]]
+ while isinstance(ctmp, list):
+ ctmp = ctmp[0]
+ cbs = ctmp.shape[0]
+ if cbs != batch_size:
+ print(
+ f"Warning: Got {cbs} conditionings but batch-size is {batch_size}"
+ )
+
+ elif isinstance(conditioning, list):
+ for ctmp in conditioning:
+ if ctmp.shape[0] != 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=S, ddim_eta=eta, verbose=verbose)
+ # sampling
+ C, H, W = shape
+ size = (batch_size, C, H, W)
+ print(f"Data shape for DDIM sampling is {size}, eta {eta}")
+
+ samples, intermediates = self.ddim_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,
+ dynamic_threshold=dynamic_threshold,
+ ucg_schedule=ucg_schedule,
+ )
+ return samples, intermediates
+
+ @torch.no_grad()
+ def ddim_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.0,
+ noise_dropout=0.0,
+ score_corrector=None,
+ corrector_kwargs=None,
+ unconditional_guidance_scale=1.0,
+ unconditional_conditioning=None,
+ dynamic_threshold=None,
+ ucg_schedule=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]
+
+ intermediates = {"x_inter": [img], "pred_x0": [img]}
+ 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]
+ print(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)
+
+ 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.0 - mask) * img
+
+ if ucg_schedule is not None:
+ assert len(ucg_schedule) == len(time_range)
+ unconditional_guidance_scale = ucg_schedule[i]
+
+ 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,
+ score_corrector=score_corrector,
+ corrector_kwargs=corrector_kwargs,
+ unconditional_guidance_scale=unconditional_guidance_scale,
+ unconditional_conditioning=unconditional_conditioning,
+ dynamic_threshold=dynamic_threshold,
+ )
+ img, pred_x0 = outs
+ if callback:
+ callback(None, i, None, None)
+ if img_callback:
+ img_callback(pred_x0, i)
+
+ if index % log_every_t == 0 or index == total_steps - 1:
+ intermediates["x_inter"].append(img)
+ intermediates["pred_x0"].append(pred_x0)
+
+ return img, intermediates
+
+ @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,
+ score_corrector=None,
+ corrector_kwargs=None,
+ unconditional_guidance_scale=1.0,
+ unconditional_conditioning=None,
+ dynamic_threshold=None,
+ ):
+ b, *_, device = *x.shape, x.device
+
+ if unconditional_conditioning is None or unconditional_guidance_scale == 1.0:
+ model_output = self.model.apply_model(x, t, c)
+ else:
+ model_t = self.model.apply_model(x, t, c)
+ model_uncond = self.model.apply_model(x, t, unconditional_conditioning)
+ model_output = model_uncond + unconditional_guidance_scale * (
+ model_t - model_uncond
+ )
+
+ if self.model.parameterization == "v":
+ e_t = self.model.predict_eps_from_z_and_v(x, t, model_output)
+ else:
+ e_t = model_output
+
+ if score_corrector is not None:
+ assert self.model.parameterization == "eps", "not implemented"
+ e_t = score_corrector.modify_score(
+ self.model, e_t, x, t, c, **corrector_kwargs
+ )
+
+ 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
+ if self.model.parameterization != "v":
+ pred_x0 = (x - sqrt_one_minus_at * e_t) / a_t.sqrt()
+ else:
+ pred_x0 = self.model.predict_start_from_z_and_v(x, t, model_output)
+
+ if quantize_denoised:
+ pred_x0, _, *_ = self.model.first_stage_model.quantize(pred_x0)
+
+ if dynamic_threshold is not None:
+ raise NotImplementedError()
+
+ # 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
+
+ @torch.no_grad()
+ def encode(
+ self,
+ x0,
+ c,
+ t_enc,
+ use_original_steps=False,
+ return_intermediates=None,
+ unconditional_guidance_scale=1.0,
+ unconditional_conditioning=None,
+ callback=None,
+ ):
+ timesteps = (
+ np.arange(self.ddpm_num_timesteps)
+ if use_original_steps
+ else self.ddim_timesteps
+ )
+ num_reference_steps = timesteps.shape[0]
+
+ assert t_enc <= num_reference_steps
+ num_steps = t_enc
+
+ if use_original_steps:
+ alphas_next = self.alphas_cumprod[:num_steps]
+ alphas = self.alphas_cumprod_prev[:num_steps]
+ else:
+ alphas_next = self.ddim_alphas[:num_steps]
+ alphas = torch.tensor(self.ddim_alphas_prev[:num_steps])
+
+ x_next = x0
+ intermediates = []
+ inter_steps = []
+ for i in tqdm(range(num_steps), desc="Encoding Image"):
+ t = torch.full(
+ (x0.shape[0],), timesteps[i], device=self.model.device, dtype=torch.long
+ )
+ if unconditional_guidance_scale == 1.0:
+ noise_pred = self.model.apply_model(x_next, t, c)
+ else:
+ assert unconditional_conditioning is not None
+ e_t_uncond, noise_pred = torch.chunk(
+ self.model.apply_model(
+ torch.cat((x_next, x_next)),
+ torch.cat((t, t)),
+ torch.cat((unconditional_conditioning, c)),
+ ),
+ 2,
+ )
+ noise_pred = e_t_uncond + unconditional_guidance_scale * (
+ noise_pred - e_t_uncond
+ )
+
+ xt_weighted = (alphas_next[i] / alphas[i]).sqrt() * x_next
+ weighted_noise_pred = (
+ alphas_next[i].sqrt()
+ * ((1 / alphas_next[i] - 1).sqrt() - (1 / alphas[i] - 1).sqrt())
+ * noise_pred
+ )
+ x_next = xt_weighted + weighted_noise_pred
+ if (
+ return_intermediates
+ and i % (num_steps // return_intermediates) == 0
+ and i < num_steps - 1
+ ):
+ intermediates.append(x_next)
+ inter_steps.append(i)
+ elif return_intermediates and i >= num_steps - 2:
+ intermediates.append(x_next)
+ inter_steps.append(i)
+ if callback:
+ callback(i)
+
+ out = {"x_encoded": x_next, "intermediate_steps": inter_steps}
+ if return_intermediates:
+ out.update({"intermediates": intermediates})
+ return x_next, out
+
+ @torch.no_grad()
+ def stochastic_encode(self, x0, t, use_original_steps=False, noise=None):
+ # fast, but does not allow for exact reconstruction
+ # t serves as an index to gather the correct alphas
+ if use_original_steps:
+ sqrt_alphas_cumprod = self.sqrt_alphas_cumprod
+ sqrt_one_minus_alphas_cumprod = self.sqrt_one_minus_alphas_cumprod
+ else:
+ sqrt_alphas_cumprod = torch.sqrt(self.ddim_alphas)
+ sqrt_one_minus_alphas_cumprod = self.ddim_sqrt_one_minus_alphas
+
+ if noise is None:
+ noise = torch.randn_like(x0)
+ return (
+ extract_into_tensor(sqrt_alphas_cumprod, t, x0.shape) * x0
+ + extract_into_tensor(sqrt_one_minus_alphas_cumprod, t, x0.shape) * noise
+ )
+
+ @torch.no_grad()
+ def decode(
+ self,
+ x_latent,
+ cond,
+ t_start,
+ unconditional_guidance_scale=1.0,
+ unconditional_conditioning=None,
+ use_original_steps=False,
+ callback=None,
+ ):
+ timesteps = (
+ np.arange(self.ddpm_num_timesteps)
+ if use_original_steps
+ else self.ddim_timesteps
+ )
+ timesteps = timesteps[:t_start]
+
+ time_range = np.flip(timesteps)
+ total_steps = timesteps.shape[0]
+ print(f"Running DDIM Sampling with {total_steps} timesteps")
+
+ iterator = tqdm(time_range, desc="Decoding image", total=total_steps)
+ x_dec = x_latent
+ for i, step in enumerate(iterator):
+ index = total_steps - i - 1
+ ts = torch.full(
+ (x_latent.shape[0],), step, device=x_latent.device, dtype=torch.long
+ )
+ x_dec, _ = self.p_sample_ddim(
+ x_dec,
+ cond,
+ ts,
+ index=index,
+ use_original_steps=use_original_steps,
+ unconditional_guidance_scale=unconditional_guidance_scale,
+ unconditional_conditioning=unconditional_conditioning,
+ )
+ if callback:
+ callback(i)
+ return x_dec
diff --git a/iopaint/model/anytext/cldm/embedding_manager.py b/iopaint/model/anytext/cldm/embedding_manager.py
new file mode 100644
index 0000000000000000000000000000000000000000..6ccf8a983d75151f7c858532ffcb9eba254125ee
--- /dev/null
+++ b/iopaint/model/anytext/cldm/embedding_manager.py
@@ -0,0 +1,165 @@
+'''
+Copyright (c) Alibaba, Inc. and its affiliates.
+'''
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from functools import partial
+from iopaint.model.anytext.ldm.modules.diffusionmodules.util import conv_nd, linear
+
+
+def get_clip_token_for_string(tokenizer, string):
+ batch_encoding = tokenizer(string, truncation=True, max_length=77, return_length=True,
+ return_overflowing_tokens=False, padding="max_length", return_tensors="pt")
+ tokens = batch_encoding["input_ids"]
+ assert torch.count_nonzero(tokens - 49407) == 2, f"String '{string}' maps to more than a single token. Please use another string"
+ return tokens[0, 1]
+
+
+def get_bert_token_for_string(tokenizer, string):
+ token = tokenizer(string)
+ assert torch.count_nonzero(token) == 3, f"String '{string}' maps to more than a single token. Please use another string"
+ token = token[0, 1]
+ return token
+
+
+def get_clip_vision_emb(encoder, processor, img):
+ _img = img.repeat(1, 3, 1, 1)*255
+ inputs = processor(images=_img, return_tensors="pt")
+ inputs['pixel_values'] = inputs['pixel_values'].to(img.device)
+ outputs = encoder(**inputs)
+ emb = outputs.image_embeds
+ return emb
+
+
+def get_recog_emb(encoder, img_list):
+ _img_list = [(img.repeat(1, 3, 1, 1)*255)[0] for img in img_list]
+ encoder.predictor.eval()
+ _, preds_neck = encoder.pred_imglist(_img_list, show_debug=False)
+ return preds_neck
+
+
+def pad_H(x):
+ _, _, H, W = x.shape
+ p_top = (W - H) // 2
+ p_bot = W - H - p_top
+ return F.pad(x, (0, 0, p_top, p_bot))
+
+
+class EncodeNet(nn.Module):
+ def __init__(self, in_channels, out_channels):
+ super(EncodeNet, self).__init__()
+ chan = 16
+ n_layer = 4 # downsample
+
+ self.conv1 = conv_nd(2, in_channels, chan, 3, padding=1)
+ self.conv_list = nn.ModuleList([])
+ _c = chan
+ for i in range(n_layer):
+ self.conv_list.append(conv_nd(2, _c, _c*2, 3, padding=1, stride=2))
+ _c *= 2
+ self.conv2 = conv_nd(2, _c, out_channels, 3, padding=1)
+ self.avgpool = nn.AdaptiveAvgPool2d(1)
+ self.act = nn.SiLU()
+
+ def forward(self, x):
+ x = self.act(self.conv1(x))
+ for layer in self.conv_list:
+ x = self.act(layer(x))
+ x = self.act(self.conv2(x))
+ x = self.avgpool(x)
+ x = x.view(x.size(0), -1)
+ return x
+
+
+class EmbeddingManager(nn.Module):
+ def __init__(
+ self,
+ embedder,
+ valid=True,
+ glyph_channels=20,
+ position_channels=1,
+ placeholder_string='*',
+ add_pos=False,
+ emb_type='ocr',
+ **kwargs
+ ):
+ super().__init__()
+ if hasattr(embedder, 'tokenizer'): # using Stable Diffusion's CLIP encoder
+ get_token_for_string = partial(get_clip_token_for_string, embedder.tokenizer)
+ token_dim = 768
+ if hasattr(embedder, 'vit'):
+ assert emb_type == 'vit'
+ self.get_vision_emb = partial(get_clip_vision_emb, embedder.vit, embedder.processor)
+ self.get_recog_emb = None
+ else: # using LDM's BERT encoder
+ get_token_for_string = partial(get_bert_token_for_string, embedder.tknz_fn)
+ token_dim = 1280
+ self.token_dim = token_dim
+ self.emb_type = emb_type
+
+ self.add_pos = add_pos
+ if add_pos:
+ self.position_encoder = EncodeNet(position_channels, token_dim)
+ if emb_type == 'ocr':
+ self.proj = linear(40*64, token_dim)
+ if emb_type == 'conv':
+ self.glyph_encoder = EncodeNet(glyph_channels, token_dim)
+
+ self.placeholder_token = get_token_for_string(placeholder_string)
+
+ def encode_text(self, text_info):
+ if self.get_recog_emb is None and self.emb_type == 'ocr':
+ self.get_recog_emb = partial(get_recog_emb, self.recog)
+
+ gline_list = []
+ pos_list = []
+ for i in range(len(text_info['n_lines'])): # sample index in a batch
+ n_lines = text_info['n_lines'][i]
+ for j in range(n_lines): # line
+ gline_list += [text_info['gly_line'][j][i:i+1]]
+ if self.add_pos:
+ pos_list += [text_info['positions'][j][i:i+1]]
+
+ if len(gline_list) > 0:
+ if self.emb_type == 'ocr':
+ recog_emb = self.get_recog_emb(gline_list)
+ enc_glyph = self.proj(recog_emb.reshape(recog_emb.shape[0], -1))
+ elif self.emb_type == 'vit':
+ enc_glyph = self.get_vision_emb(pad_H(torch.cat(gline_list, dim=0)))
+ elif self.emb_type == 'conv':
+ enc_glyph = self.glyph_encoder(pad_H(torch.cat(gline_list, dim=0)))
+ if self.add_pos:
+ enc_pos = self.position_encoder(torch.cat(gline_list, dim=0))
+ enc_glyph = enc_glyph+enc_pos
+
+ self.text_embs_all = []
+ n_idx = 0
+ for i in range(len(text_info['n_lines'])): # sample index in a batch
+ n_lines = text_info['n_lines'][i]
+ text_embs = []
+ for j in range(n_lines): # line
+ text_embs += [enc_glyph[n_idx:n_idx+1]]
+ n_idx += 1
+ self.text_embs_all += [text_embs]
+
+ def forward(
+ self,
+ tokenized_text,
+ embedded_text,
+ ):
+ b, device = tokenized_text.shape[0], tokenized_text.device
+ for i in range(b):
+ idx = tokenized_text[i] == self.placeholder_token.to(device)
+ if sum(idx) > 0:
+ if i >= len(self.text_embs_all):
+ print('truncation for log images...')
+ break
+ text_emb = torch.cat(self.text_embs_all[i], dim=0)
+ if sum(idx) != len(text_emb):
+ print('truncation for long caption...')
+ embedded_text[i][idx] = text_emb[:sum(idx)]
+ return embedded_text
+
+ def embedding_parameters(self):
+ return self.parameters()
diff --git a/iopaint/model/anytext/cldm/hack.py b/iopaint/model/anytext/cldm/hack.py
new file mode 100644
index 0000000000000000000000000000000000000000..05afe5f2513b6191f0e7b2f33f6ff80518a215d3
--- /dev/null
+++ b/iopaint/model/anytext/cldm/hack.py
@@ -0,0 +1,111 @@
+import torch
+import einops
+
+import iopaint.model.anytext.ldm.modules.encoders.modules
+import iopaint.model.anytext.ldm.modules.attention
+
+from transformers import logging
+from iopaint.model.anytext.ldm.modules.attention import default
+
+
+def disable_verbosity():
+ logging.set_verbosity_error()
+ print('logging improved.')
+ return
+
+
+def enable_sliced_attention():
+ iopaint.model.anytext.ldm.modules.attention.CrossAttention.forward = _hacked_sliced_attentin_forward
+ print('Enabled sliced_attention.')
+ return
+
+
+def hack_everything(clip_skip=0):
+ disable_verbosity()
+ iopaint.model.anytext.ldm.modules.encoders.modules.FrozenCLIPEmbedder.forward = _hacked_clip_forward
+ iopaint.model.anytext.ldm.modules.encoders.modules.FrozenCLIPEmbedder.clip_skip = clip_skip
+ print('Enabled clip hacks.')
+ return
+
+
+# Written by Lvmin
+def _hacked_clip_forward(self, text):
+ PAD = self.tokenizer.pad_token_id
+ EOS = self.tokenizer.eos_token_id
+ BOS = self.tokenizer.bos_token_id
+
+ def tokenize(t):
+ return self.tokenizer(t, truncation=False, add_special_tokens=False)["input_ids"]
+
+ def transformer_encode(t):
+ if self.clip_skip > 1:
+ rt = self.transformer(input_ids=t, output_hidden_states=True)
+ return self.transformer.text_model.final_layer_norm(rt.hidden_states[-self.clip_skip])
+ else:
+ return self.transformer(input_ids=t, output_hidden_states=False).last_hidden_state
+
+ def split(x):
+ return x[75 * 0: 75 * 1], x[75 * 1: 75 * 2], x[75 * 2: 75 * 3]
+
+ def pad(x, p, i):
+ return x[:i] if len(x) >= i else x + [p] * (i - len(x))
+
+ raw_tokens_list = tokenize(text)
+ tokens_list = []
+
+ for raw_tokens in raw_tokens_list:
+ raw_tokens_123 = split(raw_tokens)
+ raw_tokens_123 = [[BOS] + raw_tokens_i + [EOS] for raw_tokens_i in raw_tokens_123]
+ raw_tokens_123 = [pad(raw_tokens_i, PAD, 77) for raw_tokens_i in raw_tokens_123]
+ tokens_list.append(raw_tokens_123)
+
+ tokens_list = torch.IntTensor(tokens_list).to(self.device)
+
+ feed = einops.rearrange(tokens_list, 'b f i -> (b f) i')
+ y = transformer_encode(feed)
+ z = einops.rearrange(y, '(b f) i c -> b (f i) c', f=3)
+
+ return z
+
+
+# Stolen from https://github.com/basujindal/stable-diffusion/blob/main/optimizedSD/splitAttention.py
+def _hacked_sliced_attentin_forward(self, x, context=None, mask=None):
+ h = self.heads
+
+ q = self.to_q(x)
+ context = default(context, x)
+ k = self.to_k(context)
+ v = self.to_v(context)
+ del context, x
+
+ q, k, v = map(lambda t: einops.rearrange(t, 'b n (h d) -> (b h) n d', h=h), (q, k, v))
+
+ limit = k.shape[0]
+ att_step = 1
+ q_chunks = list(torch.tensor_split(q, limit // att_step, dim=0))
+ k_chunks = list(torch.tensor_split(k, limit // att_step, dim=0))
+ v_chunks = list(torch.tensor_split(v, limit // att_step, dim=0))
+
+ q_chunks.reverse()
+ k_chunks.reverse()
+ v_chunks.reverse()
+ sim = torch.zeros(q.shape[0], q.shape[1], v.shape[2], device=q.device)
+ del k, q, v
+ for i in range(0, limit, att_step):
+ q_buffer = q_chunks.pop()
+ k_buffer = k_chunks.pop()
+ v_buffer = v_chunks.pop()
+ sim_buffer = torch.einsum('b i d, b j d -> b i j', q_buffer, k_buffer) * self.scale
+
+ del k_buffer, q_buffer
+ # attention, what we cannot get enough of, by chunks
+
+ sim_buffer = sim_buffer.softmax(dim=-1)
+
+ sim_buffer = torch.einsum('b i j, b j d -> b i d', sim_buffer, v_buffer)
+ del v_buffer
+ sim[i:i + att_step, :, :] = sim_buffer
+
+ del sim_buffer
+ sim = einops.rearrange(sim, '(b h) n d -> b n (h d)', h=h)
+ return self.to_out(sim)
diff --git a/iopaint/model/anytext/cldm/model.py b/iopaint/model/anytext/cldm/model.py
new file mode 100644
index 0000000000000000000000000000000000000000..688f2edd1d04aedf500988a46c141a489dddf5e6
--- /dev/null
+++ b/iopaint/model/anytext/cldm/model.py
@@ -0,0 +1,40 @@
+import os
+import torch
+
+from omegaconf import OmegaConf
+from iopaint.model.anytext.ldm.util import instantiate_from_config
+
+
+def get_state_dict(d):
+ return d.get("state_dict", d)
+
+
+def load_state_dict(ckpt_path, location="cpu"):
+ _, extension = os.path.splitext(ckpt_path)
+ if extension.lower() == ".safetensors":
+ import safetensors.torch
+
+ state_dict = safetensors.torch.load_file(ckpt_path, device=location)
+ else:
+ state_dict = get_state_dict(
+ torch.load(ckpt_path, map_location=torch.device(location))
+ )
+ state_dict = get_state_dict(state_dict)
+ print(f"Loaded state_dict from [{ckpt_path}]")
+ return state_dict
+
+
+def create_model(config_path, device, cond_stage_path=None, use_fp16=False):
+ config = OmegaConf.load(config_path)
+ # if cond_stage_path:
+ # config.model.params.cond_stage_config.params.version = (
+ # cond_stage_path # use pre-downloaded ckpts, in case blocked
+ # )
+ config.model.params.cond_stage_config.params.device = str(device)
+ if use_fp16:
+ config.model.params.use_fp16 = True
+ config.model.params.control_stage_config.params.use_fp16 = True
+ config.model.params.unet_config.params.use_fp16 = True
+ model = instantiate_from_config(config.model).cpu()
+ print(f"Loaded model config from [{config_path}]")
+ return model
diff --git a/iopaint/model/anytext/ldm/models/diffusion/ddim.py b/iopaint/model/anytext/ldm/models/diffusion/ddim.py
new file mode 100644
index 0000000000000000000000000000000000000000..f8bbaff872ed59066cecf478f04f28857d77a305
--- /dev/null
+++ b/iopaint/model/anytext/ldm/models/diffusion/ddim.py
@@ -0,0 +1,354 @@
+"""SAMPLING ONLY."""
+
+import torch
+import numpy as np
+from tqdm import tqdm
+
+from iopaint.model.anytext.ldm.modules.diffusionmodules.util import make_ddim_sampling_parameters, make_ddim_timesteps, noise_like, extract_into_tensor
+
+
+class DDIMSampler(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):
+ if type(attr) == torch.Tensor:
+ if attr.device != torch.device("cuda"):
+ attr = attr.to(torch.device("cuda"))
+ setattr(self, name, attr)
+
+ def make_schedule(self, ddim_num_steps, ddim_discretize="uniform", ddim_eta=0., verbose=True):
+ 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,
+ S,
+ 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=True,
+ 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, ...
+ dynamic_threshold=None,
+ ucg_schedule=None,
+ **kwargs
+ ):
+ if conditioning is not None:
+ if isinstance(conditioning, dict):
+ ctmp = conditioning[list(conditioning.keys())[0]]
+ while isinstance(ctmp, list): ctmp = ctmp[0]
+ cbs = ctmp.shape[0]
+ # cbs = len(ctmp[0])
+ if cbs != batch_size:
+ print(f"Warning: Got {cbs} conditionings but batch-size is {batch_size}")
+
+ elif isinstance(conditioning, list):
+ for ctmp in conditioning:
+ if ctmp.shape[0] != 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=S, ddim_eta=eta, verbose=verbose)
+ # sampling
+ C, H, W = shape
+ size = (batch_size, C, H, W)
+ print(f'Data shape for DDIM sampling is {size}, eta {eta}')
+
+ samples, intermediates = self.ddim_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,
+ dynamic_threshold=dynamic_threshold,
+ ucg_schedule=ucg_schedule
+ )
+ return samples, intermediates
+
+ @torch.no_grad()
+ def ddim_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, dynamic_threshold=None,
+ ucg_schedule=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]
+
+ intermediates = {'x_inter': [img], 'pred_x0': [img], "index": [10000]}
+ 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]
+ print(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)
+
+ 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
+
+ if ucg_schedule is not None:
+ assert len(ucg_schedule) == len(time_range)
+ unconditional_guidance_scale = ucg_schedule[i]
+
+ 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, score_corrector=score_corrector,
+ corrector_kwargs=corrector_kwargs,
+ unconditional_guidance_scale=unconditional_guidance_scale,
+ unconditional_conditioning=unconditional_conditioning,
+ dynamic_threshold=dynamic_threshold)
+ img, pred_x0 = outs
+ if callback:
+ callback(i)
+ if img_callback:
+ img_callback(pred_x0, i)
+
+ if index % log_every_t == 0 or index == total_steps - 1:
+ intermediates['x_inter'].append(img)
+ intermediates['pred_x0'].append(pred_x0)
+ intermediates['index'].append(index)
+
+ return img, intermediates
+
+ @torch.no_grad()
+ def p_sample_ddim(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,
+ dynamic_threshold=None):
+ b, *_, device = *x.shape, x.device
+
+ if unconditional_conditioning is None or unconditional_guidance_scale == 1.:
+ model_output = self.model.apply_model(x, t, c)
+ else:
+ x_in = torch.cat([x] * 2)
+ t_in = torch.cat([t] * 2)
+ if isinstance(c, dict):
+ assert isinstance(unconditional_conditioning, dict)
+ c_in = dict()
+ for k in c:
+ if isinstance(c[k], list):
+ c_in[k] = [torch.cat([
+ unconditional_conditioning[k][i],
+ c[k][i]]) for i in range(len(c[k]))]
+ elif isinstance(c[k], dict):
+ c_in[k] = dict()
+ for key in c[k]:
+ if isinstance(c[k][key], list):
+ if not isinstance(c[k][key][0], torch.Tensor):
+ continue
+ c_in[k][key] = [torch.cat([
+ unconditional_conditioning[k][key][i],
+ c[k][key][i]]) for i in range(len(c[k][key]))]
+ else:
+ c_in[k][key] = torch.cat([
+ unconditional_conditioning[k][key],
+ c[k][key]])
+
+ else:
+ c_in[k] = torch.cat([
+ unconditional_conditioning[k],
+ c[k]])
+ elif isinstance(c, list):
+ c_in = list()
+ assert isinstance(unconditional_conditioning, list)
+ for i in range(len(c)):
+ c_in.append(torch.cat([unconditional_conditioning[i], c[i]]))
+ else:
+ c_in = torch.cat([unconditional_conditioning, c])
+ model_uncond, model_t = self.model.apply_model(x_in, t_in, c_in).chunk(2)
+ model_output = model_uncond + unconditional_guidance_scale * (model_t - model_uncond)
+
+ if self.model.parameterization == "v":
+ e_t = self.model.predict_eps_from_z_and_v(x, t, model_output)
+ else:
+ e_t = model_output
+
+ if score_corrector is not None:
+ assert self.model.parameterization == "eps", 'not implemented'
+ e_t = score_corrector.modify_score(self.model, e_t, x, t, c, **corrector_kwargs)
+
+ 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
+ if self.model.parameterization != "v":
+ pred_x0 = (x - sqrt_one_minus_at * e_t) / a_t.sqrt()
+ else:
+ pred_x0 = self.model.predict_start_from_z_and_v(x, t, model_output)
+
+ if quantize_denoised:
+ pred_x0, _, *_ = self.model.first_stage_model.quantize(pred_x0)
+
+ if dynamic_threshold is not None:
+ raise NotImplementedError()
+
+ # 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
+
+ @torch.no_grad()
+ def encode(self, x0, c, t_enc, use_original_steps=False, return_intermediates=None,
+ unconditional_guidance_scale=1.0, unconditional_conditioning=None, callback=None):
+ num_reference_steps = self.ddpm_num_timesteps if use_original_steps else self.ddim_timesteps.shape[0]
+
+ assert t_enc <= num_reference_steps
+ num_steps = t_enc
+
+ if use_original_steps:
+ alphas_next = self.alphas_cumprod[:num_steps]
+ alphas = self.alphas_cumprod_prev[:num_steps]
+ else:
+ alphas_next = self.ddim_alphas[:num_steps]
+ alphas = torch.tensor(self.ddim_alphas_prev[:num_steps])
+
+ x_next = x0
+ intermediates = []
+ inter_steps = []
+ for i in tqdm(range(num_steps), desc='Encoding Image'):
+ t = torch.full((x0.shape[0],), i, device=self.model.device, dtype=torch.long)
+ if unconditional_guidance_scale == 1.:
+ noise_pred = self.model.apply_model(x_next, t, c)
+ else:
+ assert unconditional_conditioning is not None
+ e_t_uncond, noise_pred = torch.chunk(
+ self.model.apply_model(torch.cat((x_next, x_next)), torch.cat((t, t)),
+ torch.cat((unconditional_conditioning, c))), 2)
+ noise_pred = e_t_uncond + unconditional_guidance_scale * (noise_pred - e_t_uncond)
+
+ xt_weighted = (alphas_next[i] / alphas[i]).sqrt() * x_next
+ weighted_noise_pred = alphas_next[i].sqrt() * (
+ (1 / alphas_next[i] - 1).sqrt() - (1 / alphas[i] - 1).sqrt()) * noise_pred
+ x_next = xt_weighted + weighted_noise_pred
+ if return_intermediates and i % (
+ num_steps // return_intermediates) == 0 and i < num_steps - 1:
+ intermediates.append(x_next)
+ inter_steps.append(i)
+ elif return_intermediates and i >= num_steps - 2:
+ intermediates.append(x_next)
+ inter_steps.append(i)
+ if callback: callback(i)
+
+ out = {'x_encoded': x_next, 'intermediate_steps': inter_steps}
+ if return_intermediates:
+ out.update({'intermediates': intermediates})
+ return x_next, out
+
+ @torch.no_grad()
+ def stochastic_encode(self, x0, t, use_original_steps=False, noise=None):
+ # fast, but does not allow for exact reconstruction
+ # t serves as an index to gather the correct alphas
+ if use_original_steps:
+ sqrt_alphas_cumprod = self.sqrt_alphas_cumprod
+ sqrt_one_minus_alphas_cumprod = self.sqrt_one_minus_alphas_cumprod
+ else:
+ sqrt_alphas_cumprod = torch.sqrt(self.ddim_alphas)
+ sqrt_one_minus_alphas_cumprod = self.ddim_sqrt_one_minus_alphas
+
+ if noise is None:
+ noise = torch.randn_like(x0)
+ return (extract_into_tensor(sqrt_alphas_cumprod, t, x0.shape) * x0 +
+ extract_into_tensor(sqrt_one_minus_alphas_cumprod, t, x0.shape) * noise)
+
+ @torch.no_grad()
+ def decode(self, x_latent, cond, t_start, unconditional_guidance_scale=1.0, unconditional_conditioning=None,
+ use_original_steps=False, callback=None):
+
+ timesteps = np.arange(self.ddpm_num_timesteps) if use_original_steps else self.ddim_timesteps
+ timesteps = timesteps[:t_start]
+
+ time_range = np.flip(timesteps)
+ total_steps = timesteps.shape[0]
+ print(f"Running DDIM Sampling with {total_steps} timesteps")
+
+ iterator = tqdm(time_range, desc='Decoding image', total=total_steps)
+ x_dec = x_latent
+ for i, step in enumerate(iterator):
+ index = total_steps - i - 1
+ ts = torch.full((x_latent.shape[0],), step, device=x_latent.device, dtype=torch.long)
+ x_dec, _ = self.p_sample_ddim(x_dec, cond, ts, index=index, use_original_steps=use_original_steps,
+ unconditional_guidance_scale=unconditional_guidance_scale,
+ unconditional_conditioning=unconditional_conditioning)
+ if callback: callback(i)
+ return x_dec
\ No newline at end of file
diff --git a/iopaint/model/anytext/ldm/models/diffusion/ddpm.py b/iopaint/model/anytext/ldm/models/diffusion/ddpm.py
new file mode 100644
index 0000000000000000000000000000000000000000..9f489180e99f8350aa8ca31228a17598422c89be
--- /dev/null
+++ b/iopaint/model/anytext/ldm/models/diffusion/ddpm.py
@@ -0,0 +1,2380 @@
+"""
+Part of the implementation is borrowed and modified from ControlNet, publicly available at https://github.com/lllyasviel/ControlNet/blob/main/ldm/models/diffusion/ddpm.py
+"""
+
+import torch
+import torch.nn as nn
+import numpy as np
+from torch.optim.lr_scheduler import LambdaLR
+from einops import rearrange, repeat
+from contextlib import contextmanager, nullcontext
+from functools import partial
+import itertools
+from tqdm import tqdm
+from torchvision.utils import make_grid
+from omegaconf import ListConfig
+
+from iopaint.model.anytext.ldm.util import (
+ log_txt_as_img,
+ exists,
+ default,
+ ismap,
+ isimage,
+ mean_flat,
+ count_params,
+ instantiate_from_config,
+)
+from iopaint.model.anytext.ldm.modules.ema import LitEma
+from iopaint.model.anytext.ldm.modules.distributions.distributions import (
+ normal_kl,
+ DiagonalGaussianDistribution,
+)
+from iopaint.model.anytext.ldm.models.autoencoder import IdentityFirstStage, AutoencoderKL
+from iopaint.model.anytext.ldm.modules.diffusionmodules.util import (
+ make_beta_schedule,
+ extract_into_tensor,
+ noise_like,
+)
+from iopaint.model.anytext.ldm.models.diffusion.ddim import DDIMSampler
+import cv2
+
+
+__conditioning_keys__ = {"concat": "c_concat", "crossattn": "c_crossattn", "adm": "y"}
+
+PRINT_DEBUG = False
+
+
+def print_grad(grad):
+ # print('Gradient:', grad)
+ # print(grad.shape)
+ a = grad.max()
+ b = grad.min()
+ # print(f'mean={grad.mean():.4f}, max={a:.4f}, min={b:.4f}')
+ s = 255.0 / (a - b)
+ c = 255 * (-b / (a - b))
+ grad = grad * s + c
+ # print(f'mean={grad.mean():.4f}, max={grad.max():.4f}, min={grad.min():.4f}')
+ img = grad[0].permute(1, 2, 0).detach().cpu().numpy()
+ if img.shape[0] == 512:
+ cv2.imwrite("grad-img.jpg", img)
+ elif img.shape[0] == 64:
+ cv2.imwrite("grad-latent.jpg", img)
+
+
+def disabled_train(self, mode=True):
+ """Overwrite model.train with this function to make sure train/eval mode
+ does not change anymore."""
+ return self
+
+
+def uniform_on_device(r1, r2, shape, device):
+ return (r1 - r2) * torch.rand(*shape, device=device) + r2
+
+
+class DDPM(torch.nn.Module):
+ # classic DDPM with Gaussian diffusion, in image space
+ def __init__(
+ self,
+ unet_config,
+ timesteps=1000,
+ beta_schedule="linear",
+ loss_type="l2",
+ ckpt_path=None,
+ ignore_keys=[],
+ load_only_unet=False,
+ monitor="val/loss",
+ use_ema=True,
+ first_stage_key="image",
+ image_size=256,
+ channels=3,
+ log_every_t=100,
+ clip_denoised=True,
+ linear_start=1e-4,
+ linear_end=2e-2,
+ cosine_s=8e-3,
+ given_betas=None,
+ 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,
+ conditioning_key=None,
+ parameterization="eps", # all assuming fixed variance schedules
+ scheduler_config=None,
+ use_positional_encodings=False,
+ learn_logvar=False,
+ logvar_init=0.0,
+ make_it_fit=False,
+ ucg_training=None,
+ reset_ema=False,
+ reset_num_ema_updates=False,
+ ):
+ super().__init__()
+ assert parameterization in [
+ "eps",
+ "x0",
+ "v",
+ ], 'currently only supporting "eps" and "x0" and "v"'
+ self.parameterization = parameterization
+ print(
+ f"{self.__class__.__name__}: Running in {self.parameterization}-prediction mode"
+ )
+ self.cond_stage_model = None
+ self.clip_denoised = clip_denoised
+ self.log_every_t = log_every_t
+ self.first_stage_key = first_stage_key
+ self.image_size = image_size # try conv?
+ self.channels = channels
+ self.use_positional_encodings = use_positional_encodings
+ self.model = DiffusionWrapper(unet_config, conditioning_key)
+ count_params(self.model, verbose=True)
+ self.use_ema = use_ema
+ if self.use_ema:
+ self.model_ema = LitEma(self.model)
+ print(f"Keeping EMAs of {len(list(self.model_ema.buffers()))}.")
+
+ self.use_scheduler = scheduler_config is not None
+ if self.use_scheduler:
+ self.scheduler_config = scheduler_config
+
+ self.v_posterior = v_posterior
+ self.original_elbo_weight = original_elbo_weight
+ self.l_simple_weight = l_simple_weight
+
+ if monitor is not None:
+ self.monitor = monitor
+ self.make_it_fit = make_it_fit
+ if reset_ema:
+ assert exists(ckpt_path)
+ if ckpt_path is not None:
+ self.init_from_ckpt(
+ ckpt_path, ignore_keys=ignore_keys, only_model=load_only_unet
+ )
+ if reset_ema:
+ assert self.use_ema
+ print(
+ f"Resetting ema to pure model weights. This is useful when restoring from an ema-only checkpoint."
+ )
+ self.model_ema = LitEma(self.model)
+ if reset_num_ema_updates:
+ print(
+ " +++++++++++ WARNING: RESETTING NUM_EMA UPDATES TO ZERO +++++++++++ "
+ )
+ assert self.use_ema
+ self.model_ema.reset_num_updates()
+
+ self.register_schedule(
+ given_betas=given_betas,
+ beta_schedule=beta_schedule,
+ timesteps=timesteps,
+ linear_start=linear_start,
+ linear_end=linear_end,
+ cosine_s=cosine_s,
+ )
+
+ self.loss_type = loss_type
+
+ self.learn_logvar = learn_logvar
+ logvar = torch.full(fill_value=logvar_init, size=(self.num_timesteps,))
+ if self.learn_logvar:
+ self.logvar = nn.Parameter(self.logvar, requires_grad=True)
+ else:
+ self.register_buffer("logvar", logvar)
+
+ self.ucg_training = ucg_training or dict()
+ if self.ucg_training:
+ self.ucg_prng = np.random.RandomState()
+
+ def register_schedule(
+ self,
+ given_betas=None,
+ beta_schedule="linear",
+ timesteps=1000,
+ linear_start=1e-4,
+ linear_end=2e-2,
+ cosine_s=8e-3,
+ ):
+ if exists(given_betas):
+ betas = given_betas
+ else:
+ betas = make_beta_schedule(
+ 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)
+ # np.save('1.npy', alphas_cumprod)
+ 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 = partial(torch.tensor, dtype=torch.float32)
+
+ 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))
+ )
+ elif self.parameterization == "v":
+ lvlb_weights = torch.ones_like(
+ self.betas**2
+ / (
+ 2
+ * self.posterior_variance
+ * to_torch(alphas)
+ * (1 - self.alphas_cumprod)
+ )
+ )
+ else:
+ raise NotImplementedError("mu not supported")
+ lvlb_weights[0] = lvlb_weights[1]
+ self.register_buffer("lvlb_weights", lvlb_weights, persistent=False)
+ assert not torch.isnan(self.lvlb_weights).all()
+
+ @contextmanager
+ def ema_scope(self, context=None):
+ if self.use_ema:
+ self.model_ema.store(self.model.parameters())
+ self.model_ema.copy_to(self.model)
+ if context is not None:
+ print(f"{context}: Switched to EMA weights")
+ try:
+ yield None
+ finally:
+ if self.use_ema:
+ self.model_ema.restore(self.model.parameters())
+ if context is not None:
+ print(f"{context}: Restored training weights")
+
+ @torch.no_grad()
+ def init_from_ckpt(self, path, ignore_keys=list(), only_model=False):
+ sd = torch.load(path, map_location="cpu")
+ if "state_dict" in list(sd.keys()):
+ sd = sd["state_dict"]
+ keys = list(sd.keys())
+ for k in keys:
+ for ik in ignore_keys:
+ if k.startswith(ik):
+ print("Deleting key {} from state_dict.".format(k))
+ del sd[k]
+ if self.make_it_fit:
+ n_params = len(
+ [
+ name
+ for name, _ in itertools.chain(
+ self.named_parameters(), self.named_buffers()
+ )
+ ]
+ )
+ for name, param in tqdm(
+ itertools.chain(self.named_parameters(), self.named_buffers()),
+ desc="Fitting old weights to new weights",
+ total=n_params,
+ ):
+ if not name in sd:
+ continue
+ old_shape = sd[name].shape
+ new_shape = param.shape
+ assert len(old_shape) == len(new_shape)
+ if len(new_shape) > 2:
+ # we only modify first two axes
+ assert new_shape[2:] == old_shape[2:]
+ # assumes first axis corresponds to output dim
+ if not new_shape == old_shape:
+ new_param = param.clone()
+ old_param = sd[name]
+ if len(new_shape) == 1:
+ for i in range(new_param.shape[0]):
+ new_param[i] = old_param[i % old_shape[0]]
+ elif len(new_shape) >= 2:
+ for i in range(new_param.shape[0]):
+ for j in range(new_param.shape[1]):
+ new_param[i, j] = old_param[
+ i % old_shape[0], j % old_shape[1]
+ ]
+
+ n_used_old = torch.ones(old_shape[1])
+ for j in range(new_param.shape[1]):
+ n_used_old[j % old_shape[1]] += 1
+ n_used_new = torch.zeros(new_shape[1])
+ for j in range(new_param.shape[1]):
+ n_used_new[j] = n_used_old[j % old_shape[1]]
+
+ n_used_new = n_used_new[None, :]
+ while len(n_used_new.shape) < len(new_shape):
+ n_used_new = n_used_new.unsqueeze(-1)
+ new_param /= n_used_new
+
+ sd[name] = new_param
+
+ missing, unexpected = (
+ self.load_state_dict(sd, strict=False)
+ if not only_model
+ else self.model.load_state_dict(sd, strict=False)
+ )
+ print(
+ f"Restored from {path} with {len(missing)} missing and {len(unexpected)} unexpected keys"
+ )
+ if len(missing) > 0:
+ print(f"Missing Keys:\n {missing}")
+ if len(unexpected) > 0:
+ print(f"\nUnexpected Keys:\n {unexpected}")
+
+ def q_mean_variance(self, x_start, t):
+ """
+ Get the distribution q(x_t | x_0).
+ :param x_start: the [N x C x ...] tensor of noiseless inputs.
+ :param t: the number of diffusion steps (minus 1). Here, 0 means one step.
+ :return: A tuple (mean, variance, log_variance), all of x_start's shape.
+ """
+ mean = extract_into_tensor(self.sqrt_alphas_cumprod, t, x_start.shape) * x_start
+ variance = extract_into_tensor(1.0 - self.alphas_cumprod, t, x_start.shape)
+ log_variance = extract_into_tensor(
+ self.log_one_minus_alphas_cumprod, t, x_start.shape
+ )
+ return mean, variance, log_variance
+
+ def predict_start_from_noise(self, x_t, t, noise):
+ return (
+ extract_into_tensor(self.sqrt_recip_alphas_cumprod, t, x_t.shape) * x_t
+ - extract_into_tensor(self.sqrt_recipm1_alphas_cumprod, t, x_t.shape)
+ * noise
+ )
+
+ def predict_start_from_z_and_v(self, x_t, t, v):
+ # 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. - alphas_cumprod)))
+ return (
+ extract_into_tensor(self.sqrt_alphas_cumprod, t, x_t.shape) * x_t
+ - extract_into_tensor(self.sqrt_one_minus_alphas_cumprod, t, x_t.shape) * v
+ )
+
+ def predict_eps_from_z_and_v(self, x_t, t, v):
+ return (
+ extract_into_tensor(self.sqrt_alphas_cumprod, t, x_t.shape) * v
+ + extract_into_tensor(self.sqrt_one_minus_alphas_cumprod, t, x_t.shape)
+ * x_t
+ )
+
+ def q_posterior(self, x_start, x_t, t):
+ posterior_mean = (
+ extract_into_tensor(self.posterior_mean_coef1, t, x_t.shape) * x_start
+ + extract_into_tensor(self.posterior_mean_coef2, t, x_t.shape) * x_t
+ )
+ posterior_variance = extract_into_tensor(self.posterior_variance, t, x_t.shape)
+ posterior_log_variance_clipped = extract_into_tensor(
+ self.posterior_log_variance_clipped, t, x_t.shape
+ )
+ return posterior_mean, posterior_variance, posterior_log_variance_clipped
+
+ def p_mean_variance(self, x, t, clip_denoised: bool):
+ model_out = self.model(x, t)
+ if self.parameterization == "eps":
+ x_recon = self.predict_start_from_noise(x, t=t, noise=model_out)
+ elif self.parameterization == "x0":
+ x_recon = model_out
+ if clip_denoised:
+ x_recon.clamp_(-1.0, 1.0)
+
+ model_mean, posterior_variance, posterior_log_variance = self.q_posterior(
+ x_start=x_recon, x_t=x, t=t
+ )
+ return model_mean, posterior_variance, posterior_log_variance
+
+ @torch.no_grad()
+ def p_sample(self, x, t, clip_denoised=True, repeat_noise=False):
+ b, *_, device = *x.shape, x.device
+ model_mean, _, model_log_variance = self.p_mean_variance(
+ x=x, t=t, clip_denoised=clip_denoised
+ )
+ noise = noise_like(x.shape, device, repeat_noise)
+ # no noise when t == 0
+ nonzero_mask = (1 - (t == 0).float()).reshape(b, *((1,) * (len(x.shape) - 1)))
+ return model_mean + nonzero_mask * (0.5 * model_log_variance).exp() * noise
+
+ @torch.no_grad()
+ def p_sample_loop(self, shape, return_intermediates=False):
+ device = self.betas.device
+ b = shape[0]
+ img = torch.randn(shape, device=device)
+ intermediates = [img]
+ for i in tqdm(
+ reversed(range(0, self.num_timesteps)),
+ desc="Sampling t",
+ total=self.num_timesteps,
+ ):
+ img = self.p_sample(
+ img,
+ torch.full((b,), i, device=device, dtype=torch.long),
+ clip_denoised=self.clip_denoised,
+ )
+ if i % self.log_every_t == 0 or i == self.num_timesteps - 1:
+ intermediates.append(img)
+ if return_intermediates:
+ return img, intermediates
+ return img
+
+ @torch.no_grad()
+ def sample(self, batch_size=16, return_intermediates=False):
+ image_size = self.image_size
+ channels = self.channels
+ return self.p_sample_loop(
+ (batch_size, channels, image_size, image_size),
+ return_intermediates=return_intermediates,
+ )
+
+ def q_sample(self, x_start, t, noise=None):
+ noise = default(noise, lambda: torch.randn_like(x_start))
+ return (
+ extract_into_tensor(self.sqrt_alphas_cumprod, t, x_start.shape) * x_start
+ + extract_into_tensor(self.sqrt_one_minus_alphas_cumprod, t, x_start.shape)
+ * noise
+ )
+
+ def get_v(self, x, noise, t):
+ return (
+ extract_into_tensor(self.sqrt_alphas_cumprod, t, x.shape) * noise
+ - extract_into_tensor(self.sqrt_one_minus_alphas_cumprod, t, x.shape) * x
+ )
+
+ def get_loss(self, pred, target, mean=True):
+ if self.loss_type == "l1":
+ loss = (target - pred).abs()
+ if mean:
+ loss = loss.mean()
+ elif self.loss_type == "l2":
+ if mean:
+ loss = torch.nn.functional.mse_loss(target, pred)
+ else:
+ loss = torch.nn.functional.mse_loss(target, pred, reduction="none")
+ else:
+ raise NotImplementedError("unknown loss type '{loss_type}'")
+
+ return loss
+
+ def p_losses(self, x_start, t, noise=None):
+ noise = default(noise, lambda: torch.randn_like(x_start))
+ x_noisy = self.q_sample(x_start=x_start, t=t, noise=noise)
+ model_out = self.model(x_noisy, t)
+
+ loss_dict = {}
+ if self.parameterization == "eps":
+ target = noise
+ elif self.parameterization == "x0":
+ target = x_start
+ elif self.parameterization == "v":
+ target = self.get_v(x_start, noise, t)
+ else:
+ raise NotImplementedError(
+ f"Parameterization {self.parameterization} not yet supported"
+ )
+
+ loss = self.get_loss(model_out, target, mean=False).mean(dim=[1, 2, 3])
+
+ log_prefix = "train" if self.training else "val"
+
+ loss_dict.update({f"{log_prefix}/loss_simple": loss.mean()})
+ loss_simple = loss.mean() * self.l_simple_weight
+
+ loss_vlb = (self.lvlb_weights[t] * loss).mean()
+ loss_dict.update({f"{log_prefix}/loss_vlb": loss_vlb})
+
+ loss = loss_simple + self.original_elbo_weight * loss_vlb
+
+ loss_dict.update({f"{log_prefix}/loss": loss})
+
+ return loss, loss_dict
+
+ def forward(self, x, *args, **kwargs):
+ # b, c, h, w, device, img_size, = *x.shape, x.device, self.image_size
+ # assert h == img_size and w == img_size, f'height and width of image must be {img_size}'
+ t = torch.randint(
+ 0, self.num_timesteps, (x.shape[0],), device=self.device
+ ).long()
+ return self.p_losses(x, t, *args, **kwargs)
+
+ def get_input(self, batch, k):
+ x = batch[k]
+ if len(x.shape) == 3:
+ x = x[..., None]
+ x = rearrange(x, "b h w c -> b c h w")
+ x = x.to(memory_format=torch.contiguous_format).float()
+ return x
+
+ def shared_step(self, batch):
+ x = self.get_input(batch, self.first_stage_key)
+ loss, loss_dict = self(x)
+ return loss, loss_dict
+
+ def training_step(self, batch, batch_idx):
+ for k in self.ucg_training:
+ p = self.ucg_training[k]["p"]
+ val = self.ucg_training[k]["val"]
+ if val is None:
+ val = ""
+ for i in range(len(batch[k])):
+ if self.ucg_prng.choice(2, p=[1 - p, p]):
+ batch[k][i] = val
+
+ loss, loss_dict = self.shared_step(batch)
+
+ self.log_dict(
+ loss_dict, prog_bar=True, logger=True, on_step=True, on_epoch=True
+ )
+
+ self.log(
+ "global_step",
+ self.global_step,
+ prog_bar=True,
+ logger=True,
+ on_step=True,
+ on_epoch=False,
+ )
+
+ if self.use_scheduler:
+ lr = self.optimizers().param_groups[0]["lr"]
+ self.log(
+ "lr_abs", lr, prog_bar=True, logger=True, on_step=True, on_epoch=False
+ )
+
+ return loss
+
+ @torch.no_grad()
+ def validation_step(self, batch, batch_idx):
+ _, loss_dict_no_ema = self.shared_step(batch)
+ with self.ema_scope():
+ _, loss_dict_ema = self.shared_step(batch)
+ loss_dict_ema = {key + "_ema": loss_dict_ema[key] for key in loss_dict_ema}
+ self.log_dict(
+ loss_dict_no_ema, prog_bar=False, logger=True, on_step=False, on_epoch=True
+ )
+ self.log_dict(
+ loss_dict_ema, prog_bar=False, logger=True, on_step=False, on_epoch=True
+ )
+
+ def on_train_batch_end(self, *args, **kwargs):
+ if self.use_ema:
+ self.model_ema(self.model)
+
+ def _get_rows_from_list(self, samples):
+ n_imgs_per_row = len(samples)
+ denoise_grid = rearrange(samples, "n b c h w -> b n c h w")
+ denoise_grid = rearrange(denoise_grid, "b n c h w -> (b n) c h w")
+ denoise_grid = make_grid(denoise_grid, nrow=n_imgs_per_row)
+ return denoise_grid
+
+ @torch.no_grad()
+ def log_images(self, batch, N=8, n_row=2, sample=True, return_keys=None, **kwargs):
+ log = dict()
+ x = self.get_input(batch, self.first_stage_key)
+ N = min(x.shape[0], N)
+ n_row = min(x.shape[0], n_row)
+ x = x.to(self.device)[:N]
+ log["inputs"] = x
+
+ # get diffusion row
+ diffusion_row = list()
+ x_start = x[:n_row]
+
+ for t in range(self.num_timesteps):
+ if t % self.log_every_t == 0 or t == self.num_timesteps - 1:
+ t = repeat(torch.tensor([t]), "1 -> b", b=n_row)
+ t = t.to(self.device).long()
+ noise = torch.randn_like(x_start)
+ x_noisy = self.q_sample(x_start=x_start, t=t, noise=noise)
+ diffusion_row.append(x_noisy)
+
+ log["diffusion_row"] = self._get_rows_from_list(diffusion_row)
+
+ if sample:
+ # get denoise row
+ with self.ema_scope("Plotting"):
+ samples, denoise_row = self.sample(
+ batch_size=N, return_intermediates=True
+ )
+
+ log["samples"] = samples
+ log["denoise_row"] = self._get_rows_from_list(denoise_row)
+
+ if return_keys:
+ if np.intersect1d(list(log.keys()), return_keys).shape[0] == 0:
+ return log
+ else:
+ return {key: log[key] for key in return_keys}
+ return log
+
+ def configure_optimizers(self):
+ lr = self.learning_rate
+ params = list(self.model.parameters())
+ if self.learn_logvar:
+ params = params + [self.logvar]
+ opt = torch.optim.AdamW(params, lr=lr)
+ return opt
+
+
+class LatentDiffusion(DDPM):
+ """main class"""
+
+ def __init__(
+ self,
+ first_stage_config,
+ cond_stage_config,
+ num_timesteps_cond=None,
+ cond_stage_key="image",
+ cond_stage_trainable=False,
+ concat_mode=True,
+ cond_stage_forward=None,
+ conditioning_key=None,
+ scale_factor=1.0,
+ scale_by_std=False,
+ force_null_conditioning=False,
+ *args,
+ **kwargs,
+ ):
+ self.force_null_conditioning = force_null_conditioning
+ self.num_timesteps_cond = default(num_timesteps_cond, 1)
+ self.scale_by_std = scale_by_std
+ assert self.num_timesteps_cond <= kwargs["timesteps"]
+ # for backwards compatibility after implementation of DiffusionWrapper
+ if conditioning_key is None:
+ conditioning_key = "concat" if concat_mode else "crossattn"
+ if (
+ cond_stage_config == "__is_unconditional__"
+ and not self.force_null_conditioning
+ ):
+ conditioning_key = None
+ ckpt_path = kwargs.pop("ckpt_path", None)
+ reset_ema = kwargs.pop("reset_ema", False)
+ reset_num_ema_updates = kwargs.pop("reset_num_ema_updates", False)
+ ignore_keys = kwargs.pop("ignore_keys", [])
+ super().__init__(conditioning_key=conditioning_key, *args, **kwargs)
+ self.concat_mode = concat_mode
+ self.cond_stage_trainable = cond_stage_trainable
+ self.cond_stage_key = cond_stage_key
+ try:
+ self.num_downs = len(first_stage_config.params.ddconfig.ch_mult) - 1
+ except:
+ self.num_downs = 0
+ if not scale_by_std:
+ self.scale_factor = scale_factor
+ else:
+ self.register_buffer("scale_factor", torch.tensor(scale_factor))
+ self.instantiate_first_stage(first_stage_config)
+ self.instantiate_cond_stage(cond_stage_config)
+ self.cond_stage_forward = cond_stage_forward
+ self.clip_denoised = False
+ self.bbox_tokenizer = None
+
+ self.restarted_from_ckpt = False
+ if ckpt_path is not None:
+ self.init_from_ckpt(ckpt_path, ignore_keys)
+ self.restarted_from_ckpt = True
+ if reset_ema:
+ assert self.use_ema
+ print(
+ f"Resetting ema to pure model weights. This is useful when restoring from an ema-only checkpoint."
+ )
+ self.model_ema = LitEma(self.model)
+ if reset_num_ema_updates:
+ print(
+ " +++++++++++ WARNING: RESETTING NUM_EMA UPDATES TO ZERO +++++++++++ "
+ )
+ assert self.use_ema
+ self.model_ema.reset_num_updates()
+
+ 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
+
+ @torch.no_grad()
+ def on_train_batch_start(self, batch, batch_idx, dataloader_idx):
+ # only for very first batch
+ if (
+ self.scale_by_std
+ and self.current_epoch == 0
+ and self.global_step == 0
+ and batch_idx == 0
+ and not self.restarted_from_ckpt
+ ):
+ assert (
+ self.scale_factor == 1.0
+ ), "rather not use custom rescaling and std-rescaling simultaneously"
+ # set rescale weight to 1./std of encodings
+ print("### USING STD-RESCALING ###")
+ x = super().get_input(batch, self.first_stage_key)
+ x = x.to(self.device)
+ encoder_posterior = self.encode_first_stage(x)
+ z = self.get_first_stage_encoding(encoder_posterior).detach()
+ del self.scale_factor
+ self.register_buffer("scale_factor", 1.0 / z.flatten().std())
+ print(f"setting self.scale_factor to {self.scale_factor}")
+ print("### USING STD-RESCALING ###")
+
+ 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 instantiate_first_stage(self, config):
+ model = instantiate_from_config(config)
+ self.first_stage_model = model.eval()
+ self.first_stage_model.train = disabled_train
+ for param in self.first_stage_model.parameters():
+ param.requires_grad = False
+
+ def instantiate_cond_stage(self, config):
+ if not self.cond_stage_trainable:
+ if config == "__is_first_stage__":
+ print("Using first stage also as cond stage.")
+ self.cond_stage_model = self.first_stage_model
+ elif config == "__is_unconditional__":
+ print(f"Training {self.__class__.__name__} as an unconditional model.")
+ self.cond_stage_model = None
+ # self.be_unconditional = True
+ else:
+ model = instantiate_from_config(config)
+ self.cond_stage_model = model.eval()
+ self.cond_stage_model.train = disabled_train
+ for param in self.cond_stage_model.parameters():
+ param.requires_grad = False
+ else:
+ assert config != "__is_first_stage__"
+ assert config != "__is_unconditional__"
+ model = instantiate_from_config(config)
+ self.cond_stage_model = model
+
+ def _get_denoise_row_from_list(
+ self, samples, desc="", force_no_decoder_quantization=False
+ ):
+ denoise_row = []
+ for zd in tqdm(samples, desc=desc):
+ denoise_row.append(
+ self.decode_first_stage(
+ zd.to(self.device), force_not_quantize=force_no_decoder_quantization
+ )
+ )
+ n_imgs_per_row = len(denoise_row)
+ denoise_row = torch.stack(denoise_row) # n_log_step, n_row, C, H, W
+ denoise_grid = rearrange(denoise_row, "n b c h w -> b n c h w")
+ denoise_grid = rearrange(denoise_grid, "b n c h w -> (b n) c h w")
+ denoise_grid = make_grid(denoise_grid, nrow=n_imgs_per_row)
+ return denoise_grid
+
+ def get_first_stage_encoding(self, encoder_posterior):
+ if isinstance(encoder_posterior, DiagonalGaussianDistribution):
+ z = encoder_posterior.sample()
+ elif isinstance(encoder_posterior, torch.Tensor):
+ z = encoder_posterior
+ else:
+ raise NotImplementedError(
+ f"encoder_posterior of type '{type(encoder_posterior)}' not yet implemented"
+ )
+ return self.scale_factor * z
+
+ def get_learned_conditioning(self, c):
+ if self.cond_stage_forward is None:
+ if hasattr(self.cond_stage_model, "encode") and callable(
+ self.cond_stage_model.encode
+ ):
+ c = self.cond_stage_model.encode(c)
+ if isinstance(c, DiagonalGaussianDistribution):
+ c = c.mode()
+ else:
+ c = self.cond_stage_model(c)
+ else:
+ assert hasattr(self.cond_stage_model, self.cond_stage_forward)
+ c = getattr(self.cond_stage_model, self.cond_stage_forward)(c)
+ return c
+
+ def meshgrid(self, h, w):
+ y = torch.arange(0, h).view(h, 1, 1).repeat(1, w, 1)
+ x = torch.arange(0, w).view(1, w, 1).repeat(h, 1, 1)
+
+ arr = torch.cat([y, x], dim=-1)
+ return arr
+
+ def delta_border(self, h, w):
+ """
+ :param h: height
+ :param w: width
+ :return: normalized distance to image border,
+ wtith min distance = 0 at border and max dist = 0.5 at image center
+ """
+ lower_right_corner = torch.tensor([h - 1, w - 1]).view(1, 1, 2)
+ arr = self.meshgrid(h, w) / lower_right_corner
+ dist_left_up = torch.min(arr, dim=-1, keepdims=True)[0]
+ dist_right_down = torch.min(1 - arr, dim=-1, keepdims=True)[0]
+ edge_dist = torch.min(
+ torch.cat([dist_left_up, dist_right_down], dim=-1), dim=-1
+ )[0]
+ return edge_dist
+
+ def get_weighting(self, h, w, Ly, Lx, device):
+ weighting = self.delta_border(h, w)
+ weighting = torch.clip(
+ weighting,
+ self.split_input_params["clip_min_weight"],
+ self.split_input_params["clip_max_weight"],
+ )
+ weighting = weighting.view(1, h * w, 1).repeat(1, 1, Ly * Lx).to(device)
+
+ if self.split_input_params["tie_braker"]:
+ L_weighting = self.delta_border(Ly, Lx)
+ L_weighting = torch.clip(
+ L_weighting,
+ self.split_input_params["clip_min_tie_weight"],
+ self.split_input_params["clip_max_tie_weight"],
+ )
+
+ L_weighting = L_weighting.view(1, 1, Ly * Lx).to(device)
+ weighting = weighting * L_weighting
+ return weighting
+
+ def get_fold_unfold(
+ self, x, kernel_size, stride, uf=1, df=1
+ ): # todo load once not every time, shorten code
+ """
+ :param x: img of size (bs, c, h, w)
+ :return: n img crops of size (n, bs, c, kernel_size[0], kernel_size[1])
+ """
+ bs, nc, h, w = x.shape
+
+ # number of crops in image
+ Ly = (h - kernel_size[0]) // stride[0] + 1
+ Lx = (w - kernel_size[1]) // stride[1] + 1
+
+ if uf == 1 and df == 1:
+ fold_params = dict(
+ kernel_size=kernel_size, dilation=1, padding=0, stride=stride
+ )
+ unfold = torch.nn.Unfold(**fold_params)
+
+ fold = torch.nn.Fold(output_size=x.shape[2:], **fold_params)
+
+ weighting = self.get_weighting(
+ kernel_size[0], kernel_size[1], Ly, Lx, x.device
+ ).to(x.dtype)
+ normalization = fold(weighting).view(1, 1, h, w) # normalizes the overlap
+ weighting = weighting.view((1, 1, kernel_size[0], kernel_size[1], Ly * Lx))
+
+ elif uf > 1 and df == 1:
+ fold_params = dict(
+ kernel_size=kernel_size, dilation=1, padding=0, stride=stride
+ )
+ unfold = torch.nn.Unfold(**fold_params)
+
+ fold_params2 = dict(
+ kernel_size=(kernel_size[0] * uf, kernel_size[0] * uf),
+ dilation=1,
+ padding=0,
+ stride=(stride[0] * uf, stride[1] * uf),
+ )
+ fold = torch.nn.Fold(
+ output_size=(x.shape[2] * uf, x.shape[3] * uf), **fold_params2
+ )
+
+ weighting = self.get_weighting(
+ kernel_size[0] * uf, kernel_size[1] * uf, Ly, Lx, x.device
+ ).to(x.dtype)
+ normalization = fold(weighting).view(
+ 1, 1, h * uf, w * uf
+ ) # normalizes the overlap
+ weighting = weighting.view(
+ (1, 1, kernel_size[0] * uf, kernel_size[1] * uf, Ly * Lx)
+ )
+
+ elif df > 1 and uf == 1:
+ fold_params = dict(
+ kernel_size=kernel_size, dilation=1, padding=0, stride=stride
+ )
+ unfold = torch.nn.Unfold(**fold_params)
+
+ fold_params2 = dict(
+ kernel_size=(kernel_size[0] // df, kernel_size[0] // df),
+ dilation=1,
+ padding=0,
+ stride=(stride[0] // df, stride[1] // df),
+ )
+ fold = torch.nn.Fold(
+ output_size=(x.shape[2] // df, x.shape[3] // df), **fold_params2
+ )
+
+ weighting = self.get_weighting(
+ kernel_size[0] // df, kernel_size[1] // df, Ly, Lx, x.device
+ ).to(x.dtype)
+ normalization = fold(weighting).view(
+ 1, 1, h // df, w // df
+ ) # normalizes the overlap
+ weighting = weighting.view(
+ (1, 1, kernel_size[0] // df, kernel_size[1] // df, Ly * Lx)
+ )
+
+ else:
+ raise NotImplementedError
+
+ return fold, unfold, normalization, weighting
+
+ @torch.no_grad()
+ def get_input(
+ self,
+ batch,
+ k,
+ return_first_stage_outputs=False,
+ force_c_encode=False,
+ cond_key=None,
+ return_original_cond=False,
+ bs=None,
+ return_x=False,
+ mask_k=None,
+ ):
+ x = super().get_input(batch, k)
+ if bs is not None:
+ x = x[:bs]
+ x = x.to(self.device)
+ encoder_posterior = self.encode_first_stage(x)
+ z = self.get_first_stage_encoding(encoder_posterior).detach()
+
+ if mask_k is not None:
+ mx = super().get_input(batch, mask_k)
+ if bs is not None:
+ mx = mx[:bs]
+ mx = mx.to(self.device)
+ encoder_posterior = self.encode_first_stage(mx)
+ mx = self.get_first_stage_encoding(encoder_posterior).detach()
+
+ if self.model.conditioning_key is not None and not self.force_null_conditioning:
+ if cond_key is None:
+ cond_key = self.cond_stage_key
+ if cond_key != self.first_stage_key:
+ if cond_key in ["caption", "coordinates_bbox", "txt"]:
+ xc = batch[cond_key]
+ elif cond_key in ["class_label", "cls"]:
+ xc = batch
+ else:
+ xc = super().get_input(batch, cond_key).to(self.device)
+ else:
+ xc = x
+ if not self.cond_stage_trainable or force_c_encode:
+ if isinstance(xc, dict) or isinstance(xc, list):
+ c = self.get_learned_conditioning(xc)
+ else:
+ c = self.get_learned_conditioning(xc.to(self.device))
+ else:
+ c = xc
+ if bs is not None:
+ c = c[:bs]
+
+ if self.use_positional_encodings:
+ pos_x, pos_y = self.compute_latent_shifts(batch)
+ ckey = __conditioning_keys__[self.model.conditioning_key]
+ c = {ckey: c, "pos_x": pos_x, "pos_y": pos_y}
+
+ else:
+ c = None
+ xc = None
+ if self.use_positional_encodings:
+ pos_x, pos_y = self.compute_latent_shifts(batch)
+ c = {"pos_x": pos_x, "pos_y": pos_y}
+ out = [z, c]
+ if return_first_stage_outputs:
+ xrec = self.decode_first_stage(z)
+ out.extend([x, xrec])
+ if return_x:
+ out.extend([x])
+ if return_original_cond:
+ out.append(xc)
+ if mask_k:
+ out.append(mx)
+ return out
+
+ @torch.no_grad()
+ def decode_first_stage(self, z, predict_cids=False, force_not_quantize=False):
+ if predict_cids:
+ if z.dim() == 4:
+ z = torch.argmax(z.exp(), dim=1).long()
+ z = self.first_stage_model.quantize.get_codebook_entry(z, shape=None)
+ z = rearrange(z, "b h w c -> b c h w").contiguous()
+
+ z = 1.0 / self.scale_factor * z
+ return self.first_stage_model.decode(z)
+
+ def decode_first_stage_grad(self, z, predict_cids=False, force_not_quantize=False):
+ if predict_cids:
+ if z.dim() == 4:
+ z = torch.argmax(z.exp(), dim=1).long()
+ z = self.first_stage_model.quantize.get_codebook_entry(z, shape=None)
+ z = rearrange(z, "b h w c -> b c h w").contiguous()
+
+ z = 1.0 / self.scale_factor * z
+ return self.first_stage_model.decode(z)
+
+ @torch.no_grad()
+ def encode_first_stage(self, x):
+ return self.first_stage_model.encode(x)
+
+ def shared_step(self, batch, **kwargs):
+ x, c = self.get_input(batch, self.first_stage_key)
+ loss = self(x, c)
+ return loss
+
+ def forward(self, x, c, *args, **kwargs):
+ t = torch.randint(
+ 0, self.num_timesteps, (x.shape[0],), device=self.device
+ ).long()
+ # t = torch.randint(500, 501, (x.shape[0],), device=self.device).long()
+ if self.model.conditioning_key is not None:
+ assert c is not None
+ if self.cond_stage_trainable:
+ c = self.get_learned_conditioning(c)
+ if self.shorten_cond_schedule: # TODO: drop this option
+ tc = self.cond_ids[t].to(self.device)
+ c = self.q_sample(x_start=c, t=tc, noise=torch.randn_like(c.float()))
+ return self.p_losses(x, c, t, *args, **kwargs)
+
+ def apply_model(self, x_noisy, t, cond, return_ids=False):
+ if isinstance(cond, dict):
+ # hybrid case, cond is expected to be a dict
+ pass
+ else:
+ if not isinstance(cond, list):
+ cond = [cond]
+ key = (
+ "c_concat" if self.model.conditioning_key == "concat" else "c_crossattn"
+ )
+ cond = {key: cond}
+
+ x_recon = self.model(x_noisy, t, **cond)
+
+ if isinstance(x_recon, tuple) and not return_ids:
+ return x_recon[0]
+ else:
+ return x_recon
+
+ def _predict_eps_from_xstart(self, x_t, t, pred_xstart):
+ return (
+ extract_into_tensor(self.sqrt_recip_alphas_cumprod, t, x_t.shape) * x_t
+ - pred_xstart
+ ) / extract_into_tensor(self.sqrt_recipm1_alphas_cumprod, t, x_t.shape)
+
+ def _prior_bpd(self, x_start):
+ """
+ Get the prior KL term for the variational lower-bound, measured in
+ bits-per-dim.
+ This term can't be optimized, as it only depends on the encoder.
+ :param x_start: the [N x C x ...] tensor of inputs.
+ :return: a batch of [N] KL values (in bits), one per batch element.
+ """
+ batch_size = x_start.shape[0]
+ t = torch.tensor([self.num_timesteps - 1] * batch_size, device=x_start.device)
+ qt_mean, _, qt_log_variance = self.q_mean_variance(x_start, t)
+ kl_prior = normal_kl(
+ mean1=qt_mean, logvar1=qt_log_variance, mean2=0.0, logvar2=0.0
+ )
+ return mean_flat(kl_prior) / np.log(2.0)
+
+ def p_mean_variance(
+ self,
+ x,
+ c,
+ t,
+ clip_denoised: bool,
+ return_codebook_ids=False,
+ quantize_denoised=False,
+ return_x0=False,
+ score_corrector=None,
+ corrector_kwargs=None,
+ ):
+ t_in = t
+ model_out = self.apply_model(x, t_in, c, return_ids=return_codebook_ids)
+
+ if score_corrector is not None:
+ assert self.parameterization == "eps"
+ model_out = score_corrector.modify_score(
+ self, model_out, x, t, c, **corrector_kwargs
+ )
+
+ if return_codebook_ids:
+ model_out, logits = model_out
+
+ if self.parameterization == "eps":
+ x_recon = self.predict_start_from_noise(x, t=t, noise=model_out)
+ elif self.parameterization == "x0":
+ x_recon = model_out
+ else:
+ raise NotImplementedError()
+
+ if clip_denoised:
+ x_recon.clamp_(-1.0, 1.0)
+ if quantize_denoised:
+ x_recon, _, [_, _, indices] = self.first_stage_model.quantize(x_recon)
+ model_mean, posterior_variance, posterior_log_variance = self.q_posterior(
+ x_start=x_recon, x_t=x, t=t
+ )
+ if return_codebook_ids:
+ return model_mean, posterior_variance, posterior_log_variance, logits
+ elif return_x0:
+ return model_mean, posterior_variance, posterior_log_variance, x_recon
+ else:
+ return model_mean, posterior_variance, posterior_log_variance
+
+ @torch.no_grad()
+ def p_sample(
+ self,
+ x,
+ c,
+ t,
+ clip_denoised=False,
+ repeat_noise=False,
+ return_codebook_ids=False,
+ quantize_denoised=False,
+ return_x0=False,
+ temperature=1.0,
+ noise_dropout=0.0,
+ score_corrector=None,
+ corrector_kwargs=None,
+ ):
+ b, *_, device = *x.shape, x.device
+ outputs = self.p_mean_variance(
+ x=x,
+ c=c,
+ t=t,
+ clip_denoised=clip_denoised,
+ return_codebook_ids=return_codebook_ids,
+ quantize_denoised=quantize_denoised,
+ return_x0=return_x0,
+ score_corrector=score_corrector,
+ corrector_kwargs=corrector_kwargs,
+ )
+ if return_codebook_ids:
+ raise DeprecationWarning("Support dropped.")
+ model_mean, _, model_log_variance, logits = outputs
+ elif return_x0:
+ model_mean, _, model_log_variance, x0 = outputs
+ else:
+ model_mean, _, model_log_variance = outputs
+
+ noise = noise_like(x.shape, device, repeat_noise) * temperature
+ if noise_dropout > 0.0:
+ noise = torch.nn.functional.dropout(noise, p=noise_dropout)
+ # no noise when t == 0
+ nonzero_mask = (1 - (t == 0).float()).reshape(b, *((1,) * (len(x.shape) - 1)))
+
+ if return_codebook_ids:
+ return model_mean + nonzero_mask * (
+ 0.5 * model_log_variance
+ ).exp() * noise, logits.argmax(dim=1)
+ if return_x0:
+ return (
+ model_mean + nonzero_mask * (0.5 * model_log_variance).exp() * noise,
+ x0,
+ )
+ else:
+ return model_mean + nonzero_mask * (0.5 * model_log_variance).exp() * noise
+
+ @torch.no_grad()
+ def progressive_denoising(
+ self,
+ cond,
+ shape,
+ verbose=True,
+ callback=None,
+ quantize_denoised=False,
+ img_callback=None,
+ mask=None,
+ x0=None,
+ temperature=1.0,
+ noise_dropout=0.0,
+ score_corrector=None,
+ corrector_kwargs=None,
+ batch_size=None,
+ x_T=None,
+ start_T=None,
+ log_every_t=None,
+ ):
+ if not log_every_t:
+ log_every_t = self.log_every_t
+ timesteps = self.num_timesteps
+ if batch_size is not None:
+ b = batch_size if batch_size is not None else shape[0]
+ shape = [batch_size] + list(shape)
+ else:
+ b = batch_size = shape[0]
+ if x_T is None:
+ img = torch.randn(shape, device=self.device)
+ else:
+ img = x_T
+ intermediates = []
+ if cond is not None:
+ if isinstance(cond, dict):
+ cond = {
+ key: cond[key][:batch_size]
+ if not isinstance(cond[key], list)
+ else list(map(lambda x: x[:batch_size], cond[key]))
+ for key in cond
+ }
+ else:
+ cond = (
+ [c[:batch_size] for c in cond]
+ if isinstance(cond, list)
+ else cond[:batch_size]
+ )
+
+ if start_T is not None:
+ timesteps = min(timesteps, start_T)
+ iterator = (
+ tqdm(
+ reversed(range(0, timesteps)),
+ desc="Progressive Generation",
+ total=timesteps,
+ )
+ if verbose
+ else reversed(range(0, timesteps))
+ )
+ if type(temperature) == float:
+ temperature = [temperature] * timesteps
+
+ for i in iterator:
+ ts = torch.full((b,), i, device=self.device, dtype=torch.long)
+ if self.shorten_cond_schedule:
+ assert self.model.conditioning_key != "hybrid"
+ tc = self.cond_ids[ts].to(cond.device)
+ cond = self.q_sample(x_start=cond, t=tc, noise=torch.randn_like(cond))
+
+ img, x0_partial = self.p_sample(
+ img,
+ cond,
+ ts,
+ clip_denoised=self.clip_denoised,
+ quantize_denoised=quantize_denoised,
+ return_x0=True,
+ temperature=temperature[i],
+ noise_dropout=noise_dropout,
+ score_corrector=score_corrector,
+ corrector_kwargs=corrector_kwargs,
+ )
+ if mask is not None:
+ assert x0 is not None
+ img_orig = self.q_sample(x0, ts)
+ img = img_orig * mask + (1.0 - mask) * img
+
+ if i % log_every_t == 0 or i == timesteps - 1:
+ intermediates.append(x0_partial)
+ if callback:
+ callback(i)
+ if img_callback:
+ img_callback(img, i)
+ return img, intermediates
+
+ @torch.no_grad()
+ def p_sample_loop(
+ self,
+ cond,
+ shape,
+ return_intermediates=False,
+ x_T=None,
+ verbose=True,
+ callback=None,
+ timesteps=None,
+ quantize_denoised=False,
+ mask=None,
+ x0=None,
+ img_callback=None,
+ start_T=None,
+ log_every_t=None,
+ ):
+ if not log_every_t:
+ log_every_t = self.log_every_t
+ device = self.betas.device
+ b = shape[0]
+ if x_T is None:
+ img = torch.randn(shape, device=device)
+ else:
+ img = x_T
+
+ intermediates = [img]
+ if timesteps is None:
+ timesteps = self.num_timesteps
+
+ if start_T is not None:
+ timesteps = min(timesteps, start_T)
+ iterator = (
+ tqdm(reversed(range(0, timesteps)), desc="Sampling t", total=timesteps)
+ if verbose
+ else reversed(range(0, timesteps))
+ )
+
+ if mask is not None:
+ assert x0 is not None
+ assert x0.shape[2:3] == mask.shape[2:3] # spatial size has to match
+
+ for i in iterator:
+ ts = torch.full((b,), i, device=device, dtype=torch.long)
+ if self.shorten_cond_schedule:
+ assert self.model.conditioning_key != "hybrid"
+ tc = self.cond_ids[ts].to(cond.device)
+ cond = self.q_sample(x_start=cond, t=tc, noise=torch.randn_like(cond))
+
+ img = self.p_sample(
+ img,
+ cond,
+ ts,
+ clip_denoised=self.clip_denoised,
+ quantize_denoised=quantize_denoised,
+ )
+ if mask is not None:
+ img_orig = self.q_sample(x0, ts)
+ img = img_orig * mask + (1.0 - mask) * img
+
+ if i % log_every_t == 0 or i == timesteps - 1:
+ intermediates.append(img)
+ if callback:
+ callback(i)
+ if img_callback:
+ img_callback(img, i)
+
+ if return_intermediates:
+ return img, intermediates
+ return img
+
+ @torch.no_grad()
+ def sample(
+ self,
+ cond,
+ batch_size=16,
+ return_intermediates=False,
+ x_T=None,
+ verbose=True,
+ timesteps=None,
+ quantize_denoised=False,
+ mask=None,
+ x0=None,
+ shape=None,
+ **kwargs,
+ ):
+ if shape is None:
+ shape = (batch_size, self.channels, self.image_size, self.image_size)
+ if cond is not None:
+ if isinstance(cond, dict):
+ cond = {
+ key: cond[key][:batch_size]
+ if not isinstance(cond[key], list)
+ else list(map(lambda x: x[:batch_size], cond[key]))
+ for key in cond
+ }
+ else:
+ cond = (
+ [c[:batch_size] for c in cond]
+ if isinstance(cond, list)
+ else cond[:batch_size]
+ )
+ return self.p_sample_loop(
+ cond,
+ shape,
+ return_intermediates=return_intermediates,
+ x_T=x_T,
+ verbose=verbose,
+ timesteps=timesteps,
+ quantize_denoised=quantize_denoised,
+ mask=mask,
+ x0=x0,
+ )
+
+ @torch.no_grad()
+ def sample_log(self, cond, batch_size, ddim, ddim_steps, **kwargs):
+ if ddim:
+ ddim_sampler = DDIMSampler(self)
+ shape = (self.channels, self.image_size, self.image_size)
+ samples, intermediates = ddim_sampler.sample(
+ ddim_steps, batch_size, shape, cond, verbose=False, **kwargs
+ )
+
+ else:
+ samples, intermediates = self.sample(
+ cond=cond, batch_size=batch_size, return_intermediates=True, **kwargs
+ )
+
+ return samples, intermediates
+
+ @torch.no_grad()
+ def get_unconditional_conditioning(self, batch_size, null_label=None):
+ if null_label is not None:
+ xc = null_label
+ if isinstance(xc, ListConfig):
+ xc = list(xc)
+ if isinstance(xc, dict) or isinstance(xc, list):
+ c = self.get_learned_conditioning(xc)
+ else:
+ if hasattr(xc, "to"):
+ xc = xc.to(self.device)
+ c = self.get_learned_conditioning(xc)
+ else:
+ if self.cond_stage_key in ["class_label", "cls"]:
+ xc = self.cond_stage_model.get_unconditional_conditioning(
+ batch_size, device=self.device
+ )
+ return self.get_learned_conditioning(xc)
+ else:
+ raise NotImplementedError("todo")
+ if isinstance(c, list): # in case the encoder gives us a list
+ for i in range(len(c)):
+ c[i] = repeat(c[i], "1 ... -> b ...", b=batch_size).to(self.device)
+ else:
+ c = repeat(c, "1 ... -> b ...", b=batch_size).to(self.device)
+ return c
+
+ @torch.no_grad()
+ def log_images(
+ self,
+ batch,
+ N=8,
+ n_row=4,
+ sample=True,
+ ddim_steps=50,
+ ddim_eta=0.0,
+ return_keys=None,
+ quantize_denoised=True,
+ inpaint=True,
+ plot_denoise_rows=False,
+ plot_progressive_rows=True,
+ plot_diffusion_rows=True,
+ unconditional_guidance_scale=1.0,
+ unconditional_guidance_label=None,
+ use_ema_scope=True,
+ **kwargs,
+ ):
+ ema_scope = self.ema_scope if use_ema_scope else nullcontext
+ use_ddim = ddim_steps is not None
+
+ log = dict()
+ z, c, x, xrec, xc = self.get_input(
+ batch,
+ self.first_stage_key,
+ return_first_stage_outputs=True,
+ force_c_encode=True,
+ return_original_cond=True,
+ bs=N,
+ )
+ N = min(x.shape[0], N)
+ n_row = min(x.shape[0], n_row)
+ log["inputs"] = x
+ log["reconstruction"] = xrec
+ if self.model.conditioning_key is not None:
+ if hasattr(self.cond_stage_model, "decode"):
+ xc = self.cond_stage_model.decode(c)
+ log["conditioning"] = xc
+ elif self.cond_stage_key in ["caption", "txt"]:
+ xc = log_txt_as_img(
+ (x.shape[2], x.shape[3]),
+ batch[self.cond_stage_key],
+ size=x.shape[2] // 25,
+ )
+ log["conditioning"] = xc
+ elif self.cond_stage_key in ["class_label", "cls"]:
+ try:
+ xc = log_txt_as_img(
+ (x.shape[2], x.shape[3]),
+ batch["human_label"],
+ size=x.shape[2] // 25,
+ )
+ log["conditioning"] = xc
+ except KeyError:
+ # probably no "human_label" in batch
+ pass
+ elif isimage(xc):
+ log["conditioning"] = xc
+ if ismap(xc):
+ log["original_conditioning"] = self.to_rgb(xc)
+
+ if plot_diffusion_rows:
+ # get diffusion row
+ diffusion_row = list()
+ z_start = z[:n_row]
+ for t in range(self.num_timesteps):
+ if t % self.log_every_t == 0 or t == self.num_timesteps - 1:
+ t = repeat(torch.tensor([t]), "1 -> b", b=n_row)
+ t = t.to(self.device).long()
+ noise = torch.randn_like(z_start)
+ z_noisy = self.q_sample(x_start=z_start, t=t, noise=noise)
+ diffusion_row.append(self.decode_first_stage(z_noisy))
+
+ diffusion_row = torch.stack(diffusion_row) # n_log_step, n_row, C, H, W
+ diffusion_grid = rearrange(diffusion_row, "n b c h w -> b n c h w")
+ diffusion_grid = rearrange(diffusion_grid, "b n c h w -> (b n) c h w")
+ diffusion_grid = make_grid(diffusion_grid, nrow=diffusion_row.shape[0])
+ log["diffusion_row"] = diffusion_grid
+
+ if sample:
+ # get denoise row
+ with ema_scope("Sampling"):
+ samples, z_denoise_row = self.sample_log(
+ cond=c,
+ batch_size=N,
+ ddim=use_ddim,
+ ddim_steps=ddim_steps,
+ eta=ddim_eta,
+ )
+ # samples, z_denoise_row = self.sample(cond=c, batch_size=N, return_intermediates=True)
+ x_samples = self.decode_first_stage(samples)
+ log["samples"] = x_samples
+ if plot_denoise_rows:
+ denoise_grid = self._get_denoise_row_from_list(z_denoise_row)
+ log["denoise_row"] = denoise_grid
+
+ if (
+ quantize_denoised
+ and not isinstance(self.first_stage_model, AutoencoderKL)
+ and not isinstance(self.first_stage_model, IdentityFirstStage)
+ ):
+ # also display when quantizing x0 while sampling
+ with ema_scope("Plotting Quantized Denoised"):
+ samples, z_denoise_row = self.sample_log(
+ cond=c,
+ batch_size=N,
+ ddim=use_ddim,
+ ddim_steps=ddim_steps,
+ eta=ddim_eta,
+ quantize_denoised=True,
+ )
+ # samples, z_denoise_row = self.sample(cond=c, batch_size=N, return_intermediates=True,
+ # quantize_denoised=True)
+ x_samples = self.decode_first_stage(samples.to(self.device))
+ log["samples_x0_quantized"] = x_samples
+
+ if unconditional_guidance_scale > 1.0:
+ uc = self.get_unconditional_conditioning(N, unconditional_guidance_label)
+ if self.model.conditioning_key == "crossattn-adm":
+ uc = {"c_crossattn": [uc], "c_adm": c["c_adm"]}
+ with ema_scope("Sampling with classifier-free guidance"):
+ samples_cfg, _ = self.sample_log(
+ cond=c,
+ batch_size=N,
+ ddim=use_ddim,
+ ddim_steps=ddim_steps,
+ eta=ddim_eta,
+ unconditional_guidance_scale=unconditional_guidance_scale,
+ unconditional_conditioning=uc,
+ )
+ x_samples_cfg = self.decode_first_stage(samples_cfg)
+ log[
+ f"samples_cfg_scale_{unconditional_guidance_scale:.2f}"
+ ] = x_samples_cfg
+
+ if inpaint:
+ # make a simple center square
+ b, h, w = z.shape[0], z.shape[2], z.shape[3]
+ mask = torch.ones(N, h, w).to(self.device)
+ # zeros will be filled in
+ mask[:, h // 4 : 3 * h // 4, w // 4 : 3 * w // 4] = 0.0
+ mask = mask[:, None, ...]
+ with ema_scope("Plotting Inpaint"):
+ samples, _ = self.sample_log(
+ cond=c,
+ batch_size=N,
+ ddim=use_ddim,
+ eta=ddim_eta,
+ ddim_steps=ddim_steps,
+ x0=z[:N],
+ mask=mask,
+ )
+ x_samples = self.decode_first_stage(samples.to(self.device))
+ log["samples_inpainting"] = x_samples
+ log["mask"] = mask
+
+ # outpaint
+ mask = 1.0 - mask
+ with ema_scope("Plotting Outpaint"):
+ samples, _ = self.sample_log(
+ cond=c,
+ batch_size=N,
+ ddim=use_ddim,
+ eta=ddim_eta,
+ ddim_steps=ddim_steps,
+ x0=z[:N],
+ mask=mask,
+ )
+ x_samples = self.decode_first_stage(samples.to(self.device))
+ log["samples_outpainting"] = x_samples
+
+ if plot_progressive_rows:
+ with ema_scope("Plotting Progressives"):
+ img, progressives = self.progressive_denoising(
+ c,
+ shape=(self.channels, self.image_size, self.image_size),
+ batch_size=N,
+ )
+ prog_row = self._get_denoise_row_from_list(
+ progressives, desc="Progressive Generation"
+ )
+ log["progressive_row"] = prog_row
+
+ if return_keys:
+ if np.intersect1d(list(log.keys()), return_keys).shape[0] == 0:
+ return log
+ else:
+ return {key: log[key] for key in return_keys}
+ return log
+
+ def configure_optimizers(self):
+ lr = self.learning_rate
+ params = list(self.model.parameters())
+ if self.cond_stage_trainable:
+ print(f"{self.__class__.__name__}: Also optimizing conditioner params!")
+ params = params + list(self.cond_stage_model.parameters())
+ if self.learn_logvar:
+ print("Diffusion model optimizing logvar")
+ params.append(self.logvar)
+ opt = torch.optim.AdamW(params, lr=lr)
+ if self.use_scheduler:
+ assert "target" in self.scheduler_config
+ scheduler = instantiate_from_config(self.scheduler_config)
+
+ print("Setting up LambdaLR scheduler...")
+ scheduler = [
+ {
+ "scheduler": LambdaLR(opt, lr_lambda=scheduler.schedule),
+ "interval": "step",
+ "frequency": 1,
+ }
+ ]
+ return [opt], scheduler
+ return opt
+
+ @torch.no_grad()
+ def to_rgb(self, x):
+ x = x.float()
+ if not hasattr(self, "colorize"):
+ self.colorize = torch.randn(3, x.shape[1], 1, 1).to(x)
+ x = nn.functional.conv2d(x, weight=self.colorize)
+ x = 2.0 * (x - x.min()) / (x.max() - x.min()) - 1.0
+ return x
+
+
+class DiffusionWrapper(torch.nn.Module):
+ def __init__(self, diff_model_config, conditioning_key):
+ super().__init__()
+ self.sequential_cross_attn = diff_model_config.pop(
+ "sequential_crossattn", False
+ )
+ self.diffusion_model = instantiate_from_config(diff_model_config)
+ self.conditioning_key = conditioning_key
+ assert self.conditioning_key in [
+ None,
+ "concat",
+ "crossattn",
+ "hybrid",
+ "adm",
+ "hybrid-adm",
+ "crossattn-adm",
+ ]
+
+ def forward(
+ self, x, t, c_concat: list = None, c_crossattn: list = None, c_adm=None
+ ):
+ if self.conditioning_key is None:
+ out = self.diffusion_model(x, t)
+ elif self.conditioning_key == "concat":
+ xc = torch.cat([x] + c_concat, dim=1)
+ out = self.diffusion_model(xc, t)
+ elif self.conditioning_key == "crossattn":
+ if not self.sequential_cross_attn:
+ cc = torch.cat(c_crossattn, 1)
+ else:
+ cc = c_crossattn
+ out = self.diffusion_model(x, t, context=cc)
+ elif self.conditioning_key == "hybrid":
+ xc = torch.cat([x] + c_concat, dim=1)
+ cc = torch.cat(c_crossattn, 1)
+ out = self.diffusion_model(xc, t, context=cc)
+ elif self.conditioning_key == "hybrid-adm":
+ assert c_adm is not None
+ xc = torch.cat([x] + c_concat, dim=1)
+ cc = torch.cat(c_crossattn, 1)
+ out = self.diffusion_model(xc, t, context=cc, y=c_adm)
+ elif self.conditioning_key == "crossattn-adm":
+ assert c_adm is not None
+ cc = torch.cat(c_crossattn, 1)
+ out = self.diffusion_model(x, t, context=cc, y=c_adm)
+ elif self.conditioning_key == "adm":
+ cc = c_crossattn[0]
+ out = self.diffusion_model(x, t, y=cc)
+ else:
+ raise NotImplementedError()
+
+ return out
+
+
+class LatentUpscaleDiffusion(LatentDiffusion):
+ def __init__(
+ self,
+ *args,
+ low_scale_config,
+ low_scale_key="LR",
+ noise_level_key=None,
+ **kwargs,
+ ):
+ super().__init__(*args, **kwargs)
+ # assumes that neither the cond_stage nor the low_scale_model contain trainable params
+ assert not self.cond_stage_trainable
+ self.instantiate_low_stage(low_scale_config)
+ self.low_scale_key = low_scale_key
+ self.noise_level_key = noise_level_key
+
+ def instantiate_low_stage(self, config):
+ model = instantiate_from_config(config)
+ self.low_scale_model = model.eval()
+ self.low_scale_model.train = disabled_train
+ for param in self.low_scale_model.parameters():
+ param.requires_grad = False
+
+ @torch.no_grad()
+ def get_input(self, batch, k, cond_key=None, bs=None, log_mode=False):
+ if not log_mode:
+ z, c = super().get_input(batch, k, force_c_encode=True, bs=bs)
+ else:
+ z, c, x, xrec, xc = super().get_input(
+ batch,
+ self.first_stage_key,
+ return_first_stage_outputs=True,
+ force_c_encode=True,
+ return_original_cond=True,
+ bs=bs,
+ )
+ x_low = batch[self.low_scale_key][:bs]
+ x_low = rearrange(x_low, "b h w c -> b c h w")
+ x_low = x_low.to(memory_format=torch.contiguous_format).float()
+ zx, noise_level = self.low_scale_model(x_low)
+ if self.noise_level_key is not None:
+ # get noise level from batch instead, e.g. when extracting a custom noise level for bsr
+ raise NotImplementedError("TODO")
+
+ all_conds = {"c_concat": [zx], "c_crossattn": [c], "c_adm": noise_level}
+ if log_mode:
+ # TODO: maybe disable if too expensive
+ x_low_rec = self.low_scale_model.decode(zx)
+ return z, all_conds, x, xrec, xc, x_low, x_low_rec, noise_level
+ return z, all_conds
+
+ @torch.no_grad()
+ def log_images(
+ self,
+ batch,
+ N=8,
+ n_row=4,
+ sample=True,
+ ddim_steps=200,
+ ddim_eta=1.0,
+ return_keys=None,
+ plot_denoise_rows=False,
+ plot_progressive_rows=True,
+ plot_diffusion_rows=True,
+ unconditional_guidance_scale=1.0,
+ unconditional_guidance_label=None,
+ use_ema_scope=True,
+ **kwargs,
+ ):
+ ema_scope = self.ema_scope if use_ema_scope else nullcontext
+ use_ddim = ddim_steps is not None
+
+ log = dict()
+ z, c, x, xrec, xc, x_low, x_low_rec, noise_level = self.get_input(
+ batch, self.first_stage_key, bs=N, log_mode=True
+ )
+ N = min(x.shape[0], N)
+ n_row = min(x.shape[0], n_row)
+ log["inputs"] = x
+ log["reconstruction"] = xrec
+ log["x_lr"] = x_low
+ log[
+ f"x_lr_rec_@noise_levels{'-'.join(map(lambda x: str(x), list(noise_level.cpu().numpy())))}"
+ ] = x_low_rec
+ if self.model.conditioning_key is not None:
+ if hasattr(self.cond_stage_model, "decode"):
+ xc = self.cond_stage_model.decode(c)
+ log["conditioning"] = xc
+ elif self.cond_stage_key in ["caption", "txt"]:
+ xc = log_txt_as_img(
+ (x.shape[2], x.shape[3]),
+ batch[self.cond_stage_key],
+ size=x.shape[2] // 25,
+ )
+ log["conditioning"] = xc
+ elif self.cond_stage_key in ["class_label", "cls"]:
+ xc = log_txt_as_img(
+ (x.shape[2], x.shape[3]),
+ batch["human_label"],
+ size=x.shape[2] // 25,
+ )
+ log["conditioning"] = xc
+ elif isimage(xc):
+ log["conditioning"] = xc
+ if ismap(xc):
+ log["original_conditioning"] = self.to_rgb(xc)
+
+ if plot_diffusion_rows:
+ # get diffusion row
+ diffusion_row = list()
+ z_start = z[:n_row]
+ for t in range(self.num_timesteps):
+ if t % self.log_every_t == 0 or t == self.num_timesteps - 1:
+ t = repeat(torch.tensor([t]), "1 -> b", b=n_row)
+ t = t.to(self.device).long()
+ noise = torch.randn_like(z_start)
+ z_noisy = self.q_sample(x_start=z_start, t=t, noise=noise)
+ diffusion_row.append(self.decode_first_stage(z_noisy))
+
+ diffusion_row = torch.stack(diffusion_row) # n_log_step, n_row, C, H, W
+ diffusion_grid = rearrange(diffusion_row, "n b c h w -> b n c h w")
+ diffusion_grid = rearrange(diffusion_grid, "b n c h w -> (b n) c h w")
+ diffusion_grid = make_grid(diffusion_grid, nrow=diffusion_row.shape[0])
+ log["diffusion_row"] = diffusion_grid
+
+ if sample:
+ # get denoise row
+ with ema_scope("Sampling"):
+ samples, z_denoise_row = self.sample_log(
+ cond=c,
+ batch_size=N,
+ ddim=use_ddim,
+ ddim_steps=ddim_steps,
+ eta=ddim_eta,
+ )
+ # samples, z_denoise_row = self.sample(cond=c, batch_size=N, return_intermediates=True)
+ x_samples = self.decode_first_stage(samples)
+ log["samples"] = x_samples
+ if plot_denoise_rows:
+ denoise_grid = self._get_denoise_row_from_list(z_denoise_row)
+ log["denoise_row"] = denoise_grid
+
+ if unconditional_guidance_scale > 1.0:
+ uc_tmp = self.get_unconditional_conditioning(
+ N, unconditional_guidance_label
+ )
+ # TODO explore better "unconditional" choices for the other keys
+ # maybe guide away from empty text label and highest noise level and maximally degraded zx?
+ uc = dict()
+ for k in c:
+ if k == "c_crossattn":
+ assert isinstance(c[k], list) and len(c[k]) == 1
+ uc[k] = [uc_tmp]
+ elif k == "c_adm": # todo: only run with text-based guidance?
+ assert isinstance(c[k], torch.Tensor)
+ # uc[k] = torch.ones_like(c[k]) * self.low_scale_model.max_noise_level
+ uc[k] = c[k]
+ elif isinstance(c[k], list):
+ uc[k] = [c[k][i] for i in range(len(c[k]))]
+ else:
+ uc[k] = c[k]
+
+ with ema_scope("Sampling with classifier-free guidance"):
+ samples_cfg, _ = self.sample_log(
+ cond=c,
+ batch_size=N,
+ ddim=use_ddim,
+ ddim_steps=ddim_steps,
+ eta=ddim_eta,
+ unconditional_guidance_scale=unconditional_guidance_scale,
+ unconditional_conditioning=uc,
+ )
+ x_samples_cfg = self.decode_first_stage(samples_cfg)
+ log[
+ f"samples_cfg_scale_{unconditional_guidance_scale:.2f}"
+ ] = x_samples_cfg
+
+ if plot_progressive_rows:
+ with ema_scope("Plotting Progressives"):
+ img, progressives = self.progressive_denoising(
+ c,
+ shape=(self.channels, self.image_size, self.image_size),
+ batch_size=N,
+ )
+ prog_row = self._get_denoise_row_from_list(
+ progressives, desc="Progressive Generation"
+ )
+ log["progressive_row"] = prog_row
+
+ return log
+
+
+class LatentFinetuneDiffusion(LatentDiffusion):
+ """
+ Basis for different finetunas, such as inpainting or depth2image
+ To disable finetuning mode, set finetune_keys to None
+ """
+
+ def __init__(
+ self,
+ concat_keys: tuple,
+ finetune_keys=(
+ "model.diffusion_model.input_blocks.0.0.weight",
+ "model_ema.diffusion_modelinput_blocks00weight",
+ ),
+ keep_finetune_dims=4,
+ # if model was trained without concat mode before and we would like to keep these channels
+ c_concat_log_start=None, # to log reconstruction of c_concat codes
+ c_concat_log_end=None,
+ *args,
+ **kwargs,
+ ):
+ ckpt_path = kwargs.pop("ckpt_path", None)
+ ignore_keys = kwargs.pop("ignore_keys", list())
+ super().__init__(*args, **kwargs)
+ self.finetune_keys = finetune_keys
+ self.concat_keys = concat_keys
+ self.keep_dims = keep_finetune_dims
+ self.c_concat_log_start = c_concat_log_start
+ self.c_concat_log_end = c_concat_log_end
+ if exists(self.finetune_keys):
+ assert exists(ckpt_path), "can only finetune from a given checkpoint"
+ if exists(ckpt_path):
+ self.init_from_ckpt(ckpt_path, ignore_keys)
+
+ def init_from_ckpt(self, path, ignore_keys=list(), only_model=False):
+ sd = torch.load(path, map_location="cpu")
+ if "state_dict" in list(sd.keys()):
+ sd = sd["state_dict"]
+ keys = list(sd.keys())
+ for k in keys:
+ for ik in ignore_keys:
+ if k.startswith(ik):
+ print("Deleting key {} from state_dict.".format(k))
+ del sd[k]
+
+ # make it explicit, finetune by including extra input channels
+ if exists(self.finetune_keys) and k in self.finetune_keys:
+ new_entry = None
+ for name, param in self.named_parameters():
+ if name in self.finetune_keys:
+ print(
+ f"modifying key '{name}' and keeping its original {self.keep_dims} (channels) dimensions only"
+ )
+ new_entry = torch.zeros_like(param) # zero init
+ assert exists(new_entry), "did not find matching parameter to modify"
+ new_entry[:, : self.keep_dims, ...] = sd[k]
+ sd[k] = new_entry
+
+ missing, unexpected = (
+ self.load_state_dict(sd, strict=False)
+ if not only_model
+ else self.model.load_state_dict(sd, strict=False)
+ )
+ print(
+ f"Restored from {path} with {len(missing)} missing and {len(unexpected)} unexpected keys"
+ )
+ if len(missing) > 0:
+ print(f"Missing Keys: {missing}")
+ if len(unexpected) > 0:
+ print(f"Unexpected Keys: {unexpected}")
+
+ @torch.no_grad()
+ def log_images(
+ self,
+ batch,
+ N=8,
+ n_row=4,
+ sample=True,
+ ddim_steps=200,
+ ddim_eta=1.0,
+ return_keys=None,
+ quantize_denoised=True,
+ inpaint=True,
+ plot_denoise_rows=False,
+ plot_progressive_rows=True,
+ plot_diffusion_rows=True,
+ unconditional_guidance_scale=1.0,
+ unconditional_guidance_label=None,
+ use_ema_scope=True,
+ **kwargs,
+ ):
+ ema_scope = self.ema_scope if use_ema_scope else nullcontext
+ use_ddim = ddim_steps is not None
+
+ log = dict()
+ z, c, x, xrec, xc = self.get_input(
+ batch, self.first_stage_key, bs=N, return_first_stage_outputs=True
+ )
+ c_cat, c = c["c_concat"][0], c["c_crossattn"][0]
+ N = min(x.shape[0], N)
+ n_row = min(x.shape[0], n_row)
+ log["inputs"] = x
+ log["reconstruction"] = xrec
+ if self.model.conditioning_key is not None:
+ if hasattr(self.cond_stage_model, "decode"):
+ xc = self.cond_stage_model.decode(c)
+ log["conditioning"] = xc
+ elif self.cond_stage_key in ["caption", "txt"]:
+ xc = log_txt_as_img(
+ (x.shape[2], x.shape[3]),
+ batch[self.cond_stage_key],
+ size=x.shape[2] // 25,
+ )
+ log["conditioning"] = xc
+ elif self.cond_stage_key in ["class_label", "cls"]:
+ xc = log_txt_as_img(
+ (x.shape[2], x.shape[3]),
+ batch["human_label"],
+ size=x.shape[2] // 25,
+ )
+ log["conditioning"] = xc
+ elif isimage(xc):
+ log["conditioning"] = xc
+ if ismap(xc):
+ log["original_conditioning"] = self.to_rgb(xc)
+
+ if not (self.c_concat_log_start is None and self.c_concat_log_end is None):
+ log["c_concat_decoded"] = self.decode_first_stage(
+ c_cat[:, self.c_concat_log_start : self.c_concat_log_end]
+ )
+
+ if plot_diffusion_rows:
+ # get diffusion row
+ diffusion_row = list()
+ z_start = z[:n_row]
+ for t in range(self.num_timesteps):
+ if t % self.log_every_t == 0 or t == self.num_timesteps - 1:
+ t = repeat(torch.tensor([t]), "1 -> b", b=n_row)
+ t = t.to(self.device).long()
+ noise = torch.randn_like(z_start)
+ z_noisy = self.q_sample(x_start=z_start, t=t, noise=noise)
+ diffusion_row.append(self.decode_first_stage(z_noisy))
+
+ diffusion_row = torch.stack(diffusion_row) # n_log_step, n_row, C, H, W
+ diffusion_grid = rearrange(diffusion_row, "n b c h w -> b n c h w")
+ diffusion_grid = rearrange(diffusion_grid, "b n c h w -> (b n) c h w")
+ diffusion_grid = make_grid(diffusion_grid, nrow=diffusion_row.shape[0])
+ log["diffusion_row"] = diffusion_grid
+
+ if sample:
+ # get denoise row
+ with ema_scope("Sampling"):
+ samples, z_denoise_row = self.sample_log(
+ cond={"c_concat": [c_cat], "c_crossattn": [c]},
+ batch_size=N,
+ ddim=use_ddim,
+ ddim_steps=ddim_steps,
+ eta=ddim_eta,
+ )
+ # samples, z_denoise_row = self.sample(cond=c, batch_size=N, return_intermediates=True)
+ x_samples = self.decode_first_stage(samples)
+ log["samples"] = x_samples
+ if plot_denoise_rows:
+ denoise_grid = self._get_denoise_row_from_list(z_denoise_row)
+ log["denoise_row"] = denoise_grid
+
+ if unconditional_guidance_scale > 1.0:
+ uc_cross = self.get_unconditional_conditioning(
+ N, unconditional_guidance_label
+ )
+ uc_cat = c_cat
+ uc_full = {"c_concat": [uc_cat], "c_crossattn": [uc_cross]}
+ with ema_scope("Sampling with classifier-free guidance"):
+ samples_cfg, _ = self.sample_log(
+ cond={"c_concat": [c_cat], "c_crossattn": [c]},
+ batch_size=N,
+ ddim=use_ddim,
+ ddim_steps=ddim_steps,
+ eta=ddim_eta,
+ unconditional_guidance_scale=unconditional_guidance_scale,
+ unconditional_conditioning=uc_full,
+ )
+ x_samples_cfg = self.decode_first_stage(samples_cfg)
+ log[
+ f"samples_cfg_scale_{unconditional_guidance_scale:.2f}"
+ ] = x_samples_cfg
+
+ return log
+
+
+class LatentInpaintDiffusion(LatentFinetuneDiffusion):
+ """
+ can either run as pure inpainting model (only concat mode) or with mixed conditionings,
+ e.g. mask as concat and text via cross-attn.
+ To disable finetuning mode, set finetune_keys to None
+ """
+
+ def __init__(
+ self,
+ concat_keys=("mask", "masked_image"),
+ masked_image_key="masked_image",
+ *args,
+ **kwargs,
+ ):
+ super().__init__(concat_keys, *args, **kwargs)
+ self.masked_image_key = masked_image_key
+ assert self.masked_image_key in concat_keys
+
+ @torch.no_grad()
+ def get_input(
+ self, batch, k, cond_key=None, bs=None, return_first_stage_outputs=False
+ ):
+ # note: restricted to non-trainable encoders currently
+ assert (
+ not self.cond_stage_trainable
+ ), "trainable cond stages not yet supported for inpainting"
+ z, c, x, xrec, xc = super().get_input(
+ batch,
+ self.first_stage_key,
+ return_first_stage_outputs=True,
+ force_c_encode=True,
+ return_original_cond=True,
+ bs=bs,
+ )
+
+ assert exists(self.concat_keys)
+ c_cat = list()
+ for ck in self.concat_keys:
+ cc = (
+ rearrange(batch[ck], "b h w c -> b c h w")
+ .to(memory_format=torch.contiguous_format)
+ .float()
+ )
+ if bs is not None:
+ cc = cc[:bs]
+ cc = cc.to(self.device)
+ bchw = z.shape
+ if ck != self.masked_image_key:
+ cc = torch.nn.functional.interpolate(cc, size=bchw[-2:])
+ else:
+ cc = self.get_first_stage_encoding(self.encode_first_stage(cc))
+ c_cat.append(cc)
+ c_cat = torch.cat(c_cat, dim=1)
+ all_conds = {"c_concat": [c_cat], "c_crossattn": [c]}
+ if return_first_stage_outputs:
+ return z, all_conds, x, xrec, xc
+ return z, all_conds
+
+ @torch.no_grad()
+ def log_images(self, *args, **kwargs):
+ log = super(LatentInpaintDiffusion, self).log_images(*args, **kwargs)
+ log["masked_image"] = (
+ rearrange(args[0]["masked_image"], "b h w c -> b c h w")
+ .to(memory_format=torch.contiguous_format)
+ .float()
+ )
+ return log
+
+
+class LatentDepth2ImageDiffusion(LatentFinetuneDiffusion):
+ """
+ condition on monocular depth estimation
+ """
+
+ def __init__(self, depth_stage_config, concat_keys=("midas_in",), *args, **kwargs):
+ super().__init__(concat_keys=concat_keys, *args, **kwargs)
+ self.depth_model = instantiate_from_config(depth_stage_config)
+ self.depth_stage_key = concat_keys[0]
+
+ @torch.no_grad()
+ def get_input(
+ self, batch, k, cond_key=None, bs=None, return_first_stage_outputs=False
+ ):
+ # note: restricted to non-trainable encoders currently
+ assert (
+ not self.cond_stage_trainable
+ ), "trainable cond stages not yet supported for depth2img"
+ z, c, x, xrec, xc = super().get_input(
+ batch,
+ self.first_stage_key,
+ return_first_stage_outputs=True,
+ force_c_encode=True,
+ return_original_cond=True,
+ bs=bs,
+ )
+
+ assert exists(self.concat_keys)
+ assert len(self.concat_keys) == 1
+ c_cat = list()
+ for ck in self.concat_keys:
+ cc = batch[ck]
+ if bs is not None:
+ cc = cc[:bs]
+ cc = cc.to(self.device)
+ cc = self.depth_model(cc)
+ cc = torch.nn.functional.interpolate(
+ cc,
+ size=z.shape[2:],
+ mode="bicubic",
+ align_corners=False,
+ )
+
+ depth_min, depth_max = torch.amin(
+ cc, dim=[1, 2, 3], keepdim=True
+ ), torch.amax(cc, dim=[1, 2, 3], keepdim=True)
+ cc = 2.0 * (cc - depth_min) / (depth_max - depth_min + 0.001) - 1.0
+ c_cat.append(cc)
+ c_cat = torch.cat(c_cat, dim=1)
+ all_conds = {"c_concat": [c_cat], "c_crossattn": [c]}
+ if return_first_stage_outputs:
+ return z, all_conds, x, xrec, xc
+ return z, all_conds
+
+ @torch.no_grad()
+ def log_images(self, *args, **kwargs):
+ log = super().log_images(*args, **kwargs)
+ depth = self.depth_model(args[0][self.depth_stage_key])
+ depth_min, depth_max = torch.amin(
+ depth, dim=[1, 2, 3], keepdim=True
+ ), torch.amax(depth, dim=[1, 2, 3], keepdim=True)
+ log["depth"] = 2.0 * (depth - depth_min) / (depth_max - depth_min) - 1.0
+ return log
+
+
+class LatentUpscaleFinetuneDiffusion(LatentFinetuneDiffusion):
+ """
+ condition on low-res image (and optionally on some spatial noise augmentation)
+ """
+
+ def __init__(
+ self,
+ concat_keys=("lr",),
+ reshuffle_patch_size=None,
+ low_scale_config=None,
+ low_scale_key=None,
+ *args,
+ **kwargs,
+ ):
+ super().__init__(concat_keys=concat_keys, *args, **kwargs)
+ self.reshuffle_patch_size = reshuffle_patch_size
+ self.low_scale_model = None
+ if low_scale_config is not None:
+ print("Initializing a low-scale model")
+ assert exists(low_scale_key)
+ self.instantiate_low_stage(low_scale_config)
+ self.low_scale_key = low_scale_key
+
+ def instantiate_low_stage(self, config):
+ model = instantiate_from_config(config)
+ self.low_scale_model = model.eval()
+ self.low_scale_model.train = disabled_train
+ for param in self.low_scale_model.parameters():
+ param.requires_grad = False
+
+ @torch.no_grad()
+ def get_input(
+ self, batch, k, cond_key=None, bs=None, return_first_stage_outputs=False
+ ):
+ # note: restricted to non-trainable encoders currently
+ assert (
+ not self.cond_stage_trainable
+ ), "trainable cond stages not yet supported for upscaling-ft"
+ z, c, x, xrec, xc = super().get_input(
+ batch,
+ self.first_stage_key,
+ return_first_stage_outputs=True,
+ force_c_encode=True,
+ return_original_cond=True,
+ bs=bs,
+ )
+
+ assert exists(self.concat_keys)
+ assert len(self.concat_keys) == 1
+ # optionally make spatial noise_level here
+ c_cat = list()
+ noise_level = None
+ for ck in self.concat_keys:
+ cc = batch[ck]
+ cc = rearrange(cc, "b h w c -> b c h w")
+ if exists(self.reshuffle_patch_size):
+ assert isinstance(self.reshuffle_patch_size, int)
+ cc = rearrange(
+ cc,
+ "b c (p1 h) (p2 w) -> b (p1 p2 c) h w",
+ p1=self.reshuffle_patch_size,
+ p2=self.reshuffle_patch_size,
+ )
+ if bs is not None:
+ cc = cc[:bs]
+ cc = cc.to(self.device)
+ if exists(self.low_scale_model) and ck == self.low_scale_key:
+ cc, noise_level = self.low_scale_model(cc)
+ c_cat.append(cc)
+ c_cat = torch.cat(c_cat, dim=1)
+ if exists(noise_level):
+ all_conds = {"c_concat": [c_cat], "c_crossattn": [c], "c_adm": noise_level}
+ else:
+ all_conds = {"c_concat": [c_cat], "c_crossattn": [c]}
+ if return_first_stage_outputs:
+ return z, all_conds, x, xrec, xc
+ return z, all_conds
+
+ @torch.no_grad()
+ def log_images(self, *args, **kwargs):
+ log = super().log_images(*args, **kwargs)
+ log["lr"] = rearrange(args[0]["lr"], "b h w c -> b c h w")
+ return log
diff --git a/iopaint/model/anytext/ldm/models/diffusion/dpm_solver/dpm_solver.py b/iopaint/model/anytext/ldm/models/diffusion/dpm_solver/dpm_solver.py
new file mode 100644
index 0000000000000000000000000000000000000000..095e5ba3ce0b1aa7f4b3f1e2e5d8fff7cfe6dc8c
--- /dev/null
+++ b/iopaint/model/anytext/ldm/models/diffusion/dpm_solver/dpm_solver.py
@@ -0,0 +1,1154 @@
+import torch
+import torch.nn.functional as F
+import math
+from tqdm import tqdm
+
+
+class NoiseScheduleVP:
+ def __init__(
+ self,
+ schedule='discrete',
+ betas=None,
+ alphas_cumprod=None,
+ continuous_beta_0=0.1,
+ continuous_beta_1=20.,
+ ):
+ """Create a wrapper class for the forward SDE (VP type).
+ ***
+ Update: We support discrete-time diffusion models by implementing a picewise linear interpolation for log_alpha_t.
+ We recommend to use schedule='discrete' for the discrete-time diffusion models, especially for high-resolution images.
+ ***
+ The forward SDE ensures that the condition distribution q_{t|0}(x_t | x_0) = N ( alpha_t * x_0, sigma_t^2 * I ).
+ We further define lambda_t = log(alpha_t) - log(sigma_t), which is the half-logSNR (described in the DPM-Solver paper).
+ Therefore, we implement the functions for computing alpha_t, sigma_t and lambda_t. For t in [0, T], we have:
+ log_alpha_t = self.marginal_log_mean_coeff(t)
+ sigma_t = self.marginal_std(t)
+ lambda_t = self.marginal_lambda(t)
+ Moreover, as lambda(t) is an invertible function, we also support its inverse function:
+ t = self.inverse_lambda(lambda_t)
+ ===============================================================
+ We support both discrete-time DPMs (trained on n = 0, 1, ..., N-1) and continuous-time DPMs (trained on t in [t_0, T]).
+ 1. For discrete-time DPMs:
+ For discrete-time DPMs trained on n = 0, 1, ..., N-1, we convert the discrete steps to continuous time steps by:
+ t_i = (i + 1) / N
+ e.g. for N = 1000, we have t_0 = 1e-3 and T = t_{N-1} = 1.
+ We solve the corresponding diffusion ODE from time T = 1 to time t_0 = 1e-3.
+ Args:
+ betas: A `torch.Tensor`. The beta array for the discrete-time DPM. (See the original DDPM paper for details)
+ alphas_cumprod: A `torch.Tensor`. The cumprod alphas for the discrete-time DPM. (See the original DDPM paper for details)
+ Note that we always have alphas_cumprod = cumprod(betas). Therefore, we only need to set one of `betas` and `alphas_cumprod`.
+ **Important**: Please pay special attention for the args for `alphas_cumprod`:
+ The `alphas_cumprod` is the \hat{alpha_n} arrays in the notations of DDPM. Specifically, DDPMs assume that
+ q_{t_n | 0}(x_{t_n} | x_0) = N ( \sqrt{\hat{alpha_n}} * x_0, (1 - \hat{alpha_n}) * I ).
+ Therefore, the notation \hat{alpha_n} is different from the notation alpha_t in DPM-Solver. In fact, we have
+ alpha_{t_n} = \sqrt{\hat{alpha_n}},
+ and
+ log(alpha_{t_n}) = 0.5 * log(\hat{alpha_n}).
+ 2. For continuous-time DPMs:
+ We support two types of VPSDEs: linear (DDPM) and cosine (improved-DDPM). The hyperparameters for the noise
+ schedule are the default settings in DDPM and improved-DDPM:
+ Args:
+ beta_min: A `float` number. The smallest beta for the linear schedule.
+ beta_max: A `float` number. The largest beta for the linear schedule.
+ cosine_s: A `float` number. The hyperparameter in the cosine schedule.
+ cosine_beta_max: A `float` number. The hyperparameter in the cosine schedule.
+ T: A `float` number. The ending time of the forward process.
+ ===============================================================
+ Args:
+ schedule: A `str`. The noise schedule of the forward SDE. 'discrete' for discrete-time DPMs,
+ 'linear' or 'cosine' for continuous-time DPMs.
+ Returns:
+ A wrapper object of the forward SDE (VP type).
+
+ ===============================================================
+ Example:
+ # For discrete-time DPMs, given betas (the beta array for n = 0, 1, ..., N - 1):
+ >>> ns = NoiseScheduleVP('discrete', betas=betas)
+ # For discrete-time DPMs, given alphas_cumprod (the \hat{alpha_n} array for n = 0, 1, ..., N - 1):
+ >>> ns = NoiseScheduleVP('discrete', alphas_cumprod=alphas_cumprod)
+ # For continuous-time DPMs (VPSDE), linear schedule:
+ >>> ns = NoiseScheduleVP('linear', continuous_beta_0=0.1, continuous_beta_1=20.)
+ """
+
+ if schedule not in ['discrete', 'linear', 'cosine']:
+ raise ValueError(
+ "Unsupported noise schedule {}. The schedule needs to be 'discrete' or 'linear' or 'cosine'".format(
+ schedule))
+
+ self.schedule = schedule
+ if schedule == 'discrete':
+ if betas is not None:
+ log_alphas = 0.5 * torch.log(1 - betas).cumsum(dim=0)
+ else:
+ assert alphas_cumprod is not None
+ log_alphas = 0.5 * torch.log(alphas_cumprod)
+ self.total_N = len(log_alphas)
+ self.T = 1.
+ self.t_array = torch.linspace(0., 1., self.total_N + 1)[1:].reshape((1, -1))
+ self.log_alpha_array = log_alphas.reshape((1, -1,))
+ else:
+ self.total_N = 1000
+ self.beta_0 = continuous_beta_0
+ self.beta_1 = continuous_beta_1
+ self.cosine_s = 0.008
+ self.cosine_beta_max = 999.
+ self.cosine_t_max = math.atan(self.cosine_beta_max * (1. + self.cosine_s) / math.pi) * 2. * (
+ 1. + self.cosine_s) / math.pi - self.cosine_s
+ self.cosine_log_alpha_0 = math.log(math.cos(self.cosine_s / (1. + self.cosine_s) * math.pi / 2.))
+ self.schedule = schedule
+ if schedule == 'cosine':
+ # For the cosine schedule, T = 1 will have numerical issues. So we manually set the ending time T.
+ # Note that T = 0.9946 may be not the optimal setting. However, we find it works well.
+ self.T = 0.9946
+ else:
+ self.T = 1.
+
+ def marginal_log_mean_coeff(self, t):
+ """
+ Compute log(alpha_t) of a given continuous-time label t in [0, T].
+ """
+ if self.schedule == 'discrete':
+ return interpolate_fn(t.reshape((-1, 1)), self.t_array.to(t.device),
+ self.log_alpha_array.to(t.device)).reshape((-1))
+ elif self.schedule == 'linear':
+ return -0.25 * t ** 2 * (self.beta_1 - self.beta_0) - 0.5 * t * self.beta_0
+ elif self.schedule == 'cosine':
+ log_alpha_fn = lambda s: torch.log(torch.cos((s + self.cosine_s) / (1. + self.cosine_s) * math.pi / 2.))
+ log_alpha_t = log_alpha_fn(t) - self.cosine_log_alpha_0
+ return log_alpha_t
+
+ def marginal_alpha(self, t):
+ """
+ Compute alpha_t of a given continuous-time label t in [0, T].
+ """
+ return torch.exp(self.marginal_log_mean_coeff(t))
+
+ def marginal_std(self, t):
+ """
+ Compute sigma_t of a given continuous-time label t in [0, T].
+ """
+ return torch.sqrt(1. - torch.exp(2. * self.marginal_log_mean_coeff(t)))
+
+ def marginal_lambda(self, t):
+ """
+ Compute lambda_t = log(alpha_t) - log(sigma_t) of a given continuous-time label t in [0, T].
+ """
+ log_mean_coeff = self.marginal_log_mean_coeff(t)
+ log_std = 0.5 * torch.log(1. - torch.exp(2. * log_mean_coeff))
+ return log_mean_coeff - log_std
+
+ def inverse_lambda(self, lamb):
+ """
+ Compute the continuous-time label t in [0, T] of a given half-logSNR lambda_t.
+ """
+ if self.schedule == 'linear':
+ tmp = 2. * (self.beta_1 - self.beta_0) * torch.logaddexp(-2. * lamb, torch.zeros((1,)).to(lamb))
+ Delta = self.beta_0 ** 2 + tmp
+ return tmp / (torch.sqrt(Delta) + self.beta_0) / (self.beta_1 - self.beta_0)
+ elif self.schedule == 'discrete':
+ log_alpha = -0.5 * torch.logaddexp(torch.zeros((1,)).to(lamb.device), -2. * lamb)
+ t = interpolate_fn(log_alpha.reshape((-1, 1)), torch.flip(self.log_alpha_array.to(lamb.device), [1]),
+ torch.flip(self.t_array.to(lamb.device), [1]))
+ return t.reshape((-1,))
+ else:
+ log_alpha = -0.5 * torch.logaddexp(-2. * lamb, torch.zeros((1,)).to(lamb))
+ t_fn = lambda log_alpha_t: torch.arccos(torch.exp(log_alpha_t + self.cosine_log_alpha_0)) * 2. * (
+ 1. + self.cosine_s) / math.pi - self.cosine_s
+ t = t_fn(log_alpha)
+ return t
+
+
+def model_wrapper(
+ model,
+ noise_schedule,
+ model_type="noise",
+ model_kwargs={},
+ guidance_type="uncond",
+ condition=None,
+ unconditional_condition=None,
+ guidance_scale=1.,
+ classifier_fn=None,
+ classifier_kwargs={},
+):
+ """Create a wrapper function for the noise prediction model.
+ DPM-Solver needs to solve the continuous-time diffusion ODEs. For DPMs trained on discrete-time labels, we need to
+ firstly wrap the model function to a noise prediction model that accepts the continuous time as the input.
+ We support four types of the diffusion model by setting `model_type`:
+ 1. "noise": noise prediction model. (Trained by predicting noise).
+ 2. "x_start": data prediction model. (Trained by predicting the data x_0 at time 0).
+ 3. "v": velocity prediction model. (Trained by predicting the velocity).
+ The "v" prediction is derivation detailed in Appendix D of [1], and is used in Imagen-Video [2].
+ [1] Salimans, Tim, and Jonathan Ho. "Progressive distillation for fast sampling of diffusion models."
+ arXiv preprint arXiv:2202.00512 (2022).
+ [2] Ho, Jonathan, et al. "Imagen Video: High Definition Video Generation with Diffusion Models."
+ arXiv preprint arXiv:2210.02303 (2022).
+
+ 4. "score": marginal score function. (Trained by denoising score matching).
+ Note that the score function and the noise prediction model follows a simple relationship:
+ ```
+ noise(x_t, t) = -sigma_t * score(x_t, t)
+ ```
+ We support three types of guided sampling by DPMs by setting `guidance_type`:
+ 1. "uncond": unconditional sampling by DPMs.
+ The input `model` has the following format:
+ ``
+ model(x, t_input, **model_kwargs) -> noise | x_start | v | score
+ ``
+ 2. "classifier": classifier guidance sampling [3] by DPMs and another classifier.
+ The input `model` has the following format:
+ ``
+ model(x, t_input, **model_kwargs) -> noise | x_start | v | score
+ ``
+ The input `classifier_fn` has the following format:
+ ``
+ classifier_fn(x, t_input, cond, **classifier_kwargs) -> logits(x, t_input, cond)
+ ``
+ [3] P. Dhariwal and A. Q. Nichol, "Diffusion models beat GANs on image synthesis,"
+ in Advances in Neural Information Processing Systems, vol. 34, 2021, pp. 8780-8794.
+ 3. "classifier-free": classifier-free guidance sampling by conditional DPMs.
+ The input `model` has the following format:
+ ``
+ model(x, t_input, cond, **model_kwargs) -> noise | x_start | v | score
+ ``
+ And if cond == `unconditional_condition`, the model output is the unconditional DPM output.
+ [4] Ho, Jonathan, and Tim Salimans. "Classifier-free diffusion guidance."
+ arXiv preprint arXiv:2207.12598 (2022).
+
+ The `t_input` is the time label of the model, which may be discrete-time labels (i.e. 0 to 999)
+ or continuous-time labels (i.e. epsilon to T).
+ We wrap the model function to accept only `x` and `t_continuous` as inputs, and outputs the predicted noise:
+ ``
+ def model_fn(x, t_continuous) -> noise:
+ t_input = get_model_input_time(t_continuous)
+ return noise_pred(model, x, t_input, **model_kwargs)
+ ``
+ where `t_continuous` is the continuous time labels (i.e. epsilon to T). And we use `model_fn` for DPM-Solver.
+ ===============================================================
+ Args:
+ model: A diffusion model with the corresponding format described above.
+ noise_schedule: A noise schedule object, such as NoiseScheduleVP.
+ model_type: A `str`. The parameterization type of the diffusion model.
+ "noise" or "x_start" or "v" or "score".
+ model_kwargs: A `dict`. A dict for the other inputs of the model function.
+ guidance_type: A `str`. The type of the guidance for sampling.
+ "uncond" or "classifier" or "classifier-free".
+ condition: A pytorch tensor. The condition for the guided sampling.
+ Only used for "classifier" or "classifier-free" guidance type.
+ unconditional_condition: A pytorch tensor. The condition for the unconditional sampling.
+ Only used for "classifier-free" guidance type.
+ guidance_scale: A `float`. The scale for the guided sampling.
+ classifier_fn: A classifier function. Only used for the classifier guidance.
+ classifier_kwargs: A `dict`. A dict for the other inputs of the classifier function.
+ Returns:
+ A noise prediction model that accepts the noised data and the continuous time as the inputs.
+ """
+
+ def get_model_input_time(t_continuous):
+ """
+ Convert the continuous-time `t_continuous` (in [epsilon, T]) to the model input time.
+ For discrete-time DPMs, we convert `t_continuous` in [1 / N, 1] to `t_input` in [0, 1000 * (N - 1) / N].
+ For continuous-time DPMs, we just use `t_continuous`.
+ """
+ if noise_schedule.schedule == 'discrete':
+ return (t_continuous - 1. / noise_schedule.total_N) * 1000.
+ else:
+ return t_continuous
+
+ def noise_pred_fn(x, t_continuous, cond=None):
+ if t_continuous.reshape((-1,)).shape[0] == 1:
+ t_continuous = t_continuous.expand((x.shape[0]))
+ t_input = get_model_input_time(t_continuous)
+ if cond is None:
+ output = model(x, t_input, **model_kwargs)
+ else:
+ output = model(x, t_input, cond, **model_kwargs)
+ if model_type == "noise":
+ return output
+ elif model_type == "x_start":
+ alpha_t, sigma_t = noise_schedule.marginal_alpha(t_continuous), noise_schedule.marginal_std(t_continuous)
+ dims = x.dim()
+ return (x - expand_dims(alpha_t, dims) * output) / expand_dims(sigma_t, dims)
+ elif model_type == "v":
+ alpha_t, sigma_t = noise_schedule.marginal_alpha(t_continuous), noise_schedule.marginal_std(t_continuous)
+ dims = x.dim()
+ return expand_dims(alpha_t, dims) * output + expand_dims(sigma_t, dims) * x
+ elif model_type == "score":
+ sigma_t = noise_schedule.marginal_std(t_continuous)
+ dims = x.dim()
+ return -expand_dims(sigma_t, dims) * output
+
+ def cond_grad_fn(x, t_input):
+ """
+ Compute the gradient of the classifier, i.e. nabla_{x} log p_t(cond | x_t).
+ """
+ with torch.enable_grad():
+ x_in = x.detach().requires_grad_(True)
+ log_prob = classifier_fn(x_in, t_input, condition, **classifier_kwargs)
+ return torch.autograd.grad(log_prob.sum(), x_in)[0]
+
+ def model_fn(x, t_continuous):
+ """
+ The noise predicition model function that is used for DPM-Solver.
+ """
+ if t_continuous.reshape((-1,)).shape[0] == 1:
+ t_continuous = t_continuous.expand((x.shape[0]))
+ if guidance_type == "uncond":
+ return noise_pred_fn(x, t_continuous)
+ elif guidance_type == "classifier":
+ assert classifier_fn is not None
+ t_input = get_model_input_time(t_continuous)
+ cond_grad = cond_grad_fn(x, t_input)
+ sigma_t = noise_schedule.marginal_std(t_continuous)
+ noise = noise_pred_fn(x, t_continuous)
+ return noise - guidance_scale * expand_dims(sigma_t, dims=cond_grad.dim()) * cond_grad
+ elif guidance_type == "classifier-free":
+ if guidance_scale == 1. or unconditional_condition is None:
+ return noise_pred_fn(x, t_continuous, cond=condition)
+ else:
+ x_in = torch.cat([x] * 2)
+ t_in = torch.cat([t_continuous] * 2)
+ c_in = torch.cat([unconditional_condition, condition])
+ noise_uncond, noise = noise_pred_fn(x_in, t_in, cond=c_in).chunk(2)
+ return noise_uncond + guidance_scale * (noise - noise_uncond)
+
+ assert model_type in ["noise", "x_start", "v"]
+ assert guidance_type in ["uncond", "classifier", "classifier-free"]
+ return model_fn
+
+
+class DPM_Solver:
+ def __init__(self, model_fn, noise_schedule, predict_x0=False, thresholding=False, max_val=1.):
+ """Construct a DPM-Solver.
+ We support both the noise prediction model ("predicting epsilon") and the data prediction model ("predicting x0").
+ If `predict_x0` is False, we use the solver for the noise prediction model (DPM-Solver).
+ If `predict_x0` is True, we use the solver for the data prediction model (DPM-Solver++).
+ In such case, we further support the "dynamic thresholding" in [1] when `thresholding` is True.
+ The "dynamic thresholding" can greatly improve the sample quality for pixel-space DPMs with large guidance scales.
+ Args:
+ model_fn: A noise prediction model function which accepts the continuous-time input (t in [epsilon, T]):
+ ``
+ def model_fn(x, t_continuous):
+ return noise
+ ``
+ noise_schedule: A noise schedule object, such as NoiseScheduleVP.
+ predict_x0: A `bool`. If true, use the data prediction model; else, use the noise prediction model.
+ thresholding: A `bool`. Valid when `predict_x0` is True. Whether to use the "dynamic thresholding" in [1].
+ max_val: A `float`. Valid when both `predict_x0` and `thresholding` are True. The max value for thresholding.
+
+ [1] Chitwan Saharia, William Chan, Saurabh Saxena, Lala Li, Jay Whang, Emily Denton, Seyed Kamyar Seyed Ghasemipour, Burcu Karagol Ayan, S Sara Mahdavi, Rapha Gontijo Lopes, et al. Photorealistic text-to-image diffusion models with deep language understanding. arXiv preprint arXiv:2205.11487, 2022b.
+ """
+ self.model = model_fn
+ self.noise_schedule = noise_schedule
+ self.predict_x0 = predict_x0
+ self.thresholding = thresholding
+ self.max_val = max_val
+
+ def noise_prediction_fn(self, x, t):
+ """
+ Return the noise prediction model.
+ """
+ return self.model(x, t)
+
+ def data_prediction_fn(self, x, t):
+ """
+ Return the data prediction model (with thresholding).
+ """
+ noise = self.noise_prediction_fn(x, t)
+ dims = x.dim()
+ alpha_t, sigma_t = self.noise_schedule.marginal_alpha(t), self.noise_schedule.marginal_std(t)
+ x0 = (x - expand_dims(sigma_t, dims) * noise) / expand_dims(alpha_t, dims)
+ if self.thresholding:
+ p = 0.995 # A hyperparameter in the paper of "Imagen" [1].
+ s = torch.quantile(torch.abs(x0).reshape((x0.shape[0], -1)), p, dim=1)
+ s = expand_dims(torch.maximum(s, self.max_val * torch.ones_like(s).to(s.device)), dims)
+ x0 = torch.clamp(x0, -s, s) / s
+ return x0
+
+ def model_fn(self, x, t):
+ """
+ Convert the model to the noise prediction model or the data prediction model.
+ """
+ if self.predict_x0:
+ return self.data_prediction_fn(x, t)
+ else:
+ return self.noise_prediction_fn(x, t)
+
+ def get_time_steps(self, skip_type, t_T, t_0, N, device):
+ """Compute the intermediate time steps for sampling.
+ Args:
+ skip_type: A `str`. The type for the spacing of the time steps. We support three types:
+ - 'logSNR': uniform logSNR for the time steps.
+ - 'time_uniform': uniform time for the time steps. (**Recommended for high-resolutional data**.)
+ - 'time_quadratic': quadratic time for the time steps. (Used in DDIM for low-resolutional data.)
+ t_T: A `float`. The starting time of the sampling (default is T).
+ t_0: A `float`. The ending time of the sampling (default is epsilon).
+ N: A `int`. The total number of the spacing of the time steps.
+ device: A torch device.
+ Returns:
+ A pytorch tensor of the time steps, with the shape (N + 1,).
+ """
+ if skip_type == 'logSNR':
+ lambda_T = self.noise_schedule.marginal_lambda(torch.tensor(t_T).to(device))
+ lambda_0 = self.noise_schedule.marginal_lambda(torch.tensor(t_0).to(device))
+ logSNR_steps = torch.linspace(lambda_T.cpu().item(), lambda_0.cpu().item(), N + 1).to(device)
+ return self.noise_schedule.inverse_lambda(logSNR_steps)
+ elif skip_type == 'time_uniform':
+ return torch.linspace(t_T, t_0, N + 1).to(device)
+ elif skip_type == 'time_quadratic':
+ t_order = 2
+ t = torch.linspace(t_T ** (1. / t_order), t_0 ** (1. / t_order), N + 1).pow(t_order).to(device)
+ return t
+ else:
+ raise ValueError(
+ "Unsupported skip_type {}, need to be 'logSNR' or 'time_uniform' or 'time_quadratic'".format(skip_type))
+
+ def get_orders_and_timesteps_for_singlestep_solver(self, steps, order, skip_type, t_T, t_0, device):
+ """
+ Get the order of each step for sampling by the singlestep DPM-Solver.
+ We combine both DPM-Solver-1,2,3 to use all the function evaluations, which is named as "DPM-Solver-fast".
+ Given a fixed number of function evaluations by `steps`, the sampling procedure by DPM-Solver-fast is:
+ - If order == 1:
+ We take `steps` of DPM-Solver-1 (i.e. DDIM).
+ - If order == 2:
+ - Denote K = (steps // 2). We take K or (K + 1) intermediate time steps for sampling.
+ - If steps % 2 == 0, we use K steps of DPM-Solver-2.
+ - If steps % 2 == 1, we use K steps of DPM-Solver-2 and 1 step of DPM-Solver-1.
+ - If order == 3:
+ - Denote K = (steps // 3 + 1). We take K intermediate time steps for sampling.
+ - If steps % 3 == 0, we use (K - 2) steps of DPM-Solver-3, and 1 step of DPM-Solver-2 and 1 step of DPM-Solver-1.
+ - If steps % 3 == 1, we use (K - 1) steps of DPM-Solver-3 and 1 step of DPM-Solver-1.
+ - If steps % 3 == 2, we use (K - 1) steps of DPM-Solver-3 and 1 step of DPM-Solver-2.
+ ============================================
+ Args:
+ order: A `int`. The max order for the solver (2 or 3).
+ steps: A `int`. The total number of function evaluations (NFE).
+ skip_type: A `str`. The type for the spacing of the time steps. We support three types:
+ - 'logSNR': uniform logSNR for the time steps.
+ - 'time_uniform': uniform time for the time steps. (**Recommended for high-resolutional data**.)
+ - 'time_quadratic': quadratic time for the time steps. (Used in DDIM for low-resolutional data.)
+ t_T: A `float`. The starting time of the sampling (default is T).
+ t_0: A `float`. The ending time of the sampling (default is epsilon).
+ device: A torch device.
+ Returns:
+ orders: A list of the solver order of each step.
+ """
+ if order == 3:
+ K = steps // 3 + 1
+ if steps % 3 == 0:
+ orders = [3, ] * (K - 2) + [2, 1]
+ elif steps % 3 == 1:
+ orders = [3, ] * (K - 1) + [1]
+ else:
+ orders = [3, ] * (K - 1) + [2]
+ elif order == 2:
+ if steps % 2 == 0:
+ K = steps // 2
+ orders = [2, ] * K
+ else:
+ K = steps // 2 + 1
+ orders = [2, ] * (K - 1) + [1]
+ elif order == 1:
+ K = 1
+ orders = [1, ] * steps
+ else:
+ raise ValueError("'order' must be '1' or '2' or '3'.")
+ if skip_type == 'logSNR':
+ # To reproduce the results in DPM-Solver paper
+ timesteps_outer = self.get_time_steps(skip_type, t_T, t_0, K, device)
+ else:
+ timesteps_outer = self.get_time_steps(skip_type, t_T, t_0, steps, device)[
+ torch.cumsum(torch.tensor([0, ] + orders)).to(device)]
+ return timesteps_outer, orders
+
+ def denoise_to_zero_fn(self, x, s):
+ """
+ Denoise at the final step, which is equivalent to solve the ODE from lambda_s to infty by first-order discretization.
+ """
+ return self.data_prediction_fn(x, s)
+
+ def dpm_solver_first_update(self, x, s, t, model_s=None, return_intermediate=False):
+ """
+ DPM-Solver-1 (equivalent to DDIM) from time `s` to time `t`.
+ Args:
+ x: A pytorch tensor. The initial value at time `s`.
+ s: A pytorch tensor. The starting time, with the shape (x.shape[0],).
+ t: A pytorch tensor. The ending time, with the shape (x.shape[0],).
+ model_s: A pytorch tensor. The model function evaluated at time `s`.
+ If `model_s` is None, we evaluate the model by `x` and `s`; otherwise we directly use it.
+ return_intermediate: A `bool`. If true, also return the model value at time `s`.
+ Returns:
+ x_t: A pytorch tensor. The approximated solution at time `t`.
+ """
+ ns = self.noise_schedule
+ dims = x.dim()
+ lambda_s, lambda_t = ns.marginal_lambda(s), ns.marginal_lambda(t)
+ h = lambda_t - lambda_s
+ log_alpha_s, log_alpha_t = ns.marginal_log_mean_coeff(s), ns.marginal_log_mean_coeff(t)
+ sigma_s, sigma_t = ns.marginal_std(s), ns.marginal_std(t)
+ alpha_t = torch.exp(log_alpha_t)
+
+ if self.predict_x0:
+ phi_1 = torch.expm1(-h)
+ if model_s is None:
+ model_s = self.model_fn(x, s)
+ x_t = (
+ expand_dims(sigma_t / sigma_s, dims) * x
+ - expand_dims(alpha_t * phi_1, dims) * model_s
+ )
+ if return_intermediate:
+ return x_t, {'model_s': model_s}
+ else:
+ return x_t
+ else:
+ phi_1 = torch.expm1(h)
+ if model_s is None:
+ model_s = self.model_fn(x, s)
+ x_t = (
+ expand_dims(torch.exp(log_alpha_t - log_alpha_s), dims) * x
+ - expand_dims(sigma_t * phi_1, dims) * model_s
+ )
+ if return_intermediate:
+ return x_t, {'model_s': model_s}
+ else:
+ return x_t
+
+ def singlestep_dpm_solver_second_update(self, x, s, t, r1=0.5, model_s=None, return_intermediate=False,
+ solver_type='dpm_solver'):
+ """
+ Singlestep solver DPM-Solver-2 from time `s` to time `t`.
+ Args:
+ x: A pytorch tensor. The initial value at time `s`.
+ s: A pytorch tensor. The starting time, with the shape (x.shape[0],).
+ t: A pytorch tensor. The ending time, with the shape (x.shape[0],).
+ r1: A `float`. The hyperparameter of the second-order solver.
+ model_s: A pytorch tensor. The model function evaluated at time `s`.
+ If `model_s` is None, we evaluate the model by `x` and `s`; otherwise we directly use it.
+ return_intermediate: A `bool`. If true, also return the model value at time `s` and `s1` (the intermediate time).
+ solver_type: either 'dpm_solver' or 'taylor'. The type for the high-order solvers.
+ The type slightly impacts the performance. We recommend to use 'dpm_solver' type.
+ Returns:
+ x_t: A pytorch tensor. The approximated solution at time `t`.
+ """
+ if solver_type not in ['dpm_solver', 'taylor']:
+ raise ValueError("'solver_type' must be either 'dpm_solver' or 'taylor', got {}".format(solver_type))
+ if r1 is None:
+ r1 = 0.5
+ ns = self.noise_schedule
+ dims = x.dim()
+ lambda_s, lambda_t = ns.marginal_lambda(s), ns.marginal_lambda(t)
+ h = lambda_t - lambda_s
+ lambda_s1 = lambda_s + r1 * h
+ s1 = ns.inverse_lambda(lambda_s1)
+ log_alpha_s, log_alpha_s1, log_alpha_t = ns.marginal_log_mean_coeff(s), ns.marginal_log_mean_coeff(
+ s1), ns.marginal_log_mean_coeff(t)
+ sigma_s, sigma_s1, sigma_t = ns.marginal_std(s), ns.marginal_std(s1), ns.marginal_std(t)
+ alpha_s1, alpha_t = torch.exp(log_alpha_s1), torch.exp(log_alpha_t)
+
+ if self.predict_x0:
+ phi_11 = torch.expm1(-r1 * h)
+ phi_1 = torch.expm1(-h)
+
+ if model_s is None:
+ model_s = self.model_fn(x, s)
+ x_s1 = (
+ expand_dims(sigma_s1 / sigma_s, dims) * x
+ - expand_dims(alpha_s1 * phi_11, dims) * model_s
+ )
+ model_s1 = self.model_fn(x_s1, s1)
+ if solver_type == 'dpm_solver':
+ x_t = (
+ expand_dims(sigma_t / sigma_s, dims) * x
+ - expand_dims(alpha_t * phi_1, dims) * model_s
+ - (0.5 / r1) * expand_dims(alpha_t * phi_1, dims) * (model_s1 - model_s)
+ )
+ elif solver_type == 'taylor':
+ x_t = (
+ expand_dims(sigma_t / sigma_s, dims) * x
+ - expand_dims(alpha_t * phi_1, dims) * model_s
+ + (1. / r1) * expand_dims(alpha_t * ((torch.exp(-h) - 1.) / h + 1.), dims) * (
+ model_s1 - model_s)
+ )
+ else:
+ phi_11 = torch.expm1(r1 * h)
+ phi_1 = torch.expm1(h)
+
+ if model_s is None:
+ model_s = self.model_fn(x, s)
+ x_s1 = (
+ expand_dims(torch.exp(log_alpha_s1 - log_alpha_s), dims) * x
+ - expand_dims(sigma_s1 * phi_11, dims) * model_s
+ )
+ model_s1 = self.model_fn(x_s1, s1)
+ if solver_type == 'dpm_solver':
+ x_t = (
+ expand_dims(torch.exp(log_alpha_t - log_alpha_s), dims) * x
+ - expand_dims(sigma_t * phi_1, dims) * model_s
+ - (0.5 / r1) * expand_dims(sigma_t * phi_1, dims) * (model_s1 - model_s)
+ )
+ elif solver_type == 'taylor':
+ x_t = (
+ expand_dims(torch.exp(log_alpha_t - log_alpha_s), dims) * x
+ - expand_dims(sigma_t * phi_1, dims) * model_s
+ - (1. / r1) * expand_dims(sigma_t * ((torch.exp(h) - 1.) / h - 1.), dims) * (model_s1 - model_s)
+ )
+ if return_intermediate:
+ return x_t, {'model_s': model_s, 'model_s1': model_s1}
+ else:
+ return x_t
+
+ def singlestep_dpm_solver_third_update(self, x, s, t, r1=1. / 3., r2=2. / 3., model_s=None, model_s1=None,
+ return_intermediate=False, solver_type='dpm_solver'):
+ """
+ Singlestep solver DPM-Solver-3 from time `s` to time `t`.
+ Args:
+ x: A pytorch tensor. The initial value at time `s`.
+ s: A pytorch tensor. The starting time, with the shape (x.shape[0],).
+ t: A pytorch tensor. The ending time, with the shape (x.shape[0],).
+ r1: A `float`. The hyperparameter of the third-order solver.
+ r2: A `float`. The hyperparameter of the third-order solver.
+ model_s: A pytorch tensor. The model function evaluated at time `s`.
+ If `model_s` is None, we evaluate the model by `x` and `s`; otherwise we directly use it.
+ model_s1: A pytorch tensor. The model function evaluated at time `s1` (the intermediate time given by `r1`).
+ If `model_s1` is None, we evaluate the model at `s1`; otherwise we directly use it.
+ return_intermediate: A `bool`. If true, also return the model value at time `s`, `s1` and `s2` (the intermediate times).
+ solver_type: either 'dpm_solver' or 'taylor'. The type for the high-order solvers.
+ The type slightly impacts the performance. We recommend to use 'dpm_solver' type.
+ Returns:
+ x_t: A pytorch tensor. The approximated solution at time `t`.
+ """
+ if solver_type not in ['dpm_solver', 'taylor']:
+ raise ValueError("'solver_type' must be either 'dpm_solver' or 'taylor', got {}".format(solver_type))
+ if r1 is None:
+ r1 = 1. / 3.
+ if r2 is None:
+ r2 = 2. / 3.
+ ns = self.noise_schedule
+ dims = x.dim()
+ lambda_s, lambda_t = ns.marginal_lambda(s), ns.marginal_lambda(t)
+ h = lambda_t - lambda_s
+ lambda_s1 = lambda_s + r1 * h
+ lambda_s2 = lambda_s + r2 * h
+ s1 = ns.inverse_lambda(lambda_s1)
+ s2 = ns.inverse_lambda(lambda_s2)
+ log_alpha_s, log_alpha_s1, log_alpha_s2, log_alpha_t = ns.marginal_log_mean_coeff(
+ s), ns.marginal_log_mean_coeff(s1), ns.marginal_log_mean_coeff(s2), ns.marginal_log_mean_coeff(t)
+ sigma_s, sigma_s1, sigma_s2, sigma_t = ns.marginal_std(s), ns.marginal_std(s1), ns.marginal_std(
+ s2), ns.marginal_std(t)
+ alpha_s1, alpha_s2, alpha_t = torch.exp(log_alpha_s1), torch.exp(log_alpha_s2), torch.exp(log_alpha_t)
+
+ if self.predict_x0:
+ phi_11 = torch.expm1(-r1 * h)
+ phi_12 = torch.expm1(-r2 * h)
+ phi_1 = torch.expm1(-h)
+ phi_22 = torch.expm1(-r2 * h) / (r2 * h) + 1.
+ phi_2 = phi_1 / h + 1.
+ phi_3 = phi_2 / h - 0.5
+
+ if model_s is None:
+ model_s = self.model_fn(x, s)
+ if model_s1 is None:
+ x_s1 = (
+ expand_dims(sigma_s1 / sigma_s, dims) * x
+ - expand_dims(alpha_s1 * phi_11, dims) * model_s
+ )
+ model_s1 = self.model_fn(x_s1, s1)
+ x_s2 = (
+ expand_dims(sigma_s2 / sigma_s, dims) * x
+ - expand_dims(alpha_s2 * phi_12, dims) * model_s
+ + r2 / r1 * expand_dims(alpha_s2 * phi_22, dims) * (model_s1 - model_s)
+ )
+ model_s2 = self.model_fn(x_s2, s2)
+ if solver_type == 'dpm_solver':
+ x_t = (
+ expand_dims(sigma_t / sigma_s, dims) * x
+ - expand_dims(alpha_t * phi_1, dims) * model_s
+ + (1. / r2) * expand_dims(alpha_t * phi_2, dims) * (model_s2 - model_s)
+ )
+ elif solver_type == 'taylor':
+ D1_0 = (1. / r1) * (model_s1 - model_s)
+ D1_1 = (1. / r2) * (model_s2 - model_s)
+ D1 = (r2 * D1_0 - r1 * D1_1) / (r2 - r1)
+ D2 = 2. * (D1_1 - D1_0) / (r2 - r1)
+ x_t = (
+ expand_dims(sigma_t / sigma_s, dims) * x
+ - expand_dims(alpha_t * phi_1, dims) * model_s
+ + expand_dims(alpha_t * phi_2, dims) * D1
+ - expand_dims(alpha_t * phi_3, dims) * D2
+ )
+ else:
+ phi_11 = torch.expm1(r1 * h)
+ phi_12 = torch.expm1(r2 * h)
+ phi_1 = torch.expm1(h)
+ phi_22 = torch.expm1(r2 * h) / (r2 * h) - 1.
+ phi_2 = phi_1 / h - 1.
+ phi_3 = phi_2 / h - 0.5
+
+ if model_s is None:
+ model_s = self.model_fn(x, s)
+ if model_s1 is None:
+ x_s1 = (
+ expand_dims(torch.exp(log_alpha_s1 - log_alpha_s), dims) * x
+ - expand_dims(sigma_s1 * phi_11, dims) * model_s
+ )
+ model_s1 = self.model_fn(x_s1, s1)
+ x_s2 = (
+ expand_dims(torch.exp(log_alpha_s2 - log_alpha_s), dims) * x
+ - expand_dims(sigma_s2 * phi_12, dims) * model_s
+ - r2 / r1 * expand_dims(sigma_s2 * phi_22, dims) * (model_s1 - model_s)
+ )
+ model_s2 = self.model_fn(x_s2, s2)
+ if solver_type == 'dpm_solver':
+ x_t = (
+ expand_dims(torch.exp(log_alpha_t - log_alpha_s), dims) * x
+ - expand_dims(sigma_t * phi_1, dims) * model_s
+ - (1. / r2) * expand_dims(sigma_t * phi_2, dims) * (model_s2 - model_s)
+ )
+ elif solver_type == 'taylor':
+ D1_0 = (1. / r1) * (model_s1 - model_s)
+ D1_1 = (1. / r2) * (model_s2 - model_s)
+ D1 = (r2 * D1_0 - r1 * D1_1) / (r2 - r1)
+ D2 = 2. * (D1_1 - D1_0) / (r2 - r1)
+ x_t = (
+ expand_dims(torch.exp(log_alpha_t - log_alpha_s), dims) * x
+ - expand_dims(sigma_t * phi_1, dims) * model_s
+ - expand_dims(sigma_t * phi_2, dims) * D1
+ - expand_dims(sigma_t * phi_3, dims) * D2
+ )
+
+ if return_intermediate:
+ return x_t, {'model_s': model_s, 'model_s1': model_s1, 'model_s2': model_s2}
+ else:
+ return x_t
+
+ def multistep_dpm_solver_second_update(self, x, model_prev_list, t_prev_list, t, solver_type="dpm_solver"):
+ """
+ Multistep solver DPM-Solver-2 from time `t_prev_list[-1]` to time `t`.
+ Args:
+ x: A pytorch tensor. The initial value at time `s`.
+ model_prev_list: A list of pytorch tensor. The previous computed model values.
+ t_prev_list: A list of pytorch tensor. The previous times, each time has the shape (x.shape[0],)
+ t: A pytorch tensor. The ending time, with the shape (x.shape[0],).
+ solver_type: either 'dpm_solver' or 'taylor'. The type for the high-order solvers.
+ The type slightly impacts the performance. We recommend to use 'dpm_solver' type.
+ Returns:
+ x_t: A pytorch tensor. The approximated solution at time `t`.
+ """
+ if solver_type not in ['dpm_solver', 'taylor']:
+ raise ValueError("'solver_type' must be either 'dpm_solver' or 'taylor', got {}".format(solver_type))
+ ns = self.noise_schedule
+ dims = x.dim()
+ model_prev_1, model_prev_0 = model_prev_list
+ t_prev_1, t_prev_0 = t_prev_list
+ lambda_prev_1, lambda_prev_0, lambda_t = ns.marginal_lambda(t_prev_1), ns.marginal_lambda(
+ t_prev_0), ns.marginal_lambda(t)
+ log_alpha_prev_0, log_alpha_t = ns.marginal_log_mean_coeff(t_prev_0), ns.marginal_log_mean_coeff(t)
+ sigma_prev_0, sigma_t = ns.marginal_std(t_prev_0), ns.marginal_std(t)
+ alpha_t = torch.exp(log_alpha_t)
+
+ h_0 = lambda_prev_0 - lambda_prev_1
+ h = lambda_t - lambda_prev_0
+ r0 = h_0 / h
+ D1_0 = expand_dims(1. / r0, dims) * (model_prev_0 - model_prev_1)
+ if self.predict_x0:
+ if solver_type == 'dpm_solver':
+ x_t = (
+ expand_dims(sigma_t / sigma_prev_0, dims) * x
+ - expand_dims(alpha_t * (torch.exp(-h) - 1.), dims) * model_prev_0
+ - 0.5 * expand_dims(alpha_t * (torch.exp(-h) - 1.), dims) * D1_0
+ )
+ elif solver_type == 'taylor':
+ x_t = (
+ expand_dims(sigma_t / sigma_prev_0, dims) * x
+ - expand_dims(alpha_t * (torch.exp(-h) - 1.), dims) * model_prev_0
+ + expand_dims(alpha_t * ((torch.exp(-h) - 1.) / h + 1.), dims) * D1_0
+ )
+ else:
+ if solver_type == 'dpm_solver':
+ x_t = (
+ expand_dims(torch.exp(log_alpha_t - log_alpha_prev_0), dims) * x
+ - expand_dims(sigma_t * (torch.exp(h) - 1.), dims) * model_prev_0
+ - 0.5 * expand_dims(sigma_t * (torch.exp(h) - 1.), dims) * D1_0
+ )
+ elif solver_type == 'taylor':
+ x_t = (
+ expand_dims(torch.exp(log_alpha_t - log_alpha_prev_0), dims) * x
+ - expand_dims(sigma_t * (torch.exp(h) - 1.), dims) * model_prev_0
+ - expand_dims(sigma_t * ((torch.exp(h) - 1.) / h - 1.), dims) * D1_0
+ )
+ return x_t
+
+ def multistep_dpm_solver_third_update(self, x, model_prev_list, t_prev_list, t, solver_type='dpm_solver'):
+ """
+ Multistep solver DPM-Solver-3 from time `t_prev_list[-1]` to time `t`.
+ Args:
+ x: A pytorch tensor. The initial value at time `s`.
+ model_prev_list: A list of pytorch tensor. The previous computed model values.
+ t_prev_list: A list of pytorch tensor. The previous times, each time has the shape (x.shape[0],)
+ t: A pytorch tensor. The ending time, with the shape (x.shape[0],).
+ solver_type: either 'dpm_solver' or 'taylor'. The type for the high-order solvers.
+ The type slightly impacts the performance. We recommend to use 'dpm_solver' type.
+ Returns:
+ x_t: A pytorch tensor. The approximated solution at time `t`.
+ """
+ ns = self.noise_schedule
+ dims = x.dim()
+ model_prev_2, model_prev_1, model_prev_0 = model_prev_list
+ t_prev_2, t_prev_1, t_prev_0 = t_prev_list
+ lambda_prev_2, lambda_prev_1, lambda_prev_0, lambda_t = ns.marginal_lambda(t_prev_2), ns.marginal_lambda(
+ t_prev_1), ns.marginal_lambda(t_prev_0), ns.marginal_lambda(t)
+ log_alpha_prev_0, log_alpha_t = ns.marginal_log_mean_coeff(t_prev_0), ns.marginal_log_mean_coeff(t)
+ sigma_prev_0, sigma_t = ns.marginal_std(t_prev_0), ns.marginal_std(t)
+ alpha_t = torch.exp(log_alpha_t)
+
+ h_1 = lambda_prev_1 - lambda_prev_2
+ h_0 = lambda_prev_0 - lambda_prev_1
+ h = lambda_t - lambda_prev_0
+ r0, r1 = h_0 / h, h_1 / h
+ D1_0 = expand_dims(1. / r0, dims) * (model_prev_0 - model_prev_1)
+ D1_1 = expand_dims(1. / r1, dims) * (model_prev_1 - model_prev_2)
+ D1 = D1_0 + expand_dims(r0 / (r0 + r1), dims) * (D1_0 - D1_1)
+ D2 = expand_dims(1. / (r0 + r1), dims) * (D1_0 - D1_1)
+ if self.predict_x0:
+ x_t = (
+ expand_dims(sigma_t / sigma_prev_0, dims) * x
+ - expand_dims(alpha_t * (torch.exp(-h) - 1.), dims) * model_prev_0
+ + expand_dims(alpha_t * ((torch.exp(-h) - 1.) / h + 1.), dims) * D1
+ - expand_dims(alpha_t * ((torch.exp(-h) - 1. + h) / h ** 2 - 0.5), dims) * D2
+ )
+ else:
+ x_t = (
+ expand_dims(torch.exp(log_alpha_t - log_alpha_prev_0), dims) * x
+ - expand_dims(sigma_t * (torch.exp(h) - 1.), dims) * model_prev_0
+ - expand_dims(sigma_t * ((torch.exp(h) - 1.) / h - 1.), dims) * D1
+ - expand_dims(sigma_t * ((torch.exp(h) - 1. - h) / h ** 2 - 0.5), dims) * D2
+ )
+ return x_t
+
+ def singlestep_dpm_solver_update(self, x, s, t, order, return_intermediate=False, solver_type='dpm_solver', r1=None,
+ r2=None):
+ """
+ Singlestep DPM-Solver with the order `order` from time `s` to time `t`.
+ Args:
+ x: A pytorch tensor. The initial value at time `s`.
+ s: A pytorch tensor. The starting time, with the shape (x.shape[0],).
+ t: A pytorch tensor. The ending time, with the shape (x.shape[0],).
+ order: A `int`. The order of DPM-Solver. We only support order == 1 or 2 or 3.
+ return_intermediate: A `bool`. If true, also return the model value at time `s`, `s1` and `s2` (the intermediate times).
+ solver_type: either 'dpm_solver' or 'taylor'. The type for the high-order solvers.
+ The type slightly impacts the performance. We recommend to use 'dpm_solver' type.
+ r1: A `float`. The hyperparameter of the second-order or third-order solver.
+ r2: A `float`. The hyperparameter of the third-order solver.
+ Returns:
+ x_t: A pytorch tensor. The approximated solution at time `t`.
+ """
+ if order == 1:
+ return self.dpm_solver_first_update(x, s, t, return_intermediate=return_intermediate)
+ elif order == 2:
+ return self.singlestep_dpm_solver_second_update(x, s, t, return_intermediate=return_intermediate,
+ solver_type=solver_type, r1=r1)
+ elif order == 3:
+ return self.singlestep_dpm_solver_third_update(x, s, t, return_intermediate=return_intermediate,
+ solver_type=solver_type, r1=r1, r2=r2)
+ else:
+ raise ValueError("Solver order must be 1 or 2 or 3, got {}".format(order))
+
+ def multistep_dpm_solver_update(self, x, model_prev_list, t_prev_list, t, order, solver_type='dpm_solver'):
+ """
+ Multistep DPM-Solver with the order `order` from time `t_prev_list[-1]` to time `t`.
+ Args:
+ x: A pytorch tensor. The initial value at time `s`.
+ model_prev_list: A list of pytorch tensor. The previous computed model values.
+ t_prev_list: A list of pytorch tensor. The previous times, each time has the shape (x.shape[0],)
+ t: A pytorch tensor. The ending time, with the shape (x.shape[0],).
+ order: A `int`. The order of DPM-Solver. We only support order == 1 or 2 or 3.
+ solver_type: either 'dpm_solver' or 'taylor'. The type for the high-order solvers.
+ The type slightly impacts the performance. We recommend to use 'dpm_solver' type.
+ Returns:
+ x_t: A pytorch tensor. The approximated solution at time `t`.
+ """
+ if order == 1:
+ return self.dpm_solver_first_update(x, t_prev_list[-1], t, model_s=model_prev_list[-1])
+ elif order == 2:
+ return self.multistep_dpm_solver_second_update(x, model_prev_list, t_prev_list, t, solver_type=solver_type)
+ elif order == 3:
+ return self.multistep_dpm_solver_third_update(x, model_prev_list, t_prev_list, t, solver_type=solver_type)
+ else:
+ raise ValueError("Solver order must be 1 or 2 or 3, got {}".format(order))
+
+ def dpm_solver_adaptive(self, x, order, t_T, t_0, h_init=0.05, atol=0.0078, rtol=0.05, theta=0.9, t_err=1e-5,
+ solver_type='dpm_solver'):
+ """
+ The adaptive step size solver based on singlestep DPM-Solver.
+ Args:
+ x: A pytorch tensor. The initial value at time `t_T`.
+ order: A `int`. The (higher) order of the solver. We only support order == 2 or 3.
+ t_T: A `float`. The starting time of the sampling (default is T).
+ t_0: A `float`. The ending time of the sampling (default is epsilon).
+ h_init: A `float`. The initial step size (for logSNR).
+ atol: A `float`. The absolute tolerance of the solver. For image data, the default setting is 0.0078, followed [1].
+ rtol: A `float`. The relative tolerance of the solver. The default setting is 0.05.
+ theta: A `float`. The safety hyperparameter for adapting the step size. The default setting is 0.9, followed [1].
+ t_err: A `float`. The tolerance for the time. We solve the diffusion ODE until the absolute error between the
+ current time and `t_0` is less than `t_err`. The default setting is 1e-5.
+ solver_type: either 'dpm_solver' or 'taylor'. The type for the high-order solvers.
+ The type slightly impacts the performance. We recommend to use 'dpm_solver' type.
+ Returns:
+ x_0: A pytorch tensor. The approximated solution at time `t_0`.
+ [1] A. Jolicoeur-Martineau, K. Li, R. Piché-Taillefer, T. Kachman, and I. Mitliagkas, "Gotta go fast when generating data with score-based models," arXiv preprint arXiv:2105.14080, 2021.
+ """
+ ns = self.noise_schedule
+ s = t_T * torch.ones((x.shape[0],)).to(x)
+ lambda_s = ns.marginal_lambda(s)
+ lambda_0 = ns.marginal_lambda(t_0 * torch.ones_like(s).to(x))
+ h = h_init * torch.ones_like(s).to(x)
+ x_prev = x
+ nfe = 0
+ if order == 2:
+ r1 = 0.5
+ lower_update = lambda x, s, t: self.dpm_solver_first_update(x, s, t, return_intermediate=True)
+ higher_update = lambda x, s, t, **kwargs: self.singlestep_dpm_solver_second_update(x, s, t, r1=r1,
+ solver_type=solver_type,
+ **kwargs)
+ elif order == 3:
+ r1, r2 = 1. / 3., 2. / 3.
+ lower_update = lambda x, s, t: self.singlestep_dpm_solver_second_update(x, s, t, r1=r1,
+ return_intermediate=True,
+ solver_type=solver_type)
+ higher_update = lambda x, s, t, **kwargs: self.singlestep_dpm_solver_third_update(x, s, t, r1=r1, r2=r2,
+ solver_type=solver_type,
+ **kwargs)
+ else:
+ raise ValueError("For adaptive step size solver, order must be 2 or 3, got {}".format(order))
+ while torch.abs((s - t_0)).mean() > t_err:
+ t = ns.inverse_lambda(lambda_s + h)
+ x_lower, lower_noise_kwargs = lower_update(x, s, t)
+ x_higher = higher_update(x, s, t, **lower_noise_kwargs)
+ delta = torch.max(torch.ones_like(x).to(x) * atol, rtol * torch.max(torch.abs(x_lower), torch.abs(x_prev)))
+ norm_fn = lambda v: torch.sqrt(torch.square(v.reshape((v.shape[0], -1))).mean(dim=-1, keepdim=True))
+ E = norm_fn((x_higher - x_lower) / delta).max()
+ if torch.all(E <= 1.):
+ x = x_higher
+ s = t
+ x_prev = x_lower
+ lambda_s = ns.marginal_lambda(s)
+ h = torch.min(theta * h * torch.float_power(E, -1. / order).float(), lambda_0 - lambda_s)
+ nfe += order
+ print('adaptive solver nfe', nfe)
+ return x
+
+ def sample(self, x, steps=20, t_start=None, t_end=None, order=3, skip_type='time_uniform',
+ method='singlestep', lower_order_final=True, denoise_to_zero=False, solver_type='dpm_solver',
+ atol=0.0078, rtol=0.05,
+ ):
+ """
+ Compute the sample at time `t_end` by DPM-Solver, given the initial `x` at time `t_start`.
+ =====================================================
+ We support the following algorithms for both noise prediction model and data prediction model:
+ - 'singlestep':
+ Singlestep DPM-Solver (i.e. "DPM-Solver-fast" in the paper), which combines different orders of singlestep DPM-Solver.
+ We combine all the singlestep solvers with order <= `order` to use up all the function evaluations (steps).
+ The total number of function evaluations (NFE) == `steps`.
+ Given a fixed NFE == `steps`, the sampling procedure is:
+ - If `order` == 1:
+ - Denote K = steps. We use K steps of DPM-Solver-1 (i.e. DDIM).
+ - If `order` == 2:
+ - Denote K = (steps // 2) + (steps % 2). We take K intermediate time steps for sampling.
+ - If steps % 2 == 0, we use K steps of singlestep DPM-Solver-2.
+ - If steps % 2 == 1, we use (K - 1) steps of singlestep DPM-Solver-2 and 1 step of DPM-Solver-1.
+ - If `order` == 3:
+ - Denote K = (steps // 3 + 1). We take K intermediate time steps for sampling.
+ - If steps % 3 == 0, we use (K - 2) steps of singlestep DPM-Solver-3, and 1 step of singlestep DPM-Solver-2 and 1 step of DPM-Solver-1.
+ - If steps % 3 == 1, we use (K - 1) steps of singlestep DPM-Solver-3 and 1 step of DPM-Solver-1.
+ - If steps % 3 == 2, we use (K - 1) steps of singlestep DPM-Solver-3 and 1 step of singlestep DPM-Solver-2.
+ - 'multistep':
+ Multistep DPM-Solver with the order of `order`. The total number of function evaluations (NFE) == `steps`.
+ We initialize the first `order` values by lower order multistep solvers.
+ Given a fixed NFE == `steps`, the sampling procedure is:
+ Denote K = steps.
+ - If `order` == 1:
+ - We use K steps of DPM-Solver-1 (i.e. DDIM).
+ - If `order` == 2:
+ - We firstly use 1 step of DPM-Solver-1, then use (K - 1) step of multistep DPM-Solver-2.
+ - If `order` == 3:
+ - We firstly use 1 step of DPM-Solver-1, then 1 step of multistep DPM-Solver-2, then (K - 2) step of multistep DPM-Solver-3.
+ - 'singlestep_fixed':
+ Fixed order singlestep DPM-Solver (i.e. DPM-Solver-1 or singlestep DPM-Solver-2 or singlestep DPM-Solver-3).
+ We use singlestep DPM-Solver-`order` for `order`=1 or 2 or 3, with total [`steps` // `order`] * `order` NFE.
+ - 'adaptive':
+ Adaptive step size DPM-Solver (i.e. "DPM-Solver-12" and "DPM-Solver-23" in the paper).
+ We ignore `steps` and use adaptive step size DPM-Solver with a higher order of `order`.
+ You can adjust the absolute tolerance `atol` and the relative tolerance `rtol` to balance the computatation costs
+ (NFE) and the sample quality.
+ - If `order` == 2, we use DPM-Solver-12 which combines DPM-Solver-1 and singlestep DPM-Solver-2.
+ - If `order` == 3, we use DPM-Solver-23 which combines singlestep DPM-Solver-2 and singlestep DPM-Solver-3.
+ =====================================================
+ Some advices for choosing the algorithm:
+ - For **unconditional sampling** or **guided sampling with small guidance scale** by DPMs:
+ Use singlestep DPM-Solver ("DPM-Solver-fast" in the paper) with `order = 3`.
+ e.g.
+ >>> dpm_solver = DPM_Solver(model_fn, noise_schedule, predict_x0=False)
+ >>> x_sample = dpm_solver.sample(x, steps=steps, t_start=t_start, t_end=t_end, order=3,
+ skip_type='time_uniform', method='singlestep')
+ - For **guided sampling with large guidance scale** by DPMs:
+ Use multistep DPM-Solver with `predict_x0 = True` and `order = 2`.
+ e.g.
+ >>> dpm_solver = DPM_Solver(model_fn, noise_schedule, predict_x0=True)
+ >>> x_sample = dpm_solver.sample(x, steps=steps, t_start=t_start, t_end=t_end, order=2,
+ skip_type='time_uniform', method='multistep')
+ We support three types of `skip_type`:
+ - 'logSNR': uniform logSNR for the time steps. **Recommended for low-resolutional images**
+ - 'time_uniform': uniform time for the time steps. **Recommended for high-resolutional images**.
+ - 'time_quadratic': quadratic time for the time steps.
+ =====================================================
+ Args:
+ x: A pytorch tensor. The initial value at time `t_start`
+ e.g. if `t_start` == T, then `x` is a sample from the standard normal distribution.
+ steps: A `int`. The total number of function evaluations (NFE).
+ t_start: A `float`. The starting time of the sampling.
+ If `T` is None, we use self.noise_schedule.T (default is 1.0).
+ t_end: A `float`. The ending time of the sampling.
+ If `t_end` is None, we use 1. / self.noise_schedule.total_N.
+ e.g. if total_N == 1000, we have `t_end` == 1e-3.
+ For discrete-time DPMs:
+ - We recommend `t_end` == 1. / self.noise_schedule.total_N.
+ For continuous-time DPMs:
+ - We recommend `t_end` == 1e-3 when `steps` <= 15; and `t_end` == 1e-4 when `steps` > 15.
+ order: A `int`. The order of DPM-Solver.
+ skip_type: A `str`. The type for the spacing of the time steps. 'time_uniform' or 'logSNR' or 'time_quadratic'.
+ method: A `str`. The method for sampling. 'singlestep' or 'multistep' or 'singlestep_fixed' or 'adaptive'.
+ denoise_to_zero: A `bool`. Whether to denoise to time 0 at the final step.
+ Default is `False`. If `denoise_to_zero` is `True`, the total NFE is (`steps` + 1).
+ This trick is firstly proposed by DDPM (https://arxiv.org/abs/2006.11239) and
+ score_sde (https://arxiv.org/abs/2011.13456). Such trick can improve the FID
+ for diffusion models sampling by diffusion SDEs for low-resolutional images
+ (such as CIFAR-10). However, we observed that such trick does not matter for
+ high-resolutional images. As it needs an additional NFE, we do not recommend
+ it for high-resolutional images.
+ lower_order_final: A `bool`. Whether to use lower order solvers at the final steps.
+ Only valid for `method=multistep` and `steps < 15`. We empirically find that
+ this trick is a key to stabilizing the sampling by DPM-Solver with very few steps
+ (especially for steps <= 10). So we recommend to set it to be `True`.
+ solver_type: A `str`. The taylor expansion type for the solver. `dpm_solver` or `taylor`. We recommend `dpm_solver`.
+ atol: A `float`. The absolute tolerance of the adaptive step size solver. Valid when `method` == 'adaptive'.
+ rtol: A `float`. The relative tolerance of the adaptive step size solver. Valid when `method` == 'adaptive'.
+ Returns:
+ x_end: A pytorch tensor. The approximated solution at time `t_end`.
+ """
+ t_0 = 1. / self.noise_schedule.total_N if t_end is None else t_end
+ t_T = self.noise_schedule.T if t_start is None else t_start
+ device = x.device
+ if method == 'adaptive':
+ with torch.no_grad():
+ x = self.dpm_solver_adaptive(x, order=order, t_T=t_T, t_0=t_0, atol=atol, rtol=rtol,
+ solver_type=solver_type)
+ elif method == 'multistep':
+ assert steps >= order
+ timesteps = self.get_time_steps(skip_type=skip_type, t_T=t_T, t_0=t_0, N=steps, device=device)
+ assert timesteps.shape[0] - 1 == steps
+ with torch.no_grad():
+ vec_t = timesteps[0].expand((x.shape[0]))
+ model_prev_list = [self.model_fn(x, vec_t)]
+ t_prev_list = [vec_t]
+ # Init the first `order` values by lower order multistep DPM-Solver.
+ for init_order in tqdm(range(1, order), desc="DPM init order"):
+ vec_t = timesteps[init_order].expand(x.shape[0])
+ x = self.multistep_dpm_solver_update(x, model_prev_list, t_prev_list, vec_t, init_order,
+ solver_type=solver_type)
+ model_prev_list.append(self.model_fn(x, vec_t))
+ t_prev_list.append(vec_t)
+ # Compute the remaining values by `order`-th order multistep DPM-Solver.
+ for step in tqdm(range(order, steps + 1), desc="DPM multistep"):
+ vec_t = timesteps[step].expand(x.shape[0])
+ if lower_order_final and steps < 15:
+ step_order = min(order, steps + 1 - step)
+ else:
+ step_order = order
+ x = self.multistep_dpm_solver_update(x, model_prev_list, t_prev_list, vec_t, step_order,
+ solver_type=solver_type)
+ for i in range(order - 1):
+ t_prev_list[i] = t_prev_list[i + 1]
+ model_prev_list[i] = model_prev_list[i + 1]
+ t_prev_list[-1] = vec_t
+ # We do not need to evaluate the final model value.
+ if step < steps:
+ model_prev_list[-1] = self.model_fn(x, vec_t)
+ elif method in ['singlestep', 'singlestep_fixed']:
+ if method == 'singlestep':
+ timesteps_outer, orders = self.get_orders_and_timesteps_for_singlestep_solver(steps=steps, order=order,
+ skip_type=skip_type,
+ t_T=t_T, t_0=t_0,
+ device=device)
+ elif method == 'singlestep_fixed':
+ K = steps // order
+ orders = [order, ] * K
+ timesteps_outer = self.get_time_steps(skip_type=skip_type, t_T=t_T, t_0=t_0, N=K, device=device)
+ for i, order in enumerate(orders):
+ t_T_inner, t_0_inner = timesteps_outer[i], timesteps_outer[i + 1]
+ timesteps_inner = self.get_time_steps(skip_type=skip_type, t_T=t_T_inner.item(), t_0=t_0_inner.item(),
+ N=order, device=device)
+ lambda_inner = self.noise_schedule.marginal_lambda(timesteps_inner)
+ vec_s, vec_t = t_T_inner.tile(x.shape[0]), t_0_inner.tile(x.shape[0])
+ h = lambda_inner[-1] - lambda_inner[0]
+ r1 = None if order <= 1 else (lambda_inner[1] - lambda_inner[0]) / h
+ r2 = None if order <= 2 else (lambda_inner[2] - lambda_inner[0]) / h
+ x = self.singlestep_dpm_solver_update(x, vec_s, vec_t, order, solver_type=solver_type, r1=r1, r2=r2)
+ if denoise_to_zero:
+ x = self.denoise_to_zero_fn(x, torch.ones((x.shape[0],)).to(device) * t_0)
+ return x
+
+
+#############################################################
+# other utility functions
+#############################################################
+
+def interpolate_fn(x, xp, yp):
+ """
+ A piecewise linear function y = f(x), using xp and yp as keypoints.
+ We implement f(x) in a differentiable way (i.e. applicable for autograd).
+ The function f(x) is well-defined for all x-axis. (For x beyond the bounds of xp, we use the outmost points of xp to define the linear function.)
+ Args:
+ x: PyTorch tensor with shape [N, C], where N is the batch size, C is the number of channels (we use C = 1 for DPM-Solver).
+ xp: PyTorch tensor with shape [C, K], where K is the number of keypoints.
+ yp: PyTorch tensor with shape [C, K].
+ Returns:
+ The function values f(x), with shape [N, C].
+ """
+ N, K = x.shape[0], xp.shape[1]
+ all_x = torch.cat([x.unsqueeze(2), xp.unsqueeze(0).repeat((N, 1, 1))], dim=2)
+ sorted_all_x, x_indices = torch.sort(all_x, dim=2)
+ x_idx = torch.argmin(x_indices, dim=2)
+ cand_start_idx = x_idx - 1
+ start_idx = torch.where(
+ torch.eq(x_idx, 0),
+ torch.tensor(1, device=x.device),
+ torch.where(
+ torch.eq(x_idx, K), torch.tensor(K - 2, device=x.device), cand_start_idx,
+ ),
+ )
+ end_idx = torch.where(torch.eq(start_idx, cand_start_idx), start_idx + 2, start_idx + 1)
+ start_x = torch.gather(sorted_all_x, dim=2, index=start_idx.unsqueeze(2)).squeeze(2)
+ end_x = torch.gather(sorted_all_x, dim=2, index=end_idx.unsqueeze(2)).squeeze(2)
+ start_idx2 = torch.where(
+ torch.eq(x_idx, 0),
+ torch.tensor(0, device=x.device),
+ torch.where(
+ torch.eq(x_idx, K), torch.tensor(K - 2, device=x.device), cand_start_idx,
+ ),
+ )
+ y_positions_expanded = yp.unsqueeze(0).expand(N, -1, -1)
+ start_y = torch.gather(y_positions_expanded, dim=2, index=start_idx2.unsqueeze(2)).squeeze(2)
+ end_y = torch.gather(y_positions_expanded, dim=2, index=(start_idx2 + 1).unsqueeze(2)).squeeze(2)
+ cand = start_y + (x - start_x) * (end_y - start_y) / (end_x - start_x)
+ return cand
+
+
+def expand_dims(v, dims):
+ """
+ Expand the tensor `v` to the dim `dims`.
+ Args:
+ `v`: a PyTorch tensor with shape [N].
+ `dim`: a `int`.
+ Returns:
+ a PyTorch tensor with shape [N, 1, 1, ..., 1] and the total dimension is `dims`.
+ """
+ return v[(...,) + (None,) * (dims - 1)]
\ No newline at end of file
diff --git a/iopaint/model/anytext/ldm/modules/diffusionmodules/model.py b/iopaint/model/anytext/ldm/modules/diffusionmodules/model.py
new file mode 100644
index 0000000000000000000000000000000000000000..3472824236b8fb6d8e3d08001288603ee701a62f
--- /dev/null
+++ b/iopaint/model/anytext/ldm/modules/diffusionmodules/model.py
@@ -0,0 +1,973 @@
+# pytorch_diffusion + derived encoder decoder
+import math
+
+import numpy as np
+import torch
+import torch.nn as nn
+
+
+def get_timestep_embedding(timesteps, embedding_dim):
+ """
+ This matches the implementation in Denoising Diffusion Probabilistic Models:
+ From Fairseq.
+ Build sinusoidal embeddings.
+ This matches the implementation in tensor2tensor, but differs slightly
+ from the description in Section 3.5 of "Attention Is All You Need".
+ """
+ assert len(timesteps.shape) == 1
+
+ half_dim = embedding_dim // 2
+ emb = math.log(10000) / (half_dim - 1)
+ emb = torch.exp(torch.arange(half_dim, dtype=torch.float32) * -emb)
+ emb = emb.to(device=timesteps.device)
+ emb = timesteps.float()[:, None] * emb[None, :]
+ emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=1)
+ if embedding_dim % 2 == 1: # zero pad
+ emb = torch.nn.functional.pad(emb, (0, 1, 0, 0))
+ return emb
+
+
+def nonlinearity(x):
+ # swish
+ return x * torch.sigmoid(x)
+
+
+def Normalize(in_channels, num_groups=32):
+ return torch.nn.GroupNorm(
+ num_groups=num_groups, num_channels=in_channels, eps=1e-6, affine=True
+ )
+
+
+class Upsample(nn.Module):
+ def __init__(self, in_channels, with_conv):
+ super().__init__()
+ self.with_conv = with_conv
+ if self.with_conv:
+ self.conv = torch.nn.Conv2d(
+ in_channels, in_channels, kernel_size=3, stride=1, padding=1
+ )
+
+ def forward(self, x):
+ x = torch.nn.functional.interpolate(x, scale_factor=2.0, mode="nearest")
+ if self.with_conv:
+ x = self.conv(x)
+ return x
+
+
+class Downsample(nn.Module):
+ def __init__(self, in_channels, with_conv):
+ super().__init__()
+ self.with_conv = with_conv
+ if self.with_conv:
+ # no asymmetric padding in torch conv, must do it ourselves
+ self.conv = torch.nn.Conv2d(
+ in_channels, in_channels, kernel_size=3, stride=2, padding=0
+ )
+
+ def forward(self, x):
+ if self.with_conv:
+ pad = (0, 1, 0, 1)
+ x = torch.nn.functional.pad(x, pad, mode="constant", value=0)
+ x = self.conv(x)
+ else:
+ x = torch.nn.functional.avg_pool2d(x, kernel_size=2, stride=2)
+ return x
+
+
+class ResnetBlock(nn.Module):
+ def __init__(
+ self,
+ *,
+ in_channels,
+ out_channels=None,
+ conv_shortcut=False,
+ dropout,
+ temb_channels=512,
+ ):
+ super().__init__()
+ self.in_channels = in_channels
+ out_channels = in_channels if out_channels is None else out_channels
+ self.out_channels = out_channels
+ self.use_conv_shortcut = conv_shortcut
+
+ self.norm1 = Normalize(in_channels)
+ self.conv1 = torch.nn.Conv2d(
+ in_channels, out_channels, kernel_size=3, stride=1, padding=1
+ )
+ if temb_channels > 0:
+ self.temb_proj = torch.nn.Linear(temb_channels, out_channels)
+ self.norm2 = Normalize(out_channels)
+ self.dropout = torch.nn.Dropout(dropout)
+ self.conv2 = torch.nn.Conv2d(
+ out_channels, out_channels, kernel_size=3, stride=1, padding=1
+ )
+ if self.in_channels != self.out_channels:
+ if self.use_conv_shortcut:
+ self.conv_shortcut = torch.nn.Conv2d(
+ in_channels, out_channels, kernel_size=3, stride=1, padding=1
+ )
+ else:
+ self.nin_shortcut = torch.nn.Conv2d(
+ in_channels, out_channels, kernel_size=1, stride=1, padding=0
+ )
+
+ def forward(self, x, temb):
+ h = x
+ h = self.norm1(h)
+ h = nonlinearity(h)
+ h = self.conv1(h)
+
+ if temb is not None:
+ h = h + self.temb_proj(nonlinearity(temb))[:, :, None, None]
+
+ h = self.norm2(h)
+ h = nonlinearity(h)
+ h = self.dropout(h)
+ h = self.conv2(h)
+
+ if self.in_channels != self.out_channels:
+ if self.use_conv_shortcut:
+ x = self.conv_shortcut(x)
+ else:
+ x = self.nin_shortcut(x)
+
+ return x + h
+
+
+class AttnBlock(nn.Module):
+ def __init__(self, in_channels):
+ super().__init__()
+ self.in_channels = in_channels
+
+ self.norm = Normalize(in_channels)
+ self.q = torch.nn.Conv2d(
+ in_channels, in_channels, kernel_size=1, stride=1, padding=0
+ )
+ self.k = torch.nn.Conv2d(
+ in_channels, in_channels, kernel_size=1, stride=1, padding=0
+ )
+ self.v = torch.nn.Conv2d(
+ in_channels, in_channels, kernel_size=1, stride=1, padding=0
+ )
+ self.proj_out = torch.nn.Conv2d(
+ in_channels, in_channels, kernel_size=1, stride=1, padding=0
+ )
+
+ def forward(self, x):
+ h_ = x
+ h_ = self.norm(h_)
+ q = self.q(h_)
+ k = self.k(h_)
+ v = self.v(h_)
+
+ # compute attention
+ b, c, h, w = q.shape
+ q = q.reshape(b, c, h * w)
+ q = q.permute(0, 2, 1) # b,hw,c
+ k = k.reshape(b, c, h * w) # b,c,hw
+ w_ = torch.bmm(q, k) # b,hw,hw w[b,i,j]=sum_c q[b,i,c]k[b,c,j]
+ w_ = w_ * (int(c) ** (-0.5))
+ w_ = torch.nn.functional.softmax(w_, dim=2)
+
+ # attend to values
+ v = v.reshape(b, c, h * w)
+ w_ = w_.permute(0, 2, 1) # b,hw,hw (first hw of k, second of q)
+ h_ = torch.bmm(v, w_) # b, c,hw (hw of q) h_[b,c,j] = sum_i v[b,c,i] w_[b,i,j]
+ h_ = h_.reshape(b, c, h, w)
+
+ h_ = self.proj_out(h_)
+
+ return x + h_
+
+
+class AttnBlock2_0(nn.Module):
+ def __init__(self, in_channels):
+ super().__init__()
+ self.in_channels = in_channels
+
+ self.norm = Normalize(in_channels)
+ self.q = torch.nn.Conv2d(
+ in_channels, in_channels, kernel_size=1, stride=1, padding=0
+ )
+ self.k = torch.nn.Conv2d(
+ in_channels, in_channels, kernel_size=1, stride=1, padding=0
+ )
+ self.v = torch.nn.Conv2d(
+ in_channels, in_channels, kernel_size=1, stride=1, padding=0
+ )
+ self.proj_out = torch.nn.Conv2d(
+ in_channels, in_channels, kernel_size=1, stride=1, padding=0
+ )
+
+ def forward(self, x):
+ h_ = x
+ h_ = self.norm(h_)
+ # output: [1, 512, 64, 64]
+ q = self.q(h_)
+ k = self.k(h_)
+ v = self.v(h_)
+
+ # compute attention
+ b, c, h, w = q.shape
+
+ # q = q.reshape(b, c, h * w).transpose()
+ # q = q.permute(0, 2, 1) # b,hw,c
+ # k = k.reshape(b, c, h * w) # b,c,hw
+ q = q.transpose(1, 2)
+ k = k.transpose(1, 2)
+ v = v.transpose(1, 2)
+ # (batch, num_heads, seq_len, head_dim)
+ hidden_states = torch.nn.functional.scaled_dot_product_attention(
+ q, k, v, attn_mask=None, dropout_p=0.0, is_causal=False
+ )
+ hidden_states = hidden_states.transpose(1, 2)
+ hidden_states = hidden_states.to(q.dtype)
+
+ h_ = self.proj_out(hidden_states)
+
+ return x + h_
+
+
+def make_attn(in_channels, attn_type="vanilla", attn_kwargs=None):
+ assert attn_type in [
+ "vanilla",
+ "vanilla-xformers",
+ "memory-efficient-cross-attn",
+ "linear",
+ "none",
+ ], f"attn_type {attn_type} unknown"
+ assert attn_kwargs is None
+ if hasattr(torch.nn.functional, "scaled_dot_product_attention"):
+ # print(f"Using torch.nn.functional.scaled_dot_product_attention")
+ return AttnBlock2_0(in_channels)
+ return AttnBlock(in_channels)
+
+
+class Model(nn.Module):
+ def __init__(
+ self,
+ *,
+ ch,
+ out_ch,
+ ch_mult=(1, 2, 4, 8),
+ num_res_blocks,
+ attn_resolutions,
+ dropout=0.0,
+ resamp_with_conv=True,
+ in_channels,
+ resolution,
+ use_timestep=True,
+ use_linear_attn=False,
+ attn_type="vanilla",
+ ):
+ super().__init__()
+ if use_linear_attn:
+ attn_type = "linear"
+ self.ch = ch
+ self.temb_ch = self.ch * 4
+ self.num_resolutions = len(ch_mult)
+ self.num_res_blocks = num_res_blocks
+ self.resolution = resolution
+ self.in_channels = in_channels
+
+ self.use_timestep = use_timestep
+ if self.use_timestep:
+ # timestep embedding
+ self.temb = nn.Module()
+ self.temb.dense = nn.ModuleList(
+ [
+ torch.nn.Linear(self.ch, self.temb_ch),
+ torch.nn.Linear(self.temb_ch, self.temb_ch),
+ ]
+ )
+
+ # downsampling
+ self.conv_in = torch.nn.Conv2d(
+ in_channels, self.ch, kernel_size=3, stride=1, padding=1
+ )
+
+ curr_res = resolution
+ in_ch_mult = (1,) + tuple(ch_mult)
+ self.down = nn.ModuleList()
+ for i_level in range(self.num_resolutions):
+ block = nn.ModuleList()
+ attn = nn.ModuleList()
+ block_in = ch * in_ch_mult[i_level]
+ block_out = ch * ch_mult[i_level]
+ for i_block in range(self.num_res_blocks):
+ block.append(
+ ResnetBlock(
+ in_channels=block_in,
+ out_channels=block_out,
+ temb_channels=self.temb_ch,
+ dropout=dropout,
+ )
+ )
+ block_in = block_out
+ if curr_res in attn_resolutions:
+ attn.append(make_attn(block_in, attn_type=attn_type))
+ down = nn.Module()
+ down.block = block
+ down.attn = attn
+ if i_level != self.num_resolutions - 1:
+ down.downsample = Downsample(block_in, resamp_with_conv)
+ curr_res = curr_res // 2
+ self.down.append(down)
+
+ # middle
+ self.mid = nn.Module()
+ self.mid.block_1 = ResnetBlock(
+ in_channels=block_in,
+ out_channels=block_in,
+ temb_channels=self.temb_ch,
+ dropout=dropout,
+ )
+ self.mid.attn_1 = make_attn(block_in, attn_type=attn_type)
+ self.mid.block_2 = ResnetBlock(
+ in_channels=block_in,
+ out_channels=block_in,
+ temb_channels=self.temb_ch,
+ dropout=dropout,
+ )
+
+ # upsampling
+ self.up = nn.ModuleList()
+ for i_level in reversed(range(self.num_resolutions)):
+ block = nn.ModuleList()
+ attn = nn.ModuleList()
+ block_out = ch * ch_mult[i_level]
+ skip_in = ch * ch_mult[i_level]
+ for i_block in range(self.num_res_blocks + 1):
+ if i_block == self.num_res_blocks:
+ skip_in = ch * in_ch_mult[i_level]
+ block.append(
+ ResnetBlock(
+ in_channels=block_in + skip_in,
+ out_channels=block_out,
+ temb_channels=self.temb_ch,
+ dropout=dropout,
+ )
+ )
+ block_in = block_out
+ if curr_res in attn_resolutions:
+ attn.append(make_attn(block_in, attn_type=attn_type))
+ up = nn.Module()
+ up.block = block
+ up.attn = attn
+ if i_level != 0:
+ up.upsample = Upsample(block_in, resamp_with_conv)
+ curr_res = curr_res * 2
+ self.up.insert(0, up) # prepend to get consistent order
+
+ # end
+ self.norm_out = Normalize(block_in)
+ self.conv_out = torch.nn.Conv2d(
+ block_in, out_ch, kernel_size=3, stride=1, padding=1
+ )
+
+ def forward(self, x, t=None, context=None):
+ # assert x.shape[2] == x.shape[3] == self.resolution
+ if context is not None:
+ # assume aligned context, cat along channel axis
+ x = torch.cat((x, context), dim=1)
+ if self.use_timestep:
+ # timestep embedding
+ assert t is not None
+ temb = get_timestep_embedding(t, self.ch)
+ temb = self.temb.dense[0](temb)
+ temb = nonlinearity(temb)
+ temb = self.temb.dense[1](temb)
+ else:
+ temb = None
+
+ # downsampling
+ hs = [self.conv_in(x)]
+ for i_level in range(self.num_resolutions):
+ for i_block in range(self.num_res_blocks):
+ h = self.down[i_level].block[i_block](hs[-1], temb)
+ if len(self.down[i_level].attn) > 0:
+ h = self.down[i_level].attn[i_block](h)
+ hs.append(h)
+ if i_level != self.num_resolutions - 1:
+ hs.append(self.down[i_level].downsample(hs[-1]))
+
+ # middle
+ h = hs[-1]
+ h = self.mid.block_1(h, temb)
+ h = self.mid.attn_1(h)
+ h = self.mid.block_2(h, temb)
+
+ # upsampling
+ for i_level in reversed(range(self.num_resolutions)):
+ for i_block in range(self.num_res_blocks + 1):
+ h = self.up[i_level].block[i_block](
+ torch.cat([h, hs.pop()], dim=1), temb
+ )
+ if len(self.up[i_level].attn) > 0:
+ h = self.up[i_level].attn[i_block](h)
+ if i_level != 0:
+ h = self.up[i_level].upsample(h)
+
+ # end
+ h = self.norm_out(h)
+ h = nonlinearity(h)
+ h = self.conv_out(h)
+ return h
+
+ def get_last_layer(self):
+ return self.conv_out.weight
+
+
+class Encoder(nn.Module):
+ def __init__(
+ self,
+ *,
+ ch,
+ out_ch,
+ ch_mult=(1, 2, 4, 8),
+ num_res_blocks,
+ attn_resolutions,
+ dropout=0.0,
+ resamp_with_conv=True,
+ in_channels,
+ resolution,
+ z_channels,
+ double_z=True,
+ use_linear_attn=False,
+ attn_type="vanilla",
+ **ignore_kwargs,
+ ):
+ super().__init__()
+ if use_linear_attn:
+ attn_type = "linear"
+ self.ch = ch
+ self.temb_ch = 0
+ self.num_resolutions = len(ch_mult)
+ self.num_res_blocks = num_res_blocks
+ self.resolution = resolution
+ self.in_channels = in_channels
+
+ # downsampling
+ self.conv_in = torch.nn.Conv2d(
+ in_channels, self.ch, kernel_size=3, stride=1, padding=1
+ )
+
+ curr_res = resolution
+ in_ch_mult = (1,) + tuple(ch_mult)
+ self.in_ch_mult = in_ch_mult
+ self.down = nn.ModuleList()
+ for i_level in range(self.num_resolutions):
+ block = nn.ModuleList()
+ attn = nn.ModuleList()
+ block_in = ch * in_ch_mult[i_level]
+ block_out = ch * ch_mult[i_level]
+ for i_block in range(self.num_res_blocks):
+ block.append(
+ ResnetBlock(
+ in_channels=block_in,
+ out_channels=block_out,
+ temb_channels=self.temb_ch,
+ dropout=dropout,
+ )
+ )
+ block_in = block_out
+ if curr_res in attn_resolutions:
+ attn.append(make_attn(block_in, attn_type=attn_type))
+ down = nn.Module()
+ down.block = block
+ down.attn = attn
+ if i_level != self.num_resolutions - 1:
+ down.downsample = Downsample(block_in, resamp_with_conv)
+ curr_res = curr_res // 2
+ self.down.append(down)
+
+ # middle
+ self.mid = nn.Module()
+ self.mid.block_1 = ResnetBlock(
+ in_channels=block_in,
+ out_channels=block_in,
+ temb_channels=self.temb_ch,
+ dropout=dropout,
+ )
+ self.mid.attn_1 = make_attn(block_in, attn_type=attn_type)
+ self.mid.block_2 = ResnetBlock(
+ in_channels=block_in,
+ out_channels=block_in,
+ temb_channels=self.temb_ch,
+ dropout=dropout,
+ )
+
+ # end
+ self.norm_out = Normalize(block_in)
+ self.conv_out = torch.nn.Conv2d(
+ block_in,
+ 2 * z_channels if double_z else z_channels,
+ kernel_size=3,
+ stride=1,
+ padding=1,
+ )
+
+ def forward(self, x):
+ # timestep embedding
+ temb = None
+
+ # downsampling
+ hs = [self.conv_in(x)]
+ for i_level in range(self.num_resolutions):
+ for i_block in range(self.num_res_blocks):
+ h = self.down[i_level].block[i_block](hs[-1], temb)
+ if len(self.down[i_level].attn) > 0:
+ h = self.down[i_level].attn[i_block](h)
+ hs.append(h)
+ if i_level != self.num_resolutions - 1:
+ hs.append(self.down[i_level].downsample(hs[-1]))
+
+ # middle
+ h = hs[-1]
+ h = self.mid.block_1(h, temb)
+ h = self.mid.attn_1(h)
+ h = self.mid.block_2(h, temb)
+
+ # end
+ h = self.norm_out(h)
+ h = nonlinearity(h)
+ h = self.conv_out(h)
+ return h
+
+
+class Decoder(nn.Module):
+ def __init__(
+ self,
+ *,
+ ch,
+ out_ch,
+ ch_mult=(1, 2, 4, 8),
+ num_res_blocks,
+ attn_resolutions,
+ dropout=0.0,
+ resamp_with_conv=True,
+ in_channels,
+ resolution,
+ z_channels,
+ give_pre_end=False,
+ tanh_out=False,
+ use_linear_attn=False,
+ attn_type="vanilla",
+ **ignorekwargs,
+ ):
+ super().__init__()
+ if use_linear_attn:
+ attn_type = "linear"
+ self.ch = ch
+ self.temb_ch = 0
+ self.num_resolutions = len(ch_mult)
+ self.num_res_blocks = num_res_blocks
+ self.resolution = resolution
+ self.in_channels = in_channels
+ self.give_pre_end = give_pre_end
+ self.tanh_out = tanh_out
+
+ # compute in_ch_mult, block_in and curr_res at lowest res
+ in_ch_mult = (1,) + tuple(ch_mult)
+ block_in = ch * ch_mult[self.num_resolutions - 1]
+ curr_res = resolution // 2 ** (self.num_resolutions - 1)
+ self.z_shape = (1, z_channels, curr_res, curr_res)
+ print(
+ "Working with z of shape {} = {} dimensions.".format(
+ self.z_shape, np.prod(self.z_shape)
+ )
+ )
+
+ # z to block_in
+ self.conv_in = torch.nn.Conv2d(
+ z_channels, block_in, kernel_size=3, stride=1, padding=1
+ )
+
+ # middle
+ self.mid = nn.Module()
+ self.mid.block_1 = ResnetBlock(
+ in_channels=block_in,
+ out_channels=block_in,
+ temb_channels=self.temb_ch,
+ dropout=dropout,
+ )
+ self.mid.attn_1 = make_attn(block_in, attn_type=attn_type)
+ self.mid.block_2 = ResnetBlock(
+ in_channels=block_in,
+ out_channels=block_in,
+ temb_channels=self.temb_ch,
+ dropout=dropout,
+ )
+
+ # upsampling
+ self.up = nn.ModuleList()
+ for i_level in reversed(range(self.num_resolutions)):
+ block = nn.ModuleList()
+ attn = nn.ModuleList()
+ block_out = ch * ch_mult[i_level]
+ for i_block in range(self.num_res_blocks + 1):
+ block.append(
+ ResnetBlock(
+ in_channels=block_in,
+ out_channels=block_out,
+ temb_channels=self.temb_ch,
+ dropout=dropout,
+ )
+ )
+ block_in = block_out
+ if curr_res in attn_resolutions:
+ attn.append(make_attn(block_in, attn_type=attn_type))
+ up = nn.Module()
+ up.block = block
+ up.attn = attn
+ if i_level != 0:
+ up.upsample = Upsample(block_in, resamp_with_conv)
+ curr_res = curr_res * 2
+ self.up.insert(0, up) # prepend to get consistent order
+
+ # end
+ self.norm_out = Normalize(block_in)
+ self.conv_out = torch.nn.Conv2d(
+ block_in, out_ch, kernel_size=3, stride=1, padding=1
+ )
+
+ def forward(self, z):
+ # assert z.shape[1:] == self.z_shape[1:]
+ self.last_z_shape = z.shape
+
+ # timestep embedding
+ temb = None
+
+ # z to block_in
+ h = self.conv_in(z)
+
+ # middle
+ h = self.mid.block_1(h, temb)
+ h = self.mid.attn_1(h)
+ h = self.mid.block_2(h, temb)
+
+ # upsampling
+ for i_level in reversed(range(self.num_resolutions)):
+ for i_block in range(self.num_res_blocks + 1):
+ h = self.up[i_level].block[i_block](h, temb)
+ if len(self.up[i_level].attn) > 0:
+ h = self.up[i_level].attn[i_block](h)
+ if i_level != 0:
+ h = self.up[i_level].upsample(h)
+
+ # end
+ if self.give_pre_end:
+ return h
+
+ h = self.norm_out(h)
+ h = nonlinearity(h)
+ h = self.conv_out(h)
+ if self.tanh_out:
+ h = torch.tanh(h)
+ return h
+
+
+class SimpleDecoder(nn.Module):
+ def __init__(self, in_channels, out_channels, *args, **kwargs):
+ super().__init__()
+ self.model = nn.ModuleList(
+ [
+ nn.Conv2d(in_channels, in_channels, 1),
+ ResnetBlock(
+ in_channels=in_channels,
+ out_channels=2 * in_channels,
+ temb_channels=0,
+ dropout=0.0,
+ ),
+ ResnetBlock(
+ in_channels=2 * in_channels,
+ out_channels=4 * in_channels,
+ temb_channels=0,
+ dropout=0.0,
+ ),
+ ResnetBlock(
+ in_channels=4 * in_channels,
+ out_channels=2 * in_channels,
+ temb_channels=0,
+ dropout=0.0,
+ ),
+ nn.Conv2d(2 * in_channels, in_channels, 1),
+ Upsample(in_channels, with_conv=True),
+ ]
+ )
+ # end
+ self.norm_out = Normalize(in_channels)
+ self.conv_out = torch.nn.Conv2d(
+ in_channels, out_channels, kernel_size=3, stride=1, padding=1
+ )
+
+ def forward(self, x):
+ for i, layer in enumerate(self.model):
+ if i in [1, 2, 3]:
+ x = layer(x, None)
+ else:
+ x = layer(x)
+
+ h = self.norm_out(x)
+ h = nonlinearity(h)
+ x = self.conv_out(h)
+ return x
+
+
+class UpsampleDecoder(nn.Module):
+ def __init__(
+ self,
+ in_channels,
+ out_channels,
+ ch,
+ num_res_blocks,
+ resolution,
+ ch_mult=(2, 2),
+ dropout=0.0,
+ ):
+ super().__init__()
+ # upsampling
+ self.temb_ch = 0
+ self.num_resolutions = len(ch_mult)
+ self.num_res_blocks = num_res_blocks
+ block_in = in_channels
+ curr_res = resolution // 2 ** (self.num_resolutions - 1)
+ self.res_blocks = nn.ModuleList()
+ self.upsample_blocks = nn.ModuleList()
+ for i_level in range(self.num_resolutions):
+ res_block = []
+ block_out = ch * ch_mult[i_level]
+ for i_block in range(self.num_res_blocks + 1):
+ res_block.append(
+ ResnetBlock(
+ in_channels=block_in,
+ out_channels=block_out,
+ temb_channels=self.temb_ch,
+ dropout=dropout,
+ )
+ )
+ block_in = block_out
+ self.res_blocks.append(nn.ModuleList(res_block))
+ if i_level != self.num_resolutions - 1:
+ self.upsample_blocks.append(Upsample(block_in, True))
+ curr_res = curr_res * 2
+
+ # end
+ self.norm_out = Normalize(block_in)
+ self.conv_out = torch.nn.Conv2d(
+ block_in, out_channels, kernel_size=3, stride=1, padding=1
+ )
+
+ def forward(self, x):
+ # upsampling
+ h = x
+ for k, i_level in enumerate(range(self.num_resolutions)):
+ for i_block in range(self.num_res_blocks + 1):
+ h = self.res_blocks[i_level][i_block](h, None)
+ if i_level != self.num_resolutions - 1:
+ h = self.upsample_blocks[k](h)
+ h = self.norm_out(h)
+ h = nonlinearity(h)
+ h = self.conv_out(h)
+ return h
+
+
+class LatentRescaler(nn.Module):
+ def __init__(self, factor, in_channels, mid_channels, out_channels, depth=2):
+ super().__init__()
+ # residual block, interpolate, residual block
+ self.factor = factor
+ self.conv_in = nn.Conv2d(
+ in_channels, mid_channels, kernel_size=3, stride=1, padding=1
+ )
+ self.res_block1 = nn.ModuleList(
+ [
+ ResnetBlock(
+ in_channels=mid_channels,
+ out_channels=mid_channels,
+ temb_channels=0,
+ dropout=0.0,
+ )
+ for _ in range(depth)
+ ]
+ )
+ self.attn = AttnBlock(mid_channels)
+ self.res_block2 = nn.ModuleList(
+ [
+ ResnetBlock(
+ in_channels=mid_channels,
+ out_channels=mid_channels,
+ temb_channels=0,
+ dropout=0.0,
+ )
+ for _ in range(depth)
+ ]
+ )
+
+ self.conv_out = nn.Conv2d(
+ mid_channels,
+ out_channels,
+ kernel_size=1,
+ )
+
+ def forward(self, x):
+ x = self.conv_in(x)
+ for block in self.res_block1:
+ x = block(x, None)
+ x = torch.nn.functional.interpolate(
+ x,
+ size=(
+ int(round(x.shape[2] * self.factor)),
+ int(round(x.shape[3] * self.factor)),
+ ),
+ )
+ x = self.attn(x)
+ for block in self.res_block2:
+ x = block(x, None)
+ x = self.conv_out(x)
+ return x
+
+
+class MergedRescaleEncoder(nn.Module):
+ def __init__(
+ self,
+ in_channels,
+ ch,
+ resolution,
+ out_ch,
+ num_res_blocks,
+ attn_resolutions,
+ dropout=0.0,
+ resamp_with_conv=True,
+ ch_mult=(1, 2, 4, 8),
+ rescale_factor=1.0,
+ rescale_module_depth=1,
+ ):
+ super().__init__()
+ intermediate_chn = ch * ch_mult[-1]
+ self.encoder = Encoder(
+ in_channels=in_channels,
+ num_res_blocks=num_res_blocks,
+ ch=ch,
+ ch_mult=ch_mult,
+ z_channels=intermediate_chn,
+ double_z=False,
+ resolution=resolution,
+ attn_resolutions=attn_resolutions,
+ dropout=dropout,
+ resamp_with_conv=resamp_with_conv,
+ out_ch=None,
+ )
+ self.rescaler = LatentRescaler(
+ factor=rescale_factor,
+ in_channels=intermediate_chn,
+ mid_channels=intermediate_chn,
+ out_channels=out_ch,
+ depth=rescale_module_depth,
+ )
+
+ def forward(self, x):
+ x = self.encoder(x)
+ x = self.rescaler(x)
+ return x
+
+
+class MergedRescaleDecoder(nn.Module):
+ def __init__(
+ self,
+ z_channels,
+ out_ch,
+ resolution,
+ num_res_blocks,
+ attn_resolutions,
+ ch,
+ ch_mult=(1, 2, 4, 8),
+ dropout=0.0,
+ resamp_with_conv=True,
+ rescale_factor=1.0,
+ rescale_module_depth=1,
+ ):
+ super().__init__()
+ tmp_chn = z_channels * ch_mult[-1]
+ self.decoder = Decoder(
+ out_ch=out_ch,
+ z_channels=tmp_chn,
+ attn_resolutions=attn_resolutions,
+ dropout=dropout,
+ resamp_with_conv=resamp_with_conv,
+ in_channels=None,
+ num_res_blocks=num_res_blocks,
+ ch_mult=ch_mult,
+ resolution=resolution,
+ ch=ch,
+ )
+ self.rescaler = LatentRescaler(
+ factor=rescale_factor,
+ in_channels=z_channels,
+ mid_channels=tmp_chn,
+ out_channels=tmp_chn,
+ depth=rescale_module_depth,
+ )
+
+ def forward(self, x):
+ x = self.rescaler(x)
+ x = self.decoder(x)
+ return x
+
+
+class Upsampler(nn.Module):
+ def __init__(self, in_size, out_size, in_channels, out_channels, ch_mult=2):
+ super().__init__()
+ assert out_size >= in_size
+ num_blocks = int(np.log2(out_size // in_size)) + 1
+ factor_up = 1.0 + (out_size % in_size)
+ print(
+ f"Building {self.__class__.__name__} with in_size: {in_size} --> out_size {out_size} and factor {factor_up}"
+ )
+ self.rescaler = LatentRescaler(
+ factor=factor_up,
+ in_channels=in_channels,
+ mid_channels=2 * in_channels,
+ out_channels=in_channels,
+ )
+ self.decoder = Decoder(
+ out_ch=out_channels,
+ resolution=out_size,
+ z_channels=in_channels,
+ num_res_blocks=2,
+ attn_resolutions=[],
+ in_channels=None,
+ ch=in_channels,
+ ch_mult=[ch_mult for _ in range(num_blocks)],
+ )
+
+ def forward(self, x):
+ x = self.rescaler(x)
+ x = self.decoder(x)
+ return x
+
+
+class Resize(nn.Module):
+ def __init__(self, in_channels=None, learned=False, mode="bilinear"):
+ super().__init__()
+ self.with_conv = learned
+ self.mode = mode
+ if self.with_conv:
+ print(
+ f"Note: {self.__class__.__name} uses learned downsampling and will ignore the fixed {mode} mode"
+ )
+ raise NotImplementedError()
+ assert in_channels is not None
+ # no asymmetric padding in torch conv, must do it ourselves
+ self.conv = torch.nn.Conv2d(
+ in_channels, in_channels, kernel_size=4, stride=2, padding=1
+ )
+
+ def forward(self, x, scale_factor=1.0):
+ if scale_factor == 1.0:
+ return x
+ else:
+ x = torch.nn.functional.interpolate(
+ x, mode=self.mode, align_corners=False, scale_factor=scale_factor
+ )
+ return x
diff --git a/iopaint/model/anytext/ldm/modules/diffusionmodules/openaimodel.py b/iopaint/model/anytext/ldm/modules/diffusionmodules/openaimodel.py
new file mode 100644
index 0000000000000000000000000000000000000000..fd3d6bed17710818aed0c3c77b5c2e85c555a82e
--- /dev/null
+++ b/iopaint/model/anytext/ldm/modules/diffusionmodules/openaimodel.py
@@ -0,0 +1,786 @@
+from abc import abstractmethod
+import math
+
+import numpy as np
+import torch as th
+import torch.nn as nn
+import torch.nn.functional as F
+
+from iopaint.model.anytext.ldm.modules.diffusionmodules.util import (
+ checkpoint,
+ conv_nd,
+ linear,
+ avg_pool_nd,
+ zero_module,
+ normalization,
+ timestep_embedding,
+)
+from iopaint.model.anytext.ldm.modules.attention import SpatialTransformer
+from iopaint.model.anytext.ldm.util import exists
+
+
+# dummy replace
+def convert_module_to_f16(x):
+ pass
+
+def convert_module_to_f32(x):
+ pass
+
+
+## go
+class AttentionPool2d(nn.Module):
+ """
+ Adapted from CLIP: https://github.com/openai/CLIP/blob/main/clip/model.py
+ """
+
+ def __init__(
+ self,
+ spacial_dim: int,
+ embed_dim: int,
+ num_heads_channels: int,
+ output_dim: int = None,
+ ):
+ super().__init__()
+ self.positional_embedding = nn.Parameter(th.randn(embed_dim, spacial_dim ** 2 + 1) / embed_dim ** 0.5)
+ self.qkv_proj = conv_nd(1, embed_dim, 3 * embed_dim, 1)
+ self.c_proj = conv_nd(1, embed_dim, output_dim or embed_dim, 1)
+ self.num_heads = embed_dim // num_heads_channels
+ self.attention = QKVAttention(self.num_heads)
+
+ def forward(self, x):
+ b, c, *_spatial = x.shape
+ x = x.reshape(b, c, -1) # NC(HW)
+ x = th.cat([x.mean(dim=-1, keepdim=True), x], dim=-1) # NC(HW+1)
+ x = x + self.positional_embedding[None, :, :].to(x.dtype) # NC(HW+1)
+ x = self.qkv_proj(x)
+ x = self.attention(x)
+ x = self.c_proj(x)
+ return x[:, :, 0]
+
+
+class TimestepBlock(nn.Module):
+ """
+ Any module where forward() takes timestep embeddings as a second argument.
+ """
+
+ @abstractmethod
+ def forward(self, x, emb):
+ """
+ Apply the module to `x` given `emb` timestep embeddings.
+ """
+
+
+class TimestepEmbedSequential(nn.Sequential, TimestepBlock):
+ """
+ A sequential module that passes timestep embeddings to the children that
+ support it as an extra input.
+ """
+
+ def forward(self, x, emb, context=None):
+ for layer in self:
+ if isinstance(layer, TimestepBlock):
+ x = layer(x, emb)
+ elif isinstance(layer, SpatialTransformer):
+ x = layer(x, context)
+ else:
+ x = layer(x)
+ return x
+
+
+class Upsample(nn.Module):
+ """
+ An upsampling layer with an optional convolution.
+ :param channels: channels in the inputs and outputs.
+ :param use_conv: a bool determining if a convolution is applied.
+ :param dims: determines if the signal is 1D, 2D, or 3D. If 3D, then
+ upsampling occurs in the inner-two dimensions.
+ """
+
+ def __init__(self, channels, use_conv, dims=2, out_channels=None, padding=1):
+ super().__init__()
+ self.channels = channels
+ self.out_channels = out_channels or channels
+ self.use_conv = use_conv
+ self.dims = dims
+ if use_conv:
+ self.conv = conv_nd(dims, self.channels, self.out_channels, 3, padding=padding)
+
+ def forward(self, x):
+ assert x.shape[1] == self.channels
+ if self.dims == 3:
+ x = F.interpolate(
+ x, (x.shape[2], x.shape[3] * 2, x.shape[4] * 2), mode="nearest"
+ )
+ else:
+ x = F.interpolate(x, scale_factor=2, mode="nearest")
+ if self.use_conv:
+ x = self.conv(x)
+ return x
+
+class TransposedUpsample(nn.Module):
+ 'Learned 2x upsampling without padding'
+ def __init__(self, channels, out_channels=None, ks=5):
+ super().__init__()
+ self.channels = channels
+ self.out_channels = out_channels or channels
+
+ self.up = nn.ConvTranspose2d(self.channels,self.out_channels,kernel_size=ks,stride=2)
+
+ def forward(self,x):
+ return self.up(x)
+
+
+class Downsample(nn.Module):
+ """
+ A downsampling layer with an optional convolution.
+ :param channels: channels in the inputs and outputs.
+ :param use_conv: a bool determining if a convolution is applied.
+ :param dims: determines if the signal is 1D, 2D, or 3D. If 3D, then
+ downsampling occurs in the inner-two dimensions.
+ """
+
+ def __init__(self, channels, use_conv, dims=2, out_channels=None,padding=1):
+ super().__init__()
+ self.channels = channels
+ self.out_channels = out_channels or channels
+ self.use_conv = use_conv
+ self.dims = dims
+ stride = 2 if dims != 3 else (1, 2, 2)
+ if use_conv:
+ self.op = conv_nd(
+ dims, self.channels, self.out_channels, 3, stride=stride, padding=padding
+ )
+ else:
+ assert self.channels == self.out_channels
+ self.op = avg_pool_nd(dims, kernel_size=stride, stride=stride)
+
+ def forward(self, x):
+ assert x.shape[1] == self.channels
+ return self.op(x)
+
+
+class ResBlock(TimestepBlock):
+ """
+ A residual block that can optionally change the number of channels.
+ :param channels: the number of input channels.
+ :param emb_channels: the number of timestep embedding channels.
+ :param dropout: the rate of dropout.
+ :param out_channels: if specified, the number of out channels.
+ :param use_conv: if True and out_channels is specified, use a spatial
+ convolution instead of a smaller 1x1 convolution to change the
+ channels in the skip connection.
+ :param dims: determines if the signal is 1D, 2D, or 3D.
+ :param use_checkpoint: if True, use gradient checkpointing on this module.
+ :param up: if True, use this block for upsampling.
+ :param down: if True, use this block for downsampling.
+ """
+
+ def __init__(
+ self,
+ channels,
+ emb_channels,
+ dropout,
+ out_channels=None,
+ use_conv=False,
+ use_scale_shift_norm=False,
+ dims=2,
+ use_checkpoint=False,
+ up=False,
+ down=False,
+ ):
+ super().__init__()
+ self.channels = channels
+ self.emb_channels = emb_channels
+ self.dropout = dropout
+ self.out_channels = out_channels or channels
+ self.use_conv = use_conv
+ self.use_checkpoint = use_checkpoint
+ self.use_scale_shift_norm = use_scale_shift_norm
+
+ self.in_layers = nn.Sequential(
+ normalization(channels),
+ nn.SiLU(),
+ conv_nd(dims, channels, self.out_channels, 3, padding=1),
+ )
+
+ self.updown = up or down
+
+ if up:
+ self.h_upd = Upsample(channels, False, dims)
+ self.x_upd = Upsample(channels, False, dims)
+ elif down:
+ self.h_upd = Downsample(channels, False, dims)
+ self.x_upd = Downsample(channels, False, dims)
+ else:
+ self.h_upd = self.x_upd = nn.Identity()
+
+ self.emb_layers = nn.Sequential(
+ nn.SiLU(),
+ linear(
+ emb_channels,
+ 2 * self.out_channels if use_scale_shift_norm else self.out_channels,
+ ),
+ )
+ self.out_layers = nn.Sequential(
+ normalization(self.out_channels),
+ nn.SiLU(),
+ nn.Dropout(p=dropout),
+ zero_module(
+ conv_nd(dims, self.out_channels, self.out_channels, 3, padding=1)
+ ),
+ )
+
+ if self.out_channels == channels:
+ self.skip_connection = nn.Identity()
+ elif use_conv:
+ self.skip_connection = conv_nd(
+ dims, channels, self.out_channels, 3, padding=1
+ )
+ else:
+ self.skip_connection = conv_nd(dims, channels, self.out_channels, 1)
+
+ def forward(self, x, emb):
+ """
+ Apply the block to a Tensor, conditioned on a timestep embedding.
+ :param x: an [N x C x ...] Tensor of features.
+ :param emb: an [N x emb_channels] Tensor of timestep embeddings.
+ :return: an [N x C x ...] Tensor of outputs.
+ """
+ return checkpoint(
+ self._forward, (x, emb), self.parameters(), self.use_checkpoint
+ )
+
+
+ def _forward(self, x, emb):
+ if self.updown:
+ in_rest, in_conv = self.in_layers[:-1], self.in_layers[-1]
+ h = in_rest(x)
+ h = self.h_upd(h)
+ x = self.x_upd(x)
+ h = in_conv(h)
+ else:
+ h = self.in_layers(x)
+ emb_out = self.emb_layers(emb).type(h.dtype)
+ while len(emb_out.shape) < len(h.shape):
+ emb_out = emb_out[..., None]
+ if self.use_scale_shift_norm:
+ out_norm, out_rest = self.out_layers[0], self.out_layers[1:]
+ scale, shift = th.chunk(emb_out, 2, dim=1)
+ h = out_norm(h) * (1 + scale) + shift
+ h = out_rest(h)
+ else:
+ h = h + emb_out
+ h = self.out_layers(h)
+ return self.skip_connection(x) + h
+
+
+class AttentionBlock(nn.Module):
+ """
+ An attention block that allows spatial positions to attend to each other.
+ Originally ported from here, but adapted to the N-d case.
+ https://github.com/hojonathanho/diffusion/blob/1e0dceb3b3495bbe19116a5e1b3596cd0706c543/diffusion_tf/models/unet.py#L66.
+ """
+
+ def __init__(
+ self,
+ channels,
+ num_heads=1,
+ num_head_channels=-1,
+ use_checkpoint=False,
+ use_new_attention_order=False,
+ ):
+ super().__init__()
+ self.channels = channels
+ if num_head_channels == -1:
+ self.num_heads = num_heads
+ else:
+ assert (
+ channels % num_head_channels == 0
+ ), f"q,k,v channels {channels} is not divisible by num_head_channels {num_head_channels}"
+ self.num_heads = channels // num_head_channels
+ self.use_checkpoint = use_checkpoint
+ self.norm = normalization(channels)
+ self.qkv = conv_nd(1, channels, channels * 3, 1)
+ if use_new_attention_order:
+ # split qkv before split heads
+ self.attention = QKVAttention(self.num_heads)
+ else:
+ # split heads before split qkv
+ self.attention = QKVAttentionLegacy(self.num_heads)
+
+ self.proj_out = zero_module(conv_nd(1, channels, channels, 1))
+
+ def forward(self, x):
+ return checkpoint(self._forward, (x,), self.parameters(), True) # TODO: check checkpoint usage, is True # TODO: fix the .half call!!!
+ #return pt_checkpoint(self._forward, x) # pytorch
+
+ def _forward(self, x):
+ b, c, *spatial = x.shape
+ x = x.reshape(b, c, -1)
+ qkv = self.qkv(self.norm(x))
+ h = self.attention(qkv)
+ h = self.proj_out(h)
+ return (x + h).reshape(b, c, *spatial)
+
+
+def count_flops_attn(model, _x, y):
+ """
+ A counter for the `thop` package to count the operations in an
+ attention operation.
+ Meant to be used like:
+ macs, params = thop.profile(
+ model,
+ inputs=(inputs, timestamps),
+ custom_ops={QKVAttention: QKVAttention.count_flops},
+ )
+ """
+ b, c, *spatial = y[0].shape
+ num_spatial = int(np.prod(spatial))
+ # We perform two matmuls with the same number of ops.
+ # The first computes the weight matrix, the second computes
+ # the combination of the value vectors.
+ matmul_ops = 2 * b * (num_spatial ** 2) * c
+ model.total_ops += th.DoubleTensor([matmul_ops])
+
+
+class QKVAttentionLegacy(nn.Module):
+ """
+ A module which performs QKV attention. Matches legacy QKVAttention + input/ouput heads shaping
+ """
+
+ def __init__(self, n_heads):
+ super().__init__()
+ self.n_heads = n_heads
+
+ def forward(self, qkv):
+ """
+ Apply QKV attention.
+ :param qkv: an [N x (H * 3 * C) x T] tensor of Qs, Ks, and Vs.
+ :return: an [N x (H * C) x T] tensor after attention.
+ """
+ bs, width, length = qkv.shape
+ assert width % (3 * self.n_heads) == 0
+ ch = width // (3 * self.n_heads)
+ q, k, v = qkv.reshape(bs * self.n_heads, ch * 3, length).split(ch, dim=1)
+ scale = 1 / math.sqrt(math.sqrt(ch))
+ weight = th.einsum(
+ "bct,bcs->bts", q * scale, k * scale
+ ) # More stable with f16 than dividing afterwards
+ weight = th.softmax(weight.float(), dim=-1).type(weight.dtype)
+ a = th.einsum("bts,bcs->bct", weight, v)
+ return a.reshape(bs, -1, length)
+
+ @staticmethod
+ def count_flops(model, _x, y):
+ return count_flops_attn(model, _x, y)
+
+
+class QKVAttention(nn.Module):
+ """
+ A module which performs QKV attention and splits in a different order.
+ """
+
+ def __init__(self, n_heads):
+ super().__init__()
+ self.n_heads = n_heads
+
+ def forward(self, qkv):
+ """
+ Apply QKV attention.
+ :param qkv: an [N x (3 * H * C) x T] tensor of Qs, Ks, and Vs.
+ :return: an [N x (H * C) x T] tensor after attention.
+ """
+ bs, width, length = qkv.shape
+ assert width % (3 * self.n_heads) == 0
+ ch = width // (3 * self.n_heads)
+ q, k, v = qkv.chunk(3, dim=1)
+ scale = 1 / math.sqrt(math.sqrt(ch))
+ weight = th.einsum(
+ "bct,bcs->bts",
+ (q * scale).view(bs * self.n_heads, ch, length),
+ (k * scale).view(bs * self.n_heads, ch, length),
+ ) # More stable with f16 than dividing afterwards
+ weight = th.softmax(weight.float(), dim=-1).type(weight.dtype)
+ a = th.einsum("bts,bcs->bct", weight, v.reshape(bs * self.n_heads, ch, length))
+ return a.reshape(bs, -1, length)
+
+ @staticmethod
+ def count_flops(model, _x, y):
+ return count_flops_attn(model, _x, y)
+
+
+class UNetModel(nn.Module):
+ """
+ The full UNet model with attention and timestep embedding.
+ :param in_channels: channels in the input Tensor.
+ :param model_channels: base channel count for the model.
+ :param out_channels: channels in the output Tensor.
+ :param num_res_blocks: number of residual blocks per downsample.
+ :param attention_resolutions: a collection of downsample rates at which
+ attention will take place. May be a set, list, or tuple.
+ For example, if this contains 4, then at 4x downsampling, attention
+ will be used.
+ :param dropout: the dropout probability.
+ :param channel_mult: channel multiplier for each level of the UNet.
+ :param conv_resample: if True, use learned convolutions for upsampling and
+ downsampling.
+ :param dims: determines if the signal is 1D, 2D, or 3D.
+ :param num_classes: if specified (as an int), then this model will be
+ class-conditional with `num_classes` classes.
+ :param use_checkpoint: use gradient checkpointing to reduce memory usage.
+ :param num_heads: the number of attention heads in each attention layer.
+ :param num_heads_channels: if specified, ignore num_heads and instead use
+ a fixed channel width per attention head.
+ :param num_heads_upsample: works with num_heads to set a different number
+ of heads for upsampling. Deprecated.
+ :param use_scale_shift_norm: use a FiLM-like conditioning mechanism.
+ :param resblock_updown: use residual blocks for up/downsampling.
+ :param use_new_attention_order: use a different attention pattern for potentially
+ increased efficiency.
+ """
+
+ def __init__(
+ self,
+ image_size,
+ in_channels,
+ model_channels,
+ out_channels,
+ num_res_blocks,
+ attention_resolutions,
+ dropout=0,
+ channel_mult=(1, 2, 4, 8),
+ conv_resample=True,
+ dims=2,
+ num_classes=None,
+ use_checkpoint=False,
+ use_fp16=False,
+ num_heads=-1,
+ num_head_channels=-1,
+ num_heads_upsample=-1,
+ use_scale_shift_norm=False,
+ resblock_updown=False,
+ use_new_attention_order=False,
+ use_spatial_transformer=False, # custom transformer support
+ transformer_depth=1, # custom transformer support
+ context_dim=None, # custom transformer support
+ n_embed=None, # custom support for prediction of discrete ids into codebook of first stage vq model
+ legacy=True,
+ disable_self_attentions=None,
+ num_attention_blocks=None,
+ disable_middle_self_attn=False,
+ use_linear_in_transformer=False,
+ ):
+ super().__init__()
+ if use_spatial_transformer:
+ assert context_dim is not None, 'Fool!! You forgot to include the dimension of your cross-attention conditioning...'
+
+ if context_dim is not None:
+ assert use_spatial_transformer, 'Fool!! You forgot to use the spatial transformer for your cross-attention conditioning...'
+ from omegaconf.listconfig import ListConfig
+ if type(context_dim) == ListConfig:
+ context_dim = list(context_dim)
+
+ if num_heads_upsample == -1:
+ num_heads_upsample = num_heads
+
+ if num_heads == -1:
+ assert num_head_channels != -1, 'Either num_heads or num_head_channels has to be set'
+
+ if num_head_channels == -1:
+ assert num_heads != -1, 'Either num_heads or num_head_channels has to be set'
+
+ self.image_size = image_size
+ self.in_channels = in_channels
+ self.model_channels = model_channels
+ self.out_channels = out_channels
+ if isinstance(num_res_blocks, int):
+ self.num_res_blocks = len(channel_mult) * [num_res_blocks]
+ else:
+ if len(num_res_blocks) != len(channel_mult):
+ raise ValueError("provide num_res_blocks either as an int (globally constant) or "
+ "as a list/tuple (per-level) with the same length as channel_mult")
+ self.num_res_blocks = num_res_blocks
+ if disable_self_attentions is not None:
+ # should be a list of booleans, indicating whether to disable self-attention in TransformerBlocks or not
+ assert len(disable_self_attentions) == len(channel_mult)
+ if num_attention_blocks is not None:
+ assert len(num_attention_blocks) == len(self.num_res_blocks)
+ assert all(map(lambda i: self.num_res_blocks[i] >= num_attention_blocks[i], range(len(num_attention_blocks))))
+ print(f"Constructor of UNetModel received num_attention_blocks={num_attention_blocks}. "
+ f"This option has LESS priority than attention_resolutions {attention_resolutions}, "
+ f"i.e., in cases where num_attention_blocks[i] > 0 but 2**i not in attention_resolutions, "
+ f"attention will still not be set.")
+ self.use_fp16 = use_fp16
+ self.attention_resolutions = attention_resolutions
+ self.dropout = dropout
+ self.channel_mult = channel_mult
+ self.conv_resample = conv_resample
+ self.num_classes = num_classes
+ self.use_checkpoint = use_checkpoint
+ self.dtype = th.float16 if use_fp16 else th.float32
+ self.num_heads = num_heads
+ self.num_head_channels = num_head_channels
+ self.num_heads_upsample = num_heads_upsample
+ self.predict_codebook_ids = n_embed is not None
+
+ time_embed_dim = model_channels * 4
+ self.time_embed = nn.Sequential(
+ linear(model_channels, time_embed_dim),
+ nn.SiLU(),
+ linear(time_embed_dim, time_embed_dim),
+ )
+
+ if self.num_classes is not None:
+ if isinstance(self.num_classes, int):
+ self.label_emb = nn.Embedding(num_classes, time_embed_dim)
+ elif self.num_classes == "continuous":
+ print("setting up linear c_adm embedding layer")
+ self.label_emb = nn.Linear(1, time_embed_dim)
+ else:
+ raise ValueError()
+
+ self.input_blocks = nn.ModuleList(
+ [
+ TimestepEmbedSequential(
+ conv_nd(dims, in_channels, model_channels, 3, padding=1)
+ )
+ ]
+ )
+ self._feature_size = model_channels
+ input_block_chans = [model_channels]
+ ch = model_channels
+ ds = 1
+ for level, mult in enumerate(channel_mult):
+ for nr in range(self.num_res_blocks[level]):
+ layers = [
+ ResBlock(
+ ch,
+ time_embed_dim,
+ dropout,
+ out_channels=mult * model_channels,
+ dims=dims,
+ use_checkpoint=use_checkpoint,
+ use_scale_shift_norm=use_scale_shift_norm,
+ )
+ ]
+ ch = mult * model_channels
+ if ds in attention_resolutions:
+ if num_head_channels == -1:
+ dim_head = ch // num_heads
+ else:
+ num_heads = ch // num_head_channels
+ dim_head = num_head_channels
+ if legacy:
+ #num_heads = 1
+ dim_head = ch // num_heads if use_spatial_transformer else num_head_channels
+ if exists(disable_self_attentions):
+ disabled_sa = disable_self_attentions[level]
+ else:
+ disabled_sa = False
+
+ if not exists(num_attention_blocks) or nr < num_attention_blocks[level]:
+ layers.append(
+ AttentionBlock(
+ ch,
+ use_checkpoint=use_checkpoint,
+ num_heads=num_heads,
+ num_head_channels=dim_head,
+ use_new_attention_order=use_new_attention_order,
+ ) if not use_spatial_transformer else SpatialTransformer(
+ ch, num_heads, dim_head, depth=transformer_depth, context_dim=context_dim,
+ disable_self_attn=disabled_sa, use_linear=use_linear_in_transformer,
+ use_checkpoint=use_checkpoint
+ )
+ )
+ self.input_blocks.append(TimestepEmbedSequential(*layers))
+ self._feature_size += ch
+ input_block_chans.append(ch)
+ if level != len(channel_mult) - 1:
+ out_ch = ch
+ self.input_blocks.append(
+ TimestepEmbedSequential(
+ ResBlock(
+ ch,
+ time_embed_dim,
+ dropout,
+ out_channels=out_ch,
+ dims=dims,
+ use_checkpoint=use_checkpoint,
+ use_scale_shift_norm=use_scale_shift_norm,
+ down=True,
+ )
+ if resblock_updown
+ else Downsample(
+ ch, conv_resample, dims=dims, out_channels=out_ch
+ )
+ )
+ )
+ ch = out_ch
+ input_block_chans.append(ch)
+ ds *= 2
+ self._feature_size += ch
+
+ if num_head_channels == -1:
+ dim_head = ch // num_heads
+ else:
+ num_heads = ch // num_head_channels
+ dim_head = num_head_channels
+ if legacy:
+ #num_heads = 1
+ dim_head = ch // num_heads if use_spatial_transformer else num_head_channels
+ self.middle_block = TimestepEmbedSequential(
+ ResBlock(
+ ch,
+ time_embed_dim,
+ dropout,
+ dims=dims,
+ use_checkpoint=use_checkpoint,
+ use_scale_shift_norm=use_scale_shift_norm,
+ ),
+ AttentionBlock(
+ ch,
+ use_checkpoint=use_checkpoint,
+ num_heads=num_heads,
+ num_head_channels=dim_head,
+ use_new_attention_order=use_new_attention_order,
+ ) if not use_spatial_transformer else SpatialTransformer( # always uses a self-attn
+ ch, num_heads, dim_head, depth=transformer_depth, context_dim=context_dim,
+ disable_self_attn=disable_middle_self_attn, use_linear=use_linear_in_transformer,
+ use_checkpoint=use_checkpoint
+ ),
+ ResBlock(
+ ch,
+ time_embed_dim,
+ dropout,
+ dims=dims,
+ use_checkpoint=use_checkpoint,
+ use_scale_shift_norm=use_scale_shift_norm,
+ ),
+ )
+ self._feature_size += ch
+
+ self.output_blocks = nn.ModuleList([])
+ for level, mult in list(enumerate(channel_mult))[::-1]:
+ for i in range(self.num_res_blocks[level] + 1):
+ ich = input_block_chans.pop()
+ layers = [
+ ResBlock(
+ ch + ich,
+ time_embed_dim,
+ dropout,
+ out_channels=model_channels * mult,
+ dims=dims,
+ use_checkpoint=use_checkpoint,
+ use_scale_shift_norm=use_scale_shift_norm,
+ )
+ ]
+ ch = model_channels * mult
+ if ds in attention_resolutions:
+ if num_head_channels == -1:
+ dim_head = ch // num_heads
+ else:
+ num_heads = ch // num_head_channels
+ dim_head = num_head_channels
+ if legacy:
+ #num_heads = 1
+ dim_head = ch // num_heads if use_spatial_transformer else num_head_channels
+ if exists(disable_self_attentions):
+ disabled_sa = disable_self_attentions[level]
+ else:
+ disabled_sa = False
+
+ if not exists(num_attention_blocks) or i < num_attention_blocks[level]:
+ layers.append(
+ AttentionBlock(
+ ch,
+ use_checkpoint=use_checkpoint,
+ num_heads=num_heads_upsample,
+ num_head_channels=dim_head,
+ use_new_attention_order=use_new_attention_order,
+ ) if not use_spatial_transformer else SpatialTransformer(
+ ch, num_heads, dim_head, depth=transformer_depth, context_dim=context_dim,
+ disable_self_attn=disabled_sa, use_linear=use_linear_in_transformer,
+ use_checkpoint=use_checkpoint
+ )
+ )
+ if level and i == self.num_res_blocks[level]:
+ out_ch = ch
+ layers.append(
+ ResBlock(
+ ch,
+ time_embed_dim,
+ dropout,
+ out_channels=out_ch,
+ dims=dims,
+ use_checkpoint=use_checkpoint,
+ use_scale_shift_norm=use_scale_shift_norm,
+ up=True,
+ )
+ if resblock_updown
+ else Upsample(ch, conv_resample, dims=dims, out_channels=out_ch)
+ )
+ ds //= 2
+ self.output_blocks.append(TimestepEmbedSequential(*layers))
+ self._feature_size += ch
+
+ self.out = nn.Sequential(
+ normalization(ch),
+ nn.SiLU(),
+ zero_module(conv_nd(dims, model_channels, out_channels, 3, padding=1)),
+ )
+ if self.predict_codebook_ids:
+ self.id_predictor = nn.Sequential(
+ normalization(ch),
+ conv_nd(dims, model_channels, n_embed, 1),
+ #nn.LogSoftmax(dim=1) # change to cross_entropy and produce non-normalized logits
+ )
+
+ def convert_to_fp16(self):
+ """
+ Convert the torso of the model to float16.
+ """
+ self.input_blocks.apply(convert_module_to_f16)
+ self.middle_block.apply(convert_module_to_f16)
+ self.output_blocks.apply(convert_module_to_f16)
+
+ def convert_to_fp32(self):
+ """
+ Convert the torso of the model to float32.
+ """
+ self.input_blocks.apply(convert_module_to_f32)
+ self.middle_block.apply(convert_module_to_f32)
+ self.output_blocks.apply(convert_module_to_f32)
+
+ def forward(self, x, timesteps=None, context=None, y=None,**kwargs):
+ """
+ Apply the model to an input batch.
+ :param x: an [N x C x ...] Tensor of inputs.
+ :param timesteps: a 1-D batch of timesteps.
+ :param context: conditioning plugged in via crossattn
+ :param y: an [N] Tensor of labels, if class-conditional.
+ :return: an [N x C x ...] Tensor of outputs.
+ """
+ assert (y is not None) == (
+ self.num_classes is not None
+ ), "must specify y if and only if the model is class-conditional"
+ hs = []
+ t_emb = timestep_embedding(timesteps, self.model_channels, repeat_only=False)
+ emb = self.time_embed(t_emb)
+
+ if self.num_classes is not None:
+ assert y.shape[0] == x.shape[0]
+ emb = emb + self.label_emb(y)
+
+ h = x.type(self.dtype)
+ for module in self.input_blocks:
+ h = module(h, emb, context)
+ hs.append(h)
+ h = self.middle_block(h, emb, context)
+ for module in self.output_blocks:
+ h = th.cat([h, hs.pop()], dim=1)
+ h = module(h, emb, context)
+ h = h.type(x.dtype)
+ if self.predict_codebook_ids:
+ return self.id_predictor(h)
+ else:
+ return self.out(h)
diff --git a/iopaint/model/anytext/ldm/modules/distributions/distributions.py b/iopaint/model/anytext/ldm/modules/distributions/distributions.py
new file mode 100644
index 0000000000000000000000000000000000000000..f2b8ef901130efc171aa69742ca0244d94d3f2e9
--- /dev/null
+++ b/iopaint/model/anytext/ldm/modules/distributions/distributions.py
@@ -0,0 +1,92 @@
+import torch
+import numpy as np
+
+
+class AbstractDistribution:
+ def sample(self):
+ raise NotImplementedError()
+
+ def mode(self):
+ raise NotImplementedError()
+
+
+class DiracDistribution(AbstractDistribution):
+ def __init__(self, value):
+ self.value = value
+
+ def sample(self):
+ return self.value
+
+ def mode(self):
+ return self.value
+
+
+class DiagonalGaussianDistribution(object):
+ def __init__(self, parameters, deterministic=False):
+ self.parameters = parameters
+ self.mean, self.logvar = torch.chunk(parameters, 2, dim=1)
+ self.logvar = torch.clamp(self.logvar, -30.0, 20.0)
+ self.deterministic = deterministic
+ self.std = torch.exp(0.5 * self.logvar)
+ self.var = torch.exp(self.logvar)
+ if self.deterministic:
+ self.var = self.std = torch.zeros_like(self.mean).to(device=self.parameters.device)
+
+ def sample(self):
+ x = self.mean + self.std * torch.randn(self.mean.shape).to(device=self.parameters.device)
+ return x
+
+ def kl(self, other=None):
+ if self.deterministic:
+ return torch.Tensor([0.])
+ else:
+ if other is None:
+ return 0.5 * torch.sum(torch.pow(self.mean, 2)
+ + self.var - 1.0 - self.logvar,
+ dim=[1, 2, 3])
+ else:
+ return 0.5 * torch.sum(
+ torch.pow(self.mean - other.mean, 2) / other.var
+ + self.var / other.var - 1.0 - self.logvar + other.logvar,
+ dim=[1, 2, 3])
+
+ def nll(self, sample, dims=[1,2,3]):
+ if self.deterministic:
+ return torch.Tensor([0.])
+ logtwopi = np.log(2.0 * np.pi)
+ return 0.5 * torch.sum(
+ logtwopi + self.logvar + torch.pow(sample - self.mean, 2) / self.var,
+ dim=dims)
+
+ def mode(self):
+ return self.mean
+
+
+def normal_kl(mean1, logvar1, mean2, logvar2):
+ """
+ source: https://github.com/openai/guided-diffusion/blob/27c20a8fab9cb472df5d6bdd6c8d11c8f430b924/guided_diffusion/losses.py#L12
+ Compute the KL divergence between two gaussians.
+ Shapes are automatically broadcasted, so batches can be compared to
+ scalars, among other use cases.
+ """
+ tensor = None
+ for obj in (mean1, logvar1, mean2, logvar2):
+ if isinstance(obj, torch.Tensor):
+ tensor = obj
+ break
+ assert tensor is not None, "at least one argument must be a Tensor"
+
+ # Force variances to be Tensors. Broadcasting helps convert scalars to
+ # Tensors, but it does not work for torch.exp().
+ logvar1, logvar2 = [
+ x if isinstance(x, torch.Tensor) else torch.tensor(x).to(tensor)
+ for x in (logvar1, logvar2)
+ ]
+
+ return 0.5 * (
+ -1.0
+ + logvar2
+ - logvar1
+ + torch.exp(logvar1 - logvar2)
+ + ((mean1 - mean2) ** 2) * torch.exp(-logvar2)
+ )
diff --git a/iopaint/model/anytext/ldm/modules/ema.py b/iopaint/model/anytext/ldm/modules/ema.py
new file mode 100644
index 0000000000000000000000000000000000000000..bded25019b9bcbcd0260f0b8185f8c7859ca58c4
--- /dev/null
+++ b/iopaint/model/anytext/ldm/modules/ema.py
@@ -0,0 +1,80 @@
+import torch
+from torch import nn
+
+
+class LitEma(nn.Module):
+ def __init__(self, model, decay=0.9999, use_num_upates=True):
+ super().__init__()
+ if decay < 0.0 or decay > 1.0:
+ raise ValueError('Decay must be between 0 and 1')
+
+ self.m_name2s_name = {}
+ self.register_buffer('decay', torch.tensor(decay, dtype=torch.float32))
+ self.register_buffer('num_updates', torch.tensor(0, dtype=torch.int) if use_num_upates
+ else torch.tensor(-1, dtype=torch.int))
+
+ for name, p in model.named_parameters():
+ if p.requires_grad:
+ # remove as '.'-character is not allowed in buffers
+ s_name = name.replace('.', '')
+ self.m_name2s_name.update({name: s_name})
+ self.register_buffer(s_name, p.clone().detach().data)
+
+ self.collected_params = []
+
+ def reset_num_updates(self):
+ del self.num_updates
+ self.register_buffer('num_updates', torch.tensor(0, dtype=torch.int))
+
+ def forward(self, model):
+ decay = self.decay
+
+ if self.num_updates >= 0:
+ self.num_updates += 1
+ decay = min(self.decay, (1 + self.num_updates) / (10 + self.num_updates))
+
+ one_minus_decay = 1.0 - decay
+
+ with torch.no_grad():
+ m_param = dict(model.named_parameters())
+ shadow_params = dict(self.named_buffers())
+
+ for key in m_param:
+ if m_param[key].requires_grad:
+ sname = self.m_name2s_name[key]
+ shadow_params[sname] = shadow_params[sname].type_as(m_param[key])
+ shadow_params[sname].sub_(one_minus_decay * (shadow_params[sname] - m_param[key]))
+ else:
+ assert not key in self.m_name2s_name
+
+ def copy_to(self, model):
+ m_param = dict(model.named_parameters())
+ shadow_params = dict(self.named_buffers())
+ for key in m_param:
+ if m_param[key].requires_grad:
+ m_param[key].data.copy_(shadow_params[self.m_name2s_name[key]].data)
+ else:
+ assert not key in self.m_name2s_name
+
+ def store(self, parameters):
+ """
+ Save the current parameters for restoring later.
+ Args:
+ parameters: Iterable of `torch.nn.Parameter`; the parameters to be
+ temporarily stored.
+ """
+ self.collected_params = [param.clone() for param in parameters]
+
+ def restore(self, parameters):
+ """
+ Restore the parameters stored with the `store` method.
+ Useful to validate the model with EMA parameters without affecting the
+ original optimization process. Store the parameters before the
+ `copy_to` method. After validation (or model saving), use this to
+ restore the former parameters.
+ Args:
+ parameters: Iterable of `torch.nn.Parameter`; the parameters to be
+ updated with the stored parameters.
+ """
+ for c_param, param in zip(self.collected_params, parameters):
+ param.data.copy_(c_param.data)
diff --git a/iopaint/model/anytext/ldm/modules/encoders/modules.py b/iopaint/model/anytext/ldm/modules/encoders/modules.py
new file mode 100644
index 0000000000000000000000000000000000000000..ceac395f935c70653224f97205c3ec37fcc84db4
--- /dev/null
+++ b/iopaint/model/anytext/ldm/modules/encoders/modules.py
@@ -0,0 +1,411 @@
+import torch
+import torch.nn as nn
+from torch.utils.checkpoint import checkpoint
+
+from transformers import (
+ T5Tokenizer,
+ T5EncoderModel,
+ CLIPTokenizer,
+ CLIPTextModel,
+ AutoProcessor,
+ CLIPVisionModelWithProjection,
+)
+
+from iopaint.model.anytext.ldm.util import count_params
+
+
+def _expand_mask(mask, dtype, tgt_len=None):
+ """
+ Expands attention_mask from `[bsz, seq_len]` to `[bsz, 1, tgt_seq_len, src_seq_len]`.
+ """
+ bsz, src_len = mask.size()
+ tgt_len = tgt_len if tgt_len is not None else src_len
+
+ expanded_mask = mask[:, None, None, :].expand(bsz, 1, tgt_len, src_len).to(dtype)
+
+ inverted_mask = 1.0 - expanded_mask
+
+ return inverted_mask.masked_fill(
+ inverted_mask.to(torch.bool), torch.finfo(dtype).min
+ )
+
+
+def _build_causal_attention_mask(bsz, seq_len, dtype):
+ # lazily create causal attention mask, with full attention between the vision tokens
+ # pytorch uses additive attention mask; fill with -inf
+ mask = torch.empty(bsz, seq_len, seq_len, dtype=dtype)
+ mask.fill_(torch.tensor(torch.finfo(dtype).min))
+ mask.triu_(1) # zero out the lower diagonal
+ mask = mask.unsqueeze(1) # expand mask
+ return mask
+
+
+class AbstractEncoder(nn.Module):
+ def __init__(self):
+ super().__init__()
+
+ def encode(self, *args, **kwargs):
+ raise NotImplementedError
+
+
+class IdentityEncoder(AbstractEncoder):
+ def encode(self, x):
+ return x
+
+
+class ClassEmbedder(nn.Module):
+ def __init__(self, embed_dim, n_classes=1000, key="class", ucg_rate=0.1):
+ super().__init__()
+ self.key = key
+ self.embedding = nn.Embedding(n_classes, embed_dim)
+ self.n_classes = n_classes
+ self.ucg_rate = ucg_rate
+
+ def forward(self, batch, key=None, disable_dropout=False):
+ if key is None:
+ key = self.key
+ # this is for use in crossattn
+ c = batch[key][:, None]
+ if self.ucg_rate > 0.0 and not disable_dropout:
+ mask = 1.0 - torch.bernoulli(torch.ones_like(c) * self.ucg_rate)
+ c = mask * c + (1 - mask) * torch.ones_like(c) * (self.n_classes - 1)
+ c = c.long()
+ c = self.embedding(c)
+ return c
+
+ def get_unconditional_conditioning(self, bs, device="cuda"):
+ uc_class = (
+ self.n_classes - 1
+ ) # 1000 classes --> 0 ... 999, one extra class for ucg (class 1000)
+ uc = torch.ones((bs,), device=device) * uc_class
+ uc = {self.key: uc}
+ return uc
+
+
+def disabled_train(self, mode=True):
+ """Overwrite model.train with this function to make sure train/eval mode
+ does not change anymore."""
+ return self
+
+
+class FrozenT5Embedder(AbstractEncoder):
+ """Uses the T5 transformer encoder for text"""
+
+ def __init__(
+ self, version="google/t5-v1_1-large", device="cuda", max_length=77, freeze=True
+ ): # others are google/t5-v1_1-xl and google/t5-v1_1-xxl
+ super().__init__()
+ self.tokenizer = T5Tokenizer.from_pretrained(version)
+ self.transformer = T5EncoderModel.from_pretrained(version)
+ self.device = device
+ self.max_length = max_length # TODO: typical value?
+ if freeze:
+ self.freeze()
+
+ def freeze(self):
+ self.transformer = self.transformer.eval()
+ # self.train = disabled_train
+ for param in self.parameters():
+ param.requires_grad = False
+
+ def forward(self, text):
+ batch_encoding = self.tokenizer(
+ text,
+ truncation=True,
+ max_length=self.max_length,
+ return_length=True,
+ return_overflowing_tokens=False,
+ padding="max_length",
+ return_tensors="pt",
+ )
+ tokens = batch_encoding["input_ids"].to(self.device)
+ outputs = self.transformer(input_ids=tokens)
+
+ z = outputs.last_hidden_state
+ return z
+
+ def encode(self, text):
+ return self(text)
+
+
+class FrozenCLIPEmbedder(AbstractEncoder):
+ """Uses the CLIP transformer encoder for text (from huggingface)"""
+
+ LAYERS = ["last", "pooled", "hidden"]
+
+ def __init__(
+ self,
+ version="openai/clip-vit-large-patch14",
+ device="cuda",
+ max_length=77,
+ freeze=True,
+ layer="last",
+ layer_idx=None,
+ ): # clip-vit-base-patch32
+ super().__init__()
+ assert layer in self.LAYERS
+ self.tokenizer = CLIPTokenizer.from_pretrained(version)
+ self.transformer = CLIPTextModel.from_pretrained(version)
+ self.device = device
+ self.max_length = max_length
+ if freeze:
+ self.freeze()
+ self.layer = layer
+ self.layer_idx = layer_idx
+ if layer == "hidden":
+ assert layer_idx is not None
+ assert 0 <= abs(layer_idx) <= 12
+
+ def freeze(self):
+ self.transformer = self.transformer.eval()
+ # self.train = disabled_train
+ for param in self.parameters():
+ param.requires_grad = False
+
+ def forward(self, text):
+ batch_encoding = self.tokenizer(
+ text,
+ truncation=True,
+ max_length=self.max_length,
+ return_length=True,
+ return_overflowing_tokens=False,
+ padding="max_length",
+ return_tensors="pt",
+ )
+ tokens = batch_encoding["input_ids"].to(self.device)
+ outputs = self.transformer(
+ input_ids=tokens, output_hidden_states=self.layer == "hidden"
+ )
+ if self.layer == "last":
+ z = outputs.last_hidden_state
+ elif self.layer == "pooled":
+ z = outputs.pooler_output[:, None, :]
+ else:
+ z = outputs.hidden_states[self.layer_idx]
+ return z
+
+ def encode(self, text):
+ return self(text)
+
+
+class FrozenCLIPT5Encoder(AbstractEncoder):
+ def __init__(
+ self,
+ clip_version="openai/clip-vit-large-patch14",
+ t5_version="google/t5-v1_1-xl",
+ device="cuda",
+ clip_max_length=77,
+ t5_max_length=77,
+ ):
+ super().__init__()
+ self.clip_encoder = FrozenCLIPEmbedder(
+ clip_version, device, max_length=clip_max_length
+ )
+ self.t5_encoder = FrozenT5Embedder(t5_version, device, max_length=t5_max_length)
+ print(
+ f"{self.clip_encoder.__class__.__name__} has {count_params(self.clip_encoder)*1.e-6:.2f} M parameters, "
+ f"{self.t5_encoder.__class__.__name__} comes with {count_params(self.t5_encoder)*1.e-6:.2f} M params."
+ )
+
+ def encode(self, text):
+ return self(text)
+
+ def forward(self, text):
+ clip_z = self.clip_encoder.encode(text)
+ t5_z = self.t5_encoder.encode(text)
+ return [clip_z, t5_z]
+
+
+class FrozenCLIPEmbedderT3(AbstractEncoder):
+ """Uses the CLIP transformer encoder for text (from Hugging Face)"""
+
+ def __init__(
+ self,
+ version="openai/clip-vit-large-patch14",
+ device="cuda",
+ max_length=77,
+ freeze=True,
+ use_vision=False,
+ ):
+ super().__init__()
+ self.tokenizer = CLIPTokenizer.from_pretrained(version)
+ self.transformer = CLIPTextModel.from_pretrained(version)
+ if use_vision:
+ self.vit = CLIPVisionModelWithProjection.from_pretrained(version)
+ self.processor = AutoProcessor.from_pretrained(version)
+ self.device = device
+ self.max_length = max_length
+ if freeze:
+ self.freeze()
+
+ def embedding_forward(
+ self,
+ input_ids=None,
+ position_ids=None,
+ inputs_embeds=None,
+ embedding_manager=None,
+ ):
+ seq_length = (
+ input_ids.shape[-1]
+ if input_ids is not None
+ else inputs_embeds.shape[-2]
+ )
+ if position_ids is None:
+ position_ids = self.position_ids[:, :seq_length]
+ if inputs_embeds is None:
+ inputs_embeds = self.token_embedding(input_ids)
+ if embedding_manager is not None:
+ inputs_embeds = embedding_manager(input_ids, inputs_embeds)
+ position_embeddings = self.position_embedding(position_ids)
+ embeddings = inputs_embeds + position_embeddings
+ return embeddings
+
+ self.transformer.text_model.embeddings.forward = embedding_forward.__get__(
+ self.transformer.text_model.embeddings
+ )
+
+ def encoder_forward(
+ self,
+ inputs_embeds,
+ attention_mask=None,
+ causal_attention_mask=None,
+ output_attentions=None,
+ output_hidden_states=None,
+ return_dict=None,
+ ):
+ output_attentions = (
+ output_attentions
+ if output_attentions is not None
+ else self.config.output_attentions
+ )
+ output_hidden_states = (
+ output_hidden_states
+ if output_hidden_states is not None
+ else self.config.output_hidden_states
+ )
+ return_dict = (
+ return_dict if return_dict is not None else self.config.use_return_dict
+ )
+ encoder_states = () if output_hidden_states else None
+ all_attentions = () if output_attentions else None
+ hidden_states = inputs_embeds
+ for idx, encoder_layer in enumerate(self.layers):
+ if output_hidden_states:
+ encoder_states = encoder_states + (hidden_states,)
+ layer_outputs = encoder_layer(
+ hidden_states,
+ attention_mask,
+ causal_attention_mask,
+ output_attentions=output_attentions,
+ )
+ hidden_states = layer_outputs[0]
+ if output_attentions:
+ all_attentions = all_attentions + (layer_outputs[1],)
+ if output_hidden_states:
+ encoder_states = encoder_states + (hidden_states,)
+ return hidden_states
+
+ self.transformer.text_model.encoder.forward = encoder_forward.__get__(
+ self.transformer.text_model.encoder
+ )
+
+ def text_encoder_forward(
+ self,
+ input_ids=None,
+ attention_mask=None,
+ position_ids=None,
+ output_attentions=None,
+ output_hidden_states=None,
+ return_dict=None,
+ embedding_manager=None,
+ ):
+ output_attentions = (
+ output_attentions
+ if output_attentions is not None
+ else self.config.output_attentions
+ )
+ output_hidden_states = (
+ output_hidden_states
+ if output_hidden_states is not None
+ else self.config.output_hidden_states
+ )
+ return_dict = (
+ return_dict if return_dict is not None else self.config.use_return_dict
+ )
+ if input_ids is None:
+ raise ValueError("You have to specify either input_ids")
+ input_shape = input_ids.size()
+ input_ids = input_ids.view(-1, input_shape[-1])
+ hidden_states = self.embeddings(
+ input_ids=input_ids,
+ position_ids=position_ids,
+ embedding_manager=embedding_manager,
+ )
+ bsz, seq_len = input_shape
+ # CLIP's text model uses causal mask, prepare it here.
+ # https://github.com/openai/CLIP/blob/cfcffb90e69f37bf2ff1e988237a0fbe41f33c04/clip/model.py#L324
+ causal_attention_mask = _build_causal_attention_mask(
+ bsz, seq_len, hidden_states.dtype
+ ).to(hidden_states.device)
+ # expand attention_mask
+ if attention_mask is not None:
+ # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len]
+ attention_mask = _expand_mask(attention_mask, hidden_states.dtype)
+ last_hidden_state = self.encoder(
+ inputs_embeds=hidden_states,
+ attention_mask=attention_mask,
+ causal_attention_mask=causal_attention_mask,
+ output_attentions=output_attentions,
+ output_hidden_states=output_hidden_states,
+ return_dict=return_dict,
+ )
+ last_hidden_state = self.final_layer_norm(last_hidden_state)
+ return last_hidden_state
+
+ self.transformer.text_model.forward = text_encoder_forward.__get__(
+ self.transformer.text_model
+ )
+
+ def transformer_forward(
+ self,
+ input_ids=None,
+ attention_mask=None,
+ position_ids=None,
+ output_attentions=None,
+ output_hidden_states=None,
+ return_dict=None,
+ embedding_manager=None,
+ ):
+ return self.text_model(
+ input_ids=input_ids,
+ attention_mask=attention_mask,
+ position_ids=position_ids,
+ output_attentions=output_attentions,
+ output_hidden_states=output_hidden_states,
+ return_dict=return_dict,
+ embedding_manager=embedding_manager,
+ )
+
+ self.transformer.forward = transformer_forward.__get__(self.transformer)
+
+ def freeze(self):
+ self.transformer = self.transformer.eval()
+ for param in self.parameters():
+ param.requires_grad = False
+
+ def forward(self, text, **kwargs):
+ batch_encoding = self.tokenizer(
+ text,
+ truncation=True,
+ max_length=self.max_length,
+ return_length=True,
+ return_overflowing_tokens=False,
+ padding="max_length",
+ return_tensors="pt",
+ )
+ tokens = batch_encoding["input_ids"].to(self.device)
+ z = self.transformer(input_ids=tokens, **kwargs)
+ return z
+
+ def encode(self, text, **kwargs):
+ return self(text, **kwargs)
diff --git a/iopaint/model/anytext/main.py b/iopaint/model/anytext/main.py
new file mode 100644
index 0000000000000000000000000000000000000000..f7b2d2ec984bc80ac8c6edef257a294ccab18875
--- /dev/null
+++ b/iopaint/model/anytext/main.py
@@ -0,0 +1,45 @@
+import cv2
+import os
+
+from anytext_pipeline import AnyTextPipeline
+from utils import save_images
+
+seed = 66273235
+# seed_everything(seed)
+
+pipe = AnyTextPipeline(
+ ckpt_path="/Users/cwq/code/github/IOPaint/iopaint/model/anytext/anytext_v1.1_fp16.ckpt",
+ font_path="/Users/cwq/code/github/AnyText/anytext/font/SourceHanSansSC-Medium.otf",
+ use_fp16=False,
+ device="mps",
+)
+
+img_save_folder = "SaveImages"
+rgb_image = cv2.imread(
+ "/Users/cwq/code/github/AnyText/anytext/example_images/ref7.jpg"
+)[..., ::-1]
+
+masked_image = cv2.imread(
+ "/Users/cwq/code/github/AnyText/anytext/example_images/edit7.png"
+)[..., ::-1]
+
+rgb_image = cv2.resize(rgb_image, (512, 512))
+masked_image = cv2.resize(masked_image, (512, 512))
+
+# results: list of rgb ndarray
+results, rtn_code, rtn_warning = pipe(
+ prompt='A cake with colorful characters that reads "EVERYDAY", best quality, extremely detailed,4k, HD, supper legible text, clear text edges, clear strokes, neat writing, no watermarks',
+ negative_prompt="low-res, bad anatomy, extra digit, fewer digits, cropped, worst quality, low quality, watermark, unreadable text, messy words, distorted text, disorganized writing, advertising picture",
+ image=rgb_image,
+ masked_image=masked_image,
+ num_inference_steps=20,
+ strength=1.0,
+ guidance_scale=9.0,
+ height=rgb_image.shape[0],
+ width=rgb_image.shape[1],
+ seed=seed,
+ sort_priority="y",
+)
+if rtn_code >= 0:
+ save_images(results, img_save_folder)
+ print(f"Done, result images are saved in: {img_save_folder}")
diff --git a/iopaint/model/anytext/ocr_recog/common.py b/iopaint/model/anytext/ocr_recog/common.py
new file mode 100644
index 0000000000000000000000000000000000000000..a328bb034a37934b7437893b5c2e42cd3504c17f
--- /dev/null
+++ b/iopaint/model/anytext/ocr_recog/common.py
@@ -0,0 +1,74 @@
+
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class Hswish(nn.Module):
+ def __init__(self, inplace=True):
+ super(Hswish, self).__init__()
+ self.inplace = inplace
+
+ def forward(self, x):
+ return x * F.relu6(x + 3., inplace=self.inplace) / 6.
+
+# out = max(0, min(1, slop*x+offset))
+# paddle.fluid.layers.hard_sigmoid(x, slope=0.2, offset=0.5, name=None)
+class Hsigmoid(nn.Module):
+ def __init__(self, inplace=True):
+ super(Hsigmoid, self).__init__()
+ self.inplace = inplace
+
+ def forward(self, x):
+ # torch: F.relu6(x + 3., inplace=self.inplace) / 6.
+ # paddle: F.relu6(1.2 * x + 3., inplace=self.inplace) / 6.
+ return F.relu6(1.2 * x + 3., inplace=self.inplace) / 6.
+
+class GELU(nn.Module):
+ def __init__(self, inplace=True):
+ super(GELU, self).__init__()
+ self.inplace = inplace
+
+ def forward(self, x):
+ return torch.nn.functional.gelu(x)
+
+
+class Swish(nn.Module):
+ def __init__(self, inplace=True):
+ super(Swish, self).__init__()
+ self.inplace = inplace
+
+ def forward(self, x):
+ if self.inplace:
+ x.mul_(torch.sigmoid(x))
+ return x
+ else:
+ return x*torch.sigmoid(x)
+
+
+class Activation(nn.Module):
+ def __init__(self, act_type, inplace=True):
+ super(Activation, self).__init__()
+ act_type = act_type.lower()
+ if act_type == 'relu':
+ self.act = nn.ReLU(inplace=inplace)
+ elif act_type == 'relu6':
+ self.act = nn.ReLU6(inplace=inplace)
+ elif act_type == 'sigmoid':
+ raise NotImplementedError
+ elif act_type == 'hard_sigmoid':
+ self.act = Hsigmoid(inplace)
+ elif act_type == 'hard_swish':
+ self.act = Hswish(inplace=inplace)
+ elif act_type == 'leakyrelu':
+ self.act = nn.LeakyReLU(inplace=inplace)
+ elif act_type == 'gelu':
+ self.act = GELU(inplace=inplace)
+ elif act_type == 'swish':
+ self.act = Swish(inplace=inplace)
+ else:
+ raise NotImplementedError
+
+ def forward(self, inputs):
+ return self.act(inputs)
\ No newline at end of file
diff --git a/iopaint/model/anytext/ocr_recog/en_dict.txt b/iopaint/model/anytext/ocr_recog/en_dict.txt
new file mode 100644
index 0000000000000000000000000000000000000000..7677d31b9d3f08eef2823c2cf051beeab1f0470b
--- /dev/null
+++ b/iopaint/model/anytext/ocr_recog/en_dict.txt
@@ -0,0 +1,95 @@
+0
+1
+2
+3
+4
+5
+6
+7
+8
+9
+:
+;
+<
+=
+>
+?
+@
+A
+B
+C
+D
+E
+F
+G
+H
+I
+J
+K
+L
+M
+N
+O
+P
+Q
+R
+S
+T
+U
+V
+W
+X
+Y
+Z
+[
+\
+]
+^
+_
+`
+a
+b
+c
+d
+e
+f
+g
+h
+i
+j
+k
+l
+m
+n
+o
+p
+q
+r
+s
+t
+u
+v
+w
+x
+y
+z
+{
+|
+}
+~
+!
+"
+#
+$
+%
+&
+'
+(
+)
+*
++
+,
+-
+.
+/
+
diff --git a/iopaint/model/controlnet.py b/iopaint/model/controlnet.py
new file mode 100644
index 0000000000000000000000000000000000000000..d52db01926c799a5eada97e75af874363c3dd1cd
--- /dev/null
+++ b/iopaint/model/controlnet.py
@@ -0,0 +1,190 @@
+import PIL.Image
+import cv2
+import torch
+from diffusers import ControlNetModel
+from loguru import logger
+from iopaint.schema import InpaintRequest, ModelType
+
+from .base import DiffusionInpaintModel
+from .helper.controlnet_preprocess import (
+ make_canny_control_image,
+ make_openpose_control_image,
+ make_depth_control_image,
+ make_inpaint_control_image,
+)
+from .helper.cpu_text_encoder import CPUTextEncoderWrapper
+from .original_sd_configs import get_config_files
+from .utils import (
+ get_scheduler,
+ handle_from_pretrained_exceptions,
+ get_torch_dtype,
+ enable_low_mem,
+ is_local_files_only,
+)
+
+
+class ControlNet(DiffusionInpaintModel):
+ name = "controlnet"
+ pad_mod = 8
+ min_size = 512
+
+ @property
+ def lcm_lora_id(self):
+ if self.model_info.model_type in [
+ ModelType.DIFFUSERS_SD,
+ ModelType.DIFFUSERS_SD_INPAINT,
+ ]:
+ return "latent-consistency/lcm-lora-sdv1-5"
+ if self.model_info.model_type in [
+ ModelType.DIFFUSERS_SDXL,
+ ModelType.DIFFUSERS_SDXL_INPAINT,
+ ]:
+ return "latent-consistency/lcm-lora-sdxl"
+ raise NotImplementedError(f"Unsupported controlnet lcm model {self.model_info}")
+
+ def init_model(self, device: torch.device, **kwargs):
+ model_info = kwargs["model_info"]
+ controlnet_method = kwargs["controlnet_method"]
+
+ self.model_info = model_info
+ self.controlnet_method = controlnet_method
+
+ model_kwargs = {
+ **kwargs.get("pipe_components", {}),
+ "local_files_only": is_local_files_only(**kwargs),
+ }
+ self.local_files_only = model_kwargs["local_files_only"]
+
+ disable_nsfw_checker = kwargs["disable_nsfw"] or kwargs.get(
+ "cpu_offload", False
+ )
+ if disable_nsfw_checker:
+ logger.info("Disable Stable Diffusion Model NSFW checker")
+ model_kwargs.update(
+ dict(
+ safety_checker=None,
+ feature_extractor=None,
+ requires_safety_checker=False,
+ )
+ )
+
+ use_gpu, torch_dtype = get_torch_dtype(device, kwargs.get("no_half", False))
+ self.torch_dtype = torch_dtype
+
+ if model_info.model_type in [
+ ModelType.DIFFUSERS_SD,
+ ModelType.DIFFUSERS_SD_INPAINT,
+ ]:
+ from diffusers import (
+ StableDiffusionControlNetInpaintPipeline as PipeClass,
+ )
+ elif model_info.model_type in [
+ ModelType.DIFFUSERS_SDXL,
+ ModelType.DIFFUSERS_SDXL_INPAINT,
+ ]:
+ from diffusers import (
+ StableDiffusionXLControlNetInpaintPipeline as PipeClass,
+ )
+
+ controlnet = ControlNetModel.from_pretrained(
+ pretrained_model_name_or_path=controlnet_method,
+ resume_download=True,
+ local_files_only=model_kwargs["local_files_only"],
+ torch_dtype=self.torch_dtype,
+ )
+ if model_info.is_single_file_diffusers:
+ if self.model_info.model_type == ModelType.DIFFUSERS_SD:
+ model_kwargs["num_in_channels"] = 4
+ else:
+ model_kwargs["num_in_channels"] = 9
+
+ self.model = PipeClass.from_single_file(
+ model_info.path,
+ controlnet=controlnet,
+ load_safety_checker=not disable_nsfw_checker,
+ torch_dtype=torch_dtype,
+ config_files=get_config_files(),
+ **model_kwargs,
+ )
+ else:
+ self.model = handle_from_pretrained_exceptions(
+ PipeClass.from_pretrained,
+ pretrained_model_name_or_path=model_info.path,
+ controlnet=controlnet,
+ variant="fp16",
+ torch_dtype=torch_dtype,
+ **model_kwargs,
+ )
+
+ enable_low_mem(self.model, kwargs.get("low_mem", False))
+
+ if kwargs.get("cpu_offload", False) and use_gpu:
+ logger.info("Enable sequential cpu offload")
+ self.model.enable_sequential_cpu_offload(gpu_id=0)
+ else:
+ self.model = self.model.to(device)
+ if kwargs["sd_cpu_textencoder"]:
+ logger.info("Run Stable Diffusion TextEncoder on CPU")
+ self.model.text_encoder = CPUTextEncoderWrapper(
+ self.model.text_encoder, torch_dtype
+ )
+
+ self.callback = kwargs.pop("callback", None)
+
+ def switch_controlnet_method(self, new_method: str):
+ self.controlnet_method = new_method
+ controlnet = ControlNetModel.from_pretrained(
+ new_method,
+ resume_download=True,
+ local_files_only=self.local_files_only,
+ torch_dtype=self.torch_dtype,
+ ).to(self.model.device)
+ self.model.controlnet = controlnet
+
+ def _get_control_image(self, image, mask):
+ if "canny" in self.controlnet_method:
+ control_image = make_canny_control_image(image)
+ elif "openpose" in self.controlnet_method:
+ control_image = make_openpose_control_image(image)
+ elif "depth" in self.controlnet_method:
+ control_image = make_depth_control_image(image)
+ elif "inpaint" in self.controlnet_method:
+ control_image = make_inpaint_control_image(image, mask)
+ else:
+ raise NotImplementedError(f"{self.controlnet_method} not implemented")
+ return control_image
+
+ def forward(self, image, mask, config: InpaintRequest):
+ """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
+ """
+ scheduler_config = self.model.scheduler.config
+ scheduler = get_scheduler(config.sd_sampler, scheduler_config)
+ self.model.scheduler = scheduler
+
+ img_h, img_w = image.shape[:2]
+ control_image = self._get_control_image(image, mask)
+ mask_image = PIL.Image.fromarray(mask[:, :, -1], mode="L")
+ image = PIL.Image.fromarray(image)
+
+ output = self.model(
+ image=image,
+ mask_image=mask_image,
+ control_image=control_image,
+ prompt=config.prompt,
+ negative_prompt=config.negative_prompt,
+ num_inference_steps=config.sd_steps,
+ guidance_scale=config.sd_guidance_scale,
+ output_type="np",
+ callback_on_step_end=self.callback,
+ height=img_h,
+ width=img_w,
+ generator=torch.manual_seed(config.sd_seed),
+ controlnet_conditioning_scale=config.controlnet_conditioning_scale,
+ ).images[0]
+
+ output = (output * 255).round().astype("uint8")
+ output = cv2.cvtColor(output, cv2.COLOR_RGB2BGR)
+ return output
diff --git a/iopaint/model/ddim_sampler.py b/iopaint/model/ddim_sampler.py
new file mode 100644
index 0000000000000000000000000000000000000000..a3f44fd4146aa3a5388bbfd4dd51a90f9afc91df
--- /dev/null
+++ b/iopaint/model/ddim_sampler.py
@@ -0,0 +1,193 @@
+import torch
+import numpy as np
+from tqdm import tqdm
+
+from .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
diff --git a/iopaint/model/fcf.py b/iopaint/model/fcf.py
new file mode 100644
index 0000000000000000000000000000000000000000..a6f2d4222a35f4136fde04f2e41b3a9126bd84d3
--- /dev/null
+++ b/iopaint/model/fcf.py
@@ -0,0 +1,1737 @@
+import os
+import random
+
+import cv2
+import torch
+import numpy as np
+import torch.fft as fft
+
+from iopaint.schema import InpaintRequest
+
+from iopaint.helper import (
+ load_model,
+ get_cache_path_by_url,
+ norm_img,
+ boxes_from_mask,
+ resize_max_size,
+ download_model,
+)
+from .base import InpaintModel
+from torch import conv2d, nn
+import torch.nn.functional as F
+
+from .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",
+)
+FCF_MODEL_MD5 = os.environ.get("FCF_MODEL_MD5", "3323152bc01bf1c56fd8aba74435a211")
+
+
+class FcF(InpaintModel):
+ name = "fcf"
+ min_size = 512
+ pad_mod = 512
+ pad_to_square = True
+ is_erase_model = 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, FCF_MODEL_MD5)
+ self.label = torch.zeros([1, self.model.c_dim], device=device)
+
+ @staticmethod
+ def download():
+ download_model(FCF_MODEL_URL, FCF_MODEL_MD5)
+
+ @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: InpaintRequest):
+ """
+ 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].copy()
+ 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: InpaintRequest):
+ """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
diff --git a/iopaint/model/helper/controlnet_preprocess.py b/iopaint/model/helper/controlnet_preprocess.py
new file mode 100644
index 0000000000000000000000000000000000000000..75c409fac12b4f035defb14a4c02d1c46f4b2ba3
--- /dev/null
+++ b/iopaint/model/helper/controlnet_preprocess.py
@@ -0,0 +1,68 @@
+import torch
+import PIL
+import cv2
+from PIL import Image
+import numpy as np
+
+from iopaint.helper import pad_img_to_modulo
+
+
+def make_canny_control_image(image: np.ndarray) -> Image:
+ canny_image = cv2.Canny(image, 100, 200)
+ canny_image = canny_image[:, :, None]
+ canny_image = np.concatenate([canny_image, canny_image, canny_image], axis=2)
+ canny_image = PIL.Image.fromarray(canny_image)
+ control_image = canny_image
+ return control_image
+
+
+def make_openpose_control_image(image: np.ndarray) -> Image:
+ from controlnet_aux import OpenposeDetector
+
+ processor = OpenposeDetector.from_pretrained("lllyasviel/ControlNet")
+ control_image = processor(image, hand_and_face=True)
+ return control_image
+
+
+def resize_image(input_image, resolution):
+ H, W, C = input_image.shape
+ H = float(H)
+ W = float(W)
+ k = float(resolution) / min(H, W)
+ H *= k
+ W *= k
+ H = int(np.round(H / 64.0)) * 64
+ W = int(np.round(W / 64.0)) * 64
+ img = cv2.resize(
+ input_image,
+ (W, H),
+ interpolation=cv2.INTER_LANCZOS4 if k > 1 else cv2.INTER_AREA,
+ )
+ return img
+
+
+def make_depth_control_image(image: np.ndarray) -> Image:
+ from controlnet_aux import MidasDetector
+
+ midas = MidasDetector.from_pretrained("lllyasviel/Annotators")
+
+ origin_height, origin_width = image.shape[:2]
+ pad_image = pad_img_to_modulo(image, mod=64, square=False, min_size=512)
+ depth_image = midas(pad_image)
+ depth_image = depth_image[0:origin_height, 0:origin_width]
+ depth_image = depth_image[:, :, None]
+ depth_image = np.concatenate([depth_image, depth_image, depth_image], axis=2)
+ control_image = PIL.Image.fromarray(depth_image)
+ return control_image
+
+
+def make_inpaint_control_image(image: np.ndarray, mask: np.ndarray) -> torch.Tensor:
+ """
+ image: [H, W, C] RGB
+ mask: [H, W, 1] 255 means area to repaint
+ """
+ image = image.astype(np.float32) / 255.0
+ image[mask[:, :, -1] > 128] = -1.0 # set as masked pixel
+ image = np.expand_dims(image, 0).transpose(0, 3, 1, 2)
+ image = torch.from_numpy(image)
+ return image
diff --git a/iopaint/model/helper/cpu_text_encoder.py b/iopaint/model/helper/cpu_text_encoder.py
new file mode 100644
index 0000000000000000000000000000000000000000..116eb48b6ae9b49d0c198282746b35877b2c4aa5
--- /dev/null
+++ b/iopaint/model/helper/cpu_text_encoder.py
@@ -0,0 +1,41 @@
+import torch
+from transformers import PreTrainedModel
+
+from ..utils import torch_gc
+
+
+class CPUTextEncoderWrapper(PreTrainedModel):
+ def __init__(self, text_encoder, torch_dtype):
+ super().__init__(text_encoder.config)
+ self.config = text_encoder.config
+ self._device = text_encoder.device
+ # cpu not support float16
+ self.text_encoder = text_encoder.to(torch.device("cpu"), non_blocking=True)
+ self.text_encoder = self.text_encoder.to(torch.float32, non_blocking=True)
+ self.torch_dtype = torch_dtype
+ del text_encoder
+ torch_gc()
+
+ def __call__(self, x, **kwargs):
+ input_device = x.device
+ original_output = self.text_encoder(x.to(self.text_encoder.device), **kwargs)
+ for k, v in original_output.items():
+ if isinstance(v, tuple):
+ original_output[k] = [
+ v[i].to(input_device).to(self.torch_dtype) for i in range(len(v))
+ ]
+ else:
+ original_output[k] = v.to(input_device).to(self.torch_dtype)
+ return original_output
+
+ @property
+ def dtype(self):
+ return self.torch_dtype
+
+ @property
+ def device(self) -> torch.device:
+ """
+ `torch.device`: The device on which the module is (assuming that all the module parameters are on the same
+ device).
+ """
+ return self._device
\ No newline at end of file
diff --git a/iopaint/model/helper/g_diffuser_bot.py b/iopaint/model/helper/g_diffuser_bot.py
new file mode 100644
index 0000000000000000000000000000000000000000..a4147af4001d677c4ac5f3ce7e6e808902a0af46
--- /dev/null
+++ b/iopaint/model/helper/g_diffuser_bot.py
@@ -0,0 +1,167 @@
+# code copy from: https://github.com/parlance-zz/g-diffuser-bot
+import cv2
+import numpy as np
+
+
+def np_img_grey_to_rgb(data):
+ if data.ndim == 3:
+ return data
+ return np.expand_dims(data, 2) * np.ones((1, 1, 3))
+
+
+def convolve(data1, data2): # fast convolution with fft
+ if data1.ndim != data2.ndim: # promote to rgb if mismatch
+ if data1.ndim < 3:
+ data1 = np_img_grey_to_rgb(data1)
+ if data2.ndim < 3:
+ data2 = np_img_grey_to_rgb(data2)
+ return ifft2(fft2(data1) * fft2(data2))
+
+
+def fft2(data):
+ if data.ndim > 2: # multiple channels
+ out_fft = np.zeros(
+ (data.shape[0], data.shape[1], data.shape[2]), dtype=np.complex128
+ )
+ for c in range(data.shape[2]):
+ c_data = data[:, :, c]
+ out_fft[:, :, c] = np.fft.fft2(np.fft.fftshift(c_data), norm="ortho")
+ out_fft[:, :, c] = np.fft.ifftshift(out_fft[:, :, c])
+ else: # single channel
+ out_fft = np.zeros((data.shape[0], data.shape[1]), dtype=np.complex128)
+ out_fft[:, :] = np.fft.fft2(np.fft.fftshift(data), norm="ortho")
+ out_fft[:, :] = np.fft.ifftshift(out_fft[:, :])
+
+ return out_fft
+
+
+def ifft2(data):
+ if data.ndim > 2: # multiple channels
+ out_ifft = np.zeros(
+ (data.shape[0], data.shape[1], data.shape[2]), dtype=np.complex128
+ )
+ for c in range(data.shape[2]):
+ c_data = data[:, :, c]
+ out_ifft[:, :, c] = np.fft.ifft2(np.fft.fftshift(c_data), norm="ortho")
+ out_ifft[:, :, c] = np.fft.ifftshift(out_ifft[:, :, c])
+ else: # single channel
+ out_ifft = np.zeros((data.shape[0], data.shape[1]), dtype=np.complex128)
+ out_ifft[:, :] = np.fft.ifft2(np.fft.fftshift(data), norm="ortho")
+ out_ifft[:, :] = np.fft.ifftshift(out_ifft[:, :])
+
+ return out_ifft
+
+
+def get_gradient_kernel(width, height, std=3.14, mode="linear"):
+ window_scale_x = float(
+ width / min(width, height)
+ ) # for non-square aspect ratios we still want a circular kernel
+ window_scale_y = float(height / min(width, height))
+ if mode == "gaussian":
+ x = (np.arange(width) / width * 2.0 - 1.0) * window_scale_x
+ kx = np.exp(-x * x * std)
+ if window_scale_x != window_scale_y:
+ y = (np.arange(height) / height * 2.0 - 1.0) * window_scale_y
+ ky = np.exp(-y * y * std)
+ else:
+ y = x
+ ky = kx
+ return np.outer(kx, ky)
+ elif mode == "linear":
+ x = (np.arange(width) / width * 2.0 - 1.0) * window_scale_x
+ if window_scale_x != window_scale_y:
+ y = (np.arange(height) / height * 2.0 - 1.0) * window_scale_y
+ else:
+ y = x
+ return np.clip(1.0 - np.sqrt(np.add.outer(x * x, y * y)) * std / 3.14, 0.0, 1.0)
+ else:
+ raise Exception("Error: Unknown mode in get_gradient_kernel: {0}".format(mode))
+
+
+def image_blur(data, std=3.14, mode="linear"):
+ width = data.shape[0]
+ height = data.shape[1]
+ kernel = get_gradient_kernel(width, height, std, mode=mode)
+ return np.real(convolve(data, kernel / np.sqrt(np.sum(kernel * kernel))))
+
+
+def soften_mask(mask_img, softness, space):
+ if softness == 0:
+ return mask_img
+ softness = min(softness, 1.0)
+ space = np.clip(space, 0.0, 1.0)
+ original_max_opacity = np.max(mask_img)
+ out_mask = mask_img <= 0.0
+ blurred_mask = image_blur(mask_img, 3.5 / softness, mode="linear")
+ blurred_mask = np.maximum(blurred_mask - np.max(blurred_mask[out_mask]), 0.0)
+ mask_img *= blurred_mask # preserve partial opacity in original input mask
+ mask_img /= np.max(mask_img) # renormalize
+ mask_img = np.clip(mask_img - space, 0.0, 1.0) # make space
+ mask_img /= np.max(mask_img) # and renormalize again
+ mask_img *= original_max_opacity # restore original max opacity
+ return mask_img
+
+
+def expand_image(
+ cv2_img, top: int, right: int, bottom: int, left: int, softness: float, space: float
+):
+ assert cv2_img.shape[2] == 3
+ origin_h, origin_w = cv2_img.shape[:2]
+ new_width = cv2_img.shape[1] + left + right
+ new_height = cv2_img.shape[0] + top + bottom
+
+ # TODO: which is better?
+ # new_img = np.random.randint(0, 255, (new_height, new_width, 3), np.uint8)
+ new_img = cv2.copyMakeBorder(
+ cv2_img, top, bottom, left, right, cv2.BORDER_REPLICATE
+ )
+ mask_img = np.zeros((new_height, new_width), np.uint8)
+ mask_img[top : top + cv2_img.shape[0], left : left + cv2_img.shape[1]] = 255
+
+ if softness > 0.0:
+ mask_img = soften_mask(mask_img / 255.0, softness / 100.0, space / 100.0)
+ mask_img = (np.clip(mask_img, 0.0, 1.0) * 255.0).astype(np.uint8)
+
+ mask_image = 255.0 - mask_img # extract mask from alpha channel and invert
+ rgb_init_image = (
+ 0.0 + new_img[:, :, 0:3]
+ ) # strip mask from init_img leaving only rgb channels
+
+ hard_mask = np.zeros_like(cv2_img[:, :, 0])
+ if top != 0:
+ hard_mask[0 : origin_h // 2, :] = 255
+ if bottom != 0:
+ hard_mask[origin_h // 2 :, :] = 255
+ if left != 0:
+ hard_mask[:, 0 : origin_w // 2] = 255
+ if right != 0:
+ hard_mask[:, origin_w // 2 :] = 255
+
+ hard_mask = cv2.copyMakeBorder(
+ hard_mask, top, bottom, left, right, cv2.BORDER_DEFAULT, value=255
+ )
+ mask_image = np.where(hard_mask > 0, mask_image, 0)
+ return rgb_init_image.astype(np.uint8), mask_image.astype(np.uint8)
+
+
+if __name__ == "__main__":
+ from pathlib import Path
+
+ current_dir = Path(__file__).parent.absolute().resolve()
+ image_path = current_dir.parent / "tests" / "bunny.jpeg"
+ init_image = cv2.imread(str(image_path))
+ init_image, mask_image = expand_image(
+ init_image,
+ top=100,
+ right=100,
+ bottom=100,
+ left=100,
+ softness=20,
+ space=20,
+ )
+ print(mask_image.dtype, mask_image.min(), mask_image.max())
+ print(init_image.dtype, init_image.min(), init_image.max())
+ mask_image = mask_image.astype(np.uint8)
+ init_image = init_image.astype(np.uint8)
+ cv2.imwrite("expanded_image.png", init_image)
+ cv2.imwrite("expanded_mask.png", mask_image)
diff --git a/iopaint/model/instruct_pix2pix.py b/iopaint/model/instruct_pix2pix.py
new file mode 100644
index 0000000000000000000000000000000000000000..fc8cd26c2de87407e828b934b1dcddf1b0971f54
--- /dev/null
+++ b/iopaint/model/instruct_pix2pix.py
@@ -0,0 +1,64 @@
+import PIL.Image
+import cv2
+import torch
+from loguru import logger
+
+from iopaint.const import INSTRUCT_PIX2PIX_NAME
+from .base import DiffusionInpaintModel
+from iopaint.schema import InpaintRequest
+from .utils import get_torch_dtype, enable_low_mem, is_local_files_only
+
+
+class InstructPix2Pix(DiffusionInpaintModel):
+ name = INSTRUCT_PIX2PIX_NAME
+ pad_mod = 8
+ min_size = 512
+
+ def init_model(self, device: torch.device, **kwargs):
+ from diffusers import StableDiffusionInstructPix2PixPipeline
+
+ use_gpu, torch_dtype = get_torch_dtype(device, kwargs.get("no_half", False))
+
+ model_kwargs = {"local_files_only": is_local_files_only(**kwargs)}
+ if kwargs["disable_nsfw"] or kwargs.get("cpu_offload", False):
+ logger.info("Disable Stable Diffusion Model NSFW checker")
+ model_kwargs.update(
+ dict(
+ safety_checker=None,
+ feature_extractor=None,
+ requires_safety_checker=False,
+ )
+ )
+
+ self.model = StableDiffusionInstructPix2PixPipeline.from_pretrained(
+ self.name, variant="fp16", torch_dtype=torch_dtype, **model_kwargs
+ )
+ enable_low_mem(self.model, kwargs.get("low_mem", False))
+
+ if kwargs.get("cpu_offload", False) and use_gpu:
+ logger.info("Enable sequential cpu offload")
+ self.model.enable_sequential_cpu_offload(gpu_id=0)
+ else:
+ self.model = self.model.to(device)
+
+ def forward(self, image, mask, config: InpaintRequest):
+ """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
+ edit = pipe(prompt, image=image, num_inference_steps=20, image_guidance_scale=1.5, guidance_scale=7).images[0]
+ """
+ output = self.model(
+ image=PIL.Image.fromarray(image),
+ prompt=config.prompt,
+ negative_prompt=config.negative_prompt,
+ num_inference_steps=config.sd_steps,
+ image_guidance_scale=config.p2p_image_guidance_scale,
+ guidance_scale=config.sd_guidance_scale,
+ output_type="np",
+ generator=torch.manual_seed(config.sd_seed),
+ ).images[0]
+
+ output = (output * 255).round().astype("uint8")
+ output = cv2.cvtColor(output, cv2.COLOR_RGB2BGR)
+ return output
diff --git a/iopaint/model/kandinsky.py b/iopaint/model/kandinsky.py
new file mode 100644
index 0000000000000000000000000000000000000000..1a0bf1c46e4c648c2b2e78c8b75d8478dc0613a9
--- /dev/null
+++ b/iopaint/model/kandinsky.py
@@ -0,0 +1,65 @@
+import PIL.Image
+import cv2
+import numpy as np
+import torch
+
+from iopaint.const import KANDINSKY22_NAME
+from .base import DiffusionInpaintModel
+from iopaint.schema import InpaintRequest
+from .utils import get_torch_dtype, enable_low_mem, is_local_files_only
+
+
+class Kandinsky(DiffusionInpaintModel):
+ pad_mod = 64
+ min_size = 512
+
+ def init_model(self, device: torch.device, **kwargs):
+ from diffusers import AutoPipelineForInpainting
+
+ use_gpu, torch_dtype = get_torch_dtype(device, kwargs.get("no_half", False))
+
+ model_kwargs = {
+ "torch_dtype": torch_dtype,
+ "local_files_only": is_local_files_only(**kwargs),
+ }
+ self.model = AutoPipelineForInpainting.from_pretrained(
+ self.name, **model_kwargs
+ ).to(device)
+ enable_low_mem(self.model, kwargs.get("low_mem", False))
+
+ self.callback = kwargs.pop("callback", None)
+
+ def forward(self, image, mask, config: InpaintRequest):
+ """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
+ """
+ self.set_scheduler(config)
+
+ generator = torch.manual_seed(config.sd_seed)
+ mask = mask.astype(np.float32) / 255
+ img_h, img_w = image.shape[:2]
+
+ # kandinsky 没有 strength
+ output = self.model(
+ prompt=config.prompt,
+ negative_prompt=config.negative_prompt,
+ image=PIL.Image.fromarray(image),
+ mask_image=mask[:, :, 0],
+ height=img_h,
+ width=img_w,
+ num_inference_steps=config.sd_steps,
+ guidance_scale=config.sd_guidance_scale,
+ output_type="np",
+ callback_on_step_end=self.callback,
+ generator=generator,
+ ).images[0]
+
+ output = (output * 255).round().astype("uint8")
+ output = cv2.cvtColor(output, cv2.COLOR_RGB2BGR)
+ return output
+
+
+class Kandinsky22(Kandinsky):
+ name = KANDINSKY22_NAME
diff --git a/iopaint/model/lama.py b/iopaint/model/lama.py
new file mode 100644
index 0000000000000000000000000000000000000000..7aba242ab93ab0dac285c667365e1c27dce4cbf9
--- /dev/null
+++ b/iopaint/model/lama.py
@@ -0,0 +1,57 @@
+import os
+
+import cv2
+import numpy as np
+import torch
+
+from iopaint.helper import (
+ norm_img,
+ get_cache_path_by_url,
+ load_jit_model,
+ download_model,
+)
+from iopaint.schema import InpaintRequest
+from .base import InpaintModel
+
+LAMA_MODEL_URL = os.environ.get(
+ "LAMA_MODEL_URL",
+ "https://github.com/Sanster/models/releases/download/add_big_lama/big-lama.pt",
+)
+LAMA_MODEL_MD5 = os.environ.get("LAMA_MODEL_MD5", "e3aa4aaa15225a33ec84f9f4bc47e500")
+
+
+class LaMa(InpaintModel):
+ name = "lama"
+ pad_mod = 8
+ is_erase_model = True
+
+ @staticmethod
+ def download():
+ download_model(LAMA_MODEL_URL, LAMA_MODEL_MD5)
+
+ def init_model(self, device, **kwargs):
+ self.model = load_jit_model(LAMA_MODEL_URL, device, LAMA_MODEL_MD5).eval()
+
+ @staticmethod
+ def is_downloaded() -> bool:
+ return os.path.exists(get_cache_path_by_url(LAMA_MODEL_URL))
+
+ def forward(self, image, mask, config: InpaintRequest):
+ """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
diff --git a/iopaint/model/ldm.py b/iopaint/model/ldm.py
new file mode 100644
index 0000000000000000000000000000000000000000..19e51a3ef4f2c27b5cca5398820dbaf5e81888a1
--- /dev/null
+++ b/iopaint/model/ldm.py
@@ -0,0 +1,336 @@
+import os
+
+import numpy as np
+import torch
+from loguru import logger
+
+from .base import InpaintModel
+from .ddim_sampler import DDIMSampler
+from .plms_sampler import PLMSSampler
+from iopaint.schema import InpaintRequest, LDMSampler
+
+torch.manual_seed(42)
+import torch.nn as nn
+from iopaint.helper import (
+ download_model,
+ norm_img,
+ get_cache_path_by_url,
+ load_jit_model,
+)
+from .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_ENCODE_MODEL_MD5 = os.environ.get(
+ "LDM_ENCODE_MODEL_MD5", "23239fc9081956a3e70de56472b3f296"
+)
+
+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_DECODE_MODEL_MD5 = os.environ.get(
+ "LDM_DECODE_MODEL_MD5", "fe419cd15a750d37a4733589d0d3585c"
+)
+
+LDM_DIFFUSION_MODEL_URL = os.environ.get(
+ "LDM_DIFFUSION_MODEL_URL",
+ "https://github.com/Sanster/models/releases/download/add_ldm/diffusion.pt",
+)
+
+LDM_DIFFUSION_MODEL_MD5 = os.environ.get(
+ "LDM_DIFFUSION_MODEL_MD5", "b0afda12bf790c03aba2a7431f11d22d"
+)
+
+
+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):
+ name = "ldm"
+ pad_mod = 32
+ is_erase_model = True
+
+ 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, LDM_DIFFUSION_MODEL_MD5
+ )
+ self.cond_stage_model_decode = load_jit_model(
+ LDM_DECODE_MODEL_URL, device, LDM_DECODE_MODEL_MD5
+ )
+ self.cond_stage_model_encode = load_jit_model(
+ LDM_ENCODE_MODEL_URL, device, LDM_ENCODE_MODEL_MD5
+ )
+ 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 download():
+ download_model(LDM_DIFFUSION_MODEL_URL, LDM_DIFFUSION_MODEL_MD5)
+ download_model(LDM_DECODE_MODEL_URL, LDM_DECODE_MODEL_MD5)
+ download_model(LDM_ENCODE_MODEL_URL, LDM_ENCODE_MODEL_MD5)
+
+ @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: InpaintRequest):
+ """
+ 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
diff --git a/iopaint/model/manga.py b/iopaint/model/manga.py
new file mode 100644
index 0000000000000000000000000000000000000000..1f58251891bdda54a0e40f373000d18bbea7014b
--- /dev/null
+++ b/iopaint/model/manga.py
@@ -0,0 +1,97 @@
+import os
+import random
+
+import cv2
+import numpy as np
+import torch
+import time
+from loguru import logger
+
+from iopaint.helper import get_cache_path_by_url, load_jit_model, download_model
+from .base import InpaintModel
+from iopaint.schema import InpaintRequest
+
+
+MANGA_INPAINTOR_MODEL_URL = os.environ.get(
+ "MANGA_INPAINTOR_MODEL_URL",
+ "https://github.com/Sanster/models/releases/download/manga/manga_inpaintor.jit",
+)
+MANGA_INPAINTOR_MODEL_MD5 = os.environ.get(
+ "MANGA_INPAINTOR_MODEL_MD5", "7d8b269c4613b6b3768af714610da86c"
+)
+
+MANGA_LINE_MODEL_URL = os.environ.get(
+ "MANGA_LINE_MODEL_URL",
+ "https://github.com/Sanster/models/releases/download/manga/erika.jit",
+)
+MANGA_LINE_MODEL_MD5 = os.environ.get(
+ "MANGA_LINE_MODEL_MD5", "0c926d5a4af8450b0d00bc5b9a095644"
+)
+
+
+class Manga(InpaintModel):
+ name = "manga"
+ pad_mod = 16
+ is_erase_model = True
+
+ def init_model(self, device, **kwargs):
+ self.inpaintor_model = load_jit_model(
+ MANGA_INPAINTOR_MODEL_URL, device, MANGA_INPAINTOR_MODEL_MD5
+ )
+ self.line_model = load_jit_model(
+ MANGA_LINE_MODEL_URL, device, MANGA_LINE_MODEL_MD5
+ )
+ self.seed = 42
+
+ @staticmethod
+ def download():
+ download_model(MANGA_INPAINTOR_MODEL_URL, MANGA_INPAINTOR_MODEL_MD5)
+ download_model(MANGA_LINE_MODEL_URL, MANGA_LINE_MODEL_MD5)
+
+ @staticmethod
+ def is_downloaded() -> bool:
+ model_paths = [
+ get_cache_path_by_url(MANGA_INPAINTOR_MODEL_URL),
+ get_cache_path_by_url(MANGA_LINE_MODEL_URL),
+ ]
+ return all([os.path.exists(it) for it in model_paths])
+
+ def forward(self, image, mask, config: InpaintRequest):
+ """
+ image: [H, W, C] RGB
+ mask: [H, W, 1]
+ return: BGR IMAGE
+ """
+ seed = self.seed
+ random.seed(seed)
+ np.random.seed(seed)
+ torch.manual_seed(seed)
+ torch.cuda.manual_seed_all(seed)
+
+ gray_img = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY)
+ gray_img = torch.from_numpy(
+ gray_img[np.newaxis, np.newaxis, :, :].astype(np.float32)
+ ).to(self.device)
+ start = time.time()
+ lines = self.line_model(gray_img)
+ torch.cuda.empty_cache()
+ lines = torch.clamp(lines, 0, 255)
+ logger.info(f"erika_model time: {time.time() - start}")
+
+ mask = torch.from_numpy(mask[np.newaxis, :, :, :]).to(self.device)
+ mask = mask.permute(0, 3, 1, 2)
+ mask = torch.where(mask > 0.5, 1.0, 0.0)
+ noise = torch.randn_like(mask)
+ ones = torch.ones_like(mask)
+
+ gray_img = gray_img / 255 * 2 - 1.0
+ lines = lines / 255 * 2 - 1.0
+
+ start = time.time()
+ inpainted_image = self.inpaintor_model(gray_img, lines, mask, noise, ones)
+ logger.info(f"image_inpaintor_model time: {time.time() - start}")
+
+ cur_res = inpainted_image[0].permute(1, 2, 0).detach().cpu().numpy()
+ cur_res = (cur_res * 127.5 + 127.5).astype(np.uint8)
+ cur_res = cv2.cvtColor(cur_res, cv2.COLOR_GRAY2BGR)
+ return cur_res
diff --git a/iopaint/model/mat.py b/iopaint/model/mat.py
new file mode 100644
index 0000000000000000000000000000000000000000..0c5360faca1326413b4f0329799537176d25fa00
--- /dev/null
+++ b/iopaint/model/mat.py
@@ -0,0 +1,1945 @@
+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 iopaint.helper import (
+ load_model,
+ get_cache_path_by_url,
+ norm_img,
+ download_model,
+)
+from iopaint.schema import InpaintRequest
+from .base import InpaintModel
+from .utils import (
+ setup_filter,
+ Conv2dLayer,
+ FullyConnectedLayer,
+ conv2d_resample,
+ bias_act,
+ upsample2d,
+ activation_funcs,
+ MinibatchStdLayer,
+ to_2tuple,
+ normalize_2nd_moment,
+ set_seed,
+)
+
+
+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.
+ torch_dtype=torch.float32,
+ ):
+ 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
+ self.torch_dtype = torch_dtype
+
+ 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
+ if self.z_dim > 0:
+ x = normalize_2nd_moment(z)
+ if self.c_dim > 0:
+ y = normalize_2nd_moment(self.embed(c))
+ 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:
+ 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:
+ x = x.unsqueeze(1).repeat([1, self.num_ws, 1])
+
+ # Apply truncation.
+ if truncation_psi != 1:
+ 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.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.to(torch.float32) + 1e-8)
+ update_mask = torch.clamp(update_mask, 0, 1) # 0 or 1
+ mask_ratio = torch.mul(mask_ratio, update_mask).to(x.dtype)
+ 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.0,
+ proj_drop=0.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, eps=torch.finfo(x.dtype).eps)
+ 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.0,
+ qkv_bias=True,
+ qk_scale=None,
+ drop=0.0,
+ attn_drop=0.0,
+ drop_path=0.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.dtype).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.0,
+ qkv_bias=True,
+ qk_scale=None,
+ drop=0.0,
+ attn_drop=0.0,
+ drop_path=0.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",
+)
+
+MAT_MODEL_MD5 = os.environ.get("MAT_MODEL_MD5", "8ca927835fa3f5e21d65ffcb165377ed")
+
+
+class MAT(InpaintModel):
+ name = "mat"
+ min_size = 512
+ pad_mod = 512
+ pad_to_square = True
+ is_erase_model = True
+
+ def init_model(self, device, **kwargs):
+ seed = 240 # pick up a random number
+ set_seed(seed)
+
+ fp16 = not kwargs.get("no_half", False)
+ use_gpu = "cuda" in str(device) and torch.cuda.is_available()
+ self.torch_dtype = torch.float16 if use_gpu and fp16 else torch.float32
+
+ G = Generator(
+ z_dim=512,
+ c_dim=0,
+ w_dim=512,
+ img_resolution=512,
+ img_channels=3,
+ mapping_kwargs={"torch_dtype": self.torch_dtype},
+ ).to(self.torch_dtype)
+ # fmt: off
+ self.model = load_model(G, MAT_MODEL_URL, device, MAT_MODEL_MD5)
+ self.z = torch.from_numpy(np.random.randn(1, G.z_dim)).to(self.torch_dtype).to(device)
+ self.label = torch.zeros([1, self.model.c_dim], device=device).to(self.torch_dtype)
+ # fmt: on
+
+ @staticmethod
+ def download():
+ download_model(MAT_MODEL_URL, MAT_MODEL_MD5)
+
+ @staticmethod
+ def is_downloaded() -> bool:
+ return os.path.exists(get_cache_path_by_url(MAT_MODEL_URL))
+
+ def forward(self, image, mask, config: InpaintRequest):
+ """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.torch_dtype).to(self.device)
+ )
+ mask = torch.from_numpy(mask).unsqueeze(0).to(self.torch_dtype).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
diff --git a/iopaint/model/mi_gan.py b/iopaint/model/mi_gan.py
new file mode 100644
index 0000000000000000000000000000000000000000..f1ce25ff145d793f3f96dee2a58593ddd6cb34b5
--- /dev/null
+++ b/iopaint/model/mi_gan.py
@@ -0,0 +1,110 @@
+import os
+
+import cv2
+import torch
+
+from iopaint.helper import (
+ load_jit_model,
+ download_model,
+ get_cache_path_by_url,
+ boxes_from_mask,
+ resize_max_size,
+ norm_img,
+)
+from .base import InpaintModel
+from iopaint.schema import InpaintRequest
+
+MIGAN_MODEL_URL = os.environ.get(
+ "MIGAN_MODEL_URL",
+ "https://github.com/Sanster/models/releases/download/migan/migan_traced.pt",
+)
+MIGAN_MODEL_MD5 = os.environ.get("MIGAN_MODEL_MD5", "76eb3b1a71c400ee3290524f7a11b89c")
+
+
+class MIGAN(InpaintModel):
+ name = "migan"
+ min_size = 512
+ pad_mod = 512
+ pad_to_square = True
+ is_erase_model = True
+
+ def init_model(self, device, **kwargs):
+ self.model = load_jit_model(MIGAN_MODEL_URL, device, MIGAN_MODEL_MD5).eval()
+
+ @staticmethod
+ def download():
+ download_model(MIGAN_MODEL_URL, MIGAN_MODEL_MD5)
+
+ @staticmethod
+ def is_downloaded() -> bool:
+ return os.path.exists(get_cache_path_by_url(MIGAN_MODEL_URL))
+
+ @torch.no_grad()
+ def __call__(self, image, mask, config: InpaintRequest):
+ """
+ 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].copy()
+ 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: InpaintRequest):
+ """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)
+ 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
diff --git a/iopaint/model/opencv2.py b/iopaint/model/opencv2.py
new file mode 100644
index 0000000000000000000000000000000000000000..de472094417b3eba910611d34327612aad9ff5cb
--- /dev/null
+++ b/iopaint/model/opencv2.py
@@ -0,0 +1,29 @@
+import cv2
+from .base import InpaintModel
+from iopaint.schema import InpaintRequest
+
+flag_map = {"INPAINT_NS": cv2.INPAINT_NS, "INPAINT_TELEA": cv2.INPAINT_TELEA}
+
+
+class OpenCV2(InpaintModel):
+ name = "cv2"
+ pad_mod = 1
+ is_erase_model = True
+
+ @staticmethod
+ def is_downloaded() -> bool:
+ return True
+
+ def forward(self, image, mask, config: InpaintRequest):
+ """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
diff --git a/iopaint/model_manager.py b/iopaint/model_manager.py
new file mode 100644
index 0000000000000000000000000000000000000000..de13612acfdfdcc823b2eeeb7eb99b4589f31333
--- /dev/null
+++ b/iopaint/model_manager.py
@@ -0,0 +1,191 @@
+from typing import List, Dict
+
+import torch
+from loguru import logger
+import numpy as np
+
+from iopaint.download import scan_models
+from iopaint.helper import switch_mps_device
+from iopaint.model import models, ControlNet, SD, SDXL
+from iopaint.model.utils import torch_gc, is_local_files_only
+from iopaint.schema import InpaintRequest, ModelInfo, ModelType
+
+
+class ModelManager:
+ def __init__(self, name: str, device: torch.device, **kwargs):
+ self.name = name
+ self.device = device
+ self.kwargs = kwargs
+ self.available_models: Dict[str, ModelInfo] = {}
+ self.scan_models()
+
+ self.enable_controlnet = kwargs.get("enable_controlnet", False)
+ controlnet_method = kwargs.get("controlnet_method", None)
+ if (
+ controlnet_method is None
+ and name in self.available_models
+ and self.available_models[name].support_controlnet
+ ):
+ controlnet_method = self.available_models[name].controlnets[0]
+ self.controlnet_method = controlnet_method
+ self.model = self.init_model(name, device, **kwargs)
+
+ @property
+ def current_model(self) -> ModelInfo:
+ return self.available_models[self.name]
+
+ def init_model(self, name: str, device, **kwargs):
+ logger.info(f"Loading model: {name}")
+ if name not in self.available_models:
+ raise NotImplementedError(
+ f"Unsupported model: {name}. Available models: {list(self.available_models.keys())}"
+ )
+
+ model_info = self.available_models[name]
+ kwargs = {
+ **kwargs,
+ "model_info": model_info,
+ "enable_controlnet": self.enable_controlnet,
+ "controlnet_method": self.controlnet_method,
+ }
+
+ if model_info.support_controlnet and self.enable_controlnet:
+ return ControlNet(device, **kwargs)
+ elif model_info.name in models:
+ return models[name](device, **kwargs)
+ else:
+ if model_info.model_type in [
+ ModelType.DIFFUSERS_SD_INPAINT,
+ ModelType.DIFFUSERS_SD,
+ ]:
+ return SD(device, **kwargs)
+
+ if model_info.model_type in [
+ ModelType.DIFFUSERS_SDXL_INPAINT,
+ ModelType.DIFFUSERS_SDXL,
+ ]:
+ return SDXL(device, **kwargs)
+
+ raise NotImplementedError(f"Unsupported model: {name}")
+
+ @torch.inference_mode()
+ def __call__(self, image, mask, config: InpaintRequest):
+ """
+
+ Args:
+ image: [H, W, C] RGB
+ mask: [H, W, 1] 255 means area to repaint
+ config:
+
+ Returns:
+ BGR image
+ """
+ self.switch_controlnet_method(config)
+ self.enable_disable_freeu(config)
+ self.enable_disable_lcm_lora(config)
+ return self.model(image, mask, config).astype(np.uint8)
+
+ def scan_models(self) -> List[ModelInfo]:
+ available_models = scan_models()
+ self.available_models = {it.name: it for it in available_models}
+ return available_models
+
+ def switch(self, new_name: str):
+ if new_name == self.name:
+ return
+
+ old_name = self.name
+ old_controlnet_method = self.controlnet_method
+ self.name = new_name
+
+ if (
+ self.available_models[new_name].support_controlnet
+ and self.controlnet_method
+ not in self.available_models[new_name].controlnets
+ ):
+ self.controlnet_method = self.available_models[new_name].controlnets[0]
+ try:
+ # TODO: enable/disable controlnet without reload model
+ del self.model
+ torch_gc()
+
+ self.model = self.init_model(
+ new_name, switch_mps_device(new_name, self.device), **self.kwargs
+ )
+ except Exception as e:
+ self.name = old_name
+ self.controlnet_method = old_controlnet_method
+ logger.info(f"Switch model from {old_name} to {new_name} failed, rollback")
+ self.model = self.init_model(
+ old_name, switch_mps_device(old_name, self.device), **self.kwargs
+ )
+ raise e
+
+ def switch_controlnet_method(self, config):
+ if not self.available_models[self.name].support_controlnet:
+ return
+
+ if (
+ self.enable_controlnet
+ and config.controlnet_method
+ and self.controlnet_method != config.controlnet_method
+ ):
+ old_controlnet_method = self.controlnet_method
+ self.controlnet_method = config.controlnet_method
+ self.model.switch_controlnet_method(config.controlnet_method)
+ logger.info(
+ f"Switch Controlnet method from {old_controlnet_method} to {config.controlnet_method}"
+ )
+ elif self.enable_controlnet != config.enable_controlnet:
+ self.enable_controlnet = config.enable_controlnet
+ self.controlnet_method = config.controlnet_method
+
+ pipe_components = {
+ "vae": self.model.model.vae,
+ "text_encoder": self.model.model.text_encoder,
+ "unet": self.model.model.unet,
+ }
+ if hasattr(self.model.model, "text_encoder_2"):
+ pipe_components["text_encoder_2"] = self.model.model.text_encoder_2
+
+ self.model = self.init_model(
+ self.name,
+ switch_mps_device(self.name, self.device),
+ pipe_components=pipe_components,
+ **self.kwargs,
+ )
+ if not config.enable_controlnet:
+ logger.info(f"Disable controlnet")
+ else:
+ logger.info(f"Enable controlnet: {config.controlnet_method}")
+
+ def enable_disable_freeu(self, config: InpaintRequest):
+ if str(self.model.device) == "mps":
+ return
+
+ if self.available_models[self.name].support_freeu:
+ if config.sd_freeu:
+ freeu_config = config.sd_freeu_config
+ self.model.model.enable_freeu(
+ s1=freeu_config.s1,
+ s2=freeu_config.s2,
+ b1=freeu_config.b1,
+ b2=freeu_config.b2,
+ )
+ else:
+ self.model.model.disable_freeu()
+
+ def enable_disable_lcm_lora(self, config: InpaintRequest):
+ if self.available_models[self.name].support_lcm_lora:
+ # TODO: change this if load other lora is supported
+ lcm_lora_loaded = bool(self.model.model.get_list_adapters())
+ if config.sd_lcm_lora:
+ if not lcm_lora_loaded:
+ self.model.model.load_lora_weights(
+ self.model.lcm_lora_id,
+ weight_name="pytorch_lora_weights.safetensors",
+ local_files_only=is_local_files_only(),
+ )
+ else:
+ if lcm_lora_loaded:
+ self.model.model.disable_lora()
diff --git a/iopaint/plugins/briarmbg.py b/iopaint/plugins/briarmbg.py
new file mode 100644
index 0000000000000000000000000000000000000000..880f5305695613058c3f2b69c2a0c6dcda3af3ac
--- /dev/null
+++ b/iopaint/plugins/briarmbg.py
@@ -0,0 +1,512 @@
+# copy from: https://huggingface.co/spaces/briaai/BRIA-RMBG-1.4/blob/main/briarmbg.py
+import cv2
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from PIL import Image
+import numpy as np
+from torchvision.transforms.functional import normalize
+
+
+class REBNCONV(nn.Module):
+ def __init__(self, in_ch=3, out_ch=3, dirate=1, stride=1):
+ super(REBNCONV, self).__init__()
+
+ self.conv_s1 = nn.Conv2d(
+ in_ch, out_ch, 3, padding=1 * dirate, dilation=1 * dirate, stride=stride
+ )
+ self.bn_s1 = nn.BatchNorm2d(out_ch)
+ self.relu_s1 = nn.ReLU(inplace=True)
+
+ def forward(self, x):
+ hx = x
+ xout = self.relu_s1(self.bn_s1(self.conv_s1(hx)))
+
+ return xout
+
+
+## upsample tensor 'src' to have the same spatial size with tensor 'tar'
+def _upsample_like(src, tar):
+ src = F.interpolate(src, size=tar.shape[2:], mode="bilinear")
+
+ return src
+
+
+### RSU-7 ###
+class RSU7(nn.Module):
+ def __init__(self, in_ch=3, mid_ch=12, out_ch=3, img_size=512):
+ super(RSU7, self).__init__()
+
+ self.in_ch = in_ch
+ self.mid_ch = mid_ch
+ self.out_ch = out_ch
+
+ self.rebnconvin = REBNCONV(in_ch, out_ch, dirate=1) ## 1 -> 1/2
+
+ self.rebnconv1 = REBNCONV(out_ch, mid_ch, dirate=1)
+ self.pool1 = nn.MaxPool2d(2, stride=2, ceil_mode=True)
+
+ self.rebnconv2 = REBNCONV(mid_ch, mid_ch, dirate=1)
+ self.pool2 = nn.MaxPool2d(2, stride=2, ceil_mode=True)
+
+ self.rebnconv3 = REBNCONV(mid_ch, mid_ch, dirate=1)
+ self.pool3 = nn.MaxPool2d(2, stride=2, ceil_mode=True)
+
+ self.rebnconv4 = REBNCONV(mid_ch, mid_ch, dirate=1)
+ self.pool4 = nn.MaxPool2d(2, stride=2, ceil_mode=True)
+
+ self.rebnconv5 = REBNCONV(mid_ch, mid_ch, dirate=1)
+ self.pool5 = nn.MaxPool2d(2, stride=2, ceil_mode=True)
+
+ self.rebnconv6 = REBNCONV(mid_ch, mid_ch, dirate=1)
+
+ self.rebnconv7 = REBNCONV(mid_ch, mid_ch, dirate=2)
+
+ self.rebnconv6d = REBNCONV(mid_ch * 2, mid_ch, dirate=1)
+ self.rebnconv5d = REBNCONV(mid_ch * 2, mid_ch, dirate=1)
+ self.rebnconv4d = REBNCONV(mid_ch * 2, mid_ch, dirate=1)
+ self.rebnconv3d = REBNCONV(mid_ch * 2, mid_ch, dirate=1)
+ self.rebnconv2d = REBNCONV(mid_ch * 2, mid_ch, dirate=1)
+ self.rebnconv1d = REBNCONV(mid_ch * 2, out_ch, dirate=1)
+
+ def forward(self, x):
+ b, c, h, w = x.shape
+
+ hx = x
+ hxin = self.rebnconvin(hx)
+
+ hx1 = self.rebnconv1(hxin)
+ hx = self.pool1(hx1)
+
+ hx2 = self.rebnconv2(hx)
+ hx = self.pool2(hx2)
+
+ hx3 = self.rebnconv3(hx)
+ hx = self.pool3(hx3)
+
+ hx4 = self.rebnconv4(hx)
+ hx = self.pool4(hx4)
+
+ hx5 = self.rebnconv5(hx)
+ hx = self.pool5(hx5)
+
+ hx6 = self.rebnconv6(hx)
+
+ hx7 = self.rebnconv7(hx6)
+
+ hx6d = self.rebnconv6d(torch.cat((hx7, hx6), 1))
+ hx6dup = _upsample_like(hx6d, hx5)
+
+ hx5d = self.rebnconv5d(torch.cat((hx6dup, hx5), 1))
+ hx5dup = _upsample_like(hx5d, hx4)
+
+ hx4d = self.rebnconv4d(torch.cat((hx5dup, hx4), 1))
+ hx4dup = _upsample_like(hx4d, hx3)
+
+ hx3d = self.rebnconv3d(torch.cat((hx4dup, hx3), 1))
+ hx3dup = _upsample_like(hx3d, hx2)
+
+ hx2d = self.rebnconv2d(torch.cat((hx3dup, hx2), 1))
+ hx2dup = _upsample_like(hx2d, hx1)
+
+ hx1d = self.rebnconv1d(torch.cat((hx2dup, hx1), 1))
+
+ return hx1d + hxin
+
+
+### RSU-6 ###
+class RSU6(nn.Module):
+ def __init__(self, in_ch=3, mid_ch=12, out_ch=3):
+ super(RSU6, self).__init__()
+
+ self.rebnconvin = REBNCONV(in_ch, out_ch, dirate=1)
+
+ self.rebnconv1 = REBNCONV(out_ch, mid_ch, dirate=1)
+ self.pool1 = nn.MaxPool2d(2, stride=2, ceil_mode=True)
+
+ self.rebnconv2 = REBNCONV(mid_ch, mid_ch, dirate=1)
+ self.pool2 = nn.MaxPool2d(2, stride=2, ceil_mode=True)
+
+ self.rebnconv3 = REBNCONV(mid_ch, mid_ch, dirate=1)
+ self.pool3 = nn.MaxPool2d(2, stride=2, ceil_mode=True)
+
+ self.rebnconv4 = REBNCONV(mid_ch, mid_ch, dirate=1)
+ self.pool4 = nn.MaxPool2d(2, stride=2, ceil_mode=True)
+
+ self.rebnconv5 = REBNCONV(mid_ch, mid_ch, dirate=1)
+
+ self.rebnconv6 = REBNCONV(mid_ch, mid_ch, dirate=2)
+
+ self.rebnconv5d = REBNCONV(mid_ch * 2, mid_ch, dirate=1)
+ self.rebnconv4d = REBNCONV(mid_ch * 2, mid_ch, dirate=1)
+ self.rebnconv3d = REBNCONV(mid_ch * 2, mid_ch, dirate=1)
+ self.rebnconv2d = REBNCONV(mid_ch * 2, mid_ch, dirate=1)
+ self.rebnconv1d = REBNCONV(mid_ch * 2, out_ch, dirate=1)
+
+ def forward(self, x):
+ hx = x
+
+ hxin = self.rebnconvin(hx)
+
+ hx1 = self.rebnconv1(hxin)
+ hx = self.pool1(hx1)
+
+ hx2 = self.rebnconv2(hx)
+ hx = self.pool2(hx2)
+
+ hx3 = self.rebnconv3(hx)
+ hx = self.pool3(hx3)
+
+ hx4 = self.rebnconv4(hx)
+ hx = self.pool4(hx4)
+
+ hx5 = self.rebnconv5(hx)
+
+ hx6 = self.rebnconv6(hx5)
+
+ hx5d = self.rebnconv5d(torch.cat((hx6, hx5), 1))
+ hx5dup = _upsample_like(hx5d, hx4)
+
+ hx4d = self.rebnconv4d(torch.cat((hx5dup, hx4), 1))
+ hx4dup = _upsample_like(hx4d, hx3)
+
+ hx3d = self.rebnconv3d(torch.cat((hx4dup, hx3), 1))
+ hx3dup = _upsample_like(hx3d, hx2)
+
+ hx2d = self.rebnconv2d(torch.cat((hx3dup, hx2), 1))
+ hx2dup = _upsample_like(hx2d, hx1)
+
+ hx1d = self.rebnconv1d(torch.cat((hx2dup, hx1), 1))
+
+ return hx1d + hxin
+
+
+### RSU-5 ###
+class RSU5(nn.Module):
+ def __init__(self, in_ch=3, mid_ch=12, out_ch=3):
+ super(RSU5, self).__init__()
+
+ self.rebnconvin = REBNCONV(in_ch, out_ch, dirate=1)
+
+ self.rebnconv1 = REBNCONV(out_ch, mid_ch, dirate=1)
+ self.pool1 = nn.MaxPool2d(2, stride=2, ceil_mode=True)
+
+ self.rebnconv2 = REBNCONV(mid_ch, mid_ch, dirate=1)
+ self.pool2 = nn.MaxPool2d(2, stride=2, ceil_mode=True)
+
+ self.rebnconv3 = REBNCONV(mid_ch, mid_ch, dirate=1)
+ self.pool3 = nn.MaxPool2d(2, stride=2, ceil_mode=True)
+
+ self.rebnconv4 = REBNCONV(mid_ch, mid_ch, dirate=1)
+
+ self.rebnconv5 = REBNCONV(mid_ch, mid_ch, dirate=2)
+
+ self.rebnconv4d = REBNCONV(mid_ch * 2, mid_ch, dirate=1)
+ self.rebnconv3d = REBNCONV(mid_ch * 2, mid_ch, dirate=1)
+ self.rebnconv2d = REBNCONV(mid_ch * 2, mid_ch, dirate=1)
+ self.rebnconv1d = REBNCONV(mid_ch * 2, out_ch, dirate=1)
+
+ def forward(self, x):
+ hx = x
+
+ hxin = self.rebnconvin(hx)
+
+ hx1 = self.rebnconv1(hxin)
+ hx = self.pool1(hx1)
+
+ hx2 = self.rebnconv2(hx)
+ hx = self.pool2(hx2)
+
+ hx3 = self.rebnconv3(hx)
+ hx = self.pool3(hx3)
+
+ hx4 = self.rebnconv4(hx)
+
+ hx5 = self.rebnconv5(hx4)
+
+ hx4d = self.rebnconv4d(torch.cat((hx5, hx4), 1))
+ hx4dup = _upsample_like(hx4d, hx3)
+
+ hx3d = self.rebnconv3d(torch.cat((hx4dup, hx3), 1))
+ hx3dup = _upsample_like(hx3d, hx2)
+
+ hx2d = self.rebnconv2d(torch.cat((hx3dup, hx2), 1))
+ hx2dup = _upsample_like(hx2d, hx1)
+
+ hx1d = self.rebnconv1d(torch.cat((hx2dup, hx1), 1))
+
+ return hx1d + hxin
+
+
+### RSU-4 ###
+class RSU4(nn.Module):
+ def __init__(self, in_ch=3, mid_ch=12, out_ch=3):
+ super(RSU4, self).__init__()
+
+ self.rebnconvin = REBNCONV(in_ch, out_ch, dirate=1)
+
+ self.rebnconv1 = REBNCONV(out_ch, mid_ch, dirate=1)
+ self.pool1 = nn.MaxPool2d(2, stride=2, ceil_mode=True)
+
+ self.rebnconv2 = REBNCONV(mid_ch, mid_ch, dirate=1)
+ self.pool2 = nn.MaxPool2d(2, stride=2, ceil_mode=True)
+
+ self.rebnconv3 = REBNCONV(mid_ch, mid_ch, dirate=1)
+
+ self.rebnconv4 = REBNCONV(mid_ch, mid_ch, dirate=2)
+
+ self.rebnconv3d = REBNCONV(mid_ch * 2, mid_ch, dirate=1)
+ self.rebnconv2d = REBNCONV(mid_ch * 2, mid_ch, dirate=1)
+ self.rebnconv1d = REBNCONV(mid_ch * 2, out_ch, dirate=1)
+
+ def forward(self, x):
+ hx = x
+
+ hxin = self.rebnconvin(hx)
+
+ hx1 = self.rebnconv1(hxin)
+ hx = self.pool1(hx1)
+
+ hx2 = self.rebnconv2(hx)
+ hx = self.pool2(hx2)
+
+ hx3 = self.rebnconv3(hx)
+
+ hx4 = self.rebnconv4(hx3)
+
+ hx3d = self.rebnconv3d(torch.cat((hx4, hx3), 1))
+ hx3dup = _upsample_like(hx3d, hx2)
+
+ hx2d = self.rebnconv2d(torch.cat((hx3dup, hx2), 1))
+ hx2dup = _upsample_like(hx2d, hx1)
+
+ hx1d = self.rebnconv1d(torch.cat((hx2dup, hx1), 1))
+
+ return hx1d + hxin
+
+
+### RSU-4F ###
+class RSU4F(nn.Module):
+ def __init__(self, in_ch=3, mid_ch=12, out_ch=3):
+ super(RSU4F, self).__init__()
+
+ self.rebnconvin = REBNCONV(in_ch, out_ch, dirate=1)
+
+ self.rebnconv1 = REBNCONV(out_ch, mid_ch, dirate=1)
+ self.rebnconv2 = REBNCONV(mid_ch, mid_ch, dirate=2)
+ self.rebnconv3 = REBNCONV(mid_ch, mid_ch, dirate=4)
+
+ self.rebnconv4 = REBNCONV(mid_ch, mid_ch, dirate=8)
+
+ self.rebnconv3d = REBNCONV(mid_ch * 2, mid_ch, dirate=4)
+ self.rebnconv2d = REBNCONV(mid_ch * 2, mid_ch, dirate=2)
+ self.rebnconv1d = REBNCONV(mid_ch * 2, out_ch, dirate=1)
+
+ def forward(self, x):
+ hx = x
+
+ hxin = self.rebnconvin(hx)
+
+ hx1 = self.rebnconv1(hxin)
+ hx2 = self.rebnconv2(hx1)
+ hx3 = self.rebnconv3(hx2)
+
+ hx4 = self.rebnconv4(hx3)
+
+ hx3d = self.rebnconv3d(torch.cat((hx4, hx3), 1))
+ hx2d = self.rebnconv2d(torch.cat((hx3d, hx2), 1))
+ hx1d = self.rebnconv1d(torch.cat((hx2d, hx1), 1))
+
+ return hx1d + hxin
+
+
+class myrebnconv(nn.Module):
+ def __init__(
+ self,
+ in_ch=3,
+ out_ch=1,
+ kernel_size=3,
+ stride=1,
+ padding=1,
+ dilation=1,
+ groups=1,
+ ):
+ super(myrebnconv, self).__init__()
+
+ self.conv = nn.Conv2d(
+ in_ch,
+ out_ch,
+ kernel_size=kernel_size,
+ stride=stride,
+ padding=padding,
+ dilation=dilation,
+ groups=groups,
+ )
+ self.bn = nn.BatchNorm2d(out_ch)
+ self.rl = nn.ReLU(inplace=True)
+
+ def forward(self, x):
+ return self.rl(self.bn(self.conv(x)))
+
+
+class BriaRMBG(nn.Module):
+ def __init__(self, in_ch=3, out_ch=1):
+ super(BriaRMBG, self).__init__()
+
+ self.conv_in = nn.Conv2d(in_ch, 64, 3, stride=2, padding=1)
+ self.pool_in = nn.MaxPool2d(2, stride=2, ceil_mode=True)
+
+ self.stage1 = RSU7(64, 32, 64)
+ self.pool12 = nn.MaxPool2d(2, stride=2, ceil_mode=True)
+
+ self.stage2 = RSU6(64, 32, 128)
+ self.pool23 = nn.MaxPool2d(2, stride=2, ceil_mode=True)
+
+ self.stage3 = RSU5(128, 64, 256)
+ self.pool34 = nn.MaxPool2d(2, stride=2, ceil_mode=True)
+
+ self.stage4 = RSU4(256, 128, 512)
+ self.pool45 = nn.MaxPool2d(2, stride=2, ceil_mode=True)
+
+ self.stage5 = RSU4F(512, 256, 512)
+ self.pool56 = nn.MaxPool2d(2, stride=2, ceil_mode=True)
+
+ self.stage6 = RSU4F(512, 256, 512)
+
+ # decoder
+ self.stage5d = RSU4F(1024, 256, 512)
+ self.stage4d = RSU4(1024, 128, 256)
+ self.stage3d = RSU5(512, 64, 128)
+ self.stage2d = RSU6(256, 32, 64)
+ self.stage1d = RSU7(128, 16, 64)
+
+ self.side1 = nn.Conv2d(64, out_ch, 3, padding=1)
+ self.side2 = nn.Conv2d(64, out_ch, 3, padding=1)
+ self.side3 = nn.Conv2d(128, out_ch, 3, padding=1)
+ self.side4 = nn.Conv2d(256, out_ch, 3, padding=1)
+ self.side5 = nn.Conv2d(512, out_ch, 3, padding=1)
+ self.side6 = nn.Conv2d(512, out_ch, 3, padding=1)
+
+ # self.outconv = nn.Conv2d(6*out_ch,out_ch,1)
+
+ def forward(self, x):
+ hx = x
+
+ hxin = self.conv_in(hx)
+ # hx = self.pool_in(hxin)
+
+ # stage 1
+ hx1 = self.stage1(hxin)
+ hx = self.pool12(hx1)
+
+ # stage 2
+ hx2 = self.stage2(hx)
+ hx = self.pool23(hx2)
+
+ # stage 3
+ hx3 = self.stage3(hx)
+ hx = self.pool34(hx3)
+
+ # stage 4
+ hx4 = self.stage4(hx)
+ hx = self.pool45(hx4)
+
+ # stage 5
+ hx5 = self.stage5(hx)
+ hx = self.pool56(hx5)
+
+ # stage 6
+ hx6 = self.stage6(hx)
+ hx6up = _upsample_like(hx6, hx5)
+
+ # -------------------- decoder --------------------
+ hx5d = self.stage5d(torch.cat((hx6up, hx5), 1))
+ hx5dup = _upsample_like(hx5d, hx4)
+
+ hx4d = self.stage4d(torch.cat((hx5dup, hx4), 1))
+ hx4dup = _upsample_like(hx4d, hx3)
+
+ hx3d = self.stage3d(torch.cat((hx4dup, hx3), 1))
+ hx3dup = _upsample_like(hx3d, hx2)
+
+ hx2d = self.stage2d(torch.cat((hx3dup, hx2), 1))
+ hx2dup = _upsample_like(hx2d, hx1)
+
+ hx1d = self.stage1d(torch.cat((hx2dup, hx1), 1))
+
+ # side output
+ d1 = self.side1(hx1d)
+ d1 = _upsample_like(d1, x)
+
+ d2 = self.side2(hx2d)
+ d2 = _upsample_like(d2, x)
+
+ d3 = self.side3(hx3d)
+ d3 = _upsample_like(d3, x)
+
+ d4 = self.side4(hx4d)
+ d4 = _upsample_like(d4, x)
+
+ d5 = self.side5(hx5d)
+ d5 = _upsample_like(d5, x)
+
+ d6 = self.side6(hx6)
+ d6 = _upsample_like(d6, x)
+
+ return [
+ F.sigmoid(d1),
+ F.sigmoid(d2),
+ F.sigmoid(d3),
+ F.sigmoid(d4),
+ F.sigmoid(d5),
+ F.sigmoid(d6),
+ ], [hx1d, hx2d, hx3d, hx4d, hx5d, hx6]
+
+
+def resize_image(image):
+ image = image.convert("RGB")
+ model_input_size = (1024, 1024)
+ image = image.resize(model_input_size, Image.BILINEAR)
+ return image
+
+
+def create_briarmbg_session():
+ from huggingface_hub import hf_hub_download
+
+ net = BriaRMBG()
+ model_path = hf_hub_download("briaai/RMBG-1.4", "model.pth")
+ net.load_state_dict(torch.load(model_path, map_location="cpu"))
+ net.eval()
+ return net
+
+
+def briarmbg_process(bgr_np_image, session, only_mask=False):
+ # prepare input
+ orig_bgr_image = Image.fromarray(bgr_np_image)
+ w, h = orig_im_size = orig_bgr_image.size
+ image = resize_image(orig_bgr_image)
+ im_np = np.array(image)
+ im_tensor = torch.tensor(im_np, dtype=torch.float32).permute(2, 0, 1)
+ im_tensor = torch.unsqueeze(im_tensor, 0)
+ im_tensor = torch.divide(im_tensor, 255.0)
+ im_tensor = normalize(im_tensor, [0.5, 0.5, 0.5], [1.0, 1.0, 1.0])
+ # inference
+ result = session(im_tensor)
+ # post process
+ result = torch.squeeze(F.interpolate(result[0][0], size=(h, w), mode="bilinear"), 0)
+ ma = torch.max(result)
+ mi = torch.min(result)
+ result = (result - mi) / (ma - mi)
+ # image to pil
+ im_array = (result * 255).cpu().data.numpy().astype(np.uint8)
+
+ mask = np.squeeze(im_array)
+ if only_mask:
+ return mask
+
+ pil_im = Image.fromarray(mask)
+ # paste the mask on the original image
+ new_im = Image.new("RGBA", pil_im.size, (0, 0, 0, 0))
+ new_im.paste(orig_bgr_image, mask=pil_im)
+ rgba_np_img = np.asarray(new_im)
+ return rgba_np_img
diff --git a/iopaint/plugins/gfpgan_plugin.py b/iopaint/plugins/gfpgan_plugin.py
new file mode 100644
index 0000000000000000000000000000000000000000..619280b0bd3f06f89d9f3d00ce10fe56bed0d47b
--- /dev/null
+++ b/iopaint/plugins/gfpgan_plugin.py
@@ -0,0 +1,74 @@
+import cv2
+import numpy as np
+from loguru import logger
+
+from iopaint.helper import download_model
+from iopaint.plugins.base_plugin import BasePlugin
+from iopaint.schema import RunPluginRequest
+
+
+class GFPGANPlugin(BasePlugin):
+ name = "GFPGAN"
+ support_gen_image = True
+
+ def __init__(self, device, upscaler=None):
+ super().__init__()
+ from .gfpganer import MyGFPGANer
+
+ url = "https://github.com/TencentARC/GFPGAN/releases/download/v1.3.0/GFPGANv1.4.pth"
+ model_md5 = "94d735072630ab734561130a47bc44f8"
+ model_path = download_model(url, model_md5)
+ logger.info(f"GFPGAN model path: {model_path}")
+
+ import facexlib
+
+ if hasattr(facexlib.detection.retinaface, "device"):
+ facexlib.detection.retinaface.device = device
+
+ # Use GFPGAN for face enhancement
+ self.face_enhancer = MyGFPGANer(
+ model_path=model_path,
+ upscale=1,
+ arch="clean",
+ channel_multiplier=2,
+ device=device,
+ bg_upsampler=upscaler.model if upscaler is not None else None,
+ )
+ self.face_enhancer.face_helper.face_det.mean_tensor.to(device)
+ self.face_enhancer.face_helper.face_det = (
+ self.face_enhancer.face_helper.face_det.to(device)
+ )
+
+ def gen_image(self, rgb_np_img, req: RunPluginRequest) -> np.ndarray:
+ weight = 0.5
+ bgr_np_img = cv2.cvtColor(rgb_np_img, cv2.COLOR_RGB2BGR)
+ logger.info(f"GFPGAN input shape: {bgr_np_img.shape}")
+ _, _, bgr_output = self.face_enhancer.enhance(
+ bgr_np_img,
+ has_aligned=False,
+ only_center_face=False,
+ paste_back=True,
+ weight=weight,
+ )
+ logger.info(f"GFPGAN output shape: {bgr_output.shape}")
+
+ # try:
+ # if scale != 2:
+ # interpolation = cv2.INTER_AREA if scale < 2 else cv2.INTER_LANCZOS4
+ # h, w = img.shape[0:2]
+ # output = cv2.resize(
+ # output,
+ # (int(w * scale / 2), int(h * scale / 2)),
+ # interpolation=interpolation,
+ # )
+ # except Exception as error:
+ # print("wrong scale input.", error)
+ return bgr_output
+
+ def check_dep(self):
+ try:
+ import gfpgan
+ except ImportError:
+ return (
+ "gfpgan is not installed, please install it first. pip install gfpgan"
+ )
diff --git a/iopaint/plugins/gfpganer.py b/iopaint/plugins/gfpganer.py
new file mode 100644
index 0000000000000000000000000000000000000000..75a575dbd182193f7a62badae139762322fb55a5
--- /dev/null
+++ b/iopaint/plugins/gfpganer.py
@@ -0,0 +1,84 @@
+import os
+
+import torch
+from facexlib.utils.face_restoration_helper import FaceRestoreHelper
+from gfpgan import GFPGANv1Clean, GFPGANer
+from torch.hub import get_dir
+
+
+class MyGFPGANer(GFPGANer):
+ """Helper for restoration with GFPGAN.
+
+ It will detect and crop faces, and then resize the faces to 512x512.
+ GFPGAN is used to restored the resized faces.
+ The background is upsampled with the bg_upsampler.
+ Finally, the faces will be pasted back to the upsample background image.
+
+ Args:
+ model_path (str): The path to the GFPGAN model. It can be urls (will first download it automatically).
+ upscale (float): The upscale of the final output. Default: 2.
+ arch (str): The GFPGAN architecture. Option: clean | original. Default: clean.
+ channel_multiplier (int): Channel multiplier for large networks of StyleGAN2. Default: 2.
+ bg_upsampler (nn.Module): The upsampler for the background. Default: None.
+ """
+
+ def __init__(
+ self,
+ model_path,
+ upscale=2,
+ arch="clean",
+ channel_multiplier=2,
+ bg_upsampler=None,
+ device=None,
+ ):
+ self.upscale = upscale
+ self.bg_upsampler = bg_upsampler
+
+ # initialize model
+ self.device = (
+ torch.device("cuda" if torch.cuda.is_available() else "cpu")
+ if device is None
+ else device
+ )
+ # initialize the GFP-GAN
+ if arch == "clean":
+ self.gfpgan = GFPGANv1Clean(
+ out_size=512,
+ num_style_feat=512,
+ channel_multiplier=channel_multiplier,
+ decoder_load_path=None,
+ fix_decoder=False,
+ num_mlp=8,
+ input_is_latent=True,
+ different_w=True,
+ narrow=1,
+ sft_half=True,
+ )
+ elif arch == "RestoreFormer":
+ from gfpgan.archs.restoreformer_arch import RestoreFormer
+
+ self.gfpgan = RestoreFormer()
+
+ hub_dir = get_dir()
+ model_dir = os.path.join(hub_dir, "checkpoints")
+
+ # initialize face helper
+ self.face_helper = FaceRestoreHelper(
+ upscale,
+ face_size=512,
+ crop_ratio=(1, 1),
+ det_model="retinaface_resnet50",
+ save_ext="png",
+ use_parse=True,
+ device=self.device,
+ model_rootpath=model_dir,
+ )
+
+ loadnet = torch.load(model_path)
+ if "params_ema" in loadnet:
+ keyname = "params_ema"
+ else:
+ keyname = "params"
+ self.gfpgan.load_state_dict(loadnet[keyname], strict=True)
+ self.gfpgan.eval()
+ self.gfpgan = self.gfpgan.to(self.device)
diff --git a/iopaint/plugins/interactive_seg.py b/iopaint/plugins/interactive_seg.py
new file mode 100644
index 0000000000000000000000000000000000000000..f19a3a89ac194b7929e4bd704931ade602453495
--- /dev/null
+++ b/iopaint/plugins/interactive_seg.py
@@ -0,0 +1,89 @@
+import hashlib
+from typing import List
+
+import numpy as np
+import torch
+from loguru import logger
+
+from iopaint.helper import download_model
+from iopaint.plugins.base_plugin import BasePlugin
+from iopaint.plugins.segment_anything import SamPredictor, sam_model_registry
+from iopaint.schema import RunPluginRequest
+
+# 从小到大
+SEGMENT_ANYTHING_MODELS = {
+ "vit_b": {
+ "url": "https://dl.fbaipublicfiles.com/segment_anything/sam_vit_b_01ec64.pth",
+ "md5": "01ec64d29a2fca3f0661936605ae66f8",
+ },
+ "vit_l": {
+ "url": "https://dl.fbaipublicfiles.com/segment_anything/sam_vit_l_0b3195.pth",
+ "md5": "0b3195507c641ddb6910d2bb5adee89c",
+ },
+ "vit_h": {
+ "url": "https://dl.fbaipublicfiles.com/segment_anything/sam_vit_h_4b8939.pth",
+ "md5": "4b8939a88964f0f4ff5f5b2642c598a6",
+ },
+ "mobile_sam": {
+ "url": "https://github.com/Sanster/models/releases/download/MobileSAM/mobile_sam.pt",
+ "md5": "f3c0d8cda613564d499310dab6c812cd",
+ },
+}
+
+
+class InteractiveSeg(BasePlugin):
+ name = "InteractiveSeg"
+ support_gen_mask = True
+
+ def __init__(self, model_name, device):
+ super().__init__()
+ self.model_name = model_name
+ self.device = device
+ self._init_session(model_name)
+
+ def _init_session(self, model_name: str):
+ model_path = download_model(
+ SEGMENT_ANYTHING_MODELS[model_name]["url"],
+ SEGMENT_ANYTHING_MODELS[model_name]["md5"],
+ )
+ logger.info(f"SegmentAnything model path: {model_path}")
+ self.predictor = SamPredictor(
+ sam_model_registry[model_name](checkpoint=model_path).to(self.device)
+ )
+ self.prev_img_md5 = None
+
+ def switch_model(self, new_model_name):
+ if self.model_name == new_model_name:
+ return
+
+ logger.info(
+ f"Switching InteractiveSeg model from {self.model_name} to {new_model_name}"
+ )
+ self._init_session(new_model_name)
+ self.model_name = new_model_name
+
+ def gen_mask(self, rgb_np_img, req: RunPluginRequest) -> np.ndarray:
+ img_md5 = hashlib.md5(req.image.encode("utf-8")).hexdigest()
+ return self.forward(rgb_np_img, req.clicks, img_md5)
+
+ @torch.inference_mode()
+ def forward(self, rgb_np_img, clicks: List[List], img_md5: str):
+ input_point = []
+ input_label = []
+ for click in clicks:
+ x = click[0]
+ y = click[1]
+ input_point.append([x, y])
+ input_label.append(click[2])
+
+ if img_md5 and img_md5 != self.prev_img_md5:
+ self.prev_img_md5 = img_md5
+ self.predictor.set_image(rgb_np_img)
+
+ masks, scores, _ = self.predictor.predict(
+ point_coords=np.array(input_point),
+ point_labels=np.array(input_label),
+ multimask_output=False,
+ )
+ mask = masks[0].astype(np.uint8) * 255
+ return mask
diff --git a/iopaint/plugins/segment_anything/build_sam.py b/iopaint/plugins/segment_anything/build_sam.py
new file mode 100644
index 0000000000000000000000000000000000000000..f8dea8e0baf5c66240a6c4c277a90e6b1f44f6e7
--- /dev/null
+++ b/iopaint/plugins/segment_anything/build_sam.py
@@ -0,0 +1,168 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+import torch
+
+from functools import partial
+
+from iopaint.plugins.segment_anything.modeling.tiny_vit_sam import TinyViT
+
+from .modeling import (
+ ImageEncoderViT,
+ MaskDecoder,
+ PromptEncoder,
+ Sam,
+ TwoWayTransformer,
+)
+
+
+def build_sam_vit_h(checkpoint=None):
+ return _build_sam(
+ encoder_embed_dim=1280,
+ encoder_depth=32,
+ encoder_num_heads=16,
+ encoder_global_attn_indexes=[7, 15, 23, 31],
+ checkpoint=checkpoint,
+ )
+
+
+build_sam = build_sam_vit_h
+
+
+def build_sam_vit_l(checkpoint=None):
+ return _build_sam(
+ encoder_embed_dim=1024,
+ encoder_depth=24,
+ encoder_num_heads=16,
+ encoder_global_attn_indexes=[5, 11, 17, 23],
+ checkpoint=checkpoint,
+ )
+
+
+def build_sam_vit_b(checkpoint=None):
+ return _build_sam(
+ encoder_embed_dim=768,
+ encoder_depth=12,
+ encoder_num_heads=12,
+ encoder_global_attn_indexes=[2, 5, 8, 11],
+ checkpoint=checkpoint,
+ )
+
+
+def build_sam_vit_t(checkpoint=None):
+ prompt_embed_dim = 256
+ image_size = 1024
+ vit_patch_size = 16
+ image_embedding_size = image_size // vit_patch_size
+ mobile_sam = Sam(
+ image_encoder=TinyViT(
+ img_size=1024,
+ in_chans=3,
+ num_classes=1000,
+ embed_dims=[64, 128, 160, 320],
+ depths=[2, 2, 6, 2],
+ num_heads=[2, 4, 5, 10],
+ window_sizes=[7, 7, 14, 7],
+ mlp_ratio=4.0,
+ drop_rate=0.0,
+ drop_path_rate=0.0,
+ use_checkpoint=False,
+ mbconv_expand_ratio=4.0,
+ local_conv_size=3,
+ layer_lr_decay=0.8,
+ ),
+ prompt_encoder=PromptEncoder(
+ embed_dim=prompt_embed_dim,
+ image_embedding_size=(image_embedding_size, image_embedding_size),
+ input_image_size=(image_size, image_size),
+ mask_in_chans=16,
+ ),
+ mask_decoder=MaskDecoder(
+ num_multimask_outputs=3,
+ transformer=TwoWayTransformer(
+ depth=2,
+ embedding_dim=prompt_embed_dim,
+ mlp_dim=2048,
+ num_heads=8,
+ ),
+ transformer_dim=prompt_embed_dim,
+ iou_head_depth=3,
+ iou_head_hidden_dim=256,
+ ),
+ pixel_mean=[123.675, 116.28, 103.53],
+ pixel_std=[58.395, 57.12, 57.375],
+ )
+
+ mobile_sam.eval()
+ if checkpoint is not None:
+ with open(checkpoint, "rb") as f:
+ state_dict = torch.load(f)
+ mobile_sam.load_state_dict(state_dict)
+ return mobile_sam
+
+
+sam_model_registry = {
+ "default": build_sam,
+ "vit_h": build_sam,
+ "vit_l": build_sam_vit_l,
+ "vit_b": build_sam_vit_b,
+ "mobile_sam": build_sam_vit_t,
+}
+
+
+def _build_sam(
+ encoder_embed_dim,
+ encoder_depth,
+ encoder_num_heads,
+ encoder_global_attn_indexes,
+ checkpoint=None,
+):
+ prompt_embed_dim = 256
+ image_size = 1024
+ vit_patch_size = 16
+ image_embedding_size = image_size // vit_patch_size
+ sam = Sam(
+ image_encoder=ImageEncoderViT(
+ depth=encoder_depth,
+ embed_dim=encoder_embed_dim,
+ img_size=image_size,
+ mlp_ratio=4,
+ norm_layer=partial(torch.nn.LayerNorm, eps=1e-6),
+ num_heads=encoder_num_heads,
+ patch_size=vit_patch_size,
+ qkv_bias=True,
+ use_rel_pos=True,
+ global_attn_indexes=encoder_global_attn_indexes,
+ window_size=14,
+ out_chans=prompt_embed_dim,
+ ),
+ prompt_encoder=PromptEncoder(
+ embed_dim=prompt_embed_dim,
+ image_embedding_size=(image_embedding_size, image_embedding_size),
+ input_image_size=(image_size, image_size),
+ mask_in_chans=16,
+ ),
+ mask_decoder=MaskDecoder(
+ num_multimask_outputs=3,
+ transformer=TwoWayTransformer(
+ depth=2,
+ embedding_dim=prompt_embed_dim,
+ mlp_dim=2048,
+ num_heads=8,
+ ),
+ transformer_dim=prompt_embed_dim,
+ iou_head_depth=3,
+ iou_head_hidden_dim=256,
+ ),
+ pixel_mean=[123.675, 116.28, 103.53],
+ pixel_std=[58.395, 57.12, 57.375],
+ )
+ sam.eval()
+ if checkpoint is not None:
+ with open(checkpoint, "rb") as f:
+ state_dict = torch.load(f)
+ sam.load_state_dict(state_dict)
+ return sam
diff --git a/iopaint/plugins/segment_anything/modeling/common.py b/iopaint/plugins/segment_anything/modeling/common.py
new file mode 100644
index 0000000000000000000000000000000000000000..2bf15236a3eb24d8526073bc4fa2b274cccb3f96
--- /dev/null
+++ b/iopaint/plugins/segment_anything/modeling/common.py
@@ -0,0 +1,43 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+import torch
+import torch.nn as nn
+
+from typing import Type
+
+
+class MLPBlock(nn.Module):
+ def __init__(
+ self,
+ embedding_dim: int,
+ mlp_dim: int,
+ act: Type[nn.Module] = nn.GELU,
+ ) -> None:
+ super().__init__()
+ self.lin1 = nn.Linear(embedding_dim, mlp_dim)
+ self.lin2 = nn.Linear(mlp_dim, embedding_dim)
+ self.act = act()
+
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
+ return self.lin2(self.act(self.lin1(x)))
+
+
+# From https://github.com/facebookresearch/detectron2/blob/main/detectron2/layers/batch_norm.py # noqa
+# Itself from https://github.com/facebookresearch/ConvNeXt/blob/d1fa8f6fef0a165b27399986cc2bdacc92777e40/models/convnext.py#L119 # noqa
+class LayerNorm2d(nn.Module):
+ def __init__(self, num_channels: int, eps: float = 1e-6) -> None:
+ super().__init__()
+ self.weight = nn.Parameter(torch.ones(num_channels))
+ self.bias = nn.Parameter(torch.zeros(num_channels))
+ self.eps = eps
+
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
+ u = x.mean(1, keepdim=True)
+ s = (x - u).pow(2).mean(1, keepdim=True)
+ x = (x - u) / torch.sqrt(s + self.eps)
+ x = self.weight[:, None, None] * x + self.bias[:, None, None]
+ return x
diff --git a/iopaint/plugins/segment_anything/modeling/image_encoder.py b/iopaint/plugins/segment_anything/modeling/image_encoder.py
new file mode 100644
index 0000000000000000000000000000000000000000..a6ad9ad2938842308e482a05c9d35ab08db9b2c3
--- /dev/null
+++ b/iopaint/plugins/segment_anything/modeling/image_encoder.py
@@ -0,0 +1,395 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from typing import Optional, Tuple, Type
+
+from .common import LayerNorm2d, MLPBlock
+
+
+# This class and its supporting functions below lightly adapted from the ViTDet backbone available at: https://github.com/facebookresearch/detectron2/blob/main/detectron2/modeling/backbone/vit.py # noqa
+class ImageEncoderViT(nn.Module):
+ def __init__(
+ self,
+ img_size: int = 1024,
+ patch_size: int = 16,
+ in_chans: int = 3,
+ embed_dim: int = 768,
+ depth: int = 12,
+ num_heads: int = 12,
+ mlp_ratio: float = 4.0,
+ out_chans: int = 256,
+ qkv_bias: bool = True,
+ norm_layer: Type[nn.Module] = nn.LayerNorm,
+ act_layer: Type[nn.Module] = nn.GELU,
+ use_abs_pos: bool = True,
+ use_rel_pos: bool = False,
+ rel_pos_zero_init: bool = True,
+ window_size: int = 0,
+ global_attn_indexes: Tuple[int, ...] = (),
+ ) -> None:
+ """
+ Args:
+ img_size (int): Input image size.
+ patch_size (int): Patch size.
+ in_chans (int): Number of input image channels.
+ embed_dim (int): Patch embedding dimension.
+ depth (int): Depth of ViT.
+ num_heads (int): Number of attention heads in each ViT block.
+ mlp_ratio (float): Ratio of mlp hidden dim to embedding dim.
+ qkv_bias (bool): If True, add a learnable bias to query, key, value.
+ norm_layer (nn.Module): Normalization layer.
+ act_layer (nn.Module): Activation layer.
+ use_abs_pos (bool): If True, use absolute positional embeddings.
+ use_rel_pos (bool): If True, add relative positional embeddings to the attention map.
+ rel_pos_zero_init (bool): If True, zero initialize relative positional parameters.
+ window_size (int): Window size for window attention blocks.
+ global_attn_indexes (list): Indexes for blocks using global attention.
+ """
+ super().__init__()
+ self.img_size = img_size
+
+ self.patch_embed = PatchEmbed(
+ kernel_size=(patch_size, patch_size),
+ stride=(patch_size, patch_size),
+ in_chans=in_chans,
+ embed_dim=embed_dim,
+ )
+
+ self.pos_embed: Optional[nn.Parameter] = None
+ if use_abs_pos:
+ # Initialize absolute positional embedding with pretrain image size.
+ self.pos_embed = nn.Parameter(
+ torch.zeros(1, img_size // patch_size, img_size // patch_size, embed_dim)
+ )
+
+ self.blocks = nn.ModuleList()
+ for i in range(depth):
+ block = Block(
+ dim=embed_dim,
+ num_heads=num_heads,
+ mlp_ratio=mlp_ratio,
+ qkv_bias=qkv_bias,
+ norm_layer=norm_layer,
+ act_layer=act_layer,
+ use_rel_pos=use_rel_pos,
+ rel_pos_zero_init=rel_pos_zero_init,
+ window_size=window_size if i not in global_attn_indexes else 0,
+ input_size=(img_size // patch_size, img_size // patch_size),
+ )
+ self.blocks.append(block)
+
+ self.neck = nn.Sequential(
+ nn.Conv2d(
+ embed_dim,
+ out_chans,
+ kernel_size=1,
+ bias=False,
+ ),
+ LayerNorm2d(out_chans),
+ nn.Conv2d(
+ out_chans,
+ out_chans,
+ kernel_size=3,
+ padding=1,
+ bias=False,
+ ),
+ LayerNorm2d(out_chans),
+ )
+
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
+ x = self.patch_embed(x)
+ if self.pos_embed is not None:
+ x = x + self.pos_embed
+
+ for blk in self.blocks:
+ x = blk(x)
+
+ x = self.neck(x.permute(0, 3, 1, 2))
+
+ return x
+
+
+class Block(nn.Module):
+ """Transformer blocks with support of window attention and residual propagation blocks"""
+
+ def __init__(
+ self,
+ dim: int,
+ num_heads: int,
+ mlp_ratio: float = 4.0,
+ qkv_bias: bool = True,
+ norm_layer: Type[nn.Module] = nn.LayerNorm,
+ act_layer: Type[nn.Module] = nn.GELU,
+ use_rel_pos: bool = False,
+ rel_pos_zero_init: bool = True,
+ window_size: int = 0,
+ input_size: Optional[Tuple[int, int]] = None,
+ ) -> None:
+ """
+ Args:
+ dim (int): Number of input channels.
+ num_heads (int): Number of attention heads in each ViT block.
+ mlp_ratio (float): Ratio of mlp hidden dim to embedding dim.
+ qkv_bias (bool): If True, add a learnable bias to query, key, value.
+ norm_layer (nn.Module): Normalization layer.
+ act_layer (nn.Module): Activation layer.
+ use_rel_pos (bool): If True, add relative positional embeddings to the attention map.
+ rel_pos_zero_init (bool): If True, zero initialize relative positional parameters.
+ window_size (int): Window size for window attention blocks. If it equals 0, then
+ use global attention.
+ input_size (int or None): Input resolution for calculating the relative positional
+ parameter size.
+ """
+ super().__init__()
+ self.norm1 = norm_layer(dim)
+ self.attn = Attention(
+ dim,
+ num_heads=num_heads,
+ qkv_bias=qkv_bias,
+ use_rel_pos=use_rel_pos,
+ rel_pos_zero_init=rel_pos_zero_init,
+ input_size=input_size if window_size == 0 else (window_size, window_size),
+ )
+
+ self.norm2 = norm_layer(dim)
+ self.mlp = MLPBlock(embedding_dim=dim, mlp_dim=int(dim * mlp_ratio), act=act_layer)
+
+ self.window_size = window_size
+
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
+ shortcut = x
+ x = self.norm1(x)
+ # Window partition
+ if self.window_size > 0:
+ H, W = x.shape[1], x.shape[2]
+ x, pad_hw = window_partition(x, self.window_size)
+
+ x = self.attn(x)
+ # Reverse window partition
+ if self.window_size > 0:
+ x = window_unpartition(x, self.window_size, pad_hw, (H, W))
+
+ x = shortcut + x
+ x = x + self.mlp(self.norm2(x))
+
+ return x
+
+
+class Attention(nn.Module):
+ """Multi-head Attention block with relative position embeddings."""
+
+ def __init__(
+ self,
+ dim: int,
+ num_heads: int = 8,
+ qkv_bias: bool = True,
+ use_rel_pos: bool = False,
+ rel_pos_zero_init: bool = True,
+ input_size: Optional[Tuple[int, int]] = None,
+ ) -> None:
+ """
+ Args:
+ dim (int): Number of input channels.
+ num_heads (int): Number of attention heads.
+ qkv_bias (bool: If True, add a learnable bias to query, key, value.
+ rel_pos (bool): If True, add relative positional embeddings to the attention map.
+ rel_pos_zero_init (bool): If True, zero initialize relative positional parameters.
+ input_size (int or None): Input resolution for calculating the relative positional
+ parameter size.
+ """
+ super().__init__()
+ self.num_heads = num_heads
+ head_dim = dim // num_heads
+ self.scale = head_dim**-0.5
+
+ self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
+ self.proj = nn.Linear(dim, dim)
+
+ self.use_rel_pos = use_rel_pos
+ if self.use_rel_pos:
+ assert (
+ input_size is not None
+ ), "Input size must be provided if using relative positional encoding."
+ # initialize relative positional embeddings
+ self.rel_pos_h = nn.Parameter(torch.zeros(2 * input_size[0] - 1, head_dim))
+ self.rel_pos_w = nn.Parameter(torch.zeros(2 * input_size[1] - 1, head_dim))
+
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
+ B, H, W, _ = x.shape
+ # qkv with shape (3, B, nHead, H * W, C)
+ qkv = self.qkv(x).reshape(B, H * W, 3, self.num_heads, -1).permute(2, 0, 3, 1, 4)
+ # q, k, v with shape (B * nHead, H * W, C)
+ q, k, v = qkv.reshape(3, B * self.num_heads, H * W, -1).unbind(0)
+
+ attn = (q * self.scale) @ k.transpose(-2, -1)
+
+ if self.use_rel_pos:
+ attn = add_decomposed_rel_pos(attn, q, self.rel_pos_h, self.rel_pos_w, (H, W), (H, W))
+
+ attn = attn.softmax(dim=-1)
+ x = (attn @ v).view(B, self.num_heads, H, W, -1).permute(0, 2, 3, 1, 4).reshape(B, H, W, -1)
+ x = self.proj(x)
+
+ return x
+
+
+def window_partition(x: torch.Tensor, window_size: int) -> Tuple[torch.Tensor, Tuple[int, int]]:
+ """
+ Partition into non-overlapping windows with padding if needed.
+ Args:
+ x (tensor): input tokens with [B, H, W, C].
+ window_size (int): window size.
+
+ Returns:
+ windows: windows after partition with [B * num_windows, window_size, window_size, C].
+ (Hp, Wp): padded height and width before partition
+ """
+ B, H, W, C = x.shape
+
+ pad_h = (window_size - H % window_size) % window_size
+ pad_w = (window_size - W % window_size) % window_size
+ if pad_h > 0 or pad_w > 0:
+ x = F.pad(x, (0, 0, 0, pad_w, 0, pad_h))
+ Hp, Wp = H + pad_h, W + pad_w
+
+ x = x.view(B, Hp // window_size, window_size, Wp // window_size, window_size, C)
+ windows = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(-1, window_size, window_size, C)
+ return windows, (Hp, Wp)
+
+
+def window_unpartition(
+ windows: torch.Tensor, window_size: int, pad_hw: Tuple[int, int], hw: Tuple[int, int]
+) -> torch.Tensor:
+ """
+ Window unpartition into original sequences and removing padding.
+ Args:
+ x (tensor): input tokens with [B * num_windows, window_size, window_size, C].
+ window_size (int): window size.
+ pad_hw (Tuple): padded height and width (Hp, Wp).
+ hw (Tuple): original height and width (H, W) before padding.
+
+ Returns:
+ x: unpartitioned sequences with [B, H, W, C].
+ """
+ Hp, Wp = pad_hw
+ H, W = hw
+ B = windows.shape[0] // (Hp * Wp // window_size // window_size)
+ x = windows.view(B, Hp // window_size, Wp // window_size, window_size, window_size, -1)
+ x = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(B, Hp, Wp, -1)
+
+ if Hp > H or Wp > W:
+ x = x[:, :H, :W, :].contiguous()
+ return x
+
+
+def get_rel_pos(q_size: int, k_size: int, rel_pos: torch.Tensor) -> torch.Tensor:
+ """
+ Get relative positional embeddings according to the relative positions of
+ query and key sizes.
+ Args:
+ q_size (int): size of query q.
+ k_size (int): size of key k.
+ rel_pos (Tensor): relative position embeddings (L, C).
+
+ Returns:
+ Extracted positional embeddings according to relative positions.
+ """
+ max_rel_dist = int(2 * max(q_size, k_size) - 1)
+ # Interpolate rel pos if needed.
+ if rel_pos.shape[0] != max_rel_dist:
+ # Interpolate rel pos.
+ rel_pos_resized = F.interpolate(
+ rel_pos.reshape(1, rel_pos.shape[0], -1).permute(0, 2, 1),
+ size=max_rel_dist,
+ mode="linear",
+ )
+ rel_pos_resized = rel_pos_resized.reshape(-1, max_rel_dist).permute(1, 0)
+ else:
+ rel_pos_resized = rel_pos
+
+ # Scale the coords with short length if shapes for q and k are different.
+ q_coords = torch.arange(q_size)[:, None] * max(k_size / q_size, 1.0)
+ k_coords = torch.arange(k_size)[None, :] * max(q_size / k_size, 1.0)
+ relative_coords = (q_coords - k_coords) + (k_size - 1) * max(q_size / k_size, 1.0)
+
+ return rel_pos_resized[relative_coords.long()]
+
+
+def add_decomposed_rel_pos(
+ attn: torch.Tensor,
+ q: torch.Tensor,
+ rel_pos_h: torch.Tensor,
+ rel_pos_w: torch.Tensor,
+ q_size: Tuple[int, int],
+ k_size: Tuple[int, int],
+) -> torch.Tensor:
+ """
+ Calculate decomposed Relative Positional Embeddings from :paper:`mvitv2`.
+ https://github.com/facebookresearch/mvit/blob/19786631e330df9f3622e5402b4a419a263a2c80/mvit/models/attention.py # noqa B950
+ Args:
+ attn (Tensor): attention map.
+ q (Tensor): query q in the attention layer with shape (B, q_h * q_w, C).
+ rel_pos_h (Tensor): relative position embeddings (Lh, C) for height axis.
+ rel_pos_w (Tensor): relative position embeddings (Lw, C) for width axis.
+ q_size (Tuple): spatial sequence size of query q with (q_h, q_w).
+ k_size (Tuple): spatial sequence size of key k with (k_h, k_w).
+
+ Returns:
+ attn (Tensor): attention map with added relative positional embeddings.
+ """
+ q_h, q_w = q_size
+ k_h, k_w = k_size
+ Rh = get_rel_pos(q_h, k_h, rel_pos_h)
+ Rw = get_rel_pos(q_w, k_w, rel_pos_w)
+
+ B, _, dim = q.shape
+ r_q = q.reshape(B, q_h, q_w, dim)
+ rel_h = torch.einsum("bhwc,hkc->bhwk", r_q, Rh)
+ rel_w = torch.einsum("bhwc,wkc->bhwk", r_q, Rw)
+
+ attn = (
+ attn.view(B, q_h, q_w, k_h, k_w) + rel_h[:, :, :, :, None] + rel_w[:, :, :, None, :]
+ ).view(B, q_h * q_w, k_h * k_w)
+
+ return attn
+
+
+class PatchEmbed(nn.Module):
+ """
+ Image to Patch Embedding.
+ """
+
+ def __init__(
+ self,
+ kernel_size: Tuple[int, int] = (16, 16),
+ stride: Tuple[int, int] = (16, 16),
+ padding: Tuple[int, int] = (0, 0),
+ in_chans: int = 3,
+ embed_dim: int = 768,
+ ) -> None:
+ """
+ Args:
+ kernel_size (Tuple): kernel size of the projection layer.
+ stride (Tuple): stride of the projection layer.
+ padding (Tuple): padding size of the projection layer.
+ in_chans (int): Number of input image channels.
+ embed_dim (int): embed_dim (int): Patch embedding dimension.
+ """
+ super().__init__()
+
+ self.proj = nn.Conv2d(
+ in_chans, embed_dim, kernel_size=kernel_size, stride=stride, padding=padding
+ )
+
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
+ x = self.proj(x)
+ # B C H W -> B H W C
+ x = x.permute(0, 2, 3, 1)
+ return x
diff --git a/iopaint/plugins/segment_anything/modeling/mask_decoder.py b/iopaint/plugins/segment_anything/modeling/mask_decoder.py
new file mode 100644
index 0000000000000000000000000000000000000000..3e86f7cc9ad95582a08ef2531c68d03fa4af8d99
--- /dev/null
+++ b/iopaint/plugins/segment_anything/modeling/mask_decoder.py
@@ -0,0 +1,176 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+import torch
+from torch import nn
+from torch.nn import functional as F
+
+from typing import List, Tuple, Type
+
+from .common import LayerNorm2d
+
+
+class MaskDecoder(nn.Module):
+ def __init__(
+ self,
+ *,
+ transformer_dim: int,
+ transformer: nn.Module,
+ num_multimask_outputs: int = 3,
+ activation: Type[nn.Module] = nn.GELU,
+ iou_head_depth: int = 3,
+ iou_head_hidden_dim: int = 256,
+ ) -> None:
+ """
+ Predicts masks given an image and prompt embeddings, using a
+ tranformer architecture.
+
+ Arguments:
+ transformer_dim (int): the channel dimension of the transformer
+ transformer (nn.Module): the transformer used to predict masks
+ num_multimask_outputs (int): the number of masks to predict
+ when disambiguating masks
+ activation (nn.Module): the type of activation to use when
+ upscaling masks
+ iou_head_depth (int): the depth of the MLP used to predict
+ mask quality
+ iou_head_hidden_dim (int): the hidden dimension of the MLP
+ used to predict mask quality
+ """
+ super().__init__()
+ self.transformer_dim = transformer_dim
+ self.transformer = transformer
+
+ self.num_multimask_outputs = num_multimask_outputs
+
+ self.iou_token = nn.Embedding(1, transformer_dim)
+ self.num_mask_tokens = num_multimask_outputs + 1
+ self.mask_tokens = nn.Embedding(self.num_mask_tokens, transformer_dim)
+
+ self.output_upscaling = nn.Sequential(
+ nn.ConvTranspose2d(transformer_dim, transformer_dim // 4, kernel_size=2, stride=2),
+ LayerNorm2d(transformer_dim // 4),
+ activation(),
+ nn.ConvTranspose2d(transformer_dim // 4, transformer_dim // 8, kernel_size=2, stride=2),
+ activation(),
+ )
+ self.output_hypernetworks_mlps = nn.ModuleList(
+ [
+ MLP(transformer_dim, transformer_dim, transformer_dim // 8, 3)
+ for i in range(self.num_mask_tokens)
+ ]
+ )
+
+ self.iou_prediction_head = MLP(
+ transformer_dim, iou_head_hidden_dim, self.num_mask_tokens, iou_head_depth
+ )
+
+ def forward(
+ self,
+ image_embeddings: torch.Tensor,
+ image_pe: torch.Tensor,
+ sparse_prompt_embeddings: torch.Tensor,
+ dense_prompt_embeddings: torch.Tensor,
+ multimask_output: bool,
+ ) -> Tuple[torch.Tensor, torch.Tensor]:
+ """
+ Predict masks given image and prompt embeddings.
+
+ Arguments:
+ image_embeddings (torch.Tensor): the embeddings from the image encoder
+ image_pe (torch.Tensor): positional encoding with the shape of image_embeddings
+ sparse_prompt_embeddings (torch.Tensor): the embeddings of the points and boxes
+ dense_prompt_embeddings (torch.Tensor): the embeddings of the mask inputs
+ multimask_output (bool): Whether to return multiple masks or a single
+ mask.
+
+ Returns:
+ torch.Tensor: batched predicted masks
+ torch.Tensor: batched predictions of mask quality
+ """
+ masks, iou_pred = self.predict_masks(
+ image_embeddings=image_embeddings,
+ image_pe=image_pe,
+ sparse_prompt_embeddings=sparse_prompt_embeddings,
+ dense_prompt_embeddings=dense_prompt_embeddings,
+ )
+
+ # Select the correct mask or masks for outptu
+ if multimask_output:
+ mask_slice = slice(1, None)
+ else:
+ mask_slice = slice(0, 1)
+ masks = masks[:, mask_slice, :, :]
+ iou_pred = iou_pred[:, mask_slice]
+
+ # Prepare output
+ return masks, iou_pred
+
+ def predict_masks(
+ self,
+ image_embeddings: torch.Tensor,
+ image_pe: torch.Tensor,
+ sparse_prompt_embeddings: torch.Tensor,
+ dense_prompt_embeddings: torch.Tensor,
+ ) -> Tuple[torch.Tensor, torch.Tensor]:
+ """Predicts masks. See 'forward' for more details."""
+ # Concatenate output tokens
+ output_tokens = torch.cat([self.iou_token.weight, self.mask_tokens.weight], dim=0)
+ output_tokens = output_tokens.unsqueeze(0).expand(sparse_prompt_embeddings.size(0), -1, -1)
+ tokens = torch.cat((output_tokens, sparse_prompt_embeddings), dim=1)
+
+ # Expand per-image data in batch direction to be per-mask
+ src = torch.repeat_interleave(image_embeddings, tokens.shape[0], dim=0)
+ src = src + dense_prompt_embeddings
+ pos_src = torch.repeat_interleave(image_pe, tokens.shape[0], dim=0)
+ b, c, h, w = src.shape
+
+ # Run the transformer
+ hs, src = self.transformer(src, pos_src, tokens)
+ iou_token_out = hs[:, 0, :]
+ mask_tokens_out = hs[:, 1 : (1 + self.num_mask_tokens), :]
+
+ # Upscale mask embeddings and predict masks using the mask tokens
+ src = src.transpose(1, 2).view(b, c, h, w)
+ upscaled_embedding = self.output_upscaling(src)
+ hyper_in_list: List[torch.Tensor] = []
+ for i in range(self.num_mask_tokens):
+ hyper_in_list.append(self.output_hypernetworks_mlps[i](mask_tokens_out[:, i, :]))
+ hyper_in = torch.stack(hyper_in_list, dim=1)
+ b, c, h, w = upscaled_embedding.shape
+ masks = (hyper_in @ upscaled_embedding.view(b, c, h * w)).view(b, -1, h, w)
+
+ # Generate mask quality predictions
+ iou_pred = self.iou_prediction_head(iou_token_out)
+
+ return masks, iou_pred
+
+
+# Lightly adapted from
+# https://github.com/facebookresearch/MaskFormer/blob/main/mask_former/modeling/transformer/transformer_predictor.py # noqa
+class MLP(nn.Module):
+ def __init__(
+ self,
+ input_dim: int,
+ hidden_dim: int,
+ output_dim: int,
+ num_layers: int,
+ sigmoid_output: bool = False,
+ ) -> None:
+ super().__init__()
+ self.num_layers = num_layers
+ h = [hidden_dim] * (num_layers - 1)
+ self.layers = nn.ModuleList(
+ nn.Linear(n, k) for n, k in zip([input_dim] + h, h + [output_dim])
+ )
+ self.sigmoid_output = sigmoid_output
+
+ def forward(self, x):
+ for i, layer in enumerate(self.layers):
+ x = F.relu(layer(x)) if i < self.num_layers - 1 else layer(x)
+ if self.sigmoid_output:
+ x = F.sigmoid(x)
+ return x
diff --git a/model/networks.py b/model/networks.py
new file mode 100644
index 0000000000000000000000000000000000000000..935792a79e3d624f4e03e7130852901461144043
--- /dev/null
+++ b/model/networks.py
@@ -0,0 +1,563 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.nn.utils import spectral_norm as spectral_norm_fn
+from torch.nn.utils import weight_norm as weight_norm_fn
+from PIL import Image
+from torchvision import transforms
+from torchvision import utils as vutils
+
+from utils.tools import extract_image_patches, flow_to_image, \
+ reduce_mean, reduce_sum, default_loader, same_padding
+
+
+class Generator(nn.Module):
+ def __init__(self, config, use_cuda, device_ids):
+ super(Generator, self).__init__()
+ self.input_dim = config['input_dim']
+ self.cnum = config['ngf']
+ self.use_cuda = use_cuda
+ self.device_ids = device_ids
+
+ self.coarse_generator = CoarseGenerator(self.input_dim, self.cnum, self.use_cuda, self.device_ids)
+ self.fine_generator = FineGenerator(self.input_dim, self.cnum, self.use_cuda, self.device_ids)
+
+ def forward(self, x, mask):
+ x_stage1 = self.coarse_generator(x, mask)
+ x_stage2, offset_flow = self.fine_generator(x, x_stage1, mask)
+ return x_stage1, x_stage2, offset_flow
+
+
+class CoarseGenerator(nn.Module):
+ def __init__(self, input_dim, cnum, use_cuda=True, device_ids=None):
+ super(CoarseGenerator, self).__init__()
+ self.use_cuda = use_cuda
+ self.device_ids = device_ids
+
+ self.conv1 = gen_conv(input_dim + 2, cnum, 5, 1, 2)
+ self.conv2_downsample = gen_conv(cnum, cnum*2, 3, 2, 1)
+ self.conv3 = gen_conv(cnum*2, cnum*2, 3, 1, 1)
+ self.conv4_downsample = gen_conv(cnum*2, cnum*4, 3, 2, 1)
+ self.conv5 = gen_conv(cnum*4, cnum*4, 3, 1, 1)
+ self.conv6 = gen_conv(cnum*4, cnum*4, 3, 1, 1)
+
+ self.conv7_atrous = gen_conv(cnum*4, cnum*4, 3, 1, 2, rate=2)
+ self.conv8_atrous = gen_conv(cnum*4, cnum*4, 3, 1, 4, rate=4)
+ self.conv9_atrous = gen_conv(cnum*4, cnum*4, 3, 1, 8, rate=8)
+ self.conv10_atrous = gen_conv(cnum*4, cnum*4, 3, 1, 16, rate=16)
+
+ self.conv11 = gen_conv(cnum*4, cnum*4, 3, 1, 1)
+ self.conv12 = gen_conv(cnum*4, cnum*4, 3, 1, 1)
+
+ self.conv13 = gen_conv(cnum*4, cnum*2, 3, 1, 1)
+ self.conv14 = gen_conv(cnum*2, cnum*2, 3, 1, 1)
+ self.conv15 = gen_conv(cnum*2, cnum, 3, 1, 1)
+ self.conv16 = gen_conv(cnum, cnum//2, 3, 1, 1)
+ self.conv17 = gen_conv(cnum//2, input_dim, 3, 1, 1, activation='none')
+
+ def forward(self, x, mask):
+ # For indicating the boundaries of images
+ ones = torch.ones(x.size(0), 1, x.size(2), x.size(3))
+ if self.use_cuda:
+ ones = ones.cuda()
+ mask = mask.cuda()
+ # 5 x 256 x 256
+ x = self.conv1(torch.cat([x, ones, mask], dim=1))
+ x = self.conv2_downsample(x)
+ # cnum*2 x 128 x 128
+ x = self.conv3(x)
+ x = self.conv4_downsample(x)
+ # cnum*4 x 64 x 64
+ x = self.conv5(x)
+ x = self.conv6(x)
+ x = self.conv7_atrous(x)
+ x = self.conv8_atrous(x)
+ x = self.conv9_atrous(x)
+ x = self.conv10_atrous(x)
+ x = self.conv11(x)
+ x = self.conv12(x)
+ x = F.interpolate(x, scale_factor=2, mode='nearest')
+ # cnum*2 x 128 x 128
+ x = self.conv13(x)
+ x = self.conv14(x)
+ x = F.interpolate(x, scale_factor=2, mode='nearest')
+ # cnum x 256 x 256
+ x = self.conv15(x)
+ x = self.conv16(x)
+ x = self.conv17(x)
+ # 3 x 256 x 256
+ x_stage1 = torch.clamp(x, -1., 1.)
+
+ return x_stage1
+
+
+class FineGenerator(nn.Module):
+ def __init__(self, input_dim, cnum, use_cuda=True, device_ids=None):
+ super(FineGenerator, self).__init__()
+ self.use_cuda = use_cuda
+ self.device_ids = device_ids
+
+ # 3 x 256 x 256
+ self.conv1 = gen_conv(input_dim + 2, cnum, 5, 1, 2)
+ self.conv2_downsample = gen_conv(cnum, cnum, 3, 2, 1)
+ # cnum*2 x 128 x 128
+ self.conv3 = gen_conv(cnum, cnum*2, 3, 1, 1)
+ self.conv4_downsample = gen_conv(cnum*2, cnum*2, 3, 2, 1)
+ # cnum*4 x 64 x 64
+ self.conv5 = gen_conv(cnum*2, cnum*4, 3, 1, 1)
+ self.conv6 = gen_conv(cnum*4, cnum*4, 3, 1, 1)
+
+ self.conv7_atrous = gen_conv(cnum*4, cnum*4, 3, 1, 2, rate=2)
+ self.conv8_atrous = gen_conv(cnum*4, cnum*4, 3, 1, 4, rate=4)
+ self.conv9_atrous = gen_conv(cnum*4, cnum*4, 3, 1, 8, rate=8)
+ self.conv10_atrous = gen_conv(cnum*4, cnum*4, 3, 1, 16, rate=16)
+
+ # attention branch
+ # 3 x 256 x 256
+ self.pmconv1 = gen_conv(input_dim + 2, cnum, 5, 1, 2)
+ self.pmconv2_downsample = gen_conv(cnum, cnum, 3, 2, 1)
+ # cnum*2 x 128 x 128
+ self.pmconv3 = gen_conv(cnum, cnum*2, 3, 1, 1)
+ self.pmconv4_downsample = gen_conv(cnum*2, cnum*4, 3, 2, 1)
+ # cnum*4 x 64 x 64
+ self.pmconv5 = gen_conv(cnum*4, cnum*4, 3, 1, 1)
+ self.pmconv6 = gen_conv(cnum*4, cnum*4, 3, 1, 1, activation='relu')
+ self.contextul_attention = ContextualAttention(ksize=3, stride=1, rate=2, fuse_k=3, softmax_scale=10,
+ fuse=True, use_cuda=self.use_cuda, device_ids=self.device_ids)
+ self.pmconv9 = gen_conv(cnum*4, cnum*4, 3, 1, 1)
+ self.pmconv10 = gen_conv(cnum*4, cnum*4, 3, 1, 1)
+ self.allconv11 = gen_conv(cnum*8, cnum*4, 3, 1, 1)
+ self.allconv12 = gen_conv(cnum*4, cnum*4, 3, 1, 1)
+ self.allconv13 = gen_conv(cnum*4, cnum*2, 3, 1, 1)
+ self.allconv14 = gen_conv(cnum*2, cnum*2, 3, 1, 1)
+ self.allconv15 = gen_conv(cnum*2, cnum, 3, 1, 1)
+ self.allconv16 = gen_conv(cnum, cnum//2, 3, 1, 1)
+ self.allconv17 = gen_conv(cnum//2, input_dim, 3, 1, 1, activation='none')
+
+ def forward(self, xin, x_stage1, mask):
+ x1_inpaint = x_stage1 * mask + xin * (1. - mask)
+ # For indicating the boundaries of images
+ ones = torch.ones(xin.size(0), 1, xin.size(2), xin.size(3))
+ if self.use_cuda:
+ ones = ones.cuda()
+ mask = mask.cuda()
+ # conv branch
+ xnow = torch.cat([x1_inpaint, ones, mask], dim=1)
+ x = self.conv1(xnow)
+ x = self.conv2_downsample(x)
+ x = self.conv3(x)
+ x = self.conv4_downsample(x)
+ x = self.conv5(x)
+ x = self.conv6(x)
+ x = self.conv7_atrous(x)
+ x = self.conv8_atrous(x)
+ x = self.conv9_atrous(x)
+ x = self.conv10_atrous(x)
+ x_hallu = x
+ # attention branch
+ x = self.pmconv1(xnow)
+ x = self.pmconv2_downsample(x)
+ x = self.pmconv3(x)
+ x = self.pmconv4_downsample(x)
+ x = self.pmconv5(x)
+ x = self.pmconv6(x)
+ x, offset_flow = self.contextul_attention(x, x, mask)
+ x = self.pmconv9(x)
+ x = self.pmconv10(x)
+ pm = x
+ x = torch.cat([x_hallu, pm], dim=1)
+ # merge two branches
+ x = self.allconv11(x)
+ x = self.allconv12(x)
+ x = F.interpolate(x, scale_factor=2, mode='nearest')
+ x = self.allconv13(x)
+ x = self.allconv14(x)
+ x = F.interpolate(x, scale_factor=2, mode='nearest')
+ x = self.allconv15(x)
+ x = self.allconv16(x)
+ x = self.allconv17(x)
+ x_stage2 = torch.clamp(x, -1., 1.)
+
+ return x_stage2, offset_flow
+
+
+class ContextualAttention(nn.Module):
+ def __init__(self, ksize=3, stride=1, rate=1, fuse_k=3, softmax_scale=10,
+ fuse=False, use_cuda=False, device_ids=None):
+ super(ContextualAttention, self).__init__()
+ self.ksize = ksize
+ self.stride = stride
+ self.rate = rate
+ self.fuse_k = fuse_k
+ self.softmax_scale = softmax_scale
+ self.fuse = fuse
+ self.use_cuda = use_cuda
+ self.device_ids = device_ids
+
+ def forward(self, f, b, mask=None):
+ """ Contextual attention layer implementation.
+ Contextual attention is first introduced in publication:
+ Generative Image Inpainting with Contextual Attention, Yu et al.
+ Args:
+ f: Input feature to match (foreground).
+ b: Input feature for match (background).
+ mask: Input mask for b, indicating patches not available.
+ ksize: Kernel size for contextual attention.
+ stride: Stride for extracting patches from b.
+ rate: Dilation for matching.
+ softmax_scale: Scaled softmax for attention.
+ Returns:
+ torch.tensor: output
+ """
+ # get shapes
+ raw_int_fs = list(f.size()) # b*c*h*w
+ raw_int_bs = list(b.size()) # b*c*h*w
+
+ # extract patches from background with stride and rate
+ kernel = 2 * self.rate
+ # raw_w is extracted for reconstruction
+ raw_w = extract_image_patches(b, ksizes=[kernel, kernel],
+ strides=[self.rate*self.stride,
+ self.rate*self.stride],
+ rates=[1, 1],
+ padding='same') # [N, C*k*k, L]
+ # raw_shape: [N, C, k, k, L]
+ raw_w = raw_w.view(raw_int_bs[0], raw_int_bs[1], kernel, kernel, -1)
+ raw_w = raw_w.permute(0, 4, 1, 2, 3) # raw_shape: [N, L, C, k, k]
+ raw_w_groups = torch.split(raw_w, 1, dim=0)
+
+ # downscaling foreground option: downscaling both foreground and
+ # background for matching and use original background for reconstruction.
+ f = F.interpolate(f, scale_factor=1./self.rate, mode='nearest')
+ b = F.interpolate(b, scale_factor=1./self.rate, mode='nearest')
+ int_fs = list(f.size()) # b*c*h*w
+ int_bs = list(b.size())
+ f_groups = torch.split(f, 1, dim=0) # split tensors along the batch dimension
+ # w shape: [N, C*k*k, L]
+ w = extract_image_patches(b, ksizes=[self.ksize, self.ksize],
+ strides=[self.stride, self.stride],
+ rates=[1, 1],
+ padding='same')
+ # w shape: [N, C, k, k, L]
+ w = w.view(int_bs[0], int_bs[1], self.ksize, self.ksize, -1)
+ w = w.permute(0, 4, 1, 2, 3) # w shape: [N, L, C, k, k]
+ w_groups = torch.split(w, 1, dim=0)
+
+ # process mask
+ if mask is None:
+ mask = torch.zeros([int_bs[0], 1, int_bs[2], int_bs[3]])
+ if self.use_cuda:
+ mask = mask.cuda()
+ else:
+ mask = F.interpolate(mask, scale_factor=1./(4*self.rate), mode='nearest')
+ int_ms = list(mask.size())
+ # m shape: [N, C*k*k, L]
+ m = extract_image_patches(mask, ksizes=[self.ksize, self.ksize],
+ strides=[self.stride, self.stride],
+ rates=[1, 1],
+ padding='same')
+ # m shape: [N, C, k, k, L]
+ m = m.view(int_ms[0], int_ms[1], self.ksize, self.ksize, -1)
+ m = m.permute(0, 4, 1, 2, 3) # m shape: [N, L, C, k, k]
+ m = m[0] # m shape: [L, C, k, k]
+ # mm shape: [L, 1, 1, 1]
+ mm = (reduce_mean(m, axis=[1, 2, 3], keepdim=True)==0.).to(torch.float32)
+ mm = mm.permute(1, 0, 2, 3) # mm shape: [1, L, 1, 1]
+
+ y = []
+ offsets = []
+ k = self.fuse_k
+ scale = self.softmax_scale # to fit the PyTorch tensor image value range
+ fuse_weight = torch.eye(k).view(1, 1, k, k) # 1*1*k*k
+ if self.use_cuda:
+ fuse_weight = fuse_weight.cuda()
+
+ for xi, wi, raw_wi in zip(f_groups, w_groups, raw_w_groups):
+ '''
+ O => output channel as a conv filter
+ I => input channel as a conv filter
+ xi : separated tensor along batch dimension of front; (B=1, C=128, H=32, W=32)
+ wi : separated patch tensor along batch dimension of back; (B=1, O=32*32, I=128, KH=3, KW=3)
+ raw_wi : separated tensor along batch dimension of back; (B=1, I=32*32, O=128, KH=4, KW=4)
+ '''
+ # conv for compare
+ escape_NaN = torch.FloatTensor([1e-4])
+ if self.use_cuda:
+ escape_NaN = escape_NaN.cuda()
+ wi = wi[0] # [L, C, k, k]
+ max_wi = torch.sqrt(reduce_sum(torch.pow(wi, 2) + escape_NaN, axis=[1, 2, 3], keepdim=True))
+ wi_normed = wi / max_wi
+ # xi shape: [1, C, H, W], yi shape: [1, L, H, W]
+ xi = same_padding(xi, [self.ksize, self.ksize], [1, 1], [1, 1]) # xi: 1*c*H*W
+ yi = F.conv2d(xi, wi_normed, stride=1) # [1, L, H, W]
+ # conv implementation for fuse scores to encourage large patches
+ if self.fuse:
+ # make all of depth to spatial resolution
+ yi = yi.view(1, 1, int_bs[2]*int_bs[3], int_fs[2]*int_fs[3]) # (B=1, I=1, H=32*32, W=32*32)
+ yi = same_padding(yi, [k, k], [1, 1], [1, 1])
+ yi = F.conv2d(yi, fuse_weight, stride=1) # (B=1, C=1, H=32*32, W=32*32)
+ yi = yi.contiguous().view(1, int_bs[2], int_bs[3], int_fs[2], int_fs[3]) # (B=1, 32, 32, 32, 32)
+ yi = yi.permute(0, 2, 1, 4, 3)
+ yi = yi.contiguous().view(1, 1, int_bs[2]*int_bs[3], int_fs[2]*int_fs[3])
+ yi = same_padding(yi, [k, k], [1, 1], [1, 1])
+ yi = F.conv2d(yi, fuse_weight, stride=1)
+ yi = yi.contiguous().view(1, int_bs[3], int_bs[2], int_fs[3], int_fs[2])
+ yi = yi.permute(0, 2, 1, 4, 3).contiguous()
+ yi = yi.view(1, int_bs[2] * int_bs[3], int_fs[2], int_fs[3]) # (B=1, C=32*32, H=32, W=32)
+ # softmax to match
+ yi = yi * mm
+ yi = F.softmax(yi*scale, dim=1)
+ yi = yi * mm # [1, L, H, W]
+
+ offset = torch.argmax(yi, dim=1, keepdim=True) # 1*1*H*W
+
+ if int_bs != int_fs:
+ # Normalize the offset value to match foreground dimension
+ times = float(int_fs[2] * int_fs[3]) / float(int_bs[2] * int_bs[3])
+ offset = ((offset + 1).float() * times - 1).to(torch.int64)
+ offset = torch.cat([offset//int_fs[3], offset%int_fs[3]], dim=1) # 1*2*H*W
+
+ # deconv for patch pasting
+ wi_center = raw_wi[0]
+ # yi = F.pad(yi, [0, 1, 0, 1]) # here may need conv_transpose same padding
+ yi = F.conv_transpose2d(yi, wi_center, stride=self.rate, padding=1) / 4. # (B=1, C=128, H=64, W=64)
+ y.append(yi)
+ offsets.append(offset)
+
+ y = torch.cat(y, dim=0) # back to the mini-batch
+ y.contiguous().view(raw_int_fs)
+
+ offsets = torch.cat(offsets, dim=0)
+ offsets = offsets.view(int_fs[0], 2, *int_fs[2:])
+
+ # case1: visualize optical flow: minus current position
+ h_add = torch.arange(int_fs[2]).view([1, 1, int_fs[2], 1]).expand(int_fs[0], -1, -1, int_fs[3])
+ w_add = torch.arange(int_fs[3]).view([1, 1, 1, int_fs[3]]).expand(int_fs[0], -1, int_fs[2], -1)
+ ref_coordinate = torch.cat([h_add, w_add], dim=1)
+ if self.use_cuda:
+ ref_coordinate = ref_coordinate.cuda()
+
+ offsets = offsets - ref_coordinate
+ # flow = pt_flow_to_image(offsets)
+
+ flow = torch.from_numpy(flow_to_image(offsets.permute(0, 2, 3, 1).cpu().data.numpy())) / 255.
+ flow = flow.permute(0, 3, 1, 2)
+ if self.use_cuda:
+ flow = flow.cuda()
+ # case2: visualize which pixels are attended
+ # flow = torch.from_numpy(highlight_flow((offsets * mask.long()).cpu().data.numpy()))
+
+ if self.rate != 1:
+ flow = F.interpolate(flow, scale_factor=self.rate*4, mode='nearest')
+
+ return y, flow
+
+
+def test_contextual_attention(args):
+ import cv2
+ import os
+ # run on cpu
+ os.environ['CUDA_VISIBLE_DEVICES'] = '2'
+
+ def float_to_uint8(img):
+ img = img * 255
+ return img.astype('uint8')
+
+ rate = 2
+ stride = 1
+ grid = rate*stride
+
+ b = default_loader(args.imageA)
+ w, h = b.size
+ b = b.resize((w//grid*grid//2, h//grid*grid//2), Image.ANTIALIAS)
+ # b = b.resize((w//grid*grid, h//grid*grid), Image.ANTIALIAS)
+ print('Size of imageA: {}'.format(b.size))
+
+ f = default_loader(args.imageB)
+ w, h = f.size
+ f = f.resize((w//grid*grid, h//grid*grid), Image.ANTIALIAS)
+ print('Size of imageB: {}'.format(f.size))
+
+ f, b = transforms.ToTensor()(f), transforms.ToTensor()(b)
+ f, b = f.unsqueeze(0), b.unsqueeze(0)
+ if torch.cuda.is_available():
+ f, b = f.cuda(), b.cuda()
+
+ contextual_attention = ContextualAttention(ksize=3, stride=stride, rate=rate, fuse=True)
+
+ if torch.cuda.is_available():
+ contextual_attention = contextual_attention.cuda()
+
+ yt, flow_t = contextual_attention(f, b)
+ vutils.save_image(yt, 'vutils' + args.imageOut, normalize=True)
+ vutils.save_image(flow_t, 'flow' + args.imageOut, normalize=True)
+ # y = tensor_img_to_npimg(yt.cpu()[0])
+ # flow = tensor_img_to_npimg(flow_t.cpu()[0])
+ # cv2.imwrite('flow' + args.imageOut, flow_t)
+
+
+class LocalDis(nn.Module):
+ def __init__(self, config, use_cuda=True, device_ids=None):
+ super(LocalDis, self).__init__()
+ self.input_dim = config['input_dim']
+ self.cnum = config['ndf']
+ self.use_cuda = use_cuda
+ self.device_ids = device_ids
+
+ self.dis_conv_module = DisConvModule(self.input_dim, self.cnum)
+ self.linear = nn.Linear(self.cnum*4*8*8, 1)
+
+ def forward(self, x):
+ x = self.dis_conv_module(x)
+ x = x.view(x.size()[0], -1)
+ x = self.linear(x)
+
+ return x
+
+
+class GlobalDis(nn.Module):
+ def __init__(self, config, use_cuda=True, device_ids=None):
+ super(GlobalDis, self).__init__()
+ self.input_dim = config['input_dim']
+ self.cnum = config['ndf']
+ self.use_cuda = use_cuda
+ self.device_ids = device_ids
+
+ self.dis_conv_module = DisConvModule(self.input_dim, self.cnum)
+ self.linear = nn.Linear(self.cnum*4*16*16, 1)
+
+ def forward(self, x):
+ x = self.dis_conv_module(x)
+ x = x.view(x.size()[0], -1)
+ x = self.linear(x)
+
+ return x
+
+
+class DisConvModule(nn.Module):
+ def __init__(self, input_dim, cnum, use_cuda=True, device_ids=None):
+ super(DisConvModule, self).__init__()
+ self.use_cuda = use_cuda
+ self.device_ids = device_ids
+
+ self.conv1 = dis_conv(input_dim, cnum, 5, 2, 2)
+ self.conv2 = dis_conv(cnum, cnum*2, 5, 2, 2)
+ self.conv3 = dis_conv(cnum*2, cnum*4, 5, 2, 2)
+ self.conv4 = dis_conv(cnum*4, cnum*4, 5, 2, 2)
+
+ def forward(self, x):
+ x = self.conv1(x)
+ x = self.conv2(x)
+ x = self.conv3(x)
+ x = self.conv4(x)
+
+ return x
+
+
+def gen_conv(input_dim, output_dim, kernel_size=3, stride=1, padding=0, rate=1,
+ activation='elu'):
+ return Conv2dBlock(input_dim, output_dim, kernel_size, stride,
+ conv_padding=padding, dilation=rate,
+ activation=activation)
+
+
+def dis_conv(input_dim, output_dim, kernel_size=5, stride=2, padding=0, rate=1,
+ activation='lrelu'):
+ return Conv2dBlock(input_dim, output_dim, kernel_size, stride,
+ conv_padding=padding, dilation=rate,
+ activation=activation)
+
+
+class Conv2dBlock(nn.Module):
+ def __init__(self, input_dim, output_dim, kernel_size, stride, padding=0,
+ conv_padding=0, dilation=1, weight_norm='none', norm='none',
+ activation='relu', pad_type='zero', transpose=False):
+ super(Conv2dBlock, self).__init__()
+ self.use_bias = True
+ # initialize padding
+ if pad_type == 'reflect':
+ self.pad = nn.ReflectionPad2d(padding)
+ elif pad_type == 'replicate':
+ self.pad = nn.ReplicationPad2d(padding)
+ elif pad_type == 'zero':
+ self.pad = nn.ZeroPad2d(padding)
+ elif pad_type == 'none':
+ self.pad = None
+ else:
+ assert 0, "Unsupported padding type: {}".format(pad_type)
+
+ # initialize normalization
+ norm_dim = output_dim
+ if norm == 'bn':
+ self.norm = nn.BatchNorm2d(norm_dim)
+ elif norm == 'in':
+ self.norm = nn.InstanceNorm2d(norm_dim)
+ elif norm == 'none':
+ self.norm = None
+ else:
+ assert 0, "Unsupported normalization: {}".format(norm)
+
+ if weight_norm == 'sn':
+ self.weight_norm = spectral_norm_fn
+ elif weight_norm == 'wn':
+ self.weight_norm = weight_norm_fn
+ elif weight_norm == 'none':
+ self.weight_norm = None
+ else:
+ assert 0, "Unsupported normalization: {}".format(weight_norm)
+
+ # initialize activation
+ if activation == 'relu':
+ self.activation = nn.ReLU(inplace=True)
+ elif activation == 'elu':
+ self.activation = nn.ELU(inplace=True)
+ elif activation == 'lrelu':
+ self.activation = nn.LeakyReLU(0.2, inplace=True)
+ elif activation == 'prelu':
+ self.activation = nn.PReLU()
+ elif activation == 'selu':
+ self.activation = nn.SELU(inplace=True)
+ elif activation == 'tanh':
+ self.activation = nn.Tanh()
+ elif activation == 'none':
+ self.activation = None
+ else:
+ assert 0, "Unsupported activation: {}".format(activation)
+
+ # initialize convolution
+ if transpose:
+ self.conv = nn.ConvTranspose2d(input_dim, output_dim,
+ kernel_size, stride,
+ padding=conv_padding,
+ output_padding=conv_padding,
+ dilation=dilation,
+ bias=self.use_bias)
+ else:
+ self.conv = nn.Conv2d(input_dim, output_dim, kernel_size, stride,
+ padding=conv_padding, dilation=dilation,
+ bias=self.use_bias)
+
+ if self.weight_norm:
+ self.conv = self.weight_norm(self.conv)
+
+ def forward(self, x):
+ if self.pad:
+ x = self.conv(self.pad(x))
+ else:
+ x = self.conv(x)
+ if self.norm:
+ x = self.norm(x)
+ if self.activation:
+ x = self.activation(x)
+ return x
+
+
+
+if __name__ == "__main__":
+ import argparse
+ parser = argparse.ArgumentParser()
+ parser.add_argument('--imageA', default='', type=str, help='Image A as background patches to reconstruct image B.')
+ parser.add_argument('--imageB', default='', type=str, help='Image B is reconstructed with image A.')
+ parser.add_argument('--imageOut', default='result.png', type=str, help='Image B is reconstructed with image A.')
+ args = parser.parse_args()
+ test_contextual_attention(args)
diff --git a/only_gradio_server.py b/only_gradio_server.py
new file mode 100644
index 0000000000000000000000000000000000000000..94b8214d893d0ac6204186b8480620e6eced92f0
--- /dev/null
+++ b/only_gradio_server.py
@@ -0,0 +1,188 @@
+import os
+import base64
+import io
+import uuid
+from ultralytics import YOLO
+import cv2
+import torch
+import numpy as np
+from PIL import Image
+from torchvision import transforms
+import imageio.v2 as imageio
+from trainer import Trainer
+from utils.tools import get_config
+import torch.nn.functional as F
+from iopaint.single_processing import batch_inpaint_cv2
+from pathlib import Path
+
+# set current working directory cache instead of default
+os.environ["TORCH_HOME"] = "./pretrained-model"
+os.environ["HUGGINGFACE_HUB_CACHE"] = "./pretrained-model"
+
+def resize_image(input_image_path, width=640, height=640):
+ """Resizes an image from image data and returns the resized image."""
+ try:
+ # Read the image using cv2.imread
+ img = cv2.imread(input_image_path, cv2.IMREAD_COLOR)
+
+ # Resize while maintaining the aspect ratio
+ shape = img.shape[:2] # current shape [height, width]
+ new_shape = (width, height) # the shape to resize to
+
+ # Scale ratio (new / old)
+ r = min(new_shape[0] / shape[0], new_shape[1] / shape[1])
+ ratio = r, r # width, height ratios
+ new_unpad = int(round(shape[1] * r)), int(round(shape[0] * r))
+
+ # Resize the image
+ im = cv2.resize(img, new_unpad, interpolation=cv2.INTER_LINEAR)
+
+ # Pad the image
+ color = (114, 114, 114) # color used for padding
+ dw, dh = new_shape[1] - new_unpad[0], new_shape[0] - new_unpad[1] # wh padding
+ # divide padding into 2 sides
+ dw /= 2
+ dh /= 2
+ # compute padding on all corners
+ top, bottom = int(round(dh - 0.1)), int(round(dh + 0.1))
+ left, right = int(round(dw - 0.1)), int(round(dw + 0.1))
+ im = cv2.copyMakeBorder(im, top, bottom, left, right, cv2.BORDER_CONSTANT, value=color) # add border
+ return im
+
+ except Exception as e:
+ print(f"Error resizing image: {e}")
+ return None # Or handle differently as needed
+
+
+def load_weights(path, device):
+ model_weights = torch.load(path)
+ return {
+ k: v.to(device)
+ for k, v in model_weights.items()
+ }
+
+
+# Function to convert image to base64
+def convert_image_to_base64(image):
+ # Convert image to bytes
+ _, buffer = cv2.imencode('.png', image)
+ # Convert bytes to base64
+ image_base64 = base64.b64encode(buffer).decode('utf-8')
+ return image_base64
+
+
+def convert_to_base64(image):
+ # Read the image file as binary data
+ image_data = image.read()
+ # Encode the binary data as base64
+ base64_encoded = base64.b64encode(image_data).decode('utf-8')
+ return base64_encoded
+
+def convert_to_base64_file(image):
+ # Convert the image to binary data
+ image_data = cv2.imencode('.png', image)[1].tobytes()
+ # Encode the binary data as base64
+ base64_encoded = base64.b64encode(image_data).decode('utf-8')
+ return base64_encoded
+
+
+def process_images(input_image, append_image, default_class="chair"):
+ # Static paths
+ config_path = Path('configs/config.yaml')
+ model_path = Path('pretrained-model/torch_model.p')
+
+ # Resize input image and get base64 data of resized image
+ img = resize_image(input_image)
+
+ if img is None:
+ return {'error': 'Failed to decode resized image'}, 419
+
+ H, W, _ = img.shape
+ x_point = 0
+ y_point = 0
+ width = 1
+ height = 1
+
+ # Load a model
+ model = YOLO('pretrained-model/yolov8m-seg.pt') # pretrained YOLOv8m-seg model
+
+ # Run batched inference on a list of images
+ results = model(img, imgsz=(W,H), conf=0.5) # chair class 56 with confidence >= 0.5
+ names = model.names
+
+ class_found = False
+ for result in results:
+ for i, label in enumerate(result.boxes.cls):
+ # Check if the label matches the chair label
+ if names[int(label)] == default_class:
+ class_found = True
+ # Convert the tensor to a numpy array
+ chair_mask_np = result.masks.data[i].numpy()
+
+ kernel = np.ones((5, 5), np.uint8) # Create a 5x5 kernel for dilation
+ chair_mask_np = cv2.dilate(chair_mask_np, kernel, iterations=2) # Apply dilation
+
+ # Find contours to get bounding box
+ contours, _ = cv2.findContours((chair_mask_np == 1).astype(np.uint8), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
+
+ # Iterate over contours to find the bounding box of each object
+ for contour in contours:
+ x, y, w, h = cv2.boundingRect(contour)
+ x_point = x
+ y_point = y
+ width = w
+ height = h
+
+ # Get the corresponding mask
+ mask = result.masks.data[i].numpy() * 255
+ dilated_mask = cv2.dilate(mask, kernel, iterations=2) # Apply dilation
+ # Resize the mask to match the dimensions of the original image
+ resized_mask = cv2.resize(dilated_mask, (img.shape[1], img.shape[0]))
+
+ # call repainting and merge function
+ output_base64 = repaitingAndMerge(append_image,str(model_path), str(config_path),width, height, x_point, y_point, img, resized_mask)
+ # Return the output base64 image in the API response
+ return output_base64
+
+ # return class not found in prediction
+ if not class_found:
+ return {'message': f'{default_class} object not found in the image'}, 200
+
+def repaitingAndMerge(append_image_path, model_path, config_path, width, height, xposition, yposition, input_base, mask_base):
+ config = get_config(config_path)
+ device = torch.device("cpu")
+ trainer = Trainer(config)
+ trainer.load_state_dict(load_weights(model_path, device), strict=False)
+ trainer.eval()
+
+ # lama inpainting start
+ print("lama inpainting start")
+ inpaint_result_np = batch_inpaint_cv2('lama', 'cpu', input_base, mask_base)
+ print("lama inpainting end")
+
+ # Create PIL Image from NumPy array
+ final_image = Image.fromarray(inpaint_result_np)
+
+ print("merge start")
+
+ # Load the append image using cv2.imread
+ append_image = cv2.imread(append_image_path, cv2.IMREAD_UNCHANGED)
+ cv2.imwrite('appneded-image.png',append_image)
+ # Resize the append image while preserving transparency
+ resized_image = cv2.resize(append_image, (width, height), interpolation=cv2.INTER_AREA)
+ # Convert the resized image to RGBA format (assuming it's in BGRA format)
+ resized_image = cv2.cvtColor(resized_image, cv2.COLOR_BGRA2RGBA)
+ # Create a PIL Image from the resized image with transparent background
+ append_image_pil = Image.fromarray(resized_image)
+
+ # Paste the append image onto the final image
+ final_image.paste(append_image_pil, (xposition, yposition), append_image_pil)
+ # Save the resulting image
+ print("merge end")
+
+ # Convert the final image to base64
+ with io.BytesIO() as output_buffer:
+ final_image.save(output_buffer, format='PNG')
+ output_numpy = np.array(final_image)
+
+ return output_numpy
diff --git a/templates/index.html b/templates/index.html
new file mode 100644
index 0000000000000000000000000000000000000000..3a6da75a50209552eb8d954ce2d43d6cdaee1451
--- /dev/null
+++ b/templates/index.html
@@ -0,0 +1,145 @@
+
+
+
+
+
+ Object to Object Replace
+
+
+
+ Object to Object Replace
+
+
+
+
+
+
+
+
diff --git a/utils/logger.py b/utils/logger.py
new file mode 100644
index 0000000000000000000000000000000000000000..3a92be661c3da689450b86fb0861542c2dc92a96
--- /dev/null
+++ b/utils/logger.py
@@ -0,0 +1,37 @@
+import os
+import sys
+import datetime
+import logging
+
+
+def date_uid():
+ """Generate a unique id based on date.
+
+ Returns:
+ str: Return uid string, e.g. '20171122171307111552'.
+
+ """
+ return str(datetime.datetime.now()).replace('-', '') \
+ .replace(' ', '').replace(':', '').replace('.', '')
+
+
+def get_logger(checkpoint_path=None):
+ """
+ Get the root logger
+ :param checkpoint_path: only specify this when the first time call it
+ :return: the root logger
+ """
+ if checkpoint_path:
+ logger = logging.getLogger()
+ formatter = logging.Formatter('%(asctime)s %(levelname)s %(message)s')
+ stream_hdlr = logging.StreamHandler(sys.stdout)
+ log_filename = date_uid()
+ file_hdlr = logging.FileHandler(os.path.join(checkpoint_path, log_filename + '.log'))
+ stream_hdlr.setFormatter(formatter)
+ file_hdlr.setFormatter(formatter)
+ logger.addHandler(stream_hdlr)
+ logger.addHandler(file_hdlr)
+ logger.setLevel(logging.INFO)
+ else:
+ logger = logging.getLogger()
+ return logger