init repo

.gitignore

@@ -0,0 +1,172 @@
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+pip-wheel-metadata/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py,cover
+.hypothesis/
+.pytest_cache/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+*/.ipynb_checkpoints/*
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+.python-version
+
+# pipenv
+#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
+#   However, in case of collaboration, if having platform-specific dependencies or dependencies
+#   having no cross-platform support, pipenv may install dependencies that don't work, or not
+#   install all needed dependencies.
+#Pipfile.lock
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule
+celerybeat.pid
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/
+
+# pytype static type analyzer
+.pytype/
+
+# Cython debug symbols
+cython_debug/
+
+# .idea
+.idea/
+/idea/
+*.ipr
+*.iml
+*.iws
+
+# macos system
+.DS_Store
+
+### project ###
+
+# /diffvg/
+# big-lama*
+
+# pytorch-lighting logs
+lightning_logs/*
+
+# Edit settings
+.editorconfig
+
+# model checkpoint
+/checkpoint/u2net/u2net.pth
+!/checkpoint/placeholder.md
+
+# ignore local results
+/workspace/
+.workspace/
+
+# ignore files
+./tmp/
+./tmp/*
+/tmp/
+/tmp_select/
+/tmp_select/*

.gitmodules

@@ -0,0 +1,3 @@
+[submodule "diffvg"]
+	path = diffvg
+	url = https://github.com/BachiLi/diffvg.git

ImageReward/ImageReward.py

@@ -0,0 +1,177 @@
+'''
+@File       :   ImageReward.py
+@Time       :   2023/01/28 19:53:00
+@Auther     :   Jiazheng Xu
+@Contact    :   [email protected]
+@Description:   ImageReward Reward model.
+* Based on CLIP code base and improved-aesthetic-predictor code base
+* https://github.com/openai/CLIP
+* https://github.com/christophschuhmann/improved-aesthetic-predictor
+'''
+
+import os
+import torch
+import torch.nn as nn
+from PIL import Image
+from .models.BLIP.blip_pretrain import BLIP_Pretrain
+from torchvision.transforms import Compose, Resize, CenterCrop, ToTensor, Normalize
+
+try:
+    from torchvision.transforms import InterpolationMode
+
+    BICUBIC = InterpolationMode.BICUBIC
+except ImportError:
+    BICUBIC = Image.BICUBIC
+
+
+def _convert_image_to_rgb(image):
+    return image.convert("RGB")
+
+
+def _transform(n_px):
+    return Compose([
+        Resize(n_px, interpolation=BICUBIC),
+        CenterCrop(n_px),
+        _convert_image_to_rgb,
+        ToTensor(),
+        Normalize((0.48145466, 0.4578275, 0.40821073), (0.26862954, 0.26130258, 0.27577711)),
+    ])
+
+
+class MLP(nn.Module):
+    def __init__(self, input_size):
+        super().__init__()
+        self.input_size = input_size
+
+        self.layers = nn.Sequential(
+            nn.Linear(self.input_size, 1024),
+            # nn.ReLU(),
+            nn.Dropout(0.2),
+            nn.Linear(1024, 128),
+            # nn.ReLU(),
+            nn.Dropout(0.2),
+            nn.Linear(128, 64),
+            # nn.ReLU(),
+            nn.Dropout(0.1),
+            nn.Linear(64, 16),
+            # nn.ReLU(),
+            nn.Linear(16, 1)
+        )
+
+        # initial MLP param
+        for name, param in self.layers.named_parameters():
+            if 'weight' in name:
+                nn.init.normal_(param, mean=0.0, std=1.0 / (self.input_size + 1))
+            if 'bias' in name:
+                nn.init.constant_(param, val=0)
+
+    def forward(self, input):
+        return self.layers(input)
+
+
+class ImageReward(nn.Module):
+    def __init__(self, med_config, device='cpu'):
+        super().__init__()
+        self.device = device
+
+        self.blip = BLIP_Pretrain(image_size=224, vit='large', med_config=med_config)
+        self.preprocess = _transform(224)
+        self.mlp = MLP(768)
+
+        self.mean = 0.16717362830052426
+        self.std = 1.0333394966054072
+
+    def score_gard(self, prompt_ids, prompt_attention_mask, image):
+
+        image_embeds = self.blip.visual_encoder(image)
+        # text encode cross attention with image
+        image_atts = torch.ones(image_embeds.size()[:-1], dtype=torch.long).to(self.device)
+        text_output = self.blip.text_encoder(prompt_ids,
+                                             attention_mask=prompt_attention_mask,
+                                             encoder_hidden_states=image_embeds,
+                                             encoder_attention_mask=image_atts,
+                                             return_dict=True,
+                                             )
+
+        txt_features = text_output.last_hidden_state[:, 0, :]  # (feature_dim)
+        rewards = self.mlp(txt_features)
+        rewards = (rewards - self.mean) / self.std
+
+        return rewards
+
+    def score(self, prompt, image):
+
+        if (type(image).__name__ == 'list'):
+            _, rewards = self.inference_rank(prompt, image)
+            return rewards
+
+        # text encode
+        text_input = self.blip.tokenizer(prompt, padding='max_length', truncation=True, max_length=35,
+                                         return_tensors="pt").to(self.device)
+
+        # image encode
+        if isinstance(image, Image.Image):
+            pil_image = image
+        elif isinstance(image, str):
+            if os.path.isfile(image):
+                pil_image = Image.open(image)
+        else:
+            raise TypeError(
+                r'This image parameter type has not been supportted yet. Please pass PIL.Image or file path str.')
+
+        image = self.preprocess(pil_image).unsqueeze(0).to(self.device)
+        image_embeds = self.blip.visual_encoder(image)
+
+        # text encode cross attention with image
+        image_atts = torch.ones(image_embeds.size()[:-1], dtype=torch.long).to(self.device)
+        text_output = self.blip.text_encoder(text_input.input_ids,
+                                             attention_mask=text_input.attention_mask,
+                                             encoder_hidden_states=image_embeds,
+                                             encoder_attention_mask=image_atts,
+                                             return_dict=True,
+                                             )
+
+        txt_features = text_output.last_hidden_state[:, 0, :].float()  # (feature_dim)
+        rewards = self.mlp(txt_features)
+        rewards = (rewards - self.mean) / self.std
+
+        return rewards.detach().cpu().numpy().item()
+
+    def inference_rank(self, prompt, generations_list):
+
+        text_input = self.blip.tokenizer(prompt, padding='max_length', truncation=True, max_length=35,
+                                         return_tensors="pt").to(self.device)
+
+        txt_set = []
+        for generation in generations_list:
+            # image encode
+            if isinstance(generation, Image.Image):
+                pil_image = generation
+            elif isinstance(generation, str):
+                if os.path.isfile(generation):
+                    pil_image = Image.open(generation)
+            else:
+                raise TypeError(
+                    r'This image parameter type has not been supportted yet. Please pass PIL.Image or file path str.')
+
+            image = self.preprocess(pil_image).unsqueeze(0).to(self.device)
+            image_embeds = self.blip.visual_encoder(image)
+
+            # text encode cross attention with image
+            image_atts = torch.ones(image_embeds.size()[:-1], dtype=torch.long).to(self.device)
+            text_output = self.blip.text_encoder(text_input.input_ids,
+                                                 attention_mask=text_input.attention_mask,
+                                                 encoder_hidden_states=image_embeds,
+                                                 encoder_attention_mask=image_atts,
+                                                 return_dict=True)
+            txt_set.append(text_output.last_hidden_state[:, 0, :])
+
+        txt_features = torch.cat(txt_set, 0).float()  # [image_num, feature_dim]
+        rewards = self.mlp(txt_features)  # [image_num, 1]
+        rewards = (rewards - self.mean) / self.std
+        rewards = torch.squeeze(rewards)
+        _, rank = torch.sort(rewards, dim=0, descending=True)
+        _, indices = torch.sort(rank, dim=0)
+        indices = indices + 1
+
+        return indices.detach().cpu().numpy().tolist(), rewards.detach().cpu().numpy().tolist()

ImageReward/ReFL.py

@@ -0,0 +1,830 @@
+'''
+@File       :   ReFL.py
+@Time       :   2023/05/01 19:36:00
+@Auther     :   Jiazheng Xu
+@Contact    :   [email protected]
+@Description:   ReFL Algorithm.
+* Based on diffusers code base
+* https://github.com/huggingface/diffusers/blob/main/examples/text_to_image/train_text_to_image.py
+'''
+
+import argparse
+import logging
+import math
+import os
+import random
+from pathlib import Path
+
+import accelerate
+import numpy as np
+import torch
+import torch.nn.functional as F
+import torch.utils.checkpoint
+import transformers
+from accelerate import Accelerator
+from accelerate.logging import get_logger
+from accelerate.utils import ProjectConfiguration, set_seed
+from datasets import load_dataset
+from huggingface_hub import create_repo, upload_folder
+from packaging import version
+from tqdm.auto import tqdm
+from transformers import CLIPTextModel, CLIPTokenizer
+
+from PIL import Image
+import ImageReward as RM
+
+from torchvision.transforms import Compose, Resize, CenterCrop, Normalize
+
+try:
+    from torchvision.transforms import InterpolationMode
+
+    BICUBIC = InterpolationMode.BICUBIC
+except ImportError:
+    BICUBIC = Image.BICUBIC
+
+import diffusers
+from diffusers import AutoencoderKL, DDPMScheduler, StableDiffusionPipeline, UNet2DConditionModel
+from diffusers.optimization import get_scheduler
+from diffusers.training_utils import EMAModel
+from diffusers.utils import check_min_version, deprecate
+from diffusers.utils.import_utils import is_xformers_available
+
+# Will error if the minimal version of diffusers is not installed. Remove at your own risks.
+check_min_version("0.16.0.dev0")
+
+logger = get_logger(__name__, log_level="INFO")
+
+DATASET_NAME_MAPPING = {
+    "refl": ("image", "text"),
+}
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(description="Simple example of a training script.")
+    parser.add_argument(
+        "--grad_scale", type=float, default=1e-3, help="Scale divided for grad loss value."
+    )
+    parser.add_argument(
+        "--input_pertubation", type=float, default=0, help="The scale of input pretubation. Recommended 0.1."
+    )
+    parser.add_argument(
+        "--revision",
+        type=str,
+        default=None,
+        required=False,
+        help="Revision of pretrained model identifier from huggingface.co/models.",
+    )
+    parser.add_argument(
+        "--dataset_name",
+        type=str,
+        default=None,
+        help=(
+            "The name of the Dataset (from the HuggingFace hub) to train on (could be your own, possibly private,"
+            " dataset). It can also be a path pointing to a local copy of a dataset in your filesystem,"
+            " or to a folder containing files that 🤗 Datasets can understand."
+        ),
+    )
+    parser.add_argument(
+        "--dataset_config_name",
+        type=str,
+        default=None,
+        help="The config of the Dataset, leave as None if there's only one config.",
+    )
+    parser.add_argument(
+        "--image_column", type=str, default="image", help="The column of the dataset containing an image."
+    )
+    parser.add_argument(
+        "--caption_column",
+        type=str,
+        default="text",
+        help="The column of the dataset containing a caption or a list of captions.",
+    )
+    parser.add_argument(
+        "--max_train_samples",
+        type=int,
+        default=None,
+        help=(
+            "For debugging purposes or quicker training, truncate the number of training examples to this "
+            "value if set."
+        ),
+    )
+    parser.add_argument(
+        "--validation_prompts",
+        type=str,
+        default=None,
+        nargs="+",
+        help=("A set of prompts evaluated every `--validation_epochs` and logged to `--report_to`."),
+    )
+    parser.add_argument(
+        "--output_dir",
+        type=str,
+        default="checkpoint/refl",
+        help="The output directory where the model predictions and checkpoints will be written.",
+    )
+    parser.add_argument(
+        "--cache_dir",
+        type=str,
+        default=None,
+        help="The directory where the downloaded models and datasets will be stored.",
+    )
+    parser.add_argument("--seed", type=int, default=None, help="A seed for reproducible training.")
+    parser.add_argument(
+        "--resolution",
+        type=int,
+        default=512,
+        help=(
+            "The resolution for input images, all the images in the train/validation dataset will be resized to this"
+            " resolution"
+        ),
+    )
+    parser.add_argument(
+        "--center_crop",
+        default=False,
+        action="store_true",
+        help=(
+            "Whether to center crop the input images to the resolution. If not set, the images will be randomly"
+            " cropped. The images will be resized to the resolution first before cropping."
+        ),
+    )
+    parser.add_argument(
+        "--random_flip",
+        action="store_true",
+        help="whether to randomly flip images horizontally",
+    )
+    parser.add_argument(
+        "--train_batch_size", type=int, default=2, help="Batch size (per device) for the training dataloader."
+    )
+    parser.add_argument("--num_train_epochs", type=int, default=100)
+    parser.add_argument(
+        "--max_train_steps",
+        type=int,
+        default=100,
+        help="Total number of training steps to perform.  If provided, overrides num_train_epochs.",
+    )
+    parser.add_argument(
+        "--gradient_accumulation_steps",
+        type=int,
+        default=4,
+        help="Number of updates steps to accumulate before performing a backward/update pass.",
+    )
+    parser.add_argument(
+        "--gradient_checkpointing",
+        action="store_true",
+        help="Whether or not to use gradient checkpointing to save memory at the expense of slower backward pass.",
+    )
+    parser.add_argument(
+        "--learning_rate",
+        type=float,
+        default=1e-5,
+        help="Initial learning rate (after the potential warmup period) to use.",
+    )
+    parser.add_argument(
+        "--scale_lr",
+        action="store_true",
+        default=False,
+        help="Scale the learning rate by the number of GPUs, gradient accumulation steps, and batch size.",
+    )
+    parser.add_argument(
+        "--lr_scheduler",
+        type=str,
+        default="constant",
+        help=(
+            'The scheduler type to use. Choose between ["linear", "cosine", "cosine_with_restarts", "polynomial",'
+            ' "constant", "constant_with_warmup"]'
+        ),
+    )
+    parser.add_argument(
+        "--lr_warmup_steps", type=int, default=0, help="Number of steps for the warmup in the lr scheduler."
+    )
+    parser.add_argument(
+        "--snr_gamma",
+        type=float,
+        default=None,
+        help="SNR weighting gamma to be used if rebalancing the loss. Recommended value is 5.0. "
+             "More details here: https://arxiv.org/abs/2303.09556.",
+    )
+    parser.add_argument(
+        "--use_8bit_adam", action="store_true", help="Whether or not to use 8-bit Adam from bitsandbytes."
+    )
+    parser.add_argument(
+        "--allow_tf32",
+        action="store_true",
+        help=(
+            "Whether or not to allow TF32 on Ampere GPUs. Can be used to speed up training. For more information, see"
+            " https://pytorch.org/docs/stable/notes/cuda.html#tensorfloat-32-tf32-on-ampere-devices"
+        ),
+    )
+    parser.add_argument("--use_ema", action="store_true", help="Whether to use EMA model.")
+    parser.add_argument(
+        "--non_ema_revision",
+        type=str,
+        default=None,
+        required=False,
+        help=(
+            "Revision of pretrained non-ema model identifier. Must be a branch, tag or git identifier of the local or"
+            " remote repository specified with --pretrained_model_name_or_path."
+        ),
+    )
+    parser.add_argument(
+        "--dataloader_num_workers",
+        type=int,
+        default=0,
+        help=(
+            "Number of subprocesses to use for data loading. 0 means that the data will be loaded in the main process."
+        ),
+    )
+    parser.add_argument("--adam_beta1", type=float, default=0.9, help="The beta1 parameter for the Adam optimizer.")
+    parser.add_argument("--adam_beta2", type=float, default=0.999, help="The beta2 parameter for the Adam optimizer.")
+    parser.add_argument("--adam_weight_decay", type=float, default=1e-2, help="Weight decay to use.")
+    parser.add_argument("--adam_epsilon", type=float, default=1e-08, help="Epsilon value for the Adam optimizer")
+    parser.add_argument("--max_grad_norm", default=1.0, type=float, help="Max gradient norm.")
+    parser.add_argument("--push_to_hub", action="store_true", help="Whether or not to push the model to the Hub.")
+    parser.add_argument("--hub_token", type=str, default=None, help="The token to use to push to the Model Hub.")
+    parser.add_argument(
+        "--hub_model_id",
+        type=str,
+        default=None,
+        help="The name of the repository to keep in sync with the local `output_dir`.",
+    )
+    parser.add_argument(
+        "--logging_dir",
+        type=str,
+        default="logs",
+        help=(
+            "[TensorBoard](https://www.tensorflow.org/tensorboard) log directory. Will default to"
+            " *output_dir/runs/**CURRENT_DATETIME_HOSTNAME***."
+        ),
+    )
+    parser.add_argument(
+        "--mixed_precision",
+        type=str,
+        default=None,
+        choices=["no", "fp16", "bf16"],
+        help=(
+            "Whether to use mixed precision. Choose between fp16 and bf16 (bfloat16). Bf16 requires PyTorch >="
+            " 1.10.and an Nvidia Ampere GPU.  Default to the value of accelerate config of the current system or the"
+            " flag passed with the `accelerate.launch` command. Use this argument to override the accelerate config."
+        ),
+    )
+    parser.add_argument(
+        "--report_to",
+        type=str,
+        default="tensorboard",
+        help=(
+            'The integration to report the results and logs to. Supported platforms are `"tensorboard"`'
+            ' (default), `"wandb"` and `"comet_ml"`. Use `"all"` to report to all integrations.'
+        ),
+    )
+    parser.add_argument("--local_rank", type=int, default=-1, help="For distributed training: local_rank")
+    parser.add_argument(
+        "--checkpointing_steps",
+        type=int,
+        default=100,
+        help=(
+            "Save a checkpoint of the training state every X updates. These checkpoints are only suitable for resuming"
+            " training using `--resume_from_checkpoint`."
+        ),
+    )
+    parser.add_argument(
+        "--checkpoints_total_limit",
+        type=int,
+        default=None,
+        help=(
+            "Max number of checkpoints to store. Passed as `total_limit` to the `Accelerator` `ProjectConfiguration`."
+            " See Accelerator::save_state https://huggingface.co/docs/accelerate/package_reference/accelerator#accelerate.Accelerator.save_state"
+            " for more docs"
+        ),
+    )
+    parser.add_argument(
+        "--resume_from_checkpoint",
+        type=str,
+        default=None,
+        help=(
+            "Whether training should be resumed from a previous checkpoint. Use a path saved by"
+            ' `--checkpointing_steps`, or `"latest"` to automatically select the last available checkpoint.'
+        ),
+    )
+    parser.add_argument(
+        "--enable_xformers_memory_efficient_attention", action="store_true", help="Whether or not to use xformers."
+    )
+    parser.add_argument("--noise_offset", type=float, default=0, help="The scale of noise offset.")
+    parser.add_argument(
+        "--validation_epochs",
+        type=int,
+        default=5,
+        help="Run validation every X epochs.",
+    )
+    parser.add_argument(
+        "--tracker_project_name",
+        type=str,
+        default="text2image-refl",
+        help=(
+            "The `project_name` argument passed to Accelerator.init_trackers for"
+            " more information see https://huggingface.co/docs/accelerate/v0.17.0/en/package_reference/accelerator#accelerate.Accelerator"
+        ),
+    )
+
+    args = parser.parse_args()
+    env_local_rank = int(os.environ.get("LOCAL_RANK", -1))
+    if env_local_rank != -1 and env_local_rank != args.local_rank:
+        args.local_rank = env_local_rank
+
+    # default to using the same revision for the non-ema model if not specified
+    if args.non_ema_revision is None:
+        args.non_ema_revision = args.revision
+
+    return args
+
+
+class Trainer(object):
+
+    def __init__(self, pretrained_model_name_or_path, train_data_dir, args):
+
+        self.pretrained_model_name_or_path = pretrained_model_name_or_path
+        self.train_data_dir = train_data_dir
+
+        # Sanity checks
+        if args.dataset_name is None and self.train_data_dir is None:
+            raise ValueError("Need either a dataset name or a training folder.")
+
+        if args.non_ema_revision is not None:
+            deprecate(
+                "non_ema_revision!=None",
+                "0.15.0",
+                message=(
+                    "Downloading 'non_ema' weights from revision branches of the Hub is deprecated. Please make sure to"
+                    " use `--variant=non_ema` instead."
+                ),
+            )
+        logging_dir = os.path.join(args.output_dir, args.logging_dir)
+
+        accelerator_project_config = ProjectConfiguration(total_limit=args.checkpoints_total_limit)
+
+        self.accelerator = Accelerator(
+            gradient_accumulation_steps=args.gradient_accumulation_steps,
+            mixed_precision=args.mixed_precision,
+            log_with=args.report_to,
+            logging_dir=logging_dir,
+            project_config=accelerator_project_config,
+        )
+
+        # Make one log on every process with the configuration for debugging.
+        logging.basicConfig(
+            format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
+            datefmt="%m/%d/%Y %H:%M:%S",
+            level=logging.INFO,
+        )
+        logger.info(self.accelerator.state, main_process_only=False)
+        if self.accelerator.is_local_main_process:
+            transformers.utils.logging.set_verbosity_warning()
+            diffusers.utils.logging.set_verbosity_info()
+        else:
+            transformers.utils.logging.set_verbosity_error()
+            diffusers.utils.logging.set_verbosity_error()
+
+        # If passed along, set the training seed now.
+        if args.seed is not None:
+            set_seed(args.seed)
+
+        # Handle the repository creation
+        if self.accelerator.is_main_process:
+            if args.output_dir is not None:
+                os.makedirs(args.output_dir, exist_ok=True)
+
+            if args.push_to_hub:
+                self.repo_id = create_repo(
+                    repo_id=args.hub_model_id or Path(args.output_dir).name, exist_ok=True, token=args.hub_token
+                ).repo_id
+
+        # Load scheduler, tokenizer and models.
+        self.noise_scheduler = DDPMScheduler.from_pretrained(self.pretrained_model_name_or_path, subfolder="scheduler")
+        tokenizer = CLIPTokenizer.from_pretrained(
+            self.pretrained_model_name_or_path, subfolder="tokenizer", revision=args.revision
+        )
+        self.text_encoder = CLIPTextModel.from_pretrained(
+            self.pretrained_model_name_or_path, subfolder="text_encoder", revision=args.revision
+        )
+        self.vae = AutoencoderKL.from_pretrained(self.pretrained_model_name_or_path, subfolder="vae",
+                                                 revision=args.revision)
+        self.unet = UNet2DConditionModel.from_pretrained(
+            self.pretrained_model_name_or_path, subfolder="unet", revision=args.non_ema_revision
+        )
+        self.reward_model = RM.load("ImageReward-v1.0", device=self.accelerator.device)
+
+        # Freeze vae and text_encoder
+        self.vae.requires_grad_(False)
+        self.text_encoder.requires_grad_(False)
+        self.reward_model.requires_grad_(False)
+
+        # Create EMA for the unet.
+        if args.use_ema:
+            self.ema_unet = UNet2DConditionModel.from_pretrained(
+                self.pretrained_model_name_or_path, subfolder="unet", revision=args.revision
+            )
+            self.ema_unet = EMAModel(self.ema_unet.parameters(), model_cls=UNet2DConditionModel,
+                                     model_config=self.ema_unet.config)
+
+        if args.enable_xformers_memory_efficient_attention:
+            if is_xformers_available():
+                import xformers
+
+                xformers_version = version.parse(xformers.__version__)
+                if xformers_version == version.parse("0.0.16"):
+                    logger.warn(
+                        "xFormers 0.0.16 cannot be used for training in some GPUs. If you observe problems during training, please update xFormers to at least 0.0.17. See https://huggingface.co/docs/diffusers/main/en/optimization/xformers for more details."
+                    )
+                self.unet.enable_xformers_memory_efficient_attention()
+            else:
+                raise ValueError("xformers is not available. Make sure it is installed correctly")
+
+        # `accelerate` 0.16.0 will have better support for customized saving
+        if version.parse(accelerate.__version__) >= version.parse("0.16.0"):
+            # create custom saving & loading hooks so that `self.accelerator.save_state(...)` serializes in a nice format
+            def save_model_hook(models, weights, output_dir):
+                if args.use_ema:
+                    self.ema_unet.save_pretrained(os.path.join(output_dir, "unet_ema"))
+
+                for i, model in enumerate(models):
+                    model.save_pretrained(os.path.join(output_dir, "unet"))
+
+                    # make sure to pop weight so that corresponding model is not saved again
+                    weights.pop()
+
+            def load_model_hook(models, input_dir):
+                if args.use_ema:
+                    load_model = EMAModel.from_pretrained(os.path.join(input_dir, "unet_ema"), UNet2DConditionModel)
+                    self.ema_unet.load_state_dict(load_model.state_dict())
+                    self.ema_unet.to(self.accelerator.device)
+                    del load_model
+
+                for i in range(len(models)):
+                    # pop models so that they are not loaded again
+                    model = models.pop()
+
+                    # load diffusers style into model
+                    load_model = UNet2DConditionModel.from_pretrained(input_dir, subfolder="unet")
+                    model.register_to_config(**load_model.config)
+
+                    model.load_state_dict(load_model.state_dict())
+                    del load_model
+
+            self.accelerator.register_save_state_pre_hook(save_model_hook)
+            self.accelerator.register_load_state_pre_hook(load_model_hook)
+
+        if args.gradient_checkpointing:
+            self.unet.enable_gradient_checkpointing()
+
+        # Enable TF32 for faster training on Ampere GPUs,
+        # cf https://pytorch.org/docs/stable/notes/cuda.html#tensorfloat-32-tf32-on-ampere-devices
+        if args.allow_tf32:
+            torch.backends.cuda.matmul.allow_tf32 = True
+
+        if args.scale_lr:
+            args.learning_rate = (
+                    args.learning_rate * args.gradient_accumulation_steps * args.train_batch_size * self.accelerator.num_processes
+            )
+
+        # Initialize the optimizer
+        if args.use_8bit_adam:
+            try:
+                import bitsandbytes as bnb
+            except ImportError:
+                raise ImportError(
+                    "Please install bitsandbytes to use 8-bit Adam. You can do so by running `pip install bitsandbytes`"
+                )
+
+            optimizer_cls = bnb.optim.AdamW8bit
+        else:
+            optimizer_cls = torch.optim.AdamW
+
+        self.optimizer = optimizer_cls(
+            self.unet.parameters(),
+            lr=args.learning_rate,
+            betas=(args.adam_beta1, args.adam_beta2),
+            weight_decay=args.adam_weight_decay,
+            eps=args.adam_epsilon,
+        )
+
+        # Get the datasets: you can either provide your own training and evaluation files (see below)
+        # or specify a Dataset from the hub (the dataset will be downloaded automatically from the datasets Hub).
+
+        # In distributed training, the load_dataset function guarantees that only one local process can concurrently
+        # download the dataset.
+        if args.dataset_name is not None:
+            # Downloading and loading a dataset from the hub.
+            dataset = load_dataset(
+                args.dataset_name,
+                args.dataset_config_name,
+                cache_dir=args.cache_dir,
+            )
+        else:
+            data_files = {}
+            data_files["train"] = self.train_data_dir
+            dataset = load_dataset(
+                "json",
+                data_files=data_files,
+                cache_dir=args.cache_dir,
+            )
+            # See more about loading custom images at
+            # https://huggingface.co/docs/datasets/v2.4.0/en/image_load#imagefolder
+
+        # Preprocessing the datasets.
+        # We need to tokenize inputs and targets.
+        column_names = dataset["train"].column_names
+
+        # Get the column names for input/target.
+        dataset_columns = DATASET_NAME_MAPPING.get(args.dataset_name, None)
+        if args.image_column is None:
+            image_column = dataset_columns[0] if dataset_columns is not None else column_names[0]
+        else:
+            image_column = args.image_column
+            if image_column not in column_names:
+                raise ValueError(
+                    f"--image_column' value '{args.image_column}' needs to be one of: {', '.join(column_names)}"
+                )
+        if args.caption_column is None:
+            caption_column = dataset_columns[1] if dataset_columns is not None else column_names[1]
+        else:
+            caption_column = args.caption_column
+            if caption_column not in column_names:
+                raise ValueError(
+                    f"--caption_column' value '{args.caption_column}' needs to be one of: {', '.join(column_names)}"
+                )
+
+        # Preprocessing the datasets.
+        # We need to tokenize input captions and transform the images.
+        def tokenize_captions(examples, is_train=True):
+            captions = []
+            for caption in examples[caption_column]:
+                if isinstance(caption, str):
+                    captions.append(caption)
+                elif isinstance(caption, (list, np.ndarray)):
+                    # take a random caption if there are multiple
+                    captions.append(random.choice(caption) if is_train else caption[0])
+                else:
+                    raise ValueError(
+                        f"Caption column `{caption_column}` should contain either strings or lists of strings."
+                    )
+            inputs = tokenizer(
+                captions, max_length=tokenizer.model_max_length, padding="max_length", truncation=True,
+                return_tensors="pt"
+            )
+            return inputs.input_ids
+
+        def preprocess_train(examples):
+            examples["input_ids"] = tokenize_captions(examples)
+            examples["rm_input_ids"] = self.reward_model.blip.tokenizer(examples[caption_column], padding='max_length',
+                                                                        truncation=True, max_length=35,
+                                                                        return_tensors="pt").input_ids
+            examples["rm_attention_mask"] = self.reward_model.blip.tokenizer(examples[caption_column],
+                                                                             padding='max_length', truncation=True,
+                                                                             max_length=35,
+                                                                             return_tensors="pt").attention_mask
+            return examples
+
+        with self.accelerator.main_process_first():
+            if args.max_train_samples is not None:
+                dataset["train"] = dataset["train"].shuffle(seed=args.seed).select(range(args.max_train_samples))
+            # Set the training transforms
+            self.train_dataset = dataset["train"].with_transform(preprocess_train)
+
+        def collate_fn(examples):
+            input_ids = torch.stack([example["input_ids"] for example in examples])
+            rm_input_ids = torch.stack([example["rm_input_ids"] for example in examples])
+            rm_attention_mask = torch.stack([example["rm_attention_mask"] for example in examples])
+            input_ids = input_ids.view(-1, input_ids.shape[-1])
+            rm_input_ids = rm_input_ids.view(-1, rm_input_ids.shape[-1])
+            rm_attention_mask = rm_attention_mask.view(-1, rm_attention_mask.shape[-1])
+            return {"input_ids": input_ids, "rm_input_ids": rm_input_ids, "rm_attention_mask": rm_attention_mask}
+
+        # DataLoaders creation:
+        self.train_dataloader = torch.utils.data.DataLoader(
+            self.train_dataset,
+            shuffle=True,
+            collate_fn=collate_fn,
+            batch_size=args.train_batch_size,
+            num_workers=args.dataloader_num_workers,
+        )
+
+        # Scheduler and math around the number of training steps.
+        overrode_max_train_steps = False
+        self.num_update_steps_per_epoch = math.ceil(len(self.train_dataloader) / args.gradient_accumulation_steps)
+        if args.max_train_steps is None:
+            args.max_train_steps = args.num_train_epochs * self.num_update_steps_per_epoch
+            overrode_max_train_steps = True
+
+        self.lr_scheduler = get_scheduler(
+            args.lr_scheduler,
+            optimizer=self.optimizer,
+            num_warmup_steps=args.lr_warmup_steps * args.gradient_accumulation_steps,
+            num_training_steps=args.max_train_steps * args.gradient_accumulation_steps,
+        )
+
+        # Prepare everything with our `self.accelerator`.
+        self.unet, self.optimizer, self.train_dataloader, self.lr_scheduler = self.accelerator.prepare(
+            self.unet, self.optimizer, self.train_dataloader, self.lr_scheduler
+        )
+
+        if args.use_ema:
+            self.ema_unet.to(self.accelerator.device)
+
+        # For mixed precision training we cast the text_encoder and vae weights to half-precision
+        # as these models are only used for inference, keeping weights in full precision is not required.
+        self.weight_dtype = torch.float32
+        if self.accelerator.mixed_precision == "fp16":
+            self.weight_dtype = torch.float16
+        elif self.accelerator.mixed_precision == "bf16":
+            self.weight_dtype = torch.bfloat16
+
+        # Move text_encode and vae to gpu and cast to self.weight_dtype
+        self.text_encoder.to(self.accelerator.device, dtype=self.weight_dtype)
+        self.vae.to(self.accelerator.device, dtype=self.weight_dtype)
+        self.reward_model.to(self.accelerator.device, dtype=self.weight_dtype)
+
+        # We need to recalculate our total training steps as the size of the training dataloader may have changed.
+        self.num_update_steps_per_epoch = math.ceil(len(self.train_dataloader) / args.gradient_accumulation_steps)
+        if overrode_max_train_steps:
+            args.max_train_steps = args.num_train_epochs * self.num_update_steps_per_epoch
+        # Afterwards we recalculate our number of training epochs
+        args.num_train_epochs = math.ceil(args.max_train_steps / self.num_update_steps_per_epoch)
+
+        # We need to initialize the trackers we use, and also store our configuration.
+        # The trackers initializes automatically on the main process.
+        if self.accelerator.is_main_process:
+            tracker_config = dict(vars(args))
+            tracker_config.pop("validation_prompts")
+            self.accelerator.init_trackers(args.tracker_project_name, tracker_config)
+
+    def train(self, args):
+
+        # Train!
+        total_batch_size = args.train_batch_size * self.accelerator.num_processes * args.gradient_accumulation_steps
+
+        logger.info("***** Running training *****")
+        logger.info(f"  Num examples = {len(self.train_dataset)}")
+        logger.info(f"  Num Epochs = {args.num_train_epochs}")
+        logger.info(f"  Instantaneous batch size per device = {args.train_batch_size}")
+        logger.info(f"  Total train batch size (w. parallel, distributed & accumulation) = {total_batch_size}")
+        logger.info(f"  Gradient Accumulation steps = {args.gradient_accumulation_steps}")
+        logger.info(f"  Total optimization steps = {args.max_train_steps}")
+        global_step = 0
+        first_epoch = 0
+
+        # Potentially load in the weights and states from a previous save
+        if args.resume_from_checkpoint:
+            if args.resume_from_checkpoint != "latest":
+                path = os.path.basename(args.resume_from_checkpoint)
+            else:
+                # Get the most recent checkpoint
+                dirs = os.listdir(args.output_dir)
+                dirs = [d for d in dirs if d.startswith("checkpoint")]
+                dirs = sorted(dirs, key=lambda x: int(x.split("-")[1]))
+                path = dirs[-1] if len(dirs) > 0 else None
+
+            if path is None:
+                self.accelerator.print(
+                    f"Checkpoint '{args.resume_from_checkpoint}' does not exist. Starting a new training run."
+                )
+                args.resume_from_checkpoint = None
+            else:
+                self.accelerator.print(f"Resuming from checkpoint {path}")
+                self.accelerator.load_state(os.path.join(args.output_dir, path))
+                global_step = int(path.split("-")[1])
+
+                resume_global_step = global_step * args.gradient_accumulation_steps
+                first_epoch = global_step // self.num_update_steps_per_epoch
+                resume_step = resume_global_step % (self.num_update_steps_per_epoch * args.gradient_accumulation_steps)
+
+        # Only show the progress bar once on each machine.
+        progress_bar = tqdm(range(global_step, args.max_train_steps),
+                            disable=not self.accelerator.is_local_main_process)
+        progress_bar.set_description("Steps")
+
+        for epoch in range(first_epoch, args.num_train_epochs):
+            self.unet.train()
+            train_loss = 0.0
+            for step, batch in enumerate(self.train_dataloader):
+                # Skip steps until we reach the resumed step
+                if args.resume_from_checkpoint and epoch == first_epoch and step < resume_step:
+                    if step % args.gradient_accumulation_steps == 0:
+                        progress_bar.update(1)
+                    continue
+
+                with self.accelerator.accumulate(self.unet):
+                    encoder_hidden_states = self.text_encoder(batch["input_ids"])[0]
+                    latents = torch.randn((args.train_batch_size, 4, 64, 64), device=self.accelerator.device)
+
+                    self.noise_scheduler.set_timesteps(40, device=self.accelerator.device)
+                    timesteps = self.noise_scheduler.timesteps
+
+                    mid_timestep = random.randint(30, 39)
+
+                    for i, t in enumerate(timesteps[:mid_timestep]):
+                        with torch.no_grad():
+                            latent_model_input = latents
+                            latent_model_input = self.noise_scheduler.scale_model_input(latent_model_input, t)
+                            noise_pred = self.unet(
+                                latent_model_input,
+                                t,
+                                encoder_hidden_states=encoder_hidden_states,
+                            ).sample
+                            latents = self.noise_scheduler.step(noise_pred, t, latents).prev_sample
+
+                    latent_model_input = latents
+                    latent_model_input = self.noise_scheduler.scale_model_input(latent_model_input,
+                                                                                timesteps[mid_timestep])
+                    noise_pred = self.unet(
+                        latent_model_input,
+                        timesteps[mid_timestep],
+                        encoder_hidden_states=encoder_hidden_states,
+                    ).sample
+                    pred_original_sample = self.noise_scheduler.step(noise_pred, timesteps[mid_timestep],
+                                                                     latents).pred_original_sample.to(self.weight_dtype)
+
+                    pred_original_sample = 1 / self.vae.config.scaling_factor * pred_original_sample
+                    image = self.vae.decode(pred_original_sample.to(self.weight_dtype)).sample
+                    image = (image / 2 + 0.5).clamp(0, 1)
+
+                    # image encode
+                    def _transform():
+                        return Compose([
+                            Resize(224, interpolation=BICUBIC),
+                            CenterCrop(224),
+                            Normalize((0.48145466, 0.4578275, 0.40821073), (0.26862954, 0.26130258, 0.27577711)),
+                        ])
+
+                    rm_preprocess = _transform()
+                    image = rm_preprocess(image).to(self.accelerator.device)
+
+                    rewards = self.reward_model.score_gard(batch["rm_input_ids"], batch["rm_attention_mask"], image)
+                    loss = F.relu(-rewards + 2)
+                    loss = loss.mean() * args.grad_scale
+
+                    # Gather the losses across all processes for logging (if we use distributed training).
+                    avg_loss = self.accelerator.gather(loss.repeat(args.train_batch_size)).mean()
+                    train_loss += avg_loss.item() / args.gradient_accumulation_steps
+
+                    # Backpropagate
+                    self.accelerator.backward(loss)
+                    if self.accelerator.sync_gradients:
+                        self.accelerator.clip_grad_norm_(self.unet.parameters(), args.max_grad_norm)
+                    self.optimizer.step()
+                    self.lr_scheduler.step()
+                    self.optimizer.zero_grad()
+
+                # Checks if the self.accelerator has performed an optimization step behind the scenes
+                if self.accelerator.sync_gradients:
+                    if args.use_ema:
+                        self.ema_unet.step(self.unet.parameters())
+                    progress_bar.update(1)
+                    global_step += 1
+                    self.accelerator.log({"train_loss": train_loss}, step=global_step)
+                    train_loss = 0.0
+
+                    if global_step % args.checkpointing_steps == 0:
+                        if self.accelerator.is_main_process:
+                            save_path = os.path.join(args.output_dir, f"checkpoint-{global_step}")
+                            self.accelerator.save_state(save_path)
+                            logger.info(f"Saved state to {save_path}")
+
+                logs = {"step_loss": loss.detach().item(), "lr": self.lr_scheduler.get_last_lr()[0]}
+                progress_bar.set_postfix(**logs)
+
+                if global_step >= args.max_train_steps:
+                    break
+
+            if self.accelerator.is_main_process:
+                if args.validation_prompts is not None and epoch % args.validation_epochs == 0:
+                    if args.use_ema:
+                        # Store the UNet parameters temporarily and load the EMA parameters to perform inference.
+                        self.ema_unet.store(self.unet.parameters())
+                        self.ema_unet.copy_to(self.unet.parameters())
+                    if args.use_ema:
+                        # Switch back to the original UNet parameters.
+                        self.ema_unet.restore(self.unet.parameters())
+
+        # Create the pipeline using the trained modules and save it.
+        self.accelerator.wait_for_everyone()
+        if self.accelerator.is_main_process:
+            self.unet = self.accelerator.unwrap_model(self.unet)
+            if args.use_ema:
+                self.ema_unet.copy_to(self.unet.parameters())
+
+            pipeline = StableDiffusionPipeline.from_pretrained(
+                self.pretrained_model_name_or_path,
+                text_encoder=self.text_encoder,
+                vae=self.vae,
+                unet=self.unet,
+                revision=args.revision,
+            )
+            pipeline.save_pretrained(args.output_dir)
+
+            if args.push_to_hub:
+                upload_folder(
+                    repo_id=self.repo_id,
+                    folder_path=args.output_dir,
+                    commit_message="End of training",
+                    ignore_patterns=["step_*", "epoch_*"],
+                )
+
+        self.accelerator.end_training()

ImageReward/__init__.py

@@ -0,0 +1,3 @@
+from .utils import *
+from .models import *
+from .ReFL import *

ImageReward/models/AestheticScore.py

@@ -0,0 +1,95 @@
+'''
+@File       :   AestheticScore.py
+@Time       :   2023/02/12 14:54:00
+@Auther     :   Jiazheng Xu
+@Contact    :   [email protected]
+@Description:   AestheticScore.
+* Based on improved-aesthetic-predictor code base
+* https://github.com/christophschuhmann/improved-aesthetic-predictor
+'''
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from PIL import Image
+import clip
+
+
+# if you changed the MLP architecture during training, change it also here:
+class MLP(nn.Module):
+    def __init__(self, input_size):
+        super().__init__()
+        self.input_size = input_size
+        self.layers = nn.Sequential(
+            nn.Linear(self.input_size, 1024),
+            # nn.ReLU(),
+            nn.Dropout(0.2),
+            nn.Linear(1024, 128),
+            # nn.ReLU(),
+            nn.Dropout(0.2),
+            nn.Linear(128, 64),
+            # nn.ReLU(),
+            nn.Dropout(0.1),
+
+            nn.Linear(64, 16),
+            # nn.ReLU(),
+
+            nn.Linear(16, 1)
+        )
+
+    def forward(self, x):
+        return self.layers(x)
+
+
+class AestheticScore(nn.Module):
+    def __init__(self, download_root, device='cpu'):
+        super().__init__()
+        self.device = device
+        self.clip_model, self.preprocess = clip.load("ViT-L/14", device=self.device, jit=False,
+                                                     download_root=download_root)
+        self.mlp = MLP(768)
+
+        if device == "cpu":
+            self.clip_model.float()
+        else:
+            clip.model.convert_weights(
+                self.clip_model)  # Actually this line is unnecessary since clip by default already on float16
+
+        # have clip.logit_scale require no grad.
+        self.clip_model.logit_scale.requires_grad_(False)
+
+    def score(self, prompt, image_path):
+
+        if (type(image_path).__name__ == 'list'):
+            _, rewards = self.inference_rank(prompt, image_path)
+            return rewards
+
+        # image encode
+        pil_image = Image.open(image_path)
+        image = self.preprocess(pil_image).unsqueeze(0).to(self.device)
+        image_features = F.normalize(self.clip_model.encode_image(image)).float()
+
+        # score
+        rewards = self.mlp(image_features)
+
+        return rewards.detach().cpu().numpy().item()
+
+    def inference_rank(self, prompt, generations_list):
+
+        img_set = []
+        for generations in generations_list:
+            # image encode
+            img_path = generations
+            pil_image = Image.open(img_path)
+            image = self.preprocess(pil_image).unsqueeze(0).to(self.device)
+            image_features = F.normalize(self.clip_model.encode_image(image))
+            img_set.append(image_features)
+
+        img_features = torch.cat(img_set, 0).float()  # [image_num, feature_dim]
+        rewards = self.mlp(img_features)
+        rewards = torch.squeeze(rewards)
+        _, rank = torch.sort(rewards, dim=0, descending=True)
+        _, indices = torch.sort(rank, dim=0)
+        indices = indices + 1
+
+        return indices.detach().cpu().numpy().tolist(), rewards.detach().cpu().numpy().tolist()

ImageReward/models/BLIP/__init__.py

@@ -0,0 +1 @@
+from .blip_pretrain import *

ImageReward/models/BLIP/blip.py

@@ -0,0 +1,70 @@
+'''
+ * Adapted from BLIP (https://github.com/salesforce/BLIP)
+'''
+
+import warnings
+warnings.filterwarnings("ignore")
+
+import torch
+import os
+from urllib.parse import urlparse
+from timm.models.hub import download_cached_file
+from transformers import BertTokenizer
+from .vit import VisionTransformer, interpolate_pos_embed
+
+
+def init_tokenizer():
+    tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
+    tokenizer.add_special_tokens({'bos_token':'[DEC]'})
+    tokenizer.add_special_tokens({'additional_special_tokens':['[ENC]']})       
+    tokenizer.enc_token_id = tokenizer.additional_special_tokens_ids[0]  
+    return tokenizer
+
+
+def create_vit(vit, image_size, use_grad_checkpointing=False, ckpt_layer=0, drop_path_rate=0):
+        
+    assert vit in ['base', 'large'], "vit parameter must be base or large"
+    if vit=='base':
+        vision_width = 768
+        visual_encoder = VisionTransformer(img_size=image_size, patch_size=16, embed_dim=vision_width, depth=12, 
+                                           num_heads=12, use_grad_checkpointing=use_grad_checkpointing, ckpt_layer=ckpt_layer,
+                                           drop_path_rate=0 or drop_path_rate
+                                          )   
+    elif vit=='large':
+        vision_width = 1024
+        visual_encoder = VisionTransformer(img_size=image_size, patch_size=16, embed_dim=vision_width, depth=24, 
+                                           num_heads=16, use_grad_checkpointing=use_grad_checkpointing, ckpt_layer=ckpt_layer,
+                                           drop_path_rate=0.1 or drop_path_rate
+                                          )   
+    return visual_encoder, vision_width
+
+
+def is_url(url_or_filename):
+    parsed = urlparse(url_or_filename)
+    return parsed.scheme in ("http", "https")
+
+def load_checkpoint(model,url_or_filename):
+    if is_url(url_or_filename):
+        cached_file = download_cached_file(url_or_filename, check_hash=False, progress=True)
+        checkpoint = torch.load(cached_file, map_location='cpu') 
+    elif os.path.isfile(url_or_filename):        
+        checkpoint = torch.load(url_or_filename, map_location='cpu') 
+    else:
+        raise RuntimeError('checkpoint url or path is invalid')
+        
+    state_dict = checkpoint['model']
+    
+    state_dict['visual_encoder.pos_embed'] = interpolate_pos_embed(state_dict['visual_encoder.pos_embed'],model.visual_encoder) 
+    if 'visual_encoder_m.pos_embed' in model.state_dict().keys():
+        state_dict['visual_encoder_m.pos_embed'] = interpolate_pos_embed(state_dict['visual_encoder_m.pos_embed'],
+                                                                         model.visual_encoder_m)    
+    for key in model.state_dict().keys():
+        if key in state_dict.keys():
+            if state_dict[key].shape!=model.state_dict()[key].shape:
+                print(key, ": ", state_dict[key].shape, ', ', model.state_dict()[key].shape)
+                del state_dict[key]
+    
+    msg = model.load_state_dict(state_dict,strict=False)
+    print('load checkpoint from %s'%url_or_filename)  
+    return model,msg
+    

ImageReward/models/BLIP/blip_pretrain.py

@@ -0,0 +1,43 @@
+'''
+ * Adapted from BLIP (https://github.com/salesforce/BLIP)
+'''
+
+import transformers
+transformers.logging.set_verbosity_error()
+
+from torch import nn
+import os
+from .med import BertConfig, BertModel
+from .blip import create_vit, init_tokenizer
+
+class BLIP_Pretrain(nn.Module):
+    def __init__(self,                 
+                 med_config = "med_config.json",  
+                 image_size = 224,
+                 vit = 'base',
+                 vit_grad_ckpt = False,
+                 vit_ckpt_layer = 0,                    
+                 embed_dim = 256,     
+                 queue_size = 57600,
+                 momentum = 0.995,
+                 ):
+        """
+        Args:
+            med_config (str): path for the mixture of encoder-decoder model's configuration file
+            image_size (int): input image size
+            vit (str): model size of vision transformer
+        """               
+        super().__init__()
+        
+        self.visual_encoder, vision_width = create_vit(vit,image_size, vit_grad_ckpt, vit_ckpt_layer, 0)
+        
+        self.tokenizer = init_tokenizer()   
+        encoder_config = BertConfig.from_json_file(med_config)
+        encoder_config.encoder_width = vision_width
+        self.text_encoder = BertModel(config=encoder_config, add_pooling_layer=False)
+
+        text_width = self.text_encoder.config.hidden_size
+        
+        self.vision_proj = nn.Linear(vision_width, embed_dim)
+        self.text_proj = nn.Linear(text_width, embed_dim)
+

ImageReward/models/BLIP/med.py

@@ -0,0 +1,947 @@
+'''
+ * Adapted from BLIP (https://github.com/salesforce/BLIP)
+ * Based on huggingface code base
+ * https://github.com/huggingface/transformers/blob/v4.15.0/src/transformers/models/bert
+'''
+
+import math
+from typing import Tuple
+
+import torch
+from torch import Tensor, device, nn
+import torch.utils.checkpoint
+from torch import nn
+from torch.nn import CrossEntropyLoss
+
+from transformers.activations import ACT2FN
+from transformers.file_utils import (
+    ModelOutput,
+)
+from transformers.modeling_outputs import (
+    BaseModelOutputWithPastAndCrossAttentions,
+    BaseModelOutputWithPoolingAndCrossAttentions,
+    CausalLMOutputWithCrossAttentions,
+    MaskedLMOutput,
+    MultipleChoiceModelOutput,
+    NextSentencePredictorOutput,
+    QuestionAnsweringModelOutput,
+    SequenceClassifierOutput,
+    TokenClassifierOutput,
+)
+from transformers.modeling_utils import (
+    PreTrainedModel,
+    apply_chunking_to_forward,
+    find_pruneable_heads_and_indices,
+    prune_linear_layer,
+)
+from transformers.utils import logging
+from transformers.models.bert.configuration_bert import BertConfig
+
+
+logger = logging.get_logger(__name__)
+
+
+class BertEmbeddings(nn.Module):
+    """Construct the embeddings from word and position embeddings."""
+
+    def __init__(self, config):
+        super().__init__()
+        self.word_embeddings = nn.Embedding(config.vocab_size, config.hidden_size, padding_idx=config.pad_token_id)
+        self.position_embeddings = nn.Embedding(config.max_position_embeddings, config.hidden_size)
+
+        # self.LayerNorm is not snake-cased to stick with TensorFlow model variable name and be able to load
+        # any TensorFlow checkpoint file
+        self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+        self.dropout = nn.Dropout(config.hidden_dropout_prob)
+
+        # position_ids (1, len position emb) is contiguous in memory and exported when serialized
+        self.register_buffer("position_ids", torch.arange(config.max_position_embeddings).expand((1, -1)))
+        self.position_embedding_type = getattr(config, "position_embedding_type", "absolute")
+        
+        self.config = config
+
+    def forward(
+        self, input_ids=None, position_ids=None, inputs_embeds=None, past_key_values_length=0
+    ):
+        if input_ids is not None:
+            input_shape = input_ids.size()
+        else:
+            input_shape = inputs_embeds.size()[:-1]
+
+        seq_length = input_shape[1]
+
+        if position_ids is None:
+            position_ids = self.position_ids[:, past_key_values_length : seq_length + past_key_values_length]
+
+        if inputs_embeds is None:
+            inputs_embeds = self.word_embeddings(input_ids)
+
+        embeddings = inputs_embeds
+
+        if self.position_embedding_type == "absolute":
+            position_embeddings = self.position_embeddings(position_ids)
+            embeddings += position_embeddings
+        embeddings = self.LayerNorm(embeddings)
+        embeddings = self.dropout(embeddings)
+        return embeddings
+
+
+class BertSelfAttention(nn.Module):
+    def __init__(self, config, is_cross_attention):
+        super().__init__()
+        self.config = config
+        if config.hidden_size % config.num_attention_heads != 0 and not hasattr(config, "embedding_size"):
+            raise ValueError(
+                "The hidden size (%d) is not a multiple of the number of attention "
+                "heads (%d)" % (config.hidden_size, config.num_attention_heads)
+            )
+        
+        self.num_attention_heads = config.num_attention_heads
+        self.attention_head_size = int(config.hidden_size / config.num_attention_heads)
+        self.all_head_size = self.num_attention_heads * self.attention_head_size
+
+        self.query = nn.Linear(config.hidden_size, self.all_head_size)
+        if is_cross_attention:
+            self.key = nn.Linear(config.encoder_width, self.all_head_size)
+            self.value = nn.Linear(config.encoder_width, self.all_head_size)
+        else:
+            self.key = nn.Linear(config.hidden_size, self.all_head_size)
+            self.value = nn.Linear(config.hidden_size, self.all_head_size)
+
+        self.dropout = nn.Dropout(config.attention_probs_dropout_prob)
+        self.position_embedding_type = getattr(config, "position_embedding_type", "absolute")
+        if self.position_embedding_type == "relative_key" or self.position_embedding_type == "relative_key_query":
+            self.max_position_embeddings = config.max_position_embeddings
+            self.distance_embedding = nn.Embedding(2 * config.max_position_embeddings - 1, self.attention_head_size)
+        self.save_attention = False   
+            
+    def save_attn_gradients(self, attn_gradients):
+        self.attn_gradients = attn_gradients
+        
+    def get_attn_gradients(self):
+        return self.attn_gradients
+    
+    def save_attention_map(self, attention_map):
+        self.attention_map = attention_map
+        
+    def get_attention_map(self):
+        return self.attention_map
+    
+    def transpose_for_scores(self, x):
+        new_x_shape = x.size()[:-1] + (self.num_attention_heads, self.attention_head_size)
+        x = x.view(*new_x_shape)
+        return x.permute(0, 2, 1, 3)
+
+    def forward(
+        self,
+        hidden_states,
+        attention_mask=None,
+        head_mask=None,
+        encoder_hidden_states=None,
+        encoder_attention_mask=None,
+        past_key_value=None,
+        output_attentions=False,
+    ):
+        mixed_query_layer = self.query(hidden_states)
+
+        # If this is instantiated as a cross-attention module, the keys
+        # and values come from an encoder; the attention mask needs to be
+        # such that the encoder's padding tokens are not attended to.
+        is_cross_attention = encoder_hidden_states is not None
+
+        if is_cross_attention:
+            key_layer = self.transpose_for_scores(self.key(encoder_hidden_states))
+            value_layer = self.transpose_for_scores(self.value(encoder_hidden_states))
+            attention_mask = encoder_attention_mask
+        elif past_key_value is not None:
+            key_layer = self.transpose_for_scores(self.key(hidden_states))
+            value_layer = self.transpose_for_scores(self.value(hidden_states))
+            key_layer = torch.cat([past_key_value[0], key_layer], dim=2)
+            value_layer = torch.cat([past_key_value[1], value_layer], dim=2)
+        else:
+            key_layer = self.transpose_for_scores(self.key(hidden_states))
+            value_layer = self.transpose_for_scores(self.value(hidden_states))
+
+        query_layer = self.transpose_for_scores(mixed_query_layer)
+
+        past_key_value = (key_layer, value_layer)
+
+        # Take the dot product between "query" and "key" to get the raw attention scores.
+        attention_scores = torch.matmul(query_layer, key_layer.transpose(-1, -2))
+
+        if self.position_embedding_type == "relative_key" or self.position_embedding_type == "relative_key_query":
+            seq_length = hidden_states.size()[1]
+            position_ids_l = torch.arange(seq_length, dtype=torch.long, device=hidden_states.device).view(-1, 1)
+            position_ids_r = torch.arange(seq_length, dtype=torch.long, device=hidden_states.device).view(1, -1)
+            distance = position_ids_l - position_ids_r
+            positional_embedding = self.distance_embedding(distance + self.max_position_embeddings - 1)
+            positional_embedding = positional_embedding.to(dtype=query_layer.dtype)  # fp16 compatibility
+
+            if self.position_embedding_type == "relative_key":
+                relative_position_scores = torch.einsum("bhld,lrd->bhlr", query_layer, positional_embedding)
+                attention_scores = attention_scores + relative_position_scores
+            elif self.position_embedding_type == "relative_key_query":
+                relative_position_scores_query = torch.einsum("bhld,lrd->bhlr", query_layer, positional_embedding)
+                relative_position_scores_key = torch.einsum("bhrd,lrd->bhlr", key_layer, positional_embedding)
+                attention_scores = attention_scores + relative_position_scores_query + relative_position_scores_key
+
+        attention_scores = attention_scores / math.sqrt(self.attention_head_size)
+        if attention_mask is not None:
+            # Apply the attention mask is (precomputed for all layers in BertModel forward() function)
+            attention_scores = attention_scores + attention_mask
+
+        # Normalize the attention scores to probabilities.
+        attention_probs = nn.Softmax(dim=-1)(attention_scores)
+        
+        if is_cross_attention and self.save_attention:
+            self.save_attention_map(attention_probs)
+            attention_probs.register_hook(self.save_attn_gradients)         
+
+        # This is actually dropping out entire tokens to attend to, which might
+        # seem a bit unusual, but is taken from the original Transformer paper.
+        attention_probs_dropped = self.dropout(attention_probs)
+
+        # Mask heads if we want to
+        if head_mask is not None:
+            attention_probs_dropped = attention_probs_dropped * head_mask
+
+        context_layer = torch.matmul(attention_probs_dropped, value_layer)
+
+        context_layer = context_layer.permute(0, 2, 1, 3).contiguous()
+        new_context_layer_shape = context_layer.size()[:-2] + (self.all_head_size,)
+        context_layer = context_layer.view(*new_context_layer_shape)
+
+        outputs = (context_layer, attention_probs) if output_attentions else (context_layer,)
+
+        outputs = outputs + (past_key_value,)
+        return outputs
+
+
+class BertSelfOutput(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.dense = nn.Linear(config.hidden_size, config.hidden_size)
+        self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+        self.dropout = nn.Dropout(config.hidden_dropout_prob)
+
+    def forward(self, hidden_states, input_tensor):
+        hidden_states = self.dense(hidden_states)
+        hidden_states = self.dropout(hidden_states)
+        hidden_states = self.LayerNorm(hidden_states + input_tensor)
+        return hidden_states
+
+
+class BertAttention(nn.Module):
+    def __init__(self, config, is_cross_attention=False):
+        super().__init__()
+        self.self = BertSelfAttention(config, is_cross_attention)
+        self.output = BertSelfOutput(config)
+        self.pruned_heads = set()
+
+    def prune_heads(self, heads):
+        if len(heads) == 0:
+            return
+        heads, index = find_pruneable_heads_and_indices(
+            heads, self.self.num_attention_heads, self.self.attention_head_size, self.pruned_heads
+        )
+
+        # Prune linear layers
+        self.self.query = prune_linear_layer(self.self.query, index)
+        self.self.key = prune_linear_layer(self.self.key, index)
+        self.self.value = prune_linear_layer(self.self.value, index)
+        self.output.dense = prune_linear_layer(self.output.dense, index, dim=1)
+
+        # Update hyper params and store pruned heads
+        self.self.num_attention_heads = self.self.num_attention_heads - len(heads)
+        self.self.all_head_size = self.self.attention_head_size * self.self.num_attention_heads
+        self.pruned_heads = self.pruned_heads.union(heads)
+
+    def forward(
+        self,
+        hidden_states,
+        attention_mask=None,
+        head_mask=None,
+        encoder_hidden_states=None,
+        encoder_attention_mask=None,
+        past_key_value=None,
+        output_attentions=False,
+    ):
+        self_outputs = self.self(
+            hidden_states,
+            attention_mask,
+            head_mask,
+            encoder_hidden_states,
+            encoder_attention_mask,
+            past_key_value,
+            output_attentions,
+        )
+        attention_output = self.output(self_outputs[0], hidden_states)
+        outputs = (attention_output,) + self_outputs[1:]  # add attentions if we output them
+        return outputs
+
+
+class BertIntermediate(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.dense = nn.Linear(config.hidden_size, config.intermediate_size)
+        if isinstance(config.hidden_act, str):
+            self.intermediate_act_fn = ACT2FN[config.hidden_act]
+        else:
+            self.intermediate_act_fn = config.hidden_act
+
+    def forward(self, hidden_states):
+        hidden_states = self.dense(hidden_states)
+        hidden_states = self.intermediate_act_fn(hidden_states)
+        return hidden_states
+
+
+class BertOutput(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.dense = nn.Linear(config.intermediate_size, config.hidden_size)
+        self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+        self.dropout = nn.Dropout(config.hidden_dropout_prob)
+
+    def forward(self, hidden_states, input_tensor):
+        hidden_states = self.dense(hidden_states)
+        hidden_states = self.dropout(hidden_states)
+        hidden_states = self.LayerNorm(hidden_states + input_tensor)
+        return hidden_states
+
+
+class BertLayer(nn.Module):
+    def __init__(self, config, layer_num):
+        super().__init__()
+        self.config = config
+        self.chunk_size_feed_forward = config.chunk_size_feed_forward
+        self.seq_len_dim = 1
+        self.attention = BertAttention(config)      
+        self.layer_num = layer_num          
+        if self.config.add_cross_attention:
+            self.crossattention = BertAttention(config, is_cross_attention=self.config.add_cross_attention)
+        self.intermediate = BertIntermediate(config)
+        self.output = BertOutput(config)
+
+    def forward(
+        self,
+        hidden_states,
+        attention_mask=None,
+        head_mask=None,
+        encoder_hidden_states=None,
+        encoder_attention_mask=None,
+        past_key_value=None,
+        output_attentions=False,
+        mode=None,
+    ):
+        # decoder uni-directional self-attention cached key/values tuple is at positions 1,2
+        self_attn_past_key_value = past_key_value[:2] if past_key_value is not None else None
+        self_attention_outputs = self.attention(
+            hidden_states,
+            attention_mask,
+            head_mask,
+            output_attentions=output_attentions,
+            past_key_value=self_attn_past_key_value,
+        )
+        attention_output = self_attention_outputs[0]
+
+        outputs = self_attention_outputs[1:-1]
+        present_key_value = self_attention_outputs[-1]
+
+        if mode=='multimodal':
+            assert encoder_hidden_states is not None, "encoder_hidden_states must be given for cross-attention layers"
+
+            cross_attention_outputs = self.crossattention(
+                attention_output,
+                attention_mask,
+                head_mask,
+                encoder_hidden_states,
+                encoder_attention_mask,
+                output_attentions=output_attentions,
+            )
+            attention_output = cross_attention_outputs[0]
+            outputs = outputs + cross_attention_outputs[1:-1]  # add cross attentions if we output attention weights                               
+        layer_output = apply_chunking_to_forward(
+            self.feed_forward_chunk, self.chunk_size_feed_forward, self.seq_len_dim, attention_output
+        )
+        outputs = (layer_output,) + outputs
+
+        outputs = outputs + (present_key_value,)
+
+        return outputs
+
+    def feed_forward_chunk(self, attention_output):
+        intermediate_output = self.intermediate(attention_output)
+        layer_output = self.output(intermediate_output, attention_output)
+        return layer_output
+
+
+class BertEncoder(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.config = config
+        self.layer = nn.ModuleList([BertLayer(config,i) for i in range(config.num_hidden_layers)])
+        self.gradient_checkpointing = False
+
+    def forward(
+        self,
+        hidden_states,
+        attention_mask=None,
+        head_mask=None,
+        encoder_hidden_states=None,
+        encoder_attention_mask=None,
+        past_key_values=None,
+        use_cache=None,
+        output_attentions=False,
+        output_hidden_states=False,
+        return_dict=True,
+        mode='multimodal',
+    ):
+        all_hidden_states = () if output_hidden_states else None
+        all_self_attentions = () if output_attentions else None
+        all_cross_attentions = () if output_attentions and self.config.add_cross_attention else None
+
+        next_decoder_cache = () if use_cache else None
+               
+        for i in range(self.config.num_hidden_layers):
+            layer_module = self.layer[i]
+            if output_hidden_states:
+                all_hidden_states = all_hidden_states + (hidden_states,)
+
+            layer_head_mask = head_mask[i] if head_mask is not None else None
+            past_key_value = past_key_values[i] if past_key_values is not None else None
+
+            if self.gradient_checkpointing and self.training:
+
+                if use_cache:
+                    logger.warn(
+                        "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..."
+                    )
+                    use_cache = False
+
+                def create_custom_forward(module):
+                    def custom_forward(*inputs):
+                        return module(*inputs, past_key_value, output_attentions)
+
+                    return custom_forward
+
+                layer_outputs = torch.utils.checkpoint.checkpoint(
+                    create_custom_forward(layer_module),
+                    hidden_states,
+                    attention_mask,
+                    layer_head_mask,
+                    encoder_hidden_states,
+                    encoder_attention_mask,
+                    mode=mode,
+                )
+            else:
+                layer_outputs = layer_module(
+                    hidden_states,
+                    attention_mask,
+                    layer_head_mask,
+                    encoder_hidden_states,
+                    encoder_attention_mask,
+                    past_key_value,
+                    output_attentions,
+                    mode=mode,
+                )
+
+            hidden_states = layer_outputs[0]
+            if use_cache:
+                next_decoder_cache += (layer_outputs[-1],)
+            if output_attentions:
+                all_self_attentions = all_self_attentions + (layer_outputs[1],)
+
+        if output_hidden_states:
+            all_hidden_states = all_hidden_states + (hidden_states,)
+
+        if not return_dict:
+            return tuple(
+                v
+                for v in [
+                    hidden_states,
+                    next_decoder_cache,
+                    all_hidden_states,
+                    all_self_attentions,
+                    all_cross_attentions,
+                ]
+                if v is not None
+            )
+        return BaseModelOutputWithPastAndCrossAttentions(
+            last_hidden_state=hidden_states,
+            past_key_values=next_decoder_cache,
+            hidden_states=all_hidden_states,
+            attentions=all_self_attentions,
+            cross_attentions=all_cross_attentions,
+        )
+
+
+class BertPooler(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.dense = nn.Linear(config.hidden_size, config.hidden_size)
+        self.activation = nn.Tanh()
+
+    def forward(self, hidden_states):
+        # We "pool" the model by simply taking the hidden state corresponding
+        # to the first token.
+        first_token_tensor = hidden_states[:, 0]
+        pooled_output = self.dense(first_token_tensor)
+        pooled_output = self.activation(pooled_output)
+        return pooled_output
+
+
+class BertPredictionHeadTransform(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.dense = nn.Linear(config.hidden_size, config.hidden_size)
+        if isinstance(config.hidden_act, str):
+            self.transform_act_fn = ACT2FN[config.hidden_act]
+        else:
+            self.transform_act_fn = config.hidden_act
+        self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+
+    def forward(self, hidden_states):
+        hidden_states = self.dense(hidden_states)
+        hidden_states = self.transform_act_fn(hidden_states)
+        hidden_states = self.LayerNorm(hidden_states)
+        return hidden_states
+
+
+class BertLMPredictionHead(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.transform = BertPredictionHeadTransform(config)
+
+        # The output weights are the same as the input embeddings, but there is
+        # an output-only bias for each token.
+        self.decoder = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+
+        self.bias = nn.Parameter(torch.zeros(config.vocab_size))
+
+        # Need a link between the two variables so that the bias is correctly resized with `resize_token_embeddings`
+        self.decoder.bias = self.bias
+
+    def forward(self, hidden_states):
+        hidden_states = self.transform(hidden_states)
+        hidden_states = self.decoder(hidden_states)
+        return hidden_states
+
+
+class BertOnlyMLMHead(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.predictions = BertLMPredictionHead(config)
+
+    def forward(self, sequence_output):
+        prediction_scores = self.predictions(sequence_output)
+        return prediction_scores
+
+
+class BertPreTrainedModel(PreTrainedModel):
+    """
+    An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained
+    models.
+    """
+
+    config_class = BertConfig
+    base_model_prefix = "bert"
+    _keys_to_ignore_on_load_missing = [r"position_ids"]
+
+    def _init_weights(self, module):
+        """ Initialize the weights """
+        if isinstance(module, (nn.Linear, nn.Embedding)):
+            # Slightly different from the TF version which uses truncated_normal for initialization
+            # cf https://github.com/pytorch/pytorch/pull/5617
+            module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
+        elif isinstance(module, nn.LayerNorm):
+            module.bias.data.zero_()
+            module.weight.data.fill_(1.0)
+        if isinstance(module, nn.Linear) and module.bias is not None:
+            module.bias.data.zero_()
+
+
+class BertModel(BertPreTrainedModel):
+    """
+    The model can behave as an encoder (with only self-attention) as well as a decoder, in which case a layer of
+    cross-attention is added between the self-attention layers, following the architecture described in `Attention is
+    all you need <https://arxiv.org/abs/1706.03762>`__ by Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit,
+    Llion Jones, Aidan N. Gomez, Lukasz Kaiser and Illia Polosukhin.
+    argument and :obj:`add_cross_attention` set to :obj:`True`; an :obj:`encoder_hidden_states` is then expected as an
+    input to the forward pass.
+    """
+
+    def __init__(self, config, add_pooling_layer=True):
+        super().__init__(config)
+        self.config = config
+
+        self.embeddings = BertEmbeddings(config)
+        
+        self.encoder = BertEncoder(config)
+
+        self.pooler = BertPooler(config) if add_pooling_layer else None
+
+        self.init_weights()
+ 
+
+    def get_input_embeddings(self):
+        return self.embeddings.word_embeddings
+
+    def set_input_embeddings(self, value):
+        self.embeddings.word_embeddings = value
+
+    def _prune_heads(self, heads_to_prune):
+        """
+        Prunes heads of the model. heads_to_prune: dict of {layer_num: list of heads to prune in this layer} See base
+        class PreTrainedModel
+        """
+        for layer, heads in heads_to_prune.items():
+            self.encoder.layer[layer].attention.prune_heads(heads)
+
+    
+    def get_extended_attention_mask(self, attention_mask: Tensor, input_shape: Tuple[int], device: device, is_decoder: bool) -> Tensor:
+        """
+        Makes broadcastable attention and causal masks so that future and masked tokens are ignored.
+
+        Arguments:
+            attention_mask (:obj:`torch.Tensor`):
+                Mask with ones indicating tokens to attend to, zeros for tokens to ignore.
+            input_shape (:obj:`Tuple[int]`):
+                The shape of the input to the model.
+            device: (:obj:`torch.device`):
+                The device of the input to the model.
+
+        Returns:
+            :obj:`torch.Tensor` The extended attention mask, with a the same dtype as :obj:`attention_mask.dtype`.
+        """
+        # We can provide a self-attention mask of dimensions [batch_size, from_seq_length, to_seq_length]
+        # ourselves in which case we just need to make it broadcastable to all heads.
+        if attention_mask.dim() == 3:
+            extended_attention_mask = attention_mask[:, None, :, :]
+        elif attention_mask.dim() == 2:
+            # Provided a padding mask of dimensions [batch_size, seq_length]
+            # - if the model is a decoder, apply a causal mask in addition to the padding mask
+            # - if the model is an encoder, make the mask broadcastable to [batch_size, num_heads, seq_length, seq_length]
+            if is_decoder:
+                batch_size, seq_length = input_shape
+
+                seq_ids = torch.arange(seq_length, device=device)
+                causal_mask = seq_ids[None, None, :].repeat(batch_size, seq_length, 1) <= seq_ids[None, :, None]
+                # in case past_key_values are used we need to add a prefix ones mask to the causal mask
+                # causal and attention masks must have same type with pytorch version < 1.3
+                causal_mask = causal_mask.to(attention_mask.dtype)
+   
+                if causal_mask.shape[1] < attention_mask.shape[1]:
+                    prefix_seq_len = attention_mask.shape[1] - causal_mask.shape[1]
+                    causal_mask = torch.cat(
+                        [
+                            torch.ones((batch_size, seq_length, prefix_seq_len), device=device, dtype=causal_mask.dtype),
+                            causal_mask,
+                        ],
+                        axis=-1,
+                    )                     
+
+                extended_attention_mask = causal_mask[:, None, :, :] * attention_mask[:, None, None, :]
+            else:
+                extended_attention_mask = attention_mask[:, None, None, :]
+        else:
+            raise ValueError(
+                "Wrong shape for input_ids (shape {}) or attention_mask (shape {})".format(
+                    input_shape, attention_mask.shape
+                )
+            )
+
+        # Since attention_mask is 1.0 for positions we want to attend and 0.0 for
+        # masked positions, this operation will create a tensor which is 0.0 for
+        # positions we want to attend and -10000.0 for masked positions.
+        # Since we are adding it to the raw scores before the softmax, this is
+        # effectively the same as removing these entirely.
+        extended_attention_mask = extended_attention_mask.to(dtype=self.dtype)  # fp16 compatibility
+        extended_attention_mask = (1.0 - extended_attention_mask) * -10000.0
+        return extended_attention_mask
+    
+    def forward(
+        self,
+        input_ids=None,
+        attention_mask=None,
+        position_ids=None,
+        head_mask=None,
+        inputs_embeds=None,
+        encoder_embeds=None,
+        encoder_hidden_states=None,
+        encoder_attention_mask=None,
+        past_key_values=None,
+        use_cache=None,
+        output_attentions=None,
+        output_hidden_states=None,
+        return_dict=None,
+        is_decoder=False,
+        mode='multimodal',
+    ):
+        r"""
+        encoder_hidden_states  (:obj:`torch.FloatTensor` of shape :obj:`(batch_size, sequence_length, hidden_size)`, `optional`):
+            Sequence of hidden-states at the output of the last layer of the encoder. Used in the cross-attention if
+            the model is configured as a decoder.
+        encoder_attention_mask (:obj:`torch.FloatTensor` of shape :obj:`(batch_size, sequence_length)`, `optional`):
+            Mask to avoid performing attention on the padding token indices of the encoder input. This mask is used in
+            the cross-attention if the model is configured as a decoder. Mask values selected in ``[0, 1]``:
+            - 1 for tokens that are **not masked**,
+            - 0 for tokens that are **masked**.
+        past_key_values (:obj:`tuple(tuple(torch.FloatTensor))` of length :obj:`config.n_layers` with each tuple having 4 tensors of shape :obj:`(batch_size, num_heads, sequence_length - 1, embed_size_per_head)`):
+            Contains precomputed key and value hidden states of the attention blocks. Can be used to speed up decoding.
+            If :obj:`past_key_values` are used, the user can optionally input only the last :obj:`decoder_input_ids`
+            (those that don't have their past key value states given to this model) of shape :obj:`(batch_size, 1)`
+            instead of all :obj:`decoder_input_ids` of shape :obj:`(batch_size, sequence_length)`.
+        use_cache (:obj:`bool`, `optional`):
+            If set to :obj:`True`, :obj:`past_key_values` key value states are returned and can be used to speed up
+            decoding (see :obj:`past_key_values`).
+        """
+        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
+        output_hidden_states = (
+            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+        )
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        if is_decoder:
+            use_cache = use_cache if use_cache is not None else self.config.use_cache
+        else:
+            use_cache = False
+
+        if input_ids is not None and inputs_embeds is not None:
+            raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")
+        elif input_ids is not None:
+            input_shape = input_ids.size()
+            batch_size, seq_length = input_shape
+            device = input_ids.device
+        elif inputs_embeds is not None:
+            input_shape = inputs_embeds.size()[:-1]
+            batch_size, seq_length = input_shape
+            device = inputs_embeds.device
+        elif encoder_embeds is not None:    
+            input_shape = encoder_embeds.size()[:-1]
+            batch_size, seq_length = input_shape 
+            device = encoder_embeds.device
+        else:
+            raise ValueError("You have to specify either input_ids or inputs_embeds or encoder_embeds")
+
+        # past_key_values_length
+        past_key_values_length = past_key_values[0][0].shape[2] if past_key_values is not None else 0
+
+        if attention_mask is None:
+            attention_mask = torch.ones(((batch_size, seq_length + past_key_values_length)), device=device)
+            
+        # We can provide a self-attention mask of dimensions [batch_size, from_seq_length, to_seq_length]
+        # ourselves in which case we just need to make it broadcastable to all heads.
+        extended_attention_mask: torch.Tensor = self.get_extended_attention_mask(attention_mask, input_shape, 
+                                                                                 device, is_decoder)
+
+        # If a 2D or 3D attention mask is provided for the cross-attention
+        # we need to make broadcastable to [batch_size, num_heads, seq_length, seq_length]
+        if encoder_hidden_states is not None:
+            if type(encoder_hidden_states) == list:
+                encoder_batch_size, encoder_sequence_length, _ = encoder_hidden_states[0].size()
+            else:
+                encoder_batch_size, encoder_sequence_length, _ = encoder_hidden_states.size()
+            encoder_hidden_shape = (encoder_batch_size, encoder_sequence_length)
+            
+            if type(encoder_attention_mask) == list:
+                encoder_extended_attention_mask = [self.invert_attention_mask(mask) for mask in encoder_attention_mask]
+            elif encoder_attention_mask is None:
+                encoder_attention_mask = torch.ones(encoder_hidden_shape, device=device)
+                encoder_extended_attention_mask = self.invert_attention_mask(encoder_attention_mask)
+            else:    
+                encoder_extended_attention_mask = self.invert_attention_mask(encoder_attention_mask)
+        else:
+            encoder_extended_attention_mask = None
+
+        # Prepare head mask if needed
+        # 1.0 in head_mask indicate we keep the head
+        # attention_probs has shape bsz x n_heads x N x N
+        # input head_mask has shape [num_heads] or [num_hidden_layers x num_heads]
+        # and head_mask is converted to shape [num_hidden_layers x batch x num_heads x seq_length x seq_length]
+        head_mask = self.get_head_mask(head_mask, self.config.num_hidden_layers)
+        
+        if encoder_embeds is None:
+            embedding_output = self.embeddings(
+                input_ids=input_ids,
+                position_ids=position_ids,
+                inputs_embeds=inputs_embeds,
+                past_key_values_length=past_key_values_length,
+            )
+        else:
+            embedding_output = encoder_embeds
+            
+        encoder_outputs = self.encoder(
+            embedding_output,
+            attention_mask=extended_attention_mask,
+            head_mask=head_mask,
+            encoder_hidden_states=encoder_hidden_states,
+            encoder_attention_mask=encoder_extended_attention_mask,
+            past_key_values=past_key_values,
+            use_cache=use_cache,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+            mode=mode,
+        )
+        sequence_output = encoder_outputs[0]
+        pooled_output = self.pooler(sequence_output) if self.pooler is not None else None
+
+        if not return_dict:
+            return (sequence_output, pooled_output) + encoder_outputs[1:]
+
+        return BaseModelOutputWithPoolingAndCrossAttentions(
+            last_hidden_state=sequence_output,
+            pooler_output=pooled_output,
+            past_key_values=encoder_outputs.past_key_values,
+            hidden_states=encoder_outputs.hidden_states,
+            attentions=encoder_outputs.attentions,
+            cross_attentions=encoder_outputs.cross_attentions,
+        )
+
+
+
+class BertLMHeadModel(BertPreTrainedModel):
+
+    _keys_to_ignore_on_load_unexpected = [r"pooler"]
+    _keys_to_ignore_on_load_missing = [r"position_ids", r"predictions.decoder.bias"]
+
+    def __init__(self, config):
+        super().__init__(config)
+
+        self.bert = BertModel(config, add_pooling_layer=False)
+        self.cls = BertOnlyMLMHead(config)
+
+        self.init_weights()
+
+    def get_output_embeddings(self):
+        return self.cls.predictions.decoder
+
+    def set_output_embeddings(self, new_embeddings):
+        self.cls.predictions.decoder = new_embeddings
+
+    def forward(
+        self,
+        input_ids=None,
+        attention_mask=None,
+        position_ids=None,
+        head_mask=None,
+        inputs_embeds=None,
+        encoder_hidden_states=None,
+        encoder_attention_mask=None,
+        labels=None,
+        past_key_values=None,
+        use_cache=None,
+        output_attentions=None,
+        output_hidden_states=None,
+        return_dict=None,
+        return_logits=False,            
+        is_decoder=True,
+        reduction='mean',
+        mode='multimodal', 
+    ):
+        r"""
+        encoder_hidden_states  (:obj:`torch.FloatTensor` of shape :obj:`(batch_size, sequence_length, hidden_size)`, `optional`):
+            Sequence of hidden-states at the output of the last layer of the encoder. Used in the cross-attention if
+            the model is configured as a decoder.
+        encoder_attention_mask (:obj:`torch.FloatTensor` of shape :obj:`(batch_size, sequence_length)`, `optional`):
+            Mask to avoid performing attention on the padding token indices of the encoder input. This mask is used in
+            the cross-attention if the model is configured as a decoder. Mask values selected in ``[0, 1]``:
+            - 1 for tokens that are **not masked**,
+            - 0 for tokens that are **masked**.
+        labels (:obj:`torch.LongTensor` of shape :obj:`(batch_size, sequence_length)`, `optional`):
+            Labels for computing the left-to-right language modeling loss (next word prediction). Indices should be in
+            ``[-100, 0, ..., config.vocab_size]`` (see ``input_ids`` docstring) Tokens with indices set to ``-100`` are
+            ignored (masked), the loss is only computed for the tokens with labels n ``[0, ..., config.vocab_size]``
+        past_key_values (:obj:`tuple(tuple(torch.FloatTensor))` of length :obj:`config.n_layers` with each tuple having 4 tensors of shape :obj:`(batch_size, num_heads, sequence_length - 1, embed_size_per_head)`):
+            Contains precomputed key and value hidden states of the attention blocks. Can be used to speed up decoding.
+            If :obj:`past_key_values` are used, the user can optionally input only the last :obj:`decoder_input_ids`
+            (those that don't have their past key value states given to this model) of shape :obj:`(batch_size, 1)`
+            instead of all :obj:`decoder_input_ids` of shape :obj:`(batch_size, sequence_length)`.
+        use_cache (:obj:`bool`, `optional`):
+            If set to :obj:`True`, :obj:`past_key_values` key value states are returned and can be used to speed up
+            decoding (see :obj:`past_key_values`).
+        Returns:
+        Example::
+            >>> from transformers import BertTokenizer, BertLMHeadModel, BertConfig
+            >>> import torch
+            >>> tokenizer = BertTokenizer.from_pretrained('bert-base-cased')
+            >>> config = BertConfig.from_pretrained("bert-base-cased")
+            >>> model = BertLMHeadModel.from_pretrained('bert-base-cased', config=config)
+            >>> inputs = tokenizer("Hello, my dog is cute", return_tensors="pt")
+            >>> outputs = model(**inputs)
+            >>> prediction_logits = outputs.logits
+        """
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+        if labels is not None:
+            use_cache = False
+
+        outputs = self.bert(
+            input_ids,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            head_mask=head_mask,
+            inputs_embeds=inputs_embeds,
+            encoder_hidden_states=encoder_hidden_states,
+            encoder_attention_mask=encoder_attention_mask,
+            past_key_values=past_key_values,
+            use_cache=use_cache,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+            is_decoder=is_decoder,
+            mode=mode,
+        )
+        
+        sequence_output = outputs[0]
+        prediction_scores = self.cls(sequence_output)
+        
+        if return_logits:
+            return prediction_scores[:, :-1, :].contiguous()  
+
+        lm_loss = None
+        if labels is not None:
+            # we are doing next-token prediction; shift prediction scores and input ids by one
+            shifted_prediction_scores = prediction_scores[:, :-1, :].contiguous()
+            labels = labels[:, 1:].contiguous()
+            loss_fct = CrossEntropyLoss(reduction=reduction, label_smoothing=0.1) 
+            lm_loss = loss_fct(shifted_prediction_scores.view(-1, self.config.vocab_size), labels.view(-1))
+            if reduction=='none':
+                lm_loss = lm_loss.view(prediction_scores.size(0),-1).sum(1)               
+
+        if not return_dict:
+            output = (prediction_scores,) + outputs[2:]
+            return ((lm_loss,) + output) if lm_loss is not None else output
+
+        return CausalLMOutputWithCrossAttentions(
+            loss=lm_loss,
+            logits=prediction_scores,
+            past_key_values=outputs.past_key_values,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+            cross_attentions=outputs.cross_attentions,
+        )
+
+    def prepare_inputs_for_generation(self, input_ids, past=None, attention_mask=None, **model_kwargs):
+        input_shape = input_ids.shape
+        # if model is used as a decoder in encoder-decoder model, the decoder attention mask is created on the fly
+        if attention_mask is None:
+            attention_mask = input_ids.new_ones(input_shape)
+
+        # cut decoder_input_ids if past is used
+        if past is not None:
+            input_ids = input_ids[:, -1:]
+
+        return {
+            "input_ids": input_ids, 
+            "attention_mask": attention_mask, 
+            "past_key_values": past,
+            "encoder_hidden_states": model_kwargs.get("encoder_hidden_states", None),
+            "encoder_attention_mask": model_kwargs.get("encoder_attention_mask", None),
+            "is_decoder": True,
+        }
+
+    def _reorder_cache(self, past, beam_idx):
+        reordered_past = ()
+        for layer_past in past:
+            reordered_past += (tuple(past_state.index_select(0, beam_idx) for past_state in layer_past),)
+        return reordered_past

ImageReward/models/BLIP/vit.py

@@ -0,0 +1,301 @@
+'''
+ * Adapted from BLIP (https://github.com/salesforce/BLIP)
+ * Based on timm code base
+ * https://github.com/rwightman/pytorch-image-models/tree/master/timm
+'''
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from functools import partial
+
+from timm.models.vision_transformer import _cfg, PatchEmbed
+from timm.models.registry import register_model
+from timm.models.layers import trunc_normal_, DropPath
+from timm.models.helpers import named_apply, adapt_input_conv
+
+from fairscale.nn.checkpoint.checkpoint_activations import checkpoint_wrapper
+
+class Mlp(nn.Module):
+    """ MLP as used in Vision Transformer, MLP-Mixer and related networks
+    """
+    def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.GELU, drop=0.):
+        super().__init__()
+        out_features = out_features or in_features
+        hidden_features = hidden_features or in_features
+        self.fc1 = nn.Linear(in_features, hidden_features)
+        self.act = act_layer()
+        self.fc2 = nn.Linear(hidden_features, out_features)
+        self.drop = nn.Dropout(drop)
+
+    def forward(self, x):
+        x = self.fc1(x)
+        x = self.act(x)
+        x = self.drop(x)
+        x = self.fc2(x)
+        x = self.drop(x)
+        return x
+
+
+class Attention(nn.Module):
+    def __init__(self, dim, num_heads=8, qkv_bias=False, qk_scale=None, attn_drop=0., proj_drop=0.):
+        super().__init__()
+        self.num_heads = num_heads
+        head_dim = dim // num_heads
+        # NOTE scale factor was wrong in my original version, can set manually to be compat with prev weights
+        self.scale = qk_scale or head_dim ** -0.5
+        self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
+        self.attn_drop = nn.Dropout(attn_drop)
+        self.proj = nn.Linear(dim, dim)
+        self.proj_drop = nn.Dropout(proj_drop)
+        self.attn_gradients = None
+        self.attention_map = None
+        
+    def save_attn_gradients(self, attn_gradients):
+        self.attn_gradients = attn_gradients
+        
+    def get_attn_gradients(self):
+        return self.attn_gradients
+    
+    def save_attention_map(self, attention_map):
+        self.attention_map = attention_map
+        
+    def get_attention_map(self):
+        return self.attention_map
+    
+    def forward(self, x, register_hook=False):
+        B, N, C = x.shape
+        qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)
+        q, k, v = qkv[0], qkv[1], qkv[2]   # make torchscript happy (cannot use tensor as tuple)
+
+        attn = (q @ k.transpose(-2, -1)) * self.scale
+        attn = attn.softmax(dim=-1)
+        attn = self.attn_drop(attn)
+                
+        if register_hook:
+            self.save_attention_map(attn)
+            attn.register_hook(self.save_attn_gradients)        
+
+        x = (attn @ v).transpose(1, 2).reshape(B, N, C)
+        x = self.proj(x)
+        x = self.proj_drop(x)
+        return x
+
+
+class Block(nn.Module):
+
+    def __init__(self, dim, num_heads, mlp_ratio=4., qkv_bias=False, qk_scale=None, drop=0., attn_drop=0.,
+                 drop_path=0., act_layer=nn.GELU, norm_layer=nn.LayerNorm, use_grad_checkpointing=False):
+        super().__init__()
+        self.norm1 = norm_layer(dim)
+        self.attn = Attention(
+            dim, num_heads=num_heads, qkv_bias=qkv_bias, qk_scale=qk_scale, attn_drop=attn_drop, proj_drop=drop)
+        # NOTE: drop path for stochastic depth, we shall see if this is better than dropout here
+        self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()
+        self.norm2 = norm_layer(dim)
+        mlp_hidden_dim = int(dim * mlp_ratio)
+        self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop)
+
+        if use_grad_checkpointing:
+            self.attn = checkpoint_wrapper(self.attn)
+            self.mlp = checkpoint_wrapper(self.mlp)
+
+    def forward(self, x, register_hook=False):
+        x = x + self.drop_path(self.attn(self.norm1(x), register_hook=register_hook))
+        x = x + self.drop_path(self.mlp(self.norm2(x)))
+        return x
+
+    
+class VisionTransformer(nn.Module):
+    """ Vision Transformer
+    A PyTorch impl of : `An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale`  -
+        https://arxiv.org/abs/2010.11929
+    """
+    def __init__(self, img_size=224, patch_size=16, in_chans=3, num_classes=1000, embed_dim=768, depth=12,
+                 num_heads=12, mlp_ratio=4., qkv_bias=True, qk_scale=None, representation_size=None,
+                 drop_rate=0., attn_drop_rate=0., drop_path_rate=0., norm_layer=None, 
+                 use_grad_checkpointing=False, ckpt_layer=0):
+        """
+        Args:
+            img_size (int, tuple): input image size
+            patch_size (int, tuple): patch size
+            in_chans (int): number of input channels
+            num_classes (int): number of classes for classification head
+            embed_dim (int): embedding dimension
+            depth (int): depth of transformer
+            num_heads (int): number of attention heads
+            mlp_ratio (int): ratio of mlp hidden dim to embedding dim
+            qkv_bias (bool): enable bias for qkv if True
+            qk_scale (float): override default qk scale of head_dim ** -0.5 if set
+            representation_size (Optional[int]): enable and set representation layer (pre-logits) to this value if set
+            drop_rate (float): dropout rate
+            attn_drop_rate (float): attention dropout rate
+            drop_path_rate (float): stochastic depth rate
+            norm_layer: (nn.Module): normalization layer
+        """
+        super().__init__()
+        self.num_features = self.embed_dim = embed_dim  # num_features for consistency with other models
+        norm_layer = norm_layer or partial(nn.LayerNorm, eps=1e-6)
+
+        self.patch_embed = PatchEmbed(
+            img_size=img_size, patch_size=patch_size, in_chans=in_chans, embed_dim=embed_dim)
+
+        num_patches = self.patch_embed.num_patches
+
+        self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
+        self.pos_embed = nn.Parameter(torch.zeros(1, num_patches + 1, embed_dim))
+        self.pos_drop = nn.Dropout(p=drop_rate)
+
+        dpr = [x.item() for x in torch.linspace(0, drop_path_rate, depth)]  # stochastic depth decay rule
+        self.blocks = nn.ModuleList([
+            Block(
+                dim=embed_dim, num_heads=num_heads, mlp_ratio=mlp_ratio, qkv_bias=qkv_bias, qk_scale=qk_scale,
+                drop=drop_rate, attn_drop=attn_drop_rate, drop_path=dpr[i], norm_layer=norm_layer,
+                use_grad_checkpointing=(use_grad_checkpointing and i>=depth-ckpt_layer)
+            )
+            for i in range(depth)])
+        self.norm = norm_layer(embed_dim)
+
+        trunc_normal_(self.pos_embed, std=.02)
+        trunc_normal_(self.cls_token, std=.02)
+        self.apply(self._init_weights)
+
+    def _init_weights(self, m):
+        if isinstance(m, nn.Linear):
+            trunc_normal_(m.weight, std=.02)
+            if isinstance(m, nn.Linear) and m.bias is not None:
+                nn.init.constant_(m.bias, 0)
+        elif isinstance(m, nn.LayerNorm):
+            nn.init.constant_(m.bias, 0)
+            nn.init.constant_(m.weight, 1.0)
+
+    @torch.jit.ignore
+    def no_weight_decay(self):
+        return {'pos_embed', 'cls_token'}
+
+    def forward(self, x, register_blk=-1):
+        B = x.shape[0]
+        x = self.patch_embed(x)
+
+        cls_tokens = self.cls_token.expand(B, -1, -1)  # stole cls_tokens impl from Phil Wang, thanks
+        x = torch.cat((cls_tokens, x), dim=1)
+  
+        x = x + self.pos_embed[:,:x.size(1),:]
+        x = self.pos_drop(x)
+
+        for i,blk in enumerate(self.blocks):
+            x = blk(x, register_blk==i)
+        x = self.norm(x)
+        
+        return x
+
+    @torch.jit.ignore()
+    def load_pretrained(self, checkpoint_path, prefix=''):
+        _load_weights(self, checkpoint_path, prefix)
+        
+
+@torch.no_grad()
+def _load_weights(model: VisionTransformer, checkpoint_path: str, prefix: str = ''):
+    """ Load weights from .npz checkpoints for official Google Brain Flax implementation
+    """
+    import numpy as np
+
+    def _n2p(w, t=True):
+        if w.ndim == 4 and w.shape[0] == w.shape[1] == w.shape[2] == 1:
+            w = w.flatten()
+        if t:
+            if w.ndim == 4:
+                w = w.transpose([3, 2, 0, 1])
+            elif w.ndim == 3:
+                w = w.transpose([2, 0, 1])
+            elif w.ndim == 2:
+                w = w.transpose([1, 0])
+        return torch.from_numpy(w)
+
+    w = np.load(checkpoint_path)
+    if not prefix and 'opt/target/embedding/kernel' in w:
+        prefix = 'opt/target/'
+
+    if hasattr(model.patch_embed, 'backbone'):
+        # hybrid
+        backbone = model.patch_embed.backbone
+        stem_only = not hasattr(backbone, 'stem')
+        stem = backbone if stem_only else backbone.stem
+        stem.conv.weight.copy_(adapt_input_conv(stem.conv.weight.shape[1], _n2p(w[f'{prefix}conv_root/kernel'])))
+        stem.norm.weight.copy_(_n2p(w[f'{prefix}gn_root/scale']))
+        stem.norm.bias.copy_(_n2p(w[f'{prefix}gn_root/bias']))
+        if not stem_only:
+            for i, stage in enumerate(backbone.stages):
+                for j, block in enumerate(stage.blocks):
+                    bp = f'{prefix}block{i + 1}/unit{j + 1}/'
+                    for r in range(3):
+                        getattr(block, f'conv{r + 1}').weight.copy_(_n2p(w[f'{bp}conv{r + 1}/kernel']))
+                        getattr(block, f'norm{r + 1}').weight.copy_(_n2p(w[f'{bp}gn{r + 1}/scale']))
+                        getattr(block, f'norm{r + 1}').bias.copy_(_n2p(w[f'{bp}gn{r + 1}/bias']))
+                    if block.downsample is not None:
+                        block.downsample.conv.weight.copy_(_n2p(w[f'{bp}conv_proj/kernel']))
+                        block.downsample.norm.weight.copy_(_n2p(w[f'{bp}gn_proj/scale']))
+                        block.downsample.norm.bias.copy_(_n2p(w[f'{bp}gn_proj/bias']))
+        embed_conv_w = _n2p(w[f'{prefix}embedding/kernel'])
+    else:
+        embed_conv_w = adapt_input_conv(
+            model.patch_embed.proj.weight.shape[1], _n2p(w[f'{prefix}embedding/kernel']))
+    model.patch_embed.proj.weight.copy_(embed_conv_w)
+    model.patch_embed.proj.bias.copy_(_n2p(w[f'{prefix}embedding/bias']))
+    model.cls_token.copy_(_n2p(w[f'{prefix}cls'], t=False))
+    pos_embed_w = _n2p(w[f'{prefix}Transformer/posembed_input/pos_embedding'], t=False)
+    if pos_embed_w.shape != model.pos_embed.shape:
+        pos_embed_w = resize_pos_embed(  # resize pos embedding when different size from pretrained weights
+            pos_embed_w, model.pos_embed, getattr(model, 'num_tokens', 1), model.patch_embed.grid_size)
+    model.pos_embed.copy_(pos_embed_w)
+    model.norm.weight.copy_(_n2p(w[f'{prefix}Transformer/encoder_norm/scale']))
+    model.norm.bias.copy_(_n2p(w[f'{prefix}Transformer/encoder_norm/bias']))
+#     if isinstance(model.head, nn.Linear) and model.head.bias.shape[0] == w[f'{prefix}head/bias'].shape[-1]:
+#         model.head.weight.copy_(_n2p(w[f'{prefix}head/kernel']))
+#         model.head.bias.copy_(_n2p(w[f'{prefix}head/bias']))
+#     if isinstance(getattr(model.pre_logits, 'fc', None), nn.Linear) and f'{prefix}pre_logits/bias' in w:
+#         model.pre_logits.fc.weight.copy_(_n2p(w[f'{prefix}pre_logits/kernel']))
+#         model.pre_logits.fc.bias.copy_(_n2p(w[f'{prefix}pre_logits/bias']))
+    for i, block in enumerate(model.blocks.children()):
+        block_prefix = f'{prefix}Transformer/encoderblock_{i}/'
+        mha_prefix = block_prefix + 'MultiHeadDotProductAttention_1/'
+        block.norm1.weight.copy_(_n2p(w[f'{block_prefix}LayerNorm_0/scale']))
+        block.norm1.bias.copy_(_n2p(w[f'{block_prefix}LayerNorm_0/bias']))
+        block.attn.qkv.weight.copy_(torch.cat([
+            _n2p(w[f'{mha_prefix}{n}/kernel'], t=False).flatten(1).T for n in ('query', 'key', 'value')]))
+        block.attn.qkv.bias.copy_(torch.cat([
+            _n2p(w[f'{mha_prefix}{n}/bias'], t=False).reshape(-1) for n in ('query', 'key', 'value')]))
+        block.attn.proj.weight.copy_(_n2p(w[f'{mha_prefix}out/kernel']).flatten(1))
+        block.attn.proj.bias.copy_(_n2p(w[f'{mha_prefix}out/bias']))
+        for r in range(2):
+            getattr(block.mlp, f'fc{r + 1}').weight.copy_(_n2p(w[f'{block_prefix}MlpBlock_3/Dense_{r}/kernel']))
+            getattr(block.mlp, f'fc{r + 1}').bias.copy_(_n2p(w[f'{block_prefix}MlpBlock_3/Dense_{r}/bias']))
+        block.norm2.weight.copy_(_n2p(w[f'{block_prefix}LayerNorm_2/scale']))
+        block.norm2.bias.copy_(_n2p(w[f'{block_prefix}LayerNorm_2/bias']))
+
+            
+def interpolate_pos_embed(pos_embed_checkpoint, visual_encoder):        
+    # interpolate position embedding
+    embedding_size = pos_embed_checkpoint.shape[-1]
+    num_patches = visual_encoder.patch_embed.num_patches
+    num_extra_tokens = visual_encoder.pos_embed.shape[-2] - num_patches
+    # height (== width) for the checkpoint position embedding
+    orig_size = int((pos_embed_checkpoint.shape[-2] - num_extra_tokens) ** 0.5)
+    # height (== width) for the new position embedding
+    new_size = int(num_patches ** 0.5)
+
+    if orig_size!=new_size:
+        # class_token and dist_token are kept unchanged
+        extra_tokens = pos_embed_checkpoint[:, :num_extra_tokens]
+        # only the position tokens are interpolated
+        pos_tokens = pos_embed_checkpoint[:, num_extra_tokens:]
+        pos_tokens = pos_tokens.reshape(-1, orig_size, orig_size, embedding_size).permute(0, 3, 1, 2)
+        pos_tokens = torch.nn.functional.interpolate(
+            pos_tokens, size=(new_size, new_size), mode='bicubic', align_corners=False)
+        pos_tokens = pos_tokens.permute(0, 2, 3, 1).flatten(1, 2)
+        new_pos_embed = torch.cat((extra_tokens, pos_tokens), dim=1)
+        print('reshape position embedding from %d to %d'%(orig_size ** 2,new_size ** 2))
+        
+        return new_pos_embed    
+    else:
+        return pos_embed_checkpoint

ImageReward/models/BLIPScore.py

@@ -0,0 +1,97 @@
+'''
+@File       :   BLIPScore.py
+@Time       :   2023/02/19 20:48:00
+@Auther     :   Jiazheng Xu
+@Contact    :   [email protected]
+@Description:   BLIPScore.
+* Based on BLIP code base
+* https://github.com/salesforce/BLIP
+'''
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from PIL import Image
+from ImageReward.models.BLIP.blip_pretrain import BLIP_Pretrain
+from torchvision.transforms import Compose, Resize, CenterCrop, ToTensor, Normalize
+
+try:
+    from torchvision.transforms import InterpolationMode
+    BICUBIC = InterpolationMode.BICUBIC
+except ImportError:
+    BICUBIC = Image.BICUBIC
+
+
+def _convert_image_to_rgb(image):
+    return image.convert("RGB")
+
+
+def _transform(n_px):
+    return Compose([
+        Resize(n_px, interpolation=BICUBIC),
+        CenterCrop(n_px),
+        _convert_image_to_rgb,
+        ToTensor(),
+        Normalize((0.48145466, 0.4578275, 0.40821073), (0.26862954, 0.26130258, 0.27577711)),
+    ])
+
+
+class BLIPScore(nn.Module):
+    def __init__(self, med_config, device='cpu'):
+        super().__init__()
+        self.device = device
+        
+        self.preprocess = _transform(224)
+        self.blip = BLIP_Pretrain(image_size=224, vit='large', med_config=med_config)
+
+
+    def score(self, prompt, image_path):
+        
+        if (type(image_path).__name__=='list'):
+            _, rewards = self.inference_rank(prompt, image_path)
+            return rewards
+            
+        # text encode
+        text_input = self.blip.tokenizer(prompt, padding='max_length', truncation=True, max_length=35, return_tensors="pt").to(self.device)
+        text_output = self.blip.text_encoder(text_input.input_ids, attention_mask = text_input.attention_mask, mode='text')  
+        txt_feature = F.normalize(self.blip.text_proj(text_output.last_hidden_state[:,0,:]))
+        
+        # image encode
+        pil_image = Image.open(image_path)
+        image = self.preprocess(pil_image).unsqueeze(0).to(self.device)
+        image_embeds = self.blip.visual_encoder(image)
+        image_features = F.normalize(self.blip.vision_proj(image_embeds[:,0,:]), dim=-1)    
+        
+        # score
+        rewards = torch.sum(torch.mul(txt_feature, image_features), dim=1, keepdim=True)
+        
+        return rewards.detach().cpu().numpy().item()
+
+
+    def inference_rank(self, prompt, generations_list):
+    
+        text_input = self.blip.tokenizer(prompt, padding='max_length', truncation=True, max_length=35, return_tensors="pt").to(self.device)
+        text_output = self.blip.text_encoder(text_input.input_ids, attention_mask = text_input.attention_mask, mode='text')  
+        txt_feature = F.normalize(self.blip.text_proj(text_output.last_hidden_state[:,0,:]))
+        
+        txt_set = []
+        img_set = []
+        for generations in generations_list:
+            # image encode
+            img_path = generations
+            pil_image = Image.open(img_path)
+            image = self.preprocess(pil_image).unsqueeze(0).to(self.device)
+            image_embeds = self.blip.visual_encoder(image)
+            image_features = F.normalize(self.blip.vision_proj(image_embeds[:,0,:]), dim=-1)    
+            img_set.append(image_features)
+            txt_set.append(txt_feature)
+            
+        txt_features = torch.cat(txt_set, 0).float() # [image_num, feature_dim]
+        img_features = torch.cat(img_set, 0).float() # [image_num, feature_dim]
+        rewards = torch.sum(torch.mul(txt_features, img_features), dim=1, keepdim=True)
+        rewards = torch.squeeze(rewards)
+        _, rank = torch.sort(rewards, dim=0, descending=True)
+        _, indices = torch.sort(rank, dim=0)
+        indices = indices + 1
+        
+        return indices.detach().cpu().numpy().tolist(), rewards.detach().cpu().numpy().tolist()

ImageReward/models/CLIPScore.py

@@ -0,0 +1,78 @@
+'''
+@File       :   CLIPScore.py
+@Time       :   2023/02/12 13:14:00
+@Auther     :   Jiazheng Xu
+@Contact    :   [email protected]
+@Description:   CLIPScore.
+* Based on CLIP code base
+* https://github.com/openai/CLIP
+'''
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from PIL import Image
+import clip
+
+class CLIPScore(nn.Module):
+    def __init__(self, download_root, device='cpu'):
+        super().__init__()
+        self.device = device
+        self.clip_model, self.preprocess = clip.load("ViT-L/14", device=self.device, jit=False, 
+                                                     download_root=download_root)
+        
+        if device == "cpu":
+            self.clip_model.float()
+        else:
+            clip.model.convert_weights(self.clip_model) # Actually this line is unnecessary since clip by default already on float16
+
+        # have clip.logit_scale require no grad.
+        self.clip_model.logit_scale.requires_grad_(False)
+
+
+    def score(self, prompt, image_path):
+        
+        if (type(image_path).__name__=='list'):
+            _, rewards = self.inference_rank(prompt, image_path)
+            return rewards
+            
+        # text encode
+        text = clip.tokenize(prompt, truncate=True).to(self.device)
+        txt_features = F.normalize(self.clip_model.encode_text(text))
+        
+        # image encode
+        pil_image = Image.open(image_path)
+        image = self.preprocess(pil_image).unsqueeze(0).to(self.device)
+        image_features = F.normalize(self.clip_model.encode_image(image))
+        
+        # score
+        rewards = torch.sum(torch.mul(txt_features, image_features), dim=1, keepdim=True)
+        
+        return rewards.detach().cpu().numpy().item()
+
+
+    def inference_rank(self, prompt, generations_list):
+        
+        text = clip.tokenize(prompt, truncate=True).to(self.device)
+        txt_feature = F.normalize(self.clip_model.encode_text(text))
+        
+        txt_set = []
+        img_set = []
+        for generations in generations_list:
+            # image encode
+            img_path = generations
+            pil_image = Image.open(img_path)
+            image = self.preprocess(pil_image).unsqueeze(0).to(self.device)
+            image_features = F.normalize(self.clip_model.encode_image(image))
+            img_set.append(image_features)
+            txt_set.append(txt_feature)
+            
+        txt_features = torch.cat(txt_set, 0).float() # [image_num, feature_dim]
+        img_features = torch.cat(img_set, 0).float() # [image_num, feature_dim]
+        rewards = torch.sum(torch.mul(txt_features, img_features), dim=1, keepdim=True)
+        rewards = torch.squeeze(rewards)
+        _, rank = torch.sort(rewards, dim=0, descending=True)
+        _, indices = torch.sort(rank, dim=0)
+        indices = indices + 1
+        
+        return indices.detach().cpu().numpy().tolist(), rewards.detach().cpu().numpy().tolist()

ImageReward/models/__init__.py

@@ -0,0 +1,4 @@
+from .AestheticScore import *
+from .BLIPScore import *
+from .CLIPScore import *
+from .BLIP import *

ImageReward/utils.py

@@ -0,0 +1,184 @@
+'''
+@File       :   utils.py
+@Time       :   2023/04/05 19:18:00
+@Auther     :   Jiazheng Xu
+@Contact    :   [email protected]
+* Based on CLIP code base
+* https://github.com/openai/CLIP
+* Checkpoint of CLIP/BLIP/Aesthetic are from:
+* https://github.com/openai/CLIP
+* https://github.com/salesforce/BLIP
+* https://github.com/christophschuhmann/improved-aesthetic-predictor
+'''
+
+import os
+import urllib
+from typing import Union, List
+import pathlib
+
+import torch
+from tqdm import tqdm
+from huggingface_hub import hf_hub_download
+
+from .ImageReward import ImageReward
+from .models.CLIPScore import CLIPScore
+from .models.BLIPScore import BLIPScore
+from .models.AestheticScore import AestheticScore
+
+_MODELS = {
+    "ImageReward-v1.0": "https://huggingface.co/THUDM/ImageReward/blob/main/ImageReward.pt",
+}
+
+
+def available_models() -> List[str]:
+    """Returns the names of available ImageReward models"""
+    return list(_MODELS.keys())
+
+
+def ImageReward_download(url: str, root: str):
+    os.makedirs(root, exist_ok=True)
+    filename = os.path.basename(url)
+    download_target = os.path.join(root, filename)
+    hf_hub_download(repo_id="THUDM/ImageReward", filename=filename, local_dir=root)
+    return download_target
+
+
+def load(name: str = "ImageReward-v1.0",
+         device: Union[str, torch.device] = "cuda" if torch.cuda.is_available() else "cpu",
+         download_root: str = None,
+         med_config_path: str = None):
+    """Load a ImageReward model
+
+    Parameters
+    ----------
+    name: str
+        A model name listed by `ImageReward.available_models()`, or the path to a model checkpoint containing the state_dict
+    device: Union[str, torch.device]
+        The device to put the loaded model
+    download_root: str
+        path to download the model files; by default, it uses "~/.cache/ImageReward"
+    med_config_path: str
+
+    Returns
+    -------
+    model : torch.nn.Module
+        The ImageReward model
+    """
+    if name in _MODELS:
+        download_root = download_root or "~/.cache/ImageReward"
+        download_root = pathlib.Path(download_root)
+        model_path = pathlib.Path(download_root) / 'ImageReward.pt'
+
+        if not model_path.exists():
+            model_path = ImageReward_download(_MODELS[name], root=download_root.as_posix())
+    elif os.path.isfile(name):
+        model_path = name
+    else:
+        raise RuntimeError(f"Model {name} not found; available models = {available_models()}")
+
+    print('-> load ImageReward model from %s' % model_path)
+    state_dict = torch.load(model_path, map_location='cpu')
+
+    # med_config
+    if med_config_path is None:
+        med_config_root = download_root or "~/.cache/ImageReward"
+        med_config_root = pathlib.Path(med_config_root)
+        med_config_path = med_config_root / 'med_config.json'
+
+        if not med_config_path.exists():
+            med_config_path = ImageReward_download("https://huggingface.co/THUDM/ImageReward/blob/main/med_config.json",
+                                                   root=med_config_root.as_posix())
+        print('-> load ImageReward med_config from %s' % med_config_path)
+
+    model = ImageReward(device=device, med_config=med_config_path).to(device)
+    msg = model.load_state_dict(state_dict, strict=False)
+    model.eval()
+
+    return model
+
+
+_SCORES = {
+    "CLIP": "https://openaipublic.azureedge.net/clip/models/b8cca3fd41ae0c99ba7e8951adf17d267cdb84cd88be6f7c2e0eca1737a03836/ViT-L-14.pt",
+    "BLIP": "https://storage.googleapis.com/sfr-vision-language-research/BLIP/models/model_large.pth",
+    "Aesthetic": "https://github.com/christophschuhmann/improved-aesthetic-predictor/raw/main/sac%2Blogos%2Bava1-l14-linearMSE.pth",
+}
+
+
+def available_scores() -> List[str]:
+    """Returns the names of available ImageReward scores"""
+    return list(_SCORES.keys())
+
+
+def _download(url: str, root: str):
+    os.makedirs(root, exist_ok=True)
+    filename = os.path.basename(url)
+
+    download_target = os.path.join(root, filename)
+
+    if os.path.exists(download_target) and not os.path.isfile(download_target):
+        raise RuntimeError(f"{download_target} exists and is not a regular file")
+
+    if os.path.isfile(download_target):
+        return download_target
+
+    with urllib.request.urlopen(url) as source, open(download_target, "wb") as output:
+        with tqdm(total=int(source.info().get("Content-Length")), ncols=80, unit='iB', unit_scale=True,
+                  unit_divisor=1024) as loop:
+            while True:
+                buffer = source.read(8192)
+                if not buffer:
+                    break
+
+                output.write(buffer)
+                loop.update(len(buffer))
+
+    return download_target
+
+
+def load_score(name: str = "CLIP", device: Union[str, torch.device] = "cuda" if torch.cuda.is_available() else "cpu",
+               download_root: str = None):
+    """Load a ImageReward model
+
+    Parameters
+    ----------
+    name : str
+        A model name listed by `ImageReward.available_models()`
+
+    device : Union[str, torch.device]
+        The device to put the loaded model
+
+    download_root: str
+        path to download the model files; by default, it uses "~/.cache/ImageReward"
+
+    Returns
+    -------
+    model : torch.nn.Module
+        The ImageReward model
+    """
+    model_download_root = download_root or os.path.expanduser("~/.cache/ImageReward")
+
+    if name in _SCORES:
+        model_path = _download(_SCORES[name], model_download_root)
+    else:
+        raise RuntimeError(f"Score {name} not found; available scores = {available_scores()}")
+
+    print('load checkpoint from %s' % model_path)
+    if name == "BLIP":
+        state_dict = torch.load(model_path, map_location='cpu')
+        med_config = ImageReward_download("https://huggingface.co/THUDM/ImageReward/blob/main/med_config.json",
+                                          model_download_root)
+        model = BLIPScore(med_config=med_config, device=device).to(device)
+        model.blip.load_state_dict(state_dict['model'], strict=False)
+    elif name == "CLIP":
+        model = CLIPScore(download_root=model_download_root, device=device).to(device)
+    elif name == "Aesthetic":
+        state_dict = torch.load(model_path, map_location='cpu')
+        model = AestheticScore(download_root=model_download_root, device=device).to(device)
+        model.mlp.load_state_dict(state_dict, strict=False)
+    else:
+        raise RuntimeError(f"Score {name} not found; available scores = {available_scores()}")
+
+    print("checkpoint loaded")
+    model.eval()
+
+    return model

Install.md

@@ -0,0 +1,66 @@
+## Installation
+
+Create a new conda environment:
+
+```shell
+conda create --name svgrender python=3.10
+conda activate svgrender
+```
+
+Install pytorch and the following libraries:
+
+```shell
+conda install pytorch==1.13.1 torchvision==0.14.1 torchaudio==0.13.1 pytorch-cuda=11.6 -c pytorch -c nvidia
+pip install hydra-core omegaconf 
+pip install freetype-py shapely svgutils
+pip install opencv-python scikit-image matplotlib visdom wandb BeautifulSoup4
+pip install triton numba
+pip install numpy scipy scikit-fmm einops timm fairscale=0.4.13
+pip install accelerate transformers safetensors datasets
+```
+
+Install LaMa:
+
+```shell
+pip install easydict scikit-learn pytorch_lightning webdataset 
+pip install albumentations==0.5.2
+pip install kornia==0.5.0
+pip install wldhx.yadisk-direct
+
+cd lama
+# download LaMa model weights
+# raw link(deprecated): curl -L $(yadisk-direct https://disk.yandex.ru/d/kHJkc7bs7mKIVA) -o big-lama.zip
+curl -O -L https://huggingface.co/xingxm/PyTorch-SVGRender-models/resolve/main/big-lama.zip
+unzip big-lama.zip
+```
+
+Install CLIP:
+
+```shell
+pip install ftfy regex tqdm
+pip install git+https://github.com/openai/CLIP.git
+```
+
+Install diffusers:
+
+```shell
+pip install diffusers==0.20.2
+```
+
+Install xformers (require `python=3.10`):
+
+```shell
+conda install xformers -c xformers
+```
+
+Install diffvg:
+
+```shell
+git clone https://github.com/BachiLi/diffvg.git
+cd diffvg
+git submodule update --init --recursive
+conda install -y -c anaconda cmake
+conda install -y -c conda-forge ffmpeg
+pip install svgwrite svgpathtools cssutils torch-tools
+python setup.py install
+```

LICENSE

@@ -0,0 +1,373 @@
+Mozilla Public License Version 2.0
+==================================
+
+1. Definitions
+--------------
+
+1.1. "Contributor"
+    means each individual or legal entity that creates, contributes to
+    the creation of, or owns Covered Software.
+
+1.2. "Contributor Version"
+    means the combination of the Contributions of others (if any) used
+    by a Contributor and that particular Contributor's Contribution.
+
+1.3. "Contribution"
+    means Covered Software of a particular Contributor.
+
+1.4. "Covered Software"
+    means Source Code Form to which the initial Contributor has attached
+    the notice in Exhibit A, the Executable Form of such Source Code
+    Form, and Modifications of such Source Code Form, in each case
+    including portions thereof.
+
+1.5. "Incompatible With Secondary Licenses"
+    means
+
+    (a) that the initial Contributor has attached the notice described
+        in Exhibit B to the Covered Software; or
+
+    (b) that the Covered Software was made available under the terms of
+        version 1.1 or earlier of the License, but not also under the
+        terms of a Secondary License.
+
+1.6. "Executable Form"
+    means any form of the work other than Source Code Form.
+
+1.7. "Larger Work"
+    means a work that combines Covered Software with other material, in
+    a separate file or files, that is not Covered Software.
+
+1.8. "License"
+    means this document.
+
+1.9. "Licensable"
+    means having the right to grant, to the maximum extent possible,
+    whether at the time of the initial grant or subsequently, any and
+    all of the rights conveyed by this License.
+
+1.10. "Modifications"
+    means any of the following:
+
+    (a) any file in Source Code Form that results from an addition to,
+        deletion from, or modification of the contents of Covered
+        Software; or
+
+    (b) any new file in Source Code Form that contains any Covered
+        Software.
+
+1.11. "Patent Claims" of a Contributor
+    means any patent claim(s), including without limitation, method,
+    process, and apparatus claims, in any patent Licensable by such
+    Contributor that would be infringed, but for the grant of the
+    License, by the making, using, selling, offering for sale, having
+    made, import, or transfer of either its Contributions or its
+    Contributor Version.
+
+1.12. "Secondary License"
+    means either the GNU General Public License, Version 2.0, the GNU
+    Lesser General Public License, Version 2.1, the GNU Affero General
+    Public License, Version 3.0, or any later versions of those
+    licenses.
+
+1.13. "Source Code Form"
+    means the form of the work preferred for making modifications.
+
+1.14. "You" (or "Your")
+    means an individual or a legal entity exercising rights under this
+    License. For legal entities, "You" includes any entity that
+    controls, is controlled by, or is under common control with You. For
+    purposes of this definition, "control" means (a) the power, direct
+    or indirect, to cause the direction or management of such entity,
+    whether by contract or otherwise, or (b) ownership of more than
+    fifty percent (50%) of the outstanding shares or beneficial
+    ownership of such entity.
+
+2. License Grants and Conditions
+--------------------------------
+
+2.1. Grants
+
+Each Contributor hereby grants You a world-wide, royalty-free,
+non-exclusive license:
+
+(a) under intellectual property rights (other than patent or trademark)
+    Licensable by such Contributor to use, reproduce, make available,
+    modify, display, perform, distribute, and otherwise exploit its
+    Contributions, either on an unmodified basis, with Modifications, or
+    as part of a Larger Work; and
+
+(b) under Patent Claims of such Contributor to make, use, sell, offer
+    for sale, have made, import, and otherwise transfer either its
+    Contributions or its Contributor Version.
+
+2.2. Effective Date
+
+The licenses granted in Section 2.1 with respect to any Contribution
+become effective for each Contribution on the date the Contributor first
+distributes such Contribution.
+
+2.3. Limitations on Grant Scope
+
+The licenses granted in this Section 2 are the only rights granted under
+this License. No additional rights or licenses will be implied from the
+distribution or licensing of Covered Software under this License.
+Notwithstanding Section 2.1(b) above, no patent license is granted by a
+Contributor:
+
+(a) for any code that a Contributor has removed from Covered Software;
+    or
+
+(b) for infringements caused by: (i) Your and any other third party's
+    modifications of Covered Software, or (ii) the combination of its
+    Contributions with other software (except as part of its Contributor
+    Version); or
+
+(c) under Patent Claims infringed by Covered Software in the absence of
+    its Contributions.
+
+This License does not grant any rights in the trademarks, service marks,
+or logos of any Contributor (except as may be necessary to comply with
+the notice requirements in Section 3.4).
+
+2.4. Subsequent Licenses
+
+No Contributor makes additional grants as a result of Your choice to
+distribute the Covered Software under a subsequent version of this
+License (see Section 10.2) or under the terms of a Secondary License (if
+permitted under the terms of Section 3.3).
+
+2.5. Representation
+
+Each Contributor represents that the Contributor believes its
+Contributions are its original creation(s) or it has sufficient rights
+to grant the rights to its Contributions conveyed by this License.
+
+2.6. Fair Use
+
+This License is not intended to limit any rights You have under
+applicable copyright doctrines of fair use, fair dealing, or other
+equivalents.
+
+2.7. Conditions
+
+Sections 3.1, 3.2, 3.3, and 3.4 are conditions of the licenses granted
+in Section 2.1.
+
+3. Responsibilities
+-------------------
+
+3.1. Distribution of Source Form
+
+All distribution of Covered Software in Source Code Form, including any
+Modifications that You create or to which You contribute, must be under
+the terms of this License. You must inform recipients that the Source
+Code Form of the Covered Software is governed by the terms of this
+License, and how they can obtain a copy of this License. You may not
+attempt to alter or restrict the recipients' rights in the Source Code
+Form.
+
+3.2. Distribution of Executable Form
+
+If You distribute Covered Software in Executable Form then:
+
+(a) such Covered Software must also be made available in Source Code
+    Form, as described in Section 3.1, and You must inform recipients of
+    the Executable Form how they can obtain a copy of such Source Code
+    Form by reasonable means in a timely manner, at a charge no more
+    than the cost of distribution to the recipient; and
+
+(b) You may distribute such Executable Form under the terms of this
+    License, or sublicense it under different terms, provided that the
+    license for the Executable Form does not attempt to limit or alter
+    the recipients' rights in the Source Code Form under this License.
+
+3.3. Distribution of a Larger Work
+
+You may create and distribute a Larger Work under terms of Your choice,
+provided that You also comply with the requirements of this License for
+the Covered Software. If the Larger Work is a combination of Covered
+Software with a work governed by one or more Secondary Licenses, and the
+Covered Software is not Incompatible With Secondary Licenses, this
+License permits You to additionally distribute such Covered Software
+under the terms of such Secondary License(s), so that the recipient of
+the Larger Work may, at their option, further distribute the Covered
+Software under the terms of either this License or such Secondary
+License(s).
+
+3.4. Notices
+
+You may not remove or alter the substance of any license notices
+(including copyright notices, patent notices, disclaimers of warranty,
+or limitations of liability) contained within the Source Code Form of
+the Covered Software, except that You may alter any license notices to
+the extent required to remedy known factual inaccuracies.
+
+3.5. Application of Additional Terms
+
+You may choose to offer, and to charge a fee for, warranty, support,
+indemnity or liability obligations to one or more recipients of Covered
+Software. However, You may do so only on Your own behalf, and not on
+behalf of any Contributor. You must make it absolutely clear that any
+such warranty, support, indemnity, or liability obligation is offered by
+You alone, and You hereby agree to indemnify every Contributor for any
+liability incurred by such Contributor as a result of warranty, support,
+indemnity or liability terms You offer. You may include additional
+disclaimers of warranty and limitations of liability specific to any
+jurisdiction.
+
+4. Inability to Comply Due to Statute or Regulation
+---------------------------------------------------
+
+If it is impossible for You to comply with any of the terms of this
+License with respect to some or all of the Covered Software due to
+statute, judicial order, or regulation then You must: (a) comply with
+the terms of this License to the maximum extent possible; and (b)
+describe the limitations and the code they affect. Such description must
+be placed in a text file included with all distributions of the Covered
+Software under this License. Except to the extent prohibited by statute
+or regulation, such description must be sufficiently detailed for a
+recipient of ordinary skill to be able to understand it.
+
+5. Termination
+--------------
+
+5.1. The rights granted under this License will terminate automatically
+if You fail to comply with any of its terms. However, if You become
+compliant, then the rights granted under this License from a particular
+Contributor are reinstated (a) provisionally, unless and until such
+Contributor explicitly and finally terminates Your grants, and (b) on an
+ongoing basis, if such Contributor fails to notify You of the
+non-compliance by some reasonable means prior to 60 days after You have
+come back into compliance. Moreover, Your grants from a particular
+Contributor are reinstated on an ongoing basis if such Contributor
+notifies You of the non-compliance by some reasonable means, this is the
+first time You have received notice of non-compliance with this License
+from such Contributor, and You become compliant prior to 30 days after
+Your receipt of the notice.
+
+5.2. If You initiate litigation against any entity by asserting a patent
+infringement claim (excluding declaratory judgment actions,
+counter-claims, and cross-claims) alleging that a Contributor Version
+directly or indirectly infringes any patent, then the rights granted to
+You by any and all Contributors for the Covered Software under Section
+2.1 of this License shall terminate.
+
+5.3. In the event of termination under Sections 5.1 or 5.2 above, all
+end user license agreements (excluding distributors and resellers) which
+have been validly granted by You or Your distributors under this License
+prior to termination shall survive termination.
+
+************************************************************************
+*                                                                      *
+*  6. Disclaimer of Warranty                                           *
+*  -------------------------                                           *
+*                                                                      *
+*  Covered Software is provided under this License on an "as is"       *
+*  basis, without warranty of any kind, either expressed, implied, or  *
+*  statutory, including, without limitation, warranties that the       *
+*  Covered Software is free of defects, merchantable, fit for a        *
+*  particular purpose or non-infringing. The entire risk as to the     *
+*  quality and performance of the Covered Software is with You.        *
+*  Should any Covered Software prove defective in any respect, You     *
+*  (not any Contributor) assume the cost of any necessary servicing,   *
+*  repair, or correction. This disclaimer of warranty constitutes an   *
+*  essential part of this License. No use of any Covered Software is   *
+*  authorized under this License except under this disclaimer.         *
+*                                                                      *
+************************************************************************
+
+************************************************************************
+*                                                                      *
+*  7. Limitation of Liability                                          *
+*  --------------------------                                          *
+*                                                                      *
+*  Under no circumstances and under no legal theory, whether tort      *
+*  (including negligence), contract, or otherwise, shall any           *
+*  Contributor, or anyone who distributes Covered Software as          *
+*  permitted above, be liable to You for any direct, indirect,         *
+*  special, incidental, or consequential damages of any character      *
+*  including, without limitation, damages for lost profits, loss of    *
+*  goodwill, work stoppage, computer failure or malfunction, or any    *
+*  and all other commercial damages or losses, even if such party      *
+*  shall have been informed of the possibility of such damages. This   *
+*  limitation of liability shall not apply to liability for death or   *
+*  personal injury resulting from such party's negligence to the       *
+*  extent applicable law prohibits such limitation. Some               *
+*  jurisdictions do not allow the exclusion or limitation of           *
+*  incidental or consequential damages, so this exclusion and          *
+*  limitation may not apply to You.                                    *
+*                                                                      *
+************************************************************************
+
+8. Litigation
+-------------
+
+Any litigation relating to this License may be brought only in the
+courts of a jurisdiction where the defendant maintains its principal
+place of business and such litigation shall be governed by laws of that
+jurisdiction, without reference to its conflict-of-law provisions.
+Nothing in this Section shall prevent a party's ability to bring
+cross-claims or counter-claims.
+
+9. Miscellaneous
+----------------
+
+This License represents the complete agreement concerning the subject
+matter hereof. If any provision of this License is held to be
+unenforceable, such provision shall be reformed only to the extent
+necessary to make it enforceable. Any law or regulation which provides
+that the language of a contract shall be construed against the drafter
+shall not be used to construe this License against a Contributor.
+
+10. Versions of the License
+---------------------------
+
+10.1. New Versions
+
+Mozilla Foundation is the license steward. Except as provided in Section
+10.3, no one other than the license steward has the right to modify or
+publish new versions of this License. Each version will be given a
+distinguishing version number.
+
+10.2. Effect of New Versions
+
+You may distribute the Covered Software under the terms of the version
+of the License under which You originally received the Covered Software,
+or under the terms of any subsequent version published by the license
+steward.
+
+10.3. Modified Versions
+
+If you create software not governed by this License, and you want to
+create a new license for such software, you may create and use a
+modified version of this License if you rename the license and remove
+any references to the name of the license steward (except to note that
+such modified license differs from this License).
+
+10.4. Distributing Source Code Form that is Incompatible With Secondary
+Licenses
+
+If You choose to distribute Source Code Form that is Incompatible With
+Secondary Licenses under the terms of this version of the License, the
+notice described in Exhibit B of this License must be attached.
+
+Exhibit A - Source Code Form License Notice
+-------------------------------------------
+
+  This Source Code Form is subject to the terms of the Mozilla Public
+  License, v. 2.0. If a copy of the MPL was not distributed with this
+  file, You can obtain one at http://mozilla.org/MPL/2.0/.
+
+If it is not possible or desirable to put the notice in a particular
+file, then You may include the notice in a location (such as a LICENSE
+file in a relevant directory) where a recipient would be likely to look
+for such a notice.
+
+You may add additional accurate notices of copyright ownership.
+
+Exhibit B - "Incompatible With Secondary Licenses" Notice
+---------------------------------------------------------
+
+  This Source Code Form is "Incompatible With Secondary Licenses", as
+  defined by the Mozilla Public License, v. 2.0.

README copy.md

@@ -0,0 +1,304 @@
+---
+title: SVGRender
+emoji: 💻
+colorFrom: gray
+colorTo: yellow
+sdk: gradio
+sdk_version: 4.20.1
+python_version: 3.10.12
+app_file: app.py
+pinned: false
+license: apache-2.0
+---
+
+<h1 id="ptsvg" align="center">Pytorch-SVGRender</h1>
+
+<p align="center">
+    <a href="https://www.python.org/"><img src="https://img.shields.io/badge/python-3.10-or?logo=python" alt="pyhton"></a>
+    <a href="http://mozilla.org/MPL/2.0/"><img src="https://img.shields.io/badge/license-MPL2.0-orange" alt="license"></a>
+    <a href="https://ximinng.github.io/PyTorch-SVGRender-project/"><img src="https://img.shields.io/badge/website-Gitpage-yellow" alt="website"></a>
+    <a href="https://pytorch-svgrender.readthedocs.io/en/latest/index.html"><img src="https://img.shields.io/badge/docs-readthedocs-purple" alt="docs"></a>
+</p>
+
+<div align="center">
+<img src="./assets/logo.png" style="width: 350px; height: 300px;" alt="Pytorch-SVGRender">
+<p><strong>Pytorch-SVGRender: </strong>The go-to library for differentiable rendering methods for SVG generation.</p>
+</div>
+<p align="center">
+    <a href="#recent-updates">Updates</a> •
+    <a href="#table-of-contents">Table of Contents</a> •
+    <a href="#installation">Installation</a> •
+    <a href="#quickstart">Quickstart</a> •
+    <a href="#faq">FAQ</a> •
+    <a href="#todo">TODO</a> •
+    <a href="#acknowledgement">Acknowledgment</a> •
+    <a href="#citation">Citation</a> •
+    <a href="#licence">Licence</a>
+</p>
+
+<h2 align="center">Recent Updates</h2>
+
+- [12/2023] 🔥 We open-sourced Pytorch-SVGRender V1.0.
+
+<h2 align="center">Table of Contents</h2>
+<p align="right"><a href="#ptsvg"><sup>▴ Back to top</sup></a></p>
+
+### 1. Image Vectorization
+
+- DiffVG: Differentiable Vector Graphics Rasterization for Editing and Learning (`SIGGRAPH 2020`)
+
+  [[Project]](https://people.csail.mit.edu/tzumao/diffvg/) [[Paper]](https://cseweb.ucsd.edu/~tzli/diffvg/diffvg.pdf) [[Code]](https://github.com/BachiLi/diffvg)
+
+  DiffVG is a differentiable rasterizer for 2D vector graphics. **This repository is heavily based on DiffVG.**
+
+- LIVE: Towards Layer-wise Image Vectorization (`CVPR 2022`)
+
+  [[Project]](https://ma-xu.github.io/LIVE/) [[Paper]](https://ma-xu.github.io/LIVE/index_files/CVPR22_LIVE_main.pdf) [[Code]](https://github.com/Picsart-AI-Research/LIVE-Layerwise-Image-Vectorization)
+
+- CLIPasso: Semantically-Aware Object Sketching (`SIGGRAPH 2022`)
+
+  [[Project]](https://clipasso.github.io/clipasso/) [[Paper]](https://arxiv.org/abs/2202.05822) [[Code]](https://github.com/yael-vinker/CLIPasso)
+
+- CLIPascene: Scene Sketching with Different Types and Levels of Abstraction (`ICCV 2023`)
+
+  [[Project]](https://clipascene.github.io/CLIPascene/) [[Paper]](https://arxiv.org/abs/2211.17256) [[Code]](https://github.com/yael-vinker/SceneSketch)
+
+### 2. Text-to-SVG Synthesis
+
+- CLIPDraw: Exploring Text-to-Drawing Synthesis through Language-Image Encoders (`NIPS 2022`)
+
+  [[Paper]](https://arxiv.org/abs/2106.14843) [[Code]](https://github.com/kvfrans/clipdraw)
+
+- StyleCLIPDraw: Coupling Content and Style in Text-to-Drawing Synthesis
+
+  [[Live]](https://slideslive.com/38970834/styleclipdraw-coupling-content-and-style-in-texttodrawing-synthesis?ref=account-folder-92044-folders) [[Paper]](https://arxiv.org/abs/2202.12362) [[Code]](https://github.com/pschaldenbrand/StyleCLIPDraw)
+
+- CLIPFont: Texture Guided Vector WordArt Generation (`BMVC 2022`)
+
+  [[Paper]](https://bmvc2022.mpi-inf.mpg.de/0543.pdf) [[Code]](https://github.com/songyiren98/CLIPFont)
+
+- VectorFusion: Text-to-SVG by Abstracting Pixel-Based Diffusion Models (`CVPR 2023`)
+
+  [[Project]](https://vectorfusion.github.io/) [[Paper]](https://openaccess.thecvf.com/content/CVPR2023/papers/Jain_VectorFusion_Text-to-SVG_by_Abstracting_Pixel-Based_Diffusion_Models_CVPR_2023_paper.pdf)
+
+- DiffSketcher: Text Guided Vector Sketch Synthesis through Latent Diffusion Models (`NIPS 2023`)
+
+  [[Project]](https://ximinng.github.io/DiffSketcher-project/) [[Live]](https://neurips.cc/virtual/2023/poster/72425) [[Paper]](https://arxiv.org/abs/2306.14685) [[Code]](https://github.com/ximinng/DiffSketcher)
+
+- Word-As-Image for Semantic Typography (`SIGGRAPH 2023`)
+
+  [[Project]](https://wordasimage.github.io/Word-As-Image-Page/) [[Paper]](https://arxiv.org/abs/2303.01818) [[Code]](https://github.com/Shiriluz/Word-As-Image)
+
+- SVGDreamer: Text Guided SVG Generation with Diffusion Model (`CVPR 2024`)
+
+  [[Project]](https://ximinng.github.io/SVGDreamer-project/) [[Paper]](https://arxiv.org/abs/2312.16476) [[code]](https://github.com/ximinng/SVGDreamer)
+
+<h2 align="center">Installation</h2>
+
+You can follow the steps below to quickly get up and running with PyTorch-SVGRender.
+These steps will let you run quick inference locally.
+
+In the top level directory run,
+
+```bash
+sh script/install.sh
+```
+
+Note: Make sure that the script file has execution **permissions** (you can give them using `chmod +x script.sh`), and
+then run the script.
+
+For more information, please refer to
+the [Install.md](https://github.com/ximinng/PyTorch-SVGRender/blob/main/Install.md).
+
+<h2 align="center">Quickstart</h2>
+<p align="right"><a href="#ptsvg"><sup>▴ Back to top</sup></a></p>
+
+**For more information, [read the docs](https://pytorch-svgrender.readthedocs.io/en/latest/index.html).**
+
+### 1. Basic Usage
+
+**DiffVG** vectorizes any raster images:
+
+```shell
+python svg_render.py x=diffvg target='./data/fallingwater.png'
+# change 'num_paths' and 'num_iter' for better results
+python svg_render.py x=diffvg target='./data/fallingwater.png' x.num_paths=512 x.num_iter=2000
+```
+
+**LIVE** vectorizes the raster emojis images (in original PNG format):
+
+```shell
+python svg_render.py x=live target='./data/simile.png'
+# change 'num_paths' and 'schedule_each' for better results
+python svg_render.py x=live target='./data/simile.png' x.num_paths=5 x.schedule_each=1
+```
+
+**CLIPasso** synthesizes vectorized sketches from images:
+
+**note:** first download the U2Net model `sh script/download_u2net.sh`.
+
+```shell
+python svg_render.py x=clipasso target='./data/horse.png'
+```
+
+**CLIPascene** synthesizes vectorized sketches from images:
+
+**note:** first download the U2Net model `sh script/download_u2net.sh`, and make sure the `./data/background` folder and
+the `./data/scene` folder exist with target images.
+
+```shell
+python svg_render.py x=clipascene target='ballerina.png'
+```
+
+**CLIPDraw** synthesizes SVGs based on text prompts:
+
+```shell
+python svg_render.py x=clipdraw "prompt='a photo of a cat'"
+```
+
+**StyleCLIPDraw** synthesizes SVG based on a text prompt and a reference image:
+
+```shell
+python svg_render.py x=styleclipdraw "prompt='a photo of a cat'" target='./data/starry.png'
+```
+
+**CLIPFont** styles vector fonts according to text prompts:
+
+```shell
+python svg_render.py x=clipfont "prompt='Starry Night by Vincent van gogh'" target='./data/alphabet1.svg'
+```
+
+---
+
+> Because the following methods rely on stable diffusion, add `diffuser.download=True` to the command the **first time** you
+run the script.
+
+**SVGDreamer** generates various styles of SVG based on text prompts. It supports the use of six vector primitives,
+including Iconography, Sketch, Pixel Art, Low-Poly, Painting, and Ink and Wash.
+
+```shell
+# primitive: iconography
+## 1. German shepherd
+python svg_render.py x=svgdreamer "prompt='A colorful German shepherd in vector art. tending on artstation.'" save_step=30 x.guidance.n_particle=6 x.guidance.vsd_n_particle=4 x.guidance.phi_n_particle=2 result_path='./svgdreamer/GermanShepherd'
+## 2. sydney opera house
+python svg_render.py x=svgdreamer "prompt='Sydney opera house. oil painting. by Van Gogh'" save_step=30 x.guidance.n_particle=6 x.guidance.vsd_n_particle=4 x.guidance.phi_n_particle=2 x.num_paths=512 result_path='./svgdreamer/Sydney'
+# primitive: low-ploy
+python svg_render.py x=svgdreamer "prompt='A picture of a bald eagle. low-ploy. polygon'" x.style='low-poly' save_step=30 x.guidance.n_particle=6 x.guidance.vsd_n_particle=4 x.guidance.phi_n_particle=2 x.guidance.num_iter=1000 result_path='./svgdreamer/eagle'
+# primitive: pixel-art
+python svg_render.py x=svgdreamer "prompt='Darth vader with lightsaber. ultrarealistic. pixelart. trending on artstation.'" x.style='pixelart' save_step=30 x.guidance.n_particle=6 x.guidance.vsd_n_particle=4 x.guidance.phi_n_particle=2 x.guidance.num_iter=1000 result_path='./svgdreamer/DarthVader'
+# primitive: painting
+python svg_render.py x=svgdreamer "prompt='self portrait of Van Gogh. oil painting. cmyk portrait. multi colored. defiant and beautiful. cmyk. expressive eyes.'" x.style='painting' save_step=50 x.guidance.n_particle=6 x.guidance.vsd_n_particle=4 x.guidance.phi_n_particle=2 x.guidance.t_schedule='randint' x.num_paths=1500 result_path='./svgdreamer/VanGogh_portrait'
+# primitive: sketch
+python svg_render.py x=svgdreamer "prompt='A free-hand drawing of A speeding Lamborghini. black and white drawing.'" x.style='sketch' save_step=30 x.guidance.n_particle=6 x.guidance.vsd_n_particle=4 x.guidance.phi_n_particle=2 x.guidance.t_schedule='randint' x.num_paths=128 result_path='./svgdreamer/Lamborghini'
+# primitive: ink and wash
+python svg_render.py x=svgdreamer "prompt='Big Wild Goose Pagoda. ink style. Minimalist abstract art grayscale watercolor.'" x.style='ink' save_step=30 x.guidance.n_particle=6 x.guidance.vsd_n_particle=4 x.guidance.phi_n_particle=2 x.guidance.t_schedule='randint' x.num_paths=128 x.width=6 result_path='./svgdreamer/BigWildGoosePagoda'
+```
+
+**VectorFusion** synthesizes SVGs in various styles based on text prompts:
+
+```shell
+# Iconography style
+python svg_render.py x=vectorfusion x.style='iconography' "prompt='a panda rowing a boat in a pond. minimal flat 2d vector icon. lineal color. trending on artstation.'"
+# PixelArt style
+python svg_render.py x=vectorfusion x.style='pixelart' "prompt='a panda rowing a boat in a pond. pixel art. trending on artstation.'"
+# Sketch style
+python svg_render.py x=vectorfusion x.style='sketch' "prompt='a panda rowing a boat in a pond. minimal 2d line drawing. trending on artstation.'"
+```
+
+Following SVGDreamer, we've added three additional styles (`Paining`, `Ink and Wash` and `low-ploy`) to VectorFusion.
+
+**DiffSketcher** synthesizes vector sketches based on text prompts:
+
+```shell
+# DiffSketcher
+python svg_render.py x=diffsketcher "prompt='a photo of Sydney opera house'" x.token_ind=5 seed=8019
+# DiffSketcher, variable stroke width
+python svg_render.py x=diffsketcher "prompt='a photo of Sydney opera house'" x.token_ind=5 x.optim_width=True seed=8019
+# DiffSketcher RGBA version
+python svg_render.py x=diffsketcher "prompt='a photo of Sydney opera house'" x.token_ind=5 x.optim_width=True x.optim_rgba=True x.optim_opacity=False seed=8019
+# DiffSketcher + style transfer
+python svg_render.py x=stylediffsketcher "prompt='The French Revolution. highly detailed. 8k. ornate. intricate. cinematic. dehazed. atmospheric. oil painting. by Van Gogh'" x.token_ind=4 x.num_paths=2000 target='./data/starry.png' seed=876809
+```
+
+**Word-As-Image** follow a text prompt to style a letter in a word:
+
+```shell
+# Inject the meaning of the word bunny into the 'Y' in the word 'BUNNY'
+python svg_render.py x=wordasimage x.word='BUNNY' prompt='BUNNY' x.optim_letter='Y'
+```
+
+### 2. SDS Loss based Approach
+
+This is achieved by utilizing a pretrained text-to-image diffusion model as a strong image prior to supervise the
+training of the PyDiffVG, enabling rendering SVG alignment with the text. This remarkable capability is fundamentally
+grounded in the use of Score Distillation Sampling (SDS). SDS acts as the core mechanism that lifts raster images from
+diffusion models to the SVG domain, enabling the training of SVG parameters without images.
+This includes the methods VectorFusion, DiffSketcher and SVGDreamer.
+
+We only compare the performance of SDS, which means that no other loss is used:
+
+```shell
+# SDS loss
+python svg_render.py x=vectorfusion "prompt='a panda rowing a boat in a pond. minimal flat 2d vector icon. lineal color. trending on artstation.'"
+# Input Augmentation SDS loss (LSDS loss)
+python svg_render.py x=vectorfusion x.style='sketch' "prompt='an elephant. minimal 2d line drawing. trending on artstation.'"
+# Input Augmentation SDS loss (ASDS loss)
+python svg_render.py x=diffsketcher "prompt='an elephant. minimal 2d line drawing. trending on artstation.'" x.token_ind=2 x.sds.grad_scale=1 x.sds.num_aug=4 x.clip.vis_loss=0 x.perceptual.coeff=0 x.opacity_delta=0.3 
+# Vectorized Particle-based Score Distillation (VPSD loss)
+python svg_render.py x=svgdreamer "prompt='a panda rowing a boat in a pond. minimal flat 2d vector icon. lineal color. trending on artstation.'" save_step=60 x.guidance.n_particle=6 x.guidance.vsd_n_particle=4 x.guidance.phi_n_particle=2 
+```
+
+<h2 align="center">FAQ</h2>
+<p align="right"><a href="#ptsvg"><sup>▴ Back to top</sup></a></p>
+
+- Q: Where can I get more scripts and visualizations?
+- A: check the [pytorch-svgrender.readthedocs.io](https://pytorch-svgrender.readthedocs.io/en/latest/index.html).
+
+- Q: An error says HuggingFace cannot find the model in the disk cache.
+- A: Add *`diffuser.download=True`* to the command for downloading model checkpoints the **first time** you run the script.
+
+<h2 align="center">TODO</h2>
+<p align="right"><a href="#ptsvg"><sup>▴ Back to top</sup></a></p>
+
+- [x] integrated SVGDreamer.
+
+<h2 align="center">Acknowledgement</h2>
+<p align="right"><a href="#ptsvg"><sup>▴ Back to top</sup></a></p>
+
+The project is built based on the following repository:
+
+[BachiLi/diffvg](https://github.com/BachiLi/diffvg),
+[huggingface/diffusers](https://github.com/huggingface/diffusers),
+[threestudio-project/threestudio](https://github.com/threestudio-project/threestudio),
+[yael-vinker/CLIPasso](https://github.com/yael-vinker/CLIPasso),
+[ximinng/DiffSketcher](https://github.com/ximinng/DiffSketcher),
+[THUDM/ImageReward](https://github.com/THUDM/ImageReward),
+[advimman/lama](https://github.com/advimman/lama)
+
+We gratefully thank the authors for their wonderful works.
+
+<h2 align="center">Citation</h2>
+<p align="right"><a href="#ptsvg"><sup>▴ Back to top</sup></a></p>
+
+If you use this code for your research, please cite the following work:
+
+```
+@article{xing2023svgdreamer,
+    title={SVGDreamer: Text Guided SVG Generation with Diffusion Model},
+    author={Xing, Ximing and Zhou, Haitao and Wang, Chuang and Zhang, Jing and Xu, Dong and Yu, Qian},
+    journal={arXiv preprint arXiv:2312.16476},
+    year={2023}
+}
+@inproceedings{xing2023diffsketcher,
+    title={DiffSketcher: Text Guided Vector Sketch Synthesis through Latent Diffusion Models},
+    author={XiMing Xing and Chuang Wang and Haitao Zhou and Jing Zhang and Qian Yu and Dong Xu},
+    booktitle={Thirty-seventh Conference on Neural Information Processing Systems (NeurIPS)},
+    year={2023},
+    url={https://openreview.net/forum?id=CY1xatvEQj}
+}
+```
+
+<h2 align="center">Licence</h2>
+<p align="right"><a href="#ptsvg"><sup>▴ Back to top</sup></a></p>
+
+This work is licensed under a **Mozilla Public License Version 2.0**.

app.py

@@ -0,0 +1,83 @@
+import os
+import subprocess
+import sys
+import tempfile
+
+import gradio as gr
+from PIL import Image
+
+sys.path.append('/home/user/app/code')
+
+# set up diffvg
+# os.system('git clone https://github.com/BachiLi/diffvg.git')
+os.chdir('diffvg')
+os.system('git submodule update --init --recursive')
+os.system('python setup.py install --user')
+sys.path.append("/home/user/.local/lib/python3.10/site-packages/diffvg-0.0.1-py3.10-linux-x86_64.egg")
+print("diffvg installed.")
+os.chdir('/home/user/app')
+
+
+def process_images(prompt, num_paths, token_index, seed, optimize_width=False, optimize_color=False):
+    with tempfile.TemporaryDirectory() as tmpdirname:
+        command = [
+            "python", "svg_render.py",
+            "x=diffsketcher",
+            f"prompt={prompt}",
+            f"x.num_paths={num_paths}",
+            f"x.token_ind={token_index}",
+            f"seed={seed}",
+            f"x.optim_width={optimize_width}",
+            f"x.optim_rgba={optimize_color}",
+            "x.optim_opacity=False",
+        ]
+        result = subprocess.run(command, check=True)
+        if result.returncode == 0:
+            output_image = Image.open(os.path.join(tmpdirname, "final_render.png"))
+    return output_image
+
+
+with gr.Blocks() as demo:
+    gr.Markdown("# DiffSketcher")
+    gr.Markdown("DiffSketcher synthesizes **vector sketches** based on **text prompts**.")
+    li = [
+        "https://raw.githubusercontent.com/ximinng/DiffSketcher/main/img/cat.svg",
+        "https://raw.githubusercontent.com/ximinng/DiffSketcher/main/img/rose.svg",
+        "https://raw.githubusercontent.com/ximinng/DiffSketcher/main/img/elephant.svg",
+        "https://raw.githubusercontent.com/ximinng/DiffSketcher/main/img/elephant_silhouette.svg",
+        "https://raw.githubusercontent.com/ximinng/DiffSketcher/main/img/horse_width.svg",
+        "https://raw.githubusercontent.com/ximinng/DiffSketcher/main/img/horse_rgba.svg",
+        "https://ximinng.github.io/PyTorch-SVGRender-project/assets/diffsketcher/Sydney_opera.svg",
+        "https://ximinng.github.io/PyTorch-SVGRender-project/assets/diffsketcher/Sydney_opera_width.svg",
+        "https://ximinng.github.io/PyTorch-SVGRender-project/assets/diffsketcher/Sydney_opera_width_color.svg",
+    ]
+    gr.Gallery(li, columns=6)
+    with gr.Row():
+        with gr.Column():
+            text = gr.Textbox(label="prompt")
+            num_paths = gr.Slider(label="path number", value=96, minimum=1, maximum=500, step=1)
+            token_index = gr.Textbox(label="token_index", info="CLIP embedding token index. Starting from 1.")
+            seed = gr.Slider(0, 10000, label="random seed", value=8019)
+            with gr.Accordion("Selectable Inputs"):
+                optimize_width = gr.Checkbox(label="optimize stroke width")
+                optimize_color = gr.Checkbox(label="optimize stroke color")
+            btn = gr.Button("Synthesize")
+        with gr.Column():
+            output = gr.Image(label="output image", height=512)
+    btn.click(process_images,
+              inputs=[text, num_paths, token_index, seed, optimize_width, optimize_color],
+              outputs=[output])
+    gr.Markdown("## Examples")
+    gr.Markdown("Here are some config examples. Feel free to try your own prompts!")
+    gr.Examples(
+        inputs=[text, num_paths, token_index, seed, optimize_width, optimize_color],
+        outputs=[output],
+        fn=process_images,
+        examples=[
+            ["A photo of Sydney opera house.", 96, 5, 8019, False, False],
+            ["A photo of Sydney opera house.", 96, 5, 8019, True, False],
+            ["A photo of Sydney opera house.", 128, 5, 8019, True, True],
+        ],
+    )
+
+demo.launch()

checkpoint/placeholder.md

@@ -0,0 +1 @@
+**place model here**

conf/config.yaml

@@ -0,0 +1,56 @@
+#-----------------#
+#  Global Config  #
+#-----------------#
+
+# optional args
+target: ~
+prompt: ~
+neg_prompt: ~ # negative prompt
+
+# Accelerate config
+state:
+  cpu: False # use cpu
+  mprec: no # mixed precision, choices: 'no', 'fp16', 'bf16'
+#  wandb: False
+#  tensorboard: False
+
+# Diffusers config
+diffuser:
+  download: True # Set this variable to True the first time it runs
+  force_download: False
+  resume_download: False
+
+# PyDiffVG config
+diffvg:
+  print_timing: False
+
+# reproduction
+seed: 951222
+# multi-run
+multirun: False
+srange: ~ # seed range, example: [100, 100]
+
+# log
+result_path: './workspace'
+save_step: 10
+eval_step: 10
+
+# visual rendering process
+mv: False # make video
+framefreq: 5 # save the image interval
+framerate: 24 # by adjusting the frame rate, you can control the playback speed of the output video
+
+# hydra setting
+hydra:
+  help:
+    # app name, override to match the name your app is known by
+    app_name: 'SVGRender'
+  run:
+    # output directory for normal runs
+    # warning: make sure that the L56-58 of '/libs/engine/model_state.py' and 'dir' are modified together
+    dir: ./${result_path}/${x.method}-${now:%Y-%m-%d-%H-%M}
+
+# default settings
+defaults:
+  - _self_
+  - x: ~

conf/x/clipascene.yaml

@@ -0,0 +1,87 @@
+method: 'clipascene'
+
+im_name: ""
+image_size: 224
+u2net_path: "./checkpoint/u2net/u2net.pth"
+
+background_layer: 2 # 2, 8, 11
+background_div: 0.35 # 0.35, 0.5, 0.85
+background_num_iter: 1501
+
+foreground_layer: 2 # 2, 8, 11
+foreground_div: 0.4 # 0.4, 0.5, 0.9
+foreground_num_iter: 600 # 1000 if foreground_layer >= 8 else 600
+
+# general
+target: null
+output_dir: null
+path_svg: "none"
+mask_object: 0
+resize_obj: 0
+fix_scale: 0
+display_logs: 0
+display: 0
+test_name: "test"
+
+# training
+num_iter: 2001
+num_stages: 1
+lr_scheduler: 0
+lr: 0.0001
+color_lr: 0.01
+width_lr: 0.0001
+color_vars_threshold: 0.0
+batch_size: 1
+save_step: 100
+eval_step: 20
+loss_mask: "none"
+dilated_mask: 0
+mask_cls: None
+mask_attention: 0
+
+# strokes params
+num_paths: 64
+width: 1.5
+control_points_per_seg: 4
+num_segments: 1
+attention_init: 1
+saliency_model: "clip"
+saliency_clip_model: "ViT-B/32"
+xdog_intersec: 1
+mask_object_attention: 0
+softmax_temp: 0.3
+mlp_train: 1
+width_optim: 0
+mlp_width_weights_path: "none"
+mlp_points_weights_path: "none"
+switch_loss: 0
+gumbel_temp: 0.2
+width_loss_weight: 0
+width_loss_type: "L1"
+optimize_points: 1
+load_points_opt_weights: 0
+gradnorm: 0
+width_weights_lst: ""
+ratio_loss: 0
+
+# loss
+percep_loss: "none"
+perceptual_weight: 0
+train_with_clip: 0
+clip_weight: 0
+start_clip: 0
+num_aug_clip: 4
+include_target_in_aug: 0
+augment_both: 1
+augemntations: "affine"
+noise_thresh: 0.5
+aug_scale_min: 0.7
+force_sparse: 0
+clip_conv_loss: 1
+clip_mask_loss: 0
+clip_conv_loss_type: "L2"
+clip_conv_layer_weights: "0,0,1.0,1.0,0"
+clip_model_name: "ViT-B/32"
+clip_fc_loss_weight: 0
+clip_text_guide: 0
+text_target: None

conf/x/clipasso.yaml

@@ -0,0 +1,48 @@
+method: 'clipasso'
+
+image_size: 224
+mask_object: False
+fix_scale: False
+path_svg: ~ # if you want to load a svg file and train from it
+
+# train
+num_iter: 2001
+num_stages: 1 # training stages, you can train x strokes, then freeze them and train another x strokes etc
+lr_schedule: False
+lr: 1
+color_lr: 0.01
+color_vars_threshold: 0.0
+
+# SVG path attr
+num_paths: 24 # number of strokes
+width: 1.5 # stroke width
+control_points_per_seg: 4
+num_segments: 1
+attention_init: 1 # if True, use the attention heads of Dino model to set the location of the initial strokes
+saliency_model: "clip"
+saliency_clip_model: "ViT-B/32"
+xdog_intersec: 1
+mask_object_attention: 0
+softmax_temp: 0.3
+u2net_path: "./checkpoint/u2net/u2net.pth"
+
+# loss
+percep_loss: "none"
+perceptual_weight: 0
+train_with_clip: 0
+clip_weight: 0
+start_clip: 0
+num_aug_clip: 4
+include_target_in_aug: 0
+augment_both: 0
+augemntations: "affine"  # can be any combination of: 'affine_noise_eraserchunks_eraser_press'
+noise_thresh: 0.5
+aug_scale_min: 0.7
+force_sparse: 0 # if True, use L1 regularization on stroke's opacity to encourage small number of strokes
+clip_conv_loss: 1
+clip_conv_loss_type: "L2"
+clip_conv_layer_weights: "0,0,1.0,1.0,0"
+clip_model_name: "RN101"
+clip_fc_loss_weight: 0.1
+clip_text_guide: 0
+text_target: None

conf/x/clipdraw.yaml

@@ -0,0 +1,20 @@
+method: 'clipdraw'
+
+image_size: 224 # canvas size
+path_svg: ~ # if you want to load a svg file and train from it
+
+# train
+num_iter: 1000
+num_stages: 1 # training stages, you can train x strokes, then freeze them and train another x strokes etc
+lr_schedule: True
+lr: 1
+width_lr: 0.1
+color_lr: 0.01
+
+# SVG path attr
+num_paths: 512 # number of strokes
+max_width: 50 # stroke width
+black_stroke_color: False
+
+# loss
+num_aug: 4

conf/x/clipfont.yaml

@@ -0,0 +1,27 @@
+method: 'clipfont'
+
+# optimizer
+lr_base:
+  point: 0.1
+  color: 0.01
+lr_decay_rate: 0.1
+decay_steps: [ 1000, 1500 ]
+lr_schedule: False
+
+# train
+num_iter: 200
+batch_size: 1
+font:
+  reinit: False
+  reinit_color: 'randn' # 'randn', 'randn_all', 'green' et al
+
+# loss
+clip:
+  model_name: "ViT-B/32"  # RN101, 'ViT-B/32', ViT-L/14
+thresh: 0.0
+num_crops: 128
+crop_size: 230
+lam_patch: 150
+lam_dir: 30
+lam_lpips: 0
+lam_l2: 0

conf/x/diffsketcher.yaml

@@ -0,0 +1,76 @@
+method: 'diffsketcher'
+
+image_size: 224 # canvas size
+path_svg: ~ # if you want to load a svg file and train from it
+mask_object: False # if the target image contains background, it's better to mask it out
+fix_scale: False  # if the target image is not squared, it is recommended to fix the scale
+
+# train
+num_iter: 2000
+num_stages: 1 # training stages, you can train x strokes, then freeze them and train another x strokes etc
+lr_schedule: False
+lr_decay_rate: 0.1
+decay_steps: [ 1000, 1500 ]
+lr: 1
+color_lr: 0.01
+color_vars_threshold: 0.0 # uncomment the code
+width_lr: 0.1
+max_width: 50 # stroke width
+
+# stroke attrs
+num_paths: 128 # number of strokes
+width: 1.5 # stroke width
+control_points_per_seg: 4
+num_segments: 1
+optim_opacity: True # if True, the stroke opacity is optimized
+optim_width: False  # if True, the stroke width is optimized
+optim_rgba: False # if True, the stroke RGBA is optimized
+opacity_delta: 0 # stroke pruning
+
+# init strokes
+attention_init: True # if True, use the attention heads of Dino model to set the location of the initial strokes
+xdog_intersec: True # initialize along the edge, mix XDoG and attn up
+softmax_temp: 0.5
+cross_attn_res: 16
+self_attn_res: 32
+max_com: 20
+mean_comp: False
+comp_idx: 0
+attn_coeff: 1.0 # attn fusion, w * cross-attn + (1-w) * self-attn
+log_cross_attn: False # True if cross attn every step
+u2net_path: "./checkpoint/u2net/u2net.pth"
+
+# ldm
+model_id: "sd15"
+ldm_speed_up: False
+enable_xformers: True
+gradient_checkpoint: False
+token_ind: 5
+use_ddim: True
+num_inference_steps: 100
+guidance_scale: 7.5 # sdxl default 5.0
+
+# ASDS loss
+sds:
+  crop_size: 512
+  augmentations: "affine"
+  guidance_scale: 100
+  grad_scale: 1e-6
+  t_range: [ 0.05, 0.95 ]
+  warmup: 2000
+
+clip:
+  model_name: "RN101"  # RN101, ViT-L/14
+  feats_loss_type: "l2" # clip visual loss type, conv layers
+  feats_loss_weights: [ 0,0,1.0,1.0,0 ] # RN based
+  #  feats_loss_weights: [ 0,0,1.0,1.0,0,0,0,0,0,0,0,0 ] # ViT based
+  fc_loss_weight: 0.1 # clip visual loss, fc layer weight
+  augmentations: "affine" # augmentation before clip visual computation
+  num_aug: 4 # num of augmentation before clip visual computation
+  vis_loss: 1 # 1 or 0 for use or disable clip visual loss
+  text_visual_coeff: 0 # cosine similarity between text and img
+
+perceptual:
+  name: "lpips" # dists
+  lpips_net: 'vgg'
+  coeff: 0.2

conf/x/diffvg.yaml

@@ -0,0 +1,18 @@
+method: 'diffvg'
+
+# train
+num_iter: 2000 # num_iter
+lr_base:
+  point: 1
+  color: 0.01
+  stroke_width: 0.1
+  stroke_color: 0.01
+lr_schedule: False # use lr_schedule
+
+# SVG path attr
+num_paths: 512 # number of paths
+max_width: 5.0 # maximum width
+path_type: 'unclosed' # or 'closed', using Closed curve or non-closed curve
+
+# loss
+loss_type: 'l2' # or 'l1', 'l2', 'lpips', 'l2+lpips', loss type

conf/x/live.yaml

@@ -0,0 +1,31 @@
+method: 'live'
+
+image_size: 240 # img size and canvas size
+
+# train
+num_iter: 500 # num_iter per path group
+num_stages: 1 # training stages, you can train x strokes, then freeze them and train another x strokes etc
+lr_base:
+  point: 1
+  color: 0.01
+  bg: 0.01
+  stroke_width: 0.1
+  stroke_color: 0.01
+lr_schedule: True # use lr_schedule
+
+# SVG path attr
+num_paths: 5 # number of strokes
+path_schedule: 'repeat'
+schedule_each: 1 # [1, 3, 5, 7]
+train_stroke: False # train stroke width and color
+trainable_bg: False # set the background to be trainable
+width: 3 # stroke width
+num_segments: 4
+segment_init: 'circle' # 'random'
+radius: 5
+coord_init: 'sparse' # 'random', 'naive', place the first control point
+
+# loss
+use_l1_loss: False
+use_distance_weighted_loss: True
+xing_loss_weight: 0.01

conf/x/styleclipdraw.yaml

@@ -0,0 +1,21 @@
+method: 'styleclipdraw'
+
+image_size: 224 # canvas size
+path_svg: ~ # if you want to load an svg file and train from it
+
+# train
+num_iter: 1000
+num_stages: 1 # training stages, you can train x strokes, then freeze them and train another x strokes etc
+lr_schedule: True # anneal learning rate
+lr: 1
+width_lr: 0.1
+color_lr: 0.01
+
+# strokes
+num_paths: 512 # number of strokes
+max_width: 50 # stroke width
+black_stroke_color: False
+style_strength: 50  # How strong the style should be. 100 (max) is a lot. 0 (min) is no style.
+
+# loss
+num_aug: 10 # Number of image augmentations

conf/x/stylediffsketcher.yaml

@@ -0,0 +1,77 @@
+method: 'stylediffsketcher'
+
+image_size: 224 # canvas size
+path_svg: ~ # if you want to load a svg file and train from it
+mask_object: False # if the target image contains background, it's better to mask it out
+fix_scale: False  # if the target image is not squared, it is recommended to fix the scale
+
+# train
+num_iter: 2000
+num_stages: 1 # training stages, you can train x strokes, then freeze them and train another x strokes etc
+lr_schedule: False
+lr_decay_rate: 0.1
+decay_steps: [ 1000, 1500 ]
+lr: 1
+color_lr: 0.01
+color_vars_threshold: 0.0 # uncomment the code
+width_lr: 0.1
+max_width: 50 # stroke width
+
+# SVG path attrs
+num_paths: 512 # number of strokes
+width: 1.5 # init stroke width
+control_points_per_seg: 4
+num_segments: 1
+optim_opacity: True # if True, the stroke opacity is optimized
+optim_width: True  # if True, the stroke width is optimized
+optim_rgba: True # if True, the stroke RGBA is optimized
+
+# init strokes
+attention_init: True # if True, use the attention heads of Dino model to set the location of the initial strokes
+xdog_intersec: True # initialize along the edge, mix XDoG and attn up
+softmax_temp: 0.4
+cross_attn_res: 16
+self_attn_res: 32
+max_com: 20
+mean_comp: False
+comp_idx: 0
+attn_coeff: 1.0 # attn fusion, w * cross-attn + (1-w) * self-attn
+log_cross_attn: False
+u2net_path: "./checkpoint/u2net/u2net.pth"
+
+# ldm
+model_id: "sd15"
+ldm_speed_up: False
+enable_xformers: True
+gradient_checkpoint: False
+token_ind: 5
+use_ddim: True
+num_inference_steps: 100
+guidance_scale: 7.5
+
+# ASDS loss
+sds:
+  crop_size: 512
+  augmentations: "affine"
+  guidance_scale: 100
+  grad_scale: 0
+  t_range: [ 0.05, 0.95 ]
+  warmup: 120
+
+clip:
+  model_name: "RN101"  # RN101, ViT-L/14
+  feats_loss_type: "l2" # clip visual loss type, conv layers
+  feats_loss_weights: [ 0,0,1.0,1.0,0 ] # RN based
+  #  feats_loss_weights: [ 0,0,1.0,1.0,0,0,0,0,0,0,0,0 ] # ViT based
+  fc_loss_weight: 0.1 # clip visual loss, fc layer weight
+  augmentations: "affine_norm" # augmentation before clip visual computation, affine_norm_trivial
+  num_aug: 4 # num of augmentation before clip visual computation
+  vis_loss: 1 # 1 or 0 for use or disable clip visual loss
+  text_visual_coeff: 0 # cosine similarity between text and img
+
+perceptual:
+  name: "lpips" # dists
+  lpips_net: 'vgg'
+  coeff: 0.2
+
+style_strength: 1  # How strong the style should be. 100 (max) is a lot. 0 (min) is no style.

conf/x/svgdreamer.yaml

@@ -0,0 +1,122 @@
+method: "svgdreamer"
+
+image_size: 600 # canvas size
+path_svg: ~  # if you want to load a svg file and train from it
+num_stages: 1 # training stages, you can train x strokes, then freeze them and train another x strokes etc
+skip_sive: True # optimize from scratch without SIVE init
+color_init: 'rand' # if skip_live=True, then use color_init to init target_img
+style: "iconography" # "iconography", "pixelart", "low-poly", "painting", "sketch", "ink"
+
+# lr and optim
+lr_stage_one: # SIVE stage
+  point: 1 # control points
+  width: 0.1 # stroke width
+  color: 0.01 # fill color and stroke color
+  bg: 0.01 # bg in render_warp
+  optim:
+    name: 'adam'
+    betas: [ 0.9, 0.9 ]
+    eps: 1e-6
+  lr_schedule: True # use lr_scheduler
+  schedule:
+    name: 'linear'
+    keep_ratio: 0.2
+    decay_ratio: 0.4
+lr_stage_two: # VPSD stage
+  point: 1
+  width: 0.1
+  color: 0.01
+  bg: 0.01
+  lr_schedule: True # use lr_scheduler
+  optim:
+    name: 'adam'
+    betas: [ 0.9, 0.9 ]
+    eps: 1e-6
+  schedule:
+    name: 'cosine'
+    warmup_steps: 10
+    warmup_start_lr: 0.02
+    warmup_end_lr: 0.8
+    cosine_end_lr: 0.4
+
+# primitives
+num_paths: 256 # number of strokes
+trainable_bg: False # set the background to be trainable
+width: 3 # stroke width
+num_segments: 4
+segment_init: 'circle' # 'random'
+radius: 20
+coord_init: 'random' # 'random', 'naive', place the first control point
+grid: 50 # divide the canvas into n grids
+path_reinit: # reinitializing paths
+  use: True
+  freq: 100 # every 50 iterations
+  stop_step: 1000 # for VPSD fine-tuning
+  opacity_threshold: 0.05
+  area_threshold: 64
+
+# diffusion
+model_id: "sd21b" # sd14, sd15, sd21, sd21b, sdxl
+ldm_speed_up: False
+enable_xformers: True
+gradient_checkpoint: False
+cpu_offload: True
+num_inference_steps: 50
+guidance_scale: 7.5 # sdxl default 5.0
+K: 4
+lora_path: ~
+
+# VPSD loss
+guidance:
+  use: True
+  type: 'vpsd'
+  n_particle: 1 # 4, 8, 16
+  vsd_n_particle: 1 # the batch size of particles
+  particle_aug: False # do data enhancement for the input particles
+  num_iter: 2000 # total iterations
+  guidance_scale: 7.5 # CFG value
+  grad_scale: 1.0 # increase or decrease the gradient
+  grad_clip_val: ~ # eg: 10, clip the gradient of VPSD
+  t_range: [ 0.02, 0.98 ]
+  # 'randint': random time steps, this may have a more authentic style.
+  # 'max_0.5_900': annealing from 0.98 to 0.5 after 900 steps, this may have a more colorful results.
+  t_schedule: 'max_0.5_1000' # or 'randint'
+  # phi model config
+  phi_single: False # if False new an unet model to estimate noise
+  phi_model: 'lora' # 'lora', 'unet_simple'
+  use_attn_scale: ${x.guidance.phi_single} # use lora_attn_scale or not
+  lora_attn_scale: 1.0 # the scale of the attn based lora layer
+  phi_guidance_scale: 1.0
+  phi_t: False # different t for phi fine-tuning
+  phi_update_step: 1 # enable multi-update phi model or not
+  phi_lr: 0.0001 # learning rate of phi model
+  phi_scheduler: 'ddim' # 'dpm-solver'
+  phi_n_particle: 1 # the batch size of phi_model
+  # ReFL config
+  phi_ReFL: False # enable reward feed back learning
+  n_phi_sample: 1 # number of samples used in ReFL
+  phi_sample_step: 200 # the phi log step
+  phi_infer_step: 50 # the phi num_inference_steps
+  # phi model optim
+  phi_optim:
+    name: 'adamw'
+    betas: [ 0.9, 0.999 ]
+    eps: 1e-8
+    weight_decay: ~ # 1e-5
+  # phi model lr learning schedule
+  phi_schedule:
+    use: False
+    name: 'cosine'
+    warmup_steps: 50
+    warmup_start_lr: 0.00001
+    warmup_end_lr: 0.0001
+    total_step: 800
+    cosine_end_lr: 0.0001
+
+# reward model
+reward_path: './checkpoint/ImageReward'
+
+# xing loss for closed-form paths
+xing_loss:
+  use: False
+  weight: 0.01

conf/x/vectorfusion.yaml

@@ -0,0 +1,85 @@
+method: "vectorfusion"
+
+image_size: 600 # canvas size
+path_svg: ~  # if you want to load a svg file and train from it
+num_stages: 1 # training stages, you can train x strokes, then freeze them and train another x strokes etc
+skip_live: False # if skip_live then training from scratch
+style: "iconography" # "iconography", "pixelart", "low-poly", "painting", "sketch", "ink"
+
+# train
+batch_size: 1
+num_iter: 500 # num_iter per path group
+# lr and optim
+lr_stage_one:
+  point: 1
+  width: 0.1
+  color: 0.01
+  bg: 0.01
+  optim:
+    name: 'adam'
+    betas: [ 0.9, 0.9 ]
+    eps: 1e-6
+  lr_schedule: True # use lr_scheduler
+  schedule:
+    name: 'linear'
+    keep_ratio: 0.2
+    decay_ratio: 0.4
+lr_stage_two:
+  point: 1
+  width: 0.1
+  color: 0.01
+  bg: 0.01
+  lr_schedule: True # use lr_scheduler
+  optim:
+    name: 'adam'
+    betas: [ 0.9, 0.9 ]
+    eps: 1e-6
+  schedule:
+    name: 'cosine'
+    warmup_steps: 50
+    warmup_start_lr: 0.02
+    warmup_end_lr: 1.0
+    cosine_end_lr: 0.4
+
+# primitives
+num_paths: 128 # number of strokes
+path_schedule: 'repeat' # 'list'
+schedule_each: 16 # [1, 3, 5, 7]
+trainable_bg: False # set the background to be trainable
+width: 3 # stroke width
+num_segments: 4
+segment_init: 'circle' # 'random'
+radius: 20
+coord_init: 'sparse' # 'random', 'naive', place the first control point
+grid: 32 # divide the canvas into n grids
+path_reinit: # reinitializing paths
+  use: True
+  freq: 50 # every 50 iterations
+  stop_step: 800 # for SDS fine-tuning
+  opacity_threshold: 0.05
+  area_threshold: 64
+
+# diffusion
+model_id: "sd15" # sd14, sd15, sd21, sd21b, sdxl
+ldm_speed_up: False
+enable_xformers: True
+gradient_checkpoint: False
+cpu_offload: True
+num_inference_steps: 50
+guidance_scale: 7.5 # sdxl default 5.0
+K: 6
+lora_path: ~
+
+# SDS
+sds:
+  im_size: 512
+  guidance_scale: 100
+  grad_scale: 1.0
+  t_range: [ 0.05, 0.95 ]
+  num_iter: 1000 # fine-tuning steps
+
+# Live loss
+use_distance_weighted_loss: True
+xing_loss_weight: 0.01
+# pixel loss
+penalty_weight: 0.05

conf/x/wordasimage.yaml

@@ -0,0 +1,46 @@
+method: "wordasimage"
+
+image_size: 600 # canvas size
+word: "BUNNY"
+optim_letter: "Y"
+prompt_suffix: "minimal flat 2d vector. lineal color. trending on artstation"
+
+# train
+num_iter: 500
+lr_schedule: True
+lr:
+  point_lr: 1
+  lr_init: 0.002
+  lr_final: 0.0008
+  lr_delay_mult: 0.1
+  lr_delay_steps: 100
+
+# font
+font: 'KaushanScript-Regular'
+font_path: "./assets/fonts/${x.font}.ttf"
+level_of_cc: 1 # 0 - original number of cc / 1 - recommended / 2 - more control points
+
+# diffusion
+model_id: "sd15"
+ldm_speed_up: False
+enable_xformers: False
+gradient_checkpoint: False
+lora_path: ~
+
+# SDS
+sds:
+  im_size: 512
+  guidance_scale: 100
+  grad_scale: 1.0
+  t_range: [ 0.05, 0.95 ]
+  num_iter: 1000
+
+tone_loss:
+  use: True
+  dist_loss_weight: 100
+  pixel_dist_kernel_blur: 201
+  pixel_dist_sigma: 30
+
+conformal:
+  use: True
+  angeles_w: 0.5