Add files using upload-large-folder tool

ad.ipynb

@@ -0,0 +1,517 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "c7dbfe5c",
+   "metadata": {},
+   "source": [
+    "## Texture Synthesis"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "8a06ff4b",
+   "metadata": {},
+   "source": [
+    "### Texture Synthesis via Optimization"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "6349220d",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import os\n",
+    "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"0\"  \n",
+    "from accelerate.utils import set_seed\n",
+    "from diffusers import AutoencoderKL, DDIMScheduler, UNet2DConditionModel\n",
+    "from pipeline_sd import ADPipeline\n",
+    "from utils import *\n",
+    "\n",
+    "model_name = \"/root/models/stable-diffusion-v1-5\"\n",
+    "vae = \"\"\n",
+    "lr = 0.05\n",
+    "iters = 1\n",
+    "seed = 42\n",
+    "width, height = 512, 512\n",
+    "weight = 0\n",
+    "batch_size = 3\n",
+    "mixed_precision = \"bf16\"\n",
+    "num_inference_steps = 300\n",
+    "enable_gradient_checkpoint = False\n",
+    "start_layer, end_layer = 10, 16\n",
+    "\n",
+    "\n",
+    "style_image = [\"./data/texture/4.jpg\"]\n",
+    "content_image = \"\"\n",
+    "\n",
+    "\n",
+    "scheduler = DDIMScheduler.from_pretrained(model_name, subfolder=\"scheduler\")\n",
+    "pipe = ADPipeline.from_pretrained(\n",
+    "    model_name, scheduler=scheduler, safety_checker=None\n",
+    ")\n",
+    "if vae != \"\":\n",
+    "    vae = AutoencoderKL.from_pretrained(vae)\n",
+    "    pipe.vae = vae\n",
+    "pipe.classifier = pipe.unet\n",
+    "set_seed(seed)\n",
+    "\n",
+    "style_image = torch.cat([load_image(path, size=(512, 512)) for path in style_image])\n",
+    "rec_style_image = pipe.latent2image(pipe.image2latent(style_image))\n",
+    "if content_image == \"\":\n",
+    "    content_image = None\n",
+    "else:\n",
+    "    content_image = load_image(content_image, size=(width, height))\n",
+    "controller = Controller(self_layers=(start_layer, end_layer))\n",
+    "result = pipe.optimize(\n",
+    "    lr=lr,\n",
+    "    batch_size=batch_size,\n",
+    "    iters=iters,\n",
+    "    width=width,\n",
+    "    height=height,\n",
+    "    weight=weight,\n",
+    "    controller=controller,\n",
+    "    style_image=style_image,\n",
+    "    content_image=content_image,\n",
+    "    mixed_precision=mixed_precision,\n",
+    "    num_inference_steps=num_inference_steps,\n",
+    "    enable_gradient_checkpoint=enable_gradient_checkpoint,\n",
+    ")\n",
+    "\n",
+    "save_image(style_image, \"style.png\")\n",
+    "save_image(result, \"output.png\")\n",
+    "show_image(\"style.png\", title=\"style image\")\n",
+    "show_image(\"output.png\", title=\"generated\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "c4ba9b1d",
+   "metadata": {},
+   "source": [
+    "### Texture Synthesis via Sample"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "67a535c1",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import os\n",
+    "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"0\"  \n",
+    "from accelerate.utils import set_seed\n",
+    "from diffusers import AutoencoderKL, DDIMScheduler\n",
+    "from pipeline_sd import ADPipeline\n",
+    "from utils import *\n",
+    "\n",
+    "model_name = \"/root/models/stable-diffusion-v1-5\"\n",
+    "vae = \"\"\n",
+    "lr = 0.05\n",
+    "iters =3\n",
+    "seed = 42\n",
+    "width, height = 512, 512\n",
+    "weight = 0.\n",
+    "mixed_precision = \"bf16\"\n",
+    "num_inference_steps = 50\n",
+    "guidance_scale = 1\n",
+    "num_images_per_prompt = 3\n",
+    "enable_gradient_checkpoint = False\n",
+    "start_layer, end_layer = 10, 16\n",
+    "\n",
+    "\n",
+    "style_image = [\"./data/texture/8.jpg\"]\n",
+    "content_image = \"\"\n",
+    "\n",
+    "scheduler = DDIMScheduler.from_pretrained(model_name, subfolder=\"scheduler\")\n",
+    "pipe = ADPipeline.from_pretrained(model_name, scheduler=scheduler, safety_checker=None)\n",
+    "if vae != \"\":\n",
+    "    vae = AutoencoderKL.from_pretrained(vae)\n",
+    "    pipe.vae = vae\n",
+    "pipe.classifier = pipe.unet\n",
+    "set_seed(seed)\n",
+    "\n",
+    "style_image = torch.cat([load_image(path, size=(512, 512)) for path in style_image])\n",
+    "rec_style_image = pipe.latent2image(pipe.image2latent(style_image))\n",
+    "if content_image == \"\":\n",
+    "    content_image = None\n",
+    "else:\n",
+    "    content_image = load_image(content_image, size=(width, height))\n",
+    "controller = Controller(self_layers=(start_layer, end_layer))\n",
+    "\n",
+    "result = pipe.sample(\n",
+    "    lr=lr,\n",
+    "    adain=False,\n",
+    "    iters=iters,\n",
+    "    width=width,\n",
+    "    height=height,\n",
+    "    weight=weight,\n",
+    "    controller=controller,\n",
+    "    style_image=style_image,\n",
+    "    content_image=content_image,\n",
+    "    prompt=\"\",\n",
+    "    negative_prompt=\"\",\n",
+    "    mixed_precision=mixed_precision,\n",
+    "    num_inference_steps=num_inference_steps,\n",
+    "    guidance_scale=guidance_scale,\n",
+    "    num_images_per_prompt=num_images_per_prompt,\n",
+    "    enable_gradient_checkpoint=enable_gradient_checkpoint,\n",
+    ")\n",
+    "\n",
+    "save_image(style_image, \"style.png\")\n",
+    "save_image(result, \"output.png\")\n",
+    "show_image(\"style.png\", title=\"style image\")\n",
+    "show_image(\"output.png\", title=\"generated\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "c69ae623",
+   "metadata": {},
+   "source": [
+    "### Texture Synthesis via MultiDiffusion"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "6d059173",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import os\n",
+    "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"0\" \n",
+    "from accelerate.utils import set_seed\n",
+    "from diffusers import AutoencoderKL, DDIMScheduler\n",
+    "from pipeline_sd import ADPipeline\n",
+    "from utils import *\n",
+    "\n",
+    "model_name = \"/root/models/stable-diffusion-v1-5\"\n",
+    "vae = \"\"\n",
+    "lr = 0.05\n",
+    "iters = 2\n",
+    "seed = 42\n",
+    "width, height = 512*2, 512\n",
+    "weight = 0.0\n",
+    "mixed_precision = \"bf16\"\n",
+    "num_inference_steps = 50\n",
+    "guidance_scale = 1\n",
+    "num_images_per_prompt = 1\n",
+    "enable_gradient_checkpoint = False\n",
+    "start_layer, end_layer = 10, 16\n",
+    "\n",
+    "\n",
+    "style_image = [\"./data/texture/17.jpg\"]\n",
+    "content_image = \"\"\n",
+    "\n",
+    "scheduler = DDIMScheduler.from_pretrained(model_name, subfolder=\"scheduler\")\n",
+    "pipe = ADPipeline.from_pretrained(model_name, scheduler=scheduler, safety_checker=None)\n",
+    "if vae != \"\":\n",
+    "    vae = AutoencoderKL.from_pretrained(vae)\n",
+    "    pipe.vae = vae\n",
+    "\n",
+    "pipe.classifier = pipe.unet\n",
+    "set_seed(seed)\n",
+    "\n",
+    "style_image = torch.cat([load_image(path, size=(512, 512)) for path in style_image])\n",
+    "if content_image == \"\":\n",
+    "    content_image = None\n",
+    "else:\n",
+    "    content_image = load_image(content_image, size=(width, height))\n",
+    "controller = Controller(self_layers=(start_layer, end_layer))\n",
+    "\n",
+    "result = pipe.panorama(\n",
+    "    lr=lr,\n",
+    "    iters=iters,\n",
+    "    width=width,\n",
+    "    height=height,\n",
+    "    weight=weight,\n",
+    "    controller=controller,\n",
+    "    style_image=style_image,\n",
+    "    content_image=content_image,\n",
+    "    prompt=\"\",\n",
+    "    negative_prompt=\"\",\n",
+    "    stride=8,\n",
+    "    view_batch_size=8,\n",
+    "    mixed_precision=mixed_precision,\n",
+    "    num_inference_steps=num_inference_steps,\n",
+    "    guidance_scale=guidance_scale,\n",
+    "    num_images_per_prompt=num_images_per_prompt,\n",
+    "    enable_gradient_checkpoint=enable_gradient_checkpoint,\n",
+    ")\n",
+    "\n",
+    "save_image(style_image, \"style.png\")\n",
+    "save_image(result, \"output.png\")\n",
+    "show_image(\"style.png\", title=\"style image\")\n",
+    "show_image(\"output.png\", title=\"generated\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "25e4e702",
+   "metadata": {},
+   "source": [
+    "## Style/Appearance Transfer"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "4badee8f",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import os\n",
+    "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"0\" \n",
+    "from accelerate.utils import set_seed\n",
+    "from diffusers import AutoencoderKL, DDIMScheduler\n",
+    "from pipeline_sd import ADPipeline\n",
+    "from utils import *\n",
+    "\n",
+    "model_name = \"/root/models/stable-diffusion-v1-5\"\n",
+    "vae = \"\"\n",
+    "lr = 0.05\n",
+    "iters = 1\n",
+    "seed = 42\n",
+    "width = 512\n",
+    "height = 512\n",
+    "weight = 0.25\n",
+    "batch_size = 1\n",
+    "mixed_precision = \"bf16\"\n",
+    "num_inference_steps = 200\n",
+    "guidance_scale = 1\n",
+    "num_images_per_prompt = 1\n",
+    "enable_gradient_checkpoint = False\n",
+    "start_layer, end_layer = 10, 16\n",
+    "\n",
+    "\n",
+    "style_image = [\"./data/style/12.jpg\"]\n",
+    "content_image = \"./data/content/deer.jpg\"\n",
+    "\n",
+    "\n",
+    "scheduler = DDIMScheduler.from_pretrained(model_name, subfolder=\"scheduler\")\n",
+    "pipe = ADPipeline.from_pretrained(\n",
+    "    model_name, scheduler=scheduler, safety_checker=None\n",
+    ")\n",
+    "if vae != \"\":\n",
+    "    vae = AutoencoderKL.from_pretrained(vae)\n",
+    "    pipe.vae = vae\n",
+    "\n",
+    "pipe.classifier = pipe.unet\n",
+    "set_seed(seed)\n",
+    "\n",
+    "style_image = torch.cat([load_image(path, size=(512, 512)) for path in style_image])\n",
+    "if content_image == \"\":\n",
+    "    content_image = None\n",
+    "else:\n",
+    "    content_image = load_image(content_image, size=(width, height))\n",
+    "controller = Controller(self_layers=(start_layer, end_layer))\n",
+    "result = pipe.optimize(\n",
+    "    lr=lr,\n",
+    "    batch_size=batch_size,\n",
+    "    iters=iters,\n",
+    "    width=width,\n",
+    "    height=height,\n",
+    "    weight=weight,\n",
+    "    controller=controller,\n",
+    "    style_image=style_image,\n",
+    "    content_image=content_image,\n",
+    "    mixed_precision=mixed_precision,\n",
+    "    num_inference_steps=num_inference_steps,\n",
+    "    enable_gradient_checkpoint=enable_gradient_checkpoint,\n",
+    ")\n",
+    "\n",
+    "save_image(style_image, \"style.png\")\n",
+    "save_image(content_image, \"content.png\")\n",
+    "save_image(result, \"output.png\")\n",
+    "show_image(\"style.png\", title=\"style image\")\n",
+    "show_image(\"content.png\", title=\"content image\")\n",
+    "show_image(\"output.png\", title=\"generated\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "088ca839",
+   "metadata": {},
+   "source": [
+    "## Style-specific T2I Generation "
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "02efdafd",
+   "metadata": {},
+   "source": [
+    "### Style-specific T2I Generation with SD1.5"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "15c9fb96",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import os\n",
+    "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"0\"  \n",
+    "from pipeline_sd import ADPipeline\n",
+    "from diffusers import DDIMScheduler, AutoencoderKL\n",
+    "import torch\n",
+    "from utils import *\n",
+    "from accelerate.utils import set_seed\n",
+    "\n",
+    "\n",
+    "model_name = \"/root/models/stable-diffusion-v1-5\"\n",
+    "vae = \"\"\n",
+    "lr = 0.015\n",
+    "iters = 3\n",
+    "seed = 42\n",
+    "mixed_precision = \"bf16\"\n",
+    "num_inference_steps = 50\n",
+    "guidance_scale = 7.5\n",
+    "num_images_per_prompt = 3\n",
+    "enable_gradient_checkpoint = False\n",
+    "start_layer, end_layer = 10, 16\n",
+    "\n",
+    "prompt = \"A deer\"\n",
+    "style_image = [\"./data/style/1.jpg\"]\n",
+    "\n",
+    "scheduler = DDIMScheduler.from_pretrained(model_name, subfolder=\"scheduler\")\n",
+    "pipe = ADPipeline.from_pretrained(\n",
+    "    model_name, scheduler=scheduler, safety_checker=None\n",
+    ")\n",
+    "if vae != \"\":\n",
+    "    vae = AutoencoderKL.from_pretrained(vae)\n",
+    "    pipe.vae = vae\n",
+    "\n",
+    "pipe.classifier = pipe.unet\n",
+    "set_seed(seed)\n",
+    "\n",
+    "style_image = torch.cat([load_image(path, size=(512, 512)) for path in style_image])\n",
+    "controller = Controller(self_layers=(start_layer, end_layer))\n",
+    "\n",
+    "result = pipe.sample(\n",
+    "    controller=controller,\n",
+    "    iters=iters,\n",
+    "    lr=lr,\n",
+    "    adain=True,\n",
+    "    height=512,\n",
+    "    width=512,\n",
+    "    mixed_precision=\"bf16\",\n",
+    "    style_image=style_image,\n",
+    "    prompt=prompt,\n",
+    "    negative_prompt=\"\",\n",
+    "    guidance_scale=guidance_scale,\n",
+    "    num_inference_steps=num_inference_steps,\n",
+    "    num_images_per_prompt=num_images_per_prompt,\n",
+    "    enable_gradient_checkpoint=enable_gradient_checkpoint\n",
+    ")\n",
+    "\n",
+    "save_image(style_image, \"style.png\")\n",
+    "save_image(result, \"output.png\")\n",
+    "show_image(\"style.png\", title=\"style image\")\n",
+    "show_image(\"output.png\", title=prompt)\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "dd75eac7",
+   "metadata": {},
+   "source": [
+    "### Style-specific T2I Generation with SDXL"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "1541fd6b",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import os\n",
+    "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"0\"  \n",
+    "from pipeline_sdxl import ADPipeline\n",
+    "from diffusers import DDIMScheduler, AutoencoderKL\n",
+    "import torch\n",
+    "from utils import *\n",
+    "from accelerate.utils import set_seed\n",
+    "\n",
+    "\n",
+    "model_name = \"/root/models/stable-diffusion-xl-base-1.0/\"\n",
+    "vae = \"\"\n",
+    "lr = 0.015\n",
+    "iters = 2\n",
+    "seed = 42\n",
+    "mixed_precision = \"bf16\"\n",
+    "num_inference_steps = 50\n",
+    "guidance_scale = 7\n",
+    "num_images_per_prompt = 5\n",
+    "enable_gradient_checkpoint = True\n",
+    "start_layer, end_layer = 64, 70\n",
+    "\n",
+    "prompt = \"A rocket\"\n",
+    "style_image = [\"./data/style/1.png\"]\n",
+    "\n",
+    "scheduler = DDIMScheduler.from_pretrained(model_name, subfolder=\"scheduler\")\n",
+    "pipe = ADPipeline.from_pretrained(\n",
+    "    model_name, scheduler=scheduler, safety_checker=None\n",
+    ")\n",
+    "if vae != \"\":\n",
+    "    vae = AutoencoderKL.from_pretrained(vae)\n",
+    "    pipe.vae = vae\n",
+    "\n",
+    "pipe.classifier = pipe.unet\n",
+    "set_seed(seed)\n",
+    "\n",
+    "style_image = torch.cat([load_image(path, size=(1024, 1024)) for path in style_image])\n",
+    "controller = Controller(self_layers=(start_layer, end_layer))\n",
+    "\n",
+    "result = pipe.sample(\n",
+    "    controller=controller,\n",
+    "    iters=iters,\n",
+    "    lr=lr,\n",
+    "    adain=True,\n",
+    "    height=1024,\n",
+    "    width=1024,\n",
+    "    mixed_precision=\"bf16\",\n",
+    "    style_image=style_image,\n",
+    "    prompt=prompt,\n",
+    "    negative_prompt=\"\",\n",
+    "    guidance_scale=guidance_scale,\n",
+    "    num_inference_steps=num_inference_steps,\n",
+    "    num_images_per_prompt=num_images_per_prompt,\n",
+    "    enable_gradient_checkpoint=enable_gradient_checkpoint\n",
+    ")\n",
+    "\n",
+    "save_image(style_image, \"style.png\")\n",
+    "save_image(result, \"output.png\")\n",
+    "show_image(\"style.png\", title=\"style image\")\n",
+    "show_image(\"output.png\", title=prompt)\n"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "ad",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.10.16"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

checkpoints/ostracoda_cyclegan/test_opt.txt

@@ -0,0 +1,42 @@
+----------------- Options ---------------
+             aspect_ratio: 1.0                           
+               batch_size: 1                             
+          checkpoints_dir: ./checkpoints                 
+                crop_size: 256                           
+                 dataroot: None                          
+             dataset_mode: single                        
+                direction: AtoB                          
+          display_winsize: 256                           
+                    epoch: latest                        
+                     eval: False                         
+                  gpu_ids: 0                             
+                init_gain: 0.02                          
+                init_type: normal                        
+                 input_nc: 3                             
+                  isTrain: False                         	[default: None]
+                load_iter: 0                             	[default: 0]
+                load_size: 256                           
+         max_dataset_size: inf                           
+                    model: test                          
+             model_suffix:                               
+               n_layers_D: 3                             
+                     name: ostracoda_cyclegan            
+                      ndf: 64                            
+                     netD: basic                         
+                     netG: resnet_9blocks                
+                      ngf: 64                            
+               no_dropout: False                         
+                  no_flip: False                         
+                     norm: instance                      
+                 num_test: 100                           
+              num_threads: 4                             
+                output_nc: 3                             
+                    phase: test                          
+               preprocess: resize_and_crop               
+              results_dir: ./results/                    
+           serial_batches: False                         
+                   suffix:                               
+                use_wandb: False                         
+                  verbose: False                         
+       wandb_project_name: CycleGAN-and-pix2pix          
+----------------- End -------------------

configs/config.yaml

@@ -0,0 +1,52 @@
+# data parameters
+dataset_name: imagenet
+data_with_subfolder: True
+
+train_data_path: traindata/train
+val_data_path: traindata/val
+resume: checkpoints/imagenet/hole_benchmark
+
+
+batch_size: 4
+image_shape: [256, 256, 3]
+mask_shape: [128, 128]
+mask_batch_same: True
+max_delta_shape: [32, 32]
+margin: [0, 0]
+discounted_mask: True
+spatial_discounting_gamma: 0.9
+random_crop: True
+mask_type: hole     # hole | mosaic
+mosaic_unit_size: 12
+
+# training parameters
+expname: benchmark
+cuda: Ture
+gpu_ids: [0]    # set the GPU ids to use, e.g. [0] or [1, 2]
+num_workers: 4
+lr: 0.0001
+beta1: 0.5
+beta2: 0.9
+n_critic: 5
+niter: 480000
+print_iter: 100
+viz_iter: 1000
+viz_max_out: 16
+snapshot_save_iter: 5000
+
+# loss weight
+coarse_l1_alpha: 1.2
+l1_loss_alpha: 1.2
+ae_loss_alpha: 1.2
+global_wgan_loss_alpha: 1.
+gan_loss_alpha: 0.001
+wgan_gp_lambda: 10
+
+# network parameters
+netG:
+  input_dim: 3
+  ngf: 32
+
+netD:
+  input_dim: 3
+  ndf: 64

cyclegan_model/data/__init__.py

@@ -0,0 +1,89 @@
+"""This package includes all the modules related to data loading and preprocessing
+
+ To add a custom dataset class called 'dummy', you need to add a file called 'dummy_dataset.py' and define a subclass 'DummyDataset' inherited from BaseDataset.
+ You need to implement four functions:
+    -- <__init__>:                      initialize the class, first call BaseDataset.__init__(self, opt).
+    -- <__len__>:                       return the size of dataset.
+    -- <__getitem__>:                   get a data point from data loader.
+    -- <modify_commandline_options>:    (optionally) add dataset-specific options and set default options.
+
+Now you can use the dataset class by specifying flag '--dataset_mode dummy'.
+See our template dataset class 'template_dataset.py' for more details.
+"""
+import importlib
+import torch.utils.data
+from cyclegan_model.data.base_dataset import BaseDataset
+
+
+def find_dataset_using_name(dataset_name):
+    """Import the module "data/[dataset_name]_dataset.py".
+
+    In the file, the class called DatasetNameDataset() will
+    be instantiated. It has to be a subclass of BaseDataset,
+    and it is case-insensitive.
+    """
+    dataset_filename = "cyclegan_model.data." + dataset_name + "_dataset"
+    datasetlib = importlib.import_module(dataset_filename)
+
+    dataset = None
+    target_dataset_name = dataset_name.replace('_', '') + 'dataset'
+    for name, cls in datasetlib.__dict__.items():
+        if name.lower() == target_dataset_name.lower() \
+           and issubclass(cls, BaseDataset):
+            dataset = cls
+
+    if dataset is None:
+        raise NotImplementedError("In %s.py, there should be a subclass of BaseDataset with class name that matches %s in lowercase." % (dataset_filename, target_dataset_name))
+
+    return dataset
+
+
+def get_option_setter(dataset_name):
+    """Return the static method <modify_commandline_options> of the dataset class."""
+    dataset_class = find_dataset_using_name(dataset_name)
+    return dataset_class.modify_commandline_options
+
+
+def create_dataset(opt):
+    """Create a dataset given the option.
+
+    This function wraps the class CustomDatasetDataLoader.
+        This is the main interface between this package and 'train.py'/'test.py'
+    """
+    data_loader = CustomDatasetDataLoader(opt)
+    dataset = data_loader.load_data()
+    return dataset
+
+
+class CustomDatasetDataLoader():
+    """Wrapper class of Dataset class that performs multi-threaded data loading"""
+
+    def __init__(self, opt):
+        """Initialize this class
+
+        Step 1: create a dataset instance given the name [dataset_mode]
+        Step 2: create a multi-threaded data loader.
+        """
+        self.opt = opt
+        dataset_class = find_dataset_using_name(opt.dataset_mode)
+        self.dataset = dataset_class(opt)
+        print("dataset [%s] was created" % type(self.dataset).__name__)
+        self.dataloader = torch.utils.data.DataLoader(
+            self.dataset,
+            batch_size=opt.batch_size,
+            shuffle=not opt.serial_batches,
+            num_workers=int(opt.num_threads))
+
+    def load_data(self):
+        return self
+
+    def __len__(self):
+        """Return the number of data in the dataset"""
+        return min(len(self.dataset), self.opt.max_dataset_size)
+
+    def __iter__(self):
+        """Return a batch of data"""
+        for i, data in enumerate(self.dataloader):
+            if i * self.opt.batch_size >= self.opt.max_dataset_size:
+                break
+            yield data

cyclegan_model/data/base_dataset.py

@@ -0,0 +1,167 @@
+"""This module implements an abstract base class (ABC) 'BaseDataset' for datasets.
+
+It also includes common transformation functions (e.g., get_transform, __scale_width), which can be later used in subclasses.
+"""
+import random
+import numpy as np
+import torch.utils.data as data
+from PIL import Image
+import torchvision.transforms as transforms
+from abc import ABC, abstractmethod
+
+
+class BaseDataset(data.Dataset, ABC):
+    """This class is an abstract base class (ABC) for datasets.
+
+    To create a subclass, you need to implement the following four functions:
+    -- <__init__>:                      initialize the class, first call BaseDataset.__init__(self, opt).
+    -- <__len__>:                       return the size of dataset.
+    -- <__getitem__>:                   get a data point.
+    -- <modify_commandline_options>:    (optionally) add dataset-specific options and set default options.
+    """
+
+    def __init__(self, opt):
+        """Initialize the class; save the options in the class
+
+        Parameters:
+            opt (Option class)-- stores all the experiment flags; needs to be a subclass of BaseOptions
+        """
+        self.opt = opt
+        self.root = opt.dataroot
+
+    @staticmethod
+    def modify_commandline_options(parser, is_train):
+        """Add new dataset-specific options, and rewrite default values for existing options.
+
+        Parameters:
+            parser          -- original option parser
+            is_train (bool) -- whether training phase or test phase. You can use this flag to add training-specific or test-specific options.
+
+        Returns:
+            the modified parser.
+        """
+        return parser
+
+    @abstractmethod
+    def __len__(self):
+        """Return the total number of images in the dataset."""
+        return 0
+
+    @abstractmethod
+    def __getitem__(self, index):
+        """Return a data point and its metadata information.
+
+        Parameters:
+            index - - a random integer for data indexing
+
+        Returns:
+            a dictionary of data with their names. It ususally contains the data itself and its metadata information.
+        """
+        pass
+
+
+def get_params(opt, size):
+    w, h = size
+    new_h = h
+    new_w = w
+    if opt.preprocess == 'resize_and_crop':
+        new_h = new_w = opt.load_size
+    elif opt.preprocess == 'scale_width_and_crop':
+        new_w = opt.load_size
+        new_h = opt.load_size * h // w
+
+    x = random.randint(0, np.maximum(0, new_w - opt.crop_size))
+    y = random.randint(0, np.maximum(0, new_h - opt.crop_size))
+
+    flip = random.random() > 0.5
+
+    return {'crop_pos': (x, y), 'flip': flip}
+
+
+def get_transform(opt, params=None, grayscale=False, method=transforms.InterpolationMode.BICUBIC, convert=True):
+    transform_list = []
+    if grayscale:
+        transform_list.append(transforms.Grayscale(1))
+    if 'resize' in opt.preprocess:
+        osize = [opt.load_size, opt.load_size]
+        transform_list.append(transforms.Resize(osize, method))
+    elif 'scale_width' in opt.preprocess:
+        transform_list.append(transforms.Lambda(lambda img: __scale_width(img, opt.load_size, opt.crop_size, method)))
+
+    if 'crop' in opt.preprocess:
+        if params is None:
+            transform_list.append(transforms.RandomCrop(opt.crop_size))
+        else:
+            transform_list.append(transforms.Lambda(lambda img: __crop(img, params['crop_pos'], opt.crop_size)))
+
+    if opt.preprocess == 'none':
+        transform_list.append(transforms.Lambda(lambda img: __make_power_2(img, base=4, method=method)))
+
+    if not opt.no_flip:
+        if params is None:
+            transform_list.append(transforms.RandomHorizontalFlip())
+        elif params['flip']:
+            transform_list.append(transforms.Lambda(lambda img: __flip(img, params['flip'])))
+
+    if convert:
+        transform_list += [transforms.ToTensor()]
+        if grayscale:
+            transform_list += [transforms.Normalize((0.5,), (0.5,))]
+        else:
+            transform_list += [transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))]
+    return transforms.Compose(transform_list)
+
+
+def __transforms2pil_resize(method):
+    mapper = {transforms.InterpolationMode.BILINEAR: Image.BILINEAR,
+              transforms.InterpolationMode.BICUBIC: Image.BICUBIC,
+              transforms.InterpolationMode.NEAREST: Image.NEAREST,
+              transforms.InterpolationMode.LANCZOS: Image.LANCZOS,}
+    return mapper[method]
+
+
+def __make_power_2(img, base, method=transforms.InterpolationMode.BICUBIC):
+    method = __transforms2pil_resize(method)
+    ow, oh = img.size
+    h = int(round(oh / base) * base)
+    w = int(round(ow / base) * base)
+    if h == oh and w == ow:
+        return img
+
+    __print_size_warning(ow, oh, w, h)
+    return img.resize((w, h), method)
+
+
+def __scale_width(img, target_size, crop_size, method=transforms.InterpolationMode.BICUBIC):
+    method = __transforms2pil_resize(method)
+    ow, oh = img.size
+    if ow == target_size and oh >= crop_size:
+        return img
+    w = target_size
+    h = int(max(target_size * oh / ow, crop_size))
+    return img.resize((w, h), method)
+
+
+def __crop(img, pos, size):
+    ow, oh = img.size
+    x1, y1 = pos
+    tw = th = size
+    if (ow > tw or oh > th):
+        return img.crop((x1, y1, x1 + tw, y1 + th))
+    return img
+
+
+def __flip(img, flip):
+    if flip:
+        return img.transpose(Image.FLIP_LEFT_RIGHT)
+    return img
+
+
+def __print_size_warning(ow, oh, w, h):
+    """Print warning information about image size(only print once)"""
+    if not hasattr(__print_size_warning, 'has_printed'):
+        print("The image size needs to be a multiple of 4. "
+              "The loaded image size was (%d, %d), so it was adjusted to "
+              "(%d, %d). This adjustment will be done to all images "
+              "whose sizes are not multiples of 4" % (ow, oh, w, h))
+        __print_size_warning.has_printed = True

cyclegan_model/data/image_folder.py

@@ -0,0 +1,65 @@
+"""A modified image folder class
+
+We modify the official PyTorch image folder (https://github.com/pytorch/vision/blob/master/torchvision/datasets/folder.py)
+so that this class can load images from both current directory and its subdirectories.
+"""
+
+import torch.utils.data as data
+
+from PIL import Image
+import os
+
+IMG_EXTENSIONS = [
+    '.jpg', '.JPG', '.jpeg', '.JPEG',
+    '.png', '.PNG', '.ppm', '.PPM', '.bmp', '.BMP',
+    '.tif', '.TIF', '.tiff', '.TIFF',
+]
+
+
+def is_image_file(filename):
+    return any(filename.endswith(extension) for extension in IMG_EXTENSIONS)
+
+
+def make_dataset(dir, max_dataset_size=float("inf")):
+    images = []
+    assert os.path.isdir(dir), '%s is not a valid directory' % dir
+
+    for root, _, fnames in sorted(os.walk(dir)):
+        for fname in fnames:
+            if is_image_file(fname):
+                path = os.path.join(root, fname)
+                images.append(path)
+    return images[:min(max_dataset_size, len(images))]
+
+
+def default_loader(path):
+    return Image.open(path).convert('RGB')
+
+
+class ImageFolder(data.Dataset):
+
+    def __init__(self, root, transform=None, return_paths=False,
+                 loader=default_loader):
+        imgs = make_dataset(root)
+        if len(imgs) == 0:
+            raise(RuntimeError("Found 0 images in: " + root + "\n"
+                               "Supported image extensions are: " + ",".join(IMG_EXTENSIONS)))
+
+        self.root = root
+        self.imgs = imgs
+        self.transform = transform
+        self.return_paths = return_paths
+        self.loader = loader
+
+    def __getitem__(self, index):
+        path = self.imgs[index]
+        img = self.loader(path)
+        if self.transform is not None:
+            img = self.transform(img)
+        if self.return_paths:
+            return img, path
+        else:
+            return img
+
+    def __len__(self):
+        return len(self.imgs)

cyclegan_model/data/single_dataset.py

@@ -0,0 +1,40 @@
+from cyclegan_model.data.base_dataset import BaseDataset, get_transform
+from cyclegan_model.data.image_folder import make_dataset
+from PIL import Image
+
+
+class SingleDataset(BaseDataset):
+    """This dataset class can load a set of images specified by the path --dataroot /path/to/data.
+
+    It can be used for generating CycleGAN results only for one side with the model option '-model test'.
+    """
+
+    def __init__(self, opt):
+        """Initialize this dataset class.
+
+        Parameters:
+            opt (Option class) -- stores all the experiment flags; needs to be a subclass of BaseOptions
+        """
+        BaseDataset.__init__(self, opt)
+        self.A_paths = sorted(make_dataset(opt.dataroot, opt.max_dataset_size))
+        input_nc = self.opt.output_nc if self.opt.direction == 'BtoA' else self.opt.input_nc
+        self.transform = get_transform(opt, grayscale=(input_nc == 1))
+
+    def __getitem__(self, index):
+        """Return a data point and its metadata information.
+
+        Parameters:
+            index - - a random integer for data indexing
+
+        Returns a dictionary that contains A and A_paths
+            A(tensor) - - an image in one domain
+            A_paths(str) - - the path of the image
+        """
+        A_path = self.A_paths[index]
+        A_img = Image.open(A_path).convert('RGB')
+        A = self.transform(A_img)
+        return {'A': A, 'A_paths': A_path}
+
+    def __len__(self):
+        """Return the total number of images in the dataset."""
+        return len(self.A_paths)

cyclegan_model/data/unaligned_dataset.py

@@ -0,0 +1,71 @@
+import os
+from cyclegan_model.data.base_dataset import BaseDataset, get_transform
+from cyclegan_model.data.image_folder import make_dataset
+from PIL import Image
+import random
+
+
+class UnalignedDataset(BaseDataset):
+    """
+    This dataset class can load unaligned/unpaired datasets.
+
+    It requires two directories to host training images from domain A '/path/to/data/trainA'
+    and from domain B '/path/to/data/trainB' respectively.
+    You can train the model with the dataset flag '--dataroot /path/to/data'.
+    Similarly, you need to prepare two directories:
+    '/path/to/data/testA' and '/path/to/data/testB' during test time.
+    """
+
+    def __init__(self, opt):
+        """Initialize this dataset class.
+
+        Parameters:
+            opt (Option class) -- stores all the experiment flags; needs to be a subclass of BaseOptions
+        """
+        BaseDataset.__init__(self, opt)
+        self.dir_A = os.path.join(opt.dataroot, opt.phase + 'A')  # create a path '/path/to/data/trainA'
+        self.dir_B = os.path.join(opt.dataroot, opt.phase + 'B')  # create a path '/path/to/data/trainB'
+
+        self.A_paths = sorted(make_dataset(self.dir_A, opt.max_dataset_size))   # load images from '/path/to/data/trainA'
+        self.B_paths = sorted(make_dataset(self.dir_B, opt.max_dataset_size))    # load images from '/path/to/data/trainB'
+        self.A_size = len(self.A_paths)  # get the size of dataset A
+        self.B_size = len(self.B_paths)  # get the size of dataset B
+        btoA = self.opt.direction == 'BtoA'
+        input_nc = self.opt.output_nc if btoA else self.opt.input_nc       # get the number of channels of input image
+        output_nc = self.opt.input_nc if btoA else self.opt.output_nc      # get the number of channels of output image
+        self.transform_A = get_transform(self.opt, grayscale=(input_nc == 1))
+        self.transform_B = get_transform(self.opt, grayscale=(output_nc == 1))
+
+    def __getitem__(self, index):
+        """Return a data point and its metadata information.
+
+        Parameters:
+            index (int)      -- a random integer for data indexing
+
+        Returns a dictionary that contains A, B, A_paths and B_paths
+            A (tensor)       -- an image in the input domain
+            B (tensor)       -- its corresponding image in the target domain
+            A_paths (str)    -- image paths
+            B_paths (str)    -- image paths
+        """
+        A_path = self.A_paths[index % self.A_size]  # make sure index is within then range
+        if self.opt.serial_batches:   # make sure index is within then range
+            index_B = index % self.B_size
+        else:   # randomize the index for domain B to avoid fixed pairs.
+            index_B = random.randint(0, self.B_size - 1)
+        B_path = self.B_paths[index_B]
+        A_img = Image.open(A_path).convert('RGB')
+        B_img = Image.open(B_path).convert('RGB')
+        # apply image transformation
+        A = self.transform_A(A_img)
+        B = self.transform_B(B_img)
+
+        return {'A': A, 'B': B, 'A_paths': A_path, 'B_paths': B_path}
+
+    def __len__(self):
+        """Return the total number of images in the dataset.
+
+        As we have two datasets with potentially different number of images,
+        we take a maximum of
+        """
+        return max(self.A_size, self.B_size)

cyclegan_model/model/__init__.py

@@ -0,0 +1,63 @@
+"""This package contains modules related to objective functions, optimizations, and network architectures.
+
+To add a custom model class called 'dummy', you need to add a file called 'dummy_model.py' and define a subclass DummyModel inherited from BaseModel.
+You need to implement the following five functions:
+    -- <__init__>:                      initialize the class; first call BaseModel.__init__(self, opt).
+    -- <set_input>:                     unpack data from dataset and apply preprocessing.
+    -- <forward>:                       produce intermediate results.
+    -- <optimize_parameters>:           calculate loss, gradients, and update network weights.
+    -- <modify_commandline_options>:    (optionally) add model-specific options and set default options.
+
+In the function <__init__>, you need to define four lists:
+    -- self.loss_names (str list):          specify the training losses that you want to plot and save.
+    -- self.model_names (str list):         define networks used in our training.
+    -- self.visual_names (str list):        specify the images that you want to display and save.
+    -- self.optimizers (optimizer list):    define and initialize optimizers. You can define one optimizer for each network. If two networks are updated at the same time, you can use itertools.chain to group them. See cycle_gan_model.py for an usage.
+
+Now you can use the model class by specifying flag '--model dummy'.
+See our template model class 'template_model.py' for more details.
+"""
+
+import importlib
+from cyclegan_model.model.base_model import BaseModel
+
+
+def find_model_using_name(model_name):
+    """Import the module "models/[model_name]_model.py".
+
+    In the file, the class called DatasetNameModel() will
+    be instantiated. It has to be a subclass of BaseModel,
+    and it is case-insensitive.
+    """
+    model_filename = "cyclegan_model.model." + model_name + "_model"
+    modellib = importlib.import_module(model_filename)
+    model = None
+    target_model_name = model_name.replace('_', '') + 'model'
+    for name, cls in modellib.__dict__.items():
+        if name.lower() == target_model_name.lower() \
+           and issubclass(cls, BaseModel):
+            model = cls
+
+    if model is None:
+        print("In %s.py, there should be a subclass of BaseModel with class name that matches %s in lowercase." % (model_filename, target_model_name))
+        exit(0)
+
+    return model
+
+
+def get_option_setter(model_name):
+    """Return the static method <modify_commandline_options> of the model class."""
+    model_class = find_model_using_name(model_name)
+    return model_class.modify_commandline_options
+
+
+def create_model(opt):
+    """Create a model given the option.
+
+    This function warps the class CustomDatasetDataLoader.
+    This is the main interface between this package and 'train.py'/'test.py'
+    """
+    model = find_model_using_name(opt.model)
+    instance = model(opt)
+    print("model [%s] was created" % type(instance).__name__)
+    return instance

cyclegan_model/model/base_model.py

@@ -0,0 +1,230 @@
+import os
+import torch
+from collections import OrderedDict
+from abc import ABC, abstractmethod
+from . import networks
+
+
+class BaseModel(ABC):
+    """This class is an abstract base class (ABC) for models.
+    To create a subclass, you need to implement the following five functions:
+        -- <__init__>:                      initialize the class; first call BaseModel.__init__(self, opt).
+        -- <set_input>:                     unpack data from dataset and apply preprocessing.
+        -- <forward>:                       produce intermediate results.
+        -- <optimize_parameters>:           calculate losses, gradients, and update network weights.
+        -- <modify_commandline_options>:    (optionally) add model-specific options and set default options.
+    """
+
+    def __init__(self, opt):
+        """Initialize the BaseModel class.
+
+        Parameters:
+            opt (Option class)-- stores all the experiment flags; needs to be a subclass of BaseOptions
+
+        When creating your custom class, you need to implement your own initialization.
+        In this function, you should first call <BaseModel.__init__(self, opt)>
+        Then, you need to define four lists:
+            -- self.loss_names (str list):          specify the training losses that you want to plot and save.
+            -- self.model_names (str list):         define networks used in our training.
+            -- self.visual_names (str list):        specify the images that you want to display and save.
+            -- self.optimizers (optimizer list):    define and initialize optimizers. You can define one optimizer for each network. If two networks are updated at the same time, you can use itertools.chain to group them. See cycle_gan_model.py for an example.
+        """
+        self.opt = opt
+        self.gpu_ids = opt.gpu_ids
+        self.isTrain = opt.isTrain
+        self.device = torch.device('cuda:{}'.format(self.gpu_ids[0])) if self.gpu_ids else torch.device('cpu')  # get device name: CPU or GPU
+        self.save_dir = os.path.join(opt.checkpoints_dir, opt.name)  # save all the checkpoints to save_dir
+        if opt.preprocess != 'scale_width':  # with [scale_width], input images might have different sizes, which hurts the performance of cudnn.benchmark.
+            torch.backends.cudnn.benchmark = True
+        self.loss_names = []
+        self.model_names = []
+        self.visual_names = []
+        self.optimizers = []
+        self.image_paths = []
+        self.metric = 0  # used for learning rate policy 'plateau'
+
+    @staticmethod
+    def modify_commandline_options(parser, is_train):
+        """Add new model-specific options, and rewrite default values for existing options.
+
+        Parameters:
+            parser          -- original option parser
+            is_train (bool) -- whether training phase or test phase. You can use this flag to add training-specific or test-specific options.
+
+        Returns:
+            the modified parser.
+        """
+        return parser
+
+    @abstractmethod
+    def set_input(self, input):
+        """Unpack input data from the dataloader and perform necessary pre-processing steps.
+
+        Parameters:
+            input (dict): includes the data itself and its metadata information.
+        """
+        pass
+
+    @abstractmethod
+    def forward(self):
+        """Run forward pass; called by both functions <optimize_parameters> and <test>."""
+        pass
+
+    @abstractmethod
+    def optimize_parameters(self):
+        """Calculate losses, gradients, and update network weights; called in every training iteration"""
+        pass
+
+    def setup(self, opt):
+        """Load and print networks; create schedulers
+
+        Parameters:
+            opt (Option class) -- stores all the experiment flags; needs to be a subclass of BaseOptions
+        """
+        if self.isTrain:
+            self.schedulers = [networks.get_scheduler(optimizer, opt) for optimizer in self.optimizers]
+        if not self.isTrain or opt.continue_train:
+            load_suffix = 'iter_%d' % opt.load_iter if opt.load_iter > 0 else opt.epoch
+            self.load_networks(load_suffix)
+        self.print_networks(opt.verbose)
+
+    def eval(self):
+        """Make models eval mode during test time"""
+        for name in self.model_names:
+            if isinstance(name, str):
+                net = getattr(self, 'net' + name)
+                net.eval()
+
+    def test(self):
+        """Forward function used in test time.
+
+        This function wraps <forward> function in no_grad() so we don't save intermediate steps for backprop
+        It also calls <compute_visuals> to produce additional visualization results
+        """
+        with torch.no_grad():
+            self.forward()
+            self.compute_visuals()
+
+    def compute_visuals(self):
+        """Calculate additional output images for visdom and HTML visualization"""
+        pass
+
+    def get_image_paths(self):
+        """ Return image paths that are used to load current data"""
+        return self.image_paths
+
+    def update_learning_rate(self):
+        """Update learning rates for all the networks; called at the end of every epoch"""
+        old_lr = self.optimizers[0].param_groups[0]['lr']
+        for scheduler in self.schedulers:
+            if self.opt.lr_policy == 'plateau':
+                scheduler.step(self.metric)
+            else:
+                scheduler.step()
+
+        lr = self.optimizers[0].param_groups[0]['lr']
+        print('learning rate %.7f -> %.7f' % (old_lr, lr))
+
+    def get_current_visuals(self):
+        """Return visualization images. train.py will display these images with visdom, and save the images to a HTML"""
+        visual_ret = OrderedDict()
+        for name in self.visual_names:
+            if isinstance(name, str):
+                visual_ret[name] = getattr(self, name)
+        return visual_ret
+
+    def get_current_losses(self):
+        """Return traning losses / errors. train.py will print out these errors on console, and save them to a file"""
+        errors_ret = OrderedDict()
+        for name in self.loss_names:
+            if isinstance(name, str):
+                errors_ret[name] = float(getattr(self, 'loss_' + name))  # float(...) works for both scalar tensor and float number
+        return errors_ret
+
+    def save_networks(self, epoch):
+        """Save all the networks to the disk.
+
+        Parameters:
+            epoch (int) -- current epoch; used in the file name '%s_net_%s.pth' % (epoch, name)
+        """
+        for name in self.model_names:
+            if isinstance(name, str):
+                save_filename = '%s_net_%s.pth' % (epoch, name)
+                save_path = os.path.join(self.save_dir, save_filename)
+                net = getattr(self, 'net' + name)
+
+                if len(self.gpu_ids) > 0 and torch.cuda.is_available():
+                    torch.save(net.module.cpu().state_dict(), save_path)
+                    net.cuda(self.gpu_ids[0])
+                else:
+                    torch.save(net.cpu().state_dict(), save_path)
+
+    def __patch_instance_norm_state_dict(self, state_dict, module, keys, i=0):
+        """Fix InstanceNorm checkpoints incompatibility (prior to 0.4)"""
+        key = keys[i]
+        if i + 1 == len(keys):  # at the end, pointing to a parameter/buffer
+            if module.__class__.__name__.startswith('InstanceNorm') and \
+                    (key == 'running_mean' or key == 'running_var'):
+                if getattr(module, key) is None:
+                    state_dict.pop('.'.join(keys))
+            if module.__class__.__name__.startswith('InstanceNorm') and \
+               (key == 'num_batches_tracked'):
+                state_dict.pop('.'.join(keys))
+        else:
+            self.__patch_instance_norm_state_dict(state_dict, getattr(module, key), keys, i + 1)
+
+    def load_networks(self, epoch):
+        """Load all the networks from the disk.
+
+        Parameters:
+            epoch (int) -- current epoch; used in the file name '%s_net_%s.pth' % (epoch, name)
+        """
+        for name in self.model_names:
+            if isinstance(name, str):
+                load_filename = '%s_net_%s.pth' % (epoch, name)
+                load_path = os.path.join(self.save_dir, load_filename)
+                net = getattr(self, 'net' + name)
+                if isinstance(net, torch.nn.DataParallel):
+                    net = net.module
+                print('loading the model from %s' % load_path)
+                # if you are using PyTorch newer than 0.4 (e.g., built from
+                # GitHub source), you can remove str() on self.device
+                state_dict = torch.load(load_path, map_location=str(self.device))
+                if hasattr(state_dict, '_metadata'):
+                    del state_dict._metadata
+
+                # patch InstanceNorm checkpoints prior to 0.4
+                for key in list(state_dict.keys()):  # need to copy keys here because we mutate in loop
+                    self.__patch_instance_norm_state_dict(state_dict, net, key.split('.'))
+                net.load_state_dict(state_dict)
+
+    def print_networks(self, verbose):
+        """Print the total number of parameters in the network and (if verbose) network architecture
+
+        Parameters:
+            verbose (bool) -- if verbose: print the network architecture
+        """
+        print('---------- Networks initialized -------------')
+        for name in self.model_names:
+            if isinstance(name, str):
+                net = getattr(self, 'net' + name)
+                num_params = 0
+                for param in net.parameters():
+                    num_params += param.numel()
+                if verbose:
+                    print(net)
+                print('[Network %s] Total number of parameters : %.3f M' % (name, num_params / 1e6))
+        print('-----------------------------------------------')
+
+    def set_requires_grad(self, nets, requires_grad=False):
+        """Set requies_grad=Fasle for all the networks to avoid unnecessary computations
+        Parameters:
+            nets (network list)   -- a list of networks
+            requires_grad (bool)  -- whether the networks require gradients or not
+        """
+        if not isinstance(nets, list):
+            nets = [nets]
+        for net in nets:
+            if net is not None:
+                for param in net.parameters():
+                    param.requires_grad = requires_grad

cyclegan_model/model/cycle_gan_model.py

@@ -0,0 +1,229 @@
+import torch
+import itertools
+from cyclegan_model.util.image_pool import ImagePool
+from .base_model import BaseModel
+from . import networks
+from .networks import cal_gradient_penalty
+
+class CycleGANModel(BaseModel):
+    """
+    This class implements the CycleGAN model, for learning image-to-image translation without paired data.
+
+    The model training requires '--dataset_mode unaligned' dataset.
+    By default, it uses a '--netG resnet_9blocks' ResNet generator,
+    a '--netD basic' discriminator (PatchGAN introduced by pix2pix),
+    and a least-square GANs objective ('--gan_mode lsgan').
+
+    CycleGAN paper: https://arxiv.org/pdf/1703.10593.pdf
+    """
+    @staticmethod
+    def modify_commandline_options(parser, is_train=True):
+        """Add new dataset-specific options, and rewrite default values for existing options.
+
+        Parameters:
+            parser          -- original option parser
+            is_train (bool) -- whether training phase or test phase. You can use this flag to add training-specific or test-specific options.
+
+        Returns:
+            the modified parser.
+
+        For CycleGAN, in addition to GAN losses, we introduce lambda_A, lambda_B, and lambda_identity for the following losses.
+        A (source domain), B (target domain).
+        Generators: G_A: A -> B; G_B: B -> A.
+        Discriminators: D_A: G_A(A) vs. B; D_B: G_B(B) vs. A.
+        Forward cycle loss:  lambda_A * ||G_B(G_A(A)) - A|| (Eqn. (2) in the paper)
+        Backward cycle loss: lambda_B * ||G_A(G_B(B)) - B|| (Eqn. (2) in the paper)
+        Identity loss (optional): lambda_identity * (||G_A(B) - B|| * lambda_B + ||G_B(A) - A|| * lambda_A) (Sec 5.2 "Photo generation from paintings" in the paper)
+        Dropout is not used in the original CycleGAN paper.
+        """
+        parser.set_defaults(no_dropout=True)  # default CycleGAN did not use dropout
+        if is_train:
+            parser.add_argument('--lambda_A', type=float, default=10.0, help='weight for cycle loss (A -> B -> A)')
+            parser.add_argument('--lambda_B', type=float, default=10.0, help='weight for cycle loss (B -> A -> B)')
+            parser.add_argument('--lambda_identity', type=float, default=0.5, help='use identity mapping. Setting lambda_identity other than 0 has an effect of scaling the weight of the identity mapping loss. For example, if the weight of the identity loss should be 10 times smaller than the weight of the reconstruction loss, please set lambda_identity = 0.1')
+
+        return parser
+
+    def __init__(self, opt):
+        """Initialize the CycleGAN class.
+
+        Parameters:
+            opt (Option class)-- stores all the experiment flags; needs to be a subclass of BaseOptions
+        """
+        BaseModel.__init__(self, opt)
+        # specify the training losses you want to print out. The training/test scripts will call <BaseModel.get_current_losses>
+        self.loss_names = ['D_A', 'G_A', 'cycle_A', 'idt_A', 'D_B', 'G_B', 'cycle_B', 'idt_B']
+        # specify the images you want to save/display. The training/test scripts will call <BaseModel.get_current_visuals>
+        visual_names_A = ['real_A', 'fake_B', 'rec_A']
+        visual_names_B = ['real_B', 'fake_A', 'rec_B']
+        if self.isTrain and self.opt.lambda_identity > 0.0:  # if identity loss is used, we also visualize idt_B=G_A(B) ad idt_A=G_A(B)
+            visual_names_A.append('idt_B')
+            visual_names_B.append('idt_A')
+
+        self.visual_names = visual_names_A + visual_names_B  # combine visualizations for A and B
+        # specify the models you want to save to the disk. The training/test scripts will call <BaseModel.save_networks> and <BaseModel.load_networks>.
+        if self.isTrain:
+            self.model_names = ['G_A', 'G_B', 'D_A', 'D_B']
+        else:  # during test time, only load Gs
+            self.model_names = ['G_A', 'G_B']
+
+        # define networks (both Generators and discriminators)
+        # The naming is different from those used in the paper.
+        # Code (vs. paper): G_A (G), G_B (F), D_A (D_Y), D_B (D_X)
+        self.netG_A = networks.define_G(opt.input_nc, opt.output_nc, opt.ngf, opt.netG, opt.norm,
+                                        not opt.no_dropout, opt.init_type, opt.init_gain, self.gpu_ids)
+        self.netG_B = networks.define_G(opt.output_nc, opt.input_nc, opt.ngf, opt.netG, opt.norm,
+                                        not opt.no_dropout, opt.init_type, opt.init_gain, self.gpu_ids)
+
+        if self.isTrain:  # define discriminators
+            self.netD_A = networks.define_D(opt.output_nc, opt.ndf, opt.netD,
+                                            opt.n_layers_D, opt.norm, opt.init_type, opt.init_gain, self.gpu_ids)
+            self.netD_B = networks.define_D(opt.input_nc, opt.ndf, opt.netD,
+                                            opt.n_layers_D, opt.norm, opt.init_type, opt.init_gain, self.gpu_ids)
+
+        if self.isTrain:
+            if opt.lambda_identity > 0.0:  # only works when input and output images have the same number of channels
+                assert(opt.input_nc == opt.output_nc)
+            self.fake_A_pool = ImagePool(opt.pool_size)  # create image buffer to store previously generated images
+            self.fake_B_pool = ImagePool(opt.pool_size)  # create image buffer to store previously generated images
+            # define loss functions
+            self.criterionGAN = networks.GANLoss(opt.gan_mode).to(self.device)  # define GAN loss.
+            self.criterionCycle = torch.nn.L1Loss()
+            self.criterionIdt = torch.nn.L1Loss()
+
+
+
+            #self.criterionCycle = lpips.LPIPS(net='alex').to(self.device)
+            #self.criterionIdt = lpips.LPIPS(net='alex').to(self.device)
+            # initialize optimizers; schedulers will be automatically created by function <BaseModel.setup>.
+            self.optimizer_G = torch.optim.Adam(itertools.chain(self.netG_A.parameters(), self.netG_B.parameters()), lr=0.0001, betas=(opt.beta1, 0.999))
+            self.optimizer_D = torch.optim.Adam(itertools.chain(self.netD_A.parameters(), self.netD_B.parameters()), lr=0.0003, betas=(opt.beta1, 0.999))
+            self.optimizers.append(self.optimizer_G)
+            self.optimizers.append(self.optimizer_D)
+
+    def set_input(self, input):
+        """Unpack input data from the dataloader and perform necessary pre-processing steps.
+
+        Parameters:
+            input (dict): include the data itself and its metadata information.
+
+        The option 'direction' can be used to swap domain A and domain B.
+        """
+        AtoB = self.opt.direction == 'AtoB'
+        self.real_A = input['A' if AtoB else 'B'].to(self.device)
+        self.real_B = input['B' if AtoB else 'A'].to(self.device)
+
+    def forward(self):
+        """Run forward pass; called by both functions <optimize_parameters> and <test>."""
+        self.fake_B = self.netG_A(self.real_A)  # G_A(A)
+        self.rec_A = self.netG_B(self.fake_B)   # G_B(G_A(A))
+        self.fake_A = self.netG_B(self.real_B)  # G_B(B)
+        self.rec_B = self.netG_A(self.fake_A)   # G_A(G_B(B))
+
+    def backward_D_basic(self, netD, real, fake):
+        """Calculate GAN loss for the discriminator
+
+        Parameters:
+            netD (network)      -- the discriminator D
+            real (tensor array) -- real images
+            fake (tensor array) -- images generated by a generator
+
+        Return the discriminator loss.
+        We also call loss_D.backward() to calculate the gradients.
+        """
+        # Real
+        pred_real = netD(real)
+        pradient_penalty,pradients=cal_gradient_penalty(netD,real,fake,self.device)
+
+
+        loss_D_real = self.criterionGAN(pred_real, True)
+        # Fake
+        pred_fake = netD(fake.detach())
+        loss_D_fake = self.criterionGAN(pred_fake, False)
+        # Combined loss and calculate gradients
+        loss_D = (loss_D_real + loss_D_fake) * 0.5
+        loss_D.backward()
+        return loss_D
+
+    def backward_D_A(self):
+        """Calculate GAN loss for discriminator D_A"""
+        fake_B = self.fake_B_pool.query(self.fake_B)
+        self.loss_D_A = self.backward_D_basic(self.netD_A, self.real_B, fake_B)
+
+    def backward_D_B(self):
+        """Calculate GAN loss for discriminator D_B"""
+        fake_A = self.fake_A_pool.query(self.fake_A)
+        self.loss_D_B = self.backward_D_basic(self.netD_B, self.real_A, fake_A)
+
+    def backward_G(self,retain_graph=False):
+        """Calculate the loss for generators G_A and G_B"""
+        lambda_idt = self.opt.lambda_identity
+        lambda_A = self.opt.lambda_A
+        lambda_B = self.opt.lambda_B
+        # Identity loss
+        if lambda_idt > 0:
+            # G_A should be identity if real_B is fed: ||G_A(B) - B||
+            self.idt_A = self.netG_A(self.real_B)
+            self.loss_idt_A = self.criterionIdt(self.idt_A, self.real_B) * lambda_B * lambda_idt
+            # G_B should be identity if real_A is fed: ||G_B(A) - A||
+            self.idt_B = self.netG_B(self.real_A)
+            self.loss_idt_B = self.criterionIdt(self.idt_B, self.real_A) * lambda_A * lambda_idt
+        else:
+            self.loss_idt_A = 0
+            self.loss_idt_B = 0
+
+        # GAN loss D_A(G_A(A))
+        self.loss_G_A = self.criterionGAN(self.netD_A(self.fake_B), True)
+        # GAN loss D_B(G_B(B))
+        self.loss_G_B = self.criterionGAN(self.netD_B(self.fake_A), True)
+        # Forward cycle loss || G_B(G_A(A)) - A||
+        self.loss_cycle_A = self.criterionCycle(self.rec_A, self.real_A) * lambda_A
+        # Backward cycle loss || G_A(G_B(B)) - B||
+        self.loss_cycle_B = self.criterionCycle(self.rec_B, self.real_B) * lambda_B
+        # combined loss and calculate gradients
+        self.loss_G = self.loss_G_A + self.loss_G_B + 0.7*self.loss_cycle_A +0.7*self.loss_cycle_B + self.loss_idt_A + self.loss_idt_B
+
+        self.loss_G.backward(retain_graph=retain_graph)
+        #self.loss_G.backward(retain_graph=True)
+
+    def optimize_parameters(self):
+        torch.autograd.set_detect_anomaly(True)
+        """Calculate losses, gradients, and update network weights; called in every training iteration"""
+        # forward
+        self.forward()      # compute fake images and reconstruction images.
+        # G_A and G_B
+        self.set_requires_grad([self.netD_A, self.netD_B], False)  # Ds require no gradients when optimizing Gs
+        '''
+        self.optimizer_G.zero_grad()  # set G_A and G_B's gradients to zero
+        self.backward_G()  # calculate gradients for G_A and G_B
+        self.optimizer_G.step()  # update G_A and G_B's weights
+        '''
+        '''
+        self.optimizer_G.zero_grad()  # set G_A and G_B's gradients to zero
+        for i in range(2):  #生成器训练两次
+            self.backward_G(retain_graph=(i < 1))  # calculate gradients for G_A
+            self.optimizer_G.step()       # update G_A and G_B's weights
+        # D_A and D_B
+        '''
+
+        self.optimizer_G.zero_grad()
+        self.backward_G(retain_graph=True)  # 保留计算图
+        grad_cache_G = [p.grad.clone() for p in itertools.chain(self.netG_A.parameters(), self.netG_B.parameters())]
+
+        # 第二次生成器训练
+        self.optimizer_G.zero_grad()
+        self.backward_G(retain_graph=False)
+
+        # 梯度融合：将两次训练的梯度相加
+        for p, cache_g in zip(itertools.chain(self.netG_A.parameters(), self.netG_B.parameters()), grad_cache_G):
+            if p.grad is not None:
+                p.grad += cache_g
+
+        # 执行参数更新
+        self.optimizer_G.step()
+
+        self.set_requires_grad([self.netD_A, self.netD_B], True)
+        self.optimizer_D.zero_grad()   # set D_A and D_B's gradients to zero
+        self.backward_D_A()      # calculate gradients for D_A
+        self.backward_D_B()      # calculate graidents for D_B
+        self.optimizer_D.step()  # update D_A and D_B's weights

cyclegan_model/model/networks.py

@@ -0,0 +1,1091 @@
+import torch
+import torch.nn as nn
+from torch.nn import init
+import functools
+from torch.optim import lr_scheduler
+import random
+import torch.nn.functional as F
+
+###############################################################################
+# Helper Functions
+###############################################################################
+
+
+class Identity(nn.Module):
+    def forward(self, x):
+        return x
+
+
+def get_norm_layer(norm_type='instance'):
+    """Return a normalization layer
+
+    Parameters:
+        norm_type (str) -- the name of the normalization layer: batch | instance | none
+
+    For BatchNorm, we use learnable affine parameters and track running statistics (mean/stddev).
+    For InstanceNorm, we do not use learnable affine parameters. We do not track running statistics.
+    """
+    if norm_type == 'batch':
+        norm_layer = functools.partial(nn.BatchNorm2d, affine=True, track_running_stats=True)
+    elif norm_type == 'instance':
+        norm_layer = functools.partial(nn.InstanceNorm2d, affine=False, track_running_stats=False)
+    elif norm_type == 'none':
+        def norm_layer(x):
+            return Identity()
+    else:
+        raise NotImplementedError('normalization layer [%s] is not found' % norm_type)
+    return norm_layer
+
+
+def get_scheduler(optimizer, opt):
+    """Return a learning rate scheduler
+
+    Parameters:
+        optimizer          -- the optimizer of the network
+        opt (option class) -- stores all the experiment flags; needs to be a subclass of BaseOptions．　
+                              opt.lr_policy is the name of learning rate policy: linear | step | plateau | cosine
+
+    For 'linear', we keep the same learning rate for the first <opt.n_epochs> epochs
+    and linearly decay the rate to zero over the next <opt.n_epochs_decay> epochs.
+    For other schedulers (step, plateau, and cosine), we use the default PyTorch schedulers.
+    See https://pytorch.org/docs/stable/optim.html for more details.
+    """
+    if opt.lr_policy == 'linear':
+        def lambda_rule(epoch):
+            lr_l = 1.0 - max(0, epoch + opt.epoch_count - opt.n_epochs) / float(opt.n_epochs_decay + 1)
+            return lr_l
+        scheduler = lr_scheduler.LambdaLR(optimizer, lr_lambda=lambda_rule)
+    elif opt.lr_policy == 'step':
+        scheduler = lr_scheduler.StepLR(optimizer, step_size=opt.lr_decay_iters, gamma=0.1)
+    elif opt.lr_policy == 'plateau':
+        scheduler = lr_scheduler.ReduceLROnPlateau(optimizer, mode='min', factor=0.2, threshold=0.01, patience=5)
+    elif opt.lr_policy == 'cosine':
+        scheduler = lr_scheduler.CosineAnnealingLR(optimizer, T_max=opt.n_epochs, eta_min=0)
+    else:
+        return NotImplementedError('learning rate policy [%s] is not implemented', opt.lr_policy)
+    return scheduler
+
+
+def init_weights(net, init_type='normal', init_gain=0.02):
+    """Initialize network weights.
+
+    Parameters:
+        net (network)   -- network to be initialized
+        init_type (str) -- the name of an initialization method: normal | xavier | kaiming | orthogonal
+        init_gain (float)    -- scaling factor for normal, xavier and orthogonal.
+
+    We use 'normal' in the original pix2pix and CycleGAN paper. But xavier and kaiming might
+    work better for some applications. Feel free to try yourself.
+    """
+    def init_func(m):  # define the initialization function
+        classname = m.__class__.__name__
+        if hasattr(m, 'weight') and (classname.find('Conv') != -1 or classname.find('Linear') != -1):
+            if init_type == 'normal':
+                init.normal_(m.weight.data, 0.0, init_gain)
+            elif init_type == 'xavier':
+                init.xavier_normal_(m.weight.data, gain=init_gain)
+            elif init_type == 'kaiming':
+                init.kaiming_normal_(m.weight.data, a=0, mode='fan_in')
+            elif init_type == 'orthogonal':
+                init.orthogonal_(m.weight.data, gain=init_gain)
+            else:
+                raise NotImplementedError('initialization method [%s] is not implemented' % init_type)
+            if hasattr(m, 'bias') and m.bias is not None:
+                init.constant_(m.bias.data, 0.0)
+        elif classname.find('BatchNorm2d') != -1:  # BatchNorm Layer's weight is not a matrix; only normal distribution applies.
+            init.normal_(m.weight.data, 1.0, init_gain)
+            init.constant_(m.bias.data, 0.0)
+
+    print('initialize network with %s' % init_type)
+    net.apply(init_func)  # apply the initialization function <init_func>
+
+
+def init_net(net, init_type='normal', init_gain=0.02, gpu_ids=[]):
+    """Initialize a network: 1. register CPU/GPU device (with multi-GPU support); 2. initialize the network weights
+    Parameters:
+        net (network)      -- the network to be initialized
+        init_type (str)    -- the name of an initialization method: normal | xavier | kaiming | orthogonal
+        gain (float)       -- scaling factor for normal, xavier and orthogonal.
+        gpu_ids (int list) -- which GPUs the network runs on: e.g., 0,1,2
+
+    Return an initialized network.
+    """
+    if len(gpu_ids) > 0:
+        assert(torch.cuda.is_available())
+        net.to(gpu_ids[0])
+        net = torch.nn.DataParallel(net, gpu_ids)  # multi-GPUs
+    init_weights(net, init_type, init_gain=init_gain)
+    return net
+
+
+def define_G(input_nc, output_nc, ngf, netG, norm='batch', use_dropout=False, init_type='normal', init_gain=0.02, gpu_ids=[]):
+    """Create a generator
+
+    Parameters:
+        input_nc (int) -- the number of channels in input images
+        output_nc (int) -- the number of channels in output images
+        ngf (int) -- the number of filters in the last conv layer
+        netG (str) -- the architecture's name: resnet_9blocks | resnet_6blocks | unet_256 | unet_128
+        norm (str) -- the name of normalization layers used in the network: batch | instance | none
+        use_dropout (bool) -- if use dropout layers.
+        init_type (str)    -- the name of our initialization method.
+        init_gain (float)  -- scaling factor for normal, xavier and orthogonal.
+        gpu_ids (int list) -- which GPUs the network runs on: e.g., 0,1,2
+
+    Returns a generator
+
+    Our current implementation provides two types of generators:
+        U-Net: [unet_128] (for 128x128 input images) and [unet_256] (for 256x256 input images)
+        The original U-Net paper: https://arxiv.org/abs/1505.04597
+
+        Resnet-based generator: [resnet_6blocks] (with 6 Resnet blocks) and [resnet_9blocks] (with 9 Resnet blocks)
+        Resnet-based generator consists of several Resnet blocks between a few downsampling/upsampling operations.
+        We adapt Torch code from Justin Johnson's neural style transfer project (https://github.com/jcjohnson/fast-neural-style).
+
+
+    The generator has been initialized by <init_net>. It uses RELU for non-linearity.
+    """
+    net = None
+    norm_layer = get_norm_layer(norm_type=norm)
+
+    if netG == 'resnet_9blocks':
+        net = ResnetGenerator(input_nc, output_nc, ngf, norm_layer=norm_layer, use_dropout=use_dropout, n_blocks=9)
+    elif netG == 'resnet_6blocks':
+        net = ResnetGenerator(input_nc, output_nc, ngf, norm_layer=norm_layer, use_dropout=use_dropout, n_blocks=6)
+    elif netG == 'unet_128':
+        net = UnetGenerator(input_nc, output_nc, 7, ngf, norm_layer=norm_layer, use_dropout=use_dropout)
+    elif netG == 'unet_256':
+        net = UnetGenerator(input_nc, output_nc, 8, ngf, norm_layer=norm_layer, use_dropout=use_dropout)
+    elif netG == 'resnet_attention':
+        net = ResnetGeneratorWithAttention(input_nc,output_nc,ngf,norm_layer=norm_layer,use_dropout=use_dropout,n_blocks=9)
+    elif netG == 'resnet_unet_attention':
+        net = Unet_SEA_ResnetGenerator(input_nc,output_nc, ngf, norm_layer, use_dropout)
+    elif netG == 'resnet_skip_attention':
+        net = ResnetGeneratorWithAttentionAndSkipConnection(input_nc,output_nc,ngf,norm_layer=norm_layer,use_dropout=use_dropout,n_blocks=12)
+    else:
+        raise NotImplementedError('Generator model name [%s] is not recognized' % netG)
+    return init_net(net, init_type, init_gain, gpu_ids)
+
+
+def define_D(input_nc, ndf, netD, n_layers_D=3, norm='batch', init_type='normal', init_gain=0.02, gpu_ids=[]):
+    """Create a discriminator
+
+    Parameters:
+        input_nc (int)     -- the number of channels in input images
+        ndf (int)          -- the number of filters in the first conv layer
+        netD (str)         -- the architecture's name: basic | n_layers | pixel
+        n_layers_D (int)   -- the number of conv layers in the discriminator; effective when netD=='n_layers'
+        norm (str)         -- the type of normalization layers used in the network.
+        init_type (str)    -- the name of the initialization method.
+        init_gain (float)  -- scaling factor for normal, xavier and orthogonal.
+        gpu_ids (int list) -- which GPUs the network runs on: e.g., 0,1,2
+
+    Returns a discriminator
+
+    Our current implementation provides three types of discriminators:
+        [basic]: 'PatchGAN' classifier described in the original pix2pix paper.
+        It can classify whether 70×70 overlapping patches are real or fake.
+        Such a patch-level discriminator architecture has fewer parameters
+        than a full-image discriminator and can work on arbitrarily-sized images
+        in a fully convolutional fashion.
+
+        [n_layers]: With this mode, you can specify the number of conv layers in the discriminator
+        with the parameter <n_layers_D> (default=3 as used in [basic] (PatchGAN).)
+
+        [pixel]: 1x1 PixelGAN discriminator can classify whether a pixel is real or not.
+        It encourages greater color diversity but has no effect on spatial statistics.
+
+    The discriminator has been initialized by <init_net>. It uses Leakly RELU for non-linearity.
+    """
+    net = None
+    norm_layer = get_norm_layer(norm_type=norm)
+
+    if netD == 'basic':  # default PatchGAN classifier
+        net = NLayerDiscriminator(input_nc, ndf, n_layers=3, norm_layer=norm_layer)
+    elif netD == 'n_layers':  # more options
+        net = NLayerDiscriminator(input_nc, ndf, n_layers_D, norm_layer=norm_layer)
+    elif netD == 'pixel':     # classify if each pixel is real or fake
+        net = PixelDiscriminator(input_nc, ndf, norm_layer=norm_layer)
+    else:
+        raise NotImplementedError('Discriminator model name [%s] is not recognized' % netD)
+    return init_net(net, init_type, init_gain, gpu_ids)
+
+
+##############################################################################
+# Classes
+##############################################################################
+class GANLoss(nn.Module):
+    """Define different GAN objectives.
+
+    The GANLoss class abstracts away the need to create the target label tensor
+    that has the same size as the input.
+    """
+
+    def __init__(self, gan_mode):
+        super(GANLoss, self).__init__()
+        target_real_label = random.randint(7, 12) * 0.1
+        target_fake_label = random.randint(0, 3) * 0.1
+        self.register_buffer('real_label', torch.tensor(target_real_label))
+        self.register_buffer('fake_label', torch.tensor(target_fake_label))
+        self.gan_mode = gan_mode
+        if gan_mode == 'lsgan':
+            self.loss = nn.MSELoss()
+        elif gan_mode == 'vanilla':
+            self.loss = nn.BCEWithLogitsLoss()
+        elif gan_mode in ['wgangp']:
+            self.loss = None
+        else:
+            raise NotImplementedError('gan mode %s not implemented' % gan_mode)
+    '''
+    def __init__(self, gan_mode, target_real_label=1.0, target_fake_label=0.0):
+        """ Initialize the GANLoss class.
+
+        Parameters:
+            gan_mode (str) - - the type of GAN objective. It currently supports vanilla, lsgan, and wgangp.
+            target_real_label (bool) - - label for a real image
+            target_fake_label (bool) - - label of a fake image
+
+        Note: Do not use sigmoid as the last layer of Discriminator.
+        LSGAN needs no sigmoid. vanilla GANs will handle it with BCEWithLogitsLoss.
+        """
+        super(GANLoss, self).__init__()
+        self.register_buffer('real_label', torch.tensor(target_real_label))
+        self.register_buffer('fake_label', torch.tensor(target_fake_label))
+        self.gan_mode = gan_mode
+        if gan_mode == 'lsgan':
+            self.loss = nn.MSELoss()
+        elif gan_mode == 'vanilla':
+            self.loss = nn.BCEWithLogitsLoss()
+        elif gan_mode in ['wgangp']:
+            self.loss = None
+        else:
+            raise NotImplementedError('gan mode %s not implemented' % gan_mode)
+    '''
+    def get_target_tensor(self, prediction, target_is_real):
+        """Create label tensors with the same size as the input.
+
+        Parameters:
+            prediction (tensor) - - tpyically the prediction from a discriminator
+            target_is_real (bool) - - if the ground truth label is for real images or fake images
+
+        Returns:
+            A label tensor filled with ground truth label, and with the size of the input
+        """
+
+        if target_is_real:
+            target_tensor = self.real_label
+        else:
+            target_tensor = self.fake_label
+        return target_tensor.expand_as(prediction)
+
+    def __call__(self, prediction, target_is_real):
+        """Calculate loss given Discriminator's output and grount truth labels.
+
+        Parameters:
+            prediction (tensor) - - tpyically the prediction output from a discriminator
+            target_is_real (bool) - - if the ground truth label is for real images or fake images
+
+        Returns:
+            the calculated loss.
+        """
+        if self.gan_mode in ['lsgan', 'vanilla']:
+            target_tensor = self.get_target_tensor(prediction, target_is_real)
+            loss = self.loss(prediction, target_tensor)
+        elif self.gan_mode == 'wgangp':
+            if target_is_real:
+                loss = -prediction.mean()
+            else:
+                loss = prediction.mean()
+        return loss
+
+
+def cal_gradient_penalty(netD, real_data, fake_data, device, type='mixed', constant=1.0, lambda_gp=10.0):
+    """Calculate the gradient penalty loss, used in WGAN-GP paper https://arxiv.org/abs/1704.00028
+
+    Arguments:
+        netD (network)              -- discriminator network
+        real_data (tensor array)    -- real images
+        fake_data (tensor array)    -- generated images from the generator
+        device (str)                -- GPU / CPU: from torch.device('cuda:{}'.format(self.gpu_ids[0])) if self.gpu_ids else torch.device('cpu')
+        type (str)                  -- if we mix real and fake data or not [real | fake | mixed].
+        constant (float)            -- the constant used in formula ( ||gradient||_2 - constant)^2
+        lambda_gp (float)           -- weight for this loss
+
+    Returns the gradient penalty loss
+    """
+    if lambda_gp > 0.0:
+        if type == 'real':   # either use real images, fake images, or a linear interpolation of two.
+            interpolatesv = real_data
+        elif type == 'fake':
+            interpolatesv = fake_data
+        elif type == 'mixed':
+            alpha = torch.rand(real_data.shape[0], 1, device=device)
+            alpha = alpha.expand(real_data.shape[0], real_data.nelement() // real_data.shape[0]).contiguous().view(*real_data.shape)
+            interpolatesv = alpha * real_data + ((1 - alpha) * fake_data)
+        else:
+            raise NotImplementedError('{} not implemented'.format(type))
+        interpolatesv.requires_grad_(True)
+        disc_interpolates = netD(interpolatesv)
+        gradients = torch.autograd.grad(outputs=disc_interpolates, inputs=interpolatesv,
+                                        grad_outputs=torch.ones(disc_interpolates.size()).to(device),
+                                        create_graph=True, retain_graph=True, only_inputs=True)
+        gradients = gradients[0].view(real_data.size(0), -1)  # flat the data
+        gradient_penalty = (((gradients + 1e-16).norm(2, dim=1) - constant) ** 2).mean() * lambda_gp        # added eps
+        return gradient_penalty, gradients
+    else:
+        return 0.0, None
+
+
+class ResnetGenerator(nn.Module):
+    """Resnet-based generator that consists of Resnet blocks between a few downsampling/upsampling operations.
+
+    We adapt Torch code and idea from Justin Johnson's neural style transfer project(https://github.com/jcjohnson/fast-neural-style)
+    """
+
+    def __init__(self, input_nc, output_nc, ngf=64, norm_layer=nn.BatchNorm2d, use_dropout=False, n_blocks=6, padding_type='reflect'):
+        """Construct a Resnet-based generator
+
+        Parameters:
+            input_nc (int)      -- the number of channels in input images
+            output_nc (int)     -- the number of channels in output images
+            ngf (int)           -- the number of filters in the last conv layer
+            norm_layer          -- normalization layer
+            use_dropout (bool)  -- if use dropout layers
+            n_blocks (int)      -- the number of ResNet blocks
+            padding_type (str)  -- the name of padding layer in conv layers: reflect | replicate | zero
+        """
+        assert(n_blocks >= 0)
+        super(ResnetGenerator, self).__init__()
+        if type(norm_layer) == functools.partial:
+            use_bias = norm_layer.func == nn.InstanceNorm2d
+        else:
+            use_bias = norm_layer == nn.InstanceNorm2d
+
+        model = [nn.ReflectionPad2d(3),
+                 nn.Conv2d(input_nc, ngf, kernel_size=7, padding=0, bias=use_bias),
+                 norm_layer(ngf),
+                 nn.ReLU(True)]
+
+        n_downsampling = 2
+        for i in range(n_downsampling):  # add downsampling layers
+            mult = 2 ** i
+            model += [nn.Conv2d(ngf * mult, ngf * mult * 2, kernel_size=3, stride=2, padding=1, bias=use_bias),
+                      norm_layer(ngf * mult * 2),
+                      nn.ReLU(True)]
+
+        mult = 2 ** n_downsampling
+        for i in range(n_blocks):       # add ResNet blocks
+            model += [ResnetBlock(ngf * mult, padding_type=padding_type, norm_layer=norm_layer, use_dropout=use_dropout, use_bias=use_bias)]
+
+        for i in range(n_downsampling):  # add upsampling layers
+            mult = 2 ** (n_downsampling - i)
+            model += [nn.ConvTranspose2d(ngf * mult, int(ngf * mult / 2),
+                                         kernel_size=3, stride=2,
+                                         padding=1, output_padding=1,
+                                         bias=use_bias),
+                      norm_layer(int(ngf * mult / 2)),
+                      nn.ReLU(True)]
+        model += [nn.ReflectionPad2d(3)]
+        model += [nn.Conv2d(ngf, output_nc, kernel_size=7, padding=0)]
+        model += [nn.Tanh()]
+
+        self.model = nn.Sequential(*model)
+
+    def forward(self, input):
+        """Standard forward"""
+        return self.model(input)
+
+
+class ResnetBlock(nn.Module):
+    """Define a Resnet block"""
+
+    def __init__(self, dim, padding_type, norm_layer, use_dropout, use_bias):
+        """Initialize the Resnet block
+
+        A resnet block is a conv block with skip connections
+        We construct a conv block with build_conv_block function,
+        and implement skip connections in <forward> function.
+        Original Resnet paper: https://arxiv.org/pdf/1512.03385.pdf
+        """
+        super(ResnetBlock, self).__init__()
+        self.conv_block = self.build_conv_block(dim, padding_type, norm_layer, use_dropout, use_bias)
+
+    def build_conv_block(self, dim, padding_type, norm_layer, use_dropout, use_bias):
+        """Construct a convolutional block.
+
+        Parameters:
+            dim (int)           -- the number of channels in the conv layer.
+            padding_type (str)  -- the name of padding layer: reflect | replicate | zero
+            norm_layer          -- normalization layer
+            use_dropout (bool)  -- if use dropout layers.
+            use_bias (bool)     -- if the conv layer uses bias or not
+
+        Returns a conv block (with a conv layer, a normalization layer, and a non-linearity layer (ReLU))
+        """
+        conv_block = []
+        p = 0
+        if padding_type == 'reflect':
+            conv_block += [nn.ReflectionPad2d(1)]
+        elif padding_type == 'replicate':
+            conv_block += [nn.ReplicationPad2d(1)]
+        elif padding_type == 'zero':
+            p = 1
+        else:
+            raise NotImplementedError('padding [%s] is not implemented' % padding_type)
+
+        conv_block += [nn.Conv2d(dim, dim, kernel_size=3, padding=p, bias=use_bias), norm_layer(dim), nn.ReLU(True)]
+        if use_dropout:
+            conv_block += [nn.Dropout(0.5)]
+
+        p = 0
+        if padding_type == 'reflect':
+            conv_block += [nn.ReflectionPad2d(1)]
+        elif padding_type == 'replicate':
+            conv_block += [nn.ReplicationPad2d(1)]
+        elif padding_type == 'zero':
+            p = 1
+        else:
+            raise NotImplementedError('padding [%s] is not implemented' % padding_type)
+        conv_block += [nn.Conv2d(dim, dim, kernel_size=3, padding=p, bias=use_bias), norm_layer(dim)]
+
+        return nn.Sequential(*conv_block)
+
+    def forward(self, x):
+        """Forward function (with skip connections)"""
+        out = x + self.conv_block(x)  # add skip connections
+        return out
+
+
+class UnetGenerator(nn.Module):
+    """Create a Unet-based generator"""
+
+    def __init__(self, input_nc, output_nc, num_downs, ngf=64, norm_layer=nn.BatchNorm2d, use_dropout=False):
+        """Construct a Unet generator
+        Parameters:
+            input_nc (int)  -- the number of channels in input images
+            output_nc (int) -- the number of channels in output images
+            num_downs (int) -- the number of downsamplings in UNet. For example, # if |num_downs| == 7,
+                                image of size 128x128 will become of size 1x1 # at the bottleneck
+            ngf (int)       -- the number of filters in the last conv layer
+            norm_layer      -- normalization layer
+
+        We construct the U-Net from the innermost layer to the outermost layer.
+        It is a recursive process.
+        """
+        super(UnetGenerator, self).__init__()
+        # construct unet structure
+        unet_block = UnetSkipConnectionBlock(ngf * 8, ngf * 8, input_nc=None, submodule=None, norm_layer=norm_layer, innermost=True)  # add the innermost layer
+        for i in range(num_downs - 5):          # add intermediate layers with ngf * 8 filters
+            unet_block = UnetSkipConnectionBlock(ngf * 8, ngf * 8, input_nc=None, submodule=unet_block, norm_layer=norm_layer, use_dropout=use_dropout)
+        # gradually reduce the number of filters from ngf * 8 to ngf
+        unet_block = UnetSkipConnectionBlock(ngf * 4, ngf * 8, input_nc=None, submodule=unet_block, norm_layer=norm_layer)
+        unet_block = UnetSkipConnectionBlock(ngf * 2, ngf * 4, input_nc=None, submodule=unet_block, norm_layer=norm_layer)
+        unet_block = UnetSkipConnectionBlock(ngf, ngf * 2, input_nc=None, submodule=unet_block, norm_layer=norm_layer)
+        self.model = UnetSkipConnectionBlock(output_nc, ngf, input_nc=input_nc, submodule=unet_block, outermost=True, norm_layer=norm_layer)  # add the outermost layer
+
+    def forward(self, input):
+        """Standard forward"""
+        return self.model(input)
+
+
+class UnetSkipConnectionBlock(nn.Module):
+    """Defines the Unet submodule with skip connection.
+        X -------------------identity----------------------
+        |-- downsampling -- |submodule| -- upsampling --|
+    """
+
+    def __init__(self, outer_nc, inner_nc, input_nc=None,
+                 submodule=None, outermost=False, innermost=False, norm_layer=nn.BatchNorm2d, use_dropout=False):
+        """Construct a Unet submodule with skip connections.
+
+        Parameters:
+            outer_nc (int) -- the number of filters in the outer conv layer
+            inner_nc (int) -- the number of filters in the inner conv layer
+            input_nc (int) -- the number of channels in input images/features
+            submodule (UnetSkipConnectionBlock) -- previously defined submodules
+            outermost (bool)    -- if this module is the outermost module
+            innermost (bool)    -- if this module is the innermost module
+            norm_layer          -- normalization layer
+            use_dropout (bool)  -- if use dropout layers.
+        """
+        super(UnetSkipConnectionBlock, self).__init__()
+        self.outermost = outermost
+        if type(norm_layer) == functools.partial:
+            use_bias = norm_layer.func == nn.InstanceNorm2d
+        else:
+            use_bias = norm_layer == nn.InstanceNorm2d
+        if input_nc is None:
+            input_nc = outer_nc
+        downconv = nn.Conv2d(input_nc, inner_nc, kernel_size=4,
+                             stride=2, padding=1, bias=use_bias)
+        downrelu = nn.LeakyReLU(0.2, True)
+        downnorm = norm_layer(inner_nc)
+        uprelu = nn.ReLU(True)
+        upnorm = norm_layer(outer_nc)
+
+        if outermost:
+            upconv = nn.ConvTranspose2d(inner_nc * 2, outer_nc,
+                                        kernel_size=4, stride=2,
+                                        padding=1)
+            down = [downconv]
+            up = [uprelu, upconv, nn.Tanh()]
+            model = down + [submodule] + up
+        elif innermost:
+            upconv = nn.ConvTranspose2d(inner_nc, outer_nc,
+                                        kernel_size=4, stride=2,
+                                        padding=1, bias=use_bias)
+            down = [downrelu, downconv]
+            up = [uprelu, upconv, upnorm]
+            model = down + up
+        else:
+            upconv = nn.ConvTranspose2d(inner_nc * 2, outer_nc,
+                                        kernel_size=4, stride=2,
+                                        padding=1, bias=use_bias)
+            down = [downrelu, downconv, downnorm]
+            up = [uprelu, upconv, upnorm]
+
+            if use_dropout:
+                model = down + [submodule] + up + [nn.Dropout(0.5)]
+            else:
+                model = down + [submodule] + up
+
+        self.model = nn.Sequential(*model)
+
+    def forward(self, x):
+        if self.outermost:
+            return self.model(x)
+        else:   # add skip connections
+            return torch.cat([x, self.model(x)], 1)
+
+
+class NLayerDiscriminator(nn.Module):
+    """Defines a PatchGAN discriminator"""
+
+    def __init__(self, input_nc, ndf=64, n_layers=3, norm_layer=nn.BatchNorm2d):
+        """Construct a PatchGAN discriminator
+
+        Parameters:
+            input_nc (int)  -- the number of channels in input images
+            ndf (int)       -- the number of filters in the last conv layer
+            n_layers (int)  -- the number of conv layers in the discriminator
+            norm_layer      -- normalization layer
+        """
+        super(NLayerDiscriminator, self).__init__()
+        if type(norm_layer) == functools.partial:  # no need to use bias as BatchNorm2d has affine parameters
+            use_bias = norm_layer.func == nn.InstanceNorm2d
+        else:
+            use_bias = norm_layer == nn.InstanceNorm2d
+
+        kw = 4
+        padw = 1
+        sequence = [nn.Conv2d(input_nc, ndf, kernel_size=kw, stride=2, padding=padw), nn.LeakyReLU(0.2, True)]
+        nf_mult = 1
+        nf_mult_prev = 1
+        for n in range(1, n_layers):  # gradually increase the number of filters
+            nf_mult_prev = nf_mult
+            nf_mult = min(2 ** n, 8)
+            sequence += [
+                nn.Conv2d(ndf * nf_mult_prev, ndf * nf_mult, kernel_size=kw, stride=2, padding=padw, bias=use_bias),
+                norm_layer(ndf * nf_mult),
+                nn.LeakyReLU(0.2, True)
+            ]
+
+        nf_mult_prev = nf_mult
+        nf_mult = min(2 ** n_layers, 8)
+        sequence += [
+            nn.Conv2d(ndf * nf_mult_prev, ndf * nf_mult, kernel_size=kw, stride=1, padding=padw, bias=use_bias),
+            norm_layer(ndf * nf_mult),
+            nn.LeakyReLU(0.2, True)
+        ]
+
+        sequence += [nn.Conv2d(ndf * nf_mult, 1, kernel_size=kw, stride=1, padding=padw)]  # output 1 channel prediction map
+        self.model = nn.Sequential(*sequence)
+
+    def forward(self, input):
+        """Standard forward."""
+        return self.model(input)
+
+
+class PixelDiscriminator(nn.Module):
+    """Defines a 1x1 PatchGAN discriminator (pixelGAN)"""
+
+    def __init__(self, input_nc, ndf=64, norm_layer=nn.BatchNorm2d):
+        """Construct a 1x1 PatchGAN discriminator
+
+        Parameters:
+            input_nc (int)  -- the number of channels in input images
+            ndf (int)       -- the number of filters in the last conv layer
+            norm_layer      -- normalization layer
+        """
+        super(PixelDiscriminator, self).__init__()
+        if type(norm_layer) == functools.partial:  # no need to use bias as BatchNorm2d has affine parameters
+            use_bias = norm_layer.func == nn.InstanceNorm2d
+        else:
+            use_bias = norm_layer == nn.InstanceNorm2d
+
+        self.net = [
+            nn.Conv2d(input_nc, ndf, kernel_size=1, stride=1, padding=0),
+            nn.LeakyReLU(0.2, True),
+            nn.Conv2d(ndf, ndf * 2, kernel_size=1, stride=1, padding=0, bias=use_bias),
+            norm_layer(ndf * 2),
+            nn.LeakyReLU(0.2, True),
+            nn.Conv2d(ndf * 2, 1, kernel_size=1, stride=1, padding=0, bias=use_bias)]
+
+        self.net = nn.Sequential(*self.net)
+
+    def forward(self, input):
+        """Standard forward."""
+        return self.net(input)
+
+class Self_Attention(nn.Module):
+
+    def __init__(self, in_dim, activation):
+        super(Self_Attention, self).__init__()
+        self.chanel_in = in_dim
+        self.activation = activation
+        ##  下面的query_conv，key_conv，value_conv即对应Wg,Wf,Wh
+        self.query_conv = nn.Conv2d(in_channels=in_dim, out_channels=in_dim // 8, kernel_size=1)  # 即得到C^ X C
+        self.key_conv = nn.Conv2d(in_channels=in_dim, out_channels=in_dim // 8, kernel_size=1)  # 即得到C^ X C
+        self.value_conv = nn.Conv2d(in_channels=in_dim, out_channels=in_dim, kernel_size=1)  # 即得到C X C
+        self.gamma = nn.Parameter(torch.zeros(1))  # 这里即是计算最终输出的时候的伽马值，初始化为0
+
+        self.softmax = nn.Softmax(dim=-1)
+
+    def forward(self, x):
+        m_batchsize, C, width, height = x.size()
+        ##  下面的proj_query，proj_key都是C^ X C X C X N= C^ X N
+        proj_query = self.query_conv(x).view(m_batchsize, -1, width * height).permute(0, 2, 1)  # B X CX(N),permute即为转置
+        proj_key = self.key_conv(x).view(m_batchsize, -1, width * height)  # B X C x (*W*H)
+        energy = torch.bmm(proj_query, proj_key)  # transpose check，进行点乘操作
+        attention = self.softmax(energy)  # BX (N) X (N)
+        proj_value = self.value_conv(x).view(m_batchsize, -1, width * height)  # B X C X N
+
+        out = torch.bmm(proj_value, attention.permute(0, 2, 1))
+        out = out.view(m_batchsize, C, width, height)
+
+        out = self.gamma * out + x
+        return out
+
+class ResnetBlockWithAttention(nn.Module):
+    """Define a Resnet block"""
+
+    def __init__(self, dim, padding_type, norm_layer, use_dropout, use_bias):
+        """Initialize the Resnet block
+
+        A resnet block is a conv block with skip connections
+        We construct a conv block with build_conv_block function,
+        and implement skip connections in <forward> function.
+        Original Resnet paper: https://arxiv.org/pdf/1512.03385.pdf
+        """
+        super(ResnetBlockWithAttention, self).__init__()
+        self.conv_block = self.build_conv_block(dim, padding_type, norm_layer, use_dropout, use_bias)
+        self.attention=Self_Attention(dim,'relu')
+
+    def build_conv_block(self, dim, padding_type, norm_layer, use_dropout, use_bias):
+        """Construct a convolutional block.
+
+        Parameters:
+            dim (int)           -- the number of channels in the conv layer.
+            padding_type (str)  -- the name of padding layer: reflect | replicate | zero
+            norm_layer          -- normalization layer
+            use_dropout (bool)  -- if use dropout layers.
+            use_bias (bool)     -- if the conv layer uses bias or not
+
+        Returns a conv block (with a conv layer, a normalization layer, and a non-linearity layer (ReLU))
+        """
+        conv_block = []
+        p = 0
+        if padding_type == 'reflect':
+            conv_block += [nn.ReflectionPad2d(1)]
+        elif padding_type == 'replicate':
+            conv_block += [nn.ReplicationPad2d(1)]
+        elif padding_type == 'zero':
+            p = 1
+        else:
+            raise NotImplementedError('padding [%s] is not implemented' % padding_type)
+
+        conv_block += [nn.Conv2d(dim, dim, kernel_size=3, padding=p, bias=use_bias), norm_layer(dim), nn.ReLU(True)]
+        if use_dropout:
+            conv_block += [nn.Dropout(0.5)]
+
+        p = 0
+        if padding_type == 'reflect':
+            conv_block += [nn.ReflectionPad2d(1)]
+        elif padding_type == 'replicate':
+            conv_block += [nn.ReplicationPad2d(1)]
+        elif padding_type == 'zero':
+            p = 1
+        else:
+            raise NotImplementedError('padding [%s] is not implemented' % padding_type)
+        conv_block += [nn.Conv2d(dim, dim, kernel_size=3, padding=p, bias=use_bias), norm_layer(dim)]
+
+        return nn.Sequential(*conv_block)
+
+    def forward(self, x):
+        """Forward function (with skip connections)"""
+        out = x + self.conv_block(x)+self.attention(x)  # add skip connections
+        return out
+
+
+class ResnetGeneratorWithAttentionAndSkipConnection(nn.Module):
+    def __init__(self, input_nc, output_nc, ngf=64, norm_layer=nn.BatchNorm2d, use_dropout=False, n_blocks=6,
+                 padding_type='reflect'):
+        super(ResnetGeneratorWithAttentionAndSkipConnection, self).__init__()
+        assert n_blocks >= 0
+        if type(norm_layer) == functools.partial:
+            use_bias = norm_layer.func == nn.InstanceNorm2d
+        else:
+            use_bias = norm_layer == nn.InstanceNorm2d
+
+        self.n_downsampling = n_downsampling = 2
+
+        # Initial convolution block
+        self.encoder_initial = nn.Sequential(
+            nn.ReflectionPad2d(3),
+            nn.Conv2d(input_nc, ngf, kernel_size=7, padding=0, bias=use_bias),
+            norm_layer(ngf),
+            nn.ReLU(True)
+        )
+
+        # Downsampling layers
+        self.encoder_down = nn.ModuleList()
+        for i in range(n_downsampling):
+            mult = 2 ** i
+            self.encoder_down.append(nn.Sequential(
+                nn.Conv2d(ngf * mult, ngf * mult * 2, kernel_size=3, stride=2, padding=1, bias=use_bias),
+                norm_layer(ngf * mult * 2),
+                nn.ReLU(True)
+            ))
+
+        # Middle ResNet blocks with attention
+        mult = 2 ** n_downsampling
+        self.middle = nn.Sequential()
+        for i in range(n_blocks):
+            self.middle.add_module(f'resnet_block_{i}', ResnetBlockWithAttention(
+                ngf * mult, padding_type=padding_type, norm_layer=norm_layer,
+                use_dropout=use_dropout, use_bias=use_bias))
+
+        # Upsampling layers and skip connections
+        self.decoder_up = nn.ModuleList()
+        self.skip_convs = nn.ModuleList()
+        self.fusion_convs = nn.ModuleList()
+        self.fusion_norms = nn.ModuleList()
+        self.fusion_act = nn.ModuleList()
+
+        for j in range(n_downsampling):
+            # Upsampling layer
+            mult = 2 ** (n_downsampling - j)
+            self.decoder_up.append(nn.Sequential(
+                nn.ConvTranspose2d(ngf * mult, int(ngf * mult / 2),
+                                   kernel_size=3, stride=2,
+                                   padding=1, output_padding=1,
+                                   bias=use_bias),
+                norm_layer(int(ngf * mult / 2)),
+                nn.ReLU(True))
+            )
+
+            # Skip connection processing
+            encoder_layer_index = n_downsampling - j - 1
+            encoder_mult = 2 ** (encoder_layer_index + 1)
+            decoder_mult_out = int(ngf * mult / 2)
+
+            self.skip_convs.append(nn.Conv2d(ngf * encoder_mult, decoder_mult_out, kernel_size=1, bias=use_bias))
+
+            self.fusion_convs.append(
+                nn.Conv2d(decoder_mult_out * 2, decoder_mult_out, kernel_size=3, padding=1, bias=use_bias))
+            self.fusion_norms.append(norm_layer(decoder_mult_out))
+            self.fusion_act.append(nn.ReLU(True))
+
+        # Final output layer
+        self.output = nn.Sequential(
+            nn.ReflectionPad2d(3),
+            nn.Conv2d(ngf, output_nc, kernel_size=7, padding=0),
+            nn.Tanh()
+        )
+
+    def forward(self, input):
+        # Encoder
+        x = self.encoder_initial(input)
+        skips = []
+        for down_layer in self.encoder_down:
+            x = down_layer(x)
+            skips.append(x)
+
+        # Middle blocks
+        x = self.middle(x)
+
+        # Decoder with skip connections
+        for j in range(self.n_downsampling):
+            x = self.decoder_up[j](x)
+
+            # Get corresponding skip connection
+            skip_index = self.n_downsampling - j - 1
+            skip = skips[skip_index]
+
+            # Process skip connection
+            skip = F.interpolate(skip, size=x.shape[2:], mode='bilinear', align_corners=False)
+            skip = self.skip_convs[j](skip)
+
+            # Concatenate and fuse
+            x = torch.cat([x, skip], dim=1)
+            x = self.fusion_convs[j](x)
+            x = self.fusion_norms[j](x)
+            x = self.fusion_act[j](x)
+
+        # Final output
+        return self.output(x)
+
+
+class ResnetGeneratorWithAttention(nn.Module):
+    """Resnet-based generator that consists of Resnet blocks between a few downsampling/upsampling operations.
+
+    We adapt Torch code and idea from Justin Johnson's neural style transfer project(https://github.com/jcjohnson/fast-neural-style)
+    """
+
+    def __init__(self, input_nc, output_nc, ngf=64, norm_layer=nn.BatchNorm2d, use_dropout=False, n_blocks=6, padding_type='reflect'):
+        """Construct a Resnet-based generator
+
+        Parameters:
+            input_nc (int)      -- the number of channels in input images
+            output_nc (int)     -- the number of channels in output images
+            ngf (int)           -- the number of filters in the last conv layer
+            norm_layer          -- normalization layer
+            use_dropout (bool)  -- if use dropout layers
+            n_blocks (int)      -- the number of ResNet blocks
+            padding_type (str)  -- the name of padding layer in conv layers: reflect | replicate | zero
+        """
+        assert(n_blocks >= 0)
+        super(ResnetGeneratorWithAttention, self).__init__()
+        if type(norm_layer) == functools.partial:
+            use_bias = norm_layer.func == nn.InstanceNorm2d
+        else:
+            use_bias = norm_layer == nn.InstanceNorm2d
+
+        model = [nn.ReflectionPad2d(3),
+                 nn.Conv2d(input_nc, ngf, kernel_size=7, padding=0, bias=use_bias),
+                 norm_layer(ngf),
+                 nn.ReLU(True)]
+
+        n_downsampling = 2
+        for i in range(n_downsampling):  # add downsampling layers
+            mult = 2 ** i
+            model += [nn.Conv2d(ngf * mult, ngf * mult * 2, kernel_size=3, stride=2, padding=1, bias=use_bias),
+                      norm_layer(ngf * mult * 2),
+                      nn.ReLU(True)]
+
+        mult = 2 ** n_downsampling
+        for i in range(n_blocks):       # add ResNet blocks
+
+            model += [ResnetBlockWithAttention(ngf * mult, padding_type=padding_type, norm_layer=norm_layer, use_dropout=use_dropout, use_bias=use_bias)]
+
+        for i in range(n_downsampling):  # add upsampling layers
+            mult = 2 ** (n_downsampling - i)
+            model += [nn.ConvTranspose2d(ngf * mult, int(ngf * mult / 2),
+                                         kernel_size=3, stride=2,
+                                         padding=1, output_padding=1,
+                                         bias=use_bias),
+                      norm_layer(int(ngf * mult / 2)),
+                      nn.ReLU(True)]
+        model += [nn.ReflectionPad2d(3)]
+        model += [nn.Conv2d(ngf, output_nc, kernel_size=7, padding=0)]
+        model += [nn.Tanh()]
+
+        self.model = nn.Sequential(*model)
+
+    def forward(self, input):
+        """Standard forward"""
+        return self.model(input)
+
+
+class Unet_SEA_ResnetGenerator(nn.Module):
+
+    def __init__(self, input_nc, output_nc, ngf=64, norm_layer=nn.BatchNorm2d, use_dropout=False,
+                 padding_type='reflect'):
+        super(Unet_SEA_ResnetGenerator, self).__init__()
+        if type(norm_layer) == functools.partial:
+            use_bias = norm_layer.func == nn.InstanceNorm2d
+        else:
+            use_bias = norm_layer == nn.InstanceNorm2d
+        self.pad = nn.ReflectionPad2d(3)
+        self.Down_conv1 = nn.Conv2d(input_nc, ngf, kernel_size=7, padding=0, bias=use_bias)  # 下采样第一层
+        self.conv_norm = norm_layer(input_nc)
+        self.relu = nn.ReLU(True)
+        self.Down_conv2 = nn.Conv2d(ngf, ngf * 2, kernel_size=3, stride=2, padding=1, bias=use_bias)  # 下采样第二层
+        self.SA = Self_Attention(ngf * 2, 'relu')
+        self.Down_conv3 = nn.Conv2d(ngf * 2, ngf * 4, kernel_size=3, stride=2, padding=1, bias=use_bias)  # 下采样第三层
+        self.Sa_block_1 = SEA_ResnetBlock_1(ngf * 4, padding_type=padding_type, norm_layer=norm_layer,
+                                      use_dropout=use_dropout, use_bias=use_bias)
+        self.Sa_block_2 = SEA_ResnetBlock_1(ngf * 4, padding_type=padding_type, norm_layer=norm_layer,
+                                                  use_dropout=use_dropout, use_bias=use_bias)
+        self.Sa_block_3 = SEA_ResnetBlock_1(ngf * 4, padding_type=padding_type, norm_layer=norm_layer, use_dropout=use_dropout,
+                                  use_bias=use_bias)
+        self.Up_conv1 = nn.ConvTranspose2d(ngf * 4 * 2, ngf * 2, kernel_size=3, stride=2, padding=1, output_padding=1,
+                                           bias=use_bias)
+        self.Up_conv2 = nn.ConvTranspose2d(ngf * 2 * 2, ngf, kernel_size=3, stride=2, padding=1, output_padding=1,
+                                           bias=use_bias)
+        self.Up_conv3 = nn.Conv2d(ngf * 2, output_nc, kernel_size=7, padding=0)
+        self.tan = nn.Tanh()
+
+    def forward(self, x):
+        x1 = self.relu(self.conv_norm(self.Down_conv1(self.pad(x))))
+        x2 = self.relu(self.conv_norm(self.Down_conv2(x1)))
+        x3 = self.relu(self.conv_norm(self.Down_conv3(x2)))
+        x4 = self.Sa_block_3(self.Sa_block_2(self.Sa_block_1(x3)))
+        x = torch.cat([x4, x3], 1)
+        x = self.relu(self.conv_norm(self.Up_conv1(x)))
+        x = torch.cat([x, x2], 1)
+        x = self.relu(self.conv_norm(self.Up_conv2(x)))
+        x = torch.cat([x, x1], 1)
+        x = self.tan(self.Up_conv3(self.pad(x)))
+        return x
+
+class SEA_ResnetBlock_1(nn.Module):
+
+    def __init__(self, dim, padding_type, norm_layer, use_dropout, use_bias):
+        super(SEA_ResnetBlock_1, self).__init__()
+        self.conv_block = self.build_conv_block(dim, padding_type, norm_layer, use_dropout, use_bias)
+        self.self_attention = Self_Attention(dim, 'relu')
+
+    def build_conv_block(self, dim, padding_type, norm_layer, use_dropout, use_bias):
+        conv_block = []
+        p = 0
+        if padding_type == 'reflect':
+            conv_block += [nn.ReflectionPad2d(1)]
+        elif padding_type == 'replicate':
+            conv_block += [nn.ReplicationPad2d(1)]
+        elif padding_type == 'zero':
+            p = 1
+        else:
+            raise NotImplementedError('padding [%s] is not implemented' % padding_type)
+
+        conv_block += [nn.Conv2d(dim, dim, kernel_size=3, padding=p, bias=use_bias), norm_layer(dim), nn.ReLU(True)]
+        if use_dropout:
+            conv_block += [nn.Dropout(0.5)]
+
+        p = 0
+        if padding_type == 'reflect':
+            conv_block += [nn.ReflectionPad2d(1)]
+        elif padding_type == 'replicate':
+            conv_block += [nn.ReplicationPad2d(1)]
+        elif padding_type == 'zero':
+            p = 1
+        else:
+            raise NotImplementedError('padding [%s] is not implemented' % padding_type)
+        conv_block += [nn.Conv2d(dim, dim, kernel_size=3, padding=p, bias=use_bias), norm_layer(dim)]
+
+        return nn.Sequential(*conv_block)
+
+    def forward(self, x):
+        out = self.self_attention(x) + self.conv_block(x) + x  # add skip connections
+        return out
+
+
+class Discriminator(nn.Module):
+    def __init__(self, input_nc, ndf=64, n_layers=3, norm_layer=nn.BatchNorm2d):
+        super(Discriminator, self).__init__()
+        # 256 x 256
+        self.conv1 = nn.Sequential(nn.Conv2d(input_nc, ndf, kernel_size=3, stride=1, padding=1),
+                                   nn.ELU(True),
+                                   norm_layer(ndf),
+                                   nn.Conv2d(ndf,ndf,kernel_size=3,stride=1,padding=1),
+                                   nn.ELU(True),
+                                   nn.Conv2d(ndf,ndf,kernel_size=3,stride=1,padding=1),
+                                   norm_layer(ndf),
+                                   nn.AvgPool2d(2,2))
+
+
+        # 128 x 128
+        self.conv2 = nn.Sequential(nn.Conv2d(ndf,2*ndf,kernel_size=3,stride=1,padding=1),
+                                   nn.ELU(True),
+                                   norm_layer(ndf*2),
+                                   nn.Conv2d(2*ndf,2*ndf,kernel_size=3,stride=1,padding=1),
+                                   nn.ELU(True),
+                                   norm_layer(ndf*2),
+                                   nn.AvgPool2d(2,2))
+        # 64 x 64
+        self.conv3 = nn.Sequential(nn.Conv2d(2*ndf,3*ndf,kernel_size=3,stride=1,padding=1),
+                                   nn.ELU(True),
+                                   norm_layer(ndf*3),
+                                   nn.Conv2d(3*ndf,3*ndf,kernel_size=3,stride=1,padding=1),
+                                   nn.ELU(True),
+                                   norm_layer(ndf*3),
+                                   nn.AvgPool2d(2,2))
+        # 32 x 32
+        self.conv4 = nn.Sequential(nn.Conv2d(3*ndf,4*ndf,kernel_size=3,stride=1,padding=1),
+                                   nn.ELU(True),
+                                   norm_layer(ndf*4),
+                                   nn.Conv2d(4*ndf,4*ndf,kernel_size=3,stride=1,padding=1),
+                                   nn.ELU(True),
+                                   norm_layer(ndf*4),
+                                   nn.AvgPool2d(2,2))
+        # 16 x 16
+        self.conv5 = nn.Sequential(nn.Conv2d(4*ndf,5*ndf,kernel_size=3,stride=1,padding=1),
+                                   nn.ELU(True),
+                                   norm_layer(ndf*5),
+                                   nn.Conv2d(5*ndf,5*ndf,kernel_size=3,stride=1,padding=1),
+                                   nn.ELU(True),
+                                   norm_layer(ndf*5),
+                                   nn.AvgPool2d(2,2))
+        # 8 x 8
+
+
+        self.embed1 = nn.Linear(ndf * 5 * 8 * 8, 64)
+        self.embed2 = nn.Linear(64, ndf * 8 * 8)
+
+        # 8 x 8
+        self.deconv1 = nn.Sequential(nn.Conv2d(ndf,ndf,kernel_size=3,stride=1,padding=1),
+                                     nn.ELU(True),
+                                     nn.Conv2d(ndf,ndf,kernel_size=3,stride=1,padding=1),
+                                     nn.ELU(True),
+                                     nn.Upsample(size=None,scale_factor=2,mode='bilinear',align_corners=False))
+        # 16 x 16
+        self.deconv2 = nn.Sequential(nn.Conv2d(ndf,ndf,kernel_size=3,stride=1,padding=1),
+                                     nn.ELU(True),
+                                     nn.Conv2d(ndf,ndf,kernel_size=3,stride=1,padding=1),
+                                     nn.ELU(True),
+                                     nn.Upsample(size=None,scale_factor=2,mode='bilinear',align_corners=False))
+        # 32 x 32
+        self.deconv3 = nn.Sequential(nn.Conv2d(ndf,ndf,kernel_size=3,stride=1,padding=1),
+                                     nn.ELU(True),
+                                     nn.Conv2d(ndf,ndf,kernel_size=3,stride=1,padding=1),
+                                     nn.ELU(True),
+                                     nn.Upsample(size=None,scale_factor=2,mode='bilinear',align_corners=False))
+        # 64 x 64
+        self.deconv4 = nn.Sequential(nn.Conv2d(ndf,ndf,kernel_size=3,stride=1,padding=1),
+                                     nn.ELU(True),
+                                     nn.Conv2d(ndf,ndf,kernel_size=3,stride=1,padding=1),
+                                     nn.ELU(True),
+                                     nn.Upsample(size=None,scale_factor=2,mode='bilinear',align_corners=False))
+        # 128 x 128
+        self.deconv5 = nn.Sequential(nn.Conv2d(ndf,ndf,kernel_size=3,stride=1,padding=1),
+                                     nn.ELU(True),
+                                     nn.Conv2d(ndf,ndf,kernel_size=3,stride=1,padding=1),
+                                     nn.ELU(True),
+                                     nn.Upsample(size=None,scale_factor=2,mode='bilinear',align_corners=False))
+        # 256 x 256
+        self.deconv6 = nn.Sequential(nn.Conv2d(ndf, ndf, kernel_size=3, stride=1, padding=1),
+                                     nn.ELU(True),
+                                     nn.Conv2d(ndf, ndf, kernel_size=3, stride=1, padding=1),
+                                     nn.ELU(True),
+                                     nn.Conv2d(ndf, input_nc, kernel_size=3, stride=1, padding=1),
+                                     nn.Tanh())
+        self.ndf = ndf
+
+    def forward(self, x):
+        out = self.conv1(x)
+        out = self.conv2(out)
+        out = self.conv3(out)
+        out = self.conv4(out)
+        out = self.conv5(out)
+        out = out.view(out.size(0), self.ndf * 5 * 8 * 8)
+
+        out = self.embed1(out)
+        out = self.embed2(out)
+        out = out.view(out.size(0), self.ndf, 8, 8)
+        out = self.deconv1(out)
+        out = self.deconv2(out)
+        out = self.deconv3(out)
+        out = self.deconv4(out)
+        out = self.deconv5(out)
+        out = self.deconv6(out)
+        return out

cyclegan_model/model/test_model.py

@@ -0,0 +1,70 @@
+from .base_model import BaseModel
+from . import networks
+
+
+class TestModel(BaseModel):
+    """ This TesteModel can be used to generate CycleGAN results for only one direction.
+    This model will automatically set '--dataset_mode single', which only loads the images from one collection.
+
+    See the test instruction for more details.
+    """
+    @staticmethod
+    def modify_commandline_options(parser, is_train=True):
+        """Add new dataset-specific options, and rewrite default values for existing options.
+
+        Parameters:
+            parser          -- original option parser
+            is_train (bool) -- whether training phase or test phase. You can use this flag to add training-specific or test-specific options.
+
+        Returns:
+            the modified parser.
+
+        The model can only be used during test time. It requires '--dataset_mode single'.
+        You need to specify the network using the option '--model_suffix'.
+        """
+        assert not is_train, 'TestModel cannot be used during training time'
+        parser.set_defaults(dataset_mode='single')
+        parser.add_argument('--model_suffix', type=str, default='', help='In checkpoints_dir, [epoch]_net_G[model_suffix].pth will be loaded as the generator.')
+
+        return parser
+
+    def __init__(self, opt):
+        """Initialize the pix2pix class.
+
+        Parameters:
+            opt (Option class)-- stores all the experiment flags; needs to be a subclass of BaseOptions
+        """
+        assert(not opt.isTrain)
+        BaseModel.__init__(self, opt)
+        # specify the training losses you want to print out. The training/test scripts  will call <BaseModel.get_current_losses>
+        self.loss_names = []
+        # specify the images you want to save/display. The training/test scripts  will call <BaseModel.get_current_visuals>
+        self.visual_names = ['real', 'fake']
+        # specify the models you want to save to the disk. The training/test scripts will call <BaseModel.save_networks> and <BaseModel.load_networks>
+        self.model_names = ['G' + opt.model_suffix]  # only generator is needed.
+        self.netG = networks.define_G(opt.input_nc, opt.output_nc, opt.ngf, opt.netG,
+                                      opt.norm, not opt.no_dropout, opt.init_type, opt.init_gain, self.gpu_ids)
+
+        # assigns the model to self.netG_[suffix] so that it can be loaded
+        # please see <BaseModel.load_networks>
+        setattr(self, 'netG' + opt.model_suffix, self.netG)  # store netG in self.
+
+    def set_input(self, input):
+        """Unpack input data from the dataloader and perform necessary pre-processing steps.
+
+        Parameters:
+            input: a dictionary that contains the data itself and its metadata information.
+
+        We need to use 'single_dataset' dataset mode. It only load images from one domain.
+        """
+        self.real = input['A'].to(self.device)
+        self.image_paths = input['A_paths']
+
+    def forward(self):
+        """Run forward pass."""
+        self.fake = self.netG(self.real)  # G(real)
+
+    def optimize_parameters(self):
+        """No optimization for test model."""
+        pass
+

cyclegan_model/options/base_options.py

@@ -0,0 +1,138 @@
+import argparse
+import os
+from cyclegan_model.util import util
+import torch
+import cyclegan_model.model as model
+import cyclegan_model.data as data
+
+
+class BaseOptions():
+    """This class defines options used during both training and test time.
+
+    It also implements several helper functions such as parsing, printing, and saving the options.
+    It also gathers additional options defined in <modify_commandline_options> functions in both dataset class and model class.
+    """
+
+    def __init__(self):
+        """Reset the class; indicates the class hasn't been initailized"""
+        self.initialized = False
+
+    def initialize(self, parser):
+        """Define the common options that are used in both training and test."""
+        # basic parameters
+        parser.add_argument('--dataroot',type=str,help='path to images (should have subfolders trainA, trainB, valA, valB, etc)')
+        parser.add_argument('--name', type=str, default='ostracoda_cyclegan', help='name of the experiment. It decides where to store samples and models')
+        parser.add_argument('--gpu_ids', type=str, default='0', help='gpu ids: e.g. 0  0,1,2, 0,2. use -1 for CPU')
+        parser.add_argument('--checkpoints_dir', type=str, default='./checkpoints', help='models are saved here')
+        # model parameters
+        parser.add_argument('--model', type=str, default='cycle_gan', help='chooses which model to use. [cycle_gan | pix2pix | test | colorization]')
+        parser.add_argument('--input_nc', type=int, default=3, help='# of input image channels: 3 for RGB and 1 for grayscale')
+        parser.add_argument('--output_nc', type=int, default=3, help='# of output image channels: 3 for RGB and 1 for grayscale')
+        parser.add_argument('--ngf', type=int, default=64, help='# of gen filters in the last conv layer')
+        parser.add_argument('--ndf', type=int, default=64, help='# of discrim filters in the first conv layer')
+        parser.add_argument('--netD', type=str, default='basic', help='specify discriminator architecture [basic | n_layers | pixel]. The basic model is a 70x70 PatchGAN. n_layers allows you to specify the layers in the discriminator')
+        parser.add_argument('--netG', type=str, default='resnet_9blocks', help='specify generator architecture [resnet_9blocks | resnet_6blocks | unet_256 | unet_128]')
+        parser.add_argument('--n_layers_D', type=int, default=3, help='only used if netD==n_layers')
+        parser.add_argument('--norm', type=str, default='instance', help='instance normalization or batch normalization [instance | batch | none]')
+        parser.add_argument('--init_type', type=str, default='normal', help='network initialization [normal | xavier | kaiming | orthogonal]')
+        parser.add_argument('--init_gain', type=float, default=0.02, help='scaling factor for normal, xavier and orthogonal.')
+        parser.add_argument('--no_dropout', action='store_true', help='no dropout for the generator')
+        # dataset parameters
+        parser.add_argument('--dataset_mode', type=str, default='unaligned', help='chooses how datasets are loaded. [unaligned | aligned | single | colorization]')
+        parser.add_argument('--direction', type=str, default='AtoB', help='AtoB or BtoA')
+        parser.add_argument('--serial_batches', action='store_true', help='if true, takes images in order to make batches, otherwise takes them randomly')
+        parser.add_argument('--num_threads', default=4, type=int, help='# threads for loading data')
+        parser.add_argument('--batch_size', type=int, default=1, help='input batch size')
+        parser.add_argument('--load_size', type=int, default=286, help='scale images to this size')
+        parser.add_argument('--crop_size', type=int, default=256, help='then crop to this size')
+        parser.add_argument('--max_dataset_size', type=int, default=float("inf"), help='Maximum number of samples allowed per dataset. If the dataset directory contains more than max_dataset_size, only a subset is loaded.')
+        parser.add_argument('--preprocess', type=str, default='resize_and_crop', help='scaling and cropping of images at load time [resize_and_crop | crop | scale_width | scale_width_and_crop | none]')
+        parser.add_argument('--no_flip', action='store_true', help='if specified, do not flip the images for data augmentation')
+        parser.add_argument('--display_winsize', type=int, default=256, help='display window size for both visdom and HTML')
+        # additional parameters
+        parser.add_argument('--epoch', type=str, default='latest', help='which epoch to load? set to latest to use latest cached model')
+        parser.add_argument('--load_iter', type=int, default='0', help='which iteration to load? if load_iter > 0, the code will load models by iter_[load_iter]; otherwise, the code will load models by [epoch]')
+        parser.add_argument('--verbose', action='store_true', help='if specified, print more debugging information')
+        parser.add_argument('--suffix', default='', type=str, help='customized suffix: opt.name = opt.name + suffix: e.g., {model}_{netG}_size{load_size}')
+        # wandb parameters
+        parser.add_argument('--use_wandb', action='store_true', help='if specified, then init wandb logging')
+        parser.add_argument('--wandb_project_name', type=str, default='CycleGAN-and-pix2pix', help='specify wandb project name')
+        self.initialized = True
+        return parser
+
+    def gather_options(self):
+        """Initialize our parser with basic options(only once).
+        Add additional model-specific and dataset-specific options.
+        These options are defined in the <modify_commandline_options> function
+        in model and dataset classes.
+        """
+        if not self.initialized:  # check if it has been initialized
+            parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
+            parser = self.initialize(parser)
+
+        # get the basic options
+        opt, _ = parser.parse_known_args()
+
+        # modify model-related parser options
+        model_name = opt.model
+        model_option_setter = model.get_option_setter(model_name)
+        parser = model_option_setter(parser, self.isTrain)
+        opt, _ = parser.parse_known_args()  # parse again with new defaults
+
+        # modify dataset-related parser options
+        dataset_name = opt.dataset_mode
+        dataset_option_setter = data.get_option_setter(dataset_name)
+        parser = dataset_option_setter(parser, self.isTrain)
+
+        # save and return the parser
+        self.parser = parser
+        return parser.parse_args()
+
+    def print_options(self, opt):
+        """Print and save options
+
+        It will print both current options and default values(if different).
+        It will save options into a text file / [checkpoints_dir] / opt.txt
+        """
+        message = ''
+        message += '----------------- Options ---------------\n'
+        for k, v in sorted(vars(opt).items()):
+            comment = ''
+            default = self.parser.get_default(k)
+            if v != default:
+                comment = '\t[default: %s]' % str(default)
+            message += '{:>25}: {:<30}{}\n'.format(str(k), str(v), comment)
+        message += '----------------- End -------------------'
+        print(message)
+
+        # save to the disk
+        expr_dir = os.path.join(opt.checkpoints_dir, opt.name)
+        util.mkdirs(expr_dir)
+        file_name = os.path.join(expr_dir, '{}_opt.txt'.format(opt.phase))
+        with open(file_name, 'wt') as opt_file:
+            opt_file.write(message)
+            opt_file.write('\n')
+
+    def parse(self):
+        """Parse our options, create checkpoints directory suffix, and set up gpu device."""
+        opt = self.gather_options()
+        opt.isTrain = self.isTrain   # train or test
+
+        # process opt.suffix
+        if opt.suffix:
+            suffix = ('_' + opt.suffix.format(**vars(opt))) if opt.suffix != '' else ''
+            opt.name = opt.name + suffix
+
+
+        # set gpu ids
+        str_ids = opt.gpu_ids.split(',')
+        opt.gpu_ids = []
+        for str_id in str_ids:
+            id = int(str_id)
+            if id >= 0:
+                opt.gpu_ids.append(id)
+        if len(opt.gpu_ids) > 0:
+            torch.cuda.set_device(opt.gpu_ids[0])
+
+        self.opt = opt
+        return self.opt

cyclegan_model/options/test_options.py

@@ -0,0 +1,24 @@
+from .base_options import BaseOptions
+
+
+class TestOptions(BaseOptions):
+    """This class includes test options.
+
+    It also includes shared options defined in BaseOptions.
+    """
+
+    def initialize(self, parser):
+        parser = BaseOptions.initialize(self, parser)  # define shared options
+        parser.add_argument('--results_dir', type=str, default='./results/', help='saves results here.')
+        parser.add_argument('--aspect_ratio', type=float, default=1.0, help='aspect ratio of result images')
+        parser.add_argument('--phase', type=str, default='test', help='train, val, test, etc')
+        # Dropout and Batchnorm has different behavioir during training and test.
+        parser.add_argument('--eval', action='store_true', help='use eval mode during test time.')
+        parser.add_argument('--num_test', type=int, default=100, help='how many test images to run')
+
+        # rewrite devalue values
+        # parser.set_defaults(model='test')
+        # To avoid cropping, the load_size should be the same as crop_size
+        parser.set_defaults(load_size=parser.get_default('crop_size'))
+        self.isTrain = False
+        return parser

cyclegan_model/util/image_pool.py

@@ -0,0 +1,54 @@
+import random
+import torch
+
+
+class ImagePool():
+    """This class implements an image buffer that stores previously generated images.
+
+    This buffer enables us to update discriminators using a history of generated images
+    rather than the ones produced by the latest generators.
+    """
+
+    def __init__(self, pool_size):
+        """Initialize the ImagePool class
+
+        Parameters:
+            pool_size (int) -- the size of image buffer, if pool_size=0, no buffer will be created
+        """
+        self.pool_size = pool_size
+        if self.pool_size > 0:  # create an empty pool
+            self.num_imgs = 0
+            self.images = []
+
+    def query(self, images):
+        """Return an image from the pool.
+
+        Parameters:
+            images: the latest generated images from the generator
+
+        Returns images from the buffer.
+
+        By 50/100, the buffer will return input images.
+        By 50/100, the buffer will return images previously stored in the buffer,
+        and insert the current images to the buffer.
+        """
+        if self.pool_size == 0:  # if the buffer size is 0, do nothing
+            return images
+        return_images = []
+        for image in images:
+            image = torch.unsqueeze(image.data, 0)
+            if self.num_imgs < self.pool_size:   # if the buffer is not full; keep inserting current images to the buffer
+                self.num_imgs = self.num_imgs + 1
+                self.images.append(image)
+                return_images.append(image)
+            else:
+                p = random.uniform(0, 1)
+                if p > 0.5:  # by 50% chance, the buffer will return a previously stored image, and insert the current image into the buffer
+                    random_id = random.randint(0, self.pool_size - 1)  # randint is inclusive
+                    tmp = self.images[random_id].clone()
+                    self.images[random_id] = image
+                    return_images.append(tmp)
+                else:       # by another 50% chance, the buffer will return the current image
+                    return_images.append(image)
+        return_images = torch.cat(return_images, 0)   # collect all the images and return
+        return return_images

cyclegan_model/util/util.py

@@ -0,0 +1,103 @@
+"""This module contains simple helper functions """
+from __future__ import print_function
+import torch
+import numpy as np
+from PIL import Image
+import os
+
+
+def tensor2im(input_image, imtype=np.uint8):
+    """"Converts a Tensor array into a numpy image array.
+
+    Parameters:
+        input_image (tensor) --  the input image tensor array
+        imtype (type)        --  the desired type of the converted numpy array
+    """
+    if not isinstance(input_image, np.ndarray):
+        if isinstance(input_image, torch.Tensor):  # get the data from a variable
+            image_tensor = input_image.data
+        else:
+            return input_image
+        image_numpy = image_tensor[0].cpu().float().numpy()  # convert it into a numpy array
+        if image_numpy.shape[0] == 1:  # grayscale to RGB
+            image_numpy = np.tile(image_numpy, (3, 1, 1))
+        image_numpy = (np.transpose(image_numpy, (1, 2, 0)) + 1) / 2.0 * 255.0  # post-processing: tranpose and scaling
+    else:  # if it is a numpy array, do nothing
+        image_numpy = input_image
+    return image_numpy.astype(imtype)
+
+
+def diagnose_network(net, name='network'):
+    """Calculate and print the mean of average absolute(gradients)
+
+    Parameters:
+        net (torch network) -- Torch network
+        name (str) -- the name of the network
+    """
+    mean = 0.0
+    count = 0
+    for param in net.parameters():
+        if param.grad is not None:
+            mean += torch.mean(torch.abs(param.grad.data))
+            count += 1
+    if count > 0:
+        mean = mean / count
+    print(name)
+    print(mean)
+
+
+def save_image(image_numpy, image_path, aspect_ratio=1.0):
+    """Save a numpy image to the disk
+
+    Parameters:
+        image_numpy (numpy array) -- input numpy array
+        image_path (str)          -- the path of the image
+    """
+
+    image_pil = Image.fromarray(image_numpy)
+    h, w, _ = image_numpy.shape
+
+    if aspect_ratio > 1.0:
+        image_pil = image_pil.resize((h, int(w * aspect_ratio)), Image.BICUBIC)
+    if aspect_ratio < 1.0:
+        image_pil = image_pil.resize((int(h / aspect_ratio), w), Image.BICUBIC)
+    image_pil.save(image_path)
+
+
+def print_numpy(x, val=True, shp=False):
+    """Print the mean, min, max, median, std, and size of a numpy array
+
+    Parameters:
+        val (bool) -- if print the values of the numpy array
+        shp (bool) -- if print the shape of the numpy array
+    """
+    x = x.astype(np.float64)
+    if shp:
+        print('shape,', x.shape)
+    if val:
+        x = x.flatten()
+        print('mean = %3.3f, min = %3.3f, max = %3.3f, median = %3.3f, std=%3.3f' % (
+            np.mean(x), np.min(x), np.max(x), np.median(x), np.std(x)))
+
+
+def mkdirs(paths):
+    """create empty directories if they don't exist
+
+    Parameters:
+        paths (str list) -- a list of directory paths
+    """
+    if isinstance(paths, list) and not isinstance(paths, str):
+        for path in paths:
+            mkdir(path)
+    else:
+        mkdir(paths)
+
+
+def mkdir(path):
+    """create a single empty directory if it didn't exist
+
+    Parameters:
+        path (str) -- a single directory path
+    """
+    if not os.path.exists(path):
+        os.makedirs(path)

flux_ad/main.ipynb

@@ -0,0 +1,96 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Style-specific T2I Generation with Flux.1-dev"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import os\n",
+    "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"0\"\n",
+    "from accelerate.utils import set_seed\n",
+    "from mypipeline import AttDistPipeline\n",
+    "from utils import *\n",
+    "\n",
+    "\n",
+    "model_name = \"/root/models/FLUX.1-dev\"\n",
+    "lr = 0.01\n",
+    "iters = 2\n",
+    "seed = 42\n",
+    "width = 512\n",
+    "height = 512\n",
+    "mixed_precision = \"bf16\"\n",
+    "num_inference_steps = 50\n",
+    "guidance_scale = 3.5\n",
+    "enable_gradient_checkpoint = True\n",
+    "start_layer, end_layer = 50, 57\n",
+    "start_time = 9999\n",
+    "prompt=\"A panda\"\n",
+    "\n",
+    "\n",
+    "pipe = AttDistPipeline.from_pretrained(\n",
+    "    model_name, torch_dtype=torch.float16)\n",
+    "\n",
+    "\n",
+    "memory_efficient(pipe)\n",
+    "set_seed(seed)\n",
+    "loss_fn = torch.nn.L1Loss()\n",
+    "\n",
+    "style_image = [\"..//data/style/1.jpg\"]\n",
+    "style_image = torch.cat([load_image(path, size=(512, 512)) for path in style_image])\n",
+    "\n",
+    "\n",
+    "controller = Controller(self_layers=(start_layer, end_layer))\n",
+    "\n",
+    "result = pipe.sample(\n",
+    "    lr=lr,\n",
+    "    prompt=prompt,\n",
+    "    loss_fn=loss_fn,\n",
+    "    iters=iters,\n",
+    "    width=width,\n",
+    "    height=height,\n",
+    "    start_time=start_time,\n",
+    "    controller=controller,\n",
+    "    style_image=style_image,\n",
+    "    guidance_scale=guidance_scale,\n",
+    "    mixed_precision=mixed_precision,\n",
+    "    num_inference_steps=num_inference_steps,\n",
+    "    enable_gradient_checkpoint=enable_gradient_checkpoint,\n",
+    ")\n",
+    "\n",
+    "save_image(style_image, \"style.png\")\n",
+    "save_image(result, \"output.png\")\n",
+    "show_image(\"style.png\", title=\"style image\")\n",
+    "show_image(\"output.png\", title=prompt)"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "ad",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.10.16"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}

flux_ad/mypipeline.py

@@ -0,0 +1,381 @@
+import inspect
+from accelerate import Accelerator
+from typing import Any, Callable, Dict, List, Optional, Union
+from tqdm import tqdm
+import numpy as np
+import torch
+import torch.nn.functional as F
+import utils
+from diffusers import FluxPipeline
+
+
+# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents
+def retrieve_latents(
+    encoder_output: torch.Tensor,
+    generator: Optional[torch.Generator] = None,
+    sample_mode: str = "sample",
+):
+    if hasattr(encoder_output, "latent_dist") and sample_mode == "sample":
+        return encoder_output.latent_dist.sample(generator)
+    elif hasattr(encoder_output, "latent_dist") and sample_mode == "argmax":
+        return encoder_output.latent_dist.mode()
+    elif hasattr(encoder_output, "latents"):
+        return encoder_output.latents
+    else:
+        raise AttributeError("Could not access latents of provided encoder_output")
+
+
+class AttDistPipeline(FluxPipeline):
+
+    def freeze(self):
+        self.transformer.requires_grad_(False)
+        self.text_encoder.requires_grad_(False)
+        self.text_encoder_2.requires_grad_(False)
+        self.vae.requires_grad_(False)
+
+    @torch.no_grad()
+    def image2latent(self, image):
+        dtype = next(self.vae.parameters()).dtype
+        device = self._execution_device
+        image = image.to(device=device, dtype=dtype) * 2.0 - 1.0
+        latent = retrieve_latents(self.vae.encode(image))
+        latent = (
+            latent - self.vae.config.shift_factor
+        ) * self.vae.config.scaling_factor
+        return latent
+
+    @torch.no_grad()
+    def latent2image(self, latent, height, width):
+        dtype = next(self.vae.parameters()).dtype
+        device = self._execution_device
+        latent = latent.to(device=device, dtype=dtype)
+        latents = self._unpack_latents(latent, height, width, self.vae_scale_factor)
+        latents = (
+            latents / self.vae.config.scaling_factor
+        ) + self.vae.config.shift_factor
+        image = self.vae.decode(latents, return_dict=False)[0]
+        return (image * 0.5 + 0.5).clamp(0, 1)
+
+    def sample(
+        self,
+        style_image=None,
+        controller=None,
+        loss_fn=None,
+        start_time=9999,
+        lr=0.05,
+        iters=2,
+        mixed_precision="no",
+        enable_gradient_checkpoint=False,
+        prompt: Union[str, List[str]] = None,
+        prompt_2: Optional[Union[str, List[str]]] = None,
+        height: Optional[int] = None,
+        width: Optional[int] = None,
+        num_inference_steps: int = 28,
+        # timesteps: List[int] = None,
+        guidance_scale: float = 3.5,
+        num_images_per_prompt: Optional[int] = 1,
+        generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
+        latents: Optional[torch.FloatTensor] = None,
+        prompt_embeds: Optional[torch.FloatTensor] = None,
+        pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
+        output_type: Optional[str] = "pil",
+        return_dict: bool = True,
+        joint_attention_kwargs: Optional[Dict[str, Any]] = None,
+        callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
+        callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+        max_sequence_length: int = 512,
+    ):
+        height = height or self.default_sample_size * self.vae_scale_factor
+        width = width or self.default_sample_size * self.vae_scale_factor
+        device = self._execution_device
+        accelerator = Accelerator(
+            mixed_precision=mixed_precision, gradient_accumulation_steps=1
+        )
+        weight_dtype = torch.float32
+        if accelerator.mixed_precision == "fp16":
+            weight_dtype = torch.float16
+        elif accelerator.mixed_precision == "bf16":
+            weight_dtype = torch.bfloat16
+
+        # Move unet, vae and text_encoder to device and cast to weight_dtype
+        self.transformer.to(accelerator.device, dtype=weight_dtype)
+        self.vae.to(accelerator.device, dtype=weight_dtype)
+        self.text_encoder.to(accelerator.device, dtype=weight_dtype)
+        self.text_encoder_2.to(accelerator.device, dtype=weight_dtype)
+        self.transformer = accelerator.prepare(self.transformer)
+        if enable_gradient_checkpoint:
+            self.transformer.enable_gradient_checkpointing()
+            # self.transformer.train()
+
+        (null_embeds, null_pooled_embeds, null_text_ids) = self.encode_prompt(
+            prompt="",
+            prompt_2=prompt_2,
+        )
+        (
+            prompt_embeds,
+            pooled_prompt_embeds,
+            text_ids,
+        ) = self.encode_prompt(
+            prompt=prompt,
+            prompt_2=prompt_2,
+            prompt_embeds=prompt_embeds,
+            pooled_prompt_embeds=pooled_prompt_embeds,
+            device=device,
+            num_images_per_prompt=num_images_per_prompt,
+            max_sequence_length=max_sequence_length,
+        )
+        # 4. Prepare latent variables
+        num_channels_latents = self.transformer.config.in_channels // 4
+        latents, latent_image_ids = self.prepare_latents(
+            num_images_per_prompt,
+            num_channels_latents,
+            height,
+            width,
+            null_embeds.dtype,
+            device,
+            generator,
+            latents,
+        )
+
+        print(style_image.shape)
+        height_, width_ = style_image.shape[2], style_image.shape[3]
+        style_latent = self.image2latent(style_image)
+        print(style_latent.shape)
+        print(latents.shape)
+        style_latent = self._pack_latents(style_latent, 1, num_channels_latents, style_latent.shape[2], style_latent.shape[3])
+
+        _, null_image_id = self.prepare_latents(
+            num_images_per_prompt,
+            num_channels_latents,
+            height_,
+            width_,
+            null_embeds.dtype,
+            device,
+            generator,
+            style_latent,
+        )
+
+        # 5. Prepare timesteps
+        sigmas = np.linspace(1.0, 1 / num_inference_steps, num_inference_steps)
+        image_seq_len = latents.shape[1]
+        mu = calculate_shift(
+            image_seq_len,
+            self.scheduler.config.base_image_seq_len,
+            self.scheduler.config.max_image_seq_len,
+            self.scheduler.config.base_shift,
+            self.scheduler.config.max_shift,
+        )
+        timesteps, num_inference_steps = retrieve_timesteps(
+            self.scheduler,
+            num_inference_steps,
+            device,
+            None,
+            sigmas,
+            mu=mu,
+        )
+
+        timesteps = self.scheduler.timesteps
+        print(f"timesteps: {timesteps}")
+        self._num_timesteps = len(timesteps)
+
+        cache = utils.DataCache()
+        
+        utils.register_attn_control(
+            self.transformer.transformer_blocks,
+            controller=controller,
+            cache=cache,
+        )
+        utils.register_attn_control(
+            self.transformer.single_transformer_blocks,
+            controller=controller,
+            cache=cache,
+        )
+        # handle guidance
+        if self.transformer.config.guidance_embeds:
+            guidance = torch.full(
+                [1], guidance_scale, device=device, dtype=torch.float32
+            )
+            guidance = guidance.expand(latents.shape[0])
+        else:
+            guidance = None
+
+        null_guidance = torch.full(
+                [1], 1, device=device, dtype=torch.float32
+            )
+        
+        print(controller.num_self_layers)
+
+
+        pbar = tqdm(timesteps, desc="Sample")
+        for i, t in enumerate(pbar):
+            timestep = t.expand(latents.shape[0]).to(latents.dtype)
+            with torch.no_grad():
+                noise_pred = self.transformer(
+                    hidden_states=latents,
+                    timestep=timestep / 1000,
+                    guidance=guidance,
+                    pooled_projections=pooled_prompt_embeds,
+                    encoder_hidden_states=prompt_embeds,
+                    txt_ids=text_ids,
+                    img_ids=latent_image_ids,
+                    joint_attention_kwargs=None,
+                    return_dict=False,
+                )[0]
+
+            # compute the previous noisy sample x_t -> x_t-1
+            latents = self.scheduler.step(
+                noise_pred, t, latents, return_dict=False
+            )[0]
+            if t < start_time:
+                if i < num_inference_steps - 1:
+                    timestep = timesteps[i+1:i+2]
+                    # print(timestep)
+                    noise = torch.randn_like(style_latent)
+                    # print(style_latent.shape)
+                    style_latent_ = self.scheduler.scale_noise(style_latent, timestep, noise)
+                else:
+                    timestep = torch.tensor([0], device=style_latent.device)
+                    style_latent_ = style_latent
+
+                cache.clear()
+                controller.step()
+                
+                _ = self.transformer(
+                    hidden_states=style_latent_,
+                    timestep=timestep / 1000,
+                    guidance=null_guidance,
+                    pooled_projections=null_pooled_embeds,
+                    encoder_hidden_states=null_embeds,
+                    txt_ids=null_text_ids,
+                    img_ids=null_image_id,
+                    joint_attention_kwargs=None,
+                    return_dict=False,
+                )[0]
+                _, ref_k_list, ref_v_list, _ = cache.get()
+
+
+                latents = utils.adain(latents, style_latent_)
+                latents = latents.detach()
+                optimizer = torch.optim.Adam([latents.requires_grad_()], lr=lr)
+                optimizer = accelerator.prepare(optimizer)
+
+                for _ in range(iters):
+                    cache.clear()
+                    controller.step()
+                    optimizer.zero_grad()
+                    _ = self.transformer(
+                        hidden_states=latents,
+                        timestep=timestep / 1000,
+                        guidance=null_guidance,
+                        pooled_projections=null_pooled_embeds,
+                        encoder_hidden_states=null_embeds,
+                        txt_ids=null_text_ids,
+                        img_ids=latent_image_ids,
+                        joint_attention_kwargs=None,
+                        return_dict=False,
+                    )[0]
+                    q_list, _, _, self_out_list = cache.get()
+                    ref_self_out_list = [
+                        F.scaled_dot_product_attention(
+                            q,
+                            ref_k,
+                            ref_v,
+                        )
+                        for q, ref_k, ref_v in zip(q_list, ref_k_list, ref_v_list)
+                    ]
+                    style_loss = sum(
+                        [
+                            loss_fn(self_out, ref_self_out.detach())
+                            for self_out, ref_self_out in zip(
+                                self_out_list, ref_self_out_list
+                            )
+                        ]
+                    )
+                    loss = style_loss
+                    accelerator.backward(loss)
+                    # loss.backward()
+                    optimizer.step()
+
+                    pbar.set_postfix(loss=loss.item(), time=t.item())
+        torch.cuda.empty_cache()
+        latents = latents.detach()
+        return self.latent2image(latents, height, width)
+
+
+def calculate_shift(
+    image_seq_len,
+    base_seq_len: int = 256,
+    max_seq_len: int = 4096,
+    base_shift: float = 0.5,
+    max_shift: float = 1.16,
+):
+    m = (max_shift - base_shift) / (max_seq_len - base_seq_len)
+    b = base_shift - m * base_seq_len
+    mu = image_seq_len * m + b
+    return mu
+
+
+# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps
+def retrieve_timesteps(
+    scheduler,
+    num_inference_steps: Optional[int] = None,
+    device: Optional[Union[str, torch.device]] = None,
+    timesteps: Optional[List[int]] = None,
+    sigmas: Optional[List[float]] = None,
+    **kwargs,
+):
+    r"""
+    Calls the scheduler's `set_timesteps` method and retrieves timesteps from the scheduler after the call. Handles
+    custom timesteps. Any kwargs will be supplied to `scheduler.set_timesteps`.
+
+    Args:
+        scheduler (`SchedulerMixin`):
+            The scheduler to get timesteps from.
+        num_inference_steps (`int`):
+            The number of diffusion steps used when generating samples with a pre-trained model. If used, `timesteps`
+            must be `None`.
+        device (`str` or `torch.device`, *optional*):
+            The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
+        timesteps (`List[int]`, *optional*):
+            Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed,
+            `num_inference_steps` and `sigmas` must be `None`.
+        sigmas (`List[float]`, *optional*):
+            Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed,
+            `num_inference_steps` and `timesteps` must be `None`.
+
+    Returns:
+        `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
+        second element is the number of inference steps.
+    """
+    if timesteps is not None and sigmas is not None:
+        raise ValueError(
+            "Only one of `timesteps` or `sigmas` can be passed. Please choose one to set custom values"
+        )
+    if timesteps is not None:
+        accepts_timesteps = "timesteps" in set(
+            inspect.signature(scheduler.set_timesteps).parameters.keys()
+        )
+        if not accepts_timesteps:
+            raise ValueError(
+                f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"
+                f" timestep schedules. Please check whether you are using the correct scheduler."
+            )
+        scheduler.set_timesteps(timesteps=timesteps, device=device, **kwargs)
+        timesteps = scheduler.timesteps
+        num_inference_steps = len(timesteps)
+    elif sigmas is not None:
+        accept_sigmas = "sigmas" in set(
+            inspect.signature(scheduler.set_timesteps).parameters.keys()
+        )
+        if not accept_sigmas:
+            raise ValueError(
+                f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"
+                f" sigmas schedules. Please check whether you are using the correct scheduler."
+            )
+        scheduler.set_timesteps(sigmas=sigmas, device=device, **kwargs)
+        timesteps = scheduler.timesteps
+        num_inference_steps = len(timesteps)
+    else:
+        scheduler.set_timesteps(num_inference_steps, device=device, **kwargs)
+        timesteps = scheduler.timesteps
+    return timesteps, num_inference_steps

flux_ad/utils.py

@@ -0,0 +1,232 @@
+import numpy as np
+import torch
+import torch.nn.functional as F
+from PIL import Image
+from torchvision.transforms import ToTensor
+from torchvision.utils import save_image
+import matplotlib.pyplot as plt
+
+
+def register_attn_control(unet, controller, cache=None):
+    def attn_forward(self):
+
+        def forward(
+            hidden_states,
+            encoder_hidden_states=None,
+            attention_mask=None,
+            image_rotary_emb=None,
+            *args,
+            **kwargs,
+        ):
+            batch_size, _, _ = (
+                hidden_states.shape
+                if encoder_hidden_states is None
+                else encoder_hidden_states.shape
+            )
+
+            # `sample` projections.
+            query = self.to_q(hidden_states)
+            key = self.to_k(hidden_states)
+            value = self.to_v(hidden_states)
+
+            inner_dim = key.shape[-1]
+            head_dim = inner_dim // self.heads
+
+            query = query.view(batch_size, -1, self.heads, head_dim).transpose(1, 2)
+            key = key.view(batch_size, -1, self.heads, head_dim).transpose(1, 2)
+            value = value.view(batch_size, -1, self.heads, head_dim).transpose(1, 2)
+
+            if self.norm_q is not None:
+                query = self.norm_q(query)
+            if self.norm_k is not None:
+                key = self.norm_k(key)
+
+            # the attention in FluxSingleTransformerBlock does not use `encoder_hidden_states`
+            if encoder_hidden_states is not None:
+                # `context` projections.
+                encoder_hidden_states_query_proj = self.add_q_proj(
+                    encoder_hidden_states
+                )
+                encoder_hidden_states_key_proj = self.add_k_proj(encoder_hidden_states)
+                encoder_hidden_states_value_proj = self.add_v_proj(
+                    encoder_hidden_states
+                )
+
+                encoder_hidden_states_query_proj = (
+                    encoder_hidden_states_query_proj.view(
+                        batch_size, -1, self.heads, head_dim
+                    ).transpose(1, 2)
+                )
+                encoder_hidden_states_key_proj = encoder_hidden_states_key_proj.view(
+                    batch_size, -1, self.heads, head_dim
+                ).transpose(1, 2)
+                encoder_hidden_states_value_proj = (
+                    encoder_hidden_states_value_proj.view(
+                        batch_size, -1, self.heads, head_dim
+                    ).transpose(1, 2)
+                )
+
+                if self.norm_added_q is not None:
+                    encoder_hidden_states_query_proj = self.norm_added_q(
+                        encoder_hidden_states_query_proj
+                    )
+                if self.norm_added_k is not None:
+                    encoder_hidden_states_key_proj = self.norm_added_k(
+                        encoder_hidden_states_key_proj
+                    )
+
+                # attention
+                query = torch.cat([encoder_hidden_states_query_proj, query], dim=2)
+                key = torch.cat([encoder_hidden_states_key_proj, key], dim=2)
+                value = torch.cat([encoder_hidden_states_value_proj, value], dim=2)
+
+            if image_rotary_emb is not None:
+                from diffusers.models.embeddings import apply_rotary_emb
+
+                query = apply_rotary_emb(query, image_rotary_emb)
+                key = apply_rotary_emb(key, image_rotary_emb)
+
+            hidden_states = F.scaled_dot_product_attention(
+                query, key, value, dropout_p=0.0, is_causal=False
+            )
+            if controller.cur_self_layer in controller.self_layers:
+                # print("cache added")
+                cache.add(query, key, value, hidden_states)
+                # if encoder_hidden_states is None:
+            controller.cur_self_layer += 1
+
+            hidden_states = hidden_states.transpose(1, 2).reshape(
+                batch_size, -1, self.heads * head_dim
+            )
+
+            hidden_states = hidden_states.to(query.dtype)
+
+            if encoder_hidden_states is not None:
+                encoder_hidden_states, hidden_states = (
+                    hidden_states[:, : encoder_hidden_states.shape[1]],
+                    hidden_states[:, encoder_hidden_states.shape[1] :],
+                )
+
+                # linear proj
+                hidden_states = self.to_out[0](hidden_states)
+                # dropout
+                hidden_states = self.to_out[1](hidden_states)
+                encoder_hidden_states = self.to_add_out(encoder_hidden_states)
+
+                return hidden_states, encoder_hidden_states
+            else:
+                return hidden_states
+
+        return forward
+
+    def modify_forward(net, count):
+        # print(net.named_children())
+        for name, subnet in net.named_children():
+            if net.__class__.__name__ == "Attention":  # spatial Transformer layer
+                net.forward = attn_forward(net)
+                return count + 1
+            elif hasattr(net, "children"):
+                count = modify_forward(subnet, count)
+        return count
+
+    cross_att_count = 0
+    cross_att_count += modify_forward(unet, 0)
+    controller.num_self_layers += cross_att_count
+
+
+def load_image(image_path, size=None, mode="RGB"):
+    img = Image.open(image_path).convert(mode)
+    if size is None:
+        width, height = img.size
+        new_width = (width // 64) * 64
+        new_height = (height // 64) * 64
+        size = (new_width, new_height)
+    img = img.resize(size, Image.BICUBIC)
+    return ToTensor()(img).unsqueeze(0)
+
+
+def adain(source, target, eps=1e-6):
+    source_mean, source_std = torch.mean(source, dim=1, keepdim=True), torch.std(
+        source, dim=1, keepdim=True
+    )
+    target_mean, target_std = torch.mean(
+        target, dim=(0, 1), keepdim=True
+    ), torch.std(target, dim=(0, 1), keepdim=True)
+    normalized_source = (source - source_mean) / (source_std + eps)
+    transferred_source = normalized_source * target_std + target_mean
+
+    return transferred_source
+
+
+def shuffle_tensor(tensor):
+    B, C, H, W = tensor.shape
+    flat_tensor = tensor.reshape(B, C, -1)
+    shuffled_tensor = flat_tensor[:, :, torch.randperm(H * W)].reshape(B, C, H, W)
+    return shuffled_tensor
+
+
+class Controller:
+    def step(self):
+        self.cur_self_layer = 0
+
+    def __init__(self, self_layers=(0, 16)):
+        self.num_self_layers = 0
+        self.cur_self_layer = 0
+        self.self_layers = list(range(*self_layers))
+
+
+class DataCache:
+    def __init__(self):
+        self.q = []
+        self.k = []
+        self.v = []
+        self.out = []
+
+    def clear(self):
+        self.q.clear()
+        self.k.clear()
+        self.v.clear()
+        self.out.clear()
+
+    def add(self, q, k, v, out):
+        self.q.append(q)
+        self.k.append(k)
+        self.v.append(v)
+        self.out.append(out)
+
+    def get(self):
+        return self.q.copy(), self.k.copy(), self.v.copy(), self.out.copy()
+
+
+def memory_efficient(model):
+    model.freeze()
+    try:
+        model.enable_model_cpu_offload()
+    except AttributeError:
+        print("enable_model_cpu_offload is not supported.")
+    try:
+        model.enable_vae_slicing()
+    except AttributeError:
+        print("enable_vae_slicing is not supported.")
+
+    try:
+        model.enable_vae_tiling()
+    except AttributeError:
+        print("enable_vae_tiling is not supported.")
+
+def show_image(path, title, display_height=3, title_fontsize=12):
+    img = Image.open(path)
+    img_width, img_height = img.size
+
+    aspect_ratio = img_width / img_height
+    display_width = display_height * aspect_ratio
+
+    plt.figure(figsize=(display_width, display_height))
+    plt.imshow(img)
+    plt.title(title, 
+             fontsize=title_fontsize, 
+             fontweight='bold', 
+             pad=20) 
+    plt.axis('off')    
+    plt.tight_layout() 
+    plt.show()

inpaint_model/model/networks.py

@@ -0,0 +1,562 @@
+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 inpaint_model.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)

matting/image_matting.py

@@ -0,0 +1,274 @@
+import os
+import cv2
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.autograd import Variable
+from torchvision import transforms
+from PIL import Image
+
+
+# 定义U²-Net模型
+class REBNCONV(nn.Module):
+    def __init__(self, in_ch=3, out_ch=3, dirate=1):
+        super(REBNCONV, self).__init__()
+
+        self.conv_s1 = nn.Conv2d(in_ch, out_ch, 3, padding=1 * dirate, dilation=1 * dirate)
+        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
+
+
+class RSU7(nn.Module):
+    def __init__(self, in_ch=3, mid_ch=12, out_ch=3):
+        super(RSU7, 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.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):
+        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 = F.interpolate(hx6d, scale_factor=2, mode='bilinear', align_corners=False)
+
+        hx5d = self.rebnconv5d(torch.cat((hx6dup, hx5), 1))
+        hx5dup = F.interpolate(hx5d, scale_factor=2, mode='bilinear', align_corners=False)
+
+        hx4d = self.rebnconv4d(torch.cat((hx5dup, hx4), 1))
+        hx4dup = F.interpolate(hx4d, scale_factor=2, mode='bilinear', align_corners=False)
+
+        hx3d = self.rebnconv3d(torch.cat((hx4dup, hx3), 1))
+        hx3dup = F.interpolate(hx3d, scale_factor=2, mode='bilinear', align_corners=False)
+
+        hx2d = self.rebnconv2d(torch.cat((hx3dup, hx2), 1))
+        hx2dup = F.interpolate(hx2d, scale_factor=2, mode='bilinear', align_corners=False)
+
+        hx1d = self.rebnconv1d(torch.cat((hx2dup, hx1), 1))
+
+        return hx1d + hxin
+
+
+class U2NET(nn.Module):
+    def __init__(self, in_ch=3, out_ch=1):
+        super(U2NET, self).__init__()
+
+        self.stage1 = RSU7(in_ch, 32, 64)
+        self.pool12 = nn.MaxPool2d(2, stride=2, ceil_mode=True)
+
+        self.stage2 = RSU7(64, 32, 128)
+        self.pool23 = nn.MaxPool2d(2, stride=2, ceil_mode=True)
+
+        self.stage3 = RSU7(128, 64, 256)
+        self.pool34 = nn.MaxPool2d(2, stride=2, ceil_mode=True)
+
+        self.stage4 = RSU7(256, 128, 512)
+        self.pool45 = nn.MaxPool2d(2, stride=2, ceil_mode=True)
+
+        self.stage5 = RSU7(512, 256, 512)
+        self.pool56 = nn.MaxPool2d(2, stride=2, ceil_mode=True)
+
+        self.stage6 = RSU7(512, 512, 512)
+
+        self.stage5d = RSU7(1024, 256, 512)
+        self.stage4d = RSU7(1024, 128, 256)
+        self.stage3d = RSU7(512, 64, 128)
+        self.stage2d = RSU7(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, 1)
+
+    def forward(self, x):
+        hx = x
+
+        hx1 = self.stage1(hx)
+        hx = self.pool12(hx1)
+
+        hx2 = self.stage2(hx)
+        hx = self.pool23(hx2)
+
+        hx3 = self.stage3(hx)
+        hx = self.pool34(hx3)
+
+        hx4 = self.stage4(hx)
+        hx = self.pool45(hx4)
+
+        hx5 = self.stage5(hx)
+        hx = self.pool56(hx5)
+
+        hx6 = self.stage6(hx)
+        hx6up = F.interpolate(hx6, scale_factor=2, mode='bilinear', align_corners=False)
+
+        hx5d = self.stage5d(torch.cat((hx6up, hx5), 1))
+        hx5dup = F.interpolate(hx5d, scale_factor=2, mode='bilinear', align_corners=False)
+
+        hx4d = self.stage4d(torch.cat((hx5dup, hx4), 1))
+        hx4dup = F.interpolate(hx4d, scale_factor=2, mode='bilinear', align_corners=False)
+
+        hx3d = self.stage3d(torch.cat((hx4dup, hx3), 1))
+        hx3dup = F.interpolate(hx3d, scale_factor=2, mode='bilinear', align_corners=False)
+
+        hx2d = self.stage2d(torch.cat((hx3dup, hx2), 1))
+        hx2dup = F.interpolate(hx2d, scale_factor=2, mode='bilinear', align_corners=False)
+
+        hx1d = self.stage1d(torch.cat((hx2dup, hx1), 1))
+
+        d1 = self.side1(hx1d)
+
+        d2 = self.side2(hx2d)
+        d2 = F.interpolate(d2, scale_factor=2, mode='bilinear', align_corners=False)
+
+        d3 = self.side3(hx3d)
+        d3 = F.interpolate(d3, scale_factor=4, mode='bilinear', align_corners=False)
+
+        d4 = self.side4(hx4d)
+        d4 = F.interpolate(d4, scale_factor=8, mode='bilinear', align_corners=False)
+
+        d5 = self.side5(hx5d)
+        d5 = F.interpolate(d5, scale_factor=16, mode='bilinear', align_corners=False)
+
+        d6 = self.side6(hx6)
+        d6 = F.interpolate(d6, scale_factor=32, mode='bilinear', align_corners=False)
+
+        d0 = self.outconv(torch.cat((d1, d2, d3, d4, d5, d6), 1))
+
+        return torch.sigmoid(d0), torch.sigmoid(d1), torch.sigmoid(d2), torch.sigmoid(d3), torch.sigmoid(
+            d4), torch.sigmoid(d5), torch.sigmoid(d6)
+
+
+# 预处理图像函数
+def preprocess_image(image_path):
+    img = cv2.imread(image_path)
+    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
+    img = cv2.resize(img, (320, 320))
+    img = img.astype(np.float32) / 255.0
+    img = transforms.ToTensor()(img)
+    img = transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])(img)
+    return img.unsqueeze(0)
+
+
+# 后处理mask函数
+def postprocess_mask(mask):
+    mask = mask.squeeze()
+    mask = (mask > 0.5).float()
+    mask = mask.cpu().numpy()
+    return mask
+
+
+# 创建二值图片函数
+def create_binary_image(image_path, mask, output_path):
+    original_image = cv2.imread(image_path)
+    original_image = cv2.cvtColor(original_image, cv2.COLOR_BGR2RGB)
+
+    mask_resized = cv2.resize(mask, (original_image.shape[1], original_image.shape[0]))
+
+    result = np.zeros_like(original_image)
+    for c in range(3):
+        result[:, :, c] = np.where(mask_resized == 1, original_image[:, :, c], 0)
+
+    cv2.imwrite(output_path, cv2.cvtColor(result, cv2.COLOR_RGB2BGR))
+
+
+# 加载模型函数
+def load_model(model_path):
+    net = U2NET(3, 1)
+    if torch.cuda.is_available():
+        net.load_state_dict(torch.load(model_path))
+        net.cuda()
+    else:
+        net.load_state_dict(torch.load(model_path, map_location='cpu'))
+    net.eval()
+    return net
+
+
+# 处理图像函数
+def process_images(input_folder, output_folder, model_path):
+    if not os.path.exists(output_folder):
+        os.makedirs(output_folder)
+
+    model = load_model(model_path)
+
+    for filename in os.listdir(input_folder):
+        if filename.endswith(('.png', '.jpg', '.jpeg')):
+            image_path = os.path.join(input_folder, filename)
+            output_path = os.path.join(output_folder, filename)
+
+            img = preprocess_image(image_path)
+
+            if torch.cuda.is_available():
+                inputs = Variable(img.cuda())
+            else:
+                inputs = Variable(img)
+
+            d1, d2, d3, d4, d5, d6, d7 = model(inputs)
+
+            pred_mask = d1
+            mask = postprocess_mask(pred_mask)
+
+            create_binary_image(image_path, mask, output_path)
+
+            print(f"Processed {filename}")
+
+
+if __name__ == "__main__":
+    model_path = "./u2net.pth"  # 替换为正确的模型权重文件路径
+    input_folder = r"D:\learn_torch\pytorch-CycleGAN\datasets\ostracoda\trainA"  # 输入图像文件夹
+    output_folder = r"D:\learn_torch\pytorch-CycleGAN\datasets\ostracoda\trainA_matting"  # 输出图像文件夹
+
+    process_images(input_folder, output_folder, model_path)