chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,249 @@
|
||||
# TextDiffuser-2: Unleashing the Power of Language Models for Text Rendering (ECCV 2024 Oral)
|
||||
|
||||
<a href='https://arxiv.org/abs/2311.16465'><img src='https://img.shields.io/badge/Arxiv-2311.16465-red'></a>
|
||||
<a href='https://github.com/microsoft/unilm/tree/master/textdiffuser-2'><img src='https://img.shields.io/badge/Code-aka.ms/textdiffuser2-yellow'></a>
|
||||
<a href='https://jingyechen.github.io/textdiffuser2/'><img src='https://img.shields.io/badge/Homepage-link-green'></a>
|
||||
[](https://huggingface.co/spaces/JingyeChen22/TextDiffuser-2)
|
||||
<a href='https://discord.gg/HVEcfcwdHv'><img src='https://img.shields.io/badge/Discord-Invitation-purple'></a>
|
||||
[](https://replicate.com/cjwbw/textdiffuser-2)
|
||||
|
||||
<img src="assets/readme_images/teaser.jpg" width="100%">
|
||||
|
||||
TextDiffuser-2 exhibits enhanced capability powered by language models. In addition to generating text with remarkable accuracy, TextDiffuser-2 provides plausible text layouts and demonstrates a diverse range of text styles.
|
||||
|
||||
<img src="assets/readme_images/architecture.jpg" width="100%">
|
||||
|
||||
## :star2: Highlights
|
||||
|
||||
* We propose **TextDiffuser-2** which utilizes two language models for layout planning and layout encoding, increasing the flexibility and diversity in the process of text rendering.
|
||||
|
||||
* **TextDiffuser-2** alleivate several drawbacks in previous methods, such as (1) *limited flexibility and automation*, (2) *constrained capability of layout prediction*, and (3) *Restricted style diversity*.
|
||||
|
||||
* **TextDiffuser-2** is capable of handling text-to-image, text-to-image with template, and text inpainting tasks. Moreover, TextDiffuser-2 introduces *an additional feature* - it allows for the editing of generated layouts in a conversational manner.
|
||||
|
||||
* ✨ We **release the demo** at [link](https://huggingface.co/spaces/JingyeChen22/TextDiffuser-2). Welcome to use and provide feedbacks.
|
||||
|
||||
|
||||
|
||||
|
||||
## :stopwatch: News
|
||||
|
||||
- __[2024.08.12]__: TextDiffuser-2 is chosen as oral presentation.
|
||||
|
||||
- __[2024.07.01]__: TextDiffuser-2 is accepted to ECCV 2024.
|
||||
|
||||
- __[2024.05.19]__: We release the angle/quadrilateral extensions of TextDiffuser-2. Download the checkpoint at [angle](https://drive.google.com/file/d/1zyZT_-iaZQkTO4R2oOT8VQC5WAWIiF5A/view?usp=sharing) or [quadrilateral](https://drive.google.com/file/d/16KcCetsncTfWtM7qQuTSf7Y08RBTDOcV/view?usp=sharing).
|
||||
|
||||
- __[2023.12.26]__: Code, model, and demo for the text inpainting task are all released. Welcome to play with it at [link](https://huggingface.co/spaces/JingyeChen22/TextDiffuser-2-Text-Inpainting).
|
||||
|
||||
- __[2023.12.12]__: The training and inference code for text-to-image is released. We provide the code for full-parameter training and lora training.
|
||||
|
||||
- __[2023.12.10]__: The demo is released at [link](https://huggingface.co/spaces/JingyeChen22/TextDiffuser-2).
|
||||
|
||||
- __[2023.11.20]__: The paper is available at [link](https://arxiv.org/pdf/2311.16465.pdf).
|
||||
|
||||
|
||||
|
||||
## :hammer_and_wrench: Installation
|
||||
|
||||
Clone this repo:
|
||||
```
|
||||
git clone https://github.com/microsoft/unilm/
|
||||
cd unilm/tree/master/textdiffuser-2
|
||||
```
|
||||
|
||||
Build up a new environment and install packages as follows:
|
||||
```
|
||||
conda create -n textdiffuser2 python=3.8
|
||||
conda activate textdiffuser2
|
||||
pip install -r requirements.txt
|
||||
```
|
||||
|
||||
Meanwhile, please install **torch**, **torchvision**, **xformers** that matches the version of system and cuda version (refer to this [link](https://download.pytorch.org/whl/torch_stable.html)). Please also install [flash-attention](https://github.com/Dao-AILab/flash-attention) if you want to train the layout planner using [FastChat](https://github.com/lm-sys/FastChat). We provide the list of packages used in the experiments at [link](./assets/reference_requirements.txt) for your reference.
|
||||
|
||||
For training the **text inpainting task**, please install the diffusers package using the command ```pip install https://github.com/JingyeChen/diffusers_td2.git```. Note that the U-Net architecture has been modified for receiving more input features.
|
||||
|
||||
<small> If you encounterd an error of *RuntimeError: expected scalar type float Float bu found Half* trigged by *diffusers/models/attention_processor.py*, please use [attention_processor.py](./assets/attention_processor.py) to replace the corresponding file in the installed diffusers library. </small>
|
||||
|
||||
## :floppy_disk: Checkpoint
|
||||
|
||||
We upload the checkpoints to HuggingFace🤗.
|
||||
|
||||
* The checkpoint of layout planner is at [link](https://huggingface.co/JingyeChen22/textdiffuser2_layout_planner).
|
||||
* The checkpoint of diffusion model (full parameter fine-tuning) is at [link](https://huggingface.co/JingyeChen22/textdiffuser2-full-ft).
|
||||
* The checkpoint of diffusion model (lora fine-tuning) is at [link](https://huggingface.co/JingyeChen22/textdiffuser2-lora-ft).
|
||||
|
||||
Note that we provide the checkpoint with context length 77 as it performs better results when rendering general objects.
|
||||
|
||||
## :books: Dataset
|
||||
|
||||
The data for training the layout planner is at [link](./data/layout_planner_data_5k.json).
|
||||
|
||||
We employ the MARIO-10M dataset for training TextDiffuser-2. Please follow the **Dataset** section at [TextDiffuser](https://github.com/microsoft/unilm/tree/master/textdiffuser) to download the dataset, inluding the **train_dataset_index_file**.
|
||||
|
||||
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
|
||||
...
|
||||
```
|
||||
|
||||
## :steam_locomotive: Train
|
||||
|
||||
### Train layout planner
|
||||
|
||||
```bash
|
||||
CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 torchrun --nproc_per_node=8 --master_port=20003 fastchat/train/train_mem.py \
|
||||
--model_name_or_path lmsys/vicuna-7b-v1.5 \
|
||||
--data_path data/layout_planner_data_5k.json \
|
||||
--bf16 True \
|
||||
--output_dir experiment_result \
|
||||
--num_train_epochs 6 \
|
||||
--per_device_train_batch_size 2 \
|
||||
--per_device_eval_batch_size 2 \
|
||||
--gradient_accumulation_steps 16 \
|
||||
--evaluation_strategy "no" \
|
||||
--save_strategy "steps" \
|
||||
--save_steps 500 \
|
||||
--save_total_limit 5 \
|
||||
--learning_rate 2e-5 \
|
||||
--weight_decay 0. \
|
||||
--warmup_ratio 0.03 \
|
||||
--lr_scheduler_type "cosine" \
|
||||
--logging_steps 1 \
|
||||
--fsdp "full_shard auto_wrap" \
|
||||
--fsdp_transformer_layer_cls_to_wrap 'LlamaDecoderLayer' \
|
||||
--tf32 True \
|
||||
--model_max_length 2048 \
|
||||
--gradient_checkpointing True \
|
||||
--lazy_preprocess True
|
||||
```
|
||||
|
||||
It is normal that the loss curve seems like a staircase:
|
||||
<img src="assets/readme_images/m1_loss.jpg" width="80%">
|
||||
|
||||
### Train diffusion model
|
||||
|
||||
For full-parameter training:
|
||||
```bash
|
||||
accelerate launch train_textdiffuser2_t2i_full.py \
|
||||
--pretrained_model_name_or_path="runwayml/stable-diffusion-v1-5" \
|
||||
--train_batch_size=18 \
|
||||
--gradient_accumulation_steps=4 \
|
||||
--gradient_checkpointing \
|
||||
--mixed_precision="fp16" \
|
||||
--num_train_epochs=6 \
|
||||
--learning_rate=1e-5 \
|
||||
--max_grad_norm=1 \
|
||||
--lr_scheduler="constant" \
|
||||
--lr_warmup_steps=0 \
|
||||
--output_dir="diffusion_experiment_result" \
|
||||
--enable_xformers_memory_efficient_attention \
|
||||
--dataloader_num_workers=8 \
|
||||
--index_file_path='/path/to/train_dataset_index.txt' \
|
||||
--dataset_path='/path/to/laion-ocr-select/' \
|
||||
--granularity=128 \
|
||||
--coord_mode="ltrb" \
|
||||
--max_length=77 \
|
||||
--resume_from_checkpoint="latest"
|
||||
```
|
||||
|
||||
For LoRA training:
|
||||
```bash
|
||||
accelerate launch train_textdiffuser2_t2i_lora.py \
|
||||
--pretrained_model_name_or_path="runwayml/stable-diffusion-v1-5" \
|
||||
--train_batch_size=18 \
|
||||
--gradient_accumulation_steps=4 \
|
||||
--gradient_checkpointing \
|
||||
--mixed_precision="fp16" \
|
||||
--num_train_epochs=6 \
|
||||
--learning_rate=1e-4 \
|
||||
--text_encoder_learning_rate=1e-5 \
|
||||
--lr_scheduler="constant" \
|
||||
--output_dir="diffusion_experiment_result" \
|
||||
--enable_xformers_memory_efficient_attention \
|
||||
--dataloader_num_workers=8 \
|
||||
--index_file_path='/path/to/train_dataset_index.txt' \
|
||||
--dataset_path='/path/to/laion-ocr-select/' \
|
||||
--granularity=128 \
|
||||
--coord_mode="ltrb" \
|
||||
--max_length=77 \
|
||||
--resume_from_checkpoint="latest"
|
||||
```
|
||||
|
||||
If you encounter an "out-of-memory" error, please consider reducing the batch size appropriately.
|
||||
|
||||
## :firecracker: Inference
|
||||
|
||||
For full-parameter inference:
|
||||
```bash
|
||||
accelerate launch inference_textdiffuser2_t2i_full.py \
|
||||
--pretrained_model_name_or_path="runwayml/stable-diffusion-v1-5" \
|
||||
--mixed_precision="fp16" \
|
||||
--output_dir="inference_results" \
|
||||
--enable_xformers_memory_efficient_attention \
|
||||
--resume_from_checkpoint="JingyeChen22/textdiffuser2-full-ft" \
|
||||
--granularity=128 \
|
||||
--max_length=77 \
|
||||
--coord_mode="ltrb" \
|
||||
--cfg=7.5 \
|
||||
--sample_steps=20 \
|
||||
--seed=43555 \
|
||||
--m1_model_path="JingyeChen22/textdiffuser2_layout_planner" \
|
||||
--input_format='prompt' \
|
||||
--input_prompt='a hotdog with mustard and other toppings on it'
|
||||
```
|
||||
|
||||
For LoRA inference:
|
||||
```bash
|
||||
accelerate launch inference_textdiffuser2_t2i_lora.py \
|
||||
--pretrained_model_name_or_path="runwayml/stable-diffusion-v1-5" \
|
||||
--gradient_accumulation_steps=4 \
|
||||
--gradient_checkpointing \
|
||||
--mixed_precision="fp16" \
|
||||
--output_dir="inference_results" \
|
||||
--enable_xformers_memory_efficient_attention \
|
||||
--resume_from_checkpoint="JingyeChen22/textdiffuser2-lora-ft" \
|
||||
--granularity=128 \
|
||||
--coord_mode="ltrb" \
|
||||
--cfg=7.5 \
|
||||
--sample_steps=50 \
|
||||
--seed=43555 \
|
||||
--m1_model_path="JingyeChen22/textdiffuser2_layout_planner" \
|
||||
--input_format='prompt' \
|
||||
--input_prompt='a stamp of u.s.a'
|
||||
```
|
||||
|
||||
## :joystick: Demo
|
||||
TextDiffuser-2 has been deployed on [Hugging Face](https://huggingface.co/spaces/JingyeChen22/TextDiffuser-2). Welcome to play with it! You can also run ```python gradio_demo.py``` to use the demo locally.
|
||||
|
||||
<img src="assets/readme_images/demo.jpg" width="100%">
|
||||
|
||||
|
||||
|
||||
## :love_letter: Acknowledgement
|
||||
|
||||
We sincerely thank [AK](https://huggingface.co/akhaliq) and [hysts](https://huggingface.co/hysts) for helping set up the demo. We also feel thankful for the available code/api/demo of [SDXL](https://huggingface.co/stabilityai/stable-diffusion-xl-base-1.0), [PixArt](https://pixart-alpha.github.io/), [Ideogram](https://ideogram.ai/), [DALLE-3](https://openai.com/dall-e-3), and [GlyphControl](https://huggingface.co/spaces/AIGText/GlyphControl).
|
||||
|
||||
|
||||
## :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-2, 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-2, please contact Lei Cui (lecu@microsoft.com) or Furu Wei (fuwei@microsoft.com).
|
||||
|
||||
## :herb: Citation
|
||||
If you find TextDiffuser-2 useful in your research, please consider citing:
|
||||
```
|
||||
@article{chen2023textdiffuser,
|
||||
title={TextDiffuser-2: Unleashing the Power of Language Models for Text Rendering},
|
||||
author={Chen, Jingye and Huang, Yupan and Lv, Tengchao and Cui, Lei and Chen, Qifeng and Wei, Furu},
|
||||
journal={arXiv preprint arXiv:2311.16465},
|
||||
year={2023}
|
||||
}
|
||||
```
|
||||
Binary file not shown.
File diff suppressed because it is too large
Load Diff
Binary file not shown.
|
After Width: | Height: | Size: 263 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 1.2 MiB |
Binary file not shown.
|
After Width: | Height: | Size: 59 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 317 KiB |
@@ -0,0 +1,130 @@
|
||||
accelerate==0.22.0
|
||||
aiofiles==23.2.1
|
||||
aiohttp==3.9.1
|
||||
aiosignal==1.3.1
|
||||
altair==5.2.0
|
||||
anyio==3.7.1
|
||||
async-timeout==4.0.3
|
||||
attrs==23.1.0
|
||||
certifi==2023.11.17
|
||||
charset-normalizer==3.3.2
|
||||
click==8.1.7
|
||||
cmake==3.27.9
|
||||
contourpy==1.1.1
|
||||
cycler==0.12.1
|
||||
datasets==2.11.0
|
||||
diffusers==0.24.0
|
||||
dill==0.3.6
|
||||
einops==0.7.0
|
||||
exceptiongroup==1.2.0
|
||||
fastapi==0.104.1
|
||||
ffmpy==0.3.1
|
||||
filelock==3.13.1
|
||||
flash-attn==2.1.2.post3
|
||||
fonttools==4.46.0
|
||||
frozenlist==1.4.0
|
||||
fschat==0.2.26
|
||||
fsspec==2023.12.0
|
||||
gradio==3.50.2
|
||||
gradio_client==0.6.1
|
||||
h11==0.14.0
|
||||
httpcore==1.0.2
|
||||
httpx==0.25.2
|
||||
huggingface-hub==0.19.4
|
||||
idna==3.6
|
||||
importlib-metadata==7.0.0
|
||||
importlib-resources==6.1.1
|
||||
Jinja2==3.1.2
|
||||
jsonschema==4.20.0
|
||||
jsonschema-specifications==2023.11.2
|
||||
kiwisolver==1.4.5
|
||||
lit==17.0.6
|
||||
markdown-it-py==3.0.0
|
||||
markdown2==2.4.11
|
||||
MarkupSafe==2.1.3
|
||||
matplotlib==3.7.4
|
||||
mdurl==0.1.2
|
||||
mpmath==1.3.0
|
||||
multidict==6.0.4
|
||||
multiprocess==0.70.14
|
||||
networkx==3.1
|
||||
nh3==0.2.14
|
||||
ninja==1.11.1.1
|
||||
numpy==1.24.4
|
||||
nvidia-cublas-cu11==11.10.3.66
|
||||
nvidia-cublas-cu12==12.1.3.1
|
||||
nvidia-cuda-cupti-cu11==11.7.101
|
||||
nvidia-cuda-cupti-cu12==12.1.105
|
||||
nvidia-cuda-nvrtc-cu11==11.7.99
|
||||
nvidia-cuda-nvrtc-cu12==12.1.105
|
||||
nvidia-cuda-runtime-cu11==11.7.99
|
||||
nvidia-cuda-runtime-cu12==12.1.105
|
||||
nvidia-cudnn-cu11==8.5.0.96
|
||||
nvidia-cudnn-cu12==8.9.2.26
|
||||
nvidia-cufft-cu11==10.9.0.58
|
||||
nvidia-cufft-cu12==11.0.2.54
|
||||
nvidia-curand-cu11==10.2.10.91
|
||||
nvidia-curand-cu12==10.3.2.106
|
||||
nvidia-cusolver-cu11==11.4.0.1
|
||||
nvidia-cusolver-cu12==11.4.5.107
|
||||
nvidia-cusparse-cu11==11.7.4.91
|
||||
nvidia-cusparse-cu12==12.1.0.106
|
||||
nvidia-nccl-cu11==2.14.3
|
||||
nvidia-nccl-cu12==2.18.1
|
||||
nvidia-nvjitlink-cu12==12.3.101
|
||||
nvidia-nvtx-cu11==11.7.91
|
||||
nvidia-nvtx-cu12==12.1.105
|
||||
opencv-python==4.8.1.78
|
||||
orjson==3.9.10
|
||||
packaging==23.2
|
||||
pandas==2.0.3
|
||||
peft==0.6.2
|
||||
Pillow==10.1.0
|
||||
pkgutil_resolve_name==1.3.10
|
||||
prompt-toolkit==3.0.41
|
||||
protobuf==4.25.1
|
||||
psutil==5.9.6
|
||||
pyarrow==14.0.1
|
||||
pydantic==1.10.13
|
||||
pydub==0.25.1
|
||||
Pygments==2.17.2
|
||||
pyparsing==3.1.1
|
||||
python-dateutil==2.8.2
|
||||
python-multipart==0.0.6
|
||||
pytz==2023.3.post1
|
||||
PyYAML==6.0.1
|
||||
referencing==0.31.1
|
||||
regex==2023.10.3
|
||||
requests==2.31.0
|
||||
responses==0.18.0
|
||||
rich==13.7.0
|
||||
rpds-py==0.13.2
|
||||
safetensors==0.4.1
|
||||
semantic-version==2.10.0
|
||||
sentencepiece==0.1.99
|
||||
shortuuid==1.0.11
|
||||
six==1.16.0
|
||||
sniffio==1.3.0
|
||||
starlette==0.27.0
|
||||
svgwrite==1.4.3
|
||||
sympy==1.12
|
||||
termcolor==2.4.0
|
||||
tiktoken==0.5.2
|
||||
tokenizers==0.13.3
|
||||
toolz==0.12.0
|
||||
torch @ https://download.pytorch.org/whl/cu117/torch-2.0.1%2Bcu117-cp38-cp38-linux_x86_64.whl#sha256=bec39e6fe7232f399c6a5cda5785517fec759fc0852e0c31d71a39f7bf6b23b3
|
||||
torchvision==0.15.2
|
||||
tqdm==4.66.1
|
||||
transformers==4.28.1
|
||||
triton==2.0.0
|
||||
typing_extensions==4.8.0
|
||||
tzdata==2023.3
|
||||
urllib3==2.1.0
|
||||
uvicorn==0.24.0.post1
|
||||
wavedrom==2.0.3.post3
|
||||
wcwidth==0.2.12
|
||||
websockets==11.0.3
|
||||
xformers==0.0.21
|
||||
xxhash==3.4.1
|
||||
yarl==1.9.3
|
||||
zipp==3.17.0
|
||||
@@ -0,0 +1,16 @@
|
||||
# Configuration for Cog ⚙️
|
||||
# Reference: https://github.com/replicate/cog/blob/main/docs/yaml.md
|
||||
|
||||
build:
|
||||
gpu: true
|
||||
python_version: "3.11"
|
||||
python_packages:
|
||||
- torch==2.0.1
|
||||
- torchvision==0.15.2
|
||||
- diffusers==0.24.0
|
||||
- accelerate==0.22.0
|
||||
- datasets==2.11.0
|
||||
- transformers==4.28.1
|
||||
- fschat==0.2.26
|
||||
- sentencepiece==0.1.99
|
||||
predict: "predict.py:Predictor"
|
||||
@@ -0,0 +1,33 @@
|
||||
import json
|
||||
import random
|
||||
from PIL import Image, ImageDraw, ImageFont
|
||||
|
||||
f = open(f'./layout_planner_data_5k.json')
|
||||
items = json.load(f)
|
||||
# print(len(items))
|
||||
|
||||
dic = random.sample(items, k=1)
|
||||
print(dic)
|
||||
layout = dic[0]['conversations'][1]['value']
|
||||
print(layout)
|
||||
|
||||
blank = Image.new('RGB', (256,256), (0,0,0))
|
||||
draw = ImageDraw.ImageDraw(blank)
|
||||
font = ImageFont.truetype('../assets/arial.ttf', 16)
|
||||
|
||||
for line in layout.split('\n'):
|
||||
line = line.strip()
|
||||
|
||||
if len(line) == 0:
|
||||
break
|
||||
|
||||
pred = ' '.join(line.split()[:-1])
|
||||
box = line.split()[-1]
|
||||
l, t, r, b = [int(i)*2 for i in box.split(',')] # the size of canvas is 256x256
|
||||
draw.rectangle([(l, t), (r, b)], outline ="red")
|
||||
draw.text((l, t), pred, font=font)
|
||||
|
||||
blank.save('test.jpg')
|
||||
f.close()
|
||||
|
||||
print('Visualizations are successfully saved at ./test.jpg')
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,693 @@
|
||||
import os
|
||||
import re
|
||||
import zipfile
|
||||
import torch
|
||||
import gradio as gr
|
||||
|
||||
print('hello', gr.__version__)
|
||||
|
||||
import numpy as np
|
||||
import time
|
||||
from transformers import CLIPTextModel, CLIPTokenizer
|
||||
from diffusers import AutoencoderKL, DDPMScheduler, StableDiffusionPipeline, UNet2DConditionModel, DiffusionPipeline
|
||||
from tqdm import tqdm
|
||||
from PIL import Image
|
||||
from PIL import Image, ImageDraw, ImageFont
|
||||
import random
|
||||
import copy
|
||||
|
||||
import string
|
||||
alphabet = string.digits + string.ascii_lowercase + string.ascii_uppercase + string.punctuation + ' ' # len(aphabet) = 95
|
||||
'''alphabet
|
||||
0123456789abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ!"#$%&\'()*+,-./:;<=>?@[\\]^_`{|}~
|
||||
'''
|
||||
|
||||
if not os.path.exists('images2'):
|
||||
os.system('wget https://huggingface.co/datasets/JingyeChen22/TextDiffuser/resolve/main/images2.zip')
|
||||
with zipfile.ZipFile('images2.zip', 'r') as zip_ref:
|
||||
zip_ref.extractall('.')
|
||||
|
||||
# os.system('nvidia-smi')
|
||||
os.system('ls')
|
||||
|
||||
#### import diffusion models
|
||||
text_encoder = CLIPTextModel.from_pretrained(
|
||||
'JingyeChen22/textdiffuser2-full-ft-inpainting', subfolder="text_encoder"
|
||||
).cuda().half()
|
||||
tokenizer = CLIPTokenizer.from_pretrained(
|
||||
'runwayml/stable-diffusion-v1-5', subfolder="tokenizer"
|
||||
)
|
||||
|
||||
#### additional tokens are introduced, including coordinate tokens and character tokens
|
||||
print('***************')
|
||||
print(len(tokenizer))
|
||||
for i in range(520):
|
||||
tokenizer.add_tokens(['l' + str(i) ]) # left
|
||||
tokenizer.add_tokens(['t' + str(i) ]) # top
|
||||
tokenizer.add_tokens(['r' + str(i) ]) # width
|
||||
tokenizer.add_tokens(['b' + str(i) ]) # height
|
||||
for c in alphabet:
|
||||
tokenizer.add_tokens([f'[{c}]'])
|
||||
print(len(tokenizer))
|
||||
print('***************')
|
||||
|
||||
vae = AutoencoderKL.from_pretrained('runwayml/stable-diffusion-v1-5', subfolder="vae").half().cuda()
|
||||
unet = UNet2DConditionModel.from_pretrained(
|
||||
'JingyeChen22/textdiffuser2-full-ft-inpainting', subfolder="unet"
|
||||
).half().cuda()
|
||||
text_encoder.resize_token_embeddings(len(tokenizer))
|
||||
|
||||
global_dict = {}
|
||||
#### for interactive
|
||||
# stack = []
|
||||
# state = 0
|
||||
font = ImageFont.truetype("./Arial.ttf", 20)
|
||||
|
||||
def skip_fun(i, t, guest_id):
|
||||
global_dict[guest_id]['state'] = 0
|
||||
# global state
|
||||
# state = 0
|
||||
|
||||
|
||||
def exe_undo(i, orig_i, t, guest_id):
|
||||
|
||||
global_dict[guest_id]['stack'] = []
|
||||
global_dict[guest_id]['state'] = 0
|
||||
|
||||
return copy.deepcopy(orig_i)
|
||||
|
||||
|
||||
def exe_redo(i, orig_i, t, guest_id):
|
||||
|
||||
print('redo ',orig_i)
|
||||
|
||||
if type(orig_i) == str:
|
||||
orig_i = Image.open(orig_i)
|
||||
|
||||
# global state
|
||||
# state = 0
|
||||
global_dict[guest_id]['state'] = 0
|
||||
|
||||
if len(global_dict[guest_id]['stack']) > 0:
|
||||
global_dict[guest_id]['stack'].pop()
|
||||
|
||||
image = copy.deepcopy(orig_i)
|
||||
|
||||
draw = ImageDraw.Draw(image)
|
||||
|
||||
for items in global_dict[guest_id]['stack']:
|
||||
# print('now', items)
|
||||
text_position, t = items
|
||||
if len(text_position) == 2:
|
||||
x, y = text_position
|
||||
text_color = (255, 0, 0)
|
||||
draw.text((x+2, y), t, font=font, fill=text_color)
|
||||
r = 4
|
||||
leftUpPoint = (x-r, y-r)
|
||||
rightDownPoint = (x+r, y+r)
|
||||
draw.ellipse((leftUpPoint,rightDownPoint), fill='red')
|
||||
elif len(text_position) == 4:
|
||||
x0, y0, x1, y1 = text_position
|
||||
text_color = (255, 0, 0)
|
||||
draw.text((x0+2, y0), t, font=font, fill=text_color)
|
||||
r = 4
|
||||
leftUpPoint = (x0-r, y0-r)
|
||||
rightDownPoint = (x0+r, y0+r)
|
||||
draw.ellipse((leftUpPoint,rightDownPoint), fill='red')
|
||||
draw.rectangle((x0,y0,x1,y1), outline=(255, 0, 0) )
|
||||
|
||||
print('stack', global_dict[guest_id]['stack'])
|
||||
return image
|
||||
|
||||
def get_pixels(i, orig_i, radio, t, guest_id, evt: gr.SelectData):
|
||||
|
||||
print('hi1 ', i)
|
||||
print('hi2 ', orig_i)
|
||||
|
||||
width, height = Image.open(i).size
|
||||
|
||||
# register
|
||||
if guest_id == '-1': # register for the first time
|
||||
seed = str(int(time.time()))
|
||||
global_dict[str(seed)] = {
|
||||
'state': 0,
|
||||
'stack': [],
|
||||
'image_id': [list(Image.open(i).resize((512,512)).getdata())] # an image has been recorded
|
||||
}
|
||||
guest_id = str(seed)
|
||||
else:
|
||||
seed = guest_id
|
||||
|
||||
if type(i) == str:
|
||||
i = Image.open(i)
|
||||
i = i.resize((512,512))
|
||||
|
||||
images = global_dict[str(seed)]['image_id']
|
||||
flag = False
|
||||
for image in images:
|
||||
if image == list(i.getdata()):
|
||||
print('find it')
|
||||
flag = True
|
||||
break
|
||||
|
||||
if not flag:
|
||||
global_dict[str(seed)]['image_id'] = [list(i.getdata())]
|
||||
global_dict[str(seed)]['stack'] = []
|
||||
global_dict[str(seed)]['state'] = 0
|
||||
orig_i = i
|
||||
else:
|
||||
|
||||
if orig_i is not None:
|
||||
orig_i = Image.open(orig_i)
|
||||
orig_i = orig_i.resize((512,512))
|
||||
else:
|
||||
orig_i = i
|
||||
global_dict[guest_id]['stack'] = []
|
||||
global_dict[guest_id]['state'] = 0
|
||||
|
||||
text_position = evt.index
|
||||
|
||||
print('hello ', text_position)
|
||||
|
||||
if radio == 'Two Points':
|
||||
|
||||
if global_dict[guest_id]['state'] == 0:
|
||||
global_dict[guest_id]['stack'].append(
|
||||
(text_position, t)
|
||||
)
|
||||
print(text_position, global_dict[guest_id]['stack'])
|
||||
global_dict[guest_id]['state'] = 1
|
||||
else:
|
||||
|
||||
(_, t) = global_dict[guest_id]['stack'].pop()
|
||||
x, y = _
|
||||
global_dict[guest_id]['stack'].append(
|
||||
((x,y,text_position[0],text_position[1]), t)
|
||||
)
|
||||
global_dict[guest_id]['state'] = 0
|
||||
|
||||
image = copy.deepcopy(orig_i)
|
||||
draw = ImageDraw.Draw(image)
|
||||
|
||||
for items in global_dict[guest_id]['stack']:
|
||||
text_position, t = items
|
||||
if len(text_position) == 2:
|
||||
x, y = text_position
|
||||
|
||||
x = int(512 * x / width)
|
||||
y = int(512 * y / height)
|
||||
|
||||
text_color = (255, 0, 0)
|
||||
draw.text((x+2, y), t, font=font, fill=text_color)
|
||||
r = 4
|
||||
leftUpPoint = (x-r, y-r)
|
||||
rightDownPoint = (x+r, y+r)
|
||||
draw.ellipse((leftUpPoint,rightDownPoint), fill='red')
|
||||
elif len(text_position) == 4:
|
||||
x0, y0, x1, y1 = text_position
|
||||
|
||||
x0 = int(512 * x0 / width)
|
||||
x1 = int(512 * x1 / width)
|
||||
y0 = int(512 * y0 / height)
|
||||
y1 = int(512 * y1 / height)
|
||||
|
||||
text_color = (255, 0, 0)
|
||||
draw.text((x0+2, y0), t, font=font, fill=text_color)
|
||||
r = 4
|
||||
leftUpPoint = (x0-r, y0-r)
|
||||
rightDownPoint = (x0+r, y0+r)
|
||||
draw.ellipse((leftUpPoint,rightDownPoint), fill='red')
|
||||
draw.rectangle((x0,y0,x1,y1), outline=(255, 0, 0) )
|
||||
|
||||
elif radio == 'Four Points':
|
||||
|
||||
if global_dict[guest_id]['state'] == 0:
|
||||
global_dict[guest_id]['stack'].append(
|
||||
(text_position, t)
|
||||
)
|
||||
print(text_position, global_dict[guest_id]['stack'])
|
||||
global_dict[guest_id]['state'] = 1
|
||||
elif global_dict[guest_id]['state'] == 1:
|
||||
(_, t) = global_dict[guest_id]['stack'].pop()
|
||||
x, y = _
|
||||
global_dict[guest_id]['stack'].append(
|
||||
((x,y,text_position[0],text_position[1]), t)
|
||||
)
|
||||
global_dict[guest_id]['state'] = 2
|
||||
elif global_dict[guest_id]['state'] == 2:
|
||||
(_, t) = global_dict[guest_id]['stack'].pop()
|
||||
x0, y0, x1, y1 = _
|
||||
global_dict[guest_id]['stack'].append(
|
||||
((x0, y0, x1, y1,text_position[0],text_position[1]), t)
|
||||
)
|
||||
global_dict[guest_id]['state'] = 3
|
||||
elif global_dict[guest_id]['state'] == 3:
|
||||
(_, t) = global_dict[guest_id]['stack'].pop()
|
||||
x0, y0, x1, y1, x2, y2 = _
|
||||
global_dict[guest_id]['stack'].append(
|
||||
((x0, y0, x1, y1, x2, y2,text_position[0],text_position[1]), t)
|
||||
)
|
||||
global_dict[guest_id]['state'] = 0
|
||||
|
||||
image = copy.deepcopy(orig_i)
|
||||
draw = ImageDraw.Draw(image)
|
||||
|
||||
for items in global_dict[guest_id]['stack']:
|
||||
text_position, t = items
|
||||
if len(text_position) == 2:
|
||||
x, y = text_position
|
||||
|
||||
x = int(512 * x / width)
|
||||
y = int(512 * y / height)
|
||||
|
||||
text_color = (255, 0, 0)
|
||||
draw.text((x+2, y), t, font=font, fill=text_color)
|
||||
r = 4
|
||||
leftUpPoint = (x-r, y-r)
|
||||
rightDownPoint = (x+r, y+r)
|
||||
draw.ellipse((leftUpPoint,rightDownPoint), fill='red')
|
||||
elif len(text_position) == 4:
|
||||
x0, y0, x1, y1 = text_position
|
||||
text_color = (255, 0, 0)
|
||||
draw.text((x0+2, y0), t, font=font, fill=text_color)
|
||||
r = 4
|
||||
leftUpPoint = (x0-r, y0-r)
|
||||
rightDownPoint = (x0+r, y0+r)
|
||||
draw.ellipse((leftUpPoint,rightDownPoint), fill='red')
|
||||
draw.line(((x0,y0),(x1,y1)), fill=(255, 0, 0) )
|
||||
elif len(text_position) == 6:
|
||||
x0, y0, x1, y1, x2, y2 = text_position
|
||||
text_color = (255, 0, 0)
|
||||
draw.text((x0+2, y0), t, font=font, fill=text_color)
|
||||
r = 4
|
||||
leftUpPoint = (x0-r, y0-r)
|
||||
rightDownPoint = (x0+r, y0+r)
|
||||
draw.ellipse((leftUpPoint,rightDownPoint), fill='red')
|
||||
draw.line(((x0,y0),(x1,y1)), fill=(255, 0, 0) )
|
||||
draw.line(((x1,y1),(x2,y2)), fill=(255, 0, 0) )
|
||||
elif len(text_position) == 8:
|
||||
x0, y0, x1, y1, x2, y2, x3, y3 = text_position
|
||||
text_color = (255, 0, 0)
|
||||
draw.text((x0+2, y0), t, font=font, fill=text_color)
|
||||
r = 4
|
||||
leftUpPoint = (x0-r, y0-r)
|
||||
rightDownPoint = (x0+r, y0+r)
|
||||
draw.ellipse((leftUpPoint,rightDownPoint), fill='red')
|
||||
draw.line(((x0,y0),(x1,y1)), fill=(255, 0, 0) )
|
||||
draw.line(((x1,y1),(x2,y2)), fill=(255, 0, 0) )
|
||||
draw.line(((x2,y2),(x3,y3)), fill=(255, 0, 0) )
|
||||
draw.line(((x3,y3),(x0,y0)), fill=(255, 0, 0) )
|
||||
|
||||
|
||||
print('stack', global_dict[guest_id]['stack'])
|
||||
|
||||
global_dict[str(seed)]['image_id'].append(list(image.getdata()))
|
||||
|
||||
return image, orig_i, seed
|
||||
|
||||
|
||||
font_layout = ImageFont.truetype('./Arial.ttf', 16)
|
||||
|
||||
def get_layout_image(ocrs):
|
||||
|
||||
blank = Image.new('RGB', (256,256), (0,0,0))
|
||||
draw = ImageDraw.ImageDraw(blank)
|
||||
|
||||
for line in ocrs.split('\n'):
|
||||
line = line.strip()
|
||||
|
||||
if len(line) == 0:
|
||||
break
|
||||
|
||||
pred = ' '.join(line.split()[:-1])
|
||||
box = line.split()[-1]
|
||||
l, t, r, b = [int(i)*2 for i in box.split(',')] # the size of canvas is 256x256
|
||||
draw.rectangle([(l, t), (r, b)], outline ="red")
|
||||
draw.text((l, t), pred, font=font_layout)
|
||||
|
||||
return blank
|
||||
|
||||
|
||||
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 test_fn(x,y):
|
||||
print('hello')
|
||||
|
||||
def text_to_image(guest_id, i, orig_i, prompt,keywords,positive_prompt,radio,slider_step,slider_guidance,slider_batch,slider_temperature,slider_natural):
|
||||
|
||||
# print(type(i))
|
||||
# exit(0)
|
||||
|
||||
print(f'[info] Prompt: {prompt} | Keywords: {keywords} | Radio: {radio} | Steps: {slider_step} | Guidance: {slider_guidance} | Natural: {slider_natural}')
|
||||
|
||||
# global stack
|
||||
# global state
|
||||
|
||||
if len(positive_prompt.strip()) != 0:
|
||||
prompt += positive_prompt
|
||||
|
||||
with torch.no_grad():
|
||||
time1 = time.time()
|
||||
user_prompt = prompt
|
||||
|
||||
if slider_natural:
|
||||
user_prompt = f'{user_prompt}'
|
||||
composed_prompt = user_prompt
|
||||
prompt = tokenizer.encode(user_prompt)
|
||||
layout_image = None
|
||||
else:
|
||||
if guest_id not in global_dict or len(global_dict[guest_id]['stack']) == 0:
|
||||
|
||||
if len(keywords.strip()) == 0:
|
||||
template = f'Given a prompt that will be used to generate an image, plan the layout of visual text for the image. The size of the image is 128x128. Therefore, all properties of the positions should not exceed 128, including the coordinates of top, left, right, and bottom. All keywords are included in the caption. You dont need to specify the details of font styles. At each line, the format should be keyword left, top, right, bottom. So let us begin. Prompt: {user_prompt}'
|
||||
else:
|
||||
keywords = keywords.split('/')
|
||||
keywords = [i.strip() for i in keywords]
|
||||
template = f'Given a prompt that will be used to generate an image, plan the layout of visual text for the image. The size of the image is 128x128. Therefore, all properties of the positions should not exceed 128, including the coordinates of top, left, right, and bottom. In addition, we also provide all keywords at random order for reference. You dont need to specify the details of font styles. At each line, the format should be keyword left, top, right, bottom. So let us begin. Prompt: {prompt}. Keywords: {str(keywords)}'
|
||||
|
||||
msg = template
|
||||
conv = get_conversation_template(m1_model_path)
|
||||
conv.append_message(conv.roles[0], msg)
|
||||
conv.append_message(conv.roles[1], None)
|
||||
prompt = conv.get_prompt()
|
||||
inputs = m1_tokenizer([prompt], return_token_type_ids=False)
|
||||
inputs = {k: torch.tensor(v).to('cuda') for k, v in inputs.items()}
|
||||
output_ids = m1_model.generate(
|
||||
**inputs,
|
||||
do_sample=True,
|
||||
temperature=slider_temperature,
|
||||
repetition_penalty=1.0,
|
||||
max_new_tokens=512,
|
||||
)
|
||||
|
||||
if m1_model.config.is_encoder_decoder:
|
||||
output_ids = output_ids[0]
|
||||
else:
|
||||
output_ids = output_ids[0][len(inputs["input_ids"][0]) :]
|
||||
outputs = m1_tokenizer.decode(
|
||||
output_ids, skip_special_tokens=True, spaces_between_special_tokens=False
|
||||
)
|
||||
print(f"[{conv.roles[0]}]\n{msg}")
|
||||
print(f"[{conv.roles[1]}]\n{outputs}")
|
||||
layout_image = get_layout_image(outputs)
|
||||
|
||||
ocrs = outputs.split('\n')
|
||||
time2 = time.time()
|
||||
print(time2-time1)
|
||||
|
||||
# user_prompt = prompt
|
||||
current_ocr = ocrs
|
||||
|
||||
|
||||
ocr_ids = []
|
||||
print('user_prompt', user_prompt)
|
||||
print('current_ocr', current_ocr)
|
||||
|
||||
|
||||
for ocr in current_ocr:
|
||||
ocr = ocr.strip()
|
||||
|
||||
if len(ocr) == 0 or '###' in ocr or '.com' in ocr:
|
||||
continue
|
||||
|
||||
items = ocr.split()
|
||||
pred = ' '.join(items[:-1])
|
||||
box = items[-1]
|
||||
|
||||
l,t,r,b = box.split(',')
|
||||
l,t,r,b = int(l), int(t), int(r), int(b)
|
||||
ocr_ids.extend(['l'+str(l), 't'+str(t), 'r'+str(r), 'b'+str(b)])
|
||||
|
||||
char_list = list(pred)
|
||||
char_list = [f'[{i}]' for i in char_list]
|
||||
ocr_ids.extend(char_list)
|
||||
ocr_ids.append(tokenizer.eos_token_id)
|
||||
|
||||
caption_ids = tokenizer(
|
||||
user_prompt, truncation=True, return_tensors="pt"
|
||||
).input_ids[0].tolist()
|
||||
|
||||
try:
|
||||
ocr_ids = tokenizer.encode(ocr_ids)
|
||||
prompt = caption_ids + ocr_ids
|
||||
except:
|
||||
prompt = caption_ids
|
||||
|
||||
user_prompt = tokenizer.decode(prompt)
|
||||
composed_prompt = tokenizer.decode(prompt)
|
||||
|
||||
else:
|
||||
user_prompt += ' <|endoftext|><|startoftext|>'
|
||||
layout_image = None
|
||||
|
||||
image_mask = Image.new('L', (512,512), 0)
|
||||
draw = ImageDraw.Draw(image_mask)
|
||||
|
||||
for items in global_dict[guest_id]['stack']:
|
||||
position, text = items
|
||||
|
||||
# feature_mask
|
||||
# masked_feature
|
||||
|
||||
if len(position) == 2:
|
||||
x, y = position
|
||||
x = x // 4
|
||||
y = y // 4
|
||||
text_str = ' '.join([f'[{c}]' for c in list(text)])
|
||||
user_prompt += f' l{x} t{y} {text_str} <|endoftext|>'
|
||||
|
||||
elif len(position) == 4:
|
||||
x0, y0, x1, y1 = position
|
||||
x0 = x0 // 4
|
||||
y0 = y0 // 4
|
||||
x1 = x1 // 4
|
||||
y1 = y1 // 4
|
||||
text_str = ' '.join([f'[{c}]' for c in list(text)])
|
||||
user_prompt += f' l{x0} t{y0} r{x1} b{y1} {text_str} <|endoftext|>'
|
||||
|
||||
draw.rectangle((x0*4, y0*4, x1*4, y1*4), fill=1)
|
||||
print('prompt ', user_prompt)
|
||||
|
||||
elif len(position) == 8: # four points
|
||||
x0, y0, x1, y1, x2, y2, x3, y3 = position
|
||||
draw.polygon([(x0, y0), (x1, y1), (x2, y2), (x3, y3)], fill=1)
|
||||
x0 = x0 // 4
|
||||
y0 = y0 // 4
|
||||
x1 = x1 // 4
|
||||
y1 = y1 // 4
|
||||
x2 = x2 // 4
|
||||
y2 = y2 // 4
|
||||
x3 = x3 // 4
|
||||
y3 = y3 // 4
|
||||
xmin = min(x0, x1, x2, x3)
|
||||
ymin = min(y0, y1, y2, y3)
|
||||
xmax = max(x0, x1, x2, x3)
|
||||
ymax = max(y0, y1, y2, y3)
|
||||
text_str = ' '.join([f'[{c}]' for c in list(text)])
|
||||
user_prompt += f' l{xmin} t{ymin} r{xmax} b{ymax} {text_str} <|endoftext|>'
|
||||
|
||||
print('prompt ', user_prompt)
|
||||
|
||||
|
||||
prompt = tokenizer.encode(user_prompt)
|
||||
composed_prompt = tokenizer.decode(prompt)
|
||||
|
||||
prompt = prompt[:77]
|
||||
while len(prompt) < 77:
|
||||
prompt.append(tokenizer.pad_token_id)
|
||||
|
||||
prompts_cond = prompt
|
||||
prompts_nocond = [tokenizer.pad_token_id]*77
|
||||
|
||||
prompts_cond = [prompts_cond] * slider_batch
|
||||
prompts_nocond = [prompts_nocond] * slider_batch
|
||||
|
||||
prompts_cond = torch.Tensor(prompts_cond).long().cuda()
|
||||
prompts_nocond = torch.Tensor(prompts_nocond).long().cuda()
|
||||
|
||||
scheduler = DDPMScheduler.from_pretrained('runwayml/stable-diffusion-v1-5', subfolder="scheduler")
|
||||
scheduler.set_timesteps(slider_step)
|
||||
noise = torch.randn((slider_batch, 4, 64, 64)).to("cuda").half()
|
||||
input = noise
|
||||
|
||||
encoder_hidden_states_cond = text_encoder(prompts_cond)[0].half()
|
||||
encoder_hidden_states_nocond = text_encoder(prompts_nocond)[0].half()
|
||||
|
||||
image_mask = torch.Tensor(np.array(image_mask)).float().half().cuda()
|
||||
image_mask = image_mask.unsqueeze(0).unsqueeze(0).repeat(slider_batch, 1, 1, 1)
|
||||
|
||||
image = Image.open(orig_i).resize((512,512))
|
||||
image_tensor = to_tensor(image).unsqueeze(0).cuda().sub_(0.5).div_(0.5)
|
||||
print(f'image_tensor.shape {image_tensor.shape}')
|
||||
masked_image = image_tensor * (1-image_mask)
|
||||
masked_feature = vae.encode(masked_image.half()).latent_dist.sample()
|
||||
masked_feature = masked_feature * vae.config.scaling_factor
|
||||
masked_feature = masked_feature.half()
|
||||
print(f'masked_feature.shape {masked_feature.shape}')
|
||||
|
||||
feature_mask = torch.nn.functional.interpolate(image_mask, size=(64,64), mode='nearest').cuda()
|
||||
|
||||
for t in tqdm(scheduler.timesteps):
|
||||
with torch.no_grad(): # classifier free guidance
|
||||
|
||||
noise_pred_cond = unet(sample=input, timestep=t, encoder_hidden_states=encoder_hidden_states_cond[:slider_batch],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[:slider_batch],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
|
||||
input = scheduler.step(noisy_residual, t, input).prev_sample
|
||||
del noise_pred_cond
|
||||
del noise_pred_uncond
|
||||
|
||||
torch.cuda.empty_cache()
|
||||
|
||||
# decode
|
||||
input = 1 / vae.config.scaling_factor * input
|
||||
images = vae.decode(input, return_dict=False)[0]
|
||||
width, height = 512, 512
|
||||
results = []
|
||||
new_image = Image.new('RGB', (2*width, 2*height))
|
||||
for index, image in enumerate(images.cpu().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')
|
||||
results.append(image)
|
||||
row = index // 2
|
||||
col = index % 2
|
||||
new_image.paste(image, (col*width, row*height))
|
||||
# os.system('nvidia-smi')
|
||||
torch.cuda.empty_cache()
|
||||
# os.system('nvidia-smi')
|
||||
return tuple(results), composed_prompt
|
||||
|
||||
with gr.Blocks() as demo:
|
||||
|
||||
gr.HTML(
|
||||
"""
|
||||
<div style="text-align: center; max-width: 1600px; margin: 20px auto;">
|
||||
<h2 style="font-weight: 900; font-size: 2.3rem; margin: 0rem">
|
||||
TextDiffuser-2: Unleashing the Power of Language Models for Text Rendering
|
||||
</h2>
|
||||
<h2 style="font-weight: 900; font-size: 1.3rem; margin: 0rem">
|
||||
(Demo for <b>Text Inpainting</b> 🖼️🖌️)
|
||||
</h2>
|
||||
<h2 style="font-weight: 460; font-size: 1.1rem; margin: 0rem">
|
||||
<a href="https://jingyechen.github.io/">Jingye Chen</a>, <a href="https://hypjudy.github.io/website/">Yupan Huang</a>, <a href="https://scholar.google.com/citations?user=0LTZGhUAAAAJ&hl=en">Tengchao Lv</a>, <a href="https://www.microsoft.com/en-us/research/people/lecu/">Lei Cui</a>, <a href="https://cqf.io/">Qifeng Chen</a>, <a href="https://thegenerality.com/">Furu Wei</a>
|
||||
</h2>
|
||||
<h2 style="font-weight: 460; font-size: 1.1rem; margin: 0rem">
|
||||
HKUST, Sun Yat-sen University, Microsoft Research
|
||||
</h2>
|
||||
<h3 style="font-weight: 450; font-size: 1rem; margin: 0rem">
|
||||
[<a href="https://arxiv.org/abs/2311.16465" style="color:blue;">arXiv</a>]
|
||||
[<a href="https://github.com/microsoft/unilm/tree/master/textdiffuser-2" style="color:blue;">Code</a>]
|
||||
[<a href="https://jingyechen.github.io/textdiffuser2/" style="color:blue;">Project Page</a>]
|
||||
[<a href="https://discord.gg/q7eHPupu" style="color:purple;">Discord</a>]
|
||||
</h3>
|
||||
<h2 style="text-align: left; font-weight: 450; font-size: 1rem; margin-top: 0.5rem; margin-bottom: 0.5rem">
|
||||
TextDiffuser-2 leverages language models to enhance text rendering, achieving greater flexibility. Different from text editing, the text inpainting task aims to add or modify text guided by users, ensuring that the inpainted text has a reasonable style (i.e., no need to match the style of the original text during modification exactly) and is coherent with backgrounds. TextDiffuser-2 offers an <b>improved user experience</b>. Specifically, users only need to type the text they wish to inpaint into the provided input box and then select key points on the Canvas.
|
||||
</h2>
|
||||
<h2 style="text-align: left; font-weight: 450; font-size: 1rem; margin-top: 0.5rem; margin-bottom: 0.5rem">
|
||||
👀 <b>Tips for using this demo</b>: <b>(1)</b> Please carefully read the disclaimer in the below. Current verison can only support English. <b>(2)</b> The <b>prompt is optional</b>. If provided, the generated image may be more accurate. <b>(3)</b> Redo is used to cancel the last keyword, and undo is used to clear all keywords. <b>(4)</b> Current version only supports input image with resolution 512x512. <b>(5)</b> You can use either two points or four points to specify the text box. Using four points can better represent the perspective boxes. <b>(6)</b> Leave "Text to be inpaintd" empty can function as the text removal task. <b>(7)</b> Classifier-free guidance is set to a small value (e.g. 1) in default. It is noticed that a larger cfg may result in chromatic aberration against the background. <b>(8)</b> You can inpaint many text regions at one time. <b>(9)</b> Thanks for reading these tips, shall we start now?
|
||||
</h2>
|
||||
<img src="https://raw.githubusercontent.com/JingyeChen/jingyechen.github.io/master/textdiffuser2/static/images/inpainting_blank.jpg" alt="textdiffuser-2">
|
||||
</div>
|
||||
""")
|
||||
|
||||
with gr.Tab("Text Inpainting"):
|
||||
with gr.Row():
|
||||
with gr.Column():
|
||||
|
||||
keywords = gr.Textbox(label="(Optional) Keywords. Should be seperated by / (e.g., keyword1/keyword2/...)", placeholder="keyword1/keyword2", visible=False)
|
||||
positive_prompt = gr.Textbox(label="(Optional) Positive prompt", value="", visible=False)
|
||||
|
||||
i = gr.Image(label="Image", type='filepath', value='https://raw.githubusercontent.com/JingyeChen/jingyechen.github.io/master/textdiffuser2/static/images/example11.jpg')
|
||||
orig_i = gr.Image(label="Placeholder", type='filepath', height=512, width=512, visible=False)
|
||||
|
||||
radio = gr.Radio(["Two Points", "Four Points"], label="Number of points to represent the text box.", value="Two Points", visible=True)
|
||||
|
||||
with gr.Row():
|
||||
t = gr.Textbox(label="Text to be inpainted", value='Test')
|
||||
prompt = gr.Textbox(label="(Optional) Prompt.")
|
||||
with gr.Row():
|
||||
redo = gr.Button(value='Redo - Cancel the last keyword')
|
||||
undo = gr.Button(value='Undo - Clear the canvas')
|
||||
# skip_button = gr.Button(value='Skip - Operate the next keyword')
|
||||
|
||||
slider_natural = gr.Checkbox(label="Natural image generation", value=False, info="The text position and content info will not be incorporated.", visible=False)
|
||||
slider_step = gr.Slider(minimum=1, maximum=50, value=20, step=1, label="Sampling step", info="The sampling step for TextDiffuser-2.")
|
||||
slider_guidance = gr.Slider(minimum=1, maximum=13, value=1, step=0.5, label="Scale of classifier-free guidance", info="The scale of cfg and is set to 1 in default. Smaller cfg produce stable results.")
|
||||
slider_batch = gr.Slider(minimum=1, maximum=6, value=4, step=1, label="Batch size", info="The number of images to be sampled.")
|
||||
slider_temperature = gr.Slider(minimum=0.1, maximum=2, value=1.4, step=0.1, label="Temperature", info="Control the diversity of layout planner. Higher value indicates more diversity.", visible=False)
|
||||
# slider_seed = gr.Slider(minimum=1, maximum=10000, label="Seed", randomize=True)
|
||||
button = gr.Button("Generate")
|
||||
|
||||
guest_id_box = gr.Textbox(label="guest_id", value=f"-1", visible=False)
|
||||
i.select(get_pixels,[i,orig_i,radio,t,guest_id_box],[i,orig_i,guest_id_box])
|
||||
redo.click(exe_redo, [i,orig_i,t,guest_id_box],[i])
|
||||
undo.click(exe_undo, [i,orig_i,t,guest_id_box],[i])
|
||||
# skip_button.click(skip_fun, [i,t,guest_id_box])
|
||||
|
||||
|
||||
with gr.Column():
|
||||
output = gr.Gallery(label='Generated image', rows=2, height=768)
|
||||
|
||||
with gr.Accordion("Intermediate results", open=False, visible=False):
|
||||
gr.Markdown("Composed prompt")
|
||||
composed_prompt = gr.Textbox(label='')
|
||||
# gr.Markdown("Layout visualization")
|
||||
# layout = gr.Image(height=256, width=256)
|
||||
|
||||
|
||||
button.click(text_to_image, inputs=[guest_id_box, i, orig_i, prompt,keywords,positive_prompt, radio,slider_step,slider_guidance,slider_batch,slider_temperature,slider_natural], outputs=[output, composed_prompt])
|
||||
|
||||
gr.Markdown("## Image Examples")
|
||||
template = None
|
||||
gr.Examples(
|
||||
[
|
||||
["https://raw.githubusercontent.com/JingyeChen/jingyechen.github.io/master/textdiffuser2/static/images/example1.jpg"],
|
||||
["https://raw.githubusercontent.com/JingyeChen/jingyechen.github.io/master/textdiffuser2/static/images/example2.jpg"],
|
||||
["https://raw.githubusercontent.com/JingyeChen/jingyechen.github.io/master/textdiffuser2/static/images/example3.jpg"],
|
||||
["https://raw.githubusercontent.com/JingyeChen/jingyechen.github.io/master/textdiffuser2/static/images/example4.jpg"],
|
||||
["https://raw.githubusercontent.com/JingyeChen/jingyechen.github.io/master/textdiffuser2/static/images/example5.jpg"],
|
||||
["https://raw.githubusercontent.com/JingyeChen/jingyechen.github.io/master/textdiffuser2/static/images/example7.jpg"],
|
||||
["https://raw.githubusercontent.com/JingyeChen/jingyechen.github.io/master/textdiffuser2/static/images/example8.jpg"],
|
||||
["https://raw.githubusercontent.com/JingyeChen/jingyechen.github.io/master/textdiffuser2/static/images/example11.jpg"],
|
||||
["https://raw.githubusercontent.com/JingyeChen/jingyechen.github.io/master/textdiffuser2/static/images/example12.jpg"],
|
||||
["https://raw.githubusercontent.com/JingyeChen/jingyechen.github.io/master/textdiffuser2/static/images/example13.jpg"],
|
||||
["https://raw.githubusercontent.com/JingyeChen/jingyechen.github.io/master/textdiffuser2/static/images/example14.jpg"],
|
||||
["https://raw.githubusercontent.com/JingyeChen/jingyechen.github.io/master/textdiffuser2/static/images/example15.jpg"],
|
||||
],
|
||||
[
|
||||
i
|
||||
],
|
||||
examples_per_page=25,
|
||||
)
|
||||
|
||||
gr.HTML(
|
||||
"""
|
||||
<div style="text-align: justify; max-width: 1100px; 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-2, 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-2, 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>
|
||||
<h3 style="font-weight: 450; font-size: 0.8rem; margin: 0rem">
|
||||
<b>Disclaimer</b>:
|
||||
Please note that the demo is intended for academic and research purposes <b>ONLY</b>. Any use of the demo for generating inappropriate content is strictly prohibited. The responsibility for any misuse or inappropriate use of the demo lies solely with the users who generated such content, and this demo shall not be held liable for any such use.
|
||||
</h3>
|
||||
</div>
|
||||
"""
|
||||
)
|
||||
|
||||
|
||||
demo.launch()
|
||||
@@ -0,0 +1,565 @@
|
||||
import os
|
||||
import re
|
||||
import zipfile
|
||||
import torch
|
||||
import gradio as gr
|
||||
|
||||
print('hello', gr.__version__)
|
||||
|
||||
|
||||
import time
|
||||
from transformers import CLIPTextModel, CLIPTokenizer
|
||||
from diffusers import AutoencoderKL, DDPMScheduler, StableDiffusionPipeline, UNet2DConditionModel, DiffusionPipeline, LCMScheduler
|
||||
from tqdm import tqdm
|
||||
from PIL import Image
|
||||
from PIL import Image, ImageDraw, ImageFont
|
||||
import random
|
||||
import copy
|
||||
|
||||
import string
|
||||
alphabet = string.digits + string.ascii_lowercase + string.ascii_uppercase + string.punctuation + ' ' # len(aphabet) = 95
|
||||
'''alphabet
|
||||
0123456789abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ!"#$%&\'()*+,-./:;<=>?@[\\]^_`{|}~
|
||||
'''
|
||||
|
||||
if not os.path.exists('images2'):
|
||||
os.system('wget https://huggingface.co/datasets/JingyeChen22/TextDiffuser/resolve/main/images2.zip')
|
||||
with zipfile.ZipFile('images2.zip', 'r') as zip_ref:
|
||||
zip_ref.extractall('.')
|
||||
|
||||
# os.system('nvidia-smi')
|
||||
os.system('ls')
|
||||
|
||||
#### import m1
|
||||
from fastchat.model import load_model, get_conversation_template
|
||||
from transformers import AutoTokenizer, AutoModelForCausalLM
|
||||
m1_model_path = 'JingyeChen22/textdiffuser2_layout_planner'
|
||||
# m1_model, m1_tokenizer = load_model(
|
||||
# m1_model_path,
|
||||
# 'cuda',
|
||||
# 1,
|
||||
# None,
|
||||
# False,
|
||||
# False,
|
||||
# revision="main",
|
||||
# debug=False,
|
||||
# )
|
||||
|
||||
m1_tokenizer = AutoTokenizer.from_pretrained(m1_model_path, use_fast=False)
|
||||
m1_model = AutoModelForCausalLM.from_pretrained(
|
||||
m1_model_path, torch_dtype=torch.float16, low_cpu_mem_usage=True
|
||||
).cuda()
|
||||
|
||||
#### import diffusion models
|
||||
text_encoder = CLIPTextModel.from_pretrained(
|
||||
'JingyeChen22/textdiffuser2-full-ft', subfolder="text_encoder"
|
||||
).cuda().half()
|
||||
tokenizer = CLIPTokenizer.from_pretrained(
|
||||
'runwayml/stable-diffusion-v1-5', subfolder="tokenizer"
|
||||
)
|
||||
|
||||
#### additional tokens are introduced, including coordinate tokens and character tokens
|
||||
print('***************')
|
||||
print(len(tokenizer))
|
||||
for i in range(520):
|
||||
tokenizer.add_tokens(['l' + str(i) ]) # left
|
||||
tokenizer.add_tokens(['t' + str(i) ]) # top
|
||||
tokenizer.add_tokens(['r' + str(i) ]) # width
|
||||
tokenizer.add_tokens(['b' + str(i) ]) # height
|
||||
for c in alphabet:
|
||||
tokenizer.add_tokens([f'[{c}]'])
|
||||
print(len(tokenizer))
|
||||
print('***************')
|
||||
|
||||
vae = AutoencoderKL.from_pretrained('runwayml/stable-diffusion-v1-5', subfolder="vae").half().cuda()
|
||||
unet = UNet2DConditionModel.from_pretrained(
|
||||
'JingyeChen22/textdiffuser2-full-ft', subfolder="unet"
|
||||
).half().cuda()
|
||||
text_encoder.resize_token_embeddings(len(tokenizer))
|
||||
|
||||
|
||||
#### load lcm components
|
||||
model_id = "lambdalabs/sd-pokemon-diffusers"
|
||||
lcm_lora_id = "latent-consistency/lcm-lora-sdv1-5"
|
||||
pipe = DiffusionPipeline.from_pretrained(model_id, unet=copy.deepcopy(unet), tokenizer=tokenizer, text_encoder=copy.deepcopy(text_encoder), torch_dtype=torch.float16)
|
||||
pipe.scheduler = LCMScheduler.from_config(pipe.scheduler.config)
|
||||
pipe.load_lora_weights(lcm_lora_id)
|
||||
pipe.to(device="cuda")
|
||||
|
||||
global_dict = {}
|
||||
#### for interactive
|
||||
# stack = []
|
||||
# state = 0
|
||||
font = ImageFont.truetype("./Arial.ttf", 32)
|
||||
|
||||
def skip_fun(i, t, guest_id):
|
||||
global_dict[guest_id]['state'] = 0
|
||||
# global state
|
||||
# state = 0
|
||||
|
||||
|
||||
def exe_undo(i, t, guest_id):
|
||||
|
||||
global_dict[guest_id]['stack'] = []
|
||||
global_dict[guest_id]['state'] = 0
|
||||
|
||||
# global stack
|
||||
# global state
|
||||
# state = 0
|
||||
# stack = []
|
||||
image = Image.open(f'./gray256.jpg')
|
||||
# print('stack', stack)
|
||||
return image
|
||||
|
||||
|
||||
def exe_redo(i, t, guest_id):
|
||||
# global state
|
||||
# state = 0
|
||||
global_dict[guest_id]['state'] = 0
|
||||
|
||||
if len(global_dict[guest_id]['stack']) > 0:
|
||||
global_dict[guest_id]['stack'].pop()
|
||||
image = Image.open(f'./gray256.jpg')
|
||||
draw = ImageDraw.Draw(image)
|
||||
|
||||
for items in global_dict[guest_id]['stack']:
|
||||
# print('now', items)
|
||||
text_position, t = items
|
||||
if len(text_position) == 2:
|
||||
x, y = text_position
|
||||
text_color = (255, 0, 0)
|
||||
draw.text((x+2, y), t, font=font, fill=text_color)
|
||||
r = 4
|
||||
leftUpPoint = (x-r, y-r)
|
||||
rightDownPoint = (x+r, y+r)
|
||||
draw.ellipse((leftUpPoint,rightDownPoint), fill='red')
|
||||
elif len(text_position) == 4:
|
||||
x0, y0, x1, y1 = text_position
|
||||
text_color = (255, 0, 0)
|
||||
draw.text((x0+2, y0), t, font=font, fill=text_color)
|
||||
r = 4
|
||||
leftUpPoint = (x0-r, y0-r)
|
||||
rightDownPoint = (x0+r, y0+r)
|
||||
draw.ellipse((leftUpPoint,rightDownPoint), fill='red')
|
||||
draw.rectangle((x0,y0,x1,y1), outline=(255, 0, 0) )
|
||||
|
||||
print('stack', global_dict[guest_id]['stack'])
|
||||
return image
|
||||
|
||||
def get_pixels(i, t, guest_id, evt: gr.SelectData):
|
||||
# global state
|
||||
|
||||
# register
|
||||
if guest_id == '-1':
|
||||
seed = str(int(time.time()))
|
||||
global_dict[str(seed)] = {
|
||||
'state': 0,
|
||||
'stack': []
|
||||
}
|
||||
guest_id = str(seed)
|
||||
else:
|
||||
seed = guest_id
|
||||
|
||||
text_position = evt.index
|
||||
|
||||
if global_dict[guest_id]['state'] == 0:
|
||||
global_dict[guest_id]['stack'].append(
|
||||
(text_position, t)
|
||||
)
|
||||
print(text_position, global_dict[guest_id]['stack'])
|
||||
global_dict[guest_id]['state'] = 1
|
||||
else:
|
||||
|
||||
(_, t) = global_dict[guest_id]['stack'].pop()
|
||||
x, y = _
|
||||
global_dict[guest_id]['stack'].append(
|
||||
((x,y,text_position[0],text_position[1]), t)
|
||||
)
|
||||
global_dict[guest_id]['state'] = 0
|
||||
|
||||
|
||||
image = Image.open(f'./gray256.jpg')
|
||||
draw = ImageDraw.Draw(image)
|
||||
|
||||
for items in global_dict[guest_id]['stack']:
|
||||
# print('now', items)
|
||||
text_position, t = items
|
||||
if len(text_position) == 2:
|
||||
x, y = text_position
|
||||
text_color = (255, 0, 0)
|
||||
draw.text((x+2, y), t, font=font, fill=text_color)
|
||||
r = 4
|
||||
leftUpPoint = (x-r, y-r)
|
||||
rightDownPoint = (x+r, y+r)
|
||||
draw.ellipse((leftUpPoint,rightDownPoint), fill='red')
|
||||
elif len(text_position) == 4:
|
||||
x0, y0, x1, y1 = text_position
|
||||
text_color = (255, 0, 0)
|
||||
draw.text((x0+2, y0), t, font=font, fill=text_color)
|
||||
r = 4
|
||||
leftUpPoint = (x0-r, y0-r)
|
||||
rightDownPoint = (x0+r, y0+r)
|
||||
draw.ellipse((leftUpPoint,rightDownPoint), fill='red')
|
||||
draw.rectangle((x0,y0,x1,y1), outline=(255, 0, 0) )
|
||||
|
||||
print('stack', global_dict[guest_id]['stack'])
|
||||
|
||||
return image, seed
|
||||
|
||||
|
||||
font_layout = ImageFont.truetype('./Arial.ttf', 16)
|
||||
|
||||
def get_layout_image(ocrs):
|
||||
|
||||
blank = Image.new('RGB', (256,256), (0,0,0))
|
||||
draw = ImageDraw.ImageDraw(blank)
|
||||
|
||||
for line in ocrs.split('\n'):
|
||||
line = line.strip()
|
||||
|
||||
if len(line) == 0:
|
||||
break
|
||||
|
||||
pred = ' '.join(line.split()[:-1])
|
||||
box = line.split()[-1]
|
||||
l, t, r, b = [int(i)*2 for i in box.split(',')] # the size of canvas is 256x256
|
||||
draw.rectangle([(l, t), (r, b)], outline ="red")
|
||||
draw.text((l, t), pred, font=font_layout)
|
||||
|
||||
return blank
|
||||
|
||||
|
||||
|
||||
def text_to_image(guest_id, prompt,keywords,positive_prompt,radio,slider_step,slider_guidance,slider_batch,slider_temperature,slider_natural):
|
||||
|
||||
print(f'[info] Prompt: {prompt} | Keywords: {keywords} | Radio: {radio} | Steps: {slider_step} | Guidance: {slider_guidance} | Natural: {slider_natural}')
|
||||
|
||||
# global stack
|
||||
# global state
|
||||
|
||||
if len(positive_prompt.strip()) != 0:
|
||||
prompt += positive_prompt
|
||||
|
||||
with torch.no_grad():
|
||||
time1 = time.time()
|
||||
user_prompt = prompt
|
||||
|
||||
if slider_natural:
|
||||
user_prompt = f'{user_prompt}'
|
||||
composed_prompt = user_prompt
|
||||
prompt = tokenizer.encode(user_prompt)
|
||||
layout_image = None
|
||||
else:
|
||||
if guest_id not in global_dict or len(global_dict[guest_id]['stack']) == 0:
|
||||
|
||||
if len(keywords.strip()) == 0:
|
||||
template = f'Given a prompt that will be used to generate an image, plan the layout of visual text for the image. The size of the image is 128x128. Therefore, all properties of the positions should not exceed 128, including the coordinates of top, left, right, and bottom. All keywords are included in the caption. You dont need to specify the details of font styles. At each line, the format should be keyword left, top, right, bottom. So let us begin. Prompt: {user_prompt}'
|
||||
else:
|
||||
keywords = keywords.split('/')
|
||||
keywords = [i.strip() for i in keywords]
|
||||
template = f'Given a prompt that will be used to generate an image, plan the layout of visual text for the image. The size of the image is 128x128. Therefore, all properties of the positions should not exceed 128, including the coordinates of top, left, right, and bottom. In addition, we also provide all keywords at random order for reference. You dont need to specify the details of font styles. At each line, the format should be keyword left, top, right, bottom. So let us begin. Prompt: {prompt}. Keywords: {str(keywords)}'
|
||||
|
||||
msg = template
|
||||
conv = get_conversation_template(m1_model_path)
|
||||
conv.append_message(conv.roles[0], msg)
|
||||
conv.append_message(conv.roles[1], None)
|
||||
prompt = conv.get_prompt()
|
||||
inputs = m1_tokenizer([prompt], return_token_type_ids=False)
|
||||
inputs = {k: torch.tensor(v).to('cuda') for k, v in inputs.items()}
|
||||
output_ids = m1_model.generate(
|
||||
**inputs,
|
||||
do_sample=True,
|
||||
temperature=slider_temperature,
|
||||
repetition_penalty=1.0,
|
||||
max_new_tokens=512,
|
||||
)
|
||||
|
||||
if m1_model.config.is_encoder_decoder:
|
||||
output_ids = output_ids[0]
|
||||
else:
|
||||
output_ids = output_ids[0][len(inputs["input_ids"][0]) :]
|
||||
outputs = m1_tokenizer.decode(
|
||||
output_ids, skip_special_tokens=True, spaces_between_special_tokens=False
|
||||
)
|
||||
print(f"[{conv.roles[0]}]\n{msg}")
|
||||
print(f"[{conv.roles[1]}]\n{outputs}")
|
||||
layout_image = get_layout_image(outputs)
|
||||
|
||||
ocrs = outputs.split('\n')
|
||||
time2 = time.time()
|
||||
print(time2-time1)
|
||||
|
||||
# user_prompt = prompt
|
||||
current_ocr = ocrs
|
||||
|
||||
|
||||
ocr_ids = []
|
||||
print('user_prompt', user_prompt)
|
||||
print('current_ocr', current_ocr)
|
||||
|
||||
|
||||
for ocr in current_ocr:
|
||||
ocr = ocr.strip()
|
||||
|
||||
if len(ocr) == 0 or '###' in ocr or '.com' in ocr:
|
||||
continue
|
||||
|
||||
items = ocr.split()
|
||||
pred = ' '.join(items[:-1])
|
||||
box = items[-1]
|
||||
|
||||
l,t,r,b = box.split(',')
|
||||
l,t,r,b = int(l), int(t), int(r), int(b)
|
||||
ocr_ids.extend(['l'+str(l), 't'+str(t), 'r'+str(r), 'b'+str(b)])
|
||||
|
||||
char_list = list(pred)
|
||||
char_list = [f'[{i}]' for i in char_list]
|
||||
ocr_ids.extend(char_list)
|
||||
ocr_ids.append(tokenizer.eos_token_id)
|
||||
|
||||
caption_ids = tokenizer(
|
||||
user_prompt, truncation=True, return_tensors="pt"
|
||||
).input_ids[0].tolist()
|
||||
|
||||
try:
|
||||
ocr_ids = tokenizer.encode(ocr_ids)
|
||||
prompt = caption_ids + ocr_ids
|
||||
except:
|
||||
prompt = caption_ids
|
||||
|
||||
user_prompt = tokenizer.decode(prompt)
|
||||
composed_prompt = tokenizer.decode(prompt)
|
||||
|
||||
else:
|
||||
user_prompt += ' <|endoftext|><|startoftext|>'
|
||||
layout_image = None
|
||||
|
||||
for items in global_dict[guest_id]['stack']:
|
||||
position, text = items
|
||||
|
||||
|
||||
if len(position) == 2:
|
||||
x, y = position
|
||||
x = x // 4
|
||||
y = y // 4
|
||||
text_str = ' '.join([f'[{c}]' for c in list(text)])
|
||||
user_prompt += f' l{x} t{y} {text_str} <|endoftext|>'
|
||||
elif len(position) == 4:
|
||||
x0, y0, x1, y1 = position
|
||||
x0 = x0 // 4
|
||||
y0 = y0 // 4
|
||||
x1 = x1 // 4
|
||||
y1 = y1 // 4
|
||||
text_str = ' '.join([f'[{c}]' for c in list(text)])
|
||||
user_prompt += f' l{x0} t{y0} r{x1} b{y1} {text_str} <|endoftext|>'
|
||||
|
||||
# composed_prompt = user_prompt
|
||||
prompt = tokenizer.encode(user_prompt)
|
||||
composed_prompt = tokenizer.decode(prompt)
|
||||
|
||||
prompt = prompt[:77]
|
||||
while len(prompt) < 77:
|
||||
prompt.append(tokenizer.pad_token_id)
|
||||
|
||||
if radio == 'TextDiffuser-2':
|
||||
|
||||
prompts_cond = prompt
|
||||
prompts_nocond = [tokenizer.pad_token_id]*77
|
||||
|
||||
prompts_cond = [prompts_cond] * slider_batch
|
||||
prompts_nocond = [prompts_nocond] * slider_batch
|
||||
|
||||
prompts_cond = torch.Tensor(prompts_cond).long().cuda()
|
||||
prompts_nocond = torch.Tensor(prompts_nocond).long().cuda()
|
||||
|
||||
scheduler = DDPMScheduler.from_pretrained('runwayml/stable-diffusion-v1-5', subfolder="scheduler")
|
||||
scheduler.set_timesteps(slider_step)
|
||||
noise = torch.randn((slider_batch, 4, 64, 64)).to("cuda").half()
|
||||
input = noise
|
||||
|
||||
encoder_hidden_states_cond = text_encoder(prompts_cond)[0].half()
|
||||
encoder_hidden_states_nocond = text_encoder(prompts_nocond)[0].half()
|
||||
|
||||
|
||||
for t in tqdm(scheduler.timesteps):
|
||||
with torch.no_grad(): # classifier free guidance
|
||||
noise_pred_cond = unet(sample=input, timestep=t, encoder_hidden_states=encoder_hidden_states_cond[:slider_batch]).sample # b, 4, 64, 64
|
||||
noise_pred_uncond = unet(sample=input, timestep=t, encoder_hidden_states=encoder_hidden_states_nocond[:slider_batch]).sample # b, 4, 64, 64
|
||||
noisy_residual = noise_pred_uncond + slider_guidance * (noise_pred_cond - noise_pred_uncond) # b, 4, 64, 64
|
||||
input = scheduler.step(noisy_residual, t, input).prev_sample
|
||||
del noise_pred_cond
|
||||
del noise_pred_uncond
|
||||
|
||||
torch.cuda.empty_cache()
|
||||
|
||||
# decode
|
||||
input = 1 / vae.config.scaling_factor * input
|
||||
images = vae.decode(input, return_dict=False)[0]
|
||||
width, height = 512, 512
|
||||
results = []
|
||||
new_image = Image.new('RGB', (2*width, 2*height))
|
||||
for index, image in enumerate(images.cpu().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')
|
||||
results.append(image)
|
||||
row = index // 2
|
||||
col = index % 2
|
||||
new_image.paste(image, (col*width, row*height))
|
||||
# os.system('nvidia-smi')
|
||||
torch.cuda.empty_cache()
|
||||
# os.system('nvidia-smi')
|
||||
return tuple(results), composed_prompt, layout_image
|
||||
|
||||
elif radio == 'TextDiffuser-2-LCM':
|
||||
generator = torch.Generator(device=pipe.device).manual_seed(random.randint(0,1000))
|
||||
image = pipe(
|
||||
prompt=user_prompt,
|
||||
generator=generator,
|
||||
# negative_prompt=negative_prompt,
|
||||
num_inference_steps=slider_step,
|
||||
guidance_scale=1,
|
||||
# num_images_per_prompt=slider_batch,
|
||||
).images
|
||||
# os.system('nvidia-smi')
|
||||
torch.cuda.empty_cache()
|
||||
# os.system('nvidia-smi')
|
||||
return tuple(image), composed_prompt, layout_image
|
||||
|
||||
with gr.Blocks() as demo:
|
||||
|
||||
|
||||
# guest_id = random.randint(0,100000000)
|
||||
# register
|
||||
|
||||
|
||||
gr.HTML(
|
||||
"""
|
||||
<div style="text-align: center; max-width: 1600px; margin: 20px auto;">
|
||||
<h2 style="font-weight: 900; font-size: 2.3rem; margin: 0rem">
|
||||
TextDiffuser-2: Unleashing the Power of Language Models for Text Rendering
|
||||
</h2>
|
||||
<h2 style="font-weight: 460; font-size: 1.1rem; margin: 0rem">
|
||||
<a href="https://jingyechen.github.io/">Jingye Chen</a>, <a href="https://hypjudy.github.io/website/">Yupan Huang</a>, <a href="https://scholar.google.com/citations?user=0LTZGhUAAAAJ&hl=en">Tengchao Lv</a>, <a href="https://www.microsoft.com/en-us/research/people/lecu/">Lei Cui</a>, <a href="https://cqf.io/">Qifeng Chen</a>, <a href="https://thegenerality.com/">Furu Wei</a>
|
||||
</h2>
|
||||
<h2 style="font-weight: 460; font-size: 1.1rem; margin: 0rem">
|
||||
HKUST, Sun Yat-sen University, Microsoft Research
|
||||
</h2>
|
||||
<h3 style="font-weight: 450; font-size: 1rem; margin: 0rem">
|
||||
[<a href="https://arxiv.org/abs/2311.16465" style="color:blue;">arXiv</a>]
|
||||
[<a href="https://github.com/microsoft/unilm/tree/master/textdiffuser-2" style="color:blue;">Code</a>]
|
||||
[<a href="https://jingyechen.github.io/textdiffuser2/" style="color:blue;">Project Page</a>]
|
||||
[<a href="https://discord.gg/q7eHPupu" style="color:purple;">Discord</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-2</b>, aiming at unleashing the power of language models for text rendering. Specifically, we <b>tame a language model into a layout planner</b> to transform user prompt into a layout using the caption-OCR pairs. The language model demonstrates flexibility and automation by inferring keywords from user prompts or incorporating user-specified keywords to determine their positions. Secondly, we <b>leverage the language model in the diffusion model as the layout encoder</b> to represent the position and content of text at the line level. This approach enables diffusion models to generate text images with broader diversity.
|
||||
</h2>
|
||||
<h2 style="text-align: left; font-weight: 450; font-size: 1rem; margin-top: 0.5rem; margin-bottom: 0.5rem">
|
||||
👀 <b>Tips for using this demo</b>: <b>(1)</b> Please carefully read the disclaimer in the below. Current verison can only support English. <b>(2)</b> The specification of keywords is optional. If provided, the language model will do its best to plan layouts using the given keywords. <b>(3)</b> If a template is given, the layout planner (M1) is not used. <b>(4)</b> Three operations, including redo, undo, and skip are provided. When using skip, only the left-top point of a keyword will be recorded, resulting in more diversity but sometimes decreasing the accuracy. <b>(5)</b> The layout planner can produce different layouts. You can increase the temperature to enhance the diversity. ✨ <b>(6)</b> We also provide the experimental demo combining <b>TextDiffuser-2</b> and <b>LCM</b>. The inference is fast using less sampling steps, although the precision in text rendering might decrease.
|
||||
</h2>
|
||||
<img src="https://raw.githubusercontent.com/JingyeChen/jingyechen.github.io/master/textdiffuser2/static/images/architecture_blank.jpg" alt="textdiffuser-2">
|
||||
</div>
|
||||
""")
|
||||
|
||||
with gr.Tab("Text-to-Image"):
|
||||
with gr.Row():
|
||||
with gr.Column(scale=1):
|
||||
prompt = gr.Textbox(label="Prompt. You can let language model automatically identify keywords, or provide them below", placeholder="A beautiful city skyline stamp of Shanghai")
|
||||
keywords = gr.Textbox(label="(Optional) Keywords. Should be seperated by / (e.g., keyword1/keyword2/...)", placeholder="keyword1/keyword2")
|
||||
positive_prompt = gr.Textbox(label="(Optional) Positive prompt", value=", digital art, very detailed, fantasy, high definition, cinematic light, dnd, trending on artstation")
|
||||
|
||||
# many encounter concurrent problem
|
||||
with gr.Accordion("(Optional) Template - Click to paint", open=False):
|
||||
with gr.Row():
|
||||
with gr.Column(scale=1):
|
||||
i = gr.Image(label="Canvas", type='filepath', value=f'./gray256.jpg', height=256, width=256)
|
||||
with gr.Column(scale=1):
|
||||
t = gr.Textbox(label="Keyword", value='input_keyword')
|
||||
redo = gr.Button(value='Redo - Cancel the last keyword')
|
||||
undo = gr.Button(value='Undo - Clear the canvas')
|
||||
skip_button = gr.Button(value='Skip - Operate the next keyword')
|
||||
|
||||
|
||||
radio = gr.Radio(["TextDiffuser-2", "TextDiffuser-2-LCM"], label="Choice of models", value="TextDiffuser-2")
|
||||
slider_natural = gr.Checkbox(label="Natural image generation", value=False, info="The text position and content info will not be incorporated.")
|
||||
slider_step = gr.Slider(minimum=1, maximum=50, value=20, step=1, label="Sampling step", info="The sampling step for TextDiffuser-2. You may decease the step to 4 when using LCM.")
|
||||
slider_guidance = gr.Slider(minimum=1, maximum=13, value=7.5, step=0.5, label="Scale of classifier-free guidance", info="The scale of cfg and is set to 7.5 in default. When using LCM, cfg is set to 1.")
|
||||
slider_batch = gr.Slider(minimum=1, maximum=6, value=4, step=1, label="Batch size", info="The number of images to be sampled.")
|
||||
slider_temperature = gr.Slider(minimum=0.1, maximum=2, value=1.4, step=0.1, label="Temperature", info="Control the diversity of layout planner. Higher value indicates more diversity.")
|
||||
# slider_seed = gr.Slider(minimum=1, maximum=10000, label="Seed", randomize=True)
|
||||
button = gr.Button("Generate")
|
||||
|
||||
|
||||
|
||||
guest_id_box = gr.Textbox(label="guest_id", value=f"-1")
|
||||
i.select(get_pixels,[i,t,guest_id_box],[i,guest_id_box])
|
||||
redo.click(exe_redo, [i,t,guest_id_box],[i])
|
||||
undo.click(exe_undo, [i,t,guest_id_box],[i])
|
||||
skip_button.click(skip_fun, [i,t,guest_id_box])
|
||||
|
||||
|
||||
with gr.Column(scale=1):
|
||||
output = gr.Gallery(label='Generated image')
|
||||
|
||||
with gr.Accordion("Intermediate results", open=False):
|
||||
gr.Markdown("Composed prompt")
|
||||
composed_prompt = gr.Textbox(label='')
|
||||
gr.Markdown("Layout visualization")
|
||||
layout = gr.Image(height=256, width=256)
|
||||
|
||||
|
||||
button.click(text_to_image, inputs=[guest_id_box, prompt,keywords,positive_prompt, radio,slider_step,slider_guidance,slider_batch,slider_temperature,slider_natural], outputs=[output, composed_prompt, layout])
|
||||
|
||||
gr.Markdown("## Prompt Examples")
|
||||
gr.Examples(
|
||||
[
|
||||
["A beautiful city skyline stamp of Shanghai", "", False],
|
||||
["The words 'KFC VIVO50' are inscribed upon the wall in a neon light effect", "KFC/VIVO50", False],
|
||||
["A logo of superman", "", False],
|
||||
["A pencil sketch of a tree with the title nothing to tree here", "", False],
|
||||
["handwritten signature of peter", "", False],
|
||||
["Delicate greeting card of happy birthday to xyz", "", False],
|
||||
["Book cover of good morning baby ", "", False],
|
||||
["The handwritten words Hello World displayed on a wall in a neon light effect", "", False],
|
||||
["Logo of winter in artistic font, made by snowflake", "", False],
|
||||
["A book cover named summer vibe", "", False],
|
||||
["Newspaper with the title Love Story", "", False],
|
||||
["A logo for the company EcoGrow, where the letters look like plants", "EcoGrow", False],
|
||||
["A poster titled 'Quails of North America', showing different kinds of quails.", "Quails/of/North/America", False],
|
||||
["A detailed portrait of a fox guardian with a shield with Kung Fu written on it, by victo ngai and justin gerard, digital art, realistic painting", "kung/fu", False],
|
||||
["A stamp of breath of the wild", "breath/of/the/wild", False],
|
||||
["Poster of the incoming movie Transformers", "Transformers", False],
|
||||
["Some apples are on a table", "", True],
|
||||
["a hotdog with mustard and other toppings on it", "", True],
|
||||
["a bathroom that has a slanted ceiling and a large bath tub", "", True],
|
||||
["a man holding a tennis racquet on a tennis court", "", True],
|
||||
["hamburger with bacon, lettuce, tomato and cheese| promotional image| hyperquality| products shot| full - color| extreme render| mouthwatering", "", True],
|
||||
],
|
||||
[
|
||||
prompt,
|
||||
keywords,
|
||||
slider_natural
|
||||
],
|
||||
examples_per_page=25
|
||||
)
|
||||
|
||||
gr.HTML(
|
||||
"""
|
||||
<div style="text-align: justify; max-width: 1100px; 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-2, 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-2, 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>
|
||||
<h3 style="font-weight: 450; font-size: 0.8rem; margin: 0rem">
|
||||
<b>Disclaimer</b>:
|
||||
Please note that the demo is intended for academic and research purposes <b>ONLY</b>. Any use of the demo for generating inappropriate content is strictly prohibited. The responsibility for any misuse or inappropriate use of the demo lies solely with the users who generated such content, and this demo shall not be held liable for any such use.
|
||||
</h3>
|
||||
</div>
|
||||
"""
|
||||
)
|
||||
|
||||
|
||||
demo.launch()
|
||||
@@ -0,0 +1,2 @@
|
||||
Book cover of summer vibe, high quality, high resolution
|
||||
Summer Vibe 20,20,100,40,20
|
||||
@@ -0,0 +1,643 @@
|
||||
# ------------------------------------------
|
||||
# TextDiffuser-2: Unleashing the Power of Language Models for Text Rendering
|
||||
# Paper Link: https://arxiv.org/abs/2311.16465
|
||||
# Code Link: https://github.com/microsoft/unilm/tree/master/textdiffuser-2
|
||||
# Copyright (c) Microsoft Corporation.
|
||||
# ------------------------------------------
|
||||
|
||||
import os
|
||||
import cv2
|
||||
import random
|
||||
import logging
|
||||
import argparse
|
||||
import numpy as np
|
||||
import time
|
||||
|
||||
from pathlib import Path
|
||||
from tqdm.auto import tqdm
|
||||
from typing import Optional
|
||||
from packaging import version
|
||||
from PIL import Image
|
||||
from huggingface_hub import HfFolder, Repository, create_repo, whoami
|
||||
|
||||
import string
|
||||
alphabet = string.digits + string.ascii_lowercase + string.ascii_uppercase + string.punctuation + ' ' # len(aphabet) = 95
|
||||
'''alphabet
|
||||
0123456789abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ!"#$%&\'()*+,-./:;<=>?@[\\]^_`{|}~
|
||||
'''
|
||||
|
||||
import datasets
|
||||
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
|
||||
|
||||
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(
|
||||
"--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(
|
||||
"--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."
|
||||
)
|
||||
#### newly added parameters
|
||||
parser.add_argument(
|
||||
"--granularity",
|
||||
type=int,
|
||||
default=128,
|
||||
help="The granularity of coordinates, ranging from 1~512."
|
||||
)
|
||||
parser.add_argument(
|
||||
"--coord_mode",
|
||||
type=str,
|
||||
default='lt',
|
||||
choices=['lt', 'center', 'ltrb'],
|
||||
help="The way to represent coordinates."
|
||||
)
|
||||
parser.add_argument(
|
||||
"--max_length",
|
||||
default=77,
|
||||
type=int,
|
||||
help="Maximum length of the composed prompt."
|
||||
)
|
||||
parser.add_argument(
|
||||
"--cfg",
|
||||
default=7,
|
||||
type=float,
|
||||
help="classifier free guidance."
|
||||
)
|
||||
parser.add_argument(
|
||||
"--sample_steps",
|
||||
default=50,
|
||||
type=int,
|
||||
help="steps for sampling for diffusion models."
|
||||
)
|
||||
parser.add_argument(
|
||||
"--input_format",
|
||||
required=True,
|
||||
type=str,
|
||||
help="specify the input format",
|
||||
choices=['prompt', 'prompts_txt_file', 'prompt_layout_txt_file']
|
||||
)
|
||||
parser.add_argument(
|
||||
"--input_prompt",
|
||||
type=str,
|
||||
)
|
||||
parser.add_argument(
|
||||
"--input_file",
|
||||
type=str,
|
||||
)
|
||||
parser.add_argument(
|
||||
"--prompts_txt_file",
|
||||
type=str,
|
||||
)
|
||||
parser.add_argument(
|
||||
"--m1_model_path",
|
||||
type=str,
|
||||
help="the checkpoint of layout planner"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--vis_num",
|
||||
type=int,
|
||||
default=16,
|
||||
help=("The number of images to be visualized."),
|
||||
)
|
||||
args = parser.parse_args()
|
||||
|
||||
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}"
|
||||
|
||||
|
||||
DATASET_NAME_MAPPING = {
|
||||
# "lambdalabs/pokemon-blip-captions": ("image", "text"),
|
||||
"MARIO-10M": ("image", "text"),
|
||||
}
|
||||
|
||||
def main():
|
||||
args = parse_args()
|
||||
logging_dir = Path(args.output_dir, args.logging_dir)
|
||||
|
||||
accelerator_project_config = ProjectConfiguration(project_dir=args.output_dir, logging_dir=logging_dir)
|
||||
|
||||
accelerator = Accelerator(
|
||||
gradient_accumulation_steps=4,
|
||||
mixed_precision=args.mixed_precision,
|
||||
log_with=args.report_to,
|
||||
project_config=accelerator_project_config,
|
||||
)
|
||||
if args.report_to == "wandb":
|
||||
if not is_wandb_available():
|
||||
raise ImportError("Make sure to install wandb if you want to use it for logging during training.")
|
||||
import wandb
|
||||
|
||||
# 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.output_dir is not None:
|
||||
os.makedirs(args.output_dir, exist_ok=True)
|
||||
|
||||
if args.push_to_hub:
|
||||
repo_id = create_repo(
|
||||
repo_id=args.hub_model_id or Path(args.output_dir).name, exist_ok=True, token=args.hub_token
|
||||
).repo_id
|
||||
# Load scheduler, tokenizer and models.
|
||||
tokenizer = CLIPTokenizer.from_pretrained(
|
||||
args.pretrained_model_name_or_path, subfolder="tokenizer"
|
||||
)
|
||||
|
||||
#### additional tokens are introduced, including coordinate tokens and character tokens
|
||||
print('***************')
|
||||
print(len(tokenizer))
|
||||
for i in range(520):
|
||||
tokenizer.add_tokens(['l' + str(i) ]) # left
|
||||
tokenizer.add_tokens(['t' + str(i) ]) # top
|
||||
tokenizer.add_tokens(['r' + str(i) ]) # width
|
||||
tokenizer.add_tokens(['b' + str(i) ]) # height
|
||||
|
||||
#### add rotation extension by appending angle tokens
|
||||
for i in range(-180,181):
|
||||
tokenizer.add_tokens(['a' + str(i) ]) # rotation
|
||||
|
||||
for c in alphabet:
|
||||
tokenizer.add_tokens([f'[{c}]'])
|
||||
print(len(tokenizer))
|
||||
print('***************')
|
||||
|
||||
if args.max_length == 77:
|
||||
text_encoder = CLIPTextModel.from_pretrained(
|
||||
args.resume_from_checkpoint, subfolder="text_encoder", ignore_mismatched_sizes=True
|
||||
)
|
||||
else:
|
||||
#### enlarge the context length of text encoder. empirically, enlarging the context length can proceed longer sequence. However, we observe that it will be hard to render general objects
|
||||
text_encoder = CLIPTextModel.from_pretrained(
|
||||
args.resume_from_checkpoint, subfolder="text_encoder", max_position_embeddings=args.max_length, ignore_mismatched_sizes=True
|
||||
)
|
||||
|
||||
text_encoder.resize_token_embeddings(len(tokenizer))
|
||||
|
||||
vae = AutoencoderKL.from_pretrained(args.pretrained_model_name_or_path, subfolder="vae")
|
||||
unet = UNet2DConditionModel.from_pretrained(
|
||||
args.resume_from_checkpoint, subfolder="unet"
|
||||
)
|
||||
# freeze parameters of models to save more memory
|
||||
# unet.requires_grad_(False)
|
||||
vae.requires_grad_(False)
|
||||
text_encoder.requires_grad_(False)
|
||||
|
||||
|
||||
|
||||
# `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):
|
||||
if args.use_ema:
|
||||
ema_unet.save_pretrained(os.path.join(output_dir, "unet_ema"))
|
||||
|
||||
for i, model in enumerate(models):
|
||||
# model.save_pretrained(os.path.join(output_dir, "unet"))
|
||||
|
||||
if i == 0:
|
||||
model.save_pretrained(os.path.join(output_dir, f"unet"))
|
||||
elif i == 1:
|
||||
model.save_pretrained(os.path.join(output_dir, f"text_encoder"))
|
||||
|
||||
# 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()
|
||||
|
||||
if i == 1:
|
||||
load_model = UNet2DConditionModel.from_pretrained(input_dir, subfolder="unet")
|
||||
model.register_to_config(**load_model.config)
|
||||
elif i == 0:
|
||||
load_model = CLIPTextModel.from_pretrained(input_dir, subfolder="text_encoder")
|
||||
# model.register_to_config(**load_model.config)
|
||||
|
||||
# # 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)
|
||||
|
||||
# For mixed precision training we cast all non-trainable weigths (vae, non-lora text_encoder and non-lora unet) to half-precision
|
||||
# as these weights are only used for inference, keeping weights in full precision is not required.
|
||||
weight_dtype = torch.float32
|
||||
if accelerator.mixed_precision == "fp16":
|
||||
weight_dtype = torch.float16
|
||||
elif accelerator.mixed_precision == "bf16":
|
||||
weight_dtype = torch.bfloat16
|
||||
|
||||
# Move unet, vae and text_encoder to device and cast to weight_dtype
|
||||
unet.to(accelerator.device, dtype=weight_dtype)
|
||||
vae.to(accelerator.device, dtype=weight_dtype)
|
||||
text_encoder.to(accelerator.device, dtype=weight_dtype)
|
||||
|
||||
|
||||
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")
|
||||
|
||||
|
||||
# # Enable TF32 for faster training on Ampere GPUs,
|
||||
# # cf https://pytorch.org/docs/stable/notes/cuda.html#tensorfloat-32-tf32-on-ampere-devices
|
||||
# if True:
|
||||
# torch.backends.cuda.matmul.allow_tf32 = True
|
||||
|
||||
# # We need to initialize the trackers we use, and also store our configuration.
|
||||
# # The trackers initializes automatically on the main process.
|
||||
# if accelerator.is_main_process:
|
||||
# accelerator.init_trackers("text2image-fine-tune", config=vars(args))
|
||||
|
||||
# Potentially load in the weights and states from a previous save
|
||||
if args.resume_from_checkpoint:
|
||||
if args.resume_from_checkpoint != "latest":
|
||||
path = os.path.basename(args.resume_from_checkpoint)
|
||||
else:
|
||||
# Get the most recent checkpoint
|
||||
dirs = os.listdir(args.output_dir)
|
||||
dirs = [d for d in dirs if d.startswith("checkpoint")]
|
||||
dirs = sorted(dirs, key=lambda x: int(x.split("-")[1]))
|
||||
path = dirs[-1] if len(dirs) > 0 else None
|
||||
|
||||
if path is None:
|
||||
accelerator.print(
|
||||
f"Checkpoint '{args.resume_from_checkpoint}' does not exist. Starting a new training run."
|
||||
)
|
||||
args.resume_from_checkpoint = None
|
||||
else:
|
||||
accelerator.print(f"Resuming from checkpoint {path}")
|
||||
# accelerator.load_state(os.path.join(args.output_dir, path))
|
||||
accelerator.load_state(args.resume_from_checkpoint)
|
||||
|
||||
if accelerator.is_main_process and os.path.exists(f'{args.output_dir}'):
|
||||
print('detect existing output_dir, removing the contained jpg/txt files ...')
|
||||
os.system(f'rm {args.output_dir}/*.jpg')
|
||||
os.system(f'rm {args.output_dir}/*.txt')
|
||||
|
||||
# user_prompt = "Book cover of summer vibe, high quality, high resolution"
|
||||
# ocrs = [
|
||||
# 'Summer Vibe 20,20,100,40'
|
||||
# ]
|
||||
|
||||
if args.input_format == 'prompt_layout_txt_file':
|
||||
lines = open(args.input_file).readlines()
|
||||
user_prompts = [lines[0].strip()]
|
||||
ocrs = [lines[1:]]
|
||||
|
||||
elif args.input_format == 'prompt' or args.input_format == 'prompts_txt_file':
|
||||
|
||||
#### prepare m1 (layout planner)
|
||||
from fastchat.model import load_model, get_conversation_template
|
||||
m1_model, m1_tokenizer = load_model(
|
||||
args.m1_model_path,
|
||||
'cuda',
|
||||
1,
|
||||
None,
|
||||
False,
|
||||
False,
|
||||
revision="main",
|
||||
debug=False,
|
||||
)
|
||||
|
||||
# prompt = 'a text image of hello world'
|
||||
prompts = []
|
||||
if args.input_format == 'prompt':
|
||||
prompts = [args.input_prompt]
|
||||
elif args.input_format == 'prompts_txt_file':
|
||||
prompts = open(args.prompts_txt_file).readlines()
|
||||
print(f'there are {len(prompts)} samples for generation')
|
||||
|
||||
ocrs = []
|
||||
user_prompts = []
|
||||
for prompt in prompts:
|
||||
user_prompt = prompt
|
||||
user_prompts.append(user_prompt)
|
||||
template = f'Given a prompt that will be used to generate an image, plan the layout of visual text for the image. The size of the image is 128x128. Therefore, all properties of the positions should not exceed 128, including the coordinates of top, left, right, and bottom. All keywords are included in the caption. You dont need to specify the details of font styles. At each line, the format should be keyword left, top, right, bottom. So let us begin. Prompt: {user_prompt}'
|
||||
msg = template
|
||||
conv = get_conversation_template(args.m1_model_path)
|
||||
conv.append_message(conv.roles[0], msg)
|
||||
conv.append_message(conv.roles[1], None)
|
||||
prompt = conv.get_prompt()
|
||||
inputs = m1_tokenizer([prompt], return_token_type_ids=False)
|
||||
inputs = {k: torch.tensor(v).to('cuda') for k, v in inputs.items()}
|
||||
output_ids = m1_model.generate(
|
||||
**inputs,
|
||||
do_sample=True,
|
||||
temperature=0.7,
|
||||
repetition_penalty=1.0,
|
||||
max_new_tokens=512,
|
||||
)
|
||||
|
||||
if m1_model.config.is_encoder_decoder:
|
||||
output_ids = output_ids[0]
|
||||
else:
|
||||
output_ids = output_ids[0][len(inputs["input_ids"][0]) :]
|
||||
outputs = m1_tokenizer.decode(
|
||||
output_ids, skip_special_tokens=True, spaces_between_special_tokens=False
|
||||
)
|
||||
print(f"[{conv.roles[0]}]\n{msg}")
|
||||
print(f"[{conv.roles[1]}]\n{outputs}")
|
||||
# ocrs = outputs.split('\n')
|
||||
ocrs.append(outputs.split('\n'))
|
||||
|
||||
|
||||
with torch.no_grad():
|
||||
|
||||
size = len(ocrs)
|
||||
print(f'the number of samples: {size}')
|
||||
|
||||
time_seed = int(time.time())
|
||||
random.seed(time_seed)
|
||||
torch.manual_seed(time_seed)
|
||||
torch.cuda.manual_seed_all(time_seed)
|
||||
|
||||
for sample_index in range(size):
|
||||
|
||||
user_prompt = user_prompts[sample_index]
|
||||
current_ocr = ocrs[sample_index]
|
||||
|
||||
ocr_ids = []
|
||||
print('user_prompt', user_prompt)
|
||||
print('current_ocr', current_ocr)
|
||||
|
||||
current_ocr = []
|
||||
for ocr in current_ocr:
|
||||
ocr = ocr.strip()
|
||||
|
||||
if len(ocr) == 0 or '###' in ocr or '.com' in ocr:
|
||||
continue
|
||||
|
||||
items = ocr.split()
|
||||
pred = ' '.join(items[:-1])
|
||||
box = items[-1]
|
||||
|
||||
l,t,r,b,a = box.split(',')
|
||||
l,t,r,b,a = int(l), int(t), int(r), int(b), int(a)
|
||||
ocr_ids.extend(['l'+str(l), 't'+str(t), 'r'+str(r), 'b'+str(b), 'a'+str(a)])
|
||||
|
||||
char_list = list(pred)
|
||||
char_list = [f'[{i}]' for i in char_list]
|
||||
ocr_ids.extend(char_list)
|
||||
ocr_ids.append(tokenizer.eos_token_id)
|
||||
|
||||
caption_ids = tokenizer(
|
||||
user_prompt, truncation=True, return_tensors="pt"
|
||||
).input_ids[0].tolist()
|
||||
|
||||
try:
|
||||
ocr_ids = tokenizer.encode(ocr_ids)
|
||||
prompt = caption_ids + ocr_ids
|
||||
except:
|
||||
prompt = caption_ids
|
||||
|
||||
prompt = prompt[:args.max_length]
|
||||
while len(prompt) < args.max_length:
|
||||
prompt.append(tokenizer.pad_token_id)
|
||||
prompts_cond = prompt
|
||||
prompts_nocond = [tokenizer.pad_token_id]*args.max_length
|
||||
|
||||
prompts_cond = [prompts_cond] * args.vis_num
|
||||
prompts_nocond = [prompts_nocond] * args.vis_num
|
||||
|
||||
prompts_cond = torch.Tensor(prompts_cond).long().cuda()
|
||||
prompts_nocond = torch.Tensor(prompts_nocond).long().cuda()
|
||||
|
||||
scheduler = DDPMScheduler.from_pretrained(args.pretrained_model_name_or_path, subfolder="scheduler")
|
||||
scheduler.set_timesteps(args.sample_steps)
|
||||
noise = torch.randn((args.vis_num, 4, 64, 64)).to("cuda")
|
||||
input = noise
|
||||
|
||||
encoder_hidden_states_cond = text_encoder(prompts_cond)[0]
|
||||
encoder_hidden_states_nocond = text_encoder(prompts_nocond)[0]
|
||||
|
||||
texts = prompts_cond
|
||||
f = open(f'{args.output_dir}/prompt_{sample_index}_{args.local_rank}.txt', 'w+')
|
||||
for text in texts:
|
||||
sentence = tokenizer.decode(text)
|
||||
f.write(sentence + '\n')
|
||||
f.close()
|
||||
|
||||
for t in tqdm(scheduler.timesteps):
|
||||
with torch.no_grad(): # classifier free guidance
|
||||
noise_pred_cond = unet(sample=input.half(), timestep=t, encoder_hidden_states=encoder_hidden_states_cond[:args.vis_num]).sample # b, 4, 64, 64
|
||||
noise_pred_uncond = unet(sample=input.half(), timestep=t, encoder_hidden_states=encoder_hidden_states_nocond[:args.vis_num]).sample # b, 4, 64, 64
|
||||
noisy_residual = noise_pred_uncond + args.cfg * (noise_pred_cond - noise_pred_uncond) # b, 4, 64, 64
|
||||
input = scheduler.step(noisy_residual, t, input).prev_sample
|
||||
# input = prev_noisy_sample
|
||||
|
||||
# decode
|
||||
input = 1 / vae.config.scaling_factor * input
|
||||
images = vae.decode(input.half(), return_dict=False)[0]
|
||||
width, height = 512, 512
|
||||
new_image = Image.new('RGB', (4*width, 4*height))
|
||||
for index, image in enumerate(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')
|
||||
row = index // 4
|
||||
col = index % 4
|
||||
new_image.paste(image, (col*width, row*height))
|
||||
new_image.save(f'{args.output_dir}/pred_img_{sample_index}_{args.local_rank}.jpg')
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,14 @@
|
||||
accelerate launch inference_textdiffuser2_t2i_full_angle.py \
|
||||
--pretrained_model_name_or_path="runwayml/stable-diffusion-v1-5" \
|
||||
--mixed_precision="fp16" \
|
||||
--output_dir="inference_results" \
|
||||
--enable_xformers_memory_efficient_attention \
|
||||
--resume_from_checkpoint="put your ckpt here" \
|
||||
--granularity=128 \
|
||||
--max_length=77 \
|
||||
--coord_mode="ltrb" \
|
||||
--cfg=7.5 \
|
||||
--sample_steps=20 \
|
||||
--seed=43555 \
|
||||
--input_format='prompt_layout_txt_file' \
|
||||
--input_file='angle_template_file.txt'
|
||||
@@ -0,0 +1,651 @@
|
||||
# ------------------------------------------
|
||||
# TextDiffuser-2: Unleashing the Power of Language Models for Text Rendering
|
||||
# Paper Link: https://arxiv.org/abs/2311.16465
|
||||
# Code Link: https://github.com/microsoft/unilm/tree/master/textdiffuser-2
|
||||
# Copyright (c) Microsoft Corporation.
|
||||
# ------------------------------------------
|
||||
|
||||
import os
|
||||
import cv2
|
||||
import random
|
||||
import logging
|
||||
import argparse
|
||||
import numpy as np
|
||||
import time
|
||||
|
||||
from pathlib import Path
|
||||
from tqdm.auto import tqdm
|
||||
from typing import Optional
|
||||
from packaging import version
|
||||
from PIL import Image
|
||||
from huggingface_hub import HfFolder, Repository, create_repo, whoami
|
||||
|
||||
import string
|
||||
alphabet = string.digits + string.ascii_lowercase + string.ascii_uppercase + string.punctuation + ' ' # len(aphabet) = 95
|
||||
'''alphabet
|
||||
0123456789abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ!"#$%&\'()*+,-./:;<=>?@[\\]^_`{|}~
|
||||
'''
|
||||
|
||||
import datasets
|
||||
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
|
||||
|
||||
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(
|
||||
"--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(
|
||||
"--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."
|
||||
)
|
||||
#### newly added parameters
|
||||
parser.add_argument(
|
||||
"--granularity",
|
||||
type=int,
|
||||
default=128,
|
||||
help="The granularity of coordinates, ranging from 1~512."
|
||||
)
|
||||
parser.add_argument(
|
||||
"--coord_mode",
|
||||
type=str,
|
||||
default='lt',
|
||||
choices=['lt', 'center', 'ltrb'],
|
||||
help="The way to represent coordinates."
|
||||
)
|
||||
parser.add_argument(
|
||||
"--max_length",
|
||||
default=77,
|
||||
type=int,
|
||||
help="Maximum length of the composed prompt."
|
||||
)
|
||||
parser.add_argument(
|
||||
"--cfg",
|
||||
default=7,
|
||||
type=float,
|
||||
help="classifier free guidance."
|
||||
)
|
||||
parser.add_argument(
|
||||
"--sample_steps",
|
||||
default=50,
|
||||
type=int,
|
||||
help="steps for sampling for diffusion models."
|
||||
)
|
||||
parser.add_argument(
|
||||
"--input_format",
|
||||
required=True,
|
||||
type=str,
|
||||
help="specify the input format",
|
||||
choices=['prompt', 'prompts_txt_file', 'prompt_layout_txt_file']
|
||||
)
|
||||
parser.add_argument(
|
||||
"--input_prompt",
|
||||
type=str,
|
||||
)
|
||||
parser.add_argument(
|
||||
"--input_file",
|
||||
type=str,
|
||||
)
|
||||
parser.add_argument(
|
||||
"--prompts_txt_file",
|
||||
type=str,
|
||||
)
|
||||
parser.add_argument(
|
||||
"--m1_model_path",
|
||||
type=str,
|
||||
help="the checkpoint of layout planner"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--vis_num",
|
||||
type=int,
|
||||
default=16,
|
||||
help=("The number of images to be visualized."),
|
||||
)
|
||||
args = parser.parse_args()
|
||||
|
||||
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}"
|
||||
|
||||
|
||||
DATASET_NAME_MAPPING = {
|
||||
# "lambdalabs/pokemon-blip-captions": ("image", "text"),
|
||||
"MARIO-10M": ("image", "text"),
|
||||
}
|
||||
|
||||
def main():
|
||||
args = parse_args()
|
||||
logging_dir = Path(args.output_dir, args.logging_dir)
|
||||
|
||||
accelerator_project_config = ProjectConfiguration(project_dir=args.output_dir, logging_dir=logging_dir)
|
||||
|
||||
accelerator = Accelerator(
|
||||
gradient_accumulation_steps=4,
|
||||
mixed_precision=args.mixed_precision,
|
||||
log_with=args.report_to,
|
||||
project_config=accelerator_project_config,
|
||||
)
|
||||
if args.report_to == "wandb":
|
||||
if not is_wandb_available():
|
||||
raise ImportError("Make sure to install wandb if you want to use it for logging during training.")
|
||||
import wandb
|
||||
|
||||
# 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.output_dir is not None:
|
||||
os.makedirs(args.output_dir, exist_ok=True)
|
||||
|
||||
if args.push_to_hub:
|
||||
repo_id = create_repo(
|
||||
repo_id=args.hub_model_id or Path(args.output_dir).name, exist_ok=True, token=args.hub_token
|
||||
).repo_id
|
||||
# Load scheduler, tokenizer and models.
|
||||
tokenizer = CLIPTokenizer.from_pretrained(
|
||||
args.pretrained_model_name_or_path, subfolder="tokenizer"
|
||||
)
|
||||
|
||||
#### additional tokens are introduced, including coordinate tokens and character tokens
|
||||
print('***************')
|
||||
print(len(tokenizer))
|
||||
for i in range(520):
|
||||
tokenizer.add_tokens(['l' + str(i) ]) # left
|
||||
tokenizer.add_tokens(['t' + str(i) ]) # top
|
||||
tokenizer.add_tokens(['r' + str(i) ]) # width
|
||||
tokenizer.add_tokens(['b' + str(i) ]) # height
|
||||
|
||||
for c in alphabet:
|
||||
tokenizer.add_tokens([f'[{c}]'])
|
||||
print(len(tokenizer))
|
||||
print('***************')
|
||||
|
||||
if args.max_length == 77:
|
||||
text_encoder = CLIPTextModel.from_pretrained(
|
||||
args.resume_from_checkpoint, subfolder="text_encoder", ignore_mismatched_sizes=True
|
||||
)
|
||||
else:
|
||||
#### enlarge the context length of text encoder. empirically, enlarging the context length can proceed longer sequence. However, we observe that it will be hard to render general objects
|
||||
text_encoder = CLIPTextModel.from_pretrained(
|
||||
args.resume_from_checkpoint, subfolder="text_encoder", max_position_embeddings=args.max_length, ignore_mismatched_sizes=True
|
||||
)
|
||||
|
||||
text_encoder.resize_token_embeddings(len(tokenizer))
|
||||
|
||||
vae = AutoencoderKL.from_pretrained(args.pretrained_model_name_or_path, subfolder="vae")
|
||||
unet = UNet2DConditionModel.from_pretrained(
|
||||
args.resume_from_checkpoint, subfolder="unet"
|
||||
)
|
||||
# freeze parameters of models to save more memory
|
||||
# unet.requires_grad_(False)
|
||||
vae.requires_grad_(False)
|
||||
text_encoder.requires_grad_(False)
|
||||
|
||||
|
||||
|
||||
# `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):
|
||||
if args.use_ema:
|
||||
ema_unet.save_pretrained(os.path.join(output_dir, "unet_ema"))
|
||||
|
||||
for i, model in enumerate(models):
|
||||
# model.save_pretrained(os.path.join(output_dir, "unet"))
|
||||
|
||||
if i == 0:
|
||||
model.save_pretrained(os.path.join(output_dir, f"unet"))
|
||||
elif i == 1:
|
||||
model.save_pretrained(os.path.join(output_dir, f"text_encoder"))
|
||||
|
||||
# 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()
|
||||
|
||||
if i == 1:
|
||||
load_model = UNet2DConditionModel.from_pretrained(input_dir, subfolder="unet")
|
||||
model.register_to_config(**load_model.config)
|
||||
elif i == 0:
|
||||
load_model = CLIPTextModel.from_pretrained(input_dir, subfolder="text_encoder")
|
||||
# model.register_to_config(**load_model.config)
|
||||
|
||||
# # 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)
|
||||
|
||||
# For mixed precision training we cast all non-trainable weigths (vae, non-lora text_encoder and non-lora unet) to half-precision
|
||||
# as these weights are only used for inference, keeping weights in full precision is not required.
|
||||
weight_dtype = torch.float32
|
||||
if accelerator.mixed_precision == "fp16":
|
||||
weight_dtype = torch.float16
|
||||
elif accelerator.mixed_precision == "bf16":
|
||||
weight_dtype = torch.bfloat16
|
||||
|
||||
# Move unet, vae and text_encoder to device and cast to weight_dtype
|
||||
unet.to(accelerator.device, dtype=weight_dtype)
|
||||
vae.to(accelerator.device, dtype=weight_dtype)
|
||||
text_encoder.to(accelerator.device, dtype=weight_dtype)
|
||||
|
||||
|
||||
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")
|
||||
|
||||
|
||||
# # Enable TF32 for faster training on Ampere GPUs,
|
||||
# # cf https://pytorch.org/docs/stable/notes/cuda.html#tensorfloat-32-tf32-on-ampere-devices
|
||||
# if True:
|
||||
# torch.backends.cuda.matmul.allow_tf32 = True
|
||||
|
||||
# # We need to initialize the trackers we use, and also store our configuration.
|
||||
# # The trackers initializes automatically on the main process.
|
||||
# if accelerator.is_main_process:
|
||||
# accelerator.init_trackers("text2image-fine-tune", config=vars(args))
|
||||
|
||||
# Potentially load in the weights and states from a previous save
|
||||
if args.resume_from_checkpoint:
|
||||
if args.resume_from_checkpoint != "latest":
|
||||
path = os.path.basename(args.resume_from_checkpoint)
|
||||
else:
|
||||
# Get the most recent checkpoint
|
||||
dirs = os.listdir(args.output_dir)
|
||||
dirs = [d for d in dirs if d.startswith("checkpoint")]
|
||||
dirs = sorted(dirs, key=lambda x: int(x.split("-")[1]))
|
||||
path = dirs[-1] if len(dirs) > 0 else None
|
||||
|
||||
if path is None:
|
||||
accelerator.print(
|
||||
f"Checkpoint '{args.resume_from_checkpoint}' does not exist. Starting a new training run."
|
||||
)
|
||||
args.resume_from_checkpoint = None
|
||||
else:
|
||||
accelerator.print(f"Resuming from checkpoint {path}")
|
||||
# accelerator.load_state(os.path.join(args.output_dir, path))
|
||||
accelerator.load_state(args.resume_from_checkpoint)
|
||||
|
||||
if accelerator.is_main_process and os.path.exists(f'{args.output_dir}'):
|
||||
print('detect existing output_dir, removing the contained jpg/txt files ...')
|
||||
os.system(f'rm {args.output_dir}/*.jpg')
|
||||
os.system(f'rm {args.output_dir}/*.txt')
|
||||
|
||||
# user_prompt = "Book cover of summer vibe, high quality, high resolution"
|
||||
# ocrs = [
|
||||
# 'Summer Vibe 20,20,100,40'
|
||||
# ]
|
||||
|
||||
if args.input_format == 'prompt_layout_txt_file':
|
||||
lines = open(args.input_file).readlines()
|
||||
user_prompts = [lines[0].strip()]
|
||||
ocrs = [lines[1:]]
|
||||
|
||||
elif args.input_format == 'prompt' or args.input_format == 'prompts_txt_file':
|
||||
|
||||
#### prepare m1 (layout planner)
|
||||
from fastchat.model import load_model, get_conversation_template
|
||||
m1_model, m1_tokenizer = load_model(
|
||||
args.m1_model_path,
|
||||
'cuda',
|
||||
1,
|
||||
None,
|
||||
False,
|
||||
False,
|
||||
revision="main",
|
||||
debug=False,
|
||||
)
|
||||
|
||||
# prompt = 'a text image of hello world'
|
||||
prompts = []
|
||||
if args.input_format == 'prompt':
|
||||
prompts = [args.input_prompt]
|
||||
elif args.input_format == 'prompts_txt_file':
|
||||
prompts = open(args.prompts_txt_file).readlines()
|
||||
print(f'there are {len(prompts)} samples for generation')
|
||||
|
||||
ocrs = []
|
||||
user_prompts = []
|
||||
for prompt in prompts:
|
||||
user_prompt = prompt
|
||||
user_prompts.append(user_prompt)
|
||||
template = f'Given a prompt that will be used to generate an image, plan the layout of visual text for the image. The size of the image is 128x128. Therefore, all properties of the positions should not exceed 128, including the coordinates of top, left, right, and bottom. All keywords are included in the caption. You dont need to specify the details of font styles. At each line, the format should be keyword left, top, right, bottom. So let us begin. Prompt: {user_prompt}'
|
||||
msg = template
|
||||
conv = get_conversation_template(args.m1_model_path)
|
||||
conv.append_message(conv.roles[0], msg)
|
||||
conv.append_message(conv.roles[1], None)
|
||||
prompt = conv.get_prompt()
|
||||
inputs = m1_tokenizer([prompt], return_token_type_ids=False)
|
||||
inputs = {k: torch.tensor(v).to('cuda') for k, v in inputs.items()}
|
||||
output_ids = m1_model.generate(
|
||||
**inputs,
|
||||
do_sample=True,
|
||||
temperature=0.7,
|
||||
repetition_penalty=1.0,
|
||||
max_new_tokens=512,
|
||||
)
|
||||
|
||||
if m1_model.config.is_encoder_decoder:
|
||||
output_ids = output_ids[0]
|
||||
else:
|
||||
output_ids = output_ids[0][len(inputs["input_ids"][0]) :]
|
||||
outputs = m1_tokenizer.decode(
|
||||
output_ids, skip_special_tokens=True, spaces_between_special_tokens=False
|
||||
)
|
||||
print(f"[{conv.roles[0]}]\n{msg}")
|
||||
print(f"[{conv.roles[1]}]\n{outputs}")
|
||||
# ocrs = outputs.split('\n')
|
||||
ocrs.append(outputs.split('\n'))
|
||||
|
||||
|
||||
with torch.no_grad():
|
||||
|
||||
size = len(ocrs)
|
||||
print(f'the number of samples: {size}')
|
||||
|
||||
time_seed = int(time.time())
|
||||
random.seed(time_seed)
|
||||
torch.manual_seed(time_seed)
|
||||
torch.cuda.manual_seed_all(time_seed)
|
||||
|
||||
for sample_index in range(size):
|
||||
|
||||
user_prompt = user_prompts[sample_index]
|
||||
current_ocr = ocrs[sample_index]
|
||||
|
||||
ocr_ids = []
|
||||
print('user_prompt', user_prompt)
|
||||
print('current_ocr', current_ocr)
|
||||
|
||||
current_ocr = []
|
||||
for ocr in current_ocr:
|
||||
ocr = ocr.strip()
|
||||
|
||||
if len(ocr) == 0 or '###' in ocr or '.com' in ocr:
|
||||
continue
|
||||
|
||||
items = ocr.split()
|
||||
pred = ' '.join(items[:-1])
|
||||
box = items[-1]
|
||||
|
||||
# l,t,r,b,a = box.split(',')
|
||||
# l,t,r,b,a = int(l), int(t), int(r), int(b), int(a)
|
||||
# ocr_ids.extend(['l'+str(l), 't'+str(t), 'r'+str(r), 'b'+str(b), 'a'+str(a)])
|
||||
box = box.split(',')
|
||||
box = [int(i) for i in box]
|
||||
x1, y1, x2, y2, x3, y3, x4, y4 = box
|
||||
# ocr_ids.extend(['l'+str(l), 't'+str(t), 'r'+str(r), 'b'+str(b), 'a'+str(a)])
|
||||
ocr_ids.extend(
|
||||
[
|
||||
'l'+str(x1), 't'+str(y1),
|
||||
'l'+str(x2), 't'+str(y2),
|
||||
'l'+str(x3), 't'+str(y3),
|
||||
'l'+str(x4), 't'+str(y4),
|
||||
]
|
||||
)
|
||||
|
||||
char_list = list(pred)
|
||||
char_list = [f'[{i}]' for i in char_list]
|
||||
ocr_ids.extend(char_list)
|
||||
ocr_ids.append(tokenizer.eos_token_id)
|
||||
|
||||
caption_ids = tokenizer(
|
||||
user_prompt, truncation=True, return_tensors="pt"
|
||||
).input_ids[0].tolist()
|
||||
|
||||
try:
|
||||
ocr_ids = tokenizer.encode(ocr_ids)
|
||||
prompt = caption_ids + ocr_ids
|
||||
except:
|
||||
prompt = caption_ids
|
||||
|
||||
prompt = prompt[:args.max_length]
|
||||
while len(prompt) < args.max_length:
|
||||
prompt.append(tokenizer.pad_token_id)
|
||||
prompts_cond = prompt
|
||||
prompts_nocond = [tokenizer.pad_token_id]*args.max_length
|
||||
|
||||
prompts_cond = [prompts_cond] * args.vis_num
|
||||
prompts_nocond = [prompts_nocond] * args.vis_num
|
||||
|
||||
prompts_cond = torch.Tensor(prompts_cond).long().cuda()
|
||||
prompts_nocond = torch.Tensor(prompts_nocond).long().cuda()
|
||||
|
||||
scheduler = DDPMScheduler.from_pretrained(args.pretrained_model_name_or_path, subfolder="scheduler")
|
||||
scheduler.set_timesteps(args.sample_steps)
|
||||
noise = torch.randn((args.vis_num, 4, 64, 64)).to("cuda")
|
||||
input = noise
|
||||
|
||||
encoder_hidden_states_cond = text_encoder(prompts_cond)[0]
|
||||
encoder_hidden_states_nocond = text_encoder(prompts_nocond)[0]
|
||||
|
||||
texts = prompts_cond
|
||||
f = open(f'{args.output_dir}/prompt_{sample_index}_{args.local_rank}.txt', 'w+')
|
||||
for text in texts:
|
||||
sentence = tokenizer.decode(text)
|
||||
f.write(sentence + '\n')
|
||||
f.close()
|
||||
|
||||
for t in tqdm(scheduler.timesteps):
|
||||
with torch.no_grad(): # classifier free guidance
|
||||
noise_pred_cond = unet(sample=input.half(), timestep=t, encoder_hidden_states=encoder_hidden_states_cond[:args.vis_num]).sample # b, 4, 64, 64
|
||||
noise_pred_uncond = unet(sample=input.half(), timestep=t, encoder_hidden_states=encoder_hidden_states_nocond[:args.vis_num]).sample # b, 4, 64, 64
|
||||
noisy_residual = noise_pred_uncond + args.cfg * (noise_pred_cond - noise_pred_uncond) # b, 4, 64, 64
|
||||
input = scheduler.step(noisy_residual, t, input).prev_sample
|
||||
# input = prev_noisy_sample
|
||||
|
||||
# decode
|
||||
input = 1 / vae.config.scaling_factor * input
|
||||
images = vae.decode(input.half(), return_dict=False)[0]
|
||||
width, height = 512, 512
|
||||
new_image = Image.new('RGB', (4*width, 4*height))
|
||||
for index, image in enumerate(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')
|
||||
row = index // 4
|
||||
col = index % 4
|
||||
new_image.paste(image, (col*width, row*height))
|
||||
new_image.save(f'{args.output_dir}/pred_img_{sample_index}_{args.local_rank}.jpg')
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,14 @@
|
||||
accelerate launch inference_textdiffuser2_t2i_full_quadrilateral.py \
|
||||
--pretrained_model_name_or_path="runwayml/stable-diffusion-v1-5" \
|
||||
--mixed_precision="fp16" \
|
||||
--output_dir="inference_results" \
|
||||
--enable_xformers_memory_efficient_attention \
|
||||
--resume_from_checkpoint="put your ckpt here" \
|
||||
--granularity=128 \
|
||||
--max_length=77 \
|
||||
--coord_mode="ltrb" \
|
||||
--cfg=7.5 \
|
||||
--sample_steps=20 \
|
||||
--seed=43555 \
|
||||
--input_format='prompt_layout_txt_file' \
|
||||
--input_file='angle_template_file.txt'
|
||||
@@ -0,0 +1,2 @@
|
||||
hat with the word money
|
||||
MONEY 30,40,80,70,90,60,30,60
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,20 @@
|
||||
accelerate launch train_textdiffuser2_t2i_full_angle.py \
|
||||
--pretrained_model_name_or_path="runwayml/stable-diffusion-v1-5" \
|
||||
--train_batch_size=18 \
|
||||
--gradient_accumulation_steps=4 \
|
||||
--gradient_checkpointing \
|
||||
--mixed_precision="fp16" \
|
||||
--num_train_epochs=6 \
|
||||
--learning_rate=1e-5 \
|
||||
--max_grad_norm=1 \
|
||||
--lr_scheduler="constant" \
|
||||
--lr_warmup_steps=0 \
|
||||
--output_dir="diffusion_experiment_result" \
|
||||
--enable_xformers_memory_efficient_attention \
|
||||
--dataloader_num_workers=8 \
|
||||
--index_file_path='/path/to/train_dataset_index.txt' \
|
||||
--dataset_path='/path/to/laion-ocr-select/' \
|
||||
--granularity=128 \
|
||||
--coord_mode="ltrb" \
|
||||
--max_length=77 \
|
||||
--resume_from_checkpoint="latest"
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,20 @@
|
||||
accelerate launch train_textdiffuser2_t2i_full_quadrilateral.py \
|
||||
--pretrained_model_name_or_path="runwayml/stable-diffusion-v1-5" \
|
||||
--train_batch_size=18 \
|
||||
--gradient_accumulation_steps=4 \
|
||||
--gradient_checkpointing \
|
||||
--mixed_precision="fp16" \
|
||||
--num_train_epochs=6 \
|
||||
--learning_rate=1e-5 \
|
||||
--max_grad_norm=1 \
|
||||
--lr_scheduler="constant" \
|
||||
--lr_warmup_steps=0 \
|
||||
--output_dir="diffusion_experiment_result" \
|
||||
--enable_xformers_memory_efficient_attention \
|
||||
--dataloader_num_workers=8 \
|
||||
--index_file_path='/path/to/train_dataset_index.txt' \
|
||||
--dataset_path='/path/to/laion-ocr-select/' \
|
||||
--granularity=128 \
|
||||
--coord_mode="ltrb" \
|
||||
--max_length=77 \
|
||||
--resume_from_checkpoint="latest"
|
||||
@@ -0,0 +1,638 @@
|
||||
# ------------------------------------------
|
||||
# TextDiffuser-2: Unleashing the Power of Language Models for Text Rendering
|
||||
# Paper Link: https://arxiv.org/abs/2311.16465
|
||||
# Code Link: https://github.com/microsoft/unilm/tree/master/textdiffuser-2
|
||||
# Copyright (c) Microsoft Corporation.
|
||||
# ------------------------------------------
|
||||
|
||||
import os
|
||||
import cv2
|
||||
import random
|
||||
import logging
|
||||
import argparse
|
||||
import numpy as np
|
||||
import time
|
||||
|
||||
from pathlib import Path
|
||||
from tqdm.auto import tqdm
|
||||
from typing import Optional
|
||||
from packaging import version
|
||||
from PIL import Image
|
||||
from huggingface_hub import HfFolder, Repository, create_repo, whoami
|
||||
|
||||
import string
|
||||
alphabet = string.digits + string.ascii_lowercase + string.ascii_uppercase + string.punctuation + ' ' # len(aphabet) = 95
|
||||
'''alphabet
|
||||
0123456789abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ!"#$%&\'()*+,-./:;<=>?@[\\]^_`{|}~
|
||||
'''
|
||||
|
||||
import datasets
|
||||
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
|
||||
|
||||
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(
|
||||
"--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(
|
||||
"--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."
|
||||
)
|
||||
#### newly added parameters
|
||||
parser.add_argument(
|
||||
"--granularity",
|
||||
type=int,
|
||||
default=128,
|
||||
help="The granularity of coordinates, ranging from 1~512."
|
||||
)
|
||||
parser.add_argument(
|
||||
"--coord_mode",
|
||||
type=str,
|
||||
default='lt',
|
||||
choices=['lt', 'center', 'ltrb'],
|
||||
help="The way to represent coordinates."
|
||||
)
|
||||
parser.add_argument(
|
||||
"--max_length",
|
||||
default=77,
|
||||
type=int,
|
||||
help="Maximum length of the composed prompt."
|
||||
)
|
||||
parser.add_argument(
|
||||
"--cfg",
|
||||
default=7,
|
||||
type=float,
|
||||
help="classifier free guidance."
|
||||
)
|
||||
parser.add_argument(
|
||||
"--sample_steps",
|
||||
default=50,
|
||||
type=int,
|
||||
help="steps for sampling for diffusion models."
|
||||
)
|
||||
parser.add_argument(
|
||||
"--input_format",
|
||||
required=True,
|
||||
type=str,
|
||||
help="specify the input format",
|
||||
choices=['prompt', 'prompts_txt_file', 'prompt_layout_txt_file']
|
||||
)
|
||||
parser.add_argument(
|
||||
"--input_prompt",
|
||||
type=str,
|
||||
)
|
||||
parser.add_argument(
|
||||
"--input_file",
|
||||
type=str,
|
||||
)
|
||||
parser.add_argument(
|
||||
"--prompts_txt_file",
|
||||
type=str,
|
||||
)
|
||||
parser.add_argument(
|
||||
"--m1_model_path",
|
||||
type=str,
|
||||
help="the checkpoint of layout planner"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--vis_num",
|
||||
type=int,
|
||||
default=16,
|
||||
help=("The number of images to be visualized."),
|
||||
)
|
||||
args = parser.parse_args()
|
||||
|
||||
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}"
|
||||
|
||||
|
||||
DATASET_NAME_MAPPING = {
|
||||
# "lambdalabs/pokemon-blip-captions": ("image", "text"),
|
||||
"MARIO-10M": ("image", "text"),
|
||||
}
|
||||
|
||||
def main():
|
||||
args = parse_args()
|
||||
logging_dir = Path(args.output_dir, args.logging_dir)
|
||||
|
||||
accelerator_project_config = ProjectConfiguration(project_dir=args.output_dir, logging_dir=logging_dir)
|
||||
|
||||
accelerator = Accelerator(
|
||||
gradient_accumulation_steps=4,
|
||||
mixed_precision=args.mixed_precision,
|
||||
log_with=args.report_to,
|
||||
project_config=accelerator_project_config,
|
||||
)
|
||||
if args.report_to == "wandb":
|
||||
if not is_wandb_available():
|
||||
raise ImportError("Make sure to install wandb if you want to use it for logging during training.")
|
||||
import wandb
|
||||
|
||||
# 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.output_dir is not None:
|
||||
os.makedirs(args.output_dir, exist_ok=True)
|
||||
|
||||
if args.push_to_hub:
|
||||
repo_id = create_repo(
|
||||
repo_id=args.hub_model_id or Path(args.output_dir).name, exist_ok=True, token=args.hub_token
|
||||
).repo_id
|
||||
# Load scheduler, tokenizer and models.
|
||||
tokenizer = CLIPTokenizer.from_pretrained(
|
||||
args.pretrained_model_name_or_path, subfolder="tokenizer"
|
||||
)
|
||||
|
||||
#### additional tokens are introduced, including coordinate tokens and character tokens
|
||||
print('***************')
|
||||
print(len(tokenizer))
|
||||
for i in range(520):
|
||||
tokenizer.add_tokens(['l' + str(i) ]) # left
|
||||
tokenizer.add_tokens(['t' + str(i) ]) # top
|
||||
tokenizer.add_tokens(['r' + str(i) ]) # width
|
||||
tokenizer.add_tokens(['b' + str(i) ]) # height
|
||||
for c in alphabet:
|
||||
tokenizer.add_tokens([f'[{c}]'])
|
||||
print(len(tokenizer))
|
||||
print('***************')
|
||||
|
||||
if args.max_length == 77:
|
||||
text_encoder = CLIPTextModel.from_pretrained(
|
||||
args.resume_from_checkpoint, subfolder="text_encoder", ignore_mismatched_sizes=True
|
||||
)
|
||||
else:
|
||||
#### enlarge the context length of text encoder. empirically, enlarging the context length can proceed longer sequence. However, we observe that it will be hard to render general objects
|
||||
text_encoder = CLIPTextModel.from_pretrained(
|
||||
args.resume_from_checkpoint, subfolder="text_encoder", max_position_embeddings=args.max_length, ignore_mismatched_sizes=True
|
||||
)
|
||||
|
||||
text_encoder.resize_token_embeddings(len(tokenizer))
|
||||
|
||||
vae = AutoencoderKL.from_pretrained(args.pretrained_model_name_or_path, subfolder="vae")
|
||||
unet = UNet2DConditionModel.from_pretrained(
|
||||
args.resume_from_checkpoint, subfolder="unet"
|
||||
)
|
||||
# freeze parameters of models to save more memory
|
||||
# unet.requires_grad_(False)
|
||||
vae.requires_grad_(False)
|
||||
text_encoder.requires_grad_(False)
|
||||
|
||||
|
||||
|
||||
# `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):
|
||||
if args.use_ema:
|
||||
ema_unet.save_pretrained(os.path.join(output_dir, "unet_ema"))
|
||||
|
||||
for i, model in enumerate(models):
|
||||
# model.save_pretrained(os.path.join(output_dir, "unet"))
|
||||
|
||||
if i == 0:
|
||||
model.save_pretrained(os.path.join(output_dir, f"unet"))
|
||||
elif i == 1:
|
||||
model.save_pretrained(os.path.join(output_dir, f"text_encoder"))
|
||||
|
||||
# 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()
|
||||
|
||||
if i == 1:
|
||||
load_model = UNet2DConditionModel.from_pretrained(input_dir, subfolder="unet")
|
||||
model.register_to_config(**load_model.config)
|
||||
elif i == 0:
|
||||
load_model = CLIPTextModel.from_pretrained(input_dir, subfolder="text_encoder")
|
||||
# model.register_to_config(**load_model.config)
|
||||
|
||||
# # 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)
|
||||
|
||||
# For mixed precision training we cast all non-trainable weigths (vae, non-lora text_encoder and non-lora unet) to half-precision
|
||||
# as these weights are only used for inference, keeping weights in full precision is not required.
|
||||
weight_dtype = torch.float32
|
||||
if accelerator.mixed_precision == "fp16":
|
||||
weight_dtype = torch.float16
|
||||
elif accelerator.mixed_precision == "bf16":
|
||||
weight_dtype = torch.bfloat16
|
||||
|
||||
# Move unet, vae and text_encoder to device and cast to weight_dtype
|
||||
unet.to(accelerator.device, dtype=weight_dtype)
|
||||
vae.to(accelerator.device, dtype=weight_dtype)
|
||||
text_encoder.to(accelerator.device, dtype=weight_dtype)
|
||||
|
||||
|
||||
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")
|
||||
|
||||
|
||||
# # Enable TF32 for faster training on Ampere GPUs,
|
||||
# # cf https://pytorch.org/docs/stable/notes/cuda.html#tensorfloat-32-tf32-on-ampere-devices
|
||||
# if True:
|
||||
# torch.backends.cuda.matmul.allow_tf32 = True
|
||||
|
||||
# # We need to initialize the trackers we use, and also store our configuration.
|
||||
# # The trackers initializes automatically on the main process.
|
||||
# if accelerator.is_main_process:
|
||||
# accelerator.init_trackers("text2image-fine-tune", config=vars(args))
|
||||
|
||||
# Potentially load in the weights and states from a previous save
|
||||
if args.resume_from_checkpoint:
|
||||
if args.resume_from_checkpoint != "latest":
|
||||
path = os.path.basename(args.resume_from_checkpoint)
|
||||
else:
|
||||
# Get the most recent checkpoint
|
||||
dirs = os.listdir(args.output_dir)
|
||||
dirs = [d for d in dirs if d.startswith("checkpoint")]
|
||||
dirs = sorted(dirs, key=lambda x: int(x.split("-")[1]))
|
||||
path = dirs[-1] if len(dirs) > 0 else None
|
||||
|
||||
if path is None:
|
||||
accelerator.print(
|
||||
f"Checkpoint '{args.resume_from_checkpoint}' does not exist. Starting a new training run."
|
||||
)
|
||||
args.resume_from_checkpoint = None
|
||||
else:
|
||||
accelerator.print(f"Resuming from checkpoint {path}")
|
||||
# accelerator.load_state(os.path.join(args.output_dir, path))
|
||||
accelerator.load_state(args.resume_from_checkpoint)
|
||||
|
||||
if accelerator.is_main_process and os.path.exists(f'{args.output_dir}'):
|
||||
print('detect existing output_dir, removing the contained jpg/txt files ...')
|
||||
os.system(f'rm {args.output_dir}/*.jpg')
|
||||
os.system(f'rm {args.output_dir}/*.txt')
|
||||
|
||||
# user_prompt = "Book cover of summer vibe, high quality, high resolution"
|
||||
# ocrs = [
|
||||
# 'Summer Vibe 20,20,100,40'
|
||||
# ]
|
||||
|
||||
if args.input_format == 'prompt_layout_txt_file':
|
||||
lines = open(args.input_file).readlines()
|
||||
user_prompts = [lines[0].strip()]
|
||||
ocrs = [lines[1:]]
|
||||
|
||||
elif args.input_format == 'prompt' or args.input_format == 'prompts_txt_file':
|
||||
|
||||
#### prepare m1 (layout planner)
|
||||
from fastchat.model import load_model, get_conversation_template
|
||||
m1_model, m1_tokenizer = load_model(
|
||||
args.m1_model_path,
|
||||
'cuda',
|
||||
1,
|
||||
None,
|
||||
False,
|
||||
False,
|
||||
revision="main",
|
||||
debug=False,
|
||||
)
|
||||
|
||||
# prompt = 'a text image of hello world'
|
||||
prompts = []
|
||||
if args.input_format == 'prompt':
|
||||
prompts = [args.input_prompt]
|
||||
elif args.input_format == 'prompts_txt_file':
|
||||
prompts = open(args.prompts_txt_file).readlines()
|
||||
print(f'there are {len(prompts)} samples for generation')
|
||||
|
||||
ocrs = []
|
||||
user_prompts = []
|
||||
for prompt in prompts:
|
||||
user_prompt = prompt
|
||||
user_prompts.append(user_prompt)
|
||||
template = f'Given a prompt that will be used to generate an image, plan the layout of visual text for the image. The size of the image is 128x128. Therefore, all properties of the positions should not exceed 128, including the coordinates of top, left, right, and bottom. All keywords are included in the caption. You dont need to specify the details of font styles. At each line, the format should be keyword left, top, right, bottom. So let us begin. Prompt: {user_prompt}'
|
||||
msg = template
|
||||
conv = get_conversation_template(args.m1_model_path)
|
||||
conv.append_message(conv.roles[0], msg)
|
||||
conv.append_message(conv.roles[1], None)
|
||||
prompt = conv.get_prompt()
|
||||
inputs = m1_tokenizer([prompt], return_token_type_ids=False)
|
||||
inputs = {k: torch.tensor(v).to('cuda') for k, v in inputs.items()}
|
||||
output_ids = m1_model.generate(
|
||||
**inputs,
|
||||
do_sample=True,
|
||||
temperature=0.7,
|
||||
repetition_penalty=1.0,
|
||||
max_new_tokens=512,
|
||||
)
|
||||
|
||||
if m1_model.config.is_encoder_decoder:
|
||||
output_ids = output_ids[0]
|
||||
else:
|
||||
output_ids = output_ids[0][len(inputs["input_ids"][0]) :]
|
||||
outputs = m1_tokenizer.decode(
|
||||
output_ids, skip_special_tokens=True, spaces_between_special_tokens=False
|
||||
)
|
||||
print(f"[{conv.roles[0]}]\n{msg}")
|
||||
print(f"[{conv.roles[1]}]\n{outputs}")
|
||||
# ocrs = outputs.split('\n')
|
||||
ocrs.append(outputs.split('\n'))
|
||||
|
||||
|
||||
with torch.no_grad():
|
||||
|
||||
size = len(ocrs)
|
||||
print(f'the number of samples: {size}')
|
||||
|
||||
time_seed = int(time.time())
|
||||
random.seed(time_seed)
|
||||
torch.manual_seed(time_seed)
|
||||
torch.cuda.manual_seed_all(time_seed)
|
||||
|
||||
for sample_index in range(size):
|
||||
|
||||
user_prompt = user_prompts[sample_index]
|
||||
current_ocr = ocrs[sample_index]
|
||||
|
||||
ocr_ids = []
|
||||
print('user_prompt', user_prompt)
|
||||
print('current_ocr', current_ocr)
|
||||
|
||||
current_ocr = []
|
||||
for ocr in current_ocr:
|
||||
ocr = ocr.strip()
|
||||
|
||||
if len(ocr) == 0 or '###' in ocr or '.com' in ocr:
|
||||
continue
|
||||
|
||||
items = ocr.split()
|
||||
pred = ' '.join(items[:-1])
|
||||
box = items[-1]
|
||||
|
||||
l,t,r,b = box.split(',')
|
||||
l,t,r,b = int(l), int(t), int(r), int(b)
|
||||
ocr_ids.extend(['l'+str(l), 't'+str(t), 'r'+str(r), 'b'+str(b)])
|
||||
|
||||
char_list = list(pred)
|
||||
char_list = [f'[{i}]' for i in char_list]
|
||||
ocr_ids.extend(char_list)
|
||||
ocr_ids.append(tokenizer.eos_token_id)
|
||||
|
||||
caption_ids = tokenizer(
|
||||
user_prompt, truncation=True, return_tensors="pt"
|
||||
).input_ids[0].tolist()
|
||||
|
||||
try:
|
||||
ocr_ids = tokenizer.encode(ocr_ids)
|
||||
prompt = caption_ids + ocr_ids
|
||||
except:
|
||||
prompt = caption_ids
|
||||
|
||||
prompt = prompt[:args.max_length]
|
||||
while len(prompt) < args.max_length:
|
||||
prompt.append(tokenizer.pad_token_id)
|
||||
prompts_cond = prompt
|
||||
prompts_nocond = [tokenizer.pad_token_id]*args.max_length
|
||||
|
||||
prompts_cond = [prompts_cond] * args.vis_num
|
||||
prompts_nocond = [prompts_nocond] * args.vis_num
|
||||
|
||||
prompts_cond = torch.Tensor(prompts_cond).long().cuda()
|
||||
prompts_nocond = torch.Tensor(prompts_nocond).long().cuda()
|
||||
|
||||
scheduler = DDPMScheduler.from_pretrained(args.pretrained_model_name_or_path, subfolder="scheduler")
|
||||
scheduler.set_timesteps(args.sample_steps)
|
||||
noise = torch.randn((args.vis_num, 4, 64, 64)).to("cuda")
|
||||
input = noise
|
||||
|
||||
encoder_hidden_states_cond = text_encoder(prompts_cond)[0]
|
||||
encoder_hidden_states_nocond = text_encoder(prompts_nocond)[0]
|
||||
|
||||
texts = prompts_cond
|
||||
f = open(f'{args.output_dir}/prompt_{sample_index}_{args.local_rank}.txt', 'w+')
|
||||
for text in texts:
|
||||
sentence = tokenizer.decode(text)
|
||||
f.write(sentence + '\n')
|
||||
f.close()
|
||||
|
||||
for t in tqdm(scheduler.timesteps):
|
||||
with torch.no_grad(): # classifier free guidance
|
||||
noise_pred_cond = unet(sample=input.half(), timestep=t, encoder_hidden_states=encoder_hidden_states_cond[:args.vis_num]).sample # b, 4, 64, 64
|
||||
noise_pred_uncond = unet(sample=input.half(), timestep=t, encoder_hidden_states=encoder_hidden_states_nocond[:args.vis_num]).sample # b, 4, 64, 64
|
||||
noisy_residual = noise_pred_uncond + args.cfg * (noise_pred_cond - noise_pred_uncond) # b, 4, 64, 64
|
||||
input = scheduler.step(noisy_residual, t, input).prev_sample
|
||||
# input = prev_noisy_sample
|
||||
|
||||
# decode
|
||||
input = 1 / vae.config.scaling_factor * input
|
||||
images = vae.decode(input.half(), return_dict=False)[0]
|
||||
width, height = 512, 512
|
||||
new_image = Image.new('RGB', (4*width, 4*height))
|
||||
for index, image in enumerate(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')
|
||||
row = index // 4
|
||||
col = index % 4
|
||||
new_image.paste(image, (col*width, row*height))
|
||||
new_image.save(f'{args.output_dir}/pred_img_{sample_index}_{args.local_rank}.jpg')
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,15 @@
|
||||
accelerate launch inference_textdiffuser2_t2i_full.py \
|
||||
--pretrained_model_name_or_path="runwayml/stable-diffusion-v1-5" \
|
||||
--mixed_precision="fp16" \
|
||||
--output_dir="inference_results" \
|
||||
--enable_xformers_memory_efficient_attention \
|
||||
--resume_from_checkpoint="JingyeChen22/textdiffuser2-full-ft" \
|
||||
--granularity=128 \
|
||||
--max_length=77 \
|
||||
--coord_mode="ltrb" \
|
||||
--cfg=7.5 \
|
||||
--sample_steps=20 \
|
||||
--seed=43555 \
|
||||
--m1_model_path="JingyeChen22/textdiffuser2_layout_planner" \
|
||||
--input_format='prompt' \
|
||||
--input_prompt='a hotdog with mustard and other toppings on it'
|
||||
@@ -0,0 +1,616 @@
|
||||
# ------------------------------------------
|
||||
# TextDiffuser-2: Unleashing the Power of Language Models for Text Rendering
|
||||
# Paper Link: https://arxiv.org/abs/2311.16465
|
||||
# Code Link: https://github.com/microsoft/unilm/tree/master/textdiffuser-2
|
||||
# Copyright (c) Microsoft Corporation.
|
||||
# ------------------------------------------
|
||||
|
||||
import argparse
|
||||
import logging
|
||||
import math
|
||||
import os
|
||||
import random
|
||||
import shutil
|
||||
from pathlib import Path
|
||||
import glob
|
||||
import time
|
||||
|
||||
import datasets
|
||||
import numpy as np
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
import torch.utils.checkpoint
|
||||
import transformers
|
||||
from accelerate import Accelerator
|
||||
from accelerate.logging import get_logger
|
||||
from accelerate.utils import ProjectConfiguration, set_seed
|
||||
from datasets import load_dataset, Dataset
|
||||
|
||||
from huggingface_hub import create_repo, upload_folder
|
||||
from packaging import version
|
||||
from torchvision import transforms
|
||||
from tqdm.auto import tqdm
|
||||
from transformers import CLIPTextModel, CLIPTokenizer
|
||||
|
||||
import diffusers
|
||||
from diffusers import AutoencoderKL, DDPMScheduler, DiffusionPipeline, UNet2DConditionModel
|
||||
from diffusers.loaders import AttnProcsLayers
|
||||
from diffusers.models.attention_processor import LoRAAttnProcessor
|
||||
from diffusers.optimization import get_scheduler
|
||||
from diffusers.utils import check_min_version, is_wandb_available
|
||||
from diffusers.utils.import_utils import is_xformers_available
|
||||
|
||||
from PIL import Image
|
||||
import string
|
||||
alphabet = string.digits + string.ascii_lowercase + string.ascii_uppercase + string.punctuation + ' ' # len(aphabet) = 95
|
||||
'''alphabet
|
||||
0123456789abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ!"#$%&\'()*+,-./:;<=>?@[\\]^_`{|}~
|
||||
'''
|
||||
|
||||
logger = get_logger(__name__, log_level="INFO")
|
||||
|
||||
|
||||
def save_model_card(repo_id: str, images=None, base_model=str, dataset_name=str, repo_folder=None):
|
||||
img_str = ""
|
||||
for i, image in enumerate(images):
|
||||
image.save(os.path.join(repo_folder, f"image_{i}.png"))
|
||||
img_str += f"\n"
|
||||
|
||||
yaml = f"""
|
||||
---
|
||||
license: creativeml-openrail-m
|
||||
base_model: {base_model}
|
||||
tags:
|
||||
- stable-diffusion
|
||||
- stable-diffusion-diffusers
|
||||
- text-to-image
|
||||
- diffusers
|
||||
- lora
|
||||
inference: true
|
||||
---
|
||||
"""
|
||||
model_card = f"""
|
||||
# LoRA text2image fine-tuning - {repo_id}
|
||||
These are LoRA adaption weights for {base_model}. The weights were fine-tuned on the {dataset_name} dataset. You can find some example images in the following. \n
|
||||
{img_str}
|
||||
"""
|
||||
with open(os.path.join(repo_folder, "README.md"), "w") as f:
|
||||
f.write(yaml + model_card)
|
||||
|
||||
|
||||
def parse_args():
|
||||
parser = argparse.ArgumentParser(description="Simple example of a training script.")
|
||||
parser.add_argument(
|
||||
"--pretrained_model_name_or_path",
|
||||
type=str,
|
||||
default=None,
|
||||
required=True,
|
||||
help="Path to pretrained model or model identifier from huggingface.co/models.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--dataset_name",
|
||||
type=str,
|
||||
default='lambdalabs/pokemon-blip-captions',
|
||||
help=(
|
||||
"The name of the Dataset (from the HuggingFace hub) to train on (could be your own, possibly private,"
|
||||
" dataset). It can also be a path pointing to a local copy of a dataset in your filesystem,"
|
||||
" or to a folder containing files that 🤗 Datasets can understand."
|
||||
),
|
||||
)
|
||||
parser.add_argument(
|
||||
"--output_dir",
|
||||
type=str,
|
||||
default="sd-model-finetuned-lora",
|
||||
help="The output directory where the model predictions and checkpoints will be written.",
|
||||
)
|
||||
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(
|
||||
"--gradient_accumulation_steps",
|
||||
type=int,
|
||||
default=1,
|
||||
help="Number of updates steps to accumulate before performing a backward/update pass.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--gradient_checkpointing",
|
||||
action="store_true",
|
||||
help="Whether or not to use gradient checkpointing to save memory at the expense of slower backward pass.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--allow_tf32",
|
||||
action="store_true",
|
||||
help=(
|
||||
"Whether or not to allow TF32 on Ampere GPUs. Can be used to speed up training. For more information, see"
|
||||
" https://pytorch.org/docs/stable/notes/cuda.html#tensorfloat-32-tf32-on-ampere-devices"
|
||||
),
|
||||
)
|
||||
parser.add_argument("--push_to_hub", action="store_true", help="Whether or not to push the model to the Hub.")
|
||||
parser.add_argument("--hub_token", type=str, default=None, help="The token to use to push to the Model Hub.")
|
||||
parser.add_argument(
|
||||
"--hub_model_id",
|
||||
type=str,
|
||||
default=None,
|
||||
help="The name of the repository to keep in sync with the local `output_dir`.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--logging_dir",
|
||||
type=str,
|
||||
default="logs",
|
||||
help=(
|
||||
"[TensorBoard](https://www.tensorflow.org/tensorboard) log directory. Will default to"
|
||||
" *output_dir/runs/**CURRENT_DATETIME_HOSTNAME***."
|
||||
),
|
||||
)
|
||||
parser.add_argument(
|
||||
"--mixed_precision",
|
||||
type=str,
|
||||
default=None,
|
||||
choices=["no", "fp16", "bf16"],
|
||||
help=(
|
||||
"Whether to use mixed precision. Choose between fp16 and bf16 (bfloat16). Bf16 requires PyTorch >="
|
||||
" 1.10.and an Nvidia Ampere GPU. Default to the value of accelerate config of the current system or the"
|
||||
" flag passed with the `accelerate.launch` command. Use this argument to override the accelerate config."
|
||||
),
|
||||
)
|
||||
parser.add_argument(
|
||||
"--report_to",
|
||||
type=str,
|
||||
default="tensorboard",
|
||||
help=(
|
||||
'The integration to report the results and logs to. Supported platforms are `"tensorboard"`'
|
||||
' (default), `"wandb"` and `"comet_ml"`. Use `"all"` to report to all integrations.'
|
||||
),
|
||||
)
|
||||
parser.add_argument("--local_rank", type=int, default=-1, help="For distributed training: local_rank")
|
||||
parser.add_argument(
|
||||
"--resume_from_checkpoint",
|
||||
type=str,
|
||||
default=None,
|
||||
help=(
|
||||
"Whether training should be resumed from a previous checkpoint. Use a path saved by"
|
||||
' `--checkpointing_steps`, or `"latest"` to automatically select the last available checkpoint.'
|
||||
),
|
||||
)
|
||||
parser.add_argument(
|
||||
"--enable_xformers_memory_efficient_attention", action="store_true", help="Whether or not to use xformers."
|
||||
)
|
||||
parser.add_argument(
|
||||
"--rank",
|
||||
type=int,
|
||||
default=4,
|
||||
help=("The dimension of the LoRA update matrices."),
|
||||
)
|
||||
parser.add_argument(
|
||||
"--vis_num",
|
||||
type=int,
|
||||
default=16,
|
||||
help=("The number of images to be visualized."),
|
||||
)
|
||||
|
||||
#### newly added parameters
|
||||
parser.add_argument(
|
||||
"--granularity",
|
||||
type=int,
|
||||
default=128,
|
||||
help="The granularity of coordinates, ranging from 1~512."
|
||||
)
|
||||
parser.add_argument(
|
||||
"--coord_mode",
|
||||
type=str,
|
||||
default='lt',
|
||||
choices=['lt', 'center', 'ltrb'],
|
||||
help="The way to represent coordinates."
|
||||
)
|
||||
parser.add_argument(
|
||||
"--max_length",
|
||||
default=77,
|
||||
type=int,
|
||||
help="Maximum length of the composed prompt."
|
||||
)
|
||||
parser.add_argument(
|
||||
"--cfg",
|
||||
default=7,
|
||||
type=float,
|
||||
help="classifier free guidance."
|
||||
)
|
||||
parser.add_argument(
|
||||
"--sample_steps",
|
||||
default=50,
|
||||
type=int,
|
||||
help="steps for sampling for diffusion models."
|
||||
)
|
||||
parser.add_argument(
|
||||
"--input_format",
|
||||
required=True,
|
||||
type=str,
|
||||
help="specify the input format",
|
||||
choices=['prompt', 'prompts_txt_file', 'prompt_layout_txt_file']
|
||||
)
|
||||
parser.add_argument(
|
||||
"--input_prompt",
|
||||
type=str,
|
||||
)
|
||||
parser.add_argument(
|
||||
"--input_file",
|
||||
type=str,
|
||||
)
|
||||
parser.add_argument(
|
||||
"--prompts_txt_file",
|
||||
type=str,
|
||||
)
|
||||
parser.add_argument(
|
||||
"--m1_model_path",
|
||||
type=str,
|
||||
help="the checkpoint of layout planner"
|
||||
)
|
||||
|
||||
args = parser.parse_args()
|
||||
env_local_rank = int(os.environ.get("LOCAL_RANK", -1))
|
||||
if env_local_rank != -1 and env_local_rank != args.local_rank:
|
||||
args.local_rank = env_local_rank
|
||||
|
||||
return args
|
||||
|
||||
|
||||
DATASET_NAME_MAPPING = {
|
||||
# "lambdalabs/pokemon-blip-captions": ("image", "text"),
|
||||
"MARIO-10M": ("image", "text"),
|
||||
}
|
||||
|
||||
def main():
|
||||
args = parse_args()
|
||||
logging_dir = Path(args.output_dir, args.logging_dir)
|
||||
|
||||
accelerator_project_config = ProjectConfiguration(project_dir=args.output_dir, logging_dir=logging_dir)
|
||||
|
||||
accelerator = Accelerator(
|
||||
gradient_accumulation_steps=args.gradient_accumulation_steps,
|
||||
mixed_precision=args.mixed_precision,
|
||||
log_with=args.report_to,
|
||||
project_config=accelerator_project_config,
|
||||
)
|
||||
if args.report_to == "wandb":
|
||||
if not is_wandb_available():
|
||||
raise ImportError("Make sure to install wandb if you want to use it for logging during training.")
|
||||
import wandb
|
||||
|
||||
# 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.output_dir is not None:
|
||||
os.makedirs(args.output_dir, exist_ok=True)
|
||||
|
||||
if args.push_to_hub:
|
||||
repo_id = create_repo(
|
||||
repo_id=args.hub_model_id or Path(args.output_dir).name, exist_ok=True, token=args.hub_token
|
||||
).repo_id
|
||||
# Load scheduler, tokenizer and models.
|
||||
tokenizer = CLIPTokenizer.from_pretrained(
|
||||
args.pretrained_model_name_or_path, subfolder="tokenizer"
|
||||
)
|
||||
|
||||
#### additional tokens are introduced, including coordinate tokens and character tokens
|
||||
print('***************')
|
||||
print(len(tokenizer))
|
||||
for i in range(520):
|
||||
tokenizer.add_tokens(['l' + str(i) ]) # left
|
||||
tokenizer.add_tokens(['t' + str(i) ]) # top
|
||||
tokenizer.add_tokens(['r' + str(i) ]) # width
|
||||
tokenizer.add_tokens(['b' + str(i) ]) # height
|
||||
for c in alphabet:
|
||||
tokenizer.add_tokens([f'[{c}]'])
|
||||
print(len(tokenizer))
|
||||
print('***************')
|
||||
|
||||
if args.max_length == 77:
|
||||
text_encoder = CLIPTextModel.from_pretrained(
|
||||
args.pretrained_model_name_or_path, subfolder="text_encoder", ignore_mismatched_sizes=True
|
||||
)
|
||||
else:
|
||||
#### enlarge the context length of text encoder. empirically, enlarging the context length can proceed longer sequence. However, we observe that it will be hard to render general objects
|
||||
text_encoder = CLIPTextModel.from_pretrained(
|
||||
args.pretrained_model_name_or_path, subfolder="text_encoder", max_position_embeddings=args.max_length, ignore_mismatched_sizes=True
|
||||
)
|
||||
|
||||
text_encoder.resize_token_embeddings(len(tokenizer))
|
||||
|
||||
vae = AutoencoderKL.from_pretrained(args.pretrained_model_name_or_path, subfolder="vae")
|
||||
unet = UNet2DConditionModel.from_pretrained(
|
||||
args.pretrained_model_name_or_path, subfolder="unet"
|
||||
)
|
||||
# freeze parameters of models to save more memory
|
||||
unet.requires_grad_(False)
|
||||
vae.requires_grad_(False)
|
||||
text_encoder.requires_grad_(True)
|
||||
|
||||
# For mixed precision training we cast all non-trainable weigths (vae, non-lora text_encoder and non-lora unet) to half-precision
|
||||
# as these weights are only used for inference, keeping weights in full precision is not required.
|
||||
weight_dtype = torch.float32
|
||||
if accelerator.mixed_precision == "fp16":
|
||||
weight_dtype = torch.float16
|
||||
elif accelerator.mixed_precision == "bf16":
|
||||
weight_dtype = torch.bfloat16
|
||||
|
||||
# Move unet, vae and text_encoder to device and cast to weight_dtype
|
||||
unet.to(accelerator.device, dtype=weight_dtype)
|
||||
vae.to(accelerator.device, dtype=weight_dtype)
|
||||
# text_encoder.to(accelerator.device, dtype=weight_dtype)
|
||||
|
||||
# now we will add new LoRA weights to the attention layers
|
||||
# It's important to realize here how many attention weights will be added and of which sizes
|
||||
# The sizes of the attention layers consist only of two different variables:
|
||||
# 1) - the "hidden_size", which is increased according to `unet.config.block_out_channels`.
|
||||
# 2) - the "cross attention size", which is set to `unet.config.cross_attention_dim`.
|
||||
|
||||
# Let's first see how many attention processors we will have to set.
|
||||
# For Stable Diffusion, it should be equal to:
|
||||
# - down blocks (2x attention layers) * (2x transformer layers) * (3x down blocks) = 12
|
||||
# - mid blocks (2x attention layers) * (1x transformer layers) * (1x mid blocks) = 2
|
||||
# - up blocks (2x attention layers) * (3x transformer layers) * (3x down blocks) = 18
|
||||
# => 32 layers
|
||||
|
||||
# Set correct lora layers
|
||||
lora_attn_procs = {}
|
||||
for name in unet.attn_processors.keys():
|
||||
cross_attention_dim = None if name.endswith("attn1.processor") else unet.config.cross_attention_dim
|
||||
if name.startswith("mid_block"):
|
||||
hidden_size = unet.config.block_out_channels[-1]
|
||||
elif name.startswith("up_blocks"):
|
||||
block_id = int(name[len("up_blocks.")])
|
||||
hidden_size = list(reversed(unet.config.block_out_channels))[block_id]
|
||||
elif name.startswith("down_blocks"):
|
||||
block_id = int(name[len("down_blocks.")])
|
||||
hidden_size = unet.config.block_out_channels[block_id]
|
||||
|
||||
lora_attn_procs[name] = LoRAAttnProcessor(
|
||||
hidden_size=hidden_size,
|
||||
cross_attention_dim=cross_attention_dim,
|
||||
rank=args.rank,
|
||||
)
|
||||
|
||||
unet.set_attn_processor(lora_attn_procs)
|
||||
|
||||
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")
|
||||
|
||||
lora_layers = AttnProcsLayers(unet.attn_processors)
|
||||
|
||||
lora_layers, text_encoder = accelerator.prepare(
|
||||
lora_layers, text_encoder
|
||||
)
|
||||
|
||||
# Enable TF32 for faster training on Ampere GPUs,
|
||||
# cf https://pytorch.org/docs/stable/notes/cuda.html#tensorfloat-32-tf32-on-ampere-devices
|
||||
if args.allow_tf32:
|
||||
torch.backends.cuda.matmul.allow_tf32 = True
|
||||
|
||||
# # We need to initialize the trackers we use, and also store our configuration.
|
||||
# # The trackers initializes automatically on the main process.
|
||||
# if accelerator.is_main_process:
|
||||
# accelerator.init_trackers("text2image-fine-tune", config=vars(args))
|
||||
|
||||
# Potentially load in the weights and states from a previous save
|
||||
if args.resume_from_checkpoint:
|
||||
if args.resume_from_checkpoint != "latest":
|
||||
path = os.path.basename(args.resume_from_checkpoint)
|
||||
else:
|
||||
# Get the most recent checkpoint
|
||||
dirs = os.listdir(args.output_dir)
|
||||
dirs = [d for d in dirs if d.startswith("checkpoint")]
|
||||
dirs = sorted(dirs, key=lambda x: int(x.split("-")[1]))
|
||||
path = dirs[-1] if len(dirs) > 0 else None
|
||||
|
||||
if path is None:
|
||||
accelerator.print(
|
||||
f"Checkpoint '{args.resume_from_checkpoint}' does not exist. Starting a new training run."
|
||||
)
|
||||
args.resume_from_checkpoint = None
|
||||
else:
|
||||
accelerator.print(f"Resuming from checkpoint {path}")
|
||||
# accelerator.load_state(os.path.join(args.output_dir, path))
|
||||
accelerator.load_state(args.resume_from_checkpoint)
|
||||
|
||||
if accelerator.is_main_process and os.path.exists(f'{args.output_dir}'):
|
||||
print('detect existing output_dir, removing the contained jpg/txt files ...')
|
||||
os.system(f'rm {args.output_dir}/*.jpg')
|
||||
os.system(f'rm {args.output_dir}/*.txt')
|
||||
|
||||
# user_prompt = "Book cover of summer vibe, high quality, high resolution"
|
||||
# ocrs = [
|
||||
# 'Summer Vibe 20,20,100,40'
|
||||
# ]
|
||||
|
||||
if args.input_format == 'prompt_layout_txt_file':
|
||||
lines = open(args.input_file).readlines()
|
||||
user_prompts = [lines[0].strip()]
|
||||
ocrs = [lines[1:]]
|
||||
|
||||
elif args.input_format == 'prompt' or args.input_format == 'prompts_txt_file':
|
||||
|
||||
#### prepare m1 (layout planner)
|
||||
from fastchat.model import load_model, get_conversation_template
|
||||
m1_model, m1_tokenizer = load_model(
|
||||
args.m1_model_path,
|
||||
'cuda',
|
||||
1,
|
||||
None,
|
||||
False,
|
||||
False,
|
||||
revision="main",
|
||||
debug=False,
|
||||
)
|
||||
|
||||
# prompt = 'a text image of hello world'
|
||||
prompts = []
|
||||
if args.input_format == 'prompt':
|
||||
prompts = [args.input_prompt]
|
||||
elif args.input_format == 'prompts_txt_file':
|
||||
prompts = open(args.prompts_txt_file).readlines()
|
||||
print(f'there are {len(prompts)} samples for generation')
|
||||
|
||||
ocrs = []
|
||||
user_prompts = []
|
||||
for prompt in prompts:
|
||||
user_prompt = prompt
|
||||
user_prompts.append(user_prompt)
|
||||
template = f'Given a prompt that will be used to generate an image, plan the layout of visual text for the image. The size of the image is 128x128. Therefore, all properties of the positions should not exceed 128, including the coordinates of top, left, right, and bottom. All keywords are included in the caption. You dont need to specify the details of font styles. At each line, the format should be keyword left, top, right, bottom. So let us begin. Prompt: {user_prompt}'
|
||||
msg = template
|
||||
conv = get_conversation_template(args.m1_model_path)
|
||||
conv.append_message(conv.roles[0], msg)
|
||||
conv.append_message(conv.roles[1], None)
|
||||
prompt = conv.get_prompt()
|
||||
inputs = m1_tokenizer([prompt], return_token_type_ids=False)
|
||||
inputs = {k: torch.tensor(v).to('cuda') for k, v in inputs.items()}
|
||||
output_ids = m1_model.generate(
|
||||
**inputs,
|
||||
do_sample=True,
|
||||
temperature=0.7,
|
||||
repetition_penalty=1.0,
|
||||
max_new_tokens=512,
|
||||
)
|
||||
|
||||
if m1_model.config.is_encoder_decoder:
|
||||
output_ids = output_ids[0]
|
||||
else:
|
||||
output_ids = output_ids[0][len(inputs["input_ids"][0]) :]
|
||||
outputs = m1_tokenizer.decode(
|
||||
output_ids, skip_special_tokens=True, spaces_between_special_tokens=False
|
||||
)
|
||||
print(f"[{conv.roles[0]}]\n{msg}")
|
||||
print(f"[{conv.roles[1]}]\n{outputs}")
|
||||
# ocrs = outputs.split('\n')
|
||||
ocrs.append(outputs.split('\n'))
|
||||
|
||||
|
||||
with torch.no_grad():
|
||||
|
||||
size = len(ocrs)
|
||||
print(f'the number of samples: {size}')
|
||||
|
||||
time_seed = int(time.time())
|
||||
random.seed(time_seed)
|
||||
torch.manual_seed(time_seed)
|
||||
torch.cuda.manual_seed_all(time_seed)
|
||||
|
||||
for sample_index in range(size):
|
||||
|
||||
user_prompt = user_prompts[sample_index]
|
||||
current_ocr = ocrs[sample_index]
|
||||
|
||||
ocr_ids = []
|
||||
print('user_prompt', user_prompt)
|
||||
print('current_ocr', current_ocr)
|
||||
|
||||
for ocr in current_ocr:
|
||||
ocr = ocr.strip()
|
||||
|
||||
if len(ocr) == 0 or '###' in ocr or '.com' in ocr:
|
||||
continue
|
||||
|
||||
items = ocr.split()
|
||||
pred = ' '.join(items[:-1])
|
||||
box = items[-1]
|
||||
|
||||
l,t,r,b = box.split(',')
|
||||
l,t,r,b = int(l), int(t), int(r), int(b)
|
||||
ocr_ids.extend(['l'+str(l), 't'+str(t), 'r'+str(r), 'b'+str(b)])
|
||||
|
||||
char_list = list(pred)
|
||||
char_list = [f'[{i}]' for i in char_list]
|
||||
ocr_ids.extend(char_list)
|
||||
ocr_ids.append(tokenizer.eos_token_id)
|
||||
|
||||
caption_ids = tokenizer(
|
||||
user_prompt, truncation=True, return_tensors="pt"
|
||||
).input_ids[0].tolist()
|
||||
|
||||
try:
|
||||
ocr_ids = tokenizer.encode(ocr_ids)
|
||||
prompt = caption_ids + ocr_ids
|
||||
except:
|
||||
prompt = caption_ids
|
||||
|
||||
prompt = prompt[:args.max_length]
|
||||
while len(prompt) < args.max_length:
|
||||
prompt.append(tokenizer.pad_token_id)
|
||||
prompts_cond = prompt
|
||||
prompts_nocond = [tokenizer.pad_token_id]*args.max_length
|
||||
|
||||
prompts_cond = [prompts_cond] * args.vis_num
|
||||
prompts_nocond = [prompts_nocond] * args.vis_num
|
||||
|
||||
prompts_cond = torch.Tensor(prompts_cond).long().cuda()
|
||||
prompts_nocond = torch.Tensor(prompts_nocond).long().cuda()
|
||||
|
||||
scheduler = DDPMScheduler.from_pretrained(args.pretrained_model_name_or_path, subfolder="scheduler")
|
||||
scheduler.set_timesteps(args.sample_steps)
|
||||
noise = torch.randn((args.vis_num, 4, 64, 64)).to("cuda")
|
||||
input = noise
|
||||
|
||||
encoder_hidden_states_cond = text_encoder(prompts_cond)[0]
|
||||
encoder_hidden_states_nocond = text_encoder(prompts_nocond)[0]
|
||||
|
||||
texts = prompts_cond
|
||||
f = open(f'{args.output_dir}/prompt_{sample_index}_{args.local_rank}.txt', 'w+')
|
||||
for text in texts:
|
||||
sentence = tokenizer.decode(text)
|
||||
f.write(sentence + '\n')
|
||||
f.close()
|
||||
|
||||
for t in tqdm(scheduler.timesteps):
|
||||
with torch.no_grad(): # classifier free guidance
|
||||
noise_pred_cond = unet(sample=input.half(), timestep=t, encoder_hidden_states=encoder_hidden_states_cond[:args.vis_num]).sample # b, 4, 64, 64
|
||||
noise_pred_uncond = unet(sample=input.half(), timestep=t, encoder_hidden_states=encoder_hidden_states_nocond[:args.vis_num]).sample # b, 4, 64, 64
|
||||
noisy_residual = noise_pred_uncond + args.cfg * (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
|
||||
|
||||
# decode
|
||||
input = 1 / vae.config.scaling_factor * input
|
||||
images = vae.decode(input.half(), return_dict=False)[0]
|
||||
width, height = 512, 512
|
||||
new_image = Image.new('RGB', (4*width, 4*height))
|
||||
for index, image in enumerate(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')
|
||||
row = index // 4
|
||||
col = index % 4
|
||||
new_image.paste(image, (col*width, row*height))
|
||||
new_image.save(f'{args.output_dir}/pred_img_{sample_index}_{args.local_rank}.jpg')
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,16 @@
|
||||
accelerate launch inference_textdiffuser2_t2i_lora.py \
|
||||
--pretrained_model_name_or_path="runwayml/stable-diffusion-v1-5" \
|
||||
--gradient_accumulation_steps=4 \
|
||||
--gradient_checkpointing \
|
||||
--mixed_precision="fp16" \
|
||||
--output_dir="inference_results" \
|
||||
--enable_xformers_memory_efficient_attention \
|
||||
--resume_from_checkpoint="JingyeChen22/textdiffuser2-lora-ft" \
|
||||
--granularity=128 \
|
||||
--coord_mode="ltrb" \
|
||||
--cfg=7.5 \
|
||||
--sample_steps=50 \
|
||||
--seed=43555 \
|
||||
--m1_model_path="/home/jingyechen/FastChat/1204_final" \
|
||||
--input_format='prompt' \
|
||||
--input_prompt='a stamp of u.s.a'
|
||||
@@ -0,0 +1,358 @@
|
||||
# Prediction interface for Cog ⚙️
|
||||
# https://github.com/replicate/cog/blob/main/docs/python.md
|
||||
|
||||
import copy
|
||||
import string
|
||||
import random
|
||||
from typing import Optional
|
||||
import torch
|
||||
from transformers import CLIPTextModel, CLIPTokenizer
|
||||
from diffusers import (
|
||||
AutoencoderKL,
|
||||
DDPMScheduler,
|
||||
StableDiffusionPipeline,
|
||||
UNet2DConditionModel,
|
||||
DiffusionPipeline,
|
||||
LCMScheduler,
|
||||
)
|
||||
from tqdm import tqdm
|
||||
from PIL import Image
|
||||
from PIL import Image, ImageDraw, ImageFont
|
||||
from fastchat.model import load_model, get_conversation_template
|
||||
from transformers import AutoTokenizer, AutoModelForCausalLM
|
||||
from cog import BasePredictor, Input, Path, BaseModel
|
||||
|
||||
|
||||
alphabet = (
|
||||
string.digits
|
||||
+ string.ascii_lowercase
|
||||
+ string.ascii_uppercase
|
||||
+ string.punctuation
|
||||
+ " "
|
||||
) # len(aphabet) = 95
|
||||
"""alphabet
|
||||
0123456789abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ!"#$%&\'()*+,-./:;<=>?@[\\]^_`{|}~
|
||||
"""
|
||||
|
||||
font_layout = ImageFont.truetype("./Arial.ttf", 16)
|
||||
|
||||
|
||||
class ModelOutput(BaseModel):
|
||||
output_images: list[Path]
|
||||
composed_prompt: str
|
||||
|
||||
|
||||
class Predictor(BasePredictor):
|
||||
def setup(self) -> None:
|
||||
"""Load the model into memory to make running multiple predictions efficient"""
|
||||
cache_dir = "model_cache"
|
||||
local_files_only = True # set to True if the models are saved in cache_dir
|
||||
|
||||
self.m1_model_path = "JingyeChen22/textdiffuser2_layout_planner"
|
||||
self.m1_tokenizer = AutoTokenizer.from_pretrained(
|
||||
self.m1_model_path,
|
||||
use_fast=False,
|
||||
cache_dir=cache_dir,
|
||||
local_files_only=local_files_only,
|
||||
)
|
||||
self.m1_model = AutoModelForCausalLM.from_pretrained(
|
||||
self.m1_model_path,
|
||||
torch_dtype=torch.float16,
|
||||
low_cpu_mem_usage=True,
|
||||
cache_dir=cache_dir,
|
||||
local_files_only=local_files_only,
|
||||
).cuda()
|
||||
|
||||
self.text_encoder = (
|
||||
CLIPTextModel.from_pretrained(
|
||||
"JingyeChen22/textdiffuser2-full-ft",
|
||||
subfolder="text_encoder",
|
||||
cache_dir=cache_dir,
|
||||
local_files_only=local_files_only,
|
||||
)
|
||||
.cuda()
|
||||
.half()
|
||||
)
|
||||
self.tokenizer = CLIPTokenizer.from_pretrained(
|
||||
"runwayml/stable-diffusion-v1-5",
|
||||
subfolder="tokenizer",
|
||||
cache_dir=cache_dir,
|
||||
local_files_only=local_files_only,
|
||||
)
|
||||
|
||||
#### additional tokens are introduced, including coordinate tokens and character tokens
|
||||
print("***************")
|
||||
print(f"tokenizer size: {len(self.tokenizer)}")
|
||||
for i in range(520):
|
||||
self.tokenizer.add_tokens(["l" + str(i)]) # left
|
||||
self.tokenizer.add_tokens(["t" + str(i)]) # top
|
||||
self.tokenizer.add_tokens(["r" + str(i)]) # width
|
||||
self.tokenizer.add_tokens(["b" + str(i)]) # height
|
||||
for c in alphabet:
|
||||
self.tokenizer.add_tokens([f"[{c}]"])
|
||||
print(f"new tokenizer size: {len(self.tokenizer)}")
|
||||
print("***************")
|
||||
|
||||
self.vae = (
|
||||
AutoencoderKL.from_pretrained(
|
||||
"runwayml/stable-diffusion-v1-5",
|
||||
subfolder="vae",
|
||||
cache_dir=cache_dir,
|
||||
local_files_only=local_files_only,
|
||||
)
|
||||
.half()
|
||||
.cuda()
|
||||
)
|
||||
self.unet = (
|
||||
UNet2DConditionModel.from_pretrained(
|
||||
"JingyeChen22/textdiffuser2-full-ft",
|
||||
subfolder="unet",
|
||||
cache_dir=cache_dir,
|
||||
local_files_only=local_files_only,
|
||||
)
|
||||
.half()
|
||||
.cuda()
|
||||
)
|
||||
self.text_encoder.resize_token_embeddings(len(self.tokenizer))
|
||||
self.scheduler = DDPMScheduler.from_pretrained(
|
||||
"runwayml/stable-diffusion-v1-5",
|
||||
subfolder="scheduler",
|
||||
cache_dir=cache_dir,
|
||||
local_files_only=local_files_only,
|
||||
)
|
||||
|
||||
#### load lcm components
|
||||
self.pipe = DiffusionPipeline.from_pretrained(
|
||||
"lambdalabs/sd-pokemon-diffusers",
|
||||
unet=copy.deepcopy(self.unet),
|
||||
tokenizer=self.tokenizer,
|
||||
text_encoder=copy.deepcopy(self.text_encoder),
|
||||
torch_dtype=torch.float16,
|
||||
cache_dir=cache_dir,
|
||||
local_files_only=local_files_only,
|
||||
)
|
||||
self.pipe.scheduler = LCMScheduler.from_config(self.pipe.scheduler.config)
|
||||
self.pipe.load_lora_weights("latent-consistency/lcm-lora-sdv1-5")
|
||||
self.pipe.to(device="cuda")
|
||||
|
||||
def predict(
|
||||
self,
|
||||
prompt: str = Input(
|
||||
description="Input Prompt. You can let language model automatically identify keywords, or provide them below.",
|
||||
default="A beautiful city skyline stamp of Shanghai",
|
||||
),
|
||||
keywords: str = Input(
|
||||
description="(Optional) Keywords. Should be seperated by / (e.g., keyword1/keyword2/...).",
|
||||
default=None,
|
||||
),
|
||||
positive_prompt: str = Input(
|
||||
description="(Optional) Positive prompt.",
|
||||
default=", digital art, very detailed, fantasy, high definition, cinematic light, dnd, trending on artstation",
|
||||
),
|
||||
use_lcm: bool = Input(
|
||||
description="Use Latent Consistent Model.", default=False
|
||||
),
|
||||
generate_natural_image: bool = Input(
|
||||
description="If set to True, the text position and content info will not be incorporated.",
|
||||
default=False,
|
||||
),
|
||||
num_images: int = Input(
|
||||
description="Number of Output images.", default=1, ge=1, le=4
|
||||
),
|
||||
num_inference_steps: int = Input(
|
||||
description="Number of denoising steps. You may decease the step to 4 when using LCM.",
|
||||
ge=1,
|
||||
le=50,
|
||||
default=20,
|
||||
),
|
||||
guidance_scale: float = Input(
|
||||
description="Scale for classifier-free guidance. The scale is set to 7.5 by default. When using LCM, guidance_scale is set to 1.",
|
||||
ge=1,
|
||||
le=20,
|
||||
default=7.5,
|
||||
),
|
||||
temperature: float = Input(
|
||||
description="Control the diversity of layout planner. Higher value indicates more diversity.",
|
||||
ge=0.1,
|
||||
le=2,
|
||||
default=1.4,
|
||||
),
|
||||
) -> ModelOutput:
|
||||
"""Run a single prediction on the model"""
|
||||
if positive_prompt is not None and not len(positive_prompt.strip()) == 0:
|
||||
prompt += positive_prompt
|
||||
|
||||
with torch.no_grad():
|
||||
user_prompt = prompt
|
||||
|
||||
if generate_natural_image:
|
||||
composed_prompt = user_prompt
|
||||
prompt = self.tokenizer.encode(user_prompt)
|
||||
else:
|
||||
if keywords is None or len(keywords.strip()) == 0:
|
||||
template = f"Given a prompt that will be used to generate an image, plan the layout of visual text for the image. The size of the image is 128x128. Therefore, all properties of the positions should not exceed 128, including the coordinates of top, left, right, and bottom. All keywords are included in the caption. You dont need to specify the details of font styles. At each line, the format should be keyword left, top, right, bottom. So let us begin. Prompt: {user_prompt}"
|
||||
else:
|
||||
keywords = keywords.split("/")
|
||||
keywords = [i.strip() for i in keywords]
|
||||
template = f"Given a prompt that will be used to generate an image, plan the layout of visual text for the image. The size of the image is 128x128. Therefore, all properties of the positions should not exceed 128, including the coordinates of top, left, right, and bottom. In addition, we also provide all keywords at random order for reference. You dont need to specify the details of font styles. At each line, the format should be keyword left, top, right, bottom. So let us begin. Prompt: {prompt}. Keywords: {str(keywords)}"
|
||||
|
||||
msg = template
|
||||
conv = get_conversation_template(self.m1_model_path)
|
||||
conv.append_message(conv.roles[0], msg)
|
||||
conv.append_message(conv.roles[1], None)
|
||||
prompt = conv.get_prompt()
|
||||
inputs = self.m1_tokenizer([prompt], return_token_type_ids=False)
|
||||
inputs = {k: torch.tensor(v).to("cuda") for k, v in inputs.items()}
|
||||
output_ids = self.m1_model.generate(
|
||||
**inputs,
|
||||
do_sample=True,
|
||||
temperature=temperature,
|
||||
repetition_penalty=1.0,
|
||||
max_new_tokens=512,
|
||||
)
|
||||
|
||||
if self.m1_model.config.is_encoder_decoder:
|
||||
output_ids = output_ids[0]
|
||||
else:
|
||||
output_ids = output_ids[0][len(inputs["input_ids"][0]) :]
|
||||
|
||||
outputs = self.m1_tokenizer.decode(
|
||||
output_ids,
|
||||
skip_special_tokens=True,
|
||||
spaces_between_special_tokens=False,
|
||||
)
|
||||
print(f"[{conv.roles[0]}]\n{msg}")
|
||||
print(f"[{conv.roles[1]}]\n{outputs}")
|
||||
|
||||
ocrs = outputs.split("\n")
|
||||
current_ocr = ocrs
|
||||
|
||||
ocr_ids = []
|
||||
print("user_prompt", user_prompt)
|
||||
print("current_ocr", current_ocr)
|
||||
|
||||
for ocr in current_ocr:
|
||||
ocr = ocr.strip()
|
||||
|
||||
if len(ocr) == 0 or "###" in ocr or ".com" in ocr:
|
||||
continue
|
||||
|
||||
items = ocr.split()
|
||||
pred = " ".join(items[:-1])
|
||||
box = items[-1]
|
||||
|
||||
l, t, r, b = box.split(",")
|
||||
l, t, r, b = int(l), int(t), int(r), int(b)
|
||||
ocr_ids.extend(
|
||||
["l" + str(l), "t" + str(t), "r" + str(r), "b" + str(b)]
|
||||
)
|
||||
|
||||
char_list = list(pred)
|
||||
char_list = [f"[{i}]" for i in char_list]
|
||||
ocr_ids.extend(char_list)
|
||||
ocr_ids.append(self.tokenizer.eos_token_id)
|
||||
|
||||
caption_ids = (
|
||||
self.tokenizer(user_prompt, truncation=True, return_tensors="pt")
|
||||
.input_ids[0]
|
||||
.tolist()
|
||||
)
|
||||
|
||||
try:
|
||||
ocr_ids = self.tokenizer.encode(ocr_ids)
|
||||
prompt = caption_ids + ocr_ids
|
||||
except:
|
||||
prompt = caption_ids
|
||||
|
||||
user_prompt = self.tokenizer.decode(prompt)
|
||||
composed_prompt = self.tokenizer.decode(prompt)
|
||||
|
||||
prompt = prompt[:77]
|
||||
while len(prompt) < 77:
|
||||
prompt.append(self.tokenizer.pad_token_id)
|
||||
|
||||
if not use_lcm:
|
||||
prompts_cond = prompt
|
||||
prompts_nocond = [self.tokenizer.pad_token_id] * 77
|
||||
|
||||
prompts_cond = [prompts_cond] * num_images
|
||||
prompts_nocond = [prompts_nocond] * num_images
|
||||
|
||||
prompts_cond = torch.Tensor(prompts_cond).long().cuda()
|
||||
prompts_nocond = torch.Tensor(prompts_nocond).long().cuda()
|
||||
|
||||
scheduler = self.scheduler
|
||||
scheduler.set_timesteps(num_inference_steps)
|
||||
noise = torch.randn((num_images, 4, 64, 64)).to("cuda").half()
|
||||
input = noise
|
||||
|
||||
encoder_hidden_states_cond = self.text_encoder(prompts_cond)[0].half()
|
||||
encoder_hidden_states_nocond = self.text_encoder(prompts_nocond)[
|
||||
0
|
||||
].half()
|
||||
|
||||
for t in tqdm(scheduler.timesteps):
|
||||
with torch.no_grad(): # classifier free guidance
|
||||
noise_pred_cond = self.unet(
|
||||
sample=input,
|
||||
timestep=t,
|
||||
encoder_hidden_states=encoder_hidden_states_cond[
|
||||
:num_images
|
||||
],
|
||||
).sample # b, 4, 64, 64
|
||||
noise_pred_uncond = self.unet(
|
||||
sample=input,
|
||||
timestep=t,
|
||||
encoder_hidden_states=encoder_hidden_states_nocond[
|
||||
:num_images
|
||||
],
|
||||
).sample # b, 4, 64, 64
|
||||
noisy_residual = noise_pred_uncond + guidance_scale * (
|
||||
noise_pred_cond - noise_pred_uncond
|
||||
) # b, 4, 64, 64
|
||||
input = scheduler.step(noisy_residual, t, input).prev_sample
|
||||
del noise_pred_cond
|
||||
del noise_pred_uncond
|
||||
|
||||
torch.cuda.empty_cache()
|
||||
|
||||
# decode
|
||||
input = 1 / self.vae.config.scaling_factor * input
|
||||
images = self.vae.decode(input, return_dict=False)[0]
|
||||
width, height = 512, 512
|
||||
results = []
|
||||
new_image = Image.new("RGB", (2 * width, 2 * height))
|
||||
for index, image in enumerate(images.cpu().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")
|
||||
results.append(image)
|
||||
row = index // 2
|
||||
col = index % 2
|
||||
new_image.paste(image, (col * width, row * height))
|
||||
else:
|
||||
generator = torch.Generator(device=self.pipe.device).manual_seed(
|
||||
random.randint(0, 1000)
|
||||
)
|
||||
results = self.pipe(
|
||||
prompt=user_prompt,
|
||||
generator=generator,
|
||||
num_inference_steps=num_inference_steps,
|
||||
guidance_scale=1,
|
||||
num_images_per_prompt=num_images,
|
||||
).images
|
||||
|
||||
torch.cuda.empty_cache()
|
||||
output_paths = []
|
||||
|
||||
for i, sample in enumerate(results):
|
||||
output_path = f"/tmp/out-{i}.png"
|
||||
sample.save(output_path)
|
||||
output_paths.append(Path(output_path))
|
||||
|
||||
return ModelOutput(
|
||||
output_images=output_paths,
|
||||
composed_prompt=composed_prompt,
|
||||
)
|
||||
@@ -0,0 +1,6 @@
|
||||
diffusers==0.24.0
|
||||
accelerate==0.22.0
|
||||
datasets==2.11.0
|
||||
transformers==4.28.1
|
||||
fschat==0.2.26
|
||||
|
||||
@@ -0,0 +1,24 @@
|
||||
CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 torchrun --nproc_per_node=8 --master_port=20003 fastchat/train/train_mem.py \
|
||||
--model_name_or_path lmsys/vicuna-7b-v1.5 \
|
||||
--data_path data/layout_planner_data_5k.json \
|
||||
--bf16 True \
|
||||
--output_dir experiment_result \
|
||||
--num_train_epochs 6 \
|
||||
--per_device_train_batch_size 2 \
|
||||
--per_device_eval_batch_size 2 \
|
||||
--gradient_accumulation_steps 16 \
|
||||
--evaluation_strategy "no" \
|
||||
--save_strategy "steps" \
|
||||
--save_steps 500 \
|
||||
--save_total_limit 5 \
|
||||
--learning_rate 2e-5 \
|
||||
--weight_decay 0. \
|
||||
--warmup_ratio 0.03 \
|
||||
--lr_scheduler_type "cosine" \
|
||||
--logging_steps 1 \
|
||||
--fsdp "full_shard auto_wrap" \
|
||||
--fsdp_transformer_layer_cls_to_wrap 'LlamaDecoderLayer' \
|
||||
--tf32 True \
|
||||
--model_max_length 2048 \
|
||||
--gradient_checkpointing True \
|
||||
--lazy_preprocess True
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,17 @@
|
||||
accelerate launch train_textdiffuser2_inpainting_full.py \
|
||||
--train_batch_size=16 \
|
||||
--gradient_accumulation_steps=4 \
|
||||
--gradient_checkpointing \
|
||||
--mixed_precision="fp16" \
|
||||
--num_train_epochs=6 \
|
||||
--learning_rate=1e-5 \
|
||||
--max_grad_norm=1 \
|
||||
--lr_scheduler="constant" \
|
||||
--lr_warmup_steps=0 \
|
||||
--output_dir="inpainting_experiment_name" \
|
||||
--enable_xformers_memory_efficient_attention \
|
||||
--dataloader_num_workers=8 \
|
||||
--resume_from_checkpoint="latest" \
|
||||
--dataset_path=/path/to/laion-ocr-select \
|
||||
--train_dataset_index_file=/path/to/train_dataset_index.txt \
|
||||
--vis_num=16
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,20 @@
|
||||
accelerate launch train_textdiffuser2_t2i_full.py \
|
||||
--pretrained_model_name_or_path="runwayml/stable-diffusion-v1-5" \
|
||||
--train_batch_size=18 \
|
||||
--gradient_accumulation_steps=4 \
|
||||
--gradient_checkpointing \
|
||||
--mixed_precision="fp16" \
|
||||
--num_train_epochs=6 \
|
||||
--learning_rate=1e-5 \
|
||||
--max_grad_norm=1 \
|
||||
--lr_scheduler="constant" \
|
||||
--lr_warmup_steps=0 \
|
||||
--output_dir="diffusion_experiment_result" \
|
||||
--enable_xformers_memory_efficient_attention \
|
||||
--dataloader_num_workers=8 \
|
||||
--index_file_path='/path/to/train_dataset_index.txt' \
|
||||
--dataset_path='/path/to/laion-ocr-select/' \
|
||||
--granularity=128 \
|
||||
--coord_mode="ltrb" \
|
||||
--max_length=77 \
|
||||
--resume_from_checkpoint="latest"
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,19 @@
|
||||
accelerate launch train_textdiffuser2_t2i_lora.py \
|
||||
--pretrained_model_name_or_path="runwayml/stable-diffusion-v1-5" \
|
||||
--train_batch_size=18 \
|
||||
--gradient_accumulation_steps=4 \
|
||||
--gradient_checkpointing \
|
||||
--mixed_precision="fp16" \
|
||||
--num_train_epochs=6 \
|
||||
--learning_rate=1e-4 \
|
||||
--text_encoder_learning_rate=1e-5 \
|
||||
--lr_scheduler="constant" \
|
||||
--output_dir="diffusion_experiment_result" \
|
||||
--enable_xformers_memory_efficient_attention \
|
||||
--dataloader_num_workers=8 \
|
||||
--index_file_path='/path/to/train_dataset_index.txt' \
|
||||
--dataset_path='/path/to/laion-ocr-select/' \
|
||||
--granularity=128 \
|
||||
--coord_mode="ltrb" \
|
||||
--max_length=77 \
|
||||
--resume_from_checkpoint="latest"
|
||||
Reference in New Issue
Block a user