chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,306 @@
|
||||
# TextDiffuser: Diffusion Models as Text Painters (NeurIPS 2023)
|
||||
|
||||
<a href='https://arxiv.org/pdf/2305.10855.pdf'><img src='https://img.shields.io/badge/Arxiv-2305.10855-red'>
|
||||
<a href='https://github.com/microsoft/unilm/tree/master/textdiffuser'><img src='https://img.shields.io/badge/Code-aka.ms/textdiffuser-yellow'>
|
||||
<a href='https://jingyechen.github.io/textdiffuser/'><img src='https://img.shields.io/badge/Project Page-link-green'>
|
||||
</a> [](https://huggingface.co/spaces/JingyeChen22/TextDiffuser)
|
||||
<a href='https://colab.research.google.com/drive/115Qw0l5dhjlTtrbywMWRwhz9IxKE4_Dg?usp=sharing'><img src='https://img.shields.io/badge/GoogleColab-link-purple'>
|
||||
|
||||
|
||||
TextDiffuser generates images with visually appealing text that is coherent with backgrounds. It is flexible and controllable to create high-quality text images using text prompts alone or together with text template images, and conduct text inpainting to reconstruct incomplete images with text.
|
||||
|
||||
<img src="assets/readme_images/introduction.jpg" width="80%">
|
||||
|
||||
## :star2: Highlights
|
||||
|
||||
* We propose **TextDiffuser**, which is a two-stage diffusion-based framework for text rendering. It generates accurate and coherent text images from text prompts or additionally with template images, as well as conducting text inpainting to reconstruct incomplete images.
|
||||
|
||||
* We release **MARIO-10M**, containing large-scale image-text pairs with OCR annotations, including text recognition, detection, and character-level segmentation masks.
|
||||
|
||||
* We construct **MARIO-Eval**, a comprehensive text rendering benchmark containing 10k prompts at [link](https://drive.google.com/file/d/1_tnWtOqC6S4_D4z8bqcBQ9xKPlsoPB0B/view?usp=sharing).
|
||||
|
||||
* We **release the demo** at [link](https://huggingface.co/spaces/JingyeChen22/TextDiffuser). Welcome to use and provide feedbacks :hugs:.
|
||||
|
||||
## :stopwatch: News
|
||||
|
||||
- __[2023.09.22]__: :tada: TextDiffuser is accepted to NeurIPS 2023.
|
||||
- __[2023.06.22]__: Evaluation script is released.
|
||||
- __[2023.06.15]__: :raised_hands: :raised_hands: :raised_hands: The Demo of TextDiffuser pre-trained with SD v2.1 is released in this [link](https://huggingface.co/spaces/JingyeChen22/TextDiffuser). Meanwhile, GoogleColab is available in this [link](https://colab.research.google.com/drive/115Qw0l5dhjlTtrbywMWRwhz9IxKE4_Dg?usp=sharing).
|
||||
- __[2023.06.08]__: Training script is released.
|
||||
- __[2023.06.07]__: MARIO-LAION is released.
|
||||
- __[2023.06.02]__: :raised_hands: :raised_hands: :raised_hands: Demo is available in this [link](https://huggingface.co/spaces/JingyeChen22/TextDiffuser).
|
||||
- __[2023.05.26]__: Upload the inference code and checkpoint.
|
||||
- __[2023.05.19]__: The paper is available at [link](https://arxiv.org/pdf/2305.10855.pdf).
|
||||
|
||||
## :hammer_and_wrench: Installation
|
||||
|
||||
Clone this repo:
|
||||
```
|
||||
git clone github_path_to/TextDiffuser
|
||||
cd TextDiffuser
|
||||
```
|
||||
|
||||
Build up a new environment and install packages as follows:
|
||||
```
|
||||
conda create -n textdiffuser python=3.8
|
||||
conda activate textdiffuser
|
||||
pip install -r requirements.txt
|
||||
```
|
||||
|
||||
Meanwhile, please install torch and torchvision that matches the version of system and cuda (refer to this [link](https://download.pytorch.org/whl/torch_stable.html)).
|
||||
|
||||
|
||||
Install Hugging Face Diffuser and replace some files:
|
||||
```
|
||||
git clone https://github.com/JingyeChen/diffusers
|
||||
cp ./assets/files/scheduling_ddpm.py ./diffusers/src/diffusers/schedulers/scheduling_ddpm.py
|
||||
cp ./assets/files/unet_2d_condition.py ./diffusers/src/diffusers/models/unet_2d_condition.py
|
||||
cp ./assets/files/modeling_utils.py ./diffusers/src/diffusers/models/modeling_utils.py
|
||||
cd diffusers && pip install -e .
|
||||
```
|
||||
|
||||
Besides, a font file is needed for layout generation. Please put your font in ```assets/font/```. We recommend to use ```Arial.ttf```.
|
||||
|
||||
|
||||
|
||||
## :floppy_disk: Checkpoint
|
||||
|
||||
The checkpoints are in [HFLink](https://huggingface.co/datasets/JingyeChen22/TextDiffuser/resolve/main/textdiffuser-ckpt-new.zip) (3.2GB). Please download it and unzip it. The file structures should be as follows:
|
||||
|
||||
|
||||
|
||||
```
|
||||
textdiffuser
|
||||
├── textdiffuser-ckpt
|
||||
│ ├── diffusion_backbone/ # for diffusion backbone
|
||||
│ ├── character_aware_loss_unet.pth # for character-aware loss
|
||||
│ ├── layout_transformer.pth # for layout transformer
|
||||
│ └── text_segmenter.pth # for character-level segmenter
|
||||
├── README.md
|
||||
```
|
||||
|
||||
## :books: Dataset
|
||||
|
||||
<img src="assets/readme_images/laion-ocr.jpg" width="80%">
|
||||
|
||||
|
||||
**MARIO-LAION**'s meta information is at [googledrive](https://drive.google.com/file/d/1gzaW8i07Fn7-zWiP2ZwtkH_E0-2qPgJx/view?usp=sharing) (40GB), containing 9,194,613 samples. Please download it and unzip it by running ```python data/maion-laion-unzip.py```. The file structures of each folder should be as follows and ```data/maion-laion-example``` is provided for reference. We also provide ```data/visualize_charseg.ipynb``` to visualize the character-level segmentation mask.
|
||||
|
||||
```
|
||||
├── 28330/
|
||||
│ ├── 283305839/
|
||||
│ │ ├── caption.txt # caption of the image
|
||||
│ │ ├── charseg.npy # character-level segmentation mask
|
||||
│ │ ├── info.json # more meta information given by laion, such as original height and width
|
||||
├── ├── └── ocr.txt # ocr detection and recognition results
|
||||
```
|
||||
|
||||
The urls of each image is at [googledrive](https://drive.google.com/file/d/1ndG-EwM66rH9CfdFVE_XbwFxTFEEAPZ7/view?usp=sharing) (794.6MB). The file structure is as follows:
|
||||
|
||||
```
|
||||
├── maion_laion_image_url/
|
||||
│ ├── mario-laion-url.txt # urls for downloading by img2dataset
|
||||
│ ├── mario-laion-index-url.txt # urls and indices for each image
|
||||
│ └── mario-laion-test-index.txt # all indices for test dataset
|
||||
```
|
||||
|
||||
Please download img2dataset wiht ```pip install img2dataset```, and download the images using the following command:
|
||||
```
|
||||
img2dataset --url_list=url.txt --output_folder=laion_ocr --thread_count=64 --resize_mode=no
|
||||
```
|
||||
|
||||
After downloading, you need to resize each image to ```512x512```. Please follow ```mario-laion-index-url.txt``` to move each image to the corresponding folders. Images with indices in ```mario-laion-test-index.txt``` are used for testing. Please note that some links may be <span style="color:red">**invalid**</span>
|
||||
since the owners remove the images from their website.
|
||||
|
||||
## :steam_locomotive: Train
|
||||
|
||||
Please use ```accelerate config``` to configure your acceleration policy at first, then modify output_dir, dataset_path, and train_dataset_index_file in ```train.sh```. The train_dataset_index_file should be a .txt file, and each line should indicate an index of a training sample.
|
||||
|
||||
```txt
|
||||
06269_062690093
|
||||
27197_271975251
|
||||
27197_271978467
|
||||
...
|
||||
```
|
||||
|
||||
Then you can use the following to run TextDiffuser:
|
||||
|
||||
```bash
|
||||
accelerate launch train.py \
|
||||
--train_batch_size=24 \
|
||||
--gradient_accumulation_steps=4 \
|
||||
--gradient_checkpointing \
|
||||
--mixed_precision="fp16" \
|
||||
--num_train_epochs=2 \
|
||||
--learning_rate=1e-5 \
|
||||
--max_grad_norm=1 \
|
||||
--lr_scheduler="constant" \
|
||||
--lr_warmup_steps=0 \
|
||||
--output_dir="experiment_name" \
|
||||
--enable_xformers_memory_efficient_attention \
|
||||
--dataloader_num_workers=4 \
|
||||
--character_aware_loss_lambda=0.01 \
|
||||
--resume_from_checkpoint="latest" \
|
||||
--drop_caption \
|
||||
--mask_all_ratio=0.5 \
|
||||
--segmentation_mask_aug \
|
||||
--dataset_path=/home/path/to/laion-ocr-unzip \
|
||||
--train_dataset_index_file=/path/to/index_file.txt \
|
||||
--vis_num=8
|
||||
```
|
||||
|
||||
If you encounter an "out-of-memory" error, please consider reducing the batch size appropriately.
|
||||
|
||||
|
||||
## :firecracker: Inference
|
||||
|
||||
TextDiffuser can be applied on: text-to-image, text-to-image-with-template, and text-inpainting.
|
||||
|
||||
### Text-to-Image
|
||||
This task is designed to generate images based on given prompts. Users are required to enclose the keywords to be drawn with single quotation marks.
|
||||
|
||||
```bash
|
||||
CUDA_VISIBLE_DEVICES=0 python inference.py \
|
||||
--mode="text-to-image" \
|
||||
--resume_from_checkpoint="textdiffuser-ckpt/diffusion_backbone" \
|
||||
--prompt="A sign that says 'Hello'" \
|
||||
--output_dir="./output" \
|
||||
--vis_num=4
|
||||
```
|
||||
|
||||
### Text-to-Image-with-Template
|
||||
This task aims to generate images based on given prompts and template images (can be printed, handwritten, or scene text images). A pre-trained character-level segmentation model is used to extract layout information from the template image.
|
||||
|
||||
```bash
|
||||
CUDA_VISIBLE_DEVICES=0 python inference.py \
|
||||
--mode="text-to-image-with-template" \
|
||||
--resume_from_checkpoint="textdiffuser-ckpt/diffusion_backbone" \
|
||||
--prompt="a poster of monkey music festival" \
|
||||
--template_image="assets/examples/text-to-image-with-template/case2.jpg" \
|
||||
--output_dir="./output" \
|
||||
--vis_num=4
|
||||
```
|
||||
|
||||
### Text-Inpainting
|
||||
This task aims to modify a given image in an inpainting manner. The provided text mask image should contain the inpainting region and the text to be drawn within the region.
|
||||
|
||||
```bash
|
||||
CUDA_VISIBLE_DEVICES=0 python inference.py \
|
||||
--mode="text-inpainting" \
|
||||
--resume_from_checkpoint="textdiffuser-ckpt/diffusion_backbone" \
|
||||
--prompt="a boy draws good morning on a board" \
|
||||
--original_image="assets/examples/text-inpainting/case2.jpg" \
|
||||
--text_mask="assets/examples/text-inpainting/case2_mask.jpg" \
|
||||
--output_dir="./output" \
|
||||
--vis_num=4
|
||||
```
|
||||
|
||||
## :chart_with_upwards_trend: Evaluation
|
||||
|
||||
For evaluation, please download [MARIOEval](https://drive.google.com/file/d/1_tnWtOqC6S4_D4z8bqcBQ9xKPlsoPB0B/view?usp=sharing) and the generation results of each methods are at [link](https://drive.google.com/file/d/1d6SWg--MXhPDkGuy5tRmTsPkcpP2NDEo/view?usp=sharing) for reference.
|
||||
. MARIOEval contains 5,414 prompts for evaluation, including the following subsets:
|
||||
|
||||
| Subset | #Sample | Subset | #Sample |
|
||||
| --- | ---: | --- | ---: |
|
||||
| LAIONEval4000 | 4,000 | ChineseDrawText | 175 |
|
||||
| TMDBEval500 | 500 | DrawBenchText | 21 |
|
||||
| OpenLibrary500 | 500 | DrawTextCreative | 218 |
|
||||
|
||||
The structure of each folder is as follows:
|
||||
|
||||
```bash
|
||||
├── LAIONEval4000/
|
||||
│ ├── images/ # ground truth images
|
||||
│ ├── render/ # layouts of keywords generated by Layout Transformer
|
||||
│ ├── LAIONEval4000.txt # prompts with keywords enclosed with quotes
|
||||
│ └── LAIONEval4000_wo_quote.txt # prompts without quotes
|
||||
```
|
||||
|
||||
Please note that the ground truth images are only available for the LAIONEval4000, TMDBEval500, and OpenLibrary500 subsets. The render images are used for evaluating ControlNet. We manually enclose keywords with quotes according to the ocr results. Please refer to the ```_wo_quote.txt``` version for original prompts.
|
||||
|
||||
To evaluate TextDiffuser, please use the following command for sampling:
|
||||
|
||||
```python
|
||||
CUDA_VISIBLE_DEVICES=0 python evaluate.py \
|
||||
--mode="text-to-image" \
|
||||
--resume_from_checkpoint="textdiffuser-ckpt/diffusion_backbone" \
|
||||
--prompt_list="/path/to/MARIOEval/TMDBEval500/TMDBEval500.txt" \
|
||||
--output_dir="/path/to/output_dir" \
|
||||
--vis_num=4
|
||||
```
|
||||
|
||||
To sample from other baseline methods (e.g, Stable Diffusion, ControlNet, and DeepFloyd), the scripts are provided in the ```./eval``` folder. We also provided the scripts for calculating FID, Clip Score, as well as the OCR metrics.
|
||||
|
||||
|
||||
| Metrics | Stable Diffusion | ContolNet | DeepFloyd | TextDiffuser (Ours) |
|
||||
| :---: | :---: | :---: | :---: | :---: |
|
||||
| FID↓ | 51.295 | 51.485 | **34.902** | 38.758 |
|
||||
| CLIPScore↑ | 0.3015 | 0.3424 | 0.3267 | **0.3436** |
|
||||
| OCR-Accuracy↑ | 0.0003 | 0.2390 | 0.0262 | **0.5609** |
|
||||
| OCR-Precision↑ | 0.0173 | 0.5211 | 0.1450 | **0.7846** |
|
||||
| OCR-Recall↑ | 0.0280 | 0.6707 | 0.2245 | **0.7802** |
|
||||
| OCR-Fmeasure↑ | 0.0214 | 0.5865 | 0.1762 | **0.7824** |
|
||||
| *OCR-Accuracy↑ | 0.0178 | 0.2705 | 0.0457 | **0.5712** |
|
||||
| *OCR-Precision↑ | 0.0192 | 0.5391 | 0.1738 | **0.7795** |
|
||||
| *OCR-Recall↑ | 0.0260 | 0.6438 | 0.2235 | **0.7498** |
|
||||
| *OCR-Fmeasure↑ | 0.0221 | 0.5868 | 0.1955 | **0.7643** |
|
||||
|
||||
Please note that OCR metrics begin with "\*" mean we use open-source [MaskTextSpotterV3](https://github.com/MhLiao/MaskTextSpotterV3) for evaluation, and without "\*" denote we use [MicroSoft OCR API](https://azure.microsoft.com/en-us/updates/computer-vision-v3-preview-6/) for evaluation. The performance of text-to-image on MARIO-Eval compared with existing methods. TextDiffuser performs
|
||||
the best regarding CLIPScore and OCR evaluation while achieving comparable performance on FID.
|
||||
|
||||
<img src="assets/readme_images/userstudy.jpg" width="90%">
|
||||
|
||||
User studies for whole-image generation and part-image generation tasks. (a) For whole-image generation, our method clearly outperforms others in both aspects of text rendering quality and image-text matching. (b) For part-image generation, our method receives high scores from human evaluators in these two aspects.
|
||||
|
||||
|
||||
## :joystick: Demo
|
||||
TextDiffuser has been deployed on [Hugging Face](https://huggingface.co/spaces/JingyeChen22/TextDiffuser). If you have advanced GPUs, you may deploy the demo locally as follows:
|
||||
|
||||
```python
|
||||
CUDA_VISIBLE_DEVICES=0 python gradio_app.py
|
||||
```
|
||||
|
||||
Then you can enjoy the demo with local browser:
|
||||
|
||||
<img src="assets/readme_images/demo.jpg" width="90%">
|
||||
|
||||
|
||||
|
||||
|
||||
## :framed_picture: Gallery
|
||||
|
||||
### Text-to-Image
|
||||
<img src="assets/readme_images/gallery_text-to-image.jpg" width="80%">
|
||||
|
||||
### Text-to-Image-with-Template
|
||||
<img src="assets/readme_images/gallery_text-to-image-with-template.jpg" width="80%">
|
||||
|
||||
### Text-Inpainting
|
||||
<img src="assets/readme_images/gallery_text-inpainting.jpg" width="80%">
|
||||
|
||||
## :love_letter: Acknowledgement
|
||||
|
||||
We sincerely thank the following projects: [Hugging Face Diffuser](https://github.com/huggingface/diffusers), [LAION](https://laion.ai/laion-400-open-dataset/), [DB](https://github.com/MhLiao/DB), [PARSeq](https://github.com/baudm/parseq), [img2dataset](https://github.com/rom1504/img2dataset).
|
||||
|
||||
Also, special thanks to the open-source diffusion project or available demo: [DALLE](https://openai.com/product/dall-e-2), [Stable Diffusion](https://github.com/CompVis/stable-diffusion), [Stable Diffusion XL](https://dreamstudio.ai/generate), [Midjourney](https://www.midjourney.com/home/?callbackUrl=%2Fapp%2F), [ControlNet](https://github.com/lllyasviel/ControlNet), [DeepFloyd](https://github.com/deep-floyd/IF).
|
||||
|
||||
|
||||
## :exclamation: Disclaimer
|
||||
Please note that the code is intended for academic and research purposes **ONLY**. Any use of the code for generating inappropriate content is **strictly prohibited**. The responsibility for any misuse or inappropriate use of the code lies solely with the users who generated such content, and this code shall not be held liable for any such use.
|
||||
|
||||
## :envelope: Contact
|
||||
|
||||
For help or issues using TextDiffuser, please email Jingye Chen (qwerty.chen@connect.ust.hk), Yupan Huang (huangyp28@mail2.sysu.edu.cn) or submit a GitHub issue.
|
||||
|
||||
For other communications related to TextDiffuser, please contact Lei Cui (lecu@microsoft.com) or Furu Wei (fuwei@microsoft.com).
|
||||
|
||||
## :herb: Citation
|
||||
If you find this code useful in your research, please consider citing:
|
||||
```
|
||||
@article{chen2023textdiffuser,
|
||||
title={TextDiffuser: Diffusion Models as Text Painters},
|
||||
author={Chen, Jingye and Huang, Yupan and Lv, Tengchao and Cui, Lei and Chen, Qifeng and Wei, Furu},
|
||||
journal={arXiv preprint arXiv:2305.10855},
|
||||
year={2023}
|
||||
}
|
||||
```
|
||||
Binary file not shown.
|
After Width: | Height: | Size: 391 KiB |
@@ -0,0 +1 @@
|
||||
a board of assyrian on unflagging fry devastates
|
||||
Binary file not shown.
|
After Width: | Height: | Size: 57 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 168 KiB |
@@ -0,0 +1 @@
|
||||
a boy draws good morning on a board
|
||||
Binary file not shown.
|
After Width: | Height: | Size: 39 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 378 KiB |
@@ -0,0 +1 @@
|
||||
a board saying Sessions
|
||||
Binary file not shown.
|
After Width: | Height: | Size: 29 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 211 KiB |
@@ -0,0 +1 @@
|
||||
a board without text
|
||||
Binary file not shown.
|
After Width: | Height: | Size: 16 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 57 KiB |
@@ -0,0 +1 @@
|
||||
a birthday cake that says happy birthday to XYZ
|
||||
Binary file not shown.
|
After Width: | Height: | Size: 58 KiB |
@@ -0,0 +1 @@
|
||||
a poster of monkey music festival
|
||||
Binary file not shown.
|
After Width: | Height: | Size: 50 KiB |
@@ -0,0 +1 @@
|
||||
a meme of are you kidding
|
||||
@@ -0,0 +1 @@
|
||||
a text image of 'Hello World'
|
||||
@@ -0,0 +1 @@
|
||||
A cover named 'Anything is possible'
|
||||
@@ -0,0 +1 @@
|
||||
'Team' hat
|
||||
@@ -0,0 +1,861 @@
|
||||
# coding=utf-8
|
||||
# Copyright 2023 The HuggingFace Inc. team.
|
||||
# Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
import inspect
|
||||
import os
|
||||
import warnings
|
||||
from functools import partial
|
||||
from typing import Callable, List, Optional, Tuple, Union
|
||||
|
||||
import torch
|
||||
from huggingface_hub import hf_hub_download
|
||||
from huggingface_hub.utils import EntryNotFoundError, RepositoryNotFoundError, RevisionNotFoundError
|
||||
from packaging import version
|
||||
from requests import HTTPError
|
||||
from torch import Tensor, device
|
||||
|
||||
from .. import __version__
|
||||
from ..utils import (
|
||||
CONFIG_NAME,
|
||||
DEPRECATED_REVISION_ARGS,
|
||||
DIFFUSERS_CACHE,
|
||||
FLAX_WEIGHTS_NAME,
|
||||
HF_HUB_OFFLINE,
|
||||
HUGGINGFACE_CO_RESOLVE_ENDPOINT,
|
||||
SAFETENSORS_WEIGHTS_NAME,
|
||||
WEIGHTS_NAME,
|
||||
is_accelerate_available,
|
||||
is_safetensors_available,
|
||||
is_torch_version,
|
||||
logging,
|
||||
)
|
||||
|
||||
|
||||
logger = logging.get_logger(__name__)
|
||||
|
||||
|
||||
if is_torch_version(">=", "1.9.0"):
|
||||
_LOW_CPU_MEM_USAGE_DEFAULT = True
|
||||
else:
|
||||
_LOW_CPU_MEM_USAGE_DEFAULT = False
|
||||
|
||||
|
||||
if is_accelerate_available():
|
||||
import accelerate
|
||||
from accelerate.utils import set_module_tensor_to_device
|
||||
from accelerate.utils.versions import is_torch_version
|
||||
|
||||
if is_safetensors_available():
|
||||
import safetensors
|
||||
|
||||
|
||||
def get_parameter_device(parameter: torch.nn.Module):
|
||||
try:
|
||||
return next(parameter.parameters()).device
|
||||
except StopIteration:
|
||||
# For torch.nn.DataParallel compatibility in PyTorch 1.5
|
||||
|
||||
def find_tensor_attributes(module: torch.nn.Module) -> List[Tuple[str, Tensor]]:
|
||||
tuples = [(k, v) for k, v in module.__dict__.items() if torch.is_tensor(v)]
|
||||
return tuples
|
||||
|
||||
gen = parameter._named_members(get_members_fn=find_tensor_attributes)
|
||||
first_tuple = next(gen)
|
||||
return first_tuple[1].device
|
||||
|
||||
|
||||
def get_parameter_dtype(parameter: torch.nn.Module):
|
||||
try:
|
||||
return next(parameter.parameters()).dtype
|
||||
except StopIteration:
|
||||
# For torch.nn.DataParallel compatibility in PyTorch 1.5
|
||||
|
||||
def find_tensor_attributes(module: torch.nn.Module) -> List[Tuple[str, Tensor]]:
|
||||
tuples = [(k, v) for k, v in module.__dict__.items() if torch.is_tensor(v)]
|
||||
return tuples
|
||||
|
||||
gen = parameter._named_members(get_members_fn=find_tensor_attributes)
|
||||
first_tuple = next(gen)
|
||||
return first_tuple[1].dtype
|
||||
|
||||
|
||||
def load_state_dict(checkpoint_file: Union[str, os.PathLike], variant: Optional[str] = None):
|
||||
"""
|
||||
Reads a checkpoint file, returning properly formatted errors if they arise.
|
||||
"""
|
||||
try:
|
||||
if os.path.basename(checkpoint_file) == _add_variant(WEIGHTS_NAME, variant):
|
||||
return torch.load(checkpoint_file, map_location="cpu")
|
||||
else:
|
||||
return safetensors.torch.load_file(checkpoint_file, device="cpu")
|
||||
except Exception as e:
|
||||
try:
|
||||
with open(checkpoint_file) as f:
|
||||
if f.read().startswith("version"):
|
||||
raise OSError(
|
||||
"You seem to have cloned a repository without having git-lfs installed. Please install "
|
||||
"git-lfs and run `git lfs install` followed by `git lfs pull` in the folder "
|
||||
"you cloned."
|
||||
)
|
||||
else:
|
||||
raise ValueError(
|
||||
f"Unable to locate the file {checkpoint_file} which is necessary to load this pretrained "
|
||||
"model. Make sure you have saved the model properly."
|
||||
) from e
|
||||
except (UnicodeDecodeError, ValueError):
|
||||
raise OSError(
|
||||
f"Unable to load weights from checkpoint file for '{checkpoint_file}' "
|
||||
f"at '{checkpoint_file}'. "
|
||||
"If you tried to load a PyTorch model from a TF 2.0 checkpoint, please set from_tf=True."
|
||||
)
|
||||
|
||||
|
||||
def _load_state_dict_into_model(model_to_load, state_dict):
|
||||
# Convert old format to new format if needed from a PyTorch state_dict
|
||||
# copy state_dict so _load_from_state_dict can modify it
|
||||
state_dict = state_dict.copy()
|
||||
error_msgs = []
|
||||
|
||||
# PyTorch's `_load_from_state_dict` does not copy parameters in a module's descendants
|
||||
# so we need to apply the function recursively.
|
||||
def load(module: torch.nn.Module, prefix=""):
|
||||
args = (state_dict, prefix, {}, True, [], [], error_msgs)
|
||||
module._load_from_state_dict(*args)
|
||||
|
||||
for name, child in module._modules.items():
|
||||
if child is not None:
|
||||
load(child, prefix + name + ".")
|
||||
|
||||
load(model_to_load)
|
||||
|
||||
return error_msgs
|
||||
|
||||
|
||||
def _add_variant(weights_name: str, variant: Optional[str] = None) -> str:
|
||||
if variant is not None:
|
||||
splits = weights_name.split(".")
|
||||
splits = splits[:-1] + [variant] + splits[-1:]
|
||||
weights_name = ".".join(splits)
|
||||
|
||||
return weights_name
|
||||
|
||||
|
||||
class ModelMixin(torch.nn.Module):
|
||||
r"""
|
||||
Base class for all models.
|
||||
|
||||
[`ModelMixin`] takes care of storing the configuration of the models and handles methods for loading, downloading
|
||||
and saving models.
|
||||
|
||||
- **config_name** ([`str`]) -- A filename under which the model should be stored when calling
|
||||
[`~models.ModelMixin.save_pretrained`].
|
||||
"""
|
||||
config_name = CONFIG_NAME
|
||||
_automatically_saved_args = ["_diffusers_version", "_class_name", "_name_or_path"]
|
||||
_supports_gradient_checkpointing = False
|
||||
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
|
||||
@property
|
||||
def is_gradient_checkpointing(self) -> bool:
|
||||
"""
|
||||
Whether gradient checkpointing is activated for this model or not.
|
||||
|
||||
Note that in other frameworks this feature can be referred to as "activation checkpointing" or "checkpoint
|
||||
activations".
|
||||
"""
|
||||
return any(hasattr(m, "gradient_checkpointing") and m.gradient_checkpointing for m in self.modules())
|
||||
|
||||
def enable_gradient_checkpointing(self):
|
||||
"""
|
||||
Activates gradient checkpointing for the current model.
|
||||
|
||||
Note that in other frameworks this feature can be referred to as "activation checkpointing" or "checkpoint
|
||||
activations".
|
||||
"""
|
||||
if not self._supports_gradient_checkpointing:
|
||||
raise ValueError(f"{self.__class__.__name__} does not support gradient checkpointing.")
|
||||
self.apply(partial(self._set_gradient_checkpointing, value=True))
|
||||
|
||||
def disable_gradient_checkpointing(self):
|
||||
"""
|
||||
Deactivates gradient checkpointing for the current model.
|
||||
|
||||
Note that in other frameworks this feature can be referred to as "activation checkpointing" or "checkpoint
|
||||
activations".
|
||||
"""
|
||||
if self._supports_gradient_checkpointing:
|
||||
self.apply(partial(self._set_gradient_checkpointing, value=False))
|
||||
|
||||
def set_use_memory_efficient_attention_xformers(
|
||||
self, valid: bool, attention_op: Optional[Callable] = None
|
||||
) -> None:
|
||||
# Recursively walk through all the children.
|
||||
# Any children which exposes the set_use_memory_efficient_attention_xformers method
|
||||
# gets the message
|
||||
def fn_recursive_set_mem_eff(module: torch.nn.Module):
|
||||
if hasattr(module, "set_use_memory_efficient_attention_xformers"):
|
||||
module.set_use_memory_efficient_attention_xformers(valid, attention_op)
|
||||
|
||||
for child in module.children():
|
||||
fn_recursive_set_mem_eff(child)
|
||||
|
||||
for module in self.children():
|
||||
if isinstance(module, torch.nn.Module):
|
||||
fn_recursive_set_mem_eff(module)
|
||||
|
||||
def enable_xformers_memory_efficient_attention(self, attention_op: Optional[Callable] = None):
|
||||
r"""
|
||||
Enable memory efficient attention as implemented in xformers.
|
||||
|
||||
When this option is enabled, you should observe lower GPU memory usage and a potential speed up at inference
|
||||
time. Speed up at training time is not guaranteed.
|
||||
|
||||
Warning: When Memory Efficient Attention and Sliced attention are both enabled, the Memory Efficient Attention
|
||||
is used.
|
||||
|
||||
Parameters:
|
||||
attention_op (`Callable`, *optional*):
|
||||
Override the default `None` operator for use as `op` argument to the
|
||||
[`memory_efficient_attention()`](https://facebookresearch.github.io/xformers/components/ops.html#xformers.ops.memory_efficient_attention)
|
||||
function of xFormers.
|
||||
|
||||
Examples:
|
||||
|
||||
```py
|
||||
>>> import torch
|
||||
>>> from diffusers import UNet2DConditionModel
|
||||
>>> from xformers.ops import MemoryEfficientAttentionFlashAttentionOp
|
||||
|
||||
>>> model = UNet2DConditionModel.from_pretrained(
|
||||
... "stabilityai/stable-diffusion-2-1", subfolder="unet", torch_dtype=torch.float16
|
||||
... )
|
||||
>>> model = model.to("cuda")
|
||||
>>> model.enable_xformers_memory_efficient_attention(attention_op=MemoryEfficientAttentionFlashAttentionOp)
|
||||
```
|
||||
"""
|
||||
self.set_use_memory_efficient_attention_xformers(True, attention_op)
|
||||
|
||||
def disable_xformers_memory_efficient_attention(self):
|
||||
r"""
|
||||
Disable memory efficient attention as implemented in xformers.
|
||||
"""
|
||||
self.set_use_memory_efficient_attention_xformers(False)
|
||||
|
||||
def save_pretrained(
|
||||
self,
|
||||
save_directory: Union[str, os.PathLike],
|
||||
is_main_process: bool = True,
|
||||
save_function: Callable = None,
|
||||
safe_serialization: bool = False,
|
||||
variant: Optional[str] = None,
|
||||
):
|
||||
"""
|
||||
Save a model and its configuration file to a directory, so that it can be re-loaded using the
|
||||
`[`~models.ModelMixin.from_pretrained`]` class method.
|
||||
|
||||
Arguments:
|
||||
save_directory (`str` or `os.PathLike`):
|
||||
Directory to which to save. Will be created if it doesn't exist.
|
||||
is_main_process (`bool`, *optional*, defaults to `True`):
|
||||
Whether the process calling this is the main process or not. Useful when in distributed training like
|
||||
TPUs and need to call this function on all processes. In this case, set `is_main_process=True` only on
|
||||
the main process to avoid race conditions.
|
||||
save_function (`Callable`):
|
||||
The function to use to save the state dictionary. Useful on distributed training like TPUs when one
|
||||
need to replace `torch.save` by another method. Can be configured with the environment variable
|
||||
`DIFFUSERS_SAVE_MODE`.
|
||||
safe_serialization (`bool`, *optional*, defaults to `False`):
|
||||
Whether to save the model using `safetensors` or the traditional PyTorch way (that uses `pickle`).
|
||||
variant (`str`, *optional*):
|
||||
If specified, weights are saved in the format pytorch_model.<variant>.bin.
|
||||
"""
|
||||
if safe_serialization and not is_safetensors_available():
|
||||
raise ImportError("`safe_serialization` requires the `safetensors library: `pip install safetensors`.")
|
||||
|
||||
if os.path.isfile(save_directory):
|
||||
logger.error(f"Provided path ({save_directory}) should be a directory, not a file")
|
||||
return
|
||||
|
||||
os.makedirs(save_directory, exist_ok=True)
|
||||
|
||||
model_to_save = self
|
||||
|
||||
# Attach architecture to the config
|
||||
# Save the config
|
||||
if is_main_process:
|
||||
model_to_save.save_config(save_directory)
|
||||
|
||||
# Save the model
|
||||
state_dict = model_to_save.state_dict()
|
||||
|
||||
weights_name = SAFETENSORS_WEIGHTS_NAME if safe_serialization else WEIGHTS_NAME
|
||||
weights_name = _add_variant(weights_name, variant)
|
||||
|
||||
# Save the model
|
||||
if safe_serialization:
|
||||
safetensors.torch.save_file(
|
||||
state_dict, os.path.join(save_directory, weights_name), metadata={"format": "pt"}
|
||||
)
|
||||
else:
|
||||
torch.save(state_dict, os.path.join(save_directory, weights_name))
|
||||
|
||||
logger.info(f"Model weights saved in {os.path.join(save_directory, weights_name)}")
|
||||
|
||||
@classmethod
|
||||
def from_pretrained(cls, pretrained_model_name_or_path: Optional[Union[str, os.PathLike]], **kwargs):
|
||||
r"""
|
||||
Instantiate a pretrained pytorch model from a pre-trained model configuration.
|
||||
|
||||
The model is set in evaluation mode by default using `model.eval()` (Dropout modules are deactivated). To train
|
||||
the model, you should first set it back in training mode with `model.train()`.
|
||||
|
||||
The warning *Weights from XXX not initialized from pretrained model* means that the weights of XXX do not come
|
||||
pretrained with the rest of the model. It is up to you to train those weights with a downstream fine-tuning
|
||||
task.
|
||||
|
||||
The warning *Weights from XXX not used in YYY* means that the layer XXX is not used by YYY, therefore those
|
||||
weights are discarded.
|
||||
|
||||
Parameters:
|
||||
pretrained_model_name_or_path (`str` or `os.PathLike`, *optional*):
|
||||
Can be either:
|
||||
|
||||
- A string, the *model id* of a pretrained model hosted inside a model repo on huggingface.co.
|
||||
Valid model ids should have an organization name, like `google/ddpm-celebahq-256`.
|
||||
- A path to a *directory* containing model weights saved using [`~ModelMixin.save_config`], e.g.,
|
||||
`./my_model_directory/`.
|
||||
|
||||
cache_dir (`Union[str, os.PathLike]`, *optional*):
|
||||
Path to a directory in which a downloaded pretrained model configuration should be cached if the
|
||||
standard cache should not be used.
|
||||
torch_dtype (`str` or `torch.dtype`, *optional*):
|
||||
Override the default `torch.dtype` and load the model under this dtype. If `"auto"` is passed the dtype
|
||||
will be automatically derived from the model's weights.
|
||||
force_download (`bool`, *optional*, defaults to `False`):
|
||||
Whether or not to force the (re-)download of the model weights and configuration files, overriding the
|
||||
cached versions if they exist.
|
||||
resume_download (`bool`, *optional*, defaults to `False`):
|
||||
Whether or not to delete incompletely received files. Will attempt to resume the download if such a
|
||||
file exists.
|
||||
proxies (`Dict[str, str]`, *optional*):
|
||||
A dictionary of proxy servers to use by protocol or endpoint, e.g., `{'http': 'foo.bar:3128',
|
||||
'http://hostname': 'foo.bar:4012'}`. The proxies are used on each request.
|
||||
output_loading_info(`bool`, *optional*, defaults to `False`):
|
||||
Whether or not to also return a dictionary containing missing keys, unexpected keys and error messages.
|
||||
local_files_only(`bool`, *optional*, defaults to `False`):
|
||||
Whether or not to only look at local files (i.e., do not try to download the model).
|
||||
use_auth_token (`str` or *bool*, *optional*):
|
||||
The token to use as HTTP bearer authorization for remote files. If `True`, will use the token generated
|
||||
when running `diffusers-cli login` (stored in `~/.huggingface`).
|
||||
revision (`str`, *optional*, defaults to `"main"`):
|
||||
The specific model version to use. It can be a branch name, a tag name, or a commit id, since we use a
|
||||
git-based system for storing models and other artifacts on huggingface.co, so `revision` can be any
|
||||
identifier allowed by git.
|
||||
from_flax (`bool`, *optional*, defaults to `False`):
|
||||
Load the model weights from a Flax checkpoint save file.
|
||||
subfolder (`str`, *optional*, defaults to `""`):
|
||||
In case the relevant files are located inside a subfolder of the model repo (either remote in
|
||||
huggingface.co or downloaded locally), you can specify the folder name here.
|
||||
|
||||
mirror (`str`, *optional*):
|
||||
Mirror source to accelerate downloads in China. If you are from China and have an accessibility
|
||||
problem, you can set this option to resolve it. Note that we do not guarantee the timeliness or safety.
|
||||
Please refer to the mirror site for more information.
|
||||
device_map (`str` or `Dict[str, Union[int, str, torch.device]]`, *optional*):
|
||||
A map that specifies where each submodule should go. It doesn't need to be refined to each
|
||||
parameter/buffer name, once a given module name is inside, every submodule of it will be sent to the
|
||||
same device.
|
||||
|
||||
To have Accelerate compute the most optimized `device_map` automatically, set `device_map="auto"`. For
|
||||
more information about each option see [designing a device
|
||||
map](https://hf.co/docs/accelerate/main/en/usage_guides/big_modeling#designing-a-device-map).
|
||||
low_cpu_mem_usage (`bool`, *optional*, defaults to `True` if torch version >= 1.9.0 else `False`):
|
||||
Speed up model loading by not initializing the weights and only loading the pre-trained weights. This
|
||||
also tries to not use more than 1x model size in CPU memory (including peak memory) while loading the
|
||||
model. This is only supported when torch version >= 1.9.0. If you are using an older version of torch,
|
||||
setting this argument to `True` will raise an error.
|
||||
variant (`str`, *optional*):
|
||||
If specified load weights from `variant` filename, *e.g.* pytorch_model.<variant>.bin. `variant` is
|
||||
ignored when using `from_flax`.
|
||||
use_safetensors (`bool`, *optional* ):
|
||||
If set to `True`, the pipeline will forcibly load the models from `safetensors` weights. If set to
|
||||
`None` (the default). The pipeline will load using `safetensors` if safetensors weights are available
|
||||
*and* if `safetensors` is installed. If the to `False` the pipeline will *not* use `safetensors`.
|
||||
|
||||
<Tip>
|
||||
|
||||
It is required to be logged in (`huggingface-cli login`) when you want to use private or [gated
|
||||
models](https://huggingface.co/docs/hub/models-gated#gated-models).
|
||||
|
||||
</Tip>
|
||||
|
||||
<Tip>
|
||||
|
||||
Activate the special ["offline-mode"](https://huggingface.co/diffusers/installation.html#offline-mode) to use
|
||||
this method in a firewalled environment.
|
||||
|
||||
</Tip>
|
||||
|
||||
"""
|
||||
cache_dir = kwargs.pop("cache_dir", DIFFUSERS_CACHE)
|
||||
ignore_mismatched_sizes = kwargs.pop("ignore_mismatched_sizes", False)
|
||||
force_download = kwargs.pop("force_download", False)
|
||||
from_flax = kwargs.pop("from_flax", False)
|
||||
resume_download = kwargs.pop("resume_download", False)
|
||||
proxies = kwargs.pop("proxies", None)
|
||||
output_loading_info = kwargs.pop("output_loading_info", False)
|
||||
local_files_only = kwargs.pop("local_files_only", HF_HUB_OFFLINE)
|
||||
use_auth_token = kwargs.pop("use_auth_token", None)
|
||||
revision = kwargs.pop("revision", None)
|
||||
torch_dtype = kwargs.pop("torch_dtype", None)
|
||||
subfolder = kwargs.pop("subfolder", None)
|
||||
device_map = kwargs.pop("device_map", None)
|
||||
low_cpu_mem_usage = kwargs.pop("low_cpu_mem_usage", _LOW_CPU_MEM_USAGE_DEFAULT)
|
||||
variant = kwargs.pop("variant", None)
|
||||
use_safetensors = kwargs.pop("use_safetensors", None)
|
||||
|
||||
if use_safetensors and not is_safetensors_available():
|
||||
raise ValueError(
|
||||
"`use_safetensors`=True but safetensors is not installed. Please install safetensors with `pip install safetenstors"
|
||||
)
|
||||
|
||||
allow_pickle = False
|
||||
if use_safetensors is None:
|
||||
use_safetensors = is_safetensors_available()
|
||||
allow_pickle = True
|
||||
|
||||
if low_cpu_mem_usage and not is_accelerate_available():
|
||||
low_cpu_mem_usage = False
|
||||
logger.warning(
|
||||
"Cannot initialize model with low cpu memory usage because `accelerate` was not found in the"
|
||||
" environment. Defaulting to `low_cpu_mem_usage=False`. It is strongly recommended to install"
|
||||
" `accelerate` for faster and less memory-intense model loading. You can do so with: \n```\npip"
|
||||
" install accelerate\n```\n."
|
||||
)
|
||||
|
||||
if device_map is not None and not is_accelerate_available():
|
||||
raise NotImplementedError(
|
||||
"Loading and dispatching requires `accelerate`. Please make sure to install accelerate or set"
|
||||
" `device_map=None`. You can install accelerate with `pip install accelerate`."
|
||||
)
|
||||
|
||||
# Check if we can handle device_map and dispatching the weights
|
||||
if device_map is not None and not is_torch_version(">=", "1.9.0"):
|
||||
raise NotImplementedError(
|
||||
"Loading and dispatching requires torch >= 1.9.0. Please either update your PyTorch version or set"
|
||||
" `device_map=None`."
|
||||
)
|
||||
|
||||
if low_cpu_mem_usage is True and not is_torch_version(">=", "1.9.0"):
|
||||
raise NotImplementedError(
|
||||
"Low memory initialization requires torch >= 1.9.0. Please either update your PyTorch version or set"
|
||||
" `low_cpu_mem_usage=False`."
|
||||
)
|
||||
|
||||
if low_cpu_mem_usage is False and device_map is not None:
|
||||
raise ValueError(
|
||||
f"You cannot set `low_cpu_mem_usage` to `False` while using device_map={device_map} for loading and"
|
||||
" dispatching. Please make sure to set `low_cpu_mem_usage=True`."
|
||||
)
|
||||
|
||||
# Load config if we don't provide a configuration
|
||||
config_path = pretrained_model_name_or_path
|
||||
|
||||
user_agent = {
|
||||
"diffusers": __version__,
|
||||
"file_type": "model",
|
||||
"framework": "pytorch",
|
||||
}
|
||||
|
||||
# load config
|
||||
config, unused_kwargs, commit_hash = cls.load_config(
|
||||
config_path,
|
||||
cache_dir=cache_dir,
|
||||
return_unused_kwargs=True,
|
||||
return_commit_hash=True,
|
||||
force_download=force_download,
|
||||
resume_download=resume_download,
|
||||
proxies=proxies,
|
||||
local_files_only=local_files_only,
|
||||
use_auth_token=use_auth_token,
|
||||
revision=revision,
|
||||
subfolder=subfolder,
|
||||
device_map=device_map,
|
||||
user_agent=user_agent,
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
# load model
|
||||
model_file = None
|
||||
if from_flax:
|
||||
model_file = _get_model_file(
|
||||
pretrained_model_name_or_path,
|
||||
weights_name=FLAX_WEIGHTS_NAME,
|
||||
cache_dir=cache_dir,
|
||||
force_download=force_download,
|
||||
resume_download=resume_download,
|
||||
proxies=proxies,
|
||||
local_files_only=local_files_only,
|
||||
use_auth_token=use_auth_token,
|
||||
revision=revision,
|
||||
subfolder=subfolder,
|
||||
user_agent=user_agent,
|
||||
commit_hash=commit_hash,
|
||||
)
|
||||
model = cls.from_config(config, **unused_kwargs)
|
||||
|
||||
# Convert the weights
|
||||
from .modeling_pytorch_flax_utils import load_flax_checkpoint_in_pytorch_model
|
||||
|
||||
model = load_flax_checkpoint_in_pytorch_model(model, model_file)
|
||||
else:
|
||||
if use_safetensors:
|
||||
try:
|
||||
model_file = _get_model_file(
|
||||
pretrained_model_name_or_path,
|
||||
weights_name=_add_variant(SAFETENSORS_WEIGHTS_NAME, variant),
|
||||
cache_dir=cache_dir,
|
||||
force_download=force_download,
|
||||
resume_download=resume_download,
|
||||
proxies=proxies,
|
||||
local_files_only=local_files_only,
|
||||
use_auth_token=use_auth_token,
|
||||
revision=revision,
|
||||
subfolder=subfolder,
|
||||
user_agent=user_agent,
|
||||
commit_hash=commit_hash,
|
||||
)
|
||||
except IOError as e:
|
||||
if not allow_pickle:
|
||||
raise e
|
||||
pass
|
||||
if model_file is None: # deactivate low_cpu_mem_usage mode
|
||||
model_file = _get_model_file(
|
||||
pretrained_model_name_or_path,
|
||||
weights_name=_add_variant(WEIGHTS_NAME, variant),
|
||||
cache_dir=cache_dir,
|
||||
force_download=force_download,
|
||||
resume_download=resume_download,
|
||||
proxies=proxies,
|
||||
local_files_only=local_files_only,
|
||||
use_auth_token=use_auth_token,
|
||||
revision=revision,
|
||||
subfolder=subfolder,
|
||||
user_agent=user_agent,
|
||||
commit_hash=commit_hash,
|
||||
)
|
||||
|
||||
model = cls.from_config(config, **unused_kwargs)
|
||||
|
||||
state_dict = load_state_dict(model_file, variant=variant)
|
||||
|
||||
model, missing_keys, unexpected_keys, mismatched_keys, error_msgs = cls._load_pretrained_model(
|
||||
model,
|
||||
state_dict,
|
||||
model_file,
|
||||
pretrained_model_name_or_path,
|
||||
ignore_mismatched_sizes=True,
|
||||
)
|
||||
|
||||
loading_info = {
|
||||
"missing_keys": missing_keys,
|
||||
"unexpected_keys": unexpected_keys,
|
||||
"mismatched_keys": mismatched_keys,
|
||||
"error_msgs": error_msgs,
|
||||
}
|
||||
|
||||
if torch_dtype is not None and not isinstance(torch_dtype, torch.dtype):
|
||||
raise ValueError(
|
||||
f"{torch_dtype} needs to be of type `torch.dtype`, e.g. `torch.float16`, but is {type(torch_dtype)}."
|
||||
)
|
||||
elif torch_dtype is not None:
|
||||
model = model.to(torch_dtype)
|
||||
|
||||
model.register_to_config(_name_or_path=pretrained_model_name_or_path)
|
||||
|
||||
# Set model in evaluation mode to deactivate DropOut modules by default
|
||||
model.eval()
|
||||
if output_loading_info:
|
||||
return model, loading_info
|
||||
|
||||
return model
|
||||
|
||||
@classmethod
|
||||
def _load_pretrained_model(
|
||||
cls,
|
||||
model,
|
||||
state_dict,
|
||||
resolved_archive_file,
|
||||
pretrained_model_name_or_path,
|
||||
ignore_mismatched_sizes=False,
|
||||
):
|
||||
# Retrieve missing & unexpected_keys
|
||||
model_state_dict = model.state_dict()
|
||||
loaded_keys = list(state_dict.keys())
|
||||
|
||||
expected_keys = list(model_state_dict.keys())
|
||||
|
||||
original_loaded_keys = loaded_keys
|
||||
|
||||
missing_keys = list(set(expected_keys) - set(loaded_keys))
|
||||
unexpected_keys = list(set(loaded_keys) - set(expected_keys))
|
||||
|
||||
# Make sure we are able to load base models as well as derived models (with heads)
|
||||
model_to_load = model
|
||||
|
||||
def _find_mismatched_keys(
|
||||
state_dict,
|
||||
model_state_dict,
|
||||
loaded_keys,
|
||||
ignore_mismatched_sizes,
|
||||
):
|
||||
mismatched_keys = []
|
||||
if ignore_mismatched_sizes:
|
||||
for checkpoint_key in loaded_keys:
|
||||
model_key = checkpoint_key
|
||||
|
||||
if (
|
||||
model_key in model_state_dict
|
||||
and state_dict[checkpoint_key].shape != model_state_dict[model_key].shape
|
||||
):
|
||||
mismatched_keys.append(
|
||||
(checkpoint_key, state_dict[checkpoint_key].shape, model_state_dict[model_key].shape)
|
||||
)
|
||||
del state_dict[checkpoint_key]
|
||||
return mismatched_keys
|
||||
|
||||
if state_dict is not None:
|
||||
# Whole checkpoint
|
||||
mismatched_keys = _find_mismatched_keys(
|
||||
state_dict,
|
||||
model_state_dict,
|
||||
original_loaded_keys,
|
||||
ignore_mismatched_sizes,
|
||||
)
|
||||
error_msgs = _load_state_dict_into_model(model_to_load, state_dict)
|
||||
|
||||
if len(error_msgs) > 0:
|
||||
error_msg = "\n\t".join(error_msgs)
|
||||
if "size mismatch" in error_msg:
|
||||
error_msg += (
|
||||
"\n\tYou may consider adding `ignore_mismatched_sizes=True` in the model `from_pretrained` method."
|
||||
)
|
||||
raise RuntimeError(f"Error(s) in loading state_dict for {model.__class__.__name__}:\n\t{error_msg}")
|
||||
|
||||
if len(unexpected_keys) > 0:
|
||||
logger.warning(
|
||||
f"Some weights of the model checkpoint at {pretrained_model_name_or_path} were not used when"
|
||||
f" initializing {model.__class__.__name__}: {unexpected_keys}\n- This IS expected if you are"
|
||||
f" initializing {model.__class__.__name__} from the checkpoint of a model trained on another task"
|
||||
" or with another architecture (e.g. initializing a BertForSequenceClassification model from a"
|
||||
" BertForPreTraining model).\n- This IS NOT expected if you are initializing"
|
||||
f" {model.__class__.__name__} from the checkpoint of a model that you expect to be exactly"
|
||||
" identical (initializing a BertForSequenceClassification model from a"
|
||||
" BertForSequenceClassification model)."
|
||||
)
|
||||
else:
|
||||
logger.info(f"All model checkpoint weights were used when initializing {model.__class__.__name__}.\n")
|
||||
if len(missing_keys) > 0:
|
||||
logger.warning(
|
||||
f"Some weights of {model.__class__.__name__} were not initialized from the model checkpoint at"
|
||||
f" {pretrained_model_name_or_path} and are newly initialized: {missing_keys}\nYou should probably"
|
||||
" TRAIN this model on a down-stream task to be able to use it for predictions and inference."
|
||||
)
|
||||
elif len(mismatched_keys) == 0:
|
||||
logger.info(
|
||||
f"All the weights of {model.__class__.__name__} were initialized from the model checkpoint at"
|
||||
f" {pretrained_model_name_or_path}.\nIf your task is similar to the task the model of the"
|
||||
f" checkpoint was trained on, you can already use {model.__class__.__name__} for predictions"
|
||||
" without further training."
|
||||
)
|
||||
if len(mismatched_keys) > 0:
|
||||
mismatched_warning = "\n".join(
|
||||
[
|
||||
f"- {key}: found shape {shape1} in the checkpoint and {shape2} in the model instantiated"
|
||||
for key, shape1, shape2 in mismatched_keys
|
||||
]
|
||||
)
|
||||
logger.warning(
|
||||
f"Some weights of {model.__class__.__name__} were not initialized from the model checkpoint at"
|
||||
f" {pretrained_model_name_or_path} and are newly initialized because the shapes did not"
|
||||
f" match:\n{mismatched_warning}\nYou should probably TRAIN this model on a down-stream task to be"
|
||||
" able to use it for predictions and inference."
|
||||
)
|
||||
|
||||
return model, missing_keys, unexpected_keys, mismatched_keys, error_msgs
|
||||
|
||||
@property
|
||||
def device(self) -> device:
|
||||
"""
|
||||
`torch.device`: The device on which the module is (assuming that all the module parameters are on the same
|
||||
device).
|
||||
"""
|
||||
return get_parameter_device(self)
|
||||
|
||||
@property
|
||||
def dtype(self) -> torch.dtype:
|
||||
"""
|
||||
`torch.dtype`: The dtype of the module (assuming that all the module parameters have the same dtype).
|
||||
"""
|
||||
return get_parameter_dtype(self)
|
||||
|
||||
def num_parameters(self, only_trainable: bool = False, exclude_embeddings: bool = False) -> int:
|
||||
"""
|
||||
Get number of (optionally, trainable or non-embeddings) parameters in the module.
|
||||
|
||||
Args:
|
||||
only_trainable (`bool`, *optional*, defaults to `False`):
|
||||
Whether or not to return only the number of trainable parameters
|
||||
|
||||
exclude_embeddings (`bool`, *optional*, defaults to `False`):
|
||||
Whether or not to return only the number of non-embeddings parameters
|
||||
|
||||
Returns:
|
||||
`int`: The number of parameters.
|
||||
"""
|
||||
|
||||
if exclude_embeddings:
|
||||
embedding_param_names = [
|
||||
f"{name}.weight"
|
||||
for name, module_type in self.named_modules()
|
||||
if isinstance(module_type, torch.nn.Embedding)
|
||||
]
|
||||
non_embedding_parameters = [
|
||||
parameter for name, parameter in self.named_parameters() if name not in embedding_param_names
|
||||
]
|
||||
return sum(p.numel() for p in non_embedding_parameters if p.requires_grad or not only_trainable)
|
||||
else:
|
||||
return sum(p.numel() for p in self.parameters() if p.requires_grad or not only_trainable)
|
||||
|
||||
|
||||
def _get_model_file(
|
||||
pretrained_model_name_or_path,
|
||||
*,
|
||||
weights_name,
|
||||
subfolder,
|
||||
cache_dir,
|
||||
force_download,
|
||||
proxies,
|
||||
resume_download,
|
||||
local_files_only,
|
||||
use_auth_token,
|
||||
user_agent,
|
||||
revision,
|
||||
commit_hash=None,
|
||||
):
|
||||
pretrained_model_name_or_path = str(pretrained_model_name_or_path)
|
||||
if os.path.isfile(pretrained_model_name_or_path):
|
||||
return pretrained_model_name_or_path
|
||||
elif os.path.isdir(pretrained_model_name_or_path):
|
||||
if os.path.isfile(os.path.join(pretrained_model_name_or_path, weights_name)):
|
||||
# Load from a PyTorch checkpoint
|
||||
model_file = os.path.join(pretrained_model_name_or_path, weights_name)
|
||||
return model_file
|
||||
elif subfolder is not None and os.path.isfile(
|
||||
os.path.join(pretrained_model_name_or_path, subfolder, weights_name)
|
||||
):
|
||||
model_file = os.path.join(pretrained_model_name_or_path, subfolder, weights_name)
|
||||
return model_file
|
||||
else:
|
||||
raise EnvironmentError(
|
||||
f"Error no file named {weights_name} found in directory {pretrained_model_name_or_path}."
|
||||
)
|
||||
else:
|
||||
# 1. First check if deprecated way of loading from branches is used
|
||||
if (
|
||||
revision in DEPRECATED_REVISION_ARGS
|
||||
and (weights_name == WEIGHTS_NAME or weights_name == SAFETENSORS_WEIGHTS_NAME)
|
||||
and version.parse(version.parse(__version__).base_version) >= version.parse("0.17.0")
|
||||
):
|
||||
try:
|
||||
model_file = hf_hub_download(
|
||||
pretrained_model_name_or_path,
|
||||
filename=_add_variant(weights_name, revision),
|
||||
cache_dir=cache_dir,
|
||||
force_download=force_download,
|
||||
proxies=proxies,
|
||||
resume_download=resume_download,
|
||||
local_files_only=local_files_only,
|
||||
use_auth_token=use_auth_token,
|
||||
user_agent=user_agent,
|
||||
subfolder=subfolder,
|
||||
revision=revision or commit_hash,
|
||||
)
|
||||
warnings.warn(
|
||||
f"Loading the variant {revision} from {pretrained_model_name_or_path} via `revision='{revision}'` is deprecated. Loading instead from `revision='main'` with `variant={revision}`. Loading model variants via `revision='{revision}'` will be removed in diffusers v1. Please use `variant='{revision}'` instead.",
|
||||
FutureWarning,
|
||||
)
|
||||
return model_file
|
||||
except: # noqa: E722
|
||||
warnings.warn(
|
||||
f"You are loading the variant {revision} from {pretrained_model_name_or_path} via `revision='{revision}'`. This behavior is deprecated and will be removed in diffusers v1. One should use `variant='{revision}'` instead. However, it appears that {pretrained_model_name_or_path} currently does not have a {_add_variant(weights_name, revision)} file in the 'main' branch of {pretrained_model_name_or_path}. \n The Diffusers team and community would be very grateful if you could open an issue: https://github.com/huggingface/diffusers/issues/new with the title '{pretrained_model_name_or_path} is missing {_add_variant(weights_name, revision)}' so that the correct variant file can be added.",
|
||||
FutureWarning,
|
||||
)
|
||||
try:
|
||||
# 2. Load model file as usual
|
||||
model_file = hf_hub_download(
|
||||
pretrained_model_name_or_path,
|
||||
filename=weights_name,
|
||||
cache_dir=cache_dir,
|
||||
force_download=force_download,
|
||||
proxies=proxies,
|
||||
resume_download=resume_download,
|
||||
local_files_only=local_files_only,
|
||||
use_auth_token=use_auth_token,
|
||||
user_agent=user_agent,
|
||||
subfolder=subfolder,
|
||||
revision=revision or commit_hash,
|
||||
)
|
||||
return model_file
|
||||
|
||||
except RepositoryNotFoundError:
|
||||
raise EnvironmentError(
|
||||
f"{pretrained_model_name_or_path} is not a local folder and is not a valid model identifier "
|
||||
"listed on 'https://huggingface.co/models'\nIf this is a private repository, make sure to pass a "
|
||||
"token having permission to this repo with `use_auth_token` or log in with `huggingface-cli "
|
||||
"login`."
|
||||
)
|
||||
except RevisionNotFoundError:
|
||||
raise EnvironmentError(
|
||||
f"{revision} is not a valid git identifier (branch name, tag name or commit id) that exists for "
|
||||
"this model name. Check the model page at "
|
||||
f"'https://huggingface.co/{pretrained_model_name_or_path}' for available revisions."
|
||||
)
|
||||
except EntryNotFoundError:
|
||||
raise EnvironmentError(
|
||||
f"{pretrained_model_name_or_path} does not appear to have a file named {weights_name}."
|
||||
)
|
||||
except HTTPError as err:
|
||||
raise EnvironmentError(
|
||||
f"There was a specific connection error when trying to load {pretrained_model_name_or_path}:\n{err}"
|
||||
)
|
||||
except ValueError:
|
||||
raise EnvironmentError(
|
||||
f"We couldn't connect to '{HUGGINGFACE_CO_RESOLVE_ENDPOINT}' to load this model, couldn't find it"
|
||||
f" in the cached files and it looks like {pretrained_model_name_or_path} is not the path to a"
|
||||
f" directory containing a file named {weights_name} or"
|
||||
" \nCheckout your internet connection or see how to run the library in"
|
||||
" offline mode at 'https://huggingface.co/docs/diffusers/installation#offline-mode'."
|
||||
)
|
||||
except EnvironmentError:
|
||||
raise EnvironmentError(
|
||||
f"Can't load the model for '{pretrained_model_name_or_path}'. If you were trying to load it from "
|
||||
"'https://huggingface.co/models', make sure you don't have a local directory with the same name. "
|
||||
f"Otherwise, make sure '{pretrained_model_name_or_path}' is the correct path to a directory "
|
||||
f"containing a file named {weights_name}"
|
||||
)
|
||||
@@ -0,0 +1,402 @@
|
||||
# Copyright 2023 UC Berkeley Team and The HuggingFace Team. All rights reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
# DISCLAIMER: This file is strongly influenced by https://github.com/ermongroup/ddim
|
||||
|
||||
import math
|
||||
from dataclasses import dataclass
|
||||
from typing import List, Optional, Tuple, Union
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
from ..configuration_utils import ConfigMixin, register_to_config
|
||||
from ..utils import BaseOutput, randn_tensor
|
||||
from .scheduling_utils import KarrasDiffusionSchedulers, SchedulerMixin
|
||||
|
||||
|
||||
@dataclass
|
||||
class DDPMSchedulerOutput(BaseOutput):
|
||||
"""
|
||||
Output class for the scheduler's step function output.
|
||||
|
||||
Args:
|
||||
prev_sample (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)` for images):
|
||||
Computed sample (x_{t-1}) of previous timestep. `prev_sample` should be used as next model input in the
|
||||
denoising loop.
|
||||
pred_original_sample (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)` for images):
|
||||
The predicted denoised sample (x_{0}) based on the model output from the current timestep.
|
||||
`pred_original_sample` can be used to preview progress or for guidance.
|
||||
"""
|
||||
|
||||
prev_sample: torch.FloatTensor
|
||||
pred_original_sample: Optional[torch.FloatTensor] = None
|
||||
|
||||
|
||||
def betas_for_alpha_bar(num_diffusion_timesteps, max_beta=0.999):
|
||||
"""
|
||||
Create a beta schedule that discretizes the given alpha_t_bar function, which defines the cumulative product of
|
||||
(1-beta) over time from t = [0,1].
|
||||
|
||||
Contains a function alpha_bar that takes an argument t and transforms it to the cumulative product of (1-beta) up
|
||||
to that part of the diffusion process.
|
||||
|
||||
|
||||
Args:
|
||||
num_diffusion_timesteps (`int`): the number of betas to produce.
|
||||
max_beta (`float`): the maximum beta to use; use values lower than 1 to
|
||||
prevent singularities.
|
||||
|
||||
Returns:
|
||||
betas (`np.ndarray`): the betas used by the scheduler to step the model outputs
|
||||
"""
|
||||
|
||||
def alpha_bar(time_step):
|
||||
return math.cos((time_step + 0.008) / 1.008 * math.pi / 2) ** 2
|
||||
|
||||
betas = []
|
||||
for i in range(num_diffusion_timesteps):
|
||||
t1 = i / num_diffusion_timesteps
|
||||
t2 = (i + 1) / num_diffusion_timesteps
|
||||
betas.append(min(1 - alpha_bar(t2) / alpha_bar(t1), max_beta))
|
||||
return torch.tensor(betas, dtype=torch.float32)
|
||||
|
||||
|
||||
class DDPMScheduler(SchedulerMixin, ConfigMixin):
|
||||
"""
|
||||
Denoising diffusion probabilistic models (DDPMs) explores the connections between denoising score matching and
|
||||
Langevin dynamics sampling.
|
||||
|
||||
[`~ConfigMixin`] takes care of storing all config attributes that are passed in the scheduler's `__init__`
|
||||
function, such as `num_train_timesteps`. They can be accessed via `scheduler.config.num_train_timesteps`.
|
||||
[`SchedulerMixin`] provides general loading and saving functionality via the [`SchedulerMixin.save_pretrained`] and
|
||||
[`~SchedulerMixin.from_pretrained`] functions.
|
||||
|
||||
For more details, see the original paper: https://arxiv.org/abs/2006.11239
|
||||
|
||||
Args:
|
||||
num_train_timesteps (`int`): number of diffusion steps used to train the model.
|
||||
beta_start (`float`): the starting `beta` value of inference.
|
||||
beta_end (`float`): the final `beta` value.
|
||||
beta_schedule (`str`):
|
||||
the beta schedule, a mapping from a beta range to a sequence of betas for stepping the model. Choose from
|
||||
`linear`, `scaled_linear`, or `squaredcos_cap_v2`.
|
||||
trained_betas (`np.ndarray`, optional):
|
||||
option to pass an array of betas directly to the constructor to bypass `beta_start`, `beta_end` etc.
|
||||
variance_type (`str`):
|
||||
options to clip the variance used when adding noise to the denoised sample. Choose from `fixed_small`,
|
||||
`fixed_small_log`, `fixed_large`, `fixed_large_log`, `learned` or `learned_range`.
|
||||
clip_sample (`bool`, default `True`):
|
||||
option to clip predicted sample for numerical stability.
|
||||
clip_sample_range (`float`, default `1.0`):
|
||||
the maximum magnitude for sample clipping. Valid only when `clip_sample=True`.
|
||||
prediction_type (`str`, default `epsilon`, optional):
|
||||
prediction type of the scheduler function, one of `epsilon` (predicting the noise of the diffusion
|
||||
process), `sample` (directly predicting the noisy sample`) or `v_prediction` (see section 2.4
|
||||
https://imagen.research.google/video/paper.pdf)
|
||||
thresholding (`bool`, default `False`):
|
||||
whether to use the "dynamic thresholding" method (introduced by Imagen, https://arxiv.org/abs/2205.11487).
|
||||
Note that the thresholding method is unsuitable for latent-space diffusion models (such as
|
||||
stable-diffusion).
|
||||
dynamic_thresholding_ratio (`float`, default `0.995`):
|
||||
the ratio for the dynamic thresholding method. Default is `0.995`, the same as Imagen
|
||||
(https://arxiv.org/abs/2205.11487). Valid only when `thresholding=True`.
|
||||
sample_max_value (`float`, default `1.0`):
|
||||
the threshold value for dynamic thresholding. Valid only when `thresholding=True`.
|
||||
"""
|
||||
|
||||
_compatibles = [e.name for e in KarrasDiffusionSchedulers]
|
||||
order = 1
|
||||
|
||||
@register_to_config
|
||||
def __init__(
|
||||
self,
|
||||
num_train_timesteps: int = 1000,
|
||||
beta_start: float = 0.0001,
|
||||
beta_end: float = 0.02,
|
||||
beta_schedule: str = "linear",
|
||||
trained_betas: Optional[Union[np.ndarray, List[float]]] = None,
|
||||
variance_type: str = "fixed_small",
|
||||
clip_sample: bool = True,
|
||||
prediction_type: str = "epsilon",
|
||||
thresholding: bool = False,
|
||||
dynamic_thresholding_ratio: float = 0.995,
|
||||
clip_sample_range: float = 1.0,
|
||||
sample_max_value: float = 1.0,
|
||||
):
|
||||
if trained_betas is not None:
|
||||
self.betas = torch.tensor(trained_betas, dtype=torch.float32)
|
||||
elif beta_schedule == "linear":
|
||||
self.betas = torch.linspace(beta_start, beta_end, num_train_timesteps, dtype=torch.float32)
|
||||
elif beta_schedule == "scaled_linear":
|
||||
# this schedule is very specific to the latent diffusion model.
|
||||
self.betas = (
|
||||
torch.linspace(beta_start**0.5, beta_end**0.5, num_train_timesteps, dtype=torch.float32) ** 2
|
||||
)
|
||||
elif beta_schedule == "squaredcos_cap_v2":
|
||||
# Glide cosine schedule
|
||||
self.betas = betas_for_alpha_bar(num_train_timesteps)
|
||||
elif beta_schedule == "sigmoid":
|
||||
# GeoDiff sigmoid schedule
|
||||
betas = torch.linspace(-6, 6, num_train_timesteps)
|
||||
self.betas = torch.sigmoid(betas) * (beta_end - beta_start) + beta_start
|
||||
else:
|
||||
raise NotImplementedError(f"{beta_schedule} does is not implemented for {self.__class__}")
|
||||
|
||||
self.alphas = 1.0 - self.betas
|
||||
self.alphas_cumprod = torch.cumprod(self.alphas, dim=0)
|
||||
self.one = torch.tensor(1.0)
|
||||
|
||||
# standard deviation of the initial noise distribution
|
||||
self.init_noise_sigma = 1.0
|
||||
|
||||
# setable values
|
||||
self.num_inference_steps = None
|
||||
self.timesteps = torch.from_numpy(np.arange(0, num_train_timesteps)[::-1].copy())
|
||||
|
||||
self.variance_type = variance_type
|
||||
|
||||
def scale_model_input(self, sample: torch.FloatTensor, timestep: Optional[int] = None) -> torch.FloatTensor:
|
||||
"""
|
||||
Ensures interchangeability with schedulers that need to scale the denoising model input depending on the
|
||||
current timestep.
|
||||
|
||||
Args:
|
||||
sample (`torch.FloatTensor`): input sample
|
||||
timestep (`int`, optional): current timestep
|
||||
|
||||
Returns:
|
||||
`torch.FloatTensor`: scaled input sample
|
||||
"""
|
||||
return sample
|
||||
|
||||
def set_timesteps(self, num_inference_steps: int, device: Union[str, torch.device] = None):
|
||||
"""
|
||||
Sets the discrete timesteps used for the diffusion chain. Supporting function to be run before inference.
|
||||
|
||||
Args:
|
||||
num_inference_steps (`int`):
|
||||
the number of diffusion steps used when generating samples with a pre-trained model.
|
||||
"""
|
||||
|
||||
if num_inference_steps > self.config.num_train_timesteps:
|
||||
raise ValueError(
|
||||
f"`num_inference_steps`: {num_inference_steps} cannot be larger than `self.config.train_timesteps`:"
|
||||
f" {self.config.num_train_timesteps} as the unet model trained with this scheduler can only handle"
|
||||
f" maximal {self.config.num_train_timesteps} timesteps."
|
||||
)
|
||||
|
||||
self.num_inference_steps = num_inference_steps
|
||||
|
||||
step_ratio = self.config.num_train_timesteps // self.num_inference_steps
|
||||
timesteps = (np.arange(0, num_inference_steps) * step_ratio).round()[::-1].copy().astype(np.int64)
|
||||
# print(timesteps)
|
||||
# exit(0)
|
||||
self.timesteps = torch.from_numpy(timesteps).to(device)
|
||||
|
||||
def _get_variance(self, t, predicted_variance=None, variance_type=None):
|
||||
num_inference_steps = self.num_inference_steps if self.num_inference_steps else self.config.num_train_timesteps
|
||||
prev_t = t - self.config.num_train_timesteps // num_inference_steps
|
||||
alpha_prod_t = self.alphas_cumprod[t]
|
||||
alpha_prod_t_prev = self.alphas_cumprod[prev_t] if prev_t >= 0 else self.one
|
||||
current_beta_t = 1 - alpha_prod_t / alpha_prod_t_prev
|
||||
|
||||
# For t > 0, compute predicted variance βt (see formula (6) and (7) from https://arxiv.org/pdf/2006.11239.pdf)
|
||||
# and sample from it to get previous sample
|
||||
# x_{t-1} ~ N(pred_prev_sample, variance) == add variance to pred_sample
|
||||
variance = (1 - alpha_prod_t_prev) / (1 - alpha_prod_t) * current_beta_t
|
||||
|
||||
if variance_type is None:
|
||||
variance_type = self.config.variance_type
|
||||
|
||||
# hacks - were probably added for training stability
|
||||
if variance_type == "fixed_small":
|
||||
variance = torch.clamp(variance, min=1e-20)
|
||||
# for rl-diffuser https://arxiv.org/abs/2205.09991
|
||||
elif variance_type == "fixed_small_log":
|
||||
variance = torch.log(torch.clamp(variance, min=1e-20))
|
||||
variance = torch.exp(0.5 * variance)
|
||||
elif variance_type == "fixed_large":
|
||||
variance = current_beta_t
|
||||
elif variance_type == "fixed_large_log":
|
||||
# Glide max_log
|
||||
variance = torch.log(current_beta_t)
|
||||
elif variance_type == "learned":
|
||||
return predicted_variance
|
||||
elif variance_type == "learned_range":
|
||||
min_log = torch.log(variance)
|
||||
max_log = torch.log(self.betas[t])
|
||||
frac = (predicted_variance + 1) / 2
|
||||
variance = frac * max_log + (1 - frac) * min_log
|
||||
|
||||
return variance
|
||||
|
||||
def _threshold_sample(self, sample: torch.FloatTensor) -> torch.FloatTensor:
|
||||
# Dynamic thresholding in https://arxiv.org/abs/2205.11487
|
||||
dynamic_max_val = (
|
||||
sample.flatten(1)
|
||||
.abs()
|
||||
.quantile(self.config.dynamic_thresholding_ratio, dim=1)
|
||||
.clamp_min(self.config.sample_max_value)
|
||||
.view(-1, *([1] * (sample.ndim - 1)))
|
||||
)
|
||||
return sample.clamp(-dynamic_max_val, dynamic_max_val) / dynamic_max_val
|
||||
|
||||
def step(
|
||||
self,
|
||||
model_output: torch.FloatTensor,
|
||||
timestep: int,
|
||||
sample: torch.FloatTensor,
|
||||
generator=None,
|
||||
return_dict: bool = True,
|
||||
) -> Union[DDPMSchedulerOutput, Tuple]:
|
||||
"""
|
||||
Predict the sample at the previous timestep by reversing the SDE. Core function to propagate the diffusion
|
||||
process from the learned model outputs (most often the predicted noise).
|
||||
|
||||
Args:
|
||||
model_output (`torch.FloatTensor`): direct output from learned diffusion model.
|
||||
timestep (`int`): current discrete timestep in the diffusion chain.
|
||||
sample (`torch.FloatTensor`):
|
||||
current instance of sample being created by diffusion process.
|
||||
generator: random number generator.
|
||||
return_dict (`bool`): option for returning tuple rather than DDPMSchedulerOutput class
|
||||
|
||||
Returns:
|
||||
[`~schedulers.scheduling_utils.DDPMSchedulerOutput`] or `tuple`:
|
||||
[`~schedulers.scheduling_utils.DDPMSchedulerOutput`] if `return_dict` is True, otherwise a `tuple`. When
|
||||
returning a tuple, the first element is the sample tensor.
|
||||
|
||||
"""
|
||||
t = timestep
|
||||
num_inference_steps = self.num_inference_steps if self.num_inference_steps else self.config.num_train_timesteps
|
||||
prev_t = timestep - self.config.num_train_timesteps // num_inference_steps
|
||||
|
||||
if model_output.shape[1] == sample.shape[1] * 2 and self.variance_type in ["learned", "learned_range"]:
|
||||
model_output, predicted_variance = torch.split(model_output, sample.shape[1], dim=1)
|
||||
else:
|
||||
predicted_variance = None
|
||||
|
||||
# 1. compute alphas, betas
|
||||
alpha_prod_t = self.alphas_cumprod[t]
|
||||
alpha_prod_t_prev = self.alphas_cumprod[prev_t] if prev_t >= 0 else self.one
|
||||
beta_prod_t = 1 - alpha_prod_t
|
||||
beta_prod_t_prev = 1 - alpha_prod_t_prev
|
||||
current_alpha_t = alpha_prod_t / alpha_prod_t_prev
|
||||
current_beta_t = 1 - current_alpha_t
|
||||
|
||||
# 2. compute predicted original sample from predicted noise also called
|
||||
# "predicted x_0" of formula (15) from https://arxiv.org/pdf/2006.11239.pdf
|
||||
if self.config.prediction_type == "epsilon":
|
||||
pred_original_sample = (sample - beta_prod_t ** (0.5) * model_output) / alpha_prod_t ** (0.5)
|
||||
elif self.config.prediction_type == "sample":
|
||||
pred_original_sample = model_output
|
||||
elif self.config.prediction_type == "v_prediction":
|
||||
pred_original_sample = (alpha_prod_t**0.5) * sample - (beta_prod_t**0.5) * model_output
|
||||
else:
|
||||
raise ValueError(
|
||||
f"prediction_type given as {self.config.prediction_type} must be one of `epsilon`, `sample` or"
|
||||
" `v_prediction` for the DDPMScheduler."
|
||||
)
|
||||
|
||||
# 3. Clip or threshold "predicted x_0"
|
||||
if self.config.clip_sample:
|
||||
pred_original_sample = pred_original_sample.clamp(
|
||||
-self.config.clip_sample_range, self.config.clip_sample_range
|
||||
)
|
||||
|
||||
if self.config.thresholding:
|
||||
pred_original_sample = self._threshold_sample(pred_original_sample)
|
||||
|
||||
# 4. Compute coefficients for pred_original_sample x_0 and current sample x_t
|
||||
# See formula (7) from https://arxiv.org/pdf/2006.11239.pdf
|
||||
pred_original_sample_coeff = (alpha_prod_t_prev ** (0.5) * current_beta_t) / beta_prod_t
|
||||
current_sample_coeff = current_alpha_t ** (0.5) * beta_prod_t_prev / beta_prod_t
|
||||
|
||||
# 5. Compute predicted previous sample µ_t
|
||||
# See formula (7) from https://arxiv.org/pdf/2006.11239.pdf
|
||||
pred_prev_sample = pred_original_sample_coeff * pred_original_sample + current_sample_coeff * sample
|
||||
|
||||
# 6. Add noise
|
||||
variance = 0
|
||||
if t > 0:
|
||||
device = model_output.device
|
||||
variance_noise = randn_tensor(
|
||||
model_output.shape, generator=generator, device=device, dtype=model_output.dtype
|
||||
)
|
||||
if self.variance_type == "fixed_small_log":
|
||||
variance = self._get_variance(t, predicted_variance=predicted_variance) * variance_noise
|
||||
elif self.variance_type == "learned_range":
|
||||
variance = self._get_variance(t, predicted_variance=predicted_variance)
|
||||
variance = torch.exp(0.5 * variance) * variance_noise
|
||||
else:
|
||||
variance = (self._get_variance(t, predicted_variance=predicted_variance) ** 0.5) * variance_noise
|
||||
|
||||
pred_prev_sample = pred_prev_sample + variance
|
||||
|
||||
if not return_dict:
|
||||
return (pred_prev_sample,)
|
||||
|
||||
return DDPMSchedulerOutput(prev_sample=pred_prev_sample, pred_original_sample=pred_original_sample)
|
||||
|
||||
def add_noise(
|
||||
self,
|
||||
original_samples: torch.FloatTensor,
|
||||
noise: torch.FloatTensor,
|
||||
timesteps: torch.IntTensor,
|
||||
) -> torch.FloatTensor:
|
||||
# Make sure alphas_cumprod and timestep have same device and dtype as original_samples
|
||||
self.alphas_cumprod = self.alphas_cumprod.to(device=original_samples.device, dtype=original_samples.dtype)
|
||||
timesteps = timesteps.to(original_samples.device)
|
||||
|
||||
sqrt_alpha_prod = self.alphas_cumprod[timesteps] ** 0.5
|
||||
sqrt_alpha_prod = sqrt_alpha_prod.flatten()
|
||||
while len(sqrt_alpha_prod.shape) < len(original_samples.shape):
|
||||
sqrt_alpha_prod = sqrt_alpha_prod.unsqueeze(-1)
|
||||
|
||||
sqrt_one_minus_alpha_prod = (1 - self.alphas_cumprod[timesteps]) ** 0.5
|
||||
sqrt_one_minus_alpha_prod = sqrt_one_minus_alpha_prod.flatten()
|
||||
while len(sqrt_one_minus_alpha_prod.shape) < len(original_samples.shape):
|
||||
sqrt_one_minus_alpha_prod = sqrt_one_minus_alpha_prod.unsqueeze(-1)
|
||||
|
||||
noisy_samples = sqrt_alpha_prod * original_samples + sqrt_one_minus_alpha_prod * noise
|
||||
return noisy_samples
|
||||
|
||||
def get_x0_from_noise(self, noise, t, x_t): # add this function
|
||||
self.alphas_cumprod = self.alphas_cumprod.to(device=noise.device, dtype=noise.dtype)
|
||||
x_0 = 1 / torch.sqrt(self.alphas_cumprod[t][:,None,None,None]) * x_t - torch.sqrt(1 / self.alphas_cumprod[t][:,None,None,None] - 1) * noise
|
||||
return x_0
|
||||
|
||||
def get_velocity(
|
||||
self, sample: torch.FloatTensor, noise: torch.FloatTensor, timesteps: torch.IntTensor
|
||||
) -> torch.FloatTensor:
|
||||
# Make sure alphas_cumprod and timestep have same device and dtype as sample
|
||||
self.alphas_cumprod = self.alphas_cumprod.to(device=sample.device, dtype=sample.dtype)
|
||||
timesteps = timesteps.to(sample.device)
|
||||
|
||||
sqrt_alpha_prod = self.alphas_cumprod[timesteps] ** 0.5
|
||||
sqrt_alpha_prod = sqrt_alpha_prod.flatten()
|
||||
while len(sqrt_alpha_prod.shape) < len(sample.shape):
|
||||
sqrt_alpha_prod = sqrt_alpha_prod.unsqueeze(-1)
|
||||
|
||||
sqrt_one_minus_alpha_prod = (1 - self.alphas_cumprod[timesteps]) ** 0.5
|
||||
sqrt_one_minus_alpha_prod = sqrt_one_minus_alpha_prod.flatten()
|
||||
while len(sqrt_one_minus_alpha_prod.shape) < len(sample.shape):
|
||||
sqrt_one_minus_alpha_prod = sqrt_one_minus_alpha_prod.unsqueeze(-1)
|
||||
|
||||
velocity = sqrt_alpha_prod * noise - sqrt_one_minus_alpha_prod * sample
|
||||
return velocity
|
||||
|
||||
def __len__(self):
|
||||
return self.config.num_train_timesteps
|
||||
@@ -0,0 +1,713 @@
|
||||
# Copyright 2023 The HuggingFace Team. All rights reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, Dict, List, Optional, Tuple, Union
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.utils.checkpoint
|
||||
|
||||
from ..configuration_utils import ConfigMixin, register_to_config
|
||||
from ..loaders import UNet2DConditionLoadersMixin
|
||||
from ..utils import BaseOutput, logging
|
||||
from .attention_processor import AttentionProcessor, AttnProcessor
|
||||
from .embeddings import GaussianFourierProjection, TimestepEmbedding, Timesteps
|
||||
from .modeling_utils import ModelMixin
|
||||
from .unet_2d_blocks import (
|
||||
CrossAttnDownBlock2D,
|
||||
CrossAttnUpBlock2D,
|
||||
DownBlock2D,
|
||||
UNetMidBlock2DCrossAttn,
|
||||
UNetMidBlock2DSimpleCrossAttn,
|
||||
UpBlock2D,
|
||||
get_down_block,
|
||||
get_up_block,
|
||||
)
|
||||
|
||||
|
||||
logger = logging.get_logger(__name__) # pylint: disable=invalid-name
|
||||
|
||||
|
||||
@dataclass
|
||||
class UNet2DConditionOutput(BaseOutput):
|
||||
"""
|
||||
Args:
|
||||
sample (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`):
|
||||
Hidden states conditioned on `encoder_hidden_states` input. Output of last layer of model.
|
||||
"""
|
||||
|
||||
sample: torch.FloatTensor
|
||||
|
||||
|
||||
class UNet2DConditionModel(ModelMixin, ConfigMixin, UNet2DConditionLoadersMixin):
|
||||
r"""
|
||||
UNet2DConditionModel is a conditional 2D UNet model that takes in a noisy sample, conditional state, and a timestep
|
||||
and returns sample shaped output.
|
||||
|
||||
This model inherits from [`ModelMixin`]. Check the superclass documentation for the generic methods the library
|
||||
implements for all the models (such as downloading or saving, etc.)
|
||||
|
||||
Parameters:
|
||||
sample_size (`int` or `Tuple[int, int]`, *optional*, defaults to `None`):
|
||||
Height and width of input/output sample.
|
||||
in_channels (`int`, *optional*, defaults to 4): The number of channels in the input sample.
|
||||
out_channels (`int`, *optional*, defaults to 4): The number of channels in the output.
|
||||
center_input_sample (`bool`, *optional*, defaults to `False`): Whether to center the input sample.
|
||||
flip_sin_to_cos (`bool`, *optional*, defaults to `False`):
|
||||
Whether to flip the sin to cos in the time embedding.
|
||||
freq_shift (`int`, *optional*, defaults to 0): The frequency shift to apply to the time embedding.
|
||||
down_block_types (`Tuple[str]`, *optional*, defaults to `("CrossAttnDownBlock2D", "CrossAttnDownBlock2D", "CrossAttnDownBlock2D", "DownBlock2D")`):
|
||||
The tuple of downsample blocks to use.
|
||||
mid_block_type (`str`, *optional*, defaults to `"UNetMidBlock2DCrossAttn"`):
|
||||
The mid block type. Choose from `UNetMidBlock2DCrossAttn` or `UNetMidBlock2DSimpleCrossAttn`, will skip the
|
||||
mid block layer if `None`.
|
||||
up_block_types (`Tuple[str]`, *optional*, defaults to `("UpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D",)`):
|
||||
The tuple of upsample blocks to use.
|
||||
only_cross_attention(`bool` or `Tuple[bool]`, *optional*, default to `False`):
|
||||
Whether to include self-attention in the basic transformer blocks, see
|
||||
[`~models.attention.BasicTransformerBlock`].
|
||||
block_out_channels (`Tuple[int]`, *optional*, defaults to `(320, 640, 1280, 1280)`):
|
||||
The tuple of output channels for each block.
|
||||
layers_per_block (`int`, *optional*, defaults to 2): The number of layers per block.
|
||||
downsample_padding (`int`, *optional*, defaults to 1): The padding to use for the downsampling convolution.
|
||||
mid_block_scale_factor (`float`, *optional*, defaults to 1.0): The scale factor to use for the mid block.
|
||||
act_fn (`str`, *optional*, defaults to `"silu"`): The activation function to use.
|
||||
norm_num_groups (`int`, *optional*, defaults to 32): The number of groups to use for the normalization.
|
||||
If `None`, it will skip the normalization and activation layers in post-processing
|
||||
norm_eps (`float`, *optional*, defaults to 1e-5): The epsilon to use for the normalization.
|
||||
cross_attention_dim (`int` or `Tuple[int]`, *optional*, defaults to 1280):
|
||||
The dimension of the cross attention features.
|
||||
attention_head_dim (`int`, *optional*, defaults to 8): The dimension of the attention heads.
|
||||
resnet_time_scale_shift (`str`, *optional*, defaults to `"default"`): Time scale shift config
|
||||
for resnet blocks, see [`~models.resnet.ResnetBlock2D`]. Choose from `default` or `scale_shift`.
|
||||
class_embed_type (`str`, *optional*, defaults to None):
|
||||
The type of class embedding to use which is ultimately summed with the time embeddings. Choose from `None`,
|
||||
`"timestep"`, `"identity"`, `"projection"`, or `"simple_projection"`.
|
||||
num_class_embeds (`int`, *optional*, defaults to None):
|
||||
Input dimension of the learnable embedding matrix to be projected to `time_embed_dim`, when performing
|
||||
class conditioning with `class_embed_type` equal to `None`.
|
||||
time_embedding_type (`str`, *optional*, default to `positional`):
|
||||
The type of position embedding to use for timesteps. Choose from `positional` or `fourier`.
|
||||
timestep_post_act (`str, *optional*, default to `None`):
|
||||
The second activation function to use in timestep embedding. Choose from `silu`, `mish` and `gelu`.
|
||||
time_cond_proj_dim (`int`, *optional*, default to `None`):
|
||||
The dimension of `cond_proj` layer in timestep embedding.
|
||||
conv_in_kernel (`int`, *optional*, default to `3`): The kernel size of `conv_in` layer.
|
||||
conv_out_kernel (`int`, *optional*, default to `3`): The kernel size of `conv_out` layer.
|
||||
projection_class_embeddings_input_dim (`int`, *optional*): The dimension of the `class_labels` input when
|
||||
using the "projection" `class_embed_type`. Required when using the "projection" `class_embed_type`.
|
||||
class_embeddings_concat (`bool`, *optional*, defaults to `False`): Whether to concatenate the time
|
||||
embeddings with the class embeddings.
|
||||
"""
|
||||
|
||||
_supports_gradient_checkpointing = True
|
||||
|
||||
@register_to_config
|
||||
def __init__(
|
||||
self,
|
||||
sample_size: Optional[int] = None,
|
||||
in_channels: int = 17,
|
||||
out_channels: int = 4,
|
||||
center_input_sample: bool = False,
|
||||
flip_sin_to_cos: bool = True,
|
||||
freq_shift: int = 0,
|
||||
down_block_types: Tuple[str] = (
|
||||
"CrossAttnDownBlock2D",
|
||||
"CrossAttnDownBlock2D",
|
||||
"CrossAttnDownBlock2D",
|
||||
"DownBlock2D",
|
||||
),
|
||||
mid_block_type: Optional[str] = "UNetMidBlock2DCrossAttn",
|
||||
up_block_types: Tuple[str] = ("UpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D"),
|
||||
only_cross_attention: Union[bool, Tuple[bool]] = False,
|
||||
block_out_channels: Tuple[int] = (320, 640, 1280, 1280),
|
||||
layers_per_block: int = 2,
|
||||
downsample_padding: int = 1,
|
||||
mid_block_scale_factor: float = 1,
|
||||
act_fn: str = "silu",
|
||||
norm_num_groups: Optional[int] = 32,
|
||||
norm_eps: float = 1e-5,
|
||||
cross_attention_dim: Union[int, Tuple[int]] = 1280,
|
||||
attention_head_dim: Union[int, Tuple[int]] = 8,
|
||||
dual_cross_attention: bool = False,
|
||||
use_linear_projection: bool = False,
|
||||
class_embed_type: Optional[str] = None,
|
||||
num_class_embeds: Optional[int] = None,
|
||||
upcast_attention: bool = False,
|
||||
resnet_time_scale_shift: str = "default",
|
||||
time_embedding_type: str = "positional",
|
||||
timestep_post_act: Optional[str] = None,
|
||||
time_cond_proj_dim: Optional[int] = None,
|
||||
conv_in_kernel: int = 3,
|
||||
conv_out_kernel: int = 3,
|
||||
projection_class_embeddings_input_dim: Optional[int] = None,
|
||||
class_embeddings_concat: bool = False,
|
||||
):
|
||||
super().__init__()
|
||||
|
||||
# char embedding layer
|
||||
self.word_embedding = nn.Embedding(128, 8)
|
||||
# convolution layer
|
||||
self.segmap_conv = nn.Sequential(
|
||||
nn.Conv2d(8, 32, 3, 1, 1),
|
||||
nn.ReLU(),
|
||||
nn.BatchNorm2d(32),
|
||||
nn.MaxPool2d(2, 2),
|
||||
nn.Conv2d(32, 64, 3, 1, 1),
|
||||
nn.ReLU(),
|
||||
nn.BatchNorm2d(64),
|
||||
nn.MaxPool2d(2, 2),
|
||||
nn.Conv2d(64, 8, 3, 1, 1),
|
||||
)
|
||||
|
||||
self.sample_size = sample_size
|
||||
|
||||
# Check inputs
|
||||
if len(down_block_types) != len(up_block_types):
|
||||
raise ValueError(
|
||||
f"Must provide the same number of `down_block_types` as `up_block_types`. `down_block_types`: {down_block_types}. `up_block_types`: {up_block_types}."
|
||||
)
|
||||
|
||||
if len(block_out_channels) != len(down_block_types):
|
||||
raise ValueError(
|
||||
f"Must provide the same number of `block_out_channels` as `down_block_types`. `block_out_channels`: {block_out_channels}. `down_block_types`: {down_block_types}."
|
||||
)
|
||||
|
||||
if not isinstance(only_cross_attention, bool) and len(only_cross_attention) != len(down_block_types):
|
||||
raise ValueError(
|
||||
f"Must provide the same number of `only_cross_attention` as `down_block_types`. `only_cross_attention`: {only_cross_attention}. `down_block_types`: {down_block_types}."
|
||||
)
|
||||
|
||||
if not isinstance(attention_head_dim, int) and len(attention_head_dim) != len(down_block_types):
|
||||
raise ValueError(
|
||||
f"Must provide the same number of `attention_head_dim` as `down_block_types`. `attention_head_dim`: {attention_head_dim}. `down_block_types`: {down_block_types}."
|
||||
)
|
||||
|
||||
if isinstance(cross_attention_dim, list) and len(cross_attention_dim) != len(down_block_types):
|
||||
raise ValueError(
|
||||
f"Must provide the same number of `cross_attention_dim` as `down_block_types`. `cross_attention_dim`: {cross_attention_dim}. `down_block_types`: {down_block_types}."
|
||||
)
|
||||
|
||||
# input
|
||||
conv_in_padding = (conv_in_kernel - 1) // 2
|
||||
self.conv_in = nn.Conv2d(
|
||||
17, block_out_channels[0], kernel_size=conv_in_kernel, padding=conv_in_padding
|
||||
) # the input channel is modified to 17 (4+8+1+4)
|
||||
|
||||
# time
|
||||
if time_embedding_type == "fourier":
|
||||
time_embed_dim = block_out_channels[0] * 2
|
||||
if time_embed_dim % 2 != 0:
|
||||
raise ValueError(f"`time_embed_dim` should be divisible by 2, but is {time_embed_dim}.")
|
||||
self.time_proj = GaussianFourierProjection(
|
||||
time_embed_dim // 2, set_W_to_weight=False, log=False, flip_sin_to_cos=flip_sin_to_cos
|
||||
)
|
||||
timestep_input_dim = time_embed_dim
|
||||
elif time_embedding_type == "positional":
|
||||
time_embed_dim = block_out_channels[0] * 4
|
||||
|
||||
self.time_proj = Timesteps(block_out_channels[0], flip_sin_to_cos, freq_shift)
|
||||
timestep_input_dim = block_out_channels[0]
|
||||
else:
|
||||
raise ValueError(
|
||||
f"{time_embedding_type} does not exist. Please make sure to use one of `fourier` or `positional`."
|
||||
)
|
||||
|
||||
self.time_embedding = TimestepEmbedding(
|
||||
timestep_input_dim,
|
||||
time_embed_dim,
|
||||
act_fn=act_fn,
|
||||
post_act_fn=timestep_post_act,
|
||||
cond_proj_dim=time_cond_proj_dim,
|
||||
)
|
||||
|
||||
# class embedding
|
||||
if class_embed_type is None and num_class_embeds is not None:
|
||||
self.class_embedding = nn.Embedding(num_class_embeds, time_embed_dim)
|
||||
elif class_embed_type == "timestep":
|
||||
self.class_embedding = TimestepEmbedding(timestep_input_dim, time_embed_dim)
|
||||
elif class_embed_type == "identity":
|
||||
self.class_embedding = nn.Identity(time_embed_dim, time_embed_dim)
|
||||
elif class_embed_type == "projection":
|
||||
if projection_class_embeddings_input_dim is None:
|
||||
raise ValueError(
|
||||
"`class_embed_type`: 'projection' requires `projection_class_embeddings_input_dim` be set"
|
||||
)
|
||||
# The projection `class_embed_type` is the same as the timestep `class_embed_type` except
|
||||
# 1. the `class_labels` inputs are not first converted to sinusoidal embeddings
|
||||
# 2. it projects from an arbitrary input dimension.
|
||||
#
|
||||
# Note that `TimestepEmbedding` is quite general, being mainly linear layers and activations.
|
||||
# When used for embedding actual timesteps, the timesteps are first converted to sinusoidal embeddings.
|
||||
# As a result, `TimestepEmbedding` can be passed arbitrary vectors.
|
||||
self.class_embedding = TimestepEmbedding(projection_class_embeddings_input_dim, time_embed_dim)
|
||||
elif class_embed_type == "simple_projection":
|
||||
if projection_class_embeddings_input_dim is None:
|
||||
raise ValueError(
|
||||
"`class_embed_type`: 'simple_projection' requires `projection_class_embeddings_input_dim` be set"
|
||||
)
|
||||
self.class_embedding = nn.Linear(projection_class_embeddings_input_dim, time_embed_dim)
|
||||
else:
|
||||
self.class_embedding = None
|
||||
|
||||
self.down_blocks = nn.ModuleList([])
|
||||
self.up_blocks = nn.ModuleList([])
|
||||
|
||||
if isinstance(only_cross_attention, bool):
|
||||
only_cross_attention = [only_cross_attention] * len(down_block_types)
|
||||
|
||||
if isinstance(attention_head_dim, int):
|
||||
attention_head_dim = (attention_head_dim,) * len(down_block_types)
|
||||
|
||||
if isinstance(cross_attention_dim, int):
|
||||
cross_attention_dim = (cross_attention_dim,) * len(down_block_types)
|
||||
|
||||
if class_embeddings_concat:
|
||||
# The time embeddings are concatenated with the class embeddings. The dimension of the
|
||||
# time embeddings passed to the down, middle, and up blocks is twice the dimension of the
|
||||
# regular time embeddings
|
||||
blocks_time_embed_dim = time_embed_dim * 2
|
||||
else:
|
||||
blocks_time_embed_dim = time_embed_dim
|
||||
|
||||
# down
|
||||
output_channel = block_out_channels[0]
|
||||
for i, down_block_type in enumerate(down_block_types):
|
||||
input_channel = output_channel
|
||||
output_channel = block_out_channels[i]
|
||||
is_final_block = i == len(block_out_channels) - 1
|
||||
|
||||
down_block = get_down_block(
|
||||
down_block_type,
|
||||
num_layers=layers_per_block,
|
||||
in_channels=input_channel,
|
||||
out_channels=output_channel,
|
||||
temb_channels=blocks_time_embed_dim,
|
||||
add_downsample=not is_final_block,
|
||||
resnet_eps=norm_eps,
|
||||
resnet_act_fn=act_fn,
|
||||
resnet_groups=norm_num_groups,
|
||||
cross_attention_dim=cross_attention_dim[i],
|
||||
attn_num_head_channels=attention_head_dim[i],
|
||||
downsample_padding=downsample_padding,
|
||||
dual_cross_attention=dual_cross_attention,
|
||||
use_linear_projection=use_linear_projection,
|
||||
only_cross_attention=only_cross_attention[i],
|
||||
upcast_attention=upcast_attention,
|
||||
resnet_time_scale_shift=resnet_time_scale_shift,
|
||||
)
|
||||
self.down_blocks.append(down_block)
|
||||
|
||||
# mid
|
||||
if mid_block_type == "UNetMidBlock2DCrossAttn":
|
||||
self.mid_block = UNetMidBlock2DCrossAttn(
|
||||
in_channels=block_out_channels[-1],
|
||||
temb_channels=blocks_time_embed_dim,
|
||||
resnet_eps=norm_eps,
|
||||
resnet_act_fn=act_fn,
|
||||
output_scale_factor=mid_block_scale_factor,
|
||||
resnet_time_scale_shift=resnet_time_scale_shift,
|
||||
cross_attention_dim=cross_attention_dim[-1],
|
||||
attn_num_head_channels=attention_head_dim[-1],
|
||||
resnet_groups=norm_num_groups,
|
||||
dual_cross_attention=dual_cross_attention,
|
||||
use_linear_projection=use_linear_projection,
|
||||
upcast_attention=upcast_attention,
|
||||
)
|
||||
elif mid_block_type == "UNetMidBlock2DSimpleCrossAttn":
|
||||
self.mid_block = UNetMidBlock2DSimpleCrossAttn(
|
||||
in_channels=block_out_channels[-1],
|
||||
temb_channels=blocks_time_embed_dim,
|
||||
resnet_eps=norm_eps,
|
||||
resnet_act_fn=act_fn,
|
||||
output_scale_factor=mid_block_scale_factor,
|
||||
cross_attention_dim=cross_attention_dim[-1],
|
||||
attn_num_head_channels=attention_head_dim[-1],
|
||||
resnet_groups=norm_num_groups,
|
||||
resnet_time_scale_shift=resnet_time_scale_shift,
|
||||
)
|
||||
elif mid_block_type is None:
|
||||
self.mid_block = None
|
||||
else:
|
||||
raise ValueError(f"unknown mid_block_type : {mid_block_type}")
|
||||
|
||||
# count how many layers upsample the images
|
||||
self.num_upsamplers = 0
|
||||
|
||||
# up
|
||||
reversed_block_out_channels = list(reversed(block_out_channels))
|
||||
reversed_attention_head_dim = list(reversed(attention_head_dim))
|
||||
reversed_cross_attention_dim = list(reversed(cross_attention_dim))
|
||||
only_cross_attention = list(reversed(only_cross_attention))
|
||||
|
||||
output_channel = reversed_block_out_channels[0]
|
||||
for i, up_block_type in enumerate(up_block_types):
|
||||
is_final_block = i == len(block_out_channels) - 1
|
||||
|
||||
prev_output_channel = output_channel
|
||||
output_channel = reversed_block_out_channels[i]
|
||||
input_channel = reversed_block_out_channels[min(i + 1, len(block_out_channels) - 1)]
|
||||
|
||||
# add upsample block for all BUT final layer
|
||||
if not is_final_block:
|
||||
add_upsample = True
|
||||
self.num_upsamplers += 1
|
||||
else:
|
||||
add_upsample = False
|
||||
|
||||
up_block = get_up_block(
|
||||
up_block_type,
|
||||
num_layers=layers_per_block + 1,
|
||||
in_channels=input_channel,
|
||||
out_channels=output_channel,
|
||||
prev_output_channel=prev_output_channel,
|
||||
temb_channels=blocks_time_embed_dim,
|
||||
add_upsample=add_upsample,
|
||||
resnet_eps=norm_eps,
|
||||
resnet_act_fn=act_fn,
|
||||
resnet_groups=norm_num_groups,
|
||||
cross_attention_dim=reversed_cross_attention_dim[i],
|
||||
attn_num_head_channels=reversed_attention_head_dim[i],
|
||||
dual_cross_attention=dual_cross_attention,
|
||||
use_linear_projection=use_linear_projection,
|
||||
only_cross_attention=only_cross_attention[i],
|
||||
upcast_attention=upcast_attention,
|
||||
resnet_time_scale_shift=resnet_time_scale_shift,
|
||||
)
|
||||
self.up_blocks.append(up_block)
|
||||
prev_output_channel = output_channel
|
||||
|
||||
# out
|
||||
if norm_num_groups is not None:
|
||||
self.conv_norm_out = nn.GroupNorm(
|
||||
num_channels=block_out_channels[0], num_groups=norm_num_groups, eps=norm_eps
|
||||
)
|
||||
self.conv_act = nn.SiLU()
|
||||
else:
|
||||
self.conv_norm_out = None
|
||||
self.conv_act = None
|
||||
|
||||
conv_out_padding = (conv_out_kernel - 1) // 2
|
||||
self.conv_out = nn.Conv2d(
|
||||
block_out_channels[0], out_channels, kernel_size=conv_out_kernel, padding=conv_out_padding
|
||||
)
|
||||
|
||||
@property
|
||||
def attn_processors(self) -> Dict[str, AttentionProcessor]:
|
||||
r"""
|
||||
Returns:
|
||||
`dict` of attention processors: A dictionary containing all attention processors used in the model with
|
||||
indexed by its weight name.
|
||||
"""
|
||||
# set recursively
|
||||
processors = {}
|
||||
|
||||
def fn_recursive_add_processors(name: str, module: torch.nn.Module, processors: Dict[str, AttentionProcessor]):
|
||||
if hasattr(module, "set_processor"):
|
||||
processors[f"{name}.processor"] = module.processor
|
||||
|
||||
for sub_name, child in module.named_children():
|
||||
fn_recursive_add_processors(f"{name}.{sub_name}", child, processors)
|
||||
|
||||
return processors
|
||||
|
||||
for name, module in self.named_children():
|
||||
fn_recursive_add_processors(name, module, processors)
|
||||
|
||||
return processors
|
||||
|
||||
def set_attn_processor(self, processor: Union[AttentionProcessor, Dict[str, AttentionProcessor]]):
|
||||
r"""
|
||||
Parameters:
|
||||
`processor (`dict` of `AttentionProcessor` or `AttentionProcessor`):
|
||||
The instantiated processor class or a dictionary of processor classes that will be set as the processor
|
||||
of **all** `Attention` layers.
|
||||
In case `processor` is a dict, the key needs to define the path to the corresponding cross attention processor. This is strongly recommended when setting trainable attention processors.:
|
||||
|
||||
"""
|
||||
count = len(self.attn_processors.keys())
|
||||
|
||||
if isinstance(processor, dict) and len(processor) != count:
|
||||
raise ValueError(
|
||||
f"A dict of processors was passed, but the number of processors {len(processor)} does not match the"
|
||||
f" number of attention layers: {count}. Please make sure to pass {count} processor classes."
|
||||
)
|
||||
|
||||
def fn_recursive_attn_processor(name: str, module: torch.nn.Module, processor):
|
||||
if hasattr(module, "set_processor"):
|
||||
if not isinstance(processor, dict):
|
||||
module.set_processor(processor)
|
||||
else:
|
||||
module.set_processor(processor.pop(f"{name}.processor"))
|
||||
|
||||
for sub_name, child in module.named_children():
|
||||
fn_recursive_attn_processor(f"{name}.{sub_name}", child, processor)
|
||||
|
||||
for name, module in self.named_children():
|
||||
fn_recursive_attn_processor(name, module, processor)
|
||||
|
||||
def set_default_attn_processor(self):
|
||||
"""
|
||||
Disables custom attention processors and sets the default attention implementation.
|
||||
"""
|
||||
self.set_attn_processor(AttnProcessor())
|
||||
|
||||
def set_attention_slice(self, slice_size):
|
||||
r"""
|
||||
Enable sliced attention computation.
|
||||
|
||||
When this option is enabled, the attention module will split the input tensor in slices, to compute attention
|
||||
in several steps. This is useful to save some memory in exchange for a small speed decrease.
|
||||
|
||||
Args:
|
||||
slice_size (`str` or `int` or `list(int)`, *optional*, defaults to `"auto"`):
|
||||
When `"auto"`, halves the input to the attention heads, so attention will be computed in two steps. If
|
||||
`"max"`, maximum amount of memory will be saved by running only one slice at a time. If a number is
|
||||
provided, uses as many slices as `attention_head_dim // slice_size`. In this case, `attention_head_dim`
|
||||
must be a multiple of `slice_size`.
|
||||
"""
|
||||
sliceable_head_dims = []
|
||||
|
||||
def fn_recursive_retrieve_sliceable_dims(module: torch.nn.Module):
|
||||
if hasattr(module, "set_attention_slice"):
|
||||
sliceable_head_dims.append(module.sliceable_head_dim)
|
||||
|
||||
for child in module.children():
|
||||
fn_recursive_retrieve_sliceable_dims(child)
|
||||
|
||||
# retrieve number of attention layers
|
||||
for module in self.children():
|
||||
fn_recursive_retrieve_sliceable_dims(module)
|
||||
|
||||
num_sliceable_layers = len(sliceable_head_dims)
|
||||
|
||||
if slice_size == "auto":
|
||||
# half the attention head size is usually a good trade-off between
|
||||
# speed and memory
|
||||
slice_size = [dim // 2 for dim in sliceable_head_dims]
|
||||
elif slice_size == "max":
|
||||
# make smallest slice possible
|
||||
slice_size = num_sliceable_layers * [1]
|
||||
|
||||
slice_size = num_sliceable_layers * [slice_size] if not isinstance(slice_size, list) else slice_size
|
||||
|
||||
if len(slice_size) != len(sliceable_head_dims):
|
||||
raise ValueError(
|
||||
f"You have provided {len(slice_size)}, but {self.config} has {len(sliceable_head_dims)} different"
|
||||
f" attention layers. Make sure to match `len(slice_size)` to be {len(sliceable_head_dims)}."
|
||||
)
|
||||
|
||||
for i in range(len(slice_size)):
|
||||
size = slice_size[i]
|
||||
dim = sliceable_head_dims[i]
|
||||
if size is not None and size > dim:
|
||||
raise ValueError(f"size {size} has to be smaller or equal to {dim}.")
|
||||
|
||||
# Recursively walk through all the children.
|
||||
# Any children which exposes the set_attention_slice method
|
||||
# gets the message
|
||||
def fn_recursive_set_attention_slice(module: torch.nn.Module, slice_size: List[int]):
|
||||
if hasattr(module, "set_attention_slice"):
|
||||
module.set_attention_slice(slice_size.pop())
|
||||
|
||||
for child in module.children():
|
||||
fn_recursive_set_attention_slice(child, slice_size)
|
||||
|
||||
reversed_slice_size = list(reversed(slice_size))
|
||||
for module in self.children():
|
||||
fn_recursive_set_attention_slice(module, reversed_slice_size)
|
||||
|
||||
def _set_gradient_checkpointing(self, module, value=False):
|
||||
if isinstance(module, (CrossAttnDownBlock2D, DownBlock2D, CrossAttnUpBlock2D, UpBlock2D)):
|
||||
module.gradient_checkpointing = value
|
||||
|
||||
def forward(
|
||||
self,
|
||||
sample: torch.FloatTensor,
|
||||
timestep: Union[torch.Tensor, float, int],
|
||||
encoder_hidden_states: torch.Tensor,
|
||||
class_labels: Optional[torch.Tensor] = None,
|
||||
timestep_cond: Optional[torch.Tensor] = None,
|
||||
attention_mask: Optional[torch.Tensor] = None,
|
||||
cross_attention_kwargs: Optional[Dict[str, Any]] = None,
|
||||
down_block_additional_residuals: Optional[Tuple[torch.Tensor]] = None,
|
||||
mid_block_additional_residual: Optional[torch.Tensor] = None,
|
||||
return_dict: bool = True,
|
||||
segmentation_mask: torch.Tensor=None, # added
|
||||
masked_feature: torch.Tensor=None, # added
|
||||
feature_mask: torch.Tensor=None, # added
|
||||
) -> Union[UNet2DConditionOutput, Tuple]:
|
||||
r"""
|
||||
Args:
|
||||
sample (`torch.FloatTensor`): (batch, channel, height, width) noisy inputs tensor
|
||||
timestep (`torch.FloatTensor` or `float` or `int`): (batch) timesteps
|
||||
encoder_hidden_states (`torch.FloatTensor`): (batch, sequence_length, feature_dim) encoder hidden states
|
||||
return_dict (`bool`, *optional*, defaults to `True`):
|
||||
Whether or not to return a [`models.unet_2d_condition.UNet2DConditionOutput`] instead of a plain tuple.
|
||||
cross_attention_kwargs (`dict`, *optional*):
|
||||
A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under
|
||||
`self.processor` in
|
||||
[diffusers.cross_attention](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/cross_attention.py).
|
||||
|
||||
Returns:
|
||||
[`~models.unet_2d_condition.UNet2DConditionOutput`] or `tuple`:
|
||||
[`~models.unet_2d_condition.UNet2DConditionOutput`] if `return_dict` is True, otherwise a `tuple`. When
|
||||
returning a tuple, the first element is the sample tensor.
|
||||
"""
|
||||
# By default samples have to be AT least a multiple of the overall upsampling factor.
|
||||
# The overall upsampling factor is equal to 2 ** (# num of upsampling layers).
|
||||
# However, the upsampling interpolation output size can be forced to fit any upsampling size
|
||||
# on the fly if necessary.
|
||||
default_overall_up_factor = 2**self.num_upsamplers
|
||||
|
||||
# upsample size should be forwarded when sample is not a multiple of `default_overall_up_factor`
|
||||
forward_upsample_size = False
|
||||
upsample_size = None
|
||||
|
||||
if any(s % default_overall_up_factor != 0 for s in sample.shape[-2:]):
|
||||
logger.info("Forward upsample size to force interpolation output size.")
|
||||
forward_upsample_size = True
|
||||
|
||||
# prepare attention_mask
|
||||
if attention_mask is not None:
|
||||
attention_mask = (1 - attention_mask.to(sample.dtype)) * -10000.0
|
||||
attention_mask = attention_mask.unsqueeze(1)
|
||||
|
||||
# 0. concat all the feature together
|
||||
sample = torch.cat([sample, feature_mask, masked_feature], dim=1)
|
||||
|
||||
# 1. time
|
||||
timesteps = timestep
|
||||
if not torch.is_tensor(timesteps):
|
||||
# TODO: this requires sync between CPU and GPU. So try to pass timesteps as tensors if you can
|
||||
# This would be a good case for the `match` statement (Python 3.10+)
|
||||
is_mps = sample.device.type == "mps"
|
||||
if isinstance(timestep, float):
|
||||
dtype = torch.float32 if is_mps else torch.float64
|
||||
else:
|
||||
dtype = torch.int32 if is_mps else torch.int64
|
||||
timesteps = torch.tensor([timesteps], dtype=dtype, device=sample.device)
|
||||
elif len(timesteps.shape) == 0:
|
||||
timesteps = timesteps[None].to(sample.device)
|
||||
|
||||
# broadcast to batch dimension in a way that's compatible with ONNX/Core ML
|
||||
timesteps = timesteps.expand(sample.shape[0])
|
||||
|
||||
t_emb = self.time_proj(timesteps)
|
||||
|
||||
# timesteps does not contain any weights and will always return f32 tensors
|
||||
# but time_embedding might actually be running in fp16. so we need to cast here.
|
||||
# there might be better ways to encapsulate this.
|
||||
t_emb = t_emb.to(dtype=self.dtype)
|
||||
|
||||
emb = self.time_embedding(t_emb, timestep_cond)
|
||||
|
||||
if self.class_embedding is not None:
|
||||
if class_labels is None:
|
||||
raise ValueError("class_labels should be provided when num_class_embeds > 0")
|
||||
|
||||
if self.config.class_embed_type == "timestep":
|
||||
class_labels = self.time_proj(class_labels)
|
||||
|
||||
class_emb = self.class_embedding(class_labels).to(dtype=self.dtype)
|
||||
|
||||
if self.config.class_embeddings_concat:
|
||||
emb = torch.cat([emb, class_emb], dim=-1)
|
||||
else:
|
||||
emb = emb + class_emb
|
||||
|
||||
# 2. pre-process
|
||||
segmentation_mask_embedding = self.word_embedding(segmentation_mask.squeeze(1).long()).permute(0,3,1,2) # get 8-d embedding from character-level segmentation mask
|
||||
segmentation_mask_embedding = self.segmap_conv(segmentation_mask_embedding) # resize the mask using cnn
|
||||
sample = torch.cat([sample, segmentation_mask_embedding], 1)
|
||||
|
||||
sample = self.conv_in(sample)
|
||||
|
||||
# 3. down
|
||||
down_block_res_samples = (sample,)
|
||||
for downsample_block in self.down_blocks:
|
||||
if hasattr(downsample_block, "has_cross_attention") and downsample_block.has_cross_attention:
|
||||
sample, res_samples = downsample_block(
|
||||
hidden_states=sample,
|
||||
temb=emb,
|
||||
encoder_hidden_states=encoder_hidden_states,
|
||||
attention_mask=attention_mask,
|
||||
cross_attention_kwargs=cross_attention_kwargs,
|
||||
)
|
||||
else:
|
||||
sample, res_samples = downsample_block(hidden_states=sample, temb=emb)
|
||||
|
||||
down_block_res_samples += res_samples
|
||||
|
||||
if down_block_additional_residuals is not None:
|
||||
new_down_block_res_samples = ()
|
||||
|
||||
for down_block_res_sample, down_block_additional_residual in zip(
|
||||
down_block_res_samples, down_block_additional_residuals
|
||||
):
|
||||
down_block_res_sample = down_block_res_sample + down_block_additional_residual
|
||||
new_down_block_res_samples += (down_block_res_sample,)
|
||||
|
||||
down_block_res_samples = new_down_block_res_samples
|
||||
|
||||
# 4. mid
|
||||
if self.mid_block is not None:
|
||||
sample = self.mid_block(
|
||||
sample,
|
||||
emb,
|
||||
encoder_hidden_states=encoder_hidden_states,
|
||||
attention_mask=attention_mask,
|
||||
cross_attention_kwargs=cross_attention_kwargs,
|
||||
)
|
||||
|
||||
if mid_block_additional_residual is not None:
|
||||
sample = sample + mid_block_additional_residual
|
||||
|
||||
# 5. up
|
||||
for i, upsample_block in enumerate(self.up_blocks):
|
||||
is_final_block = i == len(self.up_blocks) - 1
|
||||
|
||||
res_samples = down_block_res_samples[-len(upsample_block.resnets) :]
|
||||
down_block_res_samples = down_block_res_samples[: -len(upsample_block.resnets)]
|
||||
|
||||
# if we have not reached the final block and need to forward the
|
||||
# upsample size, we do it here
|
||||
if not is_final_block and forward_upsample_size:
|
||||
upsample_size = down_block_res_samples[-1].shape[2:]
|
||||
|
||||
if hasattr(upsample_block, "has_cross_attention") and upsample_block.has_cross_attention:
|
||||
sample = upsample_block(
|
||||
hidden_states=sample,
|
||||
temb=emb,
|
||||
res_hidden_states_tuple=res_samples,
|
||||
encoder_hidden_states=encoder_hidden_states,
|
||||
cross_attention_kwargs=cross_attention_kwargs,
|
||||
upsample_size=upsample_size,
|
||||
attention_mask=attention_mask,
|
||||
)
|
||||
else:
|
||||
sample = upsample_block(
|
||||
hidden_states=sample, temb=emb, res_hidden_states_tuple=res_samples, upsample_size=upsample_size
|
||||
)
|
||||
|
||||
# 6. post-process
|
||||
if self.conv_norm_out:
|
||||
sample = self.conv_norm_out(sample)
|
||||
sample = self.conv_act(sample)
|
||||
sample = self.conv_out(sample)
|
||||
|
||||
if not return_dict:
|
||||
return (sample,)
|
||||
|
||||
return UNet2DConditionOutput(sample=sample)
|
||||
Binary file not shown.
|
After Width: | Height: | Size: 125 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 714 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 1.3 MiB |
Binary file not shown.
|
After Width: | Height: | Size: 1.4 MiB |
Binary file not shown.
|
After Width: | Height: | Size: 2.4 MiB |
Binary file not shown.
|
After Width: | Height: | Size: 400 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 454 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 117 KiB |
@@ -0,0 +1 @@
|
||||
Blechschilder MOORE - CARWASH - Best Hand Job In Town
|
||||
Binary file not shown.
Binary file not shown.
|
After Width: | Height: | Size: 53 KiB |
@@ -0,0 +1 @@
|
||||
{"NSFW": "UNLIKELY", "similarity": 0.33181482553482056, "LICENSE": "?", "caption": "Blechschilder MOORE - CARWASH - Best Hand Job In Town", "url": "https://static.posters.cz/image/350/metallschilder/moore-carwash-best-hand-job-in-town-i19436.jpg", "key": "062690093", "status": "success", "error_message": null, "width": 350, "height": 275, "exif": "{}", "original_width": 350, "original_height": 275}
|
||||
@@ -0,0 +1,5 @@
|
||||
in 343,454,384,439,394,489,354,502 0.9964783191680908
|
||||
town! 387,431,492,396,504,450,399,485 0.9821552038192749
|
||||
Best 2,206,96,173,112,243,19,277 0.995647132396698
|
||||
hand 102,173,193,139,207,201,115,234 0.9959321618080139
|
||||
job 194,143,254,117,272,184,212,210 0.9975775480270386
|
||||
@@ -0,0 +1 @@
|
||||
Weekender: Win 40K Shadowspear & GW TOP TRUMP?
|
||||
Binary file not shown.
Binary file not shown.
|
After Width: | Height: | Size: 60 KiB |
@@ -0,0 +1 @@
|
||||
{"NSFW": "UNLIKELY", "similarity": 0.30257895588874817, "LICENSE": "?", "caption": "Weekender: Win 40K Shadowspear & GW TOP TRUMP?", "url": "https://images.beastsofwar.com/2019/02/Weekender-260-Front-v2-Cover-Image-1024x576.jpg", "key": "062692210", "status": "success", "error_message": null, "width": 1024, "height": 576, "exif": "{}", "original_width": 1024, "original_height": 576}
|
||||
@@ -0,0 +1,7 @@
|
||||
an 439,418,466,379,497,448,470,487 0.9474218487739563
|
||||
DOWNSPEAR 25,440,185,382,190,421,30,480 0.5379968881607056
|
||||
WIN 15,354,91,327,101,414,25,441 0.9966829419136047
|
||||
Love 185,336,192,323,211,352,204,367 0.414926141500473
|
||||
40K 99,312,183,294,189,389,105,406 0.971768319606781
|
||||
T 184,264,218,214,295,379,262,429 0.11518792808055878
|
||||
CHAT 14,20,120,20,120,102,14,102 0.9378468990325928
|
||||
@@ -0,0 +1 @@
|
||||
"Little Joe - ""A La Guerra Me Llevan"""
|
||||
Binary file not shown.
Binary file not shown.
|
After Width: | Height: | Size: 30 KiB |
@@ -0,0 +1 @@
|
||||
{"NSFW": "UNLIKELY", "similarity": 0.3536737263202667, "LICENSE": "?", "caption": "\"Little Joe - \"\"A La Guerra Me Llevan\"\"\"", "url": "https://i.ytimg.com/vi/8bx7Wx59Oxc/hqdefault.jpg", "key": "062692530", "status": "success", "error_message": null, "width": 480, "height": 360, "exif": "{}", "original_width": 480, "original_height": 360}
|
||||
@@ -0,0 +1,3 @@
|
||||
FAMILIA 185,376,186,277,437,283,436,382 0.9923758506774902
|
||||
LITTLE 37,231,40,119,243,129,241,240 0.99787437915802
|
||||
JOE 257,236,261,105,393,113,389,244 0.996442437171936
|
||||
@@ -0,0 +1 @@
|
||||
VH-EZT at Busselton Airport Tail shot
|
||||
Binary file not shown.
Binary file not shown.
|
After Width: | Height: | Size: 20 KiB |
@@ -0,0 +1 @@
|
||||
{"NSFW": "UNLIKELY", "similarity": 0.3146217465400696, "LICENSE": "?", "caption": "VH-EZT at Busselton Airport Tail shot", "url": "http://uniflying.org.au/wp-content/uploads/EZT-STYLE-2-150x150.jpg", "key": "271975131", "status": "success", "error_message": null, "width": 512, "height": 512, "original_width": 150, "original_height": 150, "exif": "{}", "sha256": "b32fe5c60a42f18fd496a461c948683cd1f52e89033a61660795a6d2cd3b93de"}
|
||||
@@ -0,0 +1 @@
|
||||
VH-EZT 143,133,501,82,532,301,175,352 0.21680155396461487
|
||||
@@ -0,0 +1 @@
|
||||
Strangers to These Shores : Race and Ethnic Relations in the United States, Parrillo, Vincent N., 0205585574
|
||||
Binary file not shown.
Binary file not shown.
|
After Width: | Height: | Size: 32 KiB |
@@ -0,0 +1 @@
|
||||
{"NSFW": "UNLIKELY", "similarity": 0.30967631936073303, "LICENSE": "?", "caption": "Strangers to These Shores : Race and Ethnic Relations in the United States, Parrillo, Vincent N., 0205585574", "url": "https://dcqohhuh7tdzn.cloudfront.net/5571/9780205585571_medium.jpg", "key": "271975251", "status": "success", "error_message": null, "width": 512, "height": 512, "original_width": 127, "original_height": 160, "exif": "{}", "sha256": "86d27d0284bfd0f6d01ce332bd157e01d43265d48e9a20646321c6ad7c070b35"}
|
||||
@@ -0,0 +1 @@
|
||||
SHORES 148,89,361,83,362,140,150,146 0.9963237643241882
|
||||
@@ -0,0 +1 @@
|
||||
Marvel MV15598 Manta Polar, Multicolor, 150x100cm
|
||||
Binary file not shown.
Binary file not shown.
|
After Width: | Height: | Size: 36 KiB |
@@ -0,0 +1 @@
|
||||
{"NSFW": "UNLIKELY", "similarity": 0.31281691789627075, "LICENSE": "?", "caption": "Marvel MV15598 Manta Polar, Multicolor, 150x100cm", "url": "https://m.media-amazon.com/images/I/51vgtDUtImL._AC_US200_.jpg", "key": "271978467", "status": "success", "error_message": null, "width": 512, "height": 512, "original_width": 200, "original_height": 200, "exif": "{}", "sha256": "039cf017effd4aeac9b2b8511f4a62ef3afc6b7af0bbee42eb0b5fb03b3896cc"}
|
||||
@@ -0,0 +1,5 @@
|
||||
unr 357,447,362,421,414,433,408,459 0.28840482234954834
|
||||
PIDERMA 71,403,409,320,432,415,94,498 0.6835125088691711
|
||||
MARVEL 211,331,301,320,307,363,217,374 0.9939236640930176
|
||||
GREAT 361,74,417,74,417,99,361,99 0.9460398554801941
|
||||
COMES 312,54,324,39,373,78,361,93 0.9603791236877441
|
||||
@@ -0,0 +1,32 @@
|
||||
import json
|
||||
import os
|
||||
import traceback
|
||||
from tqdm import tqdm
|
||||
from multiprocessing import Pool
|
||||
|
||||
ROOT_FROM = 'XXX' # the path of laion-ocr-zip
|
||||
ROOT_TO = 'XXX' # the path for saving dataset
|
||||
MULTIPROCESSING_NUM = 64
|
||||
DOWNLOAD_IMAGES = False # whether to download images from urls
|
||||
|
||||
def unzip_file(idx):
|
||||
if not os.path.exists(f'{ROOT_FROM}/{idx}.zip') or os.path.exists(f'{ROOT_TO}/{idx}'):
|
||||
return
|
||||
cmd = f'unzip -q {ROOT_FROM}/{idx}.zip -d {ROOT_TO}'
|
||||
os.system(cmd)
|
||||
|
||||
|
||||
def multiprocess_unzip_file(idxs):
|
||||
os.makedirs(ROOT_TO, exist_ok=True)
|
||||
|
||||
with Pool(processes=MULTIPROCESSING_NUM) as p:
|
||||
with tqdm(total=len(idxs), desc='total') as pbar:
|
||||
for i, _ in enumerate(p.imap_unordered(unzip_file, idxs)):
|
||||
pbar.update()
|
||||
print("multiprocess_unzip_file done!")
|
||||
|
||||
if __name__ == '__main__':
|
||||
files = os.listdir(ROOT_FROM)
|
||||
idxs = [str(idx[:-4]).zfill(5) for idx in files]
|
||||
multiprocess_unzip_file(idxs)
|
||||
print("Finished!")
|
||||
File diff suppressed because one or more lines are too long
@@ -0,0 +1,167 @@
|
||||
# ------------------------------------------
|
||||
# TextDiffuser: Diffusion Models as Text Painters
|
||||
# Paper Link: https://arxiv.org/abs/2305.10855
|
||||
# Code Link: https://github.com/microsoft/unilm/tree/master/textdiffuser
|
||||
# Copyright (c) Microsoft Corporation.
|
||||
# This file provides the inference script.
|
||||
# ------------------------------------------
|
||||
|
||||
import json
|
||||
import os
|
||||
import numpy as np
|
||||
import argparse
|
||||
from clipscore import cal_clipscore
|
||||
from fid_score import calculate_fid_given_paths
|
||||
|
||||
|
||||
def eval_clipscore(root_eval, root_res, dataset, device="cuda:0", num_images_per_prompt=4):
|
||||
with open(os.path.join(root_eval, dataset, dataset + '.txt'), 'r') as fr:
|
||||
text_list = fr.readlines()
|
||||
text_list = [_.strip() for _ in text_list]
|
||||
clip_scores = []
|
||||
scores = []
|
||||
for seed in range(num_images_per_prompt):
|
||||
if 'stablediffusion' in root_res:
|
||||
format = '.png'
|
||||
else:
|
||||
format = '.jpg'
|
||||
image_list = [os.path.join(root_res, dataset, 'images_' + str(seed),
|
||||
str(idx) + '_' + str(seed) + format) for idx in range(len(text_list))]
|
||||
image_ids = [str(idx) + '_' + str(seed) + format for idx in range(len(text_list))]
|
||||
score = cal_clipscore(image_ids=image_ids, image_paths=image_list, text_list=text_list, device=device)
|
||||
clip_score = np.mean([s['CLIPScore'] for s in score.values()])
|
||||
clip_scores.append(clip_score)
|
||||
scores.append(score)
|
||||
print("clip_score:", np.mean(clip_scores), clip_scores)
|
||||
return np.mean(clip_scores), scores
|
||||
|
||||
|
||||
def MARIOEval_evaluate_results(root, datasets_with_images, datasets, methods, gpu,
|
||||
eval_clipscore_flag=True, eval_fid_flag=True, num_images_per_prompt=4):
|
||||
root_eval = os.path.join(root, "MARIOEval")
|
||||
method_res = {}
|
||||
device = "cuda:" + str(gpu)
|
||||
for method_idx, method in enumerate(methods):
|
||||
if method_idx != gpu: # running in different gpus simultaneously to save time
|
||||
continue
|
||||
print("\nmethod:", method)
|
||||
dataset_res = {}
|
||||
root_res = os.path.join(root, 'generation', method)
|
||||
for dataset in datasets:
|
||||
print("dataset:", dataset)
|
||||
dataset_res[dataset] = {}
|
||||
if eval_clipscore_flag:
|
||||
dataset_res[dataset]['clipscore'], dataset_res[dataset]['scores'] =\
|
||||
eval_clipscore(root_eval, root_res, dataset, device, num_images_per_prompt)
|
||||
if eval_fid_flag and dataset in datasets_with_images:
|
||||
gt_path = os.path.join(root_eval, dataset, 'images')
|
||||
fids = []
|
||||
for idx in range(num_images_per_prompt):
|
||||
gen_path = os.path.join(root_res, dataset, 'images_' + str(idx))
|
||||
fids.append(calculate_fid_given_paths(paths=[gt_path, gen_path]))
|
||||
print("fid:", np.mean(fids), fids)
|
||||
dataset_res[dataset]['fid'] = np.mean(fids)
|
||||
|
||||
if eval_clipscore_flag:
|
||||
method_clipscores = []
|
||||
for seed in range(num_images_per_prompt):
|
||||
clipscore_list = []
|
||||
for dataset in dataset_res.keys():
|
||||
clipscore_list += [_['CLIPScore'] for _ in dataset_res[dataset]['scores'][seed].values()]
|
||||
method_clipscores.append(np.mean(clipscore_list))
|
||||
method_clipscore = np.mean(method_clipscores)
|
||||
dataset_res['clipscore'] = method_clipscore
|
||||
if eval_fid_flag:
|
||||
method_fids = []
|
||||
for idx in range(num_images_per_prompt):
|
||||
gt_paths = []
|
||||
gen_paths = []
|
||||
for dataset in dataset_res.keys():
|
||||
if dataset in datasets_with_images:
|
||||
gt_paths.append(os.path.join(root_eval, dataset, 'images'))
|
||||
gen_paths.append(os.path.join(root_res, dataset, 'images_' + str(idx)))
|
||||
if len(gt_paths):
|
||||
method_fids.append(calculate_fid_given_paths(paths=[gt_paths, gen_paths]))
|
||||
print("fid:", np.mean(method_fids), method_fids)
|
||||
method_fid = np.mean(method_fids)
|
||||
dataset_res['fid'] = method_fid
|
||||
|
||||
method_res[method] = dataset_res
|
||||
with open(os.path.join(root_res, 'eval.json'), 'w') as fw:
|
||||
json.dump(dataset_res, fw)
|
||||
|
||||
print(method_res)
|
||||
with open(os.path.join(root, 'generation', 'eval.json'), 'w') as fw:
|
||||
json.dump(method_res, fw)
|
||||
|
||||
|
||||
def merge_eval_results(root, methods):
|
||||
method_res = {}
|
||||
for method_idx, method in enumerate(methods):
|
||||
root_res = os.path.join(root, 'generation', method)
|
||||
with open(os.path.join(root_res, 'eval.json'), 'r') as fr:
|
||||
dataset_res = json.load(fr)
|
||||
for k, v in dataset_res.items():
|
||||
if type(v) is dict:
|
||||
del v['scores'] # too long
|
||||
method_res[method] = dataset_res
|
||||
|
||||
with open(os.path.join(root, 'generation', 'eval.json'), 'w') as fw:
|
||||
json.dump(method_res, fw)
|
||||
|
||||
|
||||
def parse_args():
|
||||
parser = argparse.ArgumentParser(description="Simple example of a training script.")
|
||||
parser.add_argument(
|
||||
"--dataset",
|
||||
type=str,
|
||||
default='TMDBEval500',
|
||||
required=False,
|
||||
choices=['TMDBEval500', 'OpenLibraryEval500', 'LAIONEval4000',
|
||||
'ChineseDrawText', 'DrawBenchText', 'DrawTextCreative']
|
||||
)
|
||||
parser.add_argument(
|
||||
"--root",
|
||||
type=str,
|
||||
default="/path/to/data/TextDiffuser/evaluation/",
|
||||
required=True,
|
||||
)
|
||||
parser.add_argument(
|
||||
"--method",
|
||||
type=str,
|
||||
default='controlnet',
|
||||
required=False,
|
||||
choices=['controlnet', 'deepfloyd', 'stablediffusion', 'textdiffuser']
|
||||
)
|
||||
parser.add_argument(
|
||||
"--gpu",
|
||||
type=int,
|
||||
default=0,
|
||||
required=False,
|
||||
)
|
||||
parser.add_argument(
|
||||
"--split",
|
||||
type=int,
|
||||
default=0,
|
||||
required=False,
|
||||
)
|
||||
parser.add_argument(
|
||||
"--total_split",
|
||||
type=int,
|
||||
default=1,
|
||||
required=False,
|
||||
)
|
||||
args = parser.parse_args()
|
||||
return args
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
args = parse_args()
|
||||
os.environ['CUDA_VISIBLE_DEVICES'] = str(args.gpu)
|
||||
datasets_with_images = ['TMDBEval500', 'OpenLibraryEval500', 'LAIONEval4000']
|
||||
datasets = datasets_with_images + ['ChineseDrawText', 'DrawBenchText', 'DrawTextCreative']
|
||||
methods = ['textdiffuser', 'controlnet', 'deepfloyd', 'stablediffusion']
|
||||
|
||||
MARIOEval_evaluate_results(args.root, datasets_with_images, datasets, methods, args.gpu,
|
||||
eval_clipscore_flag=True, eval_fid_flag=True, num_images_per_prompt=4)
|
||||
merge_eval_results(args.root, methods)
|
||||
@@ -0,0 +1,192 @@
|
||||
# ------------------------------------------
|
||||
# TextDiffuser: Diffusion Models as Text Painters
|
||||
# Paper Link: https://arxiv.org/abs/2305.10855
|
||||
# Code Link: https://github.com/microsoft/unilm/tree/master/textdiffuser
|
||||
# Copyright (c) Microsoft Corporation.
|
||||
# This file provides the inference script.
|
||||
# ------------------------------------------
|
||||
|
||||
import os
|
||||
from PIL import Image
|
||||
import numpy as np
|
||||
import torch
|
||||
from tqdm import tqdm
|
||||
import argparse
|
||||
import cv2
|
||||
import torchvision.transforms as transforms
|
||||
|
||||
to_pil_image = transforms.ToPILImage()
|
||||
|
||||
def load_stablediffusion():
|
||||
from diffusers import StableDiffusionPipeline
|
||||
pipe = StableDiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5", torch_dtype=torch.float16)
|
||||
pipe.enable_xformers_memory_efficient_attention()
|
||||
pipe.enable_model_cpu_offload()
|
||||
return pipe
|
||||
|
||||
def test_stablediffusion(prompt, save_path, num_images_per_prompt=4,
|
||||
pipe=None, generator=None):
|
||||
images = pipe(prompt, num_inference_steps=50, generator=generator, num_images_per_prompt=num_images_per_prompt).images
|
||||
for idx, image in enumerate(images):
|
||||
image.save(save_path.replace('.jpg', '_' + str(idx) + '.jpg').replace('/images/', '/images_'+ str(idx) +'/'))
|
||||
|
||||
def load_deepfloyd_if():
|
||||
from diffusers import DiffusionPipeline
|
||||
stage_1 = DiffusionPipeline.from_pretrained("DeepFloyd/IF-I-XL-v1.0", variant="fp16", torch_dtype=torch.float16)
|
||||
# stage_1.enable_xformers_memory_efficient_attention() # remove line if torch.__version__ >= 2.0.0
|
||||
stage_1.enable_model_cpu_offload()
|
||||
stage_2 = DiffusionPipeline.from_pretrained("DeepFloyd/IF-II-L-v1.0", text_encoder=None, variant="fp16",
|
||||
torch_dtype=torch.float16)
|
||||
# stage_2.enable_xformers_memory_efficient_attention() # remove line if torch.__version__ >= 2.0.0
|
||||
stage_2.enable_model_cpu_offload()
|
||||
safety_modules = {"feature_extractor": stage_1.feature_extractor, "safety_checker": stage_1.safety_checker,
|
||||
"watermarker": stage_1.watermarker}
|
||||
stage_3 = DiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-x4-upscaler", **safety_modules,
|
||||
torch_dtype=torch.float16)
|
||||
# stage_3.enable_xformers_memory_efficient_attention() # remove line if torch.__version__ >= 2.0.0
|
||||
stage_3.enable_model_cpu_offload()
|
||||
return stage_1, stage_2, stage_3
|
||||
|
||||
|
||||
def test_deepfloyd_if(stage_1, stage_2, stage_3, prompt, save_path, num_images_per_prompt=4, generator=None):
|
||||
idx = num_images_per_prompt - 1 # if the last image of a case exists, then return
|
||||
new_save_path = save_path.replace('.jpg', '_' + str(idx) + '.jpg').replace('/images/', '/images_' + str(idx) + '/')
|
||||
if os.path.exists(new_save_path):
|
||||
return
|
||||
if not stage_1 or not stage_2 or not stage_3:
|
||||
stage_1, stage_2, stage_3 = load_deepfloyd_if()
|
||||
if generator is None:
|
||||
generator = torch.manual_seed(0)
|
||||
prompt_embeds, negative_embeds = stage_1.encode_prompt(prompt)
|
||||
stage_1.set_progress_bar_config(disable=True)
|
||||
stage_2.set_progress_bar_config(disable=True)
|
||||
stage_3.set_progress_bar_config(disable=True)
|
||||
images = stage_1(prompt_embeds=prompt_embeds, negative_prompt_embeds=negative_embeds, generator=generator,
|
||||
output_type="pt", num_images_per_prompt=num_images_per_prompt).images
|
||||
for idx, image in enumerate(images):
|
||||
image = stage_2(image=image.unsqueeze(0), prompt_embeds=prompt_embeds, negative_prompt_embeds=negative_embeds,
|
||||
generator=generator, output_type="pt").images
|
||||
image = stage_3(prompt=prompt, image=image, generator=generator, noise_level=100).images
|
||||
# image = to_pil_image(image[0].cpu())
|
||||
new_save_path = save_path.replace('.jpg', '_' + str(idx) + '.jpg').replace('/images/', '/images_'+ str(idx) +'/')
|
||||
image[0].save(new_save_path)
|
||||
|
||||
|
||||
def load_controlnet_cannyedge():
|
||||
from diffusers import StableDiffusionControlNetPipeline, ControlNetModel, UniPCMultistepScheduler
|
||||
|
||||
controlnet = ControlNetModel.from_pretrained("lllyasviel/sd-controlnet-canny", torch_dtype=torch.float16)
|
||||
pipe = StableDiffusionControlNetPipeline.from_pretrained("runwayml/stable-diffusion-v1-5", controlnet=controlnet,
|
||||
safety_checker=None, torch_dtype=torch.float16)
|
||||
pipe.set_progress_bar_config(disable=True)
|
||||
pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config)
|
||||
pipe.enable_xformers_memory_efficient_attention()
|
||||
pipe.enable_model_cpu_offload()
|
||||
return pipe
|
||||
|
||||
|
||||
def test_controlnet_cannyedge(prompt, save_path, canny_path, num_images_per_prompt=4,
|
||||
pipe=None, generator=None, low_threshold=100, high_threshold=200):
|
||||
'''ref: https://github.com/huggingface/diffusers/blob/131312caba0af97da98fc498dfdca335c9692f8c/docs/source/en/api/pipelines/stable_diffusion/controlnet.mdx'''
|
||||
from diffusers.utils import load_image
|
||||
if pipe is None:
|
||||
pipe = load_controlnet_cannyedge()
|
||||
|
||||
if os.path.exists(canny_path):
|
||||
canny_path = Image.open(canny_path)
|
||||
image = load_image(canny_path)
|
||||
image = np.array(image)
|
||||
image = cv2.Canny(image, low_threshold, high_threshold)
|
||||
image = image[:, :, None]
|
||||
image = np.concatenate([image, image, image], axis=2)
|
||||
image = Image.fromarray(image)
|
||||
images = pipe(prompt, image, num_inference_steps=20, generator=generator, num_images_per_prompt=num_images_per_prompt).images
|
||||
for idx, image in enumerate(images):
|
||||
image.save(save_path.replace('.jpg', '_' + str(idx) + '.jpg').replace('/images/', '/images_'+ str(idx) +'/'))
|
||||
|
||||
|
||||
def MARIOEval_generate_results(root, dataset, method='controlnet', num_images_per_prompt=4, split=0, total_split=1):
|
||||
root_eval = os.path.join(root, "MARIOEval")
|
||||
render_path = os.path.join(root_eval, dataset, 'render')
|
||||
root_res = os.path.join(root, "generation", method)
|
||||
for idx in range(num_images_per_prompt):
|
||||
os.makedirs(os.path.join(root_res, dataset, 'images_' + str(idx)), exist_ok=True)
|
||||
generator = torch.Generator(device="cuda").manual_seed(0)
|
||||
if method == 'controlnet':
|
||||
pipe = load_controlnet_cannyedge()
|
||||
elif method == 'stablediffusion':
|
||||
pipe = load_stablediffusion()
|
||||
elif method == 'deepfloyd':
|
||||
stage_1, stage_2, stage_3 = load_deepfloyd_if()
|
||||
|
||||
with open(os.path.join(root_eval, dataset, dataset + '.txt'), 'r') as fr:
|
||||
prompts = fr.readlines()
|
||||
prompts = [_.strip() for _ in prompts]
|
||||
for idx, prompt in tqdm(enumerate(prompts)):
|
||||
if idx < split * len(prompts) / total_split or idx > (split + 1) * len(prompts) / total_split:
|
||||
continue
|
||||
if method == 'controlnet':
|
||||
test_controlnet_cannyedge(prompt=prompt, num_images_per_prompt=num_images_per_prompt,
|
||||
save_path=os.path.join(root_res, dataset, 'images', str(idx) + '.jpg'),
|
||||
canny_path=os.path.join(render_path, str(idx) + '.png'),
|
||||
pipe=pipe, generator=generator)
|
||||
elif method == 'stablediffusion':
|
||||
test_stablediffusion(prompt=prompt, num_images_per_prompt=num_images_per_prompt,
|
||||
save_path=os.path.join(root_res, dataset, 'images', str(idx) + '.jpg'),
|
||||
pipe=pipe, generator=generator)
|
||||
elif method == 'deepfloyd':
|
||||
test_deepfloyd_if(stage_1, stage_2, stage_3, num_images_per_prompt=num_images_per_prompt,
|
||||
save_path=os.path.join(root_res, dataset, 'images', str(idx) + '.jpg'),
|
||||
prompt=prompt, generator=generator)
|
||||
|
||||
def parse_args():
|
||||
parser = argparse.ArgumentParser(description="Simple example of a training script.")
|
||||
parser.add_argument(
|
||||
"--dataset",
|
||||
type=str,
|
||||
default='TMDBEval500',
|
||||
required=False,
|
||||
choices=['TMDBEval500', 'OpenLibraryEval500', 'LAIONEval4000',
|
||||
'ChineseDrawText', 'DrawBenchText', 'DrawTextCreative']
|
||||
)
|
||||
parser.add_argument(
|
||||
"--root",
|
||||
type=str,
|
||||
default="/path/to/eval",
|
||||
required=True,
|
||||
)
|
||||
parser.add_argument(
|
||||
"--method",
|
||||
type=str,
|
||||
default='controlnet',
|
||||
required=False,
|
||||
choices=['controlnet', 'deepfloyd', 'stablediffusion', 'textdiffuser']
|
||||
)
|
||||
parser.add_argument(
|
||||
"--gpu",
|
||||
type=int,
|
||||
default=0,
|
||||
required=False,
|
||||
)
|
||||
parser.add_argument(
|
||||
"--split",
|
||||
type=int,
|
||||
default=0,
|
||||
required=False,
|
||||
)
|
||||
parser.add_argument(
|
||||
"--total_split",
|
||||
type=int,
|
||||
default=1,
|
||||
required=False,
|
||||
)
|
||||
args = parser.parse_args()
|
||||
return args
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
args = parse_args()
|
||||
os.environ['CUDA_VISIBLE_DEVICES'] = str(args.gpu)
|
||||
|
||||
MARIOEval_generate_results(root=args.root, dataset=args.dataset, method=args.method,
|
||||
split=args.split, total_split=args.total_split)
|
||||
@@ -0,0 +1,10 @@
|
||||
# Evaluation
|
||||
|
||||
We provide the code for sampling from Stable Diffusion, ControlNet, DeepFloyd at ```MARIOEval_generate.py```. Since these methods rely on *diffusers* of the original version, it is recommended to create a **NEW** environment and install packages with command ```pip install requirements.txt```. It is recommended to install pytorch with version >= 2.0 to avoid the OOM error.
|
||||
|
||||
Once the generation is complete, evaluation of FID and CLIPScore can be performed using the ```MARIOEval_evaluate.py``` file. For OCR metrics, please install [MaskTextSpotterV3](https://github.com/MhLiao/MaskTextSpotterV3) to obtain the OCR result of each image and refer to ```ocr_eval.py``` for evaluation. It should be noted that the output image of DeepFloyd contains a watermark "IF" at the right-bottom corner, which needs to be masked before performing OCR.
|
||||
|
||||
```python
|
||||
if method is 'deepfloyd':
|
||||
image[-64:, -64:] = 0 # remove watermark, the input image is resized to 512x512
|
||||
```
|
||||
@@ -0,0 +1,366 @@
|
||||
# Adapted from https://github.com/jmhessel/clipscore/blob/1036465276513621f77f1c2208d742e4a430781f/clipscore.py
|
||||
'''
|
||||
Code for CLIPScore (https://arxiv.org/abs/2104.08718)
|
||||
@inproceedings{hessel2021clipscore,
|
||||
title={{CLIPScore:} A Reference-free Evaluation Metric for Image Captioning},
|
||||
author={Hessel, Jack and Holtzman, Ari and Forbes, Maxwell and Bras, Ronan Le and Choi, Yejin},
|
||||
booktitle={EMNLP},
|
||||
year={2021}
|
||||
}
|
||||
'''
|
||||
import argparse
|
||||
import clip
|
||||
from PIL import Image
|
||||
from torchvision.transforms import Compose, Resize, CenterCrop, ToTensor, Normalize
|
||||
import torch
|
||||
import tqdm
|
||||
import numpy as np
|
||||
import sklearn.preprocessing
|
||||
import collections
|
||||
import os
|
||||
import pathlib
|
||||
import json
|
||||
import warnings
|
||||
from packaging import version
|
||||
from pycocoevalcap.tokenizer.ptbtokenizer import PTBTokenizer
|
||||
from pycocoevalcap.meteor.meteor import Meteor
|
||||
from pycocoevalcap.bleu.bleu import Bleu
|
||||
from pycocoevalcap.cider.cider import Cider
|
||||
from pycocoevalcap.rouge.rouge import Rouge
|
||||
from pycocoevalcap.spice.spice import Spice
|
||||
|
||||
|
||||
def get_all_metrics(refs, cands, return_per_cap=False):
|
||||
metrics = []
|
||||
names = []
|
||||
|
||||
pycoco_eval_cap_scorers = [(Bleu(4), 'bleu'),
|
||||
(Meteor(), 'meteor'),
|
||||
(Rouge(), 'rouge'),
|
||||
(Cider(), 'cider'),
|
||||
(Spice(), 'spice')]
|
||||
|
||||
for scorer, name in pycoco_eval_cap_scorers:
|
||||
overall, per_cap = pycoco_eval(scorer, refs, cands)
|
||||
if return_per_cap:
|
||||
metrics.append(per_cap)
|
||||
else:
|
||||
metrics.append(overall)
|
||||
names.append(name)
|
||||
|
||||
metrics = dict(zip(names, metrics))
|
||||
return metrics
|
||||
|
||||
|
||||
def tokenize(refs, cands, no_op=False):
|
||||
# no_op is a debug option to see how significantly not using the PTB tokenizer
|
||||
# affects things
|
||||
tokenizer = PTBTokenizer()
|
||||
|
||||
if no_op:
|
||||
refs = {idx: [r for r in c_refs] for idx, c_refs in enumerate(refs)}
|
||||
cands = {idx: [c] for idx, c in enumerate(cands)}
|
||||
|
||||
else:
|
||||
refs = {idx: [{'caption':r} for r in c_refs] for idx, c_refs in enumerate(refs)}
|
||||
cands = {idx: [{'caption':c}] for idx, c in enumerate(cands)}
|
||||
|
||||
refs = tokenizer.tokenize(refs)
|
||||
cands = tokenizer.tokenize(cands)
|
||||
|
||||
return refs, cands
|
||||
|
||||
|
||||
def pycoco_eval(scorer, refs, cands):
|
||||
'''
|
||||
scorer is assumed to have a compute_score function.
|
||||
refs is a list of lists of strings
|
||||
cands is a list of predictions
|
||||
'''
|
||||
refs, cands = tokenize(refs, cands)
|
||||
average_score, scores = scorer.compute_score(refs, cands)
|
||||
return average_score, scores
|
||||
|
||||
|
||||
def parse_args():
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument(
|
||||
'candidates_json',
|
||||
type=str,
|
||||
help='Candidates json mapping from image_id --> candidate.')
|
||||
|
||||
parser.add_argument(
|
||||
'image_dir',
|
||||
type=str,
|
||||
help='Directory of images, with the filenames as image ids.')
|
||||
|
||||
parser.add_argument(
|
||||
'--references_json',
|
||||
default=None,
|
||||
help='Optional references json mapping from image_id --> [list of references]')
|
||||
|
||||
parser.add_argument(
|
||||
'--compute_other_ref_metrics',
|
||||
default=1,
|
||||
type=int,
|
||||
help='If references is specified, should we compute standard reference-based metrics?')
|
||||
|
||||
parser.add_argument(
|
||||
'--save_per_instance',
|
||||
default=None,
|
||||
help='if set, we will save per instance clipscores to this file')
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
if isinstance(args.save_per_instance, str) and not args.save_per_instance.endswith('.json'):
|
||||
print('if you\'re saving per-instance, please make sure the filepath ends in json.')
|
||||
quit()
|
||||
return args
|
||||
|
||||
|
||||
class CLIPCapDataset(torch.utils.data.Dataset):
|
||||
def __init__(self, data, prefix='A photo depicts'):
|
||||
self.data = data
|
||||
self.prefix = prefix
|
||||
if self.prefix[-1] != ' ':
|
||||
self.prefix += ' '
|
||||
|
||||
def __getitem__(self, idx):
|
||||
c_data = self.data[idx]
|
||||
c_data = clip.tokenize(self.prefix + c_data, truncate=True).squeeze()
|
||||
return {'caption': c_data}
|
||||
|
||||
def __len__(self):
|
||||
return len(self.data)
|
||||
|
||||
|
||||
class CLIPImageDataset(torch.utils.data.Dataset):
|
||||
def __init__(self, data):
|
||||
self.data = data
|
||||
# only 224x224 ViT-B/32 supported for now
|
||||
self.preprocess = self._transform_test(224)
|
||||
|
||||
def _transform_test(self, n_px):
|
||||
return Compose([
|
||||
Resize(n_px, interpolation=Image.BICUBIC),
|
||||
CenterCrop(n_px),
|
||||
lambda image: image.convert("RGB"),
|
||||
ToTensor(),
|
||||
Normalize((0.48145466, 0.4578275, 0.40821073), (0.26862954, 0.26130258, 0.27577711)),
|
||||
])
|
||||
|
||||
def __getitem__(self, idx):
|
||||
c_data = self.data[idx]
|
||||
image = Image.open(c_data)
|
||||
image = self.preprocess(image)
|
||||
return {'image':image}
|
||||
|
||||
def __len__(self):
|
||||
return len(self.data)
|
||||
|
||||
|
||||
def extract_all_captions(captions, model, device, batch_size=256, num_workers=8):
|
||||
data = torch.utils.data.DataLoader(
|
||||
CLIPCapDataset(captions),
|
||||
batch_size=batch_size, num_workers=num_workers, shuffle=False)
|
||||
all_text_features = []
|
||||
with torch.no_grad():
|
||||
for b in tqdm.tqdm(data):
|
||||
b = b['caption'].to(device)
|
||||
all_text_features.append(model.encode_text(b).cpu().numpy())
|
||||
all_text_features = np.vstack(all_text_features)
|
||||
return all_text_features
|
||||
|
||||
|
||||
def extract_all_images(images, model, device, batch_size=64, num_workers=8):
|
||||
data = torch.utils.data.DataLoader(
|
||||
CLIPImageDataset(images),
|
||||
batch_size=batch_size, num_workers=num_workers, shuffle=False)
|
||||
all_image_features = []
|
||||
with torch.no_grad():
|
||||
for b in tqdm.tqdm(data):
|
||||
b = b['image'].to(device)
|
||||
if device == 'cuda':
|
||||
b = b.to(torch.float16)
|
||||
all_image_features.append(model.encode_image(b).cpu().numpy())
|
||||
all_image_features = np.vstack(all_image_features)
|
||||
return all_image_features
|
||||
|
||||
|
||||
def get_clip_score(model, images, candidates, device, w=2.5):
|
||||
'''
|
||||
get standard image-text clipscore.
|
||||
images can either be:
|
||||
- a list of strings specifying filepaths for images
|
||||
- a precomputed, ordered matrix of image features
|
||||
'''
|
||||
if isinstance(images, list):
|
||||
# need to extract image features
|
||||
images = extract_all_images(images, model, device)
|
||||
|
||||
candidates = extract_all_captions(candidates, model, device)
|
||||
|
||||
#as of numpy 1.21, normalize doesn't work properly for float16
|
||||
if version.parse(np.__version__) < version.parse('1.21'):
|
||||
images = sklearn.preprocessing.normalize(images, axis=1)
|
||||
candidates = sklearn.preprocessing.normalize(candidates, axis=1)
|
||||
else:
|
||||
warnings.warn(
|
||||
'due to a numerical instability, new numpy normalization is slightly different than paper results. '
|
||||
'to exactly replicate paper results, please use numpy version less than 1.21, e.g., 1.20.3.')
|
||||
images = images / np.sqrt(np.sum(images**2, axis=1, keepdims=True))
|
||||
candidates = candidates / np.sqrt(np.sum(candidates**2, axis=1, keepdims=True))
|
||||
|
||||
per = w*np.clip(np.sum(images * candidates, axis=1), 0, None)
|
||||
return np.mean(per), per, candidates
|
||||
|
||||
|
||||
def get_refonlyclipscore(model, references, candidates, device):
|
||||
'''
|
||||
The text only side for refclipscore
|
||||
'''
|
||||
if isinstance(candidates, list):
|
||||
candidates = extract_all_captions(candidates, model, device)
|
||||
|
||||
flattened_refs = []
|
||||
flattened_refs_idxs = []
|
||||
for idx, refs in enumerate(references):
|
||||
flattened_refs.extend(refs)
|
||||
flattened_refs_idxs.extend([idx for _ in refs])
|
||||
|
||||
flattened_refs = extract_all_captions(flattened_refs, model, device)
|
||||
|
||||
if version.parse(np.__version__) < version.parse('1.21'):
|
||||
candidates = sklearn.preprocessing.normalize(candidates, axis=1)
|
||||
flattened_refs = sklearn.preprocessing.normalize(flattened_refs, axis=1)
|
||||
else:
|
||||
warnings.warn(
|
||||
'due to a numerical instability, new numpy normalization is slightly different than paper results. '
|
||||
'to exactly replicate paper results, please use numpy version less than 1.21, e.g., 1.20.3.')
|
||||
|
||||
candidates = candidates / np.sqrt(np.sum(candidates**2, axis=1, keepdims=True))
|
||||
flattened_refs = flattened_refs / np.sqrt(np.sum(flattened_refs**2, axis=1, keepdims=True))
|
||||
|
||||
cand_idx2refs = collections.defaultdict(list)
|
||||
for ref_feats, cand_idx in zip(flattened_refs, flattened_refs_idxs):
|
||||
cand_idx2refs[cand_idx].append(ref_feats)
|
||||
|
||||
assert len(cand_idx2refs) == len(candidates)
|
||||
|
||||
cand_idx2refs = {k: np.vstack(v) for k, v in cand_idx2refs.items()}
|
||||
|
||||
per = []
|
||||
for c_idx, cand in tqdm.tqdm(enumerate(candidates)):
|
||||
cur_refs = cand_idx2refs[c_idx]
|
||||
all_sims = cand.dot(cur_refs.transpose())
|
||||
per.append(np.max(all_sims))
|
||||
|
||||
return np.mean(per), per
|
||||
|
||||
|
||||
def cal_clipscore(image_ids, image_paths, text_list, device=None, references=None, scale_weight=1):
|
||||
if device is None:
|
||||
device = "cuda" if torch.cuda.is_available() else "cpu"
|
||||
model, transform = clip.load("ViT-B/32", device=device, jit=False)
|
||||
model.eval()
|
||||
|
||||
image_feats = extract_all_images(image_paths, model, device, batch_size=64, num_workers=8)
|
||||
|
||||
# get image-text clipscore
|
||||
_, per_instance_image_text, candidate_feats = get_clip_score(model, image_feats, text_list, device, w=scale_weight)
|
||||
|
||||
if references:
|
||||
# get text-text clipscore
|
||||
_, per_instance_text_text = get_refonlyclipscore(model, references, candidate_feats, device)
|
||||
# F-score
|
||||
refclipscores = 2 * per_instance_image_text * per_instance_text_text / (per_instance_image_text + per_instance_text_text)
|
||||
scores = {image_id: {'CLIPScore': float(clipscore), 'RefCLIPScore': float(refclipscore)}
|
||||
for image_id, clipscore, refclipscore in
|
||||
zip(image_ids, per_instance_image_text, refclipscores)}
|
||||
|
||||
other_metrics = get_all_metrics(references, text_list)
|
||||
for k, v in other_metrics.items():
|
||||
if k == 'bleu':
|
||||
for bidx, sc in enumerate(v):
|
||||
print('BLEU-{}: {:.4f}'.format(bidx+1, sc))
|
||||
else:
|
||||
print('{}: {:.4f}'.format(k.upper(), v))
|
||||
print('CLIPScore: {:.4f}'.format(np.mean([s['CLIPScore'] for s in scores.values()])))
|
||||
print('RefCLIPScore: {:.4f}'.format(np.mean([s['RefCLIPScore'] for s in scores.values()])))
|
||||
|
||||
else:
|
||||
scores = {image_id: {'CLIPScore': float(clipscore)}
|
||||
for image_id, clipscore in
|
||||
zip(image_ids, per_instance_image_text)}
|
||||
print('CLIPScore: {:.4f}'.format(np.mean([s['CLIPScore'] for s in scores.values()])))
|
||||
|
||||
return scores
|
||||
|
||||
def main():
|
||||
args = parse_args()
|
||||
|
||||
image_paths = [os.path.join(args.image_dir, path) for path in os.listdir(args.image_dir)
|
||||
if path.endswith(('.png', '.jpg', '.jpeg', '.tiff'))]
|
||||
image_ids = [pathlib.Path(path).stem for path in image_paths]
|
||||
|
||||
with open(args.candidates_json) as f:
|
||||
candidates = json.load(f)
|
||||
candidates = [candidates[cid] for cid in image_ids]
|
||||
|
||||
if args.references_json:
|
||||
with open(args.references_json) as f:
|
||||
references = json.load(f)
|
||||
references = [references[cid] for cid in image_ids]
|
||||
if isinstance(references[0], str):
|
||||
references = [[r] for r in references]
|
||||
|
||||
device = "cuda" if torch.cuda.is_available() else "cpu"
|
||||
if device == 'cpu':
|
||||
warnings.warn(
|
||||
'CLIP runs in full float32 on CPU. Results in paper were computed on GPU, which uses float16. '
|
||||
'If you\'re reporting results on CPU, please note this when you report.')
|
||||
model, transform = clip.load("ViT-B/32", device=device, jit=False)
|
||||
model.eval()
|
||||
|
||||
image_feats = extract_all_images(
|
||||
image_paths, model, device, batch_size=64, num_workers=8)
|
||||
|
||||
# get image-text clipscore
|
||||
_, per_instance_image_text, candidate_feats = get_clip_score(
|
||||
model, image_feats, candidates, device)
|
||||
|
||||
if args.references_json:
|
||||
# get text-text clipscore
|
||||
_, per_instance_text_text = get_refonlyclipscore(
|
||||
model, references, candidate_feats, device)
|
||||
# F-score
|
||||
refclipscores = 2 * per_instance_image_text * per_instance_text_text / (per_instance_image_text + per_instance_text_text)
|
||||
scores = {image_id: {'CLIPScore': float(clipscore), 'RefCLIPScore': float(refclipscore)}
|
||||
for image_id, clipscore, refclipscore in
|
||||
zip(image_ids, per_instance_image_text, refclipscores)}
|
||||
|
||||
else:
|
||||
scores = {image_id: {'CLIPScore': float(clipscore)}
|
||||
for image_id, clipscore in
|
||||
zip(image_ids, per_instance_image_text)}
|
||||
print('CLIPScore: {:.4f}'.format(np.mean([s['CLIPScore'] for s in scores.values()])))
|
||||
|
||||
if args.references_json:
|
||||
if args.compute_other_ref_metrics:
|
||||
other_metrics = generation_eval_utils.get_all_metrics(references, candidates)
|
||||
for k, v in other_metrics.items():
|
||||
if k == 'bleu':
|
||||
for bidx, sc in enumerate(v):
|
||||
print('BLEU-{}: {:.4f}'.format(bidx+1, sc))
|
||||
else:
|
||||
print('{}: {:.4f}'.format(k.upper(), v))
|
||||
print('CLIPScore: {:.4f}'.format(np.mean([s['CLIPScore'] for s in scores.values()])))
|
||||
print('RefCLIPScore: {:.4f}'.format(np.mean([s['RefCLIPScore'] for s in scores.values()])))
|
||||
|
||||
if args.save_per_instance:
|
||||
with open(args.save_per_instance, 'w') as f:
|
||||
f.write(json.dumps(scores))
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
main()
|
||||
@@ -0,0 +1,5 @@
|
||||
python MARIOEval_evaluate.py \
|
||||
--gpu 0 \
|
||||
--dataset TMDBEval500 \
|
||||
--root /path/to/eval \
|
||||
--method textdiffuser
|
||||
@@ -0,0 +1,333 @@
|
||||
# Adapted from https://github.com/mseitzer/pytorch-fid/blob/0a754fb8e66021700478fd365b79c2eaa316e31b/src/pytorch_fid/fid_score.py
|
||||
"""Calculates the Frechet Inception Distance (FID) to evalulate GANs
|
||||
|
||||
The FID metric calculates the distance between two distributions of images.
|
||||
Typically, we have summary statistics (mean & covariance matrix) of one
|
||||
of these distributions, while the 2nd distribution is given by a GAN.
|
||||
|
||||
When run as a stand-alone program, it compares the distribution of
|
||||
images that are stored as PNG/JPEG at a specified location with a
|
||||
distribution given by summary statistics (in pickle format).
|
||||
|
||||
The FID is calculated by assuming that X_1 and X_2 are the activations of
|
||||
the pool_3 layer of the inception net for generated samples and real world
|
||||
samples respectively.
|
||||
|
||||
See --help to see further details.
|
||||
|
||||
Code apapted from https://github.com/bioinf-jku/TTUR to use PyTorch instead
|
||||
of Tensorflow
|
||||
|
||||
Copyright 2018 Institute of Bioinformatics, JKU Linz
|
||||
|
||||
Licensed under the Apache License, Version 2.0 (the "License");
|
||||
you may not use this file except in compliance with the License.
|
||||
You may obtain a copy of the License at
|
||||
|
||||
http://www.apache.org/licenses/LICENSE-2.0
|
||||
|
||||
Unless required by applicable law or agreed to in writing, software
|
||||
distributed under the License is distributed on an "AS IS" BASIS,
|
||||
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
See the License for the specific language governing permissions and
|
||||
limitations under the License.
|
||||
"""
|
||||
import os
|
||||
import pathlib
|
||||
from argparse import ArgumentDefaultsHelpFormatter, ArgumentParser
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
import torchvision.transforms as TF
|
||||
from PIL import Image
|
||||
from scipy import linalg
|
||||
from torch.nn.functional import adaptive_avg_pool2d
|
||||
|
||||
try:
|
||||
from tqdm import tqdm
|
||||
except ImportError:
|
||||
# If tqdm is not available, provide a mock version of it
|
||||
def tqdm(x):
|
||||
return x
|
||||
|
||||
from pytorch_fid.inception import InceptionV3
|
||||
|
||||
parser = ArgumentParser(formatter_class=ArgumentDefaultsHelpFormatter)
|
||||
parser.add_argument('--batch-size', type=int, default=50,
|
||||
help='Batch size to use')
|
||||
parser.add_argument('--num-workers', type=int,
|
||||
help=('Number of processes to use for data loading. '
|
||||
'Defaults to `min(8, num_cpus)`'))
|
||||
parser.add_argument('--device', type=str, default=None,
|
||||
help='Device to use. Like cuda, cuda:0 or cpu')
|
||||
parser.add_argument('--dims', type=int, default=2048,
|
||||
choices=list(InceptionV3.BLOCK_INDEX_BY_DIM),
|
||||
help=('Dimensionality of Inception features to use. '
|
||||
'By default, uses pool3 features'))
|
||||
parser.add_argument('--save-stats', action='store_true',
|
||||
help=('Generate an npz archive from a directory of samples. '
|
||||
'The first path is used as input and the second as output.'))
|
||||
parser.add_argument('path', type=str, nargs=2,
|
||||
help=('Paths to the generated images or '
|
||||
'to .npz statistic files'))
|
||||
|
||||
IMAGE_EXTENSIONS = {'bmp', 'jpg', 'jpeg', 'pgm', 'png', 'ppm',
|
||||
'tif', 'tiff', 'webp'}
|
||||
|
||||
|
||||
class ImagePathDataset(torch.utils.data.Dataset):
|
||||
def __init__(self, files, transforms=None):
|
||||
self.files = files
|
||||
self.transforms = transforms
|
||||
|
||||
def __len__(self):
|
||||
return len(self.files)
|
||||
|
||||
def __getitem__(self, i):
|
||||
path = self.files[i]
|
||||
img = Image.open(path).convert('RGB')
|
||||
if self.transforms is not None:
|
||||
img = self.transforms(img)
|
||||
return img
|
||||
|
||||
|
||||
def get_activations(files, model, batch_size=50, dims=2048, device='cpu',
|
||||
num_workers=1):
|
||||
"""Calculates the activations of the pool_3 layer for all images.
|
||||
|
||||
Params:
|
||||
-- files : List of image files paths
|
||||
-- model : Instance of inception model
|
||||
-- batch_size : Batch size of images for the model to process at once.
|
||||
Make sure that the number of samples is a multiple of
|
||||
the batch size, otherwise some samples are ignored. This
|
||||
behavior is retained to match the original FID score
|
||||
implementation.
|
||||
-- dims : Dimensionality of features returned by Inception
|
||||
-- device : Device to run calculations
|
||||
-- num_workers : Number of parallel dataloader workers
|
||||
|
||||
Returns:
|
||||
-- A numpy array of dimension (num images, dims) that contains the
|
||||
activations of the given tensor when feeding inception with the
|
||||
query tensor.
|
||||
"""
|
||||
model.eval()
|
||||
|
||||
if batch_size > len(files):
|
||||
print(('Warning: batch size is bigger than the data size. '
|
||||
'Setting batch size to data size'))
|
||||
batch_size = len(files)
|
||||
|
||||
dataset = ImagePathDataset(files, transforms=TF.ToTensor())
|
||||
dataloader = torch.utils.data.DataLoader(dataset,
|
||||
batch_size=batch_size,
|
||||
shuffle=False,
|
||||
drop_last=False,
|
||||
num_workers=num_workers)
|
||||
|
||||
pred_arr = np.empty((len(files), dims))
|
||||
|
||||
start_idx = 0
|
||||
|
||||
for batch in tqdm(dataloader):
|
||||
batch = batch.to(device)
|
||||
|
||||
with torch.no_grad():
|
||||
pred = model(batch)[0]
|
||||
|
||||
# If model output is not scalar, apply global spatial average pooling.
|
||||
# This happens if you choose a dimensionality not equal 2048.
|
||||
if pred.size(2) != 1 or pred.size(3) != 1:
|
||||
pred = adaptive_avg_pool2d(pred, output_size=(1, 1))
|
||||
|
||||
pred = pred.squeeze(3).squeeze(2).cpu().numpy()
|
||||
|
||||
pred_arr[start_idx:start_idx + pred.shape[0]] = pred
|
||||
|
||||
start_idx = start_idx + pred.shape[0]
|
||||
|
||||
return pred_arr
|
||||
|
||||
|
||||
def calculate_frechet_distance(mu1, sigma1, mu2, sigma2, eps=1e-6):
|
||||
"""Numpy implementation of the Frechet Distance.
|
||||
The Frechet distance between two multivariate Gaussians X_1 ~ N(mu_1, C_1)
|
||||
and X_2 ~ N(mu_2, C_2) is
|
||||
d^2 = ||mu_1 - mu_2||^2 + Tr(C_1 + C_2 - 2*sqrt(C_1*C_2)).
|
||||
|
||||
Stable version by Dougal J. Sutherland.
|
||||
|
||||
Params:
|
||||
-- mu1 : Numpy array containing the activations of a layer of the
|
||||
inception net (like returned by the function 'get_predictions')
|
||||
for generated samples.
|
||||
-- mu2 : The sample mean over activations, precalculated on an
|
||||
representative data set.
|
||||
-- sigma1: The covariance matrix over activations for generated samples.
|
||||
-- sigma2: The covariance matrix over activations, precalculated on an
|
||||
representative data set.
|
||||
|
||||
Returns:
|
||||
-- : The Frechet Distance.
|
||||
"""
|
||||
|
||||
mu1 = np.atleast_1d(mu1)
|
||||
mu2 = np.atleast_1d(mu2)
|
||||
|
||||
sigma1 = np.atleast_2d(sigma1)
|
||||
sigma2 = np.atleast_2d(sigma2)
|
||||
|
||||
assert mu1.shape == mu2.shape, \
|
||||
'Training and test mean vectors have different lengths'
|
||||
assert sigma1.shape == sigma2.shape, \
|
||||
'Training and test covariances have different dimensions'
|
||||
|
||||
diff = mu1 - mu2
|
||||
|
||||
# Product might be almost singular
|
||||
covmean, _ = linalg.sqrtm(sigma1.dot(sigma2), disp=False)
|
||||
if not np.isfinite(covmean).all():
|
||||
msg = ('fid calculation produces singular product; '
|
||||
'adding %s to diagonal of cov estimates') % eps
|
||||
print(msg)
|
||||
offset = np.eye(sigma1.shape[0]) * eps
|
||||
covmean = linalg.sqrtm((sigma1 + offset).dot(sigma2 + offset))
|
||||
|
||||
# Numerical error might give slight imaginary component
|
||||
if np.iscomplexobj(covmean):
|
||||
if not np.allclose(np.diagonal(covmean).imag, 0, atol=1e-3):
|
||||
m = np.max(np.abs(covmean.imag))
|
||||
raise ValueError('Imaginary component {}'.format(m))
|
||||
covmean = covmean.real
|
||||
|
||||
tr_covmean = np.trace(covmean)
|
||||
|
||||
return (diff.dot(diff) + np.trace(sigma1)
|
||||
+ np.trace(sigma2) - 2 * tr_covmean)
|
||||
|
||||
|
||||
def calculate_activation_statistics(files, model, batch_size=50, dims=2048,
|
||||
device='cpu', num_workers=1):
|
||||
"""Calculation of the statistics used by the FID.
|
||||
Params:
|
||||
-- files : List of image files paths
|
||||
-- model : Instance of inception model
|
||||
-- batch_size : The images numpy array is split into batches with
|
||||
batch size batch_size. A reasonable batch size
|
||||
depends on the hardware.
|
||||
-- dims : Dimensionality of features returned by Inception
|
||||
-- device : Device to run calculations
|
||||
-- num_workers : Number of parallel dataloader workers
|
||||
|
||||
Returns:
|
||||
-- mu : The mean over samples of the activations of the pool_3 layer of
|
||||
the inception model.
|
||||
-- sigma : The covariance matrix of the activations of the pool_3 layer of
|
||||
the inception model.
|
||||
"""
|
||||
act = get_activations(files, model, batch_size, dims, device, num_workers)
|
||||
mu = np.mean(act, axis=0)
|
||||
sigma = np.cov(act, rowvar=False)
|
||||
return mu, sigma
|
||||
|
||||
|
||||
def compute_statistics_of_path(path, model, batch_size, dims, device,
|
||||
num_workers=1):
|
||||
if type(path) is not list and path.endswith('.npz'):
|
||||
with np.load(path) as f:
|
||||
m, s = f['mu'][:], f['sigma'][:]
|
||||
else:
|
||||
if type(path) is list:
|
||||
files = []
|
||||
for p in path:
|
||||
p = pathlib.Path(p)
|
||||
files += sorted([file for ext in IMAGE_EXTENSIONS for file in p.glob('*.{}'.format(ext))])
|
||||
files = sorted(files)
|
||||
else:
|
||||
path = pathlib.Path(path)
|
||||
files = sorted([file for ext in IMAGE_EXTENSIONS for file in path.glob('*.{}'.format(ext))])
|
||||
m, s = calculate_activation_statistics(files, model, batch_size, dims, device, num_workers)
|
||||
|
||||
return m, s
|
||||
|
||||
|
||||
def calculate_fid_given_paths(paths, batch_size=50, device="cuda:0", dims=2048, num_workers=1):
|
||||
"""Calculates the FID of two paths"""
|
||||
for p in paths:
|
||||
if type(p) is list:
|
||||
for subp in p:
|
||||
if not os.path.exists(subp):
|
||||
raise RuntimeError('Invalid path: %s' % subp)
|
||||
else:
|
||||
if not os.path.exists(p):
|
||||
raise RuntimeError('Invalid path: %s' % p)
|
||||
|
||||
block_idx = InceptionV3.BLOCK_INDEX_BY_DIM[dims]
|
||||
|
||||
model = InceptionV3([block_idx]).to(device)
|
||||
|
||||
m1, s1 = compute_statistics_of_path(paths[0], model, batch_size,
|
||||
dims, device, num_workers)
|
||||
m2, s2 = compute_statistics_of_path(paths[1], model, batch_size,
|
||||
dims, device, num_workers)
|
||||
fid_value = calculate_frechet_distance(m1, s1, m2, s2)
|
||||
|
||||
return fid_value
|
||||
|
||||
|
||||
def save_fid_stats(paths, batch_size, device, dims, num_workers=1):
|
||||
"""Calculates the FID of two paths"""
|
||||
if not os.path.exists(paths[0]):
|
||||
raise RuntimeError('Invalid path: %s' % paths[0])
|
||||
|
||||
if os.path.exists(paths[1]):
|
||||
raise RuntimeError('Existing output file: %s' % paths[1])
|
||||
|
||||
block_idx = InceptionV3.BLOCK_INDEX_BY_DIM[dims]
|
||||
|
||||
model = InceptionV3([block_idx]).to(device)
|
||||
|
||||
print(f"Saving statistics for {paths[0]}")
|
||||
|
||||
m1, s1 = compute_statistics_of_path(paths[0], model, batch_size,
|
||||
dims, device, num_workers)
|
||||
|
||||
np.savez_compressed(paths[1], mu=m1, sigma=s1)
|
||||
|
||||
|
||||
def main():
|
||||
args = parser.parse_args()
|
||||
|
||||
if args.device is None:
|
||||
device = torch.device('cuda' if (torch.cuda.is_available()) else 'cpu')
|
||||
else:
|
||||
device = torch.device(args.device)
|
||||
|
||||
if args.num_workers is None:
|
||||
try:
|
||||
num_cpus = len(os.sched_getaffinity(0))
|
||||
except AttributeError:
|
||||
# os.sched_getaffinity is not available under Windows, use
|
||||
# os.cpu_count instead (which may not return the *available* number
|
||||
# of CPUs).
|
||||
num_cpus = os.cpu_count()
|
||||
|
||||
num_workers = min(num_cpus, 8) if num_cpus is not None else 0
|
||||
else:
|
||||
num_workers = args.num_workers
|
||||
|
||||
if args.save_stats:
|
||||
save_fid_stats(args.path, args.batch_size, device, args.dims, num_workers)
|
||||
return
|
||||
|
||||
fid_value = calculate_fid_given_paths(args.path,
|
||||
args.batch_size,
|
||||
device,
|
||||
args.dims,
|
||||
num_workers)
|
||||
print('FID: ', fid_value)
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
main()
|
||||
@@ -0,0 +1,5 @@
|
||||
python MARIOEval_generate.py \
|
||||
--dataset TMDBEval500 \
|
||||
--root /path/to/eval \
|
||||
--method stablediffusion \
|
||||
--gpu 0
|
||||
@@ -0,0 +1,342 @@
|
||||
# Copied from https://github.com/mseitzer/pytorch-fid/blob/0a754fb8e66021700478fd365b79c2eaa316e31b/src/pytorch_fid/inception.py
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
import torchvision
|
||||
|
||||
try:
|
||||
from torchvision.models.utils import load_state_dict_from_url
|
||||
except ImportError:
|
||||
from torch.utils.model_zoo import load_url as load_state_dict_from_url
|
||||
|
||||
# Inception weights ported to Pytorch from
|
||||
# http://download.tensorflow.org/models/image/imagenet/inception-2015-12-05.tgz
|
||||
FID_WEIGHTS_URL = 'https://github.com/mseitzer/pytorch-fid/releases/download/fid_weights/pt_inception-2015-12-05-6726825d.pth' # noqa: E501
|
||||
|
||||
|
||||
class InceptionV3(nn.Module):
|
||||
"""Pretrained InceptionV3 network returning feature maps"""
|
||||
|
||||
# Index of default block of inception to return,
|
||||
# corresponds to output of final average pooling
|
||||
DEFAULT_BLOCK_INDEX = 3
|
||||
|
||||
# Maps feature dimensionality to their output blocks indices
|
||||
BLOCK_INDEX_BY_DIM = {
|
||||
64: 0, # First max pooling features
|
||||
192: 1, # Second max pooling featurs
|
||||
768: 2, # Pre-aux classifier features
|
||||
2048: 3 # Final average pooling features
|
||||
}
|
||||
|
||||
def __init__(self,
|
||||
output_blocks=(DEFAULT_BLOCK_INDEX,),
|
||||
resize_input=True,
|
||||
normalize_input=True,
|
||||
requires_grad=False,
|
||||
use_fid_inception=True):
|
||||
"""Build pretrained InceptionV3
|
||||
|
||||
Parameters
|
||||
----------
|
||||
output_blocks : list of int
|
||||
Indices of blocks to return features of. Possible values are:
|
||||
- 0: corresponds to output of first max pooling
|
||||
- 1: corresponds to output of second max pooling
|
||||
- 2: corresponds to output which is fed to aux classifier
|
||||
- 3: corresponds to output of final average pooling
|
||||
resize_input : bool
|
||||
If true, bilinearly resizes input to width and height 299 before
|
||||
feeding input to model. As the network without fully connected
|
||||
layers is fully convolutional, it should be able to handle inputs
|
||||
of arbitrary size, so resizing might not be strictly needed
|
||||
normalize_input : bool
|
||||
If true, scales the input from range (0, 1) to the range the
|
||||
pretrained Inception network expects, namely (-1, 1)
|
||||
requires_grad : bool
|
||||
If true, parameters of the model require gradients. Possibly useful
|
||||
for finetuning the network
|
||||
use_fid_inception : bool
|
||||
If true, uses the pretrained Inception model used in Tensorflow's
|
||||
FID implementation. If false, uses the pretrained Inception model
|
||||
available in torchvision. The FID Inception model has different
|
||||
weights and a slightly different structure from torchvision's
|
||||
Inception model. If you want to compute FID scores, you are
|
||||
strongly advised to set this parameter to true to get comparable
|
||||
results.
|
||||
"""
|
||||
super(InceptionV3, self).__init__()
|
||||
|
||||
self.resize_input = resize_input
|
||||
self.normalize_input = normalize_input
|
||||
self.output_blocks = sorted(output_blocks)
|
||||
self.last_needed_block = max(output_blocks)
|
||||
|
||||
assert self.last_needed_block <= 3, \
|
||||
'Last possible output block index is 3'
|
||||
|
||||
self.blocks = nn.ModuleList()
|
||||
|
||||
if use_fid_inception:
|
||||
inception = fid_inception_v3()
|
||||
else:
|
||||
inception = _inception_v3(weights='DEFAULT')
|
||||
|
||||
# Block 0: input to maxpool1
|
||||
block0 = [
|
||||
inception.Conv2d_1a_3x3,
|
||||
inception.Conv2d_2a_3x3,
|
||||
inception.Conv2d_2b_3x3,
|
||||
nn.MaxPool2d(kernel_size=3, stride=2)
|
||||
]
|
||||
self.blocks.append(nn.Sequential(*block0))
|
||||
|
||||
# Block 1: maxpool1 to maxpool2
|
||||
if self.last_needed_block >= 1:
|
||||
block1 = [
|
||||
inception.Conv2d_3b_1x1,
|
||||
inception.Conv2d_4a_3x3,
|
||||
nn.MaxPool2d(kernel_size=3, stride=2)
|
||||
]
|
||||
self.blocks.append(nn.Sequential(*block1))
|
||||
|
||||
# Block 2: maxpool2 to aux classifier
|
||||
if self.last_needed_block >= 2:
|
||||
block2 = [
|
||||
inception.Mixed_5b,
|
||||
inception.Mixed_5c,
|
||||
inception.Mixed_5d,
|
||||
inception.Mixed_6a,
|
||||
inception.Mixed_6b,
|
||||
inception.Mixed_6c,
|
||||
inception.Mixed_6d,
|
||||
inception.Mixed_6e,
|
||||
]
|
||||
self.blocks.append(nn.Sequential(*block2))
|
||||
|
||||
# Block 3: aux classifier to final avgpool
|
||||
if self.last_needed_block >= 3:
|
||||
block3 = [
|
||||
inception.Mixed_7a,
|
||||
inception.Mixed_7b,
|
||||
inception.Mixed_7c,
|
||||
nn.AdaptiveAvgPool2d(output_size=(1, 1))
|
||||
]
|
||||
self.blocks.append(nn.Sequential(*block3))
|
||||
|
||||
for param in self.parameters():
|
||||
param.requires_grad = requires_grad
|
||||
|
||||
def forward(self, inp):
|
||||
"""Get Inception feature maps
|
||||
|
||||
Parameters
|
||||
----------
|
||||
inp : torch.autograd.Variable
|
||||
Input tensor of shape Bx3xHxW. Values are expected to be in
|
||||
range (0, 1)
|
||||
|
||||
Returns
|
||||
-------
|
||||
List of torch.autograd.Variable, corresponding to the selected output
|
||||
block, sorted ascending by index
|
||||
"""
|
||||
outp = []
|
||||
x = inp
|
||||
|
||||
if self.resize_input:
|
||||
x = F.interpolate(x,
|
||||
size=(299, 299),
|
||||
mode='bilinear',
|
||||
align_corners=False)
|
||||
|
||||
if self.normalize_input:
|
||||
x = 2 * x - 1 # Scale from range (0, 1) to range (-1, 1)
|
||||
|
||||
for idx, block in enumerate(self.blocks):
|
||||
x = block(x)
|
||||
if idx in self.output_blocks:
|
||||
outp.append(x)
|
||||
|
||||
if idx == self.last_needed_block:
|
||||
break
|
||||
|
||||
return outp
|
||||
|
||||
|
||||
def _inception_v3(*args, **kwargs):
|
||||
"""Wraps `torchvision.models.inception_v3`"""
|
||||
try:
|
||||
version = tuple(map(int, torchvision.__version__.split('.')[:2]))
|
||||
except ValueError:
|
||||
# Just a caution against weird version strings
|
||||
version = (0,)
|
||||
|
||||
# Skips default weight inititialization if supported by torchvision
|
||||
# version. See https://github.com/mseitzer/pytorch-fid/issues/28.
|
||||
if version >= (0, 6):
|
||||
kwargs['init_weights'] = False
|
||||
|
||||
# Backwards compatibility: `weights` argument was handled by `pretrained`
|
||||
# argument prior to version 0.13.
|
||||
if version < (0, 13) and 'weights' in kwargs:
|
||||
if kwargs['weights'] == 'DEFAULT':
|
||||
kwargs['pretrained'] = True
|
||||
elif kwargs['weights'] is None:
|
||||
kwargs['pretrained'] = False
|
||||
else:
|
||||
raise ValueError(
|
||||
'weights=={} not supported in torchvision {}'.format(
|
||||
kwargs['weights'], torchvision.__version__
|
||||
)
|
||||
)
|
||||
del kwargs['weights']
|
||||
|
||||
return torchvision.models.inception_v3(*args, **kwargs)
|
||||
|
||||
|
||||
def fid_inception_v3():
|
||||
"""Build pretrained Inception model for FID computation
|
||||
|
||||
The Inception model for FID computation uses a different set of weights
|
||||
and has a slightly different structure than torchvision's Inception.
|
||||
|
||||
This method first constructs torchvision's Inception and then patches the
|
||||
necessary parts that are different in the FID Inception model.
|
||||
"""
|
||||
inception = _inception_v3(num_classes=1008,
|
||||
aux_logits=False,
|
||||
weights=None)
|
||||
inception.Mixed_5b = FIDInceptionA(192, pool_features=32)
|
||||
inception.Mixed_5c = FIDInceptionA(256, pool_features=64)
|
||||
inception.Mixed_5d = FIDInceptionA(288, pool_features=64)
|
||||
inception.Mixed_6b = FIDInceptionC(768, channels_7x7=128)
|
||||
inception.Mixed_6c = FIDInceptionC(768, channels_7x7=160)
|
||||
inception.Mixed_6d = FIDInceptionC(768, channels_7x7=160)
|
||||
inception.Mixed_6e = FIDInceptionC(768, channels_7x7=192)
|
||||
inception.Mixed_7b = FIDInceptionE_1(1280)
|
||||
inception.Mixed_7c = FIDInceptionE_2(2048)
|
||||
|
||||
state_dict = load_state_dict_from_url(FID_WEIGHTS_URL, progress=True)
|
||||
inception.load_state_dict(state_dict)
|
||||
return inception
|
||||
|
||||
|
||||
class FIDInceptionA(torchvision.models.inception.InceptionA):
|
||||
"""InceptionA block patched for FID computation"""
|
||||
def __init__(self, in_channels, pool_features):
|
||||
super(FIDInceptionA, self).__init__(in_channels, pool_features)
|
||||
|
||||
def forward(self, x):
|
||||
branch1x1 = self.branch1x1(x)
|
||||
|
||||
branch5x5 = self.branch5x5_1(x)
|
||||
branch5x5 = self.branch5x5_2(branch5x5)
|
||||
|
||||
branch3x3dbl = self.branch3x3dbl_1(x)
|
||||
branch3x3dbl = self.branch3x3dbl_2(branch3x3dbl)
|
||||
branch3x3dbl = self.branch3x3dbl_3(branch3x3dbl)
|
||||
|
||||
# Patch: Tensorflow's average pool does not use the padded zero's in
|
||||
# its average calculation
|
||||
branch_pool = F.avg_pool2d(x, kernel_size=3, stride=1, padding=1,
|
||||
count_include_pad=False)
|
||||
branch_pool = self.branch_pool(branch_pool)
|
||||
|
||||
outputs = [branch1x1, branch5x5, branch3x3dbl, branch_pool]
|
||||
return torch.cat(outputs, 1)
|
||||
|
||||
|
||||
class FIDInceptionC(torchvision.models.inception.InceptionC):
|
||||
"""InceptionC block patched for FID computation"""
|
||||
def __init__(self, in_channels, channels_7x7):
|
||||
super(FIDInceptionC, self).__init__(in_channels, channels_7x7)
|
||||
|
||||
def forward(self, x):
|
||||
branch1x1 = self.branch1x1(x)
|
||||
|
||||
branch7x7 = self.branch7x7_1(x)
|
||||
branch7x7 = self.branch7x7_2(branch7x7)
|
||||
branch7x7 = self.branch7x7_3(branch7x7)
|
||||
|
||||
branch7x7dbl = self.branch7x7dbl_1(x)
|
||||
branch7x7dbl = self.branch7x7dbl_2(branch7x7dbl)
|
||||
branch7x7dbl = self.branch7x7dbl_3(branch7x7dbl)
|
||||
branch7x7dbl = self.branch7x7dbl_4(branch7x7dbl)
|
||||
branch7x7dbl = self.branch7x7dbl_5(branch7x7dbl)
|
||||
|
||||
# Patch: Tensorflow's average pool does not use the padded zero's in
|
||||
# its average calculation
|
||||
branch_pool = F.avg_pool2d(x, kernel_size=3, stride=1, padding=1,
|
||||
count_include_pad=False)
|
||||
branch_pool = self.branch_pool(branch_pool)
|
||||
|
||||
outputs = [branch1x1, branch7x7, branch7x7dbl, branch_pool]
|
||||
return torch.cat(outputs, 1)
|
||||
|
||||
|
||||
class FIDInceptionE_1(torchvision.models.inception.InceptionE):
|
||||
"""First InceptionE block patched for FID computation"""
|
||||
def __init__(self, in_channels):
|
||||
super(FIDInceptionE_1, self).__init__(in_channels)
|
||||
|
||||
def forward(self, x):
|
||||
branch1x1 = self.branch1x1(x)
|
||||
|
||||
branch3x3 = self.branch3x3_1(x)
|
||||
branch3x3 = [
|
||||
self.branch3x3_2a(branch3x3),
|
||||
self.branch3x3_2b(branch3x3),
|
||||
]
|
||||
branch3x3 = torch.cat(branch3x3, 1)
|
||||
|
||||
branch3x3dbl = self.branch3x3dbl_1(x)
|
||||
branch3x3dbl = self.branch3x3dbl_2(branch3x3dbl)
|
||||
branch3x3dbl = [
|
||||
self.branch3x3dbl_3a(branch3x3dbl),
|
||||
self.branch3x3dbl_3b(branch3x3dbl),
|
||||
]
|
||||
branch3x3dbl = torch.cat(branch3x3dbl, 1)
|
||||
|
||||
# Patch: Tensorflow's average pool does not use the padded zero's in
|
||||
# its average calculation
|
||||
branch_pool = F.avg_pool2d(x, kernel_size=3, stride=1, padding=1,
|
||||
count_include_pad=False)
|
||||
branch_pool = self.branch_pool(branch_pool)
|
||||
|
||||
outputs = [branch1x1, branch3x3, branch3x3dbl, branch_pool]
|
||||
return torch.cat(outputs, 1)
|
||||
|
||||
|
||||
class FIDInceptionE_2(torchvision.models.inception.InceptionE):
|
||||
"""Second InceptionE block patched for FID computation"""
|
||||
def __init__(self, in_channels):
|
||||
super(FIDInceptionE_2, self).__init__(in_channels)
|
||||
|
||||
def forward(self, x):
|
||||
branch1x1 = self.branch1x1(x)
|
||||
|
||||
branch3x3 = self.branch3x3_1(x)
|
||||
branch3x3 = [
|
||||
self.branch3x3_2a(branch3x3),
|
||||
self.branch3x3_2b(branch3x3),
|
||||
]
|
||||
branch3x3 = torch.cat(branch3x3, 1)
|
||||
|
||||
branch3x3dbl = self.branch3x3dbl_1(x)
|
||||
branch3x3dbl = self.branch3x3dbl_2(branch3x3dbl)
|
||||
branch3x3dbl = [
|
||||
self.branch3x3dbl_3a(branch3x3dbl),
|
||||
self.branch3x3dbl_3b(branch3x3dbl),
|
||||
]
|
||||
branch3x3dbl = torch.cat(branch3x3dbl, 1)
|
||||
|
||||
# Patch: The FID Inception model uses max pooling instead of average
|
||||
# pooling. This is likely an error in this specific Inception
|
||||
# implementation, as other Inception models use average pooling here
|
||||
# (which matches the description in the paper).
|
||||
branch_pool = F.max_pool2d(x, kernel_size=3, stride=1, padding=1)
|
||||
branch_pool = self.branch_pool(branch_pool)
|
||||
|
||||
outputs = [branch1x1, branch3x3, branch3x3dbl, branch_pool]
|
||||
return torch.cat(outputs, 1)
|
||||
@@ -0,0 +1,98 @@
|
||||
# ------------------------------------------
|
||||
# TextDiffuser: Diffusion Models as Text Painters
|
||||
# Paper Link: https://arxiv.org/abs/2305.10855
|
||||
# Code Link: https://github.com/microsoft/unilm/tree/master/textdiffuser
|
||||
# Copyright (c) Microsoft Corporation.
|
||||
# This file provides the inference script.
|
||||
# ------------------------------------------
|
||||
|
||||
import os
|
||||
import re
|
||||
import copy
|
||||
|
||||
gts = {
|
||||
'ChineseDrawText': [],
|
||||
'DrawBenchText': [],
|
||||
'DrawTextCreative': [],
|
||||
'LAIONEval4000': [],
|
||||
'OpenLibraryEval500': [],
|
||||
'TMDBEval500': [],
|
||||
}
|
||||
|
||||
results = {
|
||||
'stablediffusion': {'cnt':0, 'p':0, 'r':0, 'f':0, 'acc':0},
|
||||
'textdiffuser': {'cnt':0, 'p':0, 'r':0, 'f':0, 'acc':0},
|
||||
'controlnet': {'cnt':0, 'p':0, 'r':0, 'f':0, 'acc':0},
|
||||
'deepfloyd': {'cnt':0, 'p':0, 'r':0, 'f':0, 'acc':0},
|
||||
}
|
||||
|
||||
def get_key_words(text: str):
|
||||
words = []
|
||||
text = text
|
||||
matches = re.findall(r"'(.*?)'", text) # find the keywords enclosed by ''
|
||||
if matches:
|
||||
for match in matches:
|
||||
words.extend(match.split())
|
||||
|
||||
return words
|
||||
|
||||
|
||||
# load gt
|
||||
files = os.listdir('/path/to/MARIOEval')
|
||||
for file in files:
|
||||
lines = open(os.path.join('/path/to/MARIOEval', file, f'{file}.txt')).readlines()
|
||||
for line in lines:
|
||||
line = line.strip().lower()
|
||||
gts[file].append(get_key_words(line))
|
||||
print(gts['ChineseDrawText'][:10])
|
||||
|
||||
|
||||
def get_p_r_acc(method, pred, gt):
|
||||
|
||||
pred = [p.strip().lower() for p in pred]
|
||||
gt = [g.strip().lower() for g in gt]
|
||||
|
||||
pred_orig = copy.deepcopy(pred)
|
||||
gt_orig = copy.deepcopy(gt)
|
||||
|
||||
pred_length = len(pred)
|
||||
gt_length = len(gt)
|
||||
|
||||
for p in pred:
|
||||
if p in gt_orig:
|
||||
pred_orig.remove(p)
|
||||
gt_orig.remove(p)
|
||||
|
||||
p = (pred_length - len(pred_orig)) / (pred_length + 1e-8)
|
||||
r = (gt_length - len(gt_orig)) / (gt_length + 1e-8)
|
||||
|
||||
pred_sorted = sorted(pred)
|
||||
gt_sorted = sorted(gt)
|
||||
if ''.join(pred_sorted) == ''.join(gt_sorted):
|
||||
acc = 1
|
||||
else:
|
||||
acc = 0
|
||||
|
||||
return p, r, acc
|
||||
|
||||
|
||||
files = os.listdir('/path/to/MaskTextSpotterV3/tools/ocr_result')
|
||||
print(len(files))
|
||||
|
||||
for file in files:
|
||||
method, dataset, prompt_index, image_index = file.strip().split('_')
|
||||
ocrs = open(os.path.join('/path/to/MaskTextSpotterV3/tools/ocr_result', file)).readlines()
|
||||
p, r, acc = get_p_r_acc(method, ocrs, gts[dataset][int(prompt_index)])
|
||||
results[method]['cnt'] += 1
|
||||
results[method]['p'] += p
|
||||
results[method]['r'] += r
|
||||
results[method]['acc'] += acc
|
||||
|
||||
for method in results.keys():
|
||||
results[method]['p'] /= results[method]['cnt']
|
||||
results[method]['r'] /= results[method]['cnt']
|
||||
results[method]['f'] = 2 * results[method]['p'] * results[method]['r'] / (results[method]['p'] + results[method]['r'] + 1e-8)
|
||||
results[method]['acc'] /= results[method]['cnt']
|
||||
|
||||
print(results)
|
||||
|
||||
@@ -0,0 +1,6 @@
|
||||
diffusers==0.16.0
|
||||
torch==2.0.1
|
||||
pycocoevalcap
|
||||
pytorch_fid
|
||||
sentencepiece
|
||||
-e git+https://github.com/openai/CLIP.git@main#egg=clip
|
||||
@@ -0,0 +1,517 @@
|
||||
# ------------------------------------------
|
||||
# TextDiffuser: Diffusion Models as Text Painters
|
||||
# Paper Link: https://arxiv.org/abs/2305.10855
|
||||
# Code Link: https://github.com/microsoft/unilm/tree/master/textdiffuser
|
||||
# Copyright (c) Microsoft Corporation.
|
||||
# This file provides the inference script.
|
||||
# ------------------------------------------
|
||||
|
||||
import os
|
||||
import cv2
|
||||
import random
|
||||
import logging
|
||||
import argparse
|
||||
import numpy as np
|
||||
|
||||
from pathlib import Path
|
||||
from tqdm.auto import tqdm
|
||||
from typing import Optional
|
||||
from packaging import version
|
||||
from termcolor import colored
|
||||
from PIL import Image, ImageDraw, ImageFont, ImageOps, ImageEnhance # import for visualization
|
||||
from huggingface_hub import HfFolder, Repository, create_repo, whoami
|
||||
|
||||
import datasets
|
||||
from datasets import load_dataset
|
||||
from datasets import disable_caching
|
||||
|
||||
import torch
|
||||
import torch.utils.checkpoint
|
||||
import torch.nn.functional as F
|
||||
from torchvision import transforms
|
||||
|
||||
import accelerate
|
||||
from accelerate import Accelerator
|
||||
from accelerate.logging import get_logger
|
||||
from accelerate.utils import ProjectConfiguration, set_seed
|
||||
|
||||
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
|
||||
|
||||
import transformers
|
||||
from transformers import CLIPTextModel, CLIPTokenizer
|
||||
|
||||
from util import segmentation_mask_visualization, make_caption_pil, combine_image, transform_mask, filter_segmentation_mask, inpainting_merge_image
|
||||
from model.layout_generator import get_layout_from_prompt
|
||||
from model.text_segmenter.unet import UNet
|
||||
|
||||
import torchsnooper
|
||||
|
||||
disable_caching()
|
||||
check_min_version("0.15.0.dev0")
|
||||
logger = get_logger(__name__, log_level="INFO")
|
||||
|
||||
|
||||
def parse_args():
|
||||
parser = argparse.ArgumentParser(description="Simple example of a training script.")
|
||||
parser.add_argument(
|
||||
"--pretrained_model_name_or_path",
|
||||
type=str,
|
||||
default='runwayml/stable-diffusion-v1-5', # no need to modify this
|
||||
help="Path to pretrained model or model identifier from huggingface.co/models. Please do not modify this.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--revision",
|
||||
type=str,
|
||||
default=None,
|
||||
required=False,
|
||||
help="Revision of pretrained model identifier from huggingface.co/models.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--mode",
|
||||
type=str,
|
||||
default=None,
|
||||
required=True,
|
||||
choices=["text-to-image", "text-to-image-with-template", "text-inpainting"],
|
||||
help="Three modes can be used.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--prompt",
|
||||
type=str,
|
||||
default="",
|
||||
required=False,
|
||||
help="The text prompts provided by users.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--prompt_list",
|
||||
type=str,
|
||||
default="",
|
||||
required=True,
|
||||
help="The list of prompts.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--template_image",
|
||||
type=str,
|
||||
default="",
|
||||
help="The template image should be given when using 【text-to-image-with-template】 mode.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--original_image",
|
||||
type=str,
|
||||
default="",
|
||||
help="The original image should be given when using 【text-inpainting】 mode.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--text_mask",
|
||||
type=str,
|
||||
default="",
|
||||
help="The text mask should be given when using 【text-inpainting】 mode.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--output_dir",
|
||||
type=str,
|
||||
default="output",
|
||||
help="The path of the generation directory.",
|
||||
)
|
||||
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=0, # set to 0 during evaluation
|
||||
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(
|
||||
"--classifier_free_scale",
|
||||
type=float,
|
||||
default=7.5, # following stable diffusion (https://github.com/CompVis/stable-diffusion)
|
||||
help="Classifier free scale following https://arxiv.org/abs/2207.12598.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--drop_caption",
|
||||
action="store_true",
|
||||
help="Whether to drop captions during training following https://arxiv.org/abs/2207.12598.."
|
||||
)
|
||||
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(
|
||||
"--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='fp16',
|
||||
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=500,
|
||||
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=5,
|
||||
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, # should be specified during inference
|
||||
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(
|
||||
"--font_path",
|
||||
type=str,
|
||||
default='Arial.ttf',
|
||||
help="The path of font for visualization."
|
||||
)
|
||||
parser.add_argument(
|
||||
"--sample_steps",
|
||||
type=int,
|
||||
default=50, # following stable diffusion (https://github.com/CompVis/stable-diffusion)
|
||||
help="Diffusion steps for sampling."
|
||||
)
|
||||
parser.add_argument(
|
||||
"--vis_num",
|
||||
type=int,
|
||||
default=9, # please decreases the number if out-of-memory error occurs
|
||||
help="Number of images to be sample. Please decrease it when encountering out of memory error."
|
||||
)
|
||||
parser.add_argument(
|
||||
"--binarization",
|
||||
action="store_true",
|
||||
help="Whether to binarize the template image."
|
||||
)
|
||||
parser.add_argument(
|
||||
"--use_pillow_segmentation_mask",
|
||||
type=bool,
|
||||
default=True,
|
||||
help="In the 【text-to-image】 mode, please specify whether to use the segmentation masks provided by PILLOW"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--character_segmenter_path",
|
||||
type=str,
|
||||
default='textdiffuser-ckpt/text_segmenter.pth',
|
||||
help="checkpoint of character-level segmenter"
|
||||
)
|
||||
args = parser.parse_args()
|
||||
|
||||
print(f'{colored("[√]", "green")} Arguments are loaded.')
|
||||
print(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
|
||||
|
||||
return args
|
||||
|
||||
|
||||
|
||||
def get_full_repo_name(model_id: str, organization: Optional[str] = None, token: Optional[str] = None):
|
||||
if token is None:
|
||||
token = HfFolder.get_token()
|
||||
if organization is None:
|
||||
username = whoami(token)["name"]
|
||||
return f"{username}/{model_id}"
|
||||
else:
|
||||
return f"{organization}/{model_id}"
|
||||
|
||||
|
||||
# @torchsnooper.snoop()
|
||||
def main():
|
||||
args = parse_args()
|
||||
# If passed along, set the training seed now.
|
||||
seed = args.seed if args.seed is not None else random.randint(0, 1000000)
|
||||
set_seed(seed)
|
||||
print(f'{colored("[√]", "green")} Seed is set to {seed}.')
|
||||
|
||||
logging_dir = os.path.join(args.output_dir, args.logging_dir)
|
||||
# sub_output_dir = f"{args.prompt}_[{args.mode.upper()}]_[SEED-{seed}]"
|
||||
|
||||
print(f'{colored("[√]", "green")} Logging dir is set to {logging_dir}.')
|
||||
|
||||
accelerator_project_config = ProjectConfiguration(total_limit=args.checkpoints_total_limit)
|
||||
|
||||
accelerator = Accelerator(
|
||||
gradient_accumulation_steps=1,
|
||||
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(accelerator.state, main_process_only=False)
|
||||
if accelerator.is_local_main_process:
|
||||
datasets.utils.logging.set_verbosity_warning()
|
||||
transformers.utils.logging.set_verbosity_warning()
|
||||
diffusers.utils.logging.set_verbosity_info()
|
||||
else:
|
||||
datasets.utils.logging.set_verbosity_error()
|
||||
transformers.utils.logging.set_verbosity_error()
|
||||
diffusers.utils.logging.set_verbosity_error()
|
||||
|
||||
# Handle the repository creation
|
||||
if accelerator.is_main_process:
|
||||
if args.push_to_hub:
|
||||
if args.hub_model_id is None:
|
||||
repo_name = get_full_repo_name(Path(args.output_dir).name, token=args.hub_token)
|
||||
else:
|
||||
repo_name = args.hub_model_id
|
||||
create_repo(repo_name, exist_ok=True, token=args.hub_token)
|
||||
repo = Repository(args.output_dir, clone_from=repo_name, token=args.hub_token)
|
||||
|
||||
with open(os.path.join(args.output_dir, ".gitignore"), "w+") as gitignore:
|
||||
if "step_*" not in gitignore:
|
||||
gitignore.write("step_*\n")
|
||||
if "epoch_*" not in gitignore:
|
||||
gitignore.write("epoch_*\n")
|
||||
elif args.output_dir is not None:
|
||||
os.makedirs(args.output_dir, exist_ok=True)
|
||||
print(args.output_dir)
|
||||
|
||||
# Load scheduler, tokenizer and models.
|
||||
tokenizer = CLIPTokenizer.from_pretrained(
|
||||
args.pretrained_model_name_or_path, subfolder="tokenizer", revision=args.revision
|
||||
)
|
||||
text_encoder = CLIPTextModel.from_pretrained(
|
||||
args.pretrained_model_name_or_path, subfolder="text_encoder", revision=args.revision
|
||||
)
|
||||
vae = AutoencoderKL.from_pretrained(args.pretrained_model_name_or_path, subfolder="vae", revision=args.revision).cuda()
|
||||
unet = UNet2DConditionModel.from_pretrained(
|
||||
args.resume_from_checkpoint, subfolder="unet", revision=None
|
||||
).cuda()
|
||||
|
||||
# Freeze vae and text_encoder
|
||||
vae.requires_grad_(False)
|
||||
text_encoder.requires_grad_(False)
|
||||
|
||||
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."
|
||||
)
|
||||
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 `accelerator.save_state(...)` serializes in a nice format
|
||||
def save_model_hook(models, weights, output_dir):
|
||||
|
||||
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):
|
||||
|
||||
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
|
||||
|
||||
accelerator.register_save_state_pre_hook(save_model_hook)
|
||||
accelerator.register_load_state_pre_hook(load_model_hook)
|
||||
|
||||
|
||||
with open(args.prompt_list) as fr:
|
||||
prompts = fr.readlines()
|
||||
prompts = [_.strip() for _ in prompts]
|
||||
|
||||
for idx in range(args.vis_num):
|
||||
os.makedirs(os.path.join(args.output_dir, 'textdiffuser', 'images_' + str(idx)), exist_ok=True)
|
||||
|
||||
for prompt_index, prompt in enumerate(prompts):
|
||||
|
||||
args.prompt = prompt
|
||||
|
||||
# setup schedulers
|
||||
scheduler = DDPMScheduler.from_pretrained(args.pretrained_model_name_or_path, subfolder="scheduler")
|
||||
scheduler.set_timesteps(args.sample_steps)
|
||||
sample_num = args.vis_num
|
||||
noise = torch.randn((sample_num, 4, 64, 64)).to("cuda") # (b, 4, 64, 64)
|
||||
input = noise # (b, 4, 64, 64)
|
||||
|
||||
captions = [args.prompt] * sample_num
|
||||
captions_nocond = [""] * sample_num
|
||||
print(f'{colored("[√]", "green")} Prompt is loaded: {args.prompt}.')
|
||||
|
||||
# encode text prompts
|
||||
inputs = tokenizer(
|
||||
captions, max_length=tokenizer.model_max_length, padding="max_length", truncation=True, return_tensors="pt"
|
||||
).input_ids # (b, 77)
|
||||
encoder_hidden_states = text_encoder(inputs)[0].cuda() # (b, 77, 768)
|
||||
print(f'{colored("[√]", "green")} encoder_hidden_states: {encoder_hidden_states.shape}.')
|
||||
|
||||
inputs_nocond = tokenizer(
|
||||
captions_nocond, max_length=tokenizer.model_max_length, padding="max_length", truncation=True, return_tensors="pt"
|
||||
).input_ids # (b, 77)
|
||||
encoder_hidden_states_nocond = text_encoder(inputs_nocond)[0].cuda() # (b, 77, 768)
|
||||
print(f'{colored("[√]", "green")} encoder_hidden_states_nocond: {encoder_hidden_states_nocond.shape}.')
|
||||
|
||||
# load character-level segmenter
|
||||
segmenter = UNet(3, 96, True).cuda()
|
||||
segmenter = torch.nn.DataParallel(segmenter)
|
||||
segmenter.load_state_dict(torch.load(args.character_segmenter_path))
|
||||
segmenter.eval()
|
||||
print(f'{colored("[√]", "green")} Text segmenter is successfully loaded.')
|
||||
|
||||
#### text-to-image ####
|
||||
if args.mode == 'text-to-image':
|
||||
render_image, segmentation_mask_from_pillow = get_layout_from_prompt(args)
|
||||
|
||||
if args.use_pillow_segmentation_mask:
|
||||
segmentation_mask = torch.Tensor(np.array(segmentation_mask_from_pillow)).cuda() # (512, 512)
|
||||
else:
|
||||
to_tensor = transforms.ToTensor()
|
||||
image_tensor = to_tensor(render_image).unsqueeze(0).cuda().sub_(0.5).div_(0.5)
|
||||
with torch.no_grad():
|
||||
segmentation_mask = segmenter(image_tensor)
|
||||
segmentation_mask = segmentation_mask.max(1)[1].squeeze(0)
|
||||
|
||||
segmentation_mask = filter_segmentation_mask(segmentation_mask)
|
||||
segmentation_mask = torch.nn.functional.interpolate(segmentation_mask.unsqueeze(0).unsqueeze(0).float(), size=(256, 256), mode='nearest')
|
||||
segmentation_mask = segmentation_mask.squeeze(1).repeat(sample_num, 1, 1).long().to('cuda') # (1, 1, 256, 256)
|
||||
print(f'{colored("[√]", "green")} character-level segmentation_mask: {segmentation_mask.shape}.')
|
||||
|
||||
feature_mask = torch.ones(sample_num, 1, 64, 64).to('cuda') # (b, 1, 64, 64)
|
||||
masked_image = torch.zeros(sample_num, 3, 512, 512).to('cuda') # (b, 3, 512, 512)
|
||||
masked_feature = vae.encode(masked_image).latent_dist.sample() # (b, 4, 64, 64)
|
||||
masked_feature = masked_feature * vae.config.scaling_factor
|
||||
print(f'{colored("[√]", "green")} feature_mask: {feature_mask.shape}.')
|
||||
print(f'{colored("[√]", "green")} masked_feature: {masked_feature.shape}.')
|
||||
|
||||
|
||||
# diffusion process
|
||||
intermediate_images = []
|
||||
for t in tqdm(scheduler.timesteps):
|
||||
with torch.no_grad():
|
||||
noise_pred_cond = unet(sample=input, timestep=t, encoder_hidden_states=encoder_hidden_states, segmentation_mask=segmentation_mask, feature_mask=feature_mask, masked_feature=masked_feature).sample # b, 4, 64, 64
|
||||
noise_pred_uncond = unet(sample=input, timestep=t, encoder_hidden_states=encoder_hidden_states_nocond, segmentation_mask=segmentation_mask, feature_mask=feature_mask, masked_feature=masked_feature).sample # b, 4, 64, 64
|
||||
noisy_residual = noise_pred_uncond + args.classifier_free_scale * (noise_pred_cond - noise_pred_uncond) # b, 4, 64, 64
|
||||
prev_noisy_sample = scheduler.step(noisy_residual, t, input).prev_sample
|
||||
input = prev_noisy_sample
|
||||
intermediate_images.append(prev_noisy_sample)
|
||||
|
||||
# decode and visualization
|
||||
input = 1 / vae.config.scaling_factor * input
|
||||
sample_images = vae.decode(input.float(), return_dict=False)[0] # (b, 3, 512, 512)
|
||||
|
||||
image_pil = render_image.resize((512,512))
|
||||
segmentation_mask = segmentation_mask[0].squeeze().cpu().numpy()
|
||||
character_mask_pil = Image.fromarray(((segmentation_mask!=0)*255).astype('uint8')).resize((512,512))
|
||||
character_mask_highlight_pil = segmentation_mask_visualization(args.font_path,segmentation_mask)
|
||||
caption_pil = make_caption_pil(args.font_path, captions)
|
||||
|
||||
# save pred_img
|
||||
pred_image_list = []
|
||||
for image in sample_images.float():
|
||||
image = (image / 2 + 0.5).clamp(0, 1).unsqueeze(0)
|
||||
image = image.cpu().permute(0, 2, 3, 1).numpy()[0]
|
||||
image = Image.fromarray((image * 255).round().astype("uint8")).convert('RGB')
|
||||
pred_image_list.append(image)
|
||||
|
||||
for image_index, image in enumerate(pred_image_list):
|
||||
image.save(f'{args.output_dir}/textdiffuser/images_{image_index}/{prompt_index}_{image_index}.jpg')
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,932 @@
|
||||
# ------------------------------------------
|
||||
# TextDiffuser: Diffusion Models as Text Painters
|
||||
# Paper Link: https://arxiv.org/abs/2305.10855
|
||||
# Code Link: https://github.com/microsoft/unilm/tree/master/textdiffuser
|
||||
# Copyright (c) Microsoft Corporation.
|
||||
# This file provides the inference script.
|
||||
# ------------------------------------------
|
||||
|
||||
import os
|
||||
import re
|
||||
import zipfile
|
||||
|
||||
if not os.path.exists('textdiffuser-ckpt'):
|
||||
os.system('wget https://huggingface.co/datasets/JingyeChen22/TextDiffuser/resolve/main/textdiffuser-ckpt-new.zip')
|
||||
with zipfile.ZipFile('textdiffuser-ckpt-new.zip', 'r') as zip_ref:
|
||||
zip_ref.extractall('.')
|
||||
|
||||
if not os.path.exists('images'):
|
||||
os.system('wget https://huggingface.co/datasets/JingyeChen22/TextDiffuser/resolve/main/images.zip')
|
||||
with zipfile.ZipFile('images.zip', 'r') as zip_ref:
|
||||
zip_ref.extractall('.')
|
||||
|
||||
import cv2
|
||||
import random
|
||||
import logging
|
||||
import argparse
|
||||
import numpy as np
|
||||
|
||||
from pathlib import Path
|
||||
from tqdm.auto import tqdm
|
||||
from typing import Optional
|
||||
from packaging import version
|
||||
from termcolor import colored
|
||||
from PIL import Image, ImageDraw, ImageFont, ImageOps, ImageEnhance # import for visualization
|
||||
from huggingface_hub import HfFolder, Repository, create_repo, whoami
|
||||
|
||||
import datasets
|
||||
from datasets import load_dataset
|
||||
from datasets import disable_caching
|
||||
|
||||
import torch
|
||||
import torch.utils.checkpoint
|
||||
import torch.nn.functional as F
|
||||
|
||||
import accelerate
|
||||
from accelerate import Accelerator
|
||||
from accelerate.logging import get_logger
|
||||
from accelerate.utils import ProjectConfiguration, set_seed
|
||||
|
||||
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
|
||||
|
||||
import transformers
|
||||
from transformers import CLIPTextModel, CLIPTokenizer
|
||||
|
||||
from util import segmentation_mask_visualization, make_caption_pil, combine_image, combine_image_gradio, transform_mask, transform_mask_pil, filter_segmentation_mask, inpainting_merge_image
|
||||
from model.layout_generator import get_layout_from_prompt
|
||||
from model.text_segmenter.unet import UNet
|
||||
|
||||
|
||||
disable_caching()
|
||||
check_min_version("0.15.0.dev0")
|
||||
logger = get_logger(__name__, log_level="INFO")
|
||||
|
||||
|
||||
def parse_args():
|
||||
parser = argparse.ArgumentParser(description="Simple example of a training script.")
|
||||
parser.add_argument(
|
||||
"--pretrained_model_name_or_path",
|
||||
type=str,
|
||||
default='runwayml/stable-diffusion-v1-5', # no need to modify this
|
||||
help="Path to pretrained model or model identifier from huggingface.co/models. Please do not modify this.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--revision",
|
||||
type=str,
|
||||
default=None,
|
||||
required=False,
|
||||
help="Revision of pretrained model identifier from huggingface.co/models.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--mode",
|
||||
type=str,
|
||||
default="text-to-image",
|
||||
# required=True,
|
||||
choices=["text-to-image", "text-to-image-with-template", "text-inpainting"],
|
||||
help="Three modes can be used.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--prompt",
|
||||
type=str,
|
||||
default="",
|
||||
# required=True,
|
||||
help="The text prompts provided by users.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--template_image",
|
||||
type=str,
|
||||
default="",
|
||||
help="The template image should be given when using 【text-to-image-with-template】 mode.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--original_image",
|
||||
type=str,
|
||||
default="",
|
||||
help="The original image should be given when using 【text-inpainting】 mode.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--text_mask",
|
||||
type=str,
|
||||
default="",
|
||||
help="The text mask should be given when using 【text-inpainting】 mode.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--output_dir",
|
||||
type=str,
|
||||
default="output",
|
||||
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(
|
||||
"--classifier_free_scale",
|
||||
type=float,
|
||||
default=7.5, # following stable diffusion (https://github.com/CompVis/stable-diffusion)
|
||||
help="Classifier free scale following https://arxiv.org/abs/2207.12598.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--drop_caption",
|
||||
action="store_true",
|
||||
help="Whether to drop captions during training following https://arxiv.org/abs/2207.12598.."
|
||||
)
|
||||
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(
|
||||
"--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='fp16',
|
||||
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=500,
|
||||
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=5,
|
||||
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='textdiffuser-ckpt/diffusion_backbone', # should be specified during inference
|
||||
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(
|
||||
"--font_path",
|
||||
type=str,
|
||||
default='Arial.ttf',
|
||||
help="The path of font for visualization."
|
||||
)
|
||||
parser.add_argument(
|
||||
"--sample_steps",
|
||||
type=int,
|
||||
default=50, # following stable diffusion (https://github.com/CompVis/stable-diffusion)
|
||||
help="Diffusion steps for sampling."
|
||||
)
|
||||
parser.add_argument(
|
||||
"--vis_num",
|
||||
type=int,
|
||||
default=4, # please decreases the number if out-of-memory error occurs
|
||||
help="Number of images to be sample. Please decrease it when encountering out of memory error."
|
||||
)
|
||||
parser.add_argument(
|
||||
"--binarization",
|
||||
action="store_true",
|
||||
help="Whether to binarize the template image."
|
||||
)
|
||||
parser.add_argument(
|
||||
"--use_pillow_segmentation_mask",
|
||||
type=bool,
|
||||
default=True,
|
||||
help="In the 【text-to-image】 mode, please specify whether to use the segmentation masks provided by PILLOW"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--character_segmenter_path",
|
||||
type=str,
|
||||
default='textdiffuser-ckpt/text_segmenter.pth',
|
||||
help="checkpoint of character-level segmenter"
|
||||
)
|
||||
args = parser.parse_args()
|
||||
|
||||
print(f'{colored("[√]", "green")} Arguments are loaded.')
|
||||
print(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
|
||||
|
||||
return args
|
||||
|
||||
|
||||
|
||||
def get_full_repo_name(model_id: str, organization: Optional[str] = None, token: Optional[str] = None):
|
||||
if token is None:
|
||||
token = HfFolder.get_token()
|
||||
if organization is None:
|
||||
username = whoami(token)["name"]
|
||||
return f"{username}/{model_id}"
|
||||
else:
|
||||
return f"{organization}/{model_id}"
|
||||
|
||||
|
||||
|
||||
args = parse_args()
|
||||
logging_dir = os.path.join(args.output_dir, args.logging_dir)
|
||||
|
||||
print(f'{colored("[√]", "green")} Logging dir is set to {logging_dir}.')
|
||||
|
||||
accelerator_project_config = ProjectConfiguration(total_limit=args.checkpoints_total_limit)
|
||||
|
||||
accelerator = Accelerator(
|
||||
gradient_accumulation_steps=1,
|
||||
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(accelerator.state, main_process_only=False)
|
||||
if accelerator.is_local_main_process:
|
||||
datasets.utils.logging.set_verbosity_warning()
|
||||
transformers.utils.logging.set_verbosity_warning()
|
||||
diffusers.utils.logging.set_verbosity_info()
|
||||
else:
|
||||
datasets.utils.logging.set_verbosity_error()
|
||||
transformers.utils.logging.set_verbosity_error()
|
||||
diffusers.utils.logging.set_verbosity_error()
|
||||
|
||||
# Handle the repository creation
|
||||
if accelerator.is_main_process:
|
||||
if args.push_to_hub:
|
||||
if args.hub_model_id is None:
|
||||
repo_name = get_full_repo_name(Path(args.output_dir).name, token=args.hub_token)
|
||||
else:
|
||||
repo_name = args.hub_model_id
|
||||
create_repo(repo_name, exist_ok=True, token=args.hub_token)
|
||||
repo = Repository(args.output_dir, clone_from=repo_name, token=args.hub_token)
|
||||
|
||||
with open(os.path.join(args.output_dir, ".gitignore"), "w+") as gitignore:
|
||||
if "step_*" not in gitignore:
|
||||
gitignore.write("step_*\n")
|
||||
if "epoch_*" not in gitignore:
|
||||
gitignore.write("epoch_*\n")
|
||||
elif args.output_dir is not None:
|
||||
os.makedirs(args.output_dir, exist_ok=True)
|
||||
print(args.output_dir)
|
||||
|
||||
# Load scheduler, tokenizer and models.
|
||||
tokenizer = CLIPTokenizer.from_pretrained(
|
||||
args.pretrained_model_name_or_path, subfolder="tokenizer", revision=args.revision
|
||||
)
|
||||
text_encoder = CLIPTextModel.from_pretrained(
|
||||
args.pretrained_model_name_or_path, subfolder="text_encoder", revision=args.revision
|
||||
)
|
||||
vae = AutoencoderKL.from_pretrained(args.pretrained_model_name_or_path, subfolder="vae", revision=args.revision).cuda()
|
||||
unet = UNet2DConditionModel.from_pretrained(
|
||||
args.resume_from_checkpoint, subfolder="unet", revision=None
|
||||
).cuda()
|
||||
|
||||
# Freeze vae and text_encoder
|
||||
vae.requires_grad_(False)
|
||||
text_encoder.requires_grad_(False)
|
||||
|
||||
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."
|
||||
)
|
||||
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 `accelerator.save_state(...)` serializes in a nice format
|
||||
def save_model_hook(models, weights, output_dir):
|
||||
|
||||
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):
|
||||
|
||||
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
|
||||
|
||||
accelerator.register_save_state_pre_hook(save_model_hook)
|
||||
accelerator.register_load_state_pre_hook(load_model_hook)
|
||||
|
||||
|
||||
# setup schedulers
|
||||
scheduler = DDPMScheduler.from_pretrained(args.pretrained_model_name_or_path, subfolder="scheduler")
|
||||
# sample_num = args.vis_num
|
||||
|
||||
def to_tensor(image):
|
||||
if isinstance(image, Image.Image):
|
||||
image = np.array(image)
|
||||
elif not isinstance(image, np.ndarray):
|
||||
raise TypeError("Error")
|
||||
|
||||
image = image.astype(np.float32) / 255.0
|
||||
image = np.transpose(image, (2, 0, 1))
|
||||
tensor = torch.from_numpy(image)
|
||||
|
||||
return tensor
|
||||
|
||||
def text_to_image(prompt,slider_step,slider_guidance,slider_batch):
|
||||
|
||||
prompt = prompt.replace('"', "'")
|
||||
prompt = re.sub(r"[^a-zA-Z0-9'\" ]+", "", prompt)
|
||||
|
||||
if slider_step>=100:
|
||||
slider_step = 100
|
||||
|
||||
args.prompt = prompt
|
||||
sample_num = slider_batch
|
||||
seed = random.randint(0, 10000000)
|
||||
set_seed(seed)
|
||||
scheduler.set_timesteps(slider_step)
|
||||
|
||||
noise = torch.randn((sample_num, 4, 64, 64)).to("cuda") # (b, 4, 64, 64)
|
||||
input = noise # (b, 4, 64, 64)
|
||||
|
||||
captions = [args.prompt] * sample_num
|
||||
captions_nocond = [""] * sample_num
|
||||
print(f'{colored("[√]", "green")} Prompt is loaded: {args.prompt}.')
|
||||
|
||||
# encode text prompts
|
||||
inputs = tokenizer(
|
||||
captions, max_length=tokenizer.model_max_length, padding="max_length", truncation=True, return_tensors="pt"
|
||||
).input_ids # (b, 77)
|
||||
encoder_hidden_states = text_encoder(inputs)[0].cuda() # (b, 77, 768)
|
||||
print(f'{colored("[√]", "green")} encoder_hidden_states: {encoder_hidden_states.shape}.')
|
||||
|
||||
inputs_nocond = tokenizer(
|
||||
captions_nocond, max_length=tokenizer.model_max_length, padding="max_length", truncation=True, return_tensors="pt"
|
||||
).input_ids # (b, 77)
|
||||
encoder_hidden_states_nocond = text_encoder(inputs_nocond)[0].cuda() # (b, 77, 768)
|
||||
print(f'{colored("[√]", "green")} encoder_hidden_states_nocond: {encoder_hidden_states_nocond.shape}.')
|
||||
|
||||
# load character-level segmenter
|
||||
segmenter = UNet(3, 96, True).cuda()
|
||||
segmenter = torch.nn.DataParallel(segmenter)
|
||||
segmenter.load_state_dict(torch.load(args.character_segmenter_path))
|
||||
segmenter.eval()
|
||||
print(f'{colored("[√]", "green")} Text segmenter is successfully loaded.')
|
||||
|
||||
#### text-to-image ####
|
||||
render_image, segmentation_mask_from_pillow = get_layout_from_prompt(args)
|
||||
|
||||
segmentation_mask = torch.Tensor(np.array(segmentation_mask_from_pillow)).cuda() # (512, 512)
|
||||
|
||||
segmentation_mask = filter_segmentation_mask(segmentation_mask)
|
||||
segmentation_mask = torch.nn.functional.interpolate(segmentation_mask.unsqueeze(0).unsqueeze(0).float(), size=(256, 256), mode='nearest')
|
||||
segmentation_mask = segmentation_mask.squeeze(1).repeat(sample_num, 1, 1).long().to('cuda') # (1, 1, 256, 256)
|
||||
print(f'{colored("[√]", "green")} character-level segmentation_mask: {segmentation_mask.shape}.')
|
||||
|
||||
feature_mask = torch.ones(sample_num, 1, 64, 64).to('cuda') # (b, 1, 64, 64)
|
||||
masked_image = torch.zeros(sample_num, 3, 512, 512).to('cuda') # (b, 3, 512, 512)
|
||||
masked_feature = vae.encode(masked_image).latent_dist.sample() # (b, 4, 64, 64)
|
||||
masked_feature = masked_feature * vae.config.scaling_factor
|
||||
print(f'{colored("[√]", "green")} feature_mask: {feature_mask.shape}.')
|
||||
print(f'{colored("[√]", "green")} masked_feature: {masked_feature.shape}.')
|
||||
|
||||
# diffusion process
|
||||
intermediate_images = []
|
||||
for t in tqdm(scheduler.timesteps):
|
||||
with torch.no_grad():
|
||||
noise_pred_cond = unet(sample=input, timestep=t, encoder_hidden_states=encoder_hidden_states, segmentation_mask=segmentation_mask, feature_mask=feature_mask, masked_feature=masked_feature).sample # b, 4, 64, 64
|
||||
noise_pred_uncond = unet(sample=input, timestep=t, encoder_hidden_states=encoder_hidden_states_nocond, segmentation_mask=segmentation_mask, feature_mask=feature_mask, masked_feature=masked_feature).sample # b, 4, 64, 64
|
||||
noisy_residual = noise_pred_uncond + slider_guidance * (noise_pred_cond - noise_pred_uncond) # b, 4, 64, 64
|
||||
prev_noisy_sample = scheduler.step(noisy_residual, t, input).prev_sample
|
||||
input = prev_noisy_sample
|
||||
intermediate_images.append(prev_noisy_sample)
|
||||
|
||||
# decode and visualization
|
||||
input = 1 / vae.config.scaling_factor * input
|
||||
sample_images = vae.decode(input.float(), return_dict=False)[0] # (b, 3, 512, 512)
|
||||
|
||||
image_pil = render_image.resize((512,512))
|
||||
segmentation_mask = segmentation_mask[0].squeeze().cpu().numpy()
|
||||
character_mask_pil = Image.fromarray(((segmentation_mask!=0)*255).astype('uint8')).resize((512,512))
|
||||
character_mask_highlight_pil = segmentation_mask_visualization(args.font_path,segmentation_mask)
|
||||
caption_pil = make_caption_pil(args.font_path, captions)
|
||||
|
||||
# save pred_img
|
||||
pred_image_list = []
|
||||
for image in sample_images.float():
|
||||
image = (image / 2 + 0.5).clamp(0, 1).unsqueeze(0)
|
||||
image = image.cpu().permute(0, 2, 3, 1).numpy()[0]
|
||||
image = Image.fromarray((image * 255).round().astype("uint8")).convert('RGB')
|
||||
pred_image_list.append(image)
|
||||
|
||||
blank_pil = combine_image_gradio(args, None, pred_image_list, image_pil, character_mask_pil, character_mask_highlight_pil, caption_pil)
|
||||
|
||||
intermediate_result = Image.new('RGB', (512*3, 512))
|
||||
intermediate_result.paste(image_pil, (0, 0))
|
||||
intermediate_result.paste(character_mask_pil, (512, 0))
|
||||
intermediate_result.paste(character_mask_highlight_pil, (512*2, 0))
|
||||
|
||||
return blank_pil, intermediate_result
|
||||
|
||||
|
||||
# load character-level segmenter
|
||||
segmenter = UNet(3, 96, True).cuda()
|
||||
segmenter = torch.nn.DataParallel(segmenter)
|
||||
segmenter.load_state_dict(torch.load(args.character_segmenter_path))
|
||||
segmenter.eval()
|
||||
print(f'{colored("[√]", "green")} Text segmenter is successfully loaded.')
|
||||
|
||||
|
||||
|
||||
|
||||
def text_to_image_with_template(prompt,template_image,slider_step,slider_guidance,slider_batch, binary):
|
||||
|
||||
if slider_step>=100:
|
||||
slider_step = 100
|
||||
|
||||
orig_template_image = template_image.resize((512,512)).convert('RGB')
|
||||
args.prompt = prompt
|
||||
sample_num = slider_batch
|
||||
# If passed along, set the training seed now.
|
||||
# seed = slider_seed
|
||||
seed = random.randint(0, 10000000)
|
||||
set_seed(seed)
|
||||
scheduler.set_timesteps(slider_step)
|
||||
|
||||
noise = torch.randn((sample_num, 4, 64, 64)).to("cuda") # (b, 4, 64, 64)
|
||||
input = noise # (b, 4, 64, 64)
|
||||
|
||||
captions = [args.prompt] * sample_num
|
||||
captions_nocond = [""] * sample_num
|
||||
print(f'{colored("[√]", "green")} Prompt is loaded: {args.prompt}.')
|
||||
|
||||
# encode text prompts
|
||||
inputs = tokenizer(
|
||||
captions, max_length=tokenizer.model_max_length, padding="max_length", truncation=True, return_tensors="pt"
|
||||
).input_ids # (b, 77)
|
||||
encoder_hidden_states = text_encoder(inputs)[0].cuda() # (b, 77, 768)
|
||||
print(f'{colored("[√]", "green")} encoder_hidden_states: {encoder_hidden_states.shape}.')
|
||||
|
||||
inputs_nocond = tokenizer(
|
||||
captions_nocond, max_length=tokenizer.model_max_length, padding="max_length", truncation=True, return_tensors="pt"
|
||||
).input_ids # (b, 77)
|
||||
encoder_hidden_states_nocond = text_encoder(inputs_nocond)[0].cuda() # (b, 77, 768)
|
||||
print(f'{colored("[√]", "green")} encoder_hidden_states_nocond: {encoder_hidden_states_nocond.shape}.')
|
||||
|
||||
#### text-to-image-with-template ####
|
||||
template_image = template_image.resize((256,256)).convert('RGB')
|
||||
|
||||
# whether binarization is needed
|
||||
print(f'{colored("[Warning]", "red")} args.binarization is set to {binary}. You may need it when using handwritten images as templates.')
|
||||
|
||||
if binary:
|
||||
gray = ImageOps.grayscale(template_image)
|
||||
binary = gray.point(lambda x: 255 if x > 96 else 0, '1')
|
||||
template_image = binary.convert('RGB')
|
||||
|
||||
# to_tensor = transforms.ToTensor()
|
||||
image_tensor = to_tensor(template_image).unsqueeze(0).cuda().sub_(0.5).div_(0.5) # (b, 3, 256, 256)
|
||||
|
||||
with torch.no_grad():
|
||||
segmentation_mask = segmenter(image_tensor) # (b, 96, 256, 256)
|
||||
segmentation_mask = segmentation_mask.max(1)[1].squeeze(0) # (256, 256)
|
||||
segmentation_mask = filter_segmentation_mask(segmentation_mask) # (256, 256)
|
||||
|
||||
segmentation_mask = torch.nn.functional.interpolate(segmentation_mask.unsqueeze(0).unsqueeze(0).float(), size=(256, 256), mode='nearest') # (b, 1, 256, 256)
|
||||
segmentation_mask = segmentation_mask.squeeze(1).repeat(sample_num, 1, 1).long().to('cuda') # (b, 1, 256, 256)
|
||||
print(f'{colored("[√]", "green")} Character-level segmentation_mask: {segmentation_mask.shape}.')
|
||||
|
||||
feature_mask = torch.ones(sample_num, 1, 64, 64).to('cuda') # (b, 1, 64, 64)
|
||||
masked_image = torch.zeros(sample_num, 3, 512, 512).to('cuda') # (b, 3, 512, 512)
|
||||
masked_feature = vae.encode(masked_image).latent_dist.sample() # (b, 4, 64, 64)
|
||||
masked_feature = masked_feature * vae.config.scaling_factor # (b, 4, 64, 64)
|
||||
|
||||
# diffusion process
|
||||
intermediate_images = []
|
||||
for t in tqdm(scheduler.timesteps):
|
||||
with torch.no_grad():
|
||||
noise_pred_cond = unet(sample=input, timestep=t, encoder_hidden_states=encoder_hidden_states, segmentation_mask=segmentation_mask, feature_mask=feature_mask, masked_feature=masked_feature).sample # b, 4, 64, 64
|
||||
noise_pred_uncond = unet(sample=input, timestep=t, encoder_hidden_states=encoder_hidden_states_nocond, segmentation_mask=segmentation_mask, feature_mask=feature_mask, masked_feature=masked_feature).sample # b, 4, 64, 64
|
||||
noisy_residual = noise_pred_uncond + slider_guidance * (noise_pred_cond - noise_pred_uncond) # b, 4, 64, 64
|
||||
prev_noisy_sample = scheduler.step(noisy_residual, t, input).prev_sample
|
||||
input = prev_noisy_sample
|
||||
intermediate_images.append(prev_noisy_sample)
|
||||
|
||||
# decode and visualization
|
||||
input = 1 / vae.config.scaling_factor * input
|
||||
sample_images = vae.decode(input.float(), return_dict=False)[0] # (b, 3, 512, 512)
|
||||
|
||||
image_pil = None
|
||||
segmentation_mask = segmentation_mask[0].squeeze().cpu().numpy()
|
||||
character_mask_pil = Image.fromarray(((segmentation_mask!=0)*255).astype('uint8')).resize((512,512))
|
||||
character_mask_highlight_pil = segmentation_mask_visualization(args.font_path,segmentation_mask)
|
||||
caption_pil = make_caption_pil(args.font_path, captions)
|
||||
|
||||
# save pred_img
|
||||
pred_image_list = []
|
||||
for image in sample_images.float():
|
||||
image = (image / 2 + 0.5).clamp(0, 1).unsqueeze(0)
|
||||
image = image.cpu().permute(0, 2, 3, 1).numpy()[0]
|
||||
image = Image.fromarray((image * 255).round().astype("uint8")).convert('RGB')
|
||||
pred_image_list.append(image)
|
||||
|
||||
blank_pil = combine_image_gradio(args, None, pred_image_list, image_pil, character_mask_pil, character_mask_highlight_pil, caption_pil)
|
||||
|
||||
intermediate_result = Image.new('RGB', (512*3, 512))
|
||||
intermediate_result.paste(orig_template_image, (0, 0))
|
||||
intermediate_result.paste(character_mask_pil, (512, 0))
|
||||
intermediate_result.paste(character_mask_highlight_pil, (512*2, 0))
|
||||
|
||||
return blank_pil, intermediate_result
|
||||
|
||||
|
||||
def text_inpainting(prompt,orig_image,mask_image,slider_step,slider_guidance,slider_batch):
|
||||
|
||||
if slider_step>=100:
|
||||
slider_step = 100
|
||||
|
||||
args.prompt = prompt
|
||||
sample_num = slider_batch
|
||||
# If passed along, set the training seed now.
|
||||
# seed = slider_seed
|
||||
seed = random.randint(0, 10000000)
|
||||
set_seed(seed)
|
||||
scheduler.set_timesteps(slider_step)
|
||||
|
||||
noise = torch.randn((sample_num, 4, 64, 64)).to("cuda") # (b, 4, 64, 64)
|
||||
input = noise # (b, 4, 64, 64)
|
||||
|
||||
captions = [args.prompt] * sample_num
|
||||
captions_nocond = [""] * sample_num
|
||||
print(f'{colored("[√]", "green")} Prompt is loaded: {args.prompt}.')
|
||||
|
||||
# encode text prompts
|
||||
inputs = tokenizer(
|
||||
captions, max_length=tokenizer.model_max_length, padding="max_length", truncation=True, return_tensors="pt"
|
||||
).input_ids # (b, 77)
|
||||
encoder_hidden_states = text_encoder(inputs)[0].cuda() # (b, 77, 768)
|
||||
print(f'{colored("[√]", "green")} encoder_hidden_states: {encoder_hidden_states.shape}.')
|
||||
|
||||
inputs_nocond = tokenizer(
|
||||
captions_nocond, max_length=tokenizer.model_max_length, padding="max_length", truncation=True, return_tensors="pt"
|
||||
).input_ids # (b, 77)
|
||||
encoder_hidden_states_nocond = text_encoder(inputs_nocond)[0].cuda() # (b, 77, 768)
|
||||
print(f'{colored("[√]", "green")} encoder_hidden_states_nocond: {encoder_hidden_states_nocond.shape}.')
|
||||
|
||||
mask_image = cv2.resize(mask_image, (512,512))
|
||||
# mask_image = mask_image.resize((512,512)).convert('RGB')
|
||||
text_mask = np.array(mask_image)
|
||||
threshold = 128
|
||||
_, text_mask = cv2.threshold(text_mask, threshold, 255, cv2.THRESH_BINARY)
|
||||
text_mask = Image.fromarray(text_mask).convert('RGB').resize((256,256))
|
||||
text_mask.save('text_mask.png')
|
||||
text_mask_tensor = to_tensor(text_mask).unsqueeze(0).cuda().sub_(0.5).div_(0.5)
|
||||
with torch.no_grad():
|
||||
segmentation_mask = segmenter(text_mask_tensor)
|
||||
|
||||
segmentation_mask = segmentation_mask.max(1)[1].squeeze(0)
|
||||
segmentation_mask = filter_segmentation_mask(segmentation_mask)
|
||||
segmentation_mask = torch.nn.functional.interpolate(segmentation_mask.unsqueeze(0).unsqueeze(0).float(), size=(256, 256), mode='nearest')
|
||||
|
||||
image_mask = transform_mask_pil(mask_image)
|
||||
image_mask = torch.from_numpy(image_mask).cuda().unsqueeze(0).unsqueeze(0)
|
||||
|
||||
orig_image = orig_image.convert('RGB').resize((512,512))
|
||||
image = orig_image
|
||||
image_tensor = to_tensor(image).unsqueeze(0).cuda().sub_(0.5).div_(0.5)
|
||||
masked_image = image_tensor * (1-image_mask)
|
||||
masked_feature = vae.encode(masked_image).latent_dist.sample().repeat(sample_num, 1, 1, 1)
|
||||
masked_feature = masked_feature * vae.config.scaling_factor
|
||||
|
||||
image_mask = torch.nn.functional.interpolate(image_mask, size=(256, 256), mode='nearest').repeat(sample_num, 1, 1, 1)
|
||||
segmentation_mask = segmentation_mask * image_mask
|
||||
feature_mask = torch.nn.functional.interpolate(image_mask, size=(64, 64), mode='nearest')
|
||||
|
||||
# diffusion process
|
||||
intermediate_images = []
|
||||
for t in tqdm(scheduler.timesteps):
|
||||
with torch.no_grad():
|
||||
noise_pred_cond = unet(sample=input, timestep=t, encoder_hidden_states=encoder_hidden_states, segmentation_mask=segmentation_mask, feature_mask=feature_mask, masked_feature=masked_feature).sample # b, 4, 64, 64
|
||||
noise_pred_uncond = unet(sample=input, timestep=t, encoder_hidden_states=encoder_hidden_states_nocond, segmentation_mask=segmentation_mask, feature_mask=feature_mask, masked_feature=masked_feature).sample # b, 4, 64, 64
|
||||
noisy_residual = noise_pred_uncond + slider_guidance * (noise_pred_cond - noise_pred_uncond) # b, 4, 64, 64
|
||||
prev_noisy_sample = scheduler.step(noisy_residual, t, input).prev_sample
|
||||
input = prev_noisy_sample
|
||||
intermediate_images.append(prev_noisy_sample)
|
||||
|
||||
# decode and visualization
|
||||
input = 1 / vae.config.scaling_factor * input
|
||||
sample_images = vae.decode(input.float(), return_dict=False)[0] # (b, 3, 512, 512)
|
||||
|
||||
image_pil = None
|
||||
segmentation_mask = segmentation_mask[0].squeeze().cpu().numpy()
|
||||
character_mask_pil = Image.fromarray(((segmentation_mask!=0)*255).astype('uint8')).resize((512,512))
|
||||
character_mask_highlight_pil = segmentation_mask_visualization(args.font_path,segmentation_mask)
|
||||
caption_pil = make_caption_pil(args.font_path, captions)
|
||||
|
||||
# save pred_img
|
||||
pred_image_list = []
|
||||
for image in sample_images.float():
|
||||
image = (image / 2 + 0.5).clamp(0, 1).unsqueeze(0)
|
||||
image = image.cpu().permute(0, 2, 3, 1).numpy()[0]
|
||||
image = Image.fromarray((image * 255).round().astype("uint8")).convert('RGB')
|
||||
|
||||
# need to merge
|
||||
|
||||
# image = inpainting_merge_image(orig_image, Image.fromarray(mask_image).convert('L'), image)
|
||||
|
||||
pred_image_list.append(image)
|
||||
|
||||
character_mask_pil.save('character_mask_pil.png')
|
||||
character_mask_highlight_pil.save('character_mask_highlight_pil.png')
|
||||
|
||||
|
||||
blank_pil = combine_image_gradio(args, None, pred_image_list, image_pil, character_mask_pil, character_mask_highlight_pil, caption_pil)
|
||||
|
||||
|
||||
background = orig_image.resize((512, 512))
|
||||
alpha = Image.new('L', background.size, int(255 * 0.2))
|
||||
background.putalpha(alpha)
|
||||
# foreground
|
||||
foreground = Image.fromarray(mask_image).convert('L').resize((512, 512))
|
||||
threshold = 200
|
||||
alpha = foreground.point(lambda x: 0 if x > threshold else 255, '1')
|
||||
foreground.putalpha(alpha)
|
||||
merge_image = Image.alpha_composite(foreground.convert('RGBA'), background.convert('RGBA')).convert('RGB')
|
||||
|
||||
intermediate_result = Image.new('RGB', (512*3, 512))
|
||||
intermediate_result.paste(merge_image, (0, 0))
|
||||
intermediate_result.paste(character_mask_pil, (512, 0))
|
||||
intermediate_result.paste(character_mask_highlight_pil, (512*2, 0))
|
||||
|
||||
return blank_pil, intermediate_result
|
||||
|
||||
import gradio as gr
|
||||
|
||||
with gr.Blocks() as demo:
|
||||
|
||||
gr.HTML(
|
||||
"""
|
||||
<div style="text-align: center; max-width: 1200px; margin: 20px auto;">
|
||||
<h1 style="font-weight: 900; font-size: 3rem; margin: 0rem">
|
||||
TextDiffuser: Diffusion Models as Text Painters
|
||||
</h1>
|
||||
<h3 style="font-weight: 450; font-size: 1rem; margin: 0rem">
|
||||
[<a href="https://arxiv.org/abs/2305.10855" style="color:blue;">arXiv</a>]
|
||||
[<a href="https://github.com/microsoft/unilm/tree/master/textdiffuser" style="color:blue;">Code</a>]
|
||||
[<a href="https://jingyechen.github.io/textdiffuser/" style="color:blue;">ProjectPage</a>]
|
||||
</h3>
|
||||
<h2 style="text-align: left; font-weight: 450; font-size: 1rem; margin-top: 0.5rem; margin-bottom: 0.5rem">
|
||||
We propose <b>TextDiffuser</b>, a flexible and controllable framework to generate images with visually appealing text that is coherent with backgrounds.
|
||||
Main features include: (a) <b><font color="#A52A2A">Text-to-Image</font></b>: The user provides a prompt and encloses the keywords with single quotes (e.g., a text image of ‘hello’). The model first determines the layout of the keywords and then draws the image based on the layout and prompt. (b) <b><font color="#A52A2A">Text-to-Image with Templates</font></b>: The user provides a prompt and a template image containing text, which can be a printed, handwritten, or scene text image. These template images can be used to determine the layout of the characters. (c) <b><font color="#A52A2A">Text Inpainting</font></b>: The user provides an image and specifies the region to be modified along with the desired text content. The model is able to modify the original text or add text to areas without text.
|
||||
</h2>
|
||||
<img src="file/images/huggingface_blank.jpg" alt="textdiffuser">
|
||||
</div>
|
||||
""")
|
||||
|
||||
with gr.Tab("Text-to-Image"):
|
||||
with gr.Row():
|
||||
with gr.Column(scale=1):
|
||||
prompt = gr.Textbox(label="Input your prompt here. Please enclose keywords with 【single quotes】, you may refer to the examples below. The current version only supports input in English characters.", placeholder="Placeholder 'Team' hat")
|
||||
slider_step = gr.Slider(minimum=1, maximum=50, value=20, step=1, label="Sampling step", info="The sampling step for TextDiffuser.")
|
||||
slider_guidance = gr.Slider(minimum=1, maximum=9, value=7.5, step=0.5, label="Scale of classifier-free guidance", info="The scale of classifier-free guidance and is set to 7.5 in default.")
|
||||
slider_batch = gr.Slider(minimum=1, maximum=4, value=4, step=1, label="Batch size", info="The number of images to be sampled.")
|
||||
# slider_seed = gr.Slider(minimum=1, maximum=10000, label="Seed", randomize=True)
|
||||
button = gr.Button("Generate")
|
||||
|
||||
with gr.Column(scale=1):
|
||||
output = gr.Image(label='Generated image')
|
||||
|
||||
with gr.Accordion("Intermediate results", open=False):
|
||||
gr.Markdown("Layout, segmentation mask, and details of segmentation mask from left to right.")
|
||||
intermediate_results = gr.Image(label='')
|
||||
|
||||
gr.Markdown("## Prompt Examples")
|
||||
gr.Examples(
|
||||
[
|
||||
["'Team' hat"],
|
||||
["Thanksgiving 'Fam' Mens T Shirt"],
|
||||
["A storefront with 'Hello World' written on it."],
|
||||
["A poster titled 'Quails of North America', showing different kinds of quails."],
|
||||
["A storefront with 'Deep Learning' written on it."],
|
||||
["An antique bottle labeled 'Energy Tonic'"],
|
||||
["A TV show poster titled 'Tango argentino'"],
|
||||
["A TV show poster with logo 'The Dry' on it"],
|
||||
["Stupid 'History' eBook Tales of Stupidity Strangeness"],
|
||||
["Photos of 'Sampa Hostel'"],
|
||||
["A cover named 'Anything is possible'"],
|
||||
["A large recipe book titled 'Recipes from Peru'."],
|
||||
["New York Skyline with 'Diffusion' written with fireworks on the sky"],
|
||||
["Books with the word 'Science' printed on them"],
|
||||
["A globe with the words 'Planet Earth' written in bold letters with continents in bright colors"],
|
||||
["A logo for the company 'EcoGrow', where the letters look like plants"],
|
||||
],
|
||||
prompt,
|
||||
examples_per_page=100
|
||||
)
|
||||
|
||||
button.click(text_to_image, inputs=[prompt,slider_step,slider_guidance,slider_batch], outputs=[output,intermediate_results])
|
||||
|
||||
with gr.Tab("Text-to-Image-with-Template"):
|
||||
with gr.Row():
|
||||
with gr.Column(scale=1):
|
||||
prompt = gr.Textbox(label='Input your prompt here.')
|
||||
template_image = gr.Image(label='Template image', type="pil")
|
||||
slider_step = gr.Slider(minimum=1, maximum=50, value=20, step=1, label="Sampling step", info="The sampling step for TextDiffuser.")
|
||||
slider_guidance = gr.Slider(minimum=1, maximum=9, value=7.5, step=0.5, label="Scale of classifier-free guidance", info="The scale of classifier-free guidance and is set to 7.5 in default.")
|
||||
slider_batch = gr.Slider(minimum=1, maximum=4, value=4, step=1, label="Batch size", info="The number of images to be sampled.")
|
||||
# binary = gr.Radio(["park", "zoo", "road"], label="Location", info="Where did they go?")
|
||||
binary = gr.Checkbox(label="Binarization", bool=True, info="Whether to binarize the template image? You may need it when using handwritten images as templates.")
|
||||
button = gr.Button("Generate")
|
||||
|
||||
with gr.Column(scale=1):
|
||||
output = gr.Image(label='Generated image')
|
||||
|
||||
with gr.Accordion("Intermediate results", open=False):
|
||||
gr.Markdown("Template image, segmentation mask, and details of segmentation mask from left to right.")
|
||||
intermediate_results = gr.Image(label='')
|
||||
|
||||
gr.Markdown("## Prompt and Template-Image Examples")
|
||||
gr.Examples(
|
||||
[
|
||||
["a hand-drawn blueprint for a time machine with the caption 'Time traveling device'", './images/text-to-image-with-template/5.jpg', False],
|
||||
["a gate of garden", './images/text-to-image-with-template/6.jpg', False],
|
||||
["a book called summer vibe written by diffusion model", './images/text-to-image-with-template/7.jpg', False],
|
||||
["a work company", './images/text-to-image-with-template/8.jpg', False],
|
||||
["a book of AI in next century written by AI robot ", './images/text-to-image-with-template/9.jpg', False],
|
||||
["A board saying having a dog named shark at the beach was a mistake", './images/text-to-image-with-template/1.jpg', False],
|
||||
["an elephant holds a newspaper that is written elephant take over the world", './images/text-to-image-with-template/2.jpg', False],
|
||||
["a mouse with a flashlight saying i am afraid of the dark", './images/text-to-image-with-template/4.jpg', False],
|
||||
["a birthday cake of happy birthday to xyz", './images/text-to-image-with-template/10.jpg', False],
|
||||
["a poster of monkey music festival", './images/text-to-image-with-template/11.jpg', False],
|
||||
["a meme of are you kidding", './images/text-to-image-with-template/12.jpg', False],
|
||||
["a 3d model of a 1980s-style computer with the text my old habit on the screen", './images/text-to-image-with-template/13.jpg', True],
|
||||
["a board of hello world", './images/text-to-image-with-template/15.jpg', True],
|
||||
["a microsoft bag", './images/text-to-image-with-template/16.jpg', True],
|
||||
["a dog holds a paper saying please adopt me", './images/text-to-image-with-template/17.jpg', False],
|
||||
["a hello world banner", './images/text-to-image-with-template/18.jpg', False],
|
||||
["a stop pizza", './images/text-to-image-with-template/19.jpg', False],
|
||||
["a dress with text do not read the next sentence", './images/text-to-image-with-template/20.jpg', False],
|
||||
],
|
||||
[prompt,template_image, binary],
|
||||
examples_per_page=100
|
||||
)
|
||||
|
||||
button.click(text_to_image_with_template, inputs=[prompt,template_image,slider_step,slider_guidance,slider_batch,binary], outputs=[output,intermediate_results])
|
||||
|
||||
with gr.Tab("Text-Inpainting"):
|
||||
with gr.Row():
|
||||
with gr.Column(scale=1):
|
||||
prompt = gr.Textbox(label='Input your prompt here.')
|
||||
with gr.Row():
|
||||
orig_image = gr.Image(label='Original image', type="pil")
|
||||
mask_image = gr.Image(label='Mask image', type="numpy")
|
||||
slider_step = gr.Slider(minimum=1, maximum=50, value=20, step=1, label="Sampling step", info="The sampling step for TextDiffuser.")
|
||||
slider_guidance = gr.Slider(minimum=1, maximum=9, value=7.5, step=0.5, label="Scale of classifier-free guidance", info="The scale of classifier-free guidance and is set to 7.5 in default.")
|
||||
slider_batch = gr.Slider(minimum=1, maximum=4, value=4, step=1, label="Batch size", info="The number of images to be sampled.")
|
||||
button = gr.Button("Generate")
|
||||
with gr.Column(scale=1):
|
||||
output = gr.Image(label='Generated image')
|
||||
with gr.Accordion("Intermediate results", open=False):
|
||||
gr.Markdown("Masked image, segmentation mask, and details of segmentation mask from left to right.")
|
||||
intermediate_results = gr.Image(label='')
|
||||
|
||||
gr.Markdown("## Prompt, Original Image, and Mask Examples")
|
||||
gr.Examples(
|
||||
[
|
||||
["eye on security protection", './images/text-inpainting/1.jpg', './images/text-inpainting/1mask.jpg'],
|
||||
["a logo of poppins", './images/text-inpainting/2.jpg', './images/text-inpainting/2mask.jpg'],
|
||||
["tips for middle space living ", './images/text-inpainting/3.jpg', './images/text-inpainting/3mask.jpg'],
|
||||
["george is a proud big sister", './images/text-inpainting/5.jpg', './images/text-inpainting/5mask.jpg'],
|
||||
["we are the great people", './images/text-inpainting/6.jpg', './images/text-inpainting/6mask.jpg'],
|
||||
["tech house interesting terrace party", './images/text-inpainting/7.jpg', './images/text-inpainting/7mask.jpg'],
|
||||
["2023", './images/text-inpainting/8.jpg', './images/text-inpainting/8mask.jpg'],
|
||||
["wear protective equipment necessary", './images/text-inpainting/9.jpg', './images/text-inpainting/9mask.jpg'],
|
||||
["a good day in the hometown", './images/text-inpainting/10.jpg', './images/text-inpainting/10mask.jpg'],
|
||||
["a boy paints good morning on a board", './images/text-inpainting/11.jpg', './images/text-inpainting/11mask.jpg'],
|
||||
["the word my gift on a basketball", './images/text-inpainting/13.jpg', './images/text-inpainting/13mask.jpg'],
|
||||
["a logo of mono", './images/text-inpainting/14.jpg', './images/text-inpainting/14mask.jpg'],
|
||||
["a board saying assyrian on unflagging fry devastates", './images/text-inpainting/15.jpg', './images/text-inpainting/15mask.jpg'],
|
||||
["a board saying session", './images/text-inpainting/16.jpg', './images/text-inpainting/16mask.jpg'],
|
||||
["rankin dork", './images/text-inpainting/17mask.jpg', './images/text-inpainting/17.jpg'],
|
||||
["a coin of mem", './images/text-inpainting/18mask.jpg', './images/text-inpainting/18.jpg'],
|
||||
["a board without text", './images/text-inpainting/19.jpg', './images/text-inpainting/19mask.jpg'],
|
||||
["a board without text", './images/text-inpainting/20.jpg', './images/text-inpainting/20mask.jpg'],
|
||||
|
||||
],
|
||||
[prompt,orig_image,mask_image],
|
||||
)
|
||||
|
||||
|
||||
button.click(text_inpainting, inputs=[prompt,orig_image,mask_image,slider_step,slider_guidance,slider_batch], outputs=[output, intermediate_results])
|
||||
|
||||
|
||||
|
||||
gr.HTML(
|
||||
"""
|
||||
<div style="text-align: justify; max-width: 1200px; margin: 20px auto;">
|
||||
<h3 style="font-weight: 450; font-size: 0.8rem; margin: 0rem">
|
||||
<b>Version</b>: 1.0
|
||||
</h3>
|
||||
<h3 style="font-weight: 450; font-size: 0.8rem; margin: 0rem">
|
||||
<b>Contact</b>:
|
||||
For help or issues using TextDiffuser, please email Jingye Chen <a href="mailto:qwerty.chen@connect.ust.hk">(qwerty.chen@connect.ust.hk)</a>, Yupan Huang <a href="mailto:huangyp28@mail2.sysu.edu.cn">(huangyp28@mail2.sysu.edu.cn)</a> or submit a GitHub issue. For other communications related to TextDiffuser, please contact Lei Cui <a href="mailto:lecu@microsoft.com">(lecu@microsoft.com)</a> or Furu Wei <a href="mailto:fuwei@microsoft.com">(fuwei@microsoft.com)</a>.
|
||||
</h3>
|
||||
</div>
|
||||
"""
|
||||
)
|
||||
|
||||
demo.launch()
|
||||
@@ -0,0 +1,609 @@
|
||||
# ------------------------------------------
|
||||
# TextDiffuser: Diffusion Models as Text Painters
|
||||
# Paper Link: https://arxiv.org/abs/2305.10855
|
||||
# Code Link: https://github.com/microsoft/unilm/tree/master/textdiffuser
|
||||
# Copyright (c) Microsoft Corporation.
|
||||
# This file provides the inference script.
|
||||
# ------------------------------------------
|
||||
|
||||
import os
|
||||
import cv2
|
||||
import random
|
||||
import logging
|
||||
import argparse
|
||||
import numpy as np
|
||||
|
||||
from pathlib import Path
|
||||
from tqdm.auto import tqdm
|
||||
from typing import Optional
|
||||
from packaging import version
|
||||
from termcolor import colored
|
||||
from PIL import Image, ImageDraw, ImageFont, ImageOps, ImageEnhance # import for visualization
|
||||
from huggingface_hub import HfFolder, Repository, create_repo, whoami
|
||||
|
||||
import datasets
|
||||
from datasets import load_dataset
|
||||
from datasets import disable_caching
|
||||
|
||||
import torch
|
||||
import torch.utils.checkpoint
|
||||
import torch.nn.functional as F
|
||||
from torchvision import transforms
|
||||
|
||||
import accelerate
|
||||
from accelerate import Accelerator
|
||||
from accelerate.logging import get_logger
|
||||
from accelerate.utils import ProjectConfiguration, set_seed
|
||||
|
||||
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
|
||||
|
||||
import transformers
|
||||
from transformers import CLIPTextModel, CLIPTokenizer
|
||||
|
||||
from util import segmentation_mask_visualization, make_caption_pil, combine_image, transform_mask, filter_segmentation_mask, inpainting_merge_image
|
||||
from model.layout_generator import get_layout_from_prompt
|
||||
from model.text_segmenter.unet import UNet
|
||||
|
||||
import torchsnooper
|
||||
|
||||
disable_caching()
|
||||
check_min_version("0.15.0.dev0")
|
||||
logger = get_logger(__name__, log_level="INFO")
|
||||
|
||||
|
||||
def parse_args():
|
||||
parser = argparse.ArgumentParser(description="Simple example of a training script.")
|
||||
parser.add_argument(
|
||||
"--pretrained_model_name_or_path",
|
||||
type=str,
|
||||
default='runwayml/stable-diffusion-v1-5', # no need to modify this
|
||||
help="Path to pretrained model or model identifier from huggingface.co/models. Please do not modify this.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--revision",
|
||||
type=str,
|
||||
default=None,
|
||||
required=False,
|
||||
help="Revision of pretrained model identifier from huggingface.co/models.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--mode",
|
||||
type=str,
|
||||
default=None,
|
||||
required=True,
|
||||
choices=["text-to-image", "text-to-image-with-template", "text-inpainting"],
|
||||
help="Three modes can be used.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--prompt",
|
||||
type=str,
|
||||
default="",
|
||||
required=True,
|
||||
help="The text prompts provided by users.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--template_image",
|
||||
type=str,
|
||||
default="",
|
||||
help="The template image should be given when using 【text-to-image-with-template】 mode.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--original_image",
|
||||
type=str,
|
||||
default="",
|
||||
help="The original image should be given when using 【text-inpainting】 mode.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--text_mask",
|
||||
type=str,
|
||||
default="",
|
||||
help="The text mask should be given when using 【text-inpainting】 mode.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--output_dir",
|
||||
type=str,
|
||||
default="output",
|
||||
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(
|
||||
"--classifier_free_scale",
|
||||
type=float,
|
||||
default=7.5, # following stable diffusion (https://github.com/CompVis/stable-diffusion)
|
||||
help="Classifier free scale following https://arxiv.org/abs/2207.12598.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--drop_caption",
|
||||
action="store_true",
|
||||
help="Whether to drop captions during training following https://arxiv.org/abs/2207.12598.."
|
||||
)
|
||||
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(
|
||||
"--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='fp16',
|
||||
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=500,
|
||||
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=5,
|
||||
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, # should be specified during inference
|
||||
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(
|
||||
"--font_path",
|
||||
type=str,
|
||||
default='assets/font/Arial.ttf',
|
||||
help="The path of font for visualization."
|
||||
)
|
||||
parser.add_argument(
|
||||
"--sample_steps",
|
||||
type=int,
|
||||
default=50, # following stable diffusion (https://github.com/CompVis/stable-diffusion)
|
||||
help="Diffusion steps for sampling."
|
||||
)
|
||||
parser.add_argument(
|
||||
"--vis_num",
|
||||
type=int,
|
||||
default=9, # please decreases the number if out-of-memory error occurs
|
||||
help="Number of images to be sample. Please decrease it when encountering out of memory error."
|
||||
)
|
||||
parser.add_argument(
|
||||
"--binarization",
|
||||
action="store_true",
|
||||
help="Whether to binarize the template image."
|
||||
)
|
||||
parser.add_argument(
|
||||
"--use_pillow_segmentation_mask",
|
||||
type=bool,
|
||||
default=True,
|
||||
help="In the 【text-to-image】 mode, please specify whether to use the segmentation masks provided by PILLOW"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--character_segmenter_path",
|
||||
type=str,
|
||||
default='textdiffuser-ckpt/text_segmenter.pth',
|
||||
help="checkpoint of character-level segmenter"
|
||||
)
|
||||
args = parser.parse_args()
|
||||
|
||||
print(f'{colored("[√]", "green")} Arguments are loaded.')
|
||||
print(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
|
||||
|
||||
return args
|
||||
|
||||
|
||||
|
||||
def get_full_repo_name(model_id: str, organization: Optional[str] = None, token: Optional[str] = None):
|
||||
if token is None:
|
||||
token = HfFolder.get_token()
|
||||
if organization is None:
|
||||
username = whoami(token)["name"]
|
||||
return f"{username}/{model_id}"
|
||||
else:
|
||||
return f"{organization}/{model_id}"
|
||||
|
||||
|
||||
# @torchsnooper.snoop()
|
||||
def main():
|
||||
args = parse_args()
|
||||
# If passed along, set the training seed now.
|
||||
seed = args.seed if args.seed is not None else random.randint(0, 1000000)
|
||||
set_seed(seed)
|
||||
print(f'{colored("[√]", "green")} Seed is set to {seed}.')
|
||||
|
||||
logging_dir = os.path.join(args.output_dir, args.logging_dir)
|
||||
sub_output_dir = f"{args.prompt}_[{args.mode.upper()}]_[SEED-{seed}]"
|
||||
|
||||
print(f'{colored("[√]", "green")} Logging dir is set to {logging_dir}.')
|
||||
|
||||
accelerator_project_config = ProjectConfiguration(total_limit=args.checkpoints_total_limit)
|
||||
|
||||
accelerator = Accelerator(
|
||||
gradient_accumulation_steps=1,
|
||||
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(accelerator.state, main_process_only=False)
|
||||
if accelerator.is_local_main_process:
|
||||
datasets.utils.logging.set_verbosity_warning()
|
||||
transformers.utils.logging.set_verbosity_warning()
|
||||
diffusers.utils.logging.set_verbosity_info()
|
||||
else:
|
||||
datasets.utils.logging.set_verbosity_error()
|
||||
transformers.utils.logging.set_verbosity_error()
|
||||
diffusers.utils.logging.set_verbosity_error()
|
||||
|
||||
# Handle the repository creation
|
||||
if accelerator.is_main_process:
|
||||
if args.push_to_hub:
|
||||
if args.hub_model_id is None:
|
||||
repo_name = get_full_repo_name(Path(args.output_dir).name, token=args.hub_token)
|
||||
else:
|
||||
repo_name = args.hub_model_id
|
||||
create_repo(repo_name, exist_ok=True, token=args.hub_token)
|
||||
repo = Repository(args.output_dir, clone_from=repo_name, token=args.hub_token)
|
||||
|
||||
with open(os.path.join(args.output_dir, ".gitignore"), "w+") as gitignore:
|
||||
if "step_*" not in gitignore:
|
||||
gitignore.write("step_*\n")
|
||||
if "epoch_*" not in gitignore:
|
||||
gitignore.write("epoch_*\n")
|
||||
elif args.output_dir is not None:
|
||||
os.makedirs(args.output_dir, exist_ok=True)
|
||||
print(args.output_dir)
|
||||
|
||||
# Load scheduler, tokenizer and models.
|
||||
tokenizer = CLIPTokenizer.from_pretrained(
|
||||
args.pretrained_model_name_or_path, subfolder="tokenizer", revision=args.revision
|
||||
)
|
||||
text_encoder = CLIPTextModel.from_pretrained(
|
||||
args.pretrained_model_name_or_path, subfolder="text_encoder", revision=args.revision
|
||||
)
|
||||
vae = AutoencoderKL.from_pretrained(args.pretrained_model_name_or_path, subfolder="vae", revision=args.revision).cuda()
|
||||
unet = UNet2DConditionModel.from_pretrained(
|
||||
args.resume_from_checkpoint, subfolder="unet", revision=None
|
||||
).cuda()
|
||||
|
||||
# Freeze vae and text_encoder
|
||||
vae.requires_grad_(False)
|
||||
text_encoder.requires_grad_(False)
|
||||
|
||||
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."
|
||||
)
|
||||
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 `accelerator.save_state(...)` serializes in a nice format
|
||||
def save_model_hook(models, weights, output_dir):
|
||||
|
||||
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):
|
||||
|
||||
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
|
||||
|
||||
accelerator.register_save_state_pre_hook(save_model_hook)
|
||||
accelerator.register_load_state_pre_hook(load_model_hook)
|
||||
|
||||
|
||||
# setup schedulers
|
||||
scheduler = DDPMScheduler.from_pretrained(args.pretrained_model_name_or_path, subfolder="scheduler")
|
||||
scheduler.set_timesteps(args.sample_steps)
|
||||
sample_num = args.vis_num
|
||||
noise = torch.randn((sample_num, 4, 64, 64)).to("cuda") # (b, 4, 64, 64)
|
||||
input = noise # (b, 4, 64, 64)
|
||||
|
||||
captions = [args.prompt] * sample_num
|
||||
captions_nocond = [""] * sample_num
|
||||
print(f'{colored("[√]", "green")} Prompt is loaded: {args.prompt}.')
|
||||
|
||||
# encode text prompts
|
||||
inputs = tokenizer(
|
||||
captions, max_length=tokenizer.model_max_length, padding="max_length", truncation=True, return_tensors="pt"
|
||||
).input_ids # (b, 77)
|
||||
encoder_hidden_states = text_encoder(inputs)[0].cuda() # (b, 77, 768)
|
||||
print(f'{colored("[√]", "green")} encoder_hidden_states: {encoder_hidden_states.shape}.')
|
||||
|
||||
inputs_nocond = tokenizer(
|
||||
captions_nocond, max_length=tokenizer.model_max_length, padding="max_length", truncation=True, return_tensors="pt"
|
||||
).input_ids # (b, 77)
|
||||
encoder_hidden_states_nocond = text_encoder(inputs_nocond)[0].cuda() # (b, 77, 768)
|
||||
print(f'{colored("[√]", "green")} encoder_hidden_states_nocond: {encoder_hidden_states_nocond.shape}.')
|
||||
|
||||
# load character-level segmenter
|
||||
segmenter = UNet(3, 96, True).cuda()
|
||||
segmenter = torch.nn.DataParallel(segmenter)
|
||||
segmenter.load_state_dict(torch.load(args.character_segmenter_path))
|
||||
segmenter.eval()
|
||||
print(f'{colored("[√]", "green")} Text segmenter is successfully loaded.')
|
||||
|
||||
#### text-to-image ####
|
||||
if args.mode == 'text-to-image':
|
||||
render_image, segmentation_mask_from_pillow = get_layout_from_prompt(args)
|
||||
|
||||
if args.use_pillow_segmentation_mask:
|
||||
segmentation_mask = torch.Tensor(np.array(segmentation_mask_from_pillow)).cuda() # (512, 512)
|
||||
else:
|
||||
to_tensor = transforms.ToTensor()
|
||||
image_tensor = to_tensor(render_image).unsqueeze(0).cuda().sub_(0.5).div_(0.5)
|
||||
with torch.no_grad():
|
||||
segmentation_mask = segmenter(image_tensor)
|
||||
segmentation_mask = segmentation_mask.max(1)[1].squeeze(0)
|
||||
|
||||
segmentation_mask = filter_segmentation_mask(segmentation_mask)
|
||||
segmentation_mask = torch.nn.functional.interpolate(segmentation_mask.unsqueeze(0).unsqueeze(0).float(), size=(256, 256), mode='nearest')
|
||||
segmentation_mask = segmentation_mask.squeeze(1).repeat(sample_num, 1, 1).long().to('cuda') # (1, 1, 256, 256)
|
||||
print(f'{colored("[√]", "green")} character-level segmentation_mask: {segmentation_mask.shape}.')
|
||||
|
||||
feature_mask = torch.ones(sample_num, 1, 64, 64).to('cuda') # (b, 1, 64, 64)
|
||||
masked_image = torch.zeros(sample_num, 3, 512, 512).to('cuda') # (b, 3, 512, 512)
|
||||
masked_feature = vae.encode(masked_image).latent_dist.sample() # (b, 4, 64, 64)
|
||||
masked_feature = masked_feature * vae.config.scaling_factor
|
||||
print(f'{colored("[√]", "green")} feature_mask: {feature_mask.shape}.')
|
||||
print(f'{colored("[√]", "green")} masked_feature: {masked_feature.shape}.')
|
||||
|
||||
|
||||
#### text-to-image-with-template ####
|
||||
if args.mode == 'text-to-image-with-template':
|
||||
template_image = Image.open(args.template_image).resize((256,256)).convert('RGB')
|
||||
|
||||
# whether binarization is needed
|
||||
print(f'{colored("[Warning]", "red")} args.binarization is set to {args.binarization}. You may need it when using handwritten images as templates.')
|
||||
if args.binarization:
|
||||
gray = ImageOps.grayscale(template_image)
|
||||
binary = gray.point(lambda x: 255 if x > 96 else 0, '1')
|
||||
template_image = binary.convert('RGB')
|
||||
|
||||
to_tensor = transforms.ToTensor()
|
||||
image_tensor = to_tensor(template_image).unsqueeze(0).cuda().sub_(0.5).div_(0.5) # (b, 3, 256, 256)
|
||||
|
||||
with torch.no_grad():
|
||||
segmentation_mask = segmenter(image_tensor) # (b, 96, 256, 256)
|
||||
segmentation_mask = segmentation_mask.max(1)[1].squeeze(0) # (256, 256)
|
||||
segmentation_mask = filter_segmentation_mask(segmentation_mask) # (256, 256)
|
||||
segmentation_mask_pil = Image.fromarray(segmentation_mask.type(torch.uint8).cpu().numpy()).convert('RGB')
|
||||
|
||||
segmentation_mask = torch.nn.functional.interpolate(segmentation_mask.unsqueeze(0).unsqueeze(0).float(), size=(256, 256), mode='nearest') # (b, 1, 256, 256)
|
||||
segmentation_mask = segmentation_mask.squeeze(1).repeat(sample_num, 1, 1).long().to('cuda') # (b, 1, 256, 256)
|
||||
print(f'{colored("[√]", "green")} Character-level segmentation_mask: {segmentation_mask.shape}.')
|
||||
|
||||
feature_mask = torch.ones(sample_num, 1, 64, 64).to('cuda') # (b, 1, 64, 64)
|
||||
masked_image = torch.zeros(sample_num, 3, 512, 512).to('cuda') # (b, 3, 512, 512)
|
||||
masked_feature = vae.encode(masked_image).latent_dist.sample() # (b, 4, 64, 64)
|
||||
masked_feature = masked_feature * vae.config.scaling_factor # (b, 4, 64, 64)
|
||||
print(f'{colored("[√]", "green")} feature_mask: {feature_mask.shape}.')
|
||||
print(f'{colored("[√]", "green")} masked_feature: {masked_feature.shape}.')
|
||||
|
||||
render_image = template_image # for visualization
|
||||
|
||||
|
||||
#### text-inpainting ####
|
||||
if args.mode == 'text-inpainting':
|
||||
text_mask = cv2.imread(args.text_mask)
|
||||
threshold = 128
|
||||
_, text_mask = cv2.threshold(text_mask, threshold, 255, cv2.THRESH_BINARY)
|
||||
text_mask = Image.fromarray(text_mask).convert('RGB').resize((256,256))
|
||||
text_mask_tensor = transforms.ToTensor()(text_mask).unsqueeze(0).cuda().sub_(0.5).div_(0.5)
|
||||
with torch.no_grad():
|
||||
segmentation_mask = segmenter(text_mask_tensor)
|
||||
|
||||
segmentation_mask = segmentation_mask.max(1)[1].squeeze(0)
|
||||
segmentation_mask = filter_segmentation_mask(segmentation_mask)
|
||||
segmentation_mask = torch.nn.functional.interpolate(segmentation_mask.unsqueeze(0).unsqueeze(0).float(), size=(256, 256), mode='nearest')
|
||||
|
||||
image_mask = transform_mask(args.text_mask)
|
||||
image_mask = torch.from_numpy(image_mask).cuda().unsqueeze(0).unsqueeze(0)
|
||||
|
||||
image = Image.open(args.original_image).convert('RGB').resize((512,512))
|
||||
image_tensor = transforms.ToTensor()(image).unsqueeze(0).cuda().sub_(0.5).div_(0.5)
|
||||
masked_image = image_tensor * (1-image_mask)
|
||||
masked_feature = vae.encode(masked_image).latent_dist.sample().repeat(sample_num, 1, 1, 1)
|
||||
masked_feature = masked_feature * vae.config.scaling_factor
|
||||
|
||||
image_mask = torch.nn.functional.interpolate(image_mask, size=(256, 256), mode='nearest').repeat(sample_num, 1, 1, 1)
|
||||
segmentation_mask = segmentation_mask * image_mask
|
||||
feature_mask = torch.nn.functional.interpolate(image_mask, size=(64, 64), mode='nearest')
|
||||
print(f'{colored("[√]", "green")} feature_mask: {feature_mask.shape}.')
|
||||
print(f'{colored("[√]", "green")} segmentation_mask: {segmentation_mask.shape}.')
|
||||
print(f'{colored("[√]", "green")} masked_feature: {masked_feature.shape}.')
|
||||
|
||||
render_image = Image.open(args.original_image)
|
||||
|
||||
|
||||
|
||||
# diffusion process
|
||||
intermediate_images = []
|
||||
for t in tqdm(scheduler.timesteps):
|
||||
with torch.no_grad():
|
||||
noise_pred_cond = unet(sample=input, timestep=t, encoder_hidden_states=encoder_hidden_states, segmentation_mask=segmentation_mask, feature_mask=feature_mask, masked_feature=masked_feature).sample # b, 4, 64, 64
|
||||
noise_pred_uncond = unet(sample=input, timestep=t, encoder_hidden_states=encoder_hidden_states_nocond, segmentation_mask=segmentation_mask, feature_mask=feature_mask, masked_feature=masked_feature).sample # b, 4, 64, 64
|
||||
noisy_residual = noise_pred_uncond + args.classifier_free_scale * (noise_pred_cond - noise_pred_uncond) # b, 4, 64, 64
|
||||
prev_noisy_sample = scheduler.step(noisy_residual, t, input).prev_sample
|
||||
input = prev_noisy_sample
|
||||
intermediate_images.append(prev_noisy_sample)
|
||||
|
||||
# decode and visualization
|
||||
input = 1 / vae.config.scaling_factor * input
|
||||
sample_images = vae.decode(input.float(), return_dict=False)[0] # (b, 3, 512, 512)
|
||||
|
||||
image_pil = render_image.resize((512,512))
|
||||
segmentation_mask = segmentation_mask[0].squeeze().cpu().numpy()
|
||||
character_mask_pil = Image.fromarray(((segmentation_mask!=0)*255).astype('uint8')).resize((512,512))
|
||||
character_mask_highlight_pil = segmentation_mask_visualization(args.font_path,segmentation_mask)
|
||||
caption_pil = make_caption_pil(args.font_path, captions)
|
||||
|
||||
# save pred_img
|
||||
pred_image_list = []
|
||||
for image in sample_images.float():
|
||||
image = (image / 2 + 0.5).clamp(0, 1).unsqueeze(0)
|
||||
image = image.cpu().permute(0, 2, 3, 1).numpy()[0]
|
||||
image = Image.fromarray((image * 255).round().astype("uint8")).convert('RGB')
|
||||
pred_image_list.append(image)
|
||||
|
||||
os.makedirs(f'{args.output_dir}/{sub_output_dir}', exist_ok=True)
|
||||
|
||||
# save additional info
|
||||
if args.mode == 'text-to-image':
|
||||
image_pil.save(os.path.join(args.output_dir, sub_output_dir, 'render_text_image.png'))
|
||||
enhancer = ImageEnhance.Brightness(segmentation_mask_from_pillow)
|
||||
im_brightness = enhancer.enhance(5)
|
||||
im_brightness.save(os.path.join(args.output_dir, sub_output_dir, 'segmentation_mask_from_pillow.png'))
|
||||
if args.mode == 'text-to-image-with-template':
|
||||
template_image.save(os.path.join(args.output_dir, sub_output_dir, 'template.png'))
|
||||
enhancer = ImageEnhance.Brightness(segmentation_mask_pil)
|
||||
im_brightness = enhancer.enhance(5)
|
||||
im_brightness.save(os.path.join(args.output_dir, sub_output_dir, 'segmentation_mask_from_template.png'))
|
||||
if args.mode == 'text-inpainting':
|
||||
character_mask_highlight_pil = character_mask_pil
|
||||
# background
|
||||
background = Image.open(args.original_image).resize((512, 512))
|
||||
alpha = Image.new('L', background.size, int(255 * 0.2))
|
||||
background.putalpha(alpha)
|
||||
# foreground
|
||||
foreground = Image.open(args.text_mask).convert('L').resize((512, 512))
|
||||
threshold = 200
|
||||
alpha = foreground.point(lambda x: 0 if x > threshold else 255, '1')
|
||||
foreground.putalpha(alpha)
|
||||
character_mask_pil = Image.alpha_composite(foreground.convert('RGBA'), background.convert('RGBA')).convert('RGB')
|
||||
# merge
|
||||
pred_image_list_new = []
|
||||
for pred_image in pred_image_list:
|
||||
pred_image = inpainting_merge_image(Image.open(args.original_image), Image.open(args.text_mask).convert('L'), pred_image)
|
||||
pred_image_list_new.append(pred_image)
|
||||
pred_image_list = pred_image_list_new
|
||||
|
||||
|
||||
combine_image(args, sub_output_dir, pred_image_list, image_pil, character_mask_pil, character_mask_highlight_pil, caption_pil)
|
||||
|
||||
|
||||
# create a soft link
|
||||
if os.path.exists(os.path.join(args.output_dir, 'latest')):
|
||||
os.unlink(os.path.join(args.output_dir, 'latest'))
|
||||
os.symlink(os.path.abspath(os.path.join(args.output_dir, sub_output_dir)), os.path.abspath(os.path.join(args.output_dir, 'latest/')))
|
||||
|
||||
|
||||
color_sub_output_dir = colored(sub_output_dir, 'green')
|
||||
print(f'{colored("[√]", "green")} Save successfully. Please check the output at {color_sub_output_dir} OR the latest folder')
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,223 @@
|
||||
# ------------------------------------------
|
||||
# TextDiffuser: Diffusion Models as Text Painters
|
||||
# Paper Link: https://arxiv.org/abs/2305.10855
|
||||
# Code Link: https://github.com/microsoft/unilm/tree/master/textdiffuser
|
||||
# Copyright (c) Microsoft Corporation.
|
||||
# This file aims to predict the layout of keywords in user prompts.
|
||||
# ------------------------------------------
|
||||
|
||||
import warnings
|
||||
warnings.filterwarnings("ignore", category=DeprecationWarning)
|
||||
|
||||
import re
|
||||
import numpy as np
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from transformers import CLIPTokenizer
|
||||
from PIL import Image, ImageDraw, ImageFont
|
||||
from util import get_width, get_key_words, adjust_overlap_box, shrink_box, adjust_font_size, alphabet_dic
|
||||
from model.layout_transformer import LayoutTransformer, TextConditioner
|
||||
from termcolor import colored
|
||||
|
||||
# import layout transformer
|
||||
model = LayoutTransformer().cuda().eval()
|
||||
model.load_state_dict(torch.load('textdiffuser-ckpt/layout_transformer.pth'))
|
||||
|
||||
# import text encoder and tokenizer
|
||||
text_encoder = TextConditioner().cuda().eval()
|
||||
tokenizer = CLIPTokenizer.from_pretrained('openai/clip-vit-large-patch14')
|
||||
|
||||
|
||||
def process_caption(font_path, caption, keywords):
|
||||
# remove punctuations. please remove this statement if you want to paint punctuations
|
||||
caption = re.sub(u"([^\u0041-\u005a\u0061-\u007a\u0030-\u0039])", " ", caption)
|
||||
|
||||
# tokenize it into ids and get length
|
||||
caption_words = tokenizer([caption], truncation=True, max_length=77, return_length=True, return_overflowing_tokens=False, padding="max_length", return_tensors="pt")
|
||||
caption_words_ids = caption_words['input_ids'] # (1, 77)
|
||||
length = caption_words['length'] # (1, )
|
||||
|
||||
# convert id to words
|
||||
words = tokenizer.convert_ids_to_tokens(caption_words_ids.view(-1).tolist())
|
||||
words = [i.replace('</w>', '') for i in words]
|
||||
words_valid = words[:int(length)]
|
||||
|
||||
# store the box coordinates and state of each token
|
||||
info_array = np.zeros((77,5)) # (77, 5)
|
||||
|
||||
# split the caption into words and convert them into lower case
|
||||
caption_split = caption.split()
|
||||
caption_split = [i.lower() for i in caption_split]
|
||||
|
||||
start_dic = {} # get the start index of each word
|
||||
state_list = [] # 0: start, 1: middle, 2: special token
|
||||
word_match_list = [] # the index of the word in the caption
|
||||
current_caption_index = 0
|
||||
current_match = ''
|
||||
for i in range(length):
|
||||
|
||||
# the first and last token are special tokens
|
||||
if i == 0 or i == length-1:
|
||||
state_list.append(2)
|
||||
word_match_list.append(127)
|
||||
continue
|
||||
|
||||
if current_match == '':
|
||||
state_list.append(0)
|
||||
start_dic[current_caption_index] = i
|
||||
else:
|
||||
state_list.append(1)
|
||||
|
||||
current_match += words_valid[i]
|
||||
word_match_list.append(current_caption_index)
|
||||
if current_match == caption_split[current_caption_index]:
|
||||
current_match = ''
|
||||
current_caption_index += 1
|
||||
|
||||
while len(state_list) < 77:
|
||||
state_list.append(127)
|
||||
while len(word_match_list) < 77:
|
||||
word_match_list.append(127)
|
||||
|
||||
length_list = []
|
||||
width_list =[]
|
||||
for i in range(len(word_match_list)):
|
||||
if word_match_list[i] == 127:
|
||||
length_list.append(0)
|
||||
width_list.append(0)
|
||||
else:
|
||||
length_list.append(len(caption.split()[word_match_list[i]]))
|
||||
width_list.append(get_width(font_path, caption.split()[word_match_list[i]]))
|
||||
|
||||
while len(length_list) < 77:
|
||||
length_list.append(127)
|
||||
width_list.append(0)
|
||||
|
||||
length_list = torch.Tensor(length_list).long() # (77, )
|
||||
width_list = torch.Tensor(width_list).long() # (77, )
|
||||
|
||||
boxes = []
|
||||
duplicate_dict = {} # some words may appear more than once
|
||||
for keyword in keywords:
|
||||
keyword = keyword.lower()
|
||||
if keyword in caption_split:
|
||||
if keyword not in duplicate_dict:
|
||||
duplicate_dict[keyword] = caption_split.index(keyword)
|
||||
index = caption_split.index(keyword)
|
||||
else:
|
||||
if duplicate_dict[keyword]+1 < len(caption_split) and keyword in caption_split[duplicate_dict[keyword]+1:]:
|
||||
index = duplicate_dict[keyword] + caption_split[duplicate_dict[keyword]+1:].index(keyword)
|
||||
duplicate_dict[keyword] = index
|
||||
else:
|
||||
continue
|
||||
|
||||
index = caption_split.index(keyword)
|
||||
index = start_dic[index]
|
||||
info_array[index][0] = 1
|
||||
|
||||
box = [0,0,0,0]
|
||||
boxes.append(list(box))
|
||||
info_array[index][1:] = box
|
||||
|
||||
boxes_length = len(boxes)
|
||||
if boxes_length > 8:
|
||||
boxes = boxes[:8]
|
||||
while len(boxes) < 8:
|
||||
boxes.append([0,0,0,0])
|
||||
|
||||
return caption, length_list, width_list, torch.from_numpy(info_array), words, torch.Tensor(state_list).long(), torch.Tensor(word_match_list).long(), torch.Tensor(boxes), boxes_length
|
||||
|
||||
|
||||
def get_layout_from_prompt(args):
|
||||
|
||||
# prompt = args.prompt
|
||||
font_path = args.font_path
|
||||
keywords = get_key_words(args.prompt)
|
||||
|
||||
print(f'{colored("[!]", "red")} Detected keywords: {keywords} from prompt {args.prompt}')
|
||||
|
||||
text_embedding, mask = text_encoder(args.prompt) # (1, 77 768) / (1, 77)
|
||||
|
||||
# process all relevant info
|
||||
caption, length_list, width_list, target, words, state_list, word_match_list, boxes, boxes_length = process_caption(font_path, args.prompt, keywords)
|
||||
target = target.cuda().unsqueeze(0) # (77, 5)
|
||||
width_list = width_list.cuda().unsqueeze(0) # (77, )
|
||||
length_list = length_list.cuda().unsqueeze(0) # (77, )
|
||||
state_list = state_list.cuda().unsqueeze(0) # (77, )
|
||||
word_match_list = word_match_list.cuda().unsqueeze(0) # (77, )
|
||||
|
||||
padding = torch.zeros(1, 1, 4).cuda()
|
||||
boxes = boxes.unsqueeze(0).cuda()
|
||||
right_shifted_boxes = torch.cat([padding, boxes[:,0:-1,:]],1) # (1, 8, 4)
|
||||
|
||||
# inference
|
||||
return_boxes= []
|
||||
with torch.no_grad():
|
||||
for box_index in range(boxes_length):
|
||||
|
||||
if box_index == 0:
|
||||
encoder_embedding = None
|
||||
|
||||
output, encoder_embedding = model(text_embedding, length_list, width_list, mask, state_list, word_match_list, target, right_shifted_boxes, train=False, encoder_embedding=encoder_embedding)
|
||||
output = torch.clamp(output, min=0, max=1) # (1, 8, 4)
|
||||
|
||||
# add overlap detection
|
||||
output = adjust_overlap_box(output, box_index) # (1, 8, 4)
|
||||
|
||||
right_shifted_boxes[:,box_index+1,:] = output[:,box_index,:]
|
||||
xmin, ymin, xmax, ymax = output[0, box_index, :].tolist()
|
||||
return_boxes.append([xmin, ymin, xmax, ymax])
|
||||
|
||||
|
||||
# print the location of keywords
|
||||
print(f'index\tkeyword\tx_min\ty_min\tx_max\ty_max')
|
||||
for index, keyword in enumerate(keywords):
|
||||
x_min = int(return_boxes[index][0] * 512)
|
||||
y_min = int(return_boxes[index][1] * 512)
|
||||
x_max = int(return_boxes[index][2] * 512)
|
||||
y_max = int(return_boxes[index][3] * 512)
|
||||
print(f'{index}\t{keyword}\t{x_min}\t{y_min}\t{x_max}\t{y_max}')
|
||||
|
||||
|
||||
# paint the layout
|
||||
render_image = Image.new('RGB', (512, 512), (255, 255, 255))
|
||||
draw = ImageDraw.Draw(render_image)
|
||||
segmentation_mask = Image.new("L", (512,512), 0)
|
||||
segmentation_mask_draw = ImageDraw.Draw(segmentation_mask)
|
||||
|
||||
for index, box in enumerate(return_boxes):
|
||||
box = [int(i*512) for i in box]
|
||||
xmin, ymin, xmax, ymax = box
|
||||
|
||||
width = xmax - xmin
|
||||
height = ymax - ymin
|
||||
text = keywords[index]
|
||||
|
||||
font_size = adjust_font_size(args, width, height, draw, text)
|
||||
font = ImageFont.truetype(args.font_path, font_size)
|
||||
|
||||
# draw.rectangle([xmin, ymin, xmax,ymax], outline=(255,0,0))
|
||||
draw.text((xmin, ymin), text, font=font, fill=(0, 0, 0))
|
||||
|
||||
boxes = []
|
||||
for i, char in enumerate(text):
|
||||
|
||||
# paint character-level segmentation masks
|
||||
# https://github.com/python-pillow/Pillow/issues/3921
|
||||
bottom_1 = font.getsize(text[i])[1]
|
||||
right, bottom_2 = font.getsize(text[:i+1])
|
||||
bottom = bottom_1 if bottom_1 < bottom_2 else bottom_2
|
||||
width, height = font.getmask(char).size
|
||||
right += xmin
|
||||
bottom += ymin
|
||||
top = bottom - height
|
||||
left = right - width
|
||||
|
||||
char_box = (left, top, right, bottom)
|
||||
boxes.append(char_box)
|
||||
|
||||
char_index = alphabet_dic[char]
|
||||
segmentation_mask_draw.rectangle(shrink_box(char_box, scale_factor = 0.9), fill=char_index)
|
||||
|
||||
print(f'{colored("[√]", "green")} Layout is successfully generated')
|
||||
return render_image, segmentation_mask
|
||||
@@ -0,0 +1,107 @@
|
||||
# ------------------------------------------
|
||||
# TextDiffuser: Diffusion Models as Text Painters
|
||||
# Paper Link: https://arxiv.org/abs/2305.10855
|
||||
# Code Link: https://github.com/microsoft/unilm/tree/master/textdiffuser
|
||||
# Copyright (c) Microsoft Corporation.
|
||||
# This file define the Layout Transformer for predicting the layout of keywords.
|
||||
# ------------------------------------------
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from transformers import CLIPTokenizer, CLIPTextModel
|
||||
|
||||
class TextConditioner(nn.Module):
|
||||
|
||||
def __init__(self):
|
||||
super(TextConditioner, self).__init__()
|
||||
self.transformer = CLIPTextModel.from_pretrained('openai/clip-vit-large-patch14')
|
||||
self.tokenizer = CLIPTokenizer.from_pretrained('openai/clip-vit-large-patch14')
|
||||
|
||||
# fix
|
||||
self.transformer.eval()
|
||||
for param in self.transformer.parameters():
|
||||
param.requires_grad = False
|
||||
|
||||
def forward(self, prompt_list):
|
||||
batch_encoding = self.tokenizer(prompt_list, truncation=True, max_length=77, return_length=True, return_overflowing_tokens=False, padding="max_length", return_tensors="pt")
|
||||
text_embedding = self.transformer(batch_encoding["input_ids"].cuda())
|
||||
return text_embedding.last_hidden_state.cuda(), batch_encoding["attention_mask"].cuda() # 1, 77, 768 / 1, 768
|
||||
|
||||
|
||||
class LayoutTransformer(nn.Module):
|
||||
|
||||
def __init__(self, layer_number=2):
|
||||
super(LayoutTransformer, self).__init__()
|
||||
|
||||
self.encoder_layer = nn.TransformerEncoderLayer(d_model=512, nhead=8)
|
||||
self.transformer = torch.nn.TransformerEncoder(
|
||||
self.encoder_layer, num_layers=layer_number
|
||||
)
|
||||
|
||||
self.decoder_layer = nn.TransformerDecoderLayer(d_model=512, nhead=8)
|
||||
self.decoder_transformer = torch.nn.TransformerDecoder(
|
||||
self.decoder_layer, num_layers=layer_number
|
||||
)
|
||||
|
||||
self.mask_embedding = nn.Embedding(2,512)
|
||||
self.length_embedding = nn.Embedding(256,512)
|
||||
self.width_embedding = nn.Embedding(256,512)
|
||||
self.position_embedding = nn.Embedding(256,512)
|
||||
self.state_embedding = nn.Embedding(256,512)
|
||||
self.match_embedding = nn.Embedding(256,512)
|
||||
|
||||
self.x_embedding = nn.Embedding(512,512)
|
||||
self.y_embedding = nn.Embedding(512,512)
|
||||
self.w_embedding = nn.Embedding(512,512)
|
||||
self.h_embedding = nn.Embedding(512,512)
|
||||
|
||||
self.encoder_target_embedding = nn.Embedding(256,512)
|
||||
|
||||
self.input_layer = nn.Sequential(
|
||||
nn.Linear(768, 512),
|
||||
nn.ReLU(),
|
||||
nn.Linear(512, 512),
|
||||
)
|
||||
|
||||
self.output_layer = nn.Sequential(
|
||||
nn.Linear(512, 128),
|
||||
nn.ReLU(),
|
||||
nn.Linear(128, 4),
|
||||
)
|
||||
|
||||
def forward(self, x, length, width, mask, state, match, target, right_shifted_boxes, train=False, encoder_embedding=None):
|
||||
|
||||
# detect whether the encoder_embedding is cached
|
||||
if encoder_embedding is None:
|
||||
# augmentation
|
||||
if train:
|
||||
width = width + torch.randint(-3, 3, (width.shape[0], width.shape[1])).cuda()
|
||||
|
||||
x = self.input_layer(x) # (1, 77, 512)
|
||||
width_embedding = self.width_embedding(torch.clamp(width, 0, 255).long()) # (1, 77, 512)
|
||||
encoder_target_embedding = self.encoder_target_embedding(target[:,:,0].long()) # (1, 77, 512)
|
||||
pe_embedding = self.position_embedding(torch.arange(77).cuda()).unsqueeze(0) # (1, 77, 512)
|
||||
total_embedding = x + width_embedding + pe_embedding + encoder_target_embedding # combine all the embeddings (1, 77, 512)
|
||||
total_embedding = total_embedding.permute(1,0,2) # (77, 1, 512)
|
||||
encoder_embedding = self.transformer(total_embedding) # (77, 1, 512)
|
||||
|
||||
right_shifted_boxes_resize = (right_shifted_boxes * 512).long() # (1, 8, 4)
|
||||
right_shifted_boxes_resize = torch.clamp(right_shifted_boxes_resize, 0, 511) # (1, 8, 4)
|
||||
|
||||
# decoder pe
|
||||
pe_decoder = torch.arange(8).cuda() # (8, )
|
||||
pe_embedding_decoder = self.position_embedding(pe_decoder).unsqueeze(0) # (1, 8, 512)
|
||||
decoder_input = pe_embedding_decoder + self.x_embedding(right_shifted_boxes_resize[:,:,0]) + self.y_embedding(right_shifted_boxes_resize[:,:,1]) + self.w_embedding(right_shifted_boxes_resize[:,:,2]) + self.h_embedding(right_shifted_boxes_resize[:,:,3]) # (1, 8, 512)
|
||||
decoder_input = decoder_input.permute(1,0,2) # (8, 1, 512)
|
||||
|
||||
# generate triangular mask
|
||||
mask = nn.Transformer.generate_square_subsequent_mask(8) # (8, 8)
|
||||
mask = mask.cuda() # (8, 8)
|
||||
decoder_result = self.decoder_transformer(decoder_input, encoder_embedding, tgt_mask=mask) # (8, 1, 512)
|
||||
decoder_result = decoder_result.permute(1,0,2) # (1, 8, 512)
|
||||
|
||||
box_prediction = self.output_layer(decoder_result) # (1, 8, 4)
|
||||
return box_prediction, encoder_embedding
|
||||
|
||||
|
||||
|
||||
@@ -0,0 +1,53 @@
|
||||
# ------------------------------------------
|
||||
# TextDiffuser: Diffusion Models as Text Painters
|
||||
# Paper Link: https://arxiv.org/abs/2305.10855
|
||||
# Code Link: https://github.com/microsoft/unilm/tree/master/textdiffuser
|
||||
# Copyright (c) Microsoft Corporation.
|
||||
# This file define the architecture of unet.
|
||||
# ------------------------------------------
|
||||
|
||||
import torch.nn.functional as F
|
||||
from model.text_segmenter.unet_parts import *
|
||||
|
||||
|
||||
class UNet(nn.Module):
|
||||
def __init__(self, n_channels, n_classes, bilinear=True):
|
||||
super(UNet, self).__init__()
|
||||
self.n_channels = n_channels
|
||||
self.n_classes = n_classes
|
||||
self.bilinear = bilinear
|
||||
|
||||
self.inc = DoubleConv(n_channels, 64)
|
||||
self.down1 = Down(64, 128)
|
||||
self.down2 = Down(128, 256)
|
||||
self.down3 = Down(256, 512)
|
||||
factor = 2 if bilinear else 1
|
||||
self.down4 = Down(512, 1024 // factor)
|
||||
self.up1 = Up(1024, 512 // factor, bilinear)
|
||||
self.up2 = Up(512, 256 // factor, bilinear)
|
||||
self.up3 = Up(256, 128 // factor, bilinear)
|
||||
self.up4 = Up(128, 64, bilinear)
|
||||
self.outc = OutConv(64, n_classes)
|
||||
|
||||
def forward(self, x):
|
||||
x1 = self.inc(x)
|
||||
x2 = self.down1(x1)
|
||||
x3 = self.down2(x2)
|
||||
x4 = self.down3(x3)
|
||||
x5 = self.down4(x4)
|
||||
x = self.up1(x5, x4)
|
||||
x = self.up2(x, x3)
|
||||
x = self.up3(x, x2)
|
||||
x = self.up4(x, x1)
|
||||
logits = self.outc(x)
|
||||
# logits = torch.sigmoid(logits)
|
||||
return logits
|
||||
|
||||
if __name__ == '__main__':
|
||||
net = UNet(39,39,True)
|
||||
|
||||
net = net.cuda()
|
||||
|
||||
image = torch.Tensor(32,39,64,64).cuda()
|
||||
result = net(image)
|
||||
print(result.shape)
|
||||
@@ -0,0 +1,82 @@
|
||||
# ------------------------------------------
|
||||
# TextDiffuser: Diffusion Models as Text Painters
|
||||
# Paper Link: https://arxiv.org/abs/2305.10855
|
||||
# Code Link: https://github.com/microsoft/unilm/tree/master/textdiffuser
|
||||
# Copyright (c) Microsoft Corporation.
|
||||
# This file define the architecture of unet.
|
||||
# ------------------------------------------
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
|
||||
class DoubleConv(nn.Module):
|
||||
"""(convolution => [BN] => ReLU) * 2"""
|
||||
|
||||
def __init__(self, in_channels, out_channels, mid_channels=None):
|
||||
super().__init__()
|
||||
if not mid_channels:
|
||||
mid_channels = out_channels
|
||||
self.double_conv = nn.Sequential(
|
||||
nn.Conv2d(in_channels, mid_channels, kernel_size=3, padding=1),
|
||||
nn.BatchNorm2d(mid_channels),
|
||||
nn.ReLU(inplace=True),
|
||||
nn.Conv2d(mid_channels, out_channels, kernel_size=3, padding=1),
|
||||
nn.BatchNorm2d(out_channels),
|
||||
nn.ReLU(inplace=True)
|
||||
)
|
||||
|
||||
def forward(self, x):
|
||||
return self.double_conv(x)
|
||||
|
||||
|
||||
class Down(nn.Module):
|
||||
"""Downscaling with maxpool then double conv"""
|
||||
|
||||
def __init__(self, in_channels, out_channels):
|
||||
super().__init__()
|
||||
self.maxpool_conv = nn.Sequential(
|
||||
nn.MaxPool2d(2),
|
||||
DoubleConv(in_channels, out_channels)
|
||||
)
|
||||
|
||||
def forward(self, x):
|
||||
return self.maxpool_conv(x)
|
||||
|
||||
|
||||
class Up(nn.Module):
|
||||
"""Upscaling then double conv"""
|
||||
|
||||
def __init__(self, in_channels, out_channels, bilinear=True):
|
||||
super().__init__()
|
||||
|
||||
# if bilinear, use the normal convolutions to reduce the number of channels
|
||||
if bilinear:
|
||||
self.up = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)
|
||||
self.conv = DoubleConv(in_channels, out_channels, in_channels // 2)
|
||||
else:
|
||||
self.up = nn.ConvTranspose2d(in_channels , in_channels // 2, kernel_size=2, stride=2)
|
||||
self.conv = DoubleConv(in_channels, out_channels)
|
||||
|
||||
|
||||
def forward(self, x1, x2):
|
||||
x1 = self.up(x1)
|
||||
# input is CHW
|
||||
diffY = x2.size()[2] - x1.size()[2]
|
||||
diffX = x2.size()[3] - x1.size()[3]
|
||||
|
||||
x1 = F.pad(x1, [diffX // 2, diffX - diffX // 2,
|
||||
diffY // 2, diffY - diffY // 2])
|
||||
|
||||
x = torch.cat([x2, x1], dim=1)
|
||||
return self.conv(x)
|
||||
|
||||
|
||||
class OutConv(nn.Module):
|
||||
def __init__(self, in_channels, out_channels):
|
||||
super(OutConv, self).__init__()
|
||||
self.conv = nn.Conv2d(in_channels, out_channels, kernel_size=1)
|
||||
|
||||
def forward(self, x):
|
||||
return self.conv(x)
|
||||
@@ -0,0 +1,12 @@
|
||||
datasets==2.11.0
|
||||
numpy==1.24.2
|
||||
opencv-python==4.1.2.30
|
||||
pillow==5.2.0
|
||||
tokenizers==0.13.3
|
||||
transformers==4.27.4
|
||||
xformers==0.0.16
|
||||
accelerate==0.18.0
|
||||
triton==2.0.0.post1
|
||||
termcolor==2.3.0
|
||||
tinydb
|
||||
flask
|
||||
Executable
+8
@@ -0,0 +1,8 @@
|
||||
CUDA_VISIBLE_DEVICES=0 python inference.py \
|
||||
--mode="text-inpainting" \
|
||||
--resume_from_checkpoint="textdiffuser-ckpt/diffusion_backbone" \
|
||||
--prompt="a boy draws good morning on a board" \
|
||||
--original_image="assets/examples/text-inpainting/case2.jpg" \
|
||||
--text_mask="assets/examples/text-inpainting/case2_mask.jpg" \
|
||||
--output_dir="./output" \
|
||||
--vis_num=4
|
||||
Executable
+7
@@ -0,0 +1,7 @@
|
||||
CUDA_VISIBLE_DEVICES=0 python inference.py \
|
||||
--mode="text-to-image-with-template" \
|
||||
--resume_from_checkpoint="textdiffuser-ckpt/diffusion_backbone" \
|
||||
--prompt="a poster of monkey music festival" \
|
||||
--template_image="assets/examples/text-to-image-with-template/case2.jpg" \
|
||||
--output_dir="./output" \
|
||||
--vis_num=4
|
||||
Executable
+6
@@ -0,0 +1,6 @@
|
||||
CUDA_VISIBLE_DEVICES=0 python inference.py \
|
||||
--mode="text-to-image" \
|
||||
--resume_from_checkpoint="textdiffuser-ckpt/diffusion_backbone" \
|
||||
--prompt="A sign that says 'Hello'" \
|
||||
--output_dir="./output" \
|
||||
--vis_num=4
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,21 @@
|
||||
accelerate launch train.py \
|
||||
--train_batch_size=24 \
|
||||
--gradient_accumulation_steps=4 \
|
||||
--gradient_checkpointing \
|
||||
--mixed_precision="fp16" \
|
||||
--num_train_epochs=2 \
|
||||
--learning_rate=1e-5 \
|
||||
--max_grad_norm=1 \
|
||||
--lr_scheduler="constant" \
|
||||
--lr_warmup_steps=0 \
|
||||
--output_dir="experiment_name" \
|
||||
--enable_xformers_memory_efficient_attention \
|
||||
--dataloader_num_workers=4 \
|
||||
--character_aware_loss_lambda=0.01 \
|
||||
--resume_from_checkpoint="latest" \
|
||||
--drop_caption \
|
||||
--mask_all_ratio=0.5 \
|
||||
--segmentation_mask_aug \
|
||||
--dataset_path=/home/path/to/laion-ocr-unzip \
|
||||
--train_dataset_index_file=/path/to/index_file.txt \
|
||||
--vis_num=8
|
||||
@@ -0,0 +1,349 @@
|
||||
# ------------------------------------------
|
||||
# TextDiffuser: Diffusion Models as Text Painters
|
||||
# Paper Link: https://arxiv.org/abs/2305.10855
|
||||
# Code Link: https://github.com/microsoft/unilm/tree/master/textdiffuser
|
||||
# Copyright (c) Microsoft Corporation.
|
||||
# This file defines a set of commonly used utility functions.
|
||||
# ------------------------------------------
|
||||
|
||||
import os
|
||||
import re
|
||||
import cv2
|
||||
import math
|
||||
import shutil
|
||||
import string
|
||||
import textwrap
|
||||
import numpy as np
|
||||
from PIL import Image, ImageFont, ImageDraw, ImageOps
|
||||
|
||||
from typing import *
|
||||
|
||||
# define alphabet and alphabet_dic
|
||||
alphabet = string.digits + string.ascii_lowercase + string.ascii_uppercase + string.punctuation + ' ' # len(aphabet) = 95
|
||||
alphabet_dic = {}
|
||||
for index, c in enumerate(alphabet):
|
||||
alphabet_dic[c] = index + 1 # the index 0 stands for non-character
|
||||
|
||||
|
||||
|
||||
def transform_mask_pil(mask_root):
|
||||
"""
|
||||
This function extracts the mask area and text area from the images.
|
||||
|
||||
Args:
|
||||
mask_root (str): The path of mask image.
|
||||
* The white area is the unmasked area
|
||||
* The gray area is the masked area
|
||||
* The white area is the text area
|
||||
"""
|
||||
img = np.array(mask_root)
|
||||
img = cv2.resize(img, (512, 512), interpolation=cv2.INTER_NEAREST)
|
||||
gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
|
||||
ret, binary = cv2.threshold(gray, 250, 255, cv2.THRESH_BINARY) # pixel value is set to 0 or 255 according to the threshold
|
||||
return 1 - (binary.astype(np.float32) / 255)
|
||||
|
||||
|
||||
def transform_mask(mask_root: str):
|
||||
"""
|
||||
This function extracts the mask area and text area from the images.
|
||||
|
||||
Args:
|
||||
mask_root (str): The path of mask image.
|
||||
* The white area is the unmasked area
|
||||
* The gray area is the masked area
|
||||
* The white area is the text area
|
||||
"""
|
||||
img = cv2.imread(mask_root)
|
||||
img = cv2.resize(img, (512, 512), interpolation=cv2.INTER_NEAREST)
|
||||
gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
|
||||
ret, binary = cv2.threshold(gray, 250, 255, cv2.THRESH_BINARY) # pixel value is set to 0 or 255 according to the threshold
|
||||
return 1 - (binary.astype(np.float32) / 255)
|
||||
|
||||
|
||||
def segmentation_mask_visualization(font_path: str, segmentation_mask: np.array):
|
||||
"""
|
||||
This function visualizes the segmentaiton masks with characters.
|
||||
|
||||
Args:
|
||||
font_path (str): The path of font. We recommand to use Arial.ttf
|
||||
segmentation_mask (np.array): The character-level segmentation mask.
|
||||
"""
|
||||
segmentation_mask = cv2.resize(segmentation_mask, (64, 64), interpolation=cv2.INTER_NEAREST)
|
||||
font = ImageFont.truetype(font_path, 8)
|
||||
blank = Image.new('RGB', (512,512), (0,0,0))
|
||||
d = ImageDraw.Draw(blank)
|
||||
for i in range(64):
|
||||
for j in range(64):
|
||||
if int(segmentation_mask[i][j]) == 0 or int(segmentation_mask[i][j])-1 >= len(alphabet):
|
||||
continue
|
||||
else:
|
||||
d.text((j*8, i*8), alphabet[int(segmentation_mask[i][j])-1], font=font, fill=(0, 255, 0))
|
||||
return blank
|
||||
|
||||
|
||||
def make_caption_pil(font_path: str, captions: List[str]):
|
||||
"""
|
||||
This function converts captions into pil images.
|
||||
|
||||
Args:
|
||||
font_path (str): The path of font. We recommand to use Arial.ttf
|
||||
captions (List[str]): List of captions.
|
||||
"""
|
||||
caption_pil_list = []
|
||||
font = ImageFont.truetype(font_path, 18)
|
||||
|
||||
for caption in captions:
|
||||
border_size = 2
|
||||
img = Image.new('RGB', (512-4,48-4), (255,255,255))
|
||||
img = ImageOps.expand(img, border=(border_size, border_size, border_size, border_size), fill=(127, 127, 127))
|
||||
draw = ImageDraw.Draw(img)
|
||||
border_size = 2
|
||||
text = caption
|
||||
lines = textwrap.wrap(text, width=40)
|
||||
x, y = 4, 4
|
||||
line_height = font.getsize('A')[1] + 4
|
||||
|
||||
start = 0
|
||||
for line in lines:
|
||||
draw.text((x, y+start), line, font=font, fill=(200, 127, 0))
|
||||
y += line_height
|
||||
|
||||
caption_pil_list.append(img)
|
||||
return caption_pil_list
|
||||
|
||||
|
||||
def filter_segmentation_mask(segmentation_mask: np.array):
|
||||
"""
|
||||
This function removes some noisy predictions of segmentation masks.
|
||||
|
||||
Args:
|
||||
segmentation_mask (np.array): The character-level segmentation mask.
|
||||
"""
|
||||
segmentation_mask[segmentation_mask==alphabet_dic['-']] = 0
|
||||
segmentation_mask[segmentation_mask==alphabet_dic[' ']] = 0
|
||||
return segmentation_mask
|
||||
|
||||
|
||||
|
||||
def combine_image(args, sub_output_dir: str, pred_image_list: List, image_pil: Image, character_mask_pil: Image, character_mask_highlight_pil: Image, caption_pil_list: List):
|
||||
"""
|
||||
This function combines all the outputs and useful inputs together.
|
||||
|
||||
Args:
|
||||
args (argparse.ArgumentParser): The arguments.
|
||||
pred_image_list (List): List of predicted images.
|
||||
image_pil (Image): The original image.
|
||||
character_mask_pil (Image): The character-level segmentation mask.
|
||||
character_mask_highlight_pil (Image): The character-level segmentation mask highlighting character regions with green color.
|
||||
caption_pil_list (List): List of captions.
|
||||
"""
|
||||
|
||||
# # create a "latest" folder to store the results
|
||||
# if os.path.exists(f'{args.output_dir}/latest'):
|
||||
# shutil.rmtree(f'{args.output_dir}/latest')
|
||||
# os.mkdir(f'{args.output_dir}/latest')
|
||||
|
||||
# save each predicted image
|
||||
# os.makedirs(f'{args.output_dir}/{sub_output_dir}', exist_ok=True)
|
||||
for index, img in enumerate(pred_image_list):
|
||||
img.save(f'{args.output_dir}/{sub_output_dir}/{index}.jpg')
|
||||
# img.save(f'{args.output_dir}/latest/{index}.jpg')
|
||||
|
||||
length = len(pred_image_list)
|
||||
lines = math.ceil(length / 3)
|
||||
|
||||
blank = Image.new('RGB', (512*3, 512*(lines+1)+48*lines), (0,0,0))
|
||||
blank.paste(image_pil,(0,0))
|
||||
blank.paste(character_mask_pil,(512,0))
|
||||
blank.paste(character_mask_highlight_pil,(512*2,0))
|
||||
|
||||
for i in range(length):
|
||||
row, col = i // 3, i % 3
|
||||
blank.paste(pred_image_list[i],(512*col,512*(row+1)+48*row))
|
||||
blank.paste(caption_pil_list[i],(512*col,512*(row+1)+48*row+512))
|
||||
|
||||
blank.save(f'{args.output_dir}/{sub_output_dir}/combine.jpg')
|
||||
# blank.save(f'{args.output_dir}/latest/combine.jpg')
|
||||
|
||||
return blank.convert('RGB')
|
||||
|
||||
|
||||
def combine_image_gradio(args, sub_output_dir: str, pred_image_list: List, image_pil: Image, character_mask_pil: Image, character_mask_highlight_pil: Image, caption_pil_list: List):
|
||||
"""
|
||||
This function combines all the outputs and useful inputs together.
|
||||
|
||||
Args:
|
||||
args (argparse.ArgumentParser): The arguments.
|
||||
pred_image_list (List): List of predicted images.
|
||||
image_pil (Image): The original image.
|
||||
character_mask_pil (Image): The character-level segmentation mask.
|
||||
character_mask_highlight_pil (Image): The character-level segmentation mask highlighting character regions with green color.
|
||||
caption_pil_list (List): List of captions.
|
||||
"""
|
||||
|
||||
size = len(pred_image_list)
|
||||
|
||||
if size == 1:
|
||||
return pred_image_list[0]
|
||||
elif size == 2:
|
||||
blank = Image.new('RGB', (512*2, 512), (0,0,0))
|
||||
blank.paste(pred_image_list[0],(0,0))
|
||||
blank.paste(pred_image_list[1],(512,0))
|
||||
elif size == 3:
|
||||
blank = Image.new('RGB', (512*3, 512), (0,0,0))
|
||||
blank.paste(pred_image_list[0],(0,0))
|
||||
blank.paste(pred_image_list[1],(512,0))
|
||||
blank.paste(pred_image_list[2],(1024,0))
|
||||
elif size == 4:
|
||||
blank = Image.new('RGB', (512*2, 512*2), (0,0,0))
|
||||
blank.paste(pred_image_list[0],(0,0))
|
||||
blank.paste(pred_image_list[1],(512,0))
|
||||
blank.paste(pred_image_list[2],(0,512))
|
||||
blank.paste(pred_image_list[3],(512,512))
|
||||
|
||||
|
||||
return blank
|
||||
|
||||
def get_width(font_path, text):
|
||||
"""
|
||||
This function calculates the width of the text.
|
||||
|
||||
Args:
|
||||
font_path (str): user prompt.
|
||||
text (str): user prompt.
|
||||
"""
|
||||
font = ImageFont.truetype(font_path, 24)
|
||||
width, _ = font.getsize(text)
|
||||
return width
|
||||
|
||||
|
||||
|
||||
def get_key_words(text: str):
|
||||
"""
|
||||
This function detect keywords (enclosed by quotes) from user prompts. The keywords are used to guide the layout generation.
|
||||
|
||||
Args:
|
||||
text (str): user prompt.
|
||||
"""
|
||||
|
||||
words = []
|
||||
text = text
|
||||
matches = re.findall(r"'(.*?)'", text) # find the keywords enclosed by ''
|
||||
if matches:
|
||||
for match in matches:
|
||||
words.extend(match.split())
|
||||
|
||||
if len(words) >= 8:
|
||||
return []
|
||||
|
||||
return words
|
||||
|
||||
|
||||
def adjust_overlap_box(box_output, current_index):
|
||||
"""
|
||||
This function adjust the overlapping boxes.
|
||||
|
||||
Args:
|
||||
box_output (List): List of predicted boxes.
|
||||
current_index (int): the index of current box.
|
||||
"""
|
||||
|
||||
if current_index == 0:
|
||||
return box_output
|
||||
else:
|
||||
# judge whether it contains overlap with the last output
|
||||
last_box = box_output[0, current_index-1, :]
|
||||
xmin_last, ymin_last, xmax_last, ymax_last = last_box
|
||||
|
||||
current_box = box_output[0, current_index, :]
|
||||
xmin, ymin, xmax, ymax = current_box
|
||||
|
||||
if xmin_last <= xmin <= xmax_last and ymin_last <= ymin <= ymax_last:
|
||||
print('adjust overlapping')
|
||||
distance_x = xmax_last - xmin
|
||||
distance_y = ymax_last - ymin
|
||||
if distance_x <= distance_y:
|
||||
# avoid overlap
|
||||
new_x_min = xmax_last + 0.025
|
||||
new_x_max = xmax - xmin + xmax_last + 0.025
|
||||
box_output[0,current_index,0] = new_x_min
|
||||
box_output[0,current_index,2] = new_x_max
|
||||
else:
|
||||
new_y_min = ymax_last + 0.025
|
||||
new_y_max = ymax - ymin + ymax_last + 0.025
|
||||
box_output[0,current_index,1] = new_y_min
|
||||
box_output[0,current_index,3] = new_y_max
|
||||
|
||||
elif xmin_last <= xmin <= xmax_last and ymin_last <= ymax <= ymax_last:
|
||||
print('adjust overlapping')
|
||||
new_x_min = xmax_last + 0.05
|
||||
new_x_max = xmax - xmin + xmax_last + 0.05
|
||||
box_output[0,current_index,0] = new_x_min
|
||||
box_output[0,current_index,2] = new_x_max
|
||||
|
||||
return box_output
|
||||
|
||||
|
||||
def shrink_box(box, scale_factor = 0.9):
|
||||
"""
|
||||
This function shrinks the box.
|
||||
|
||||
Args:
|
||||
box (List): List of predicted boxes.
|
||||
scale_factor (float): The scale factor of shrinking.
|
||||
"""
|
||||
|
||||
x1, y1, x2, y2 = box
|
||||
x1_new = x1 + (x2 - x1) * (1 - scale_factor) / 2
|
||||
y1_new = y1 + (y2 - y1) * (1 - scale_factor) / 2
|
||||
x2_new = x2 - (x2 - x1) * (1 - scale_factor) / 2
|
||||
y2_new = y2 - (y2 - y1) * (1 - scale_factor) / 2
|
||||
return (x1_new, y1_new, x2_new, y2_new)
|
||||
|
||||
|
||||
def adjust_font_size(args, width, height, draw, text):
|
||||
"""
|
||||
This function adjusts the font size.
|
||||
|
||||
Args:
|
||||
args (argparse.ArgumentParser): The arguments.
|
||||
width (int): The width of the text.
|
||||
height (int): The height of the text.
|
||||
draw (ImageDraw): The ImageDraw object.
|
||||
text (str): The text.
|
||||
"""
|
||||
|
||||
size_start = height
|
||||
while True:
|
||||
font = ImageFont.truetype(args.font_path, size_start)
|
||||
text_width, _ = draw.textsize(text, font=font)
|
||||
if text_width >= width:
|
||||
size_start = size_start - 1
|
||||
else:
|
||||
return size_start
|
||||
|
||||
|
||||
def inpainting_merge_image(original_image, mask_image, inpainting_image):
|
||||
"""
|
||||
This function merges the original image, mask image and inpainting image.
|
||||
|
||||
Args:
|
||||
original_image (PIL.Image): The original image.
|
||||
mask_image (PIL.Image): The mask images.
|
||||
inpainting_image (PIL.Image): The inpainting images.
|
||||
"""
|
||||
|
||||
original_image = original_image.resize((512, 512))
|
||||
mask_image = mask_image.resize((512, 512))
|
||||
inpainting_image = inpainting_image.resize((512, 512))
|
||||
mask_image.convert('L')
|
||||
threshold = 250
|
||||
table = []
|
||||
for i in range(256):
|
||||
if i < threshold:
|
||||
table.append(1)
|
||||
else:
|
||||
table.append(0)
|
||||
mask_image = mask_image.point(table, "1")
|
||||
merged_image = Image.composite(inpainting_image, original_image, mask_image)
|
||||
return merged_image
|
||||
Reference in New Issue
Block a user