Liang Qu
commited on
Commit
·
f2de1ca
1
Parent(s):
6597142
Initial commit.
Browse files- README.md +63 -2
- app.py +147 -59
- images/style001.png +0 -0
- images/style002.png +0 -0
- images/style003.png +0 -0
- images/style004.png +0 -0
- images/style005.png +0 -0
- images/style006.png +0 -0
- images/style007.png +0 -0
- images/style008.png +0 -0
- images/style009.png +0 -0
- images/style010.png +0 -0
- images/style011.png +0 -0
- images/style012.png +0 -0
- images/style013.png +0 -0
- images/style014.png +0 -0
- images/style015.png +0 -0
- images/style016.png +0 -0
- images/style017.png +0 -0
- images/style018.png +0 -0
- images/style019.png +0 -0
- images/style020.png +0 -0
- images/style021.png +0 -0
- images/style022.png +0 -0
- images/style023.png +0 -0
- images/style024.png +0 -0
- images/style025.png +0 -0
- images/style026.png +0 -0
- images/style027.png +0 -0
- images/style028.png +0 -0
- images/style029.png +0 -0
- requirements.txt +9 -1
- utils.py +108 -0
README.md
CHANGED
@@ -8,7 +8,68 @@ sdk_version: 5.25.0
|
|
8 |
app_file: app.py
|
9 |
pinned: false
|
10 |
license: openrail
|
11 |
-
short_description: Unwritten Chinese Charecters
|
12 |
---
|
13 |
|
14 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
8 |
app_file: app.py
|
9 |
pinned: false
|
10 |
license: openrail
|
11 |
+
short_description: Unwritten Chinese Charecters in Style
|
12 |
---
|
13 |
|
14 |
+
# What is this?
|
15 |
+
|
16 |
+
Generate New Characters by combining parts in creative ways. Write them in a controlled style.
|
17 |
+
|
18 |
+
- Inspired by
|
19 |
+
- Lin Yutang's [Ming-Kwai typewriter](https://en.wikipedia.org/wiki/Chinese_typewriter#MingKwai_design)
|
20 |
+
- Wu Yue's [Glyffuser](https://yue-here.com/posts/glyffuser/)
|
21 |
+
|
22 |
+
# Why
|
23 |
+
|
24 |
+
- Fun to generate valid but unseen characters. (Never in a dictionary, nor Unicode).
|
25 |
+
- Implements Lin Yutang's ideas with AI/ML, without the mechanical marvel :-/ or limitations :-)
|
26 |
+
- Extends a font to support new charsets, and beyond to non-existent chars.
|
27 |
+
- Adds variation/diversity/personality to generated images. No boring duplicates from the same char.
|
28 |
+
- Other [Creative Uses](#creative-uses)
|
29 |
+
|
30 |
+
# How to use this app
|
31 |
+
- Combine components or radicals in the following way
|
32 |
+
- Specify the 'Structure' and 'Components', in a [Polish Notation](https://en.wikipedia.org/wiki/Polish_notation) fashion - Good for tree structures
|
33 |
+
- ⿰: 'LR' Left-Rigth
|
34 |
+
- ⿱: 'TB' Top-Bottom
|
35 |
+
- ⿸: 'TL' Top-Left
|
36 |
+
- ⿹: 'TR' Top-Right
|
37 |
+
- ⿺: 'BL' Bottom-Left
|
38 |
+
- ⿴: 'OI' Outer-Inner
|
39 |
+
- ⿻: 'OV' Overlap
|
40 |
+
- ⿲: 'LMR' Left-Middle-Right
|
41 |
+
- ⿳: 'TMB' Top-Middle-Bottom
|
42 |
+
- ⿵: 'BT' Bottom Open Enclosure
|
43 |
+
- ⿶: 'CT' Top Open Enclosure
|
44 |
+
- ⿷: 'RT' Right Open Enclosure
|
45 |
+
- Select a 'Style' by clicking the sample images
|
46 |
+
- Hit the 'Generate' button
|
47 |
+
- Repeat
|
48 |
+
|
49 |
+
# Usage Tips
|
50 |
+
- Simple structures work best (⿰ ⿱ ⿴ etc.)
|
51 |
+
- "Known radicals at seen positions" work best (釒on left better than right, but may also surprise you in a good way)
|
52 |
+
- Noto font family (sans and serif) gives the best results, as there are many training examples
|
53 |
+
- Cursive and handwritten styles usually give good results, as they are more tolerant
|
54 |
+
- Fonts supporting less chars are challenging
|
55 |
+
- Current model was trained with 300k samples for only 20 epochs
|
56 |
+
- Training will continue if this app gets attention or likes
|
57 |
+
|
58 |
+
- For dictionary chars, [decompose](https://github.com/cburgmer/cjklib/blob/master/cjklib/data/characterdecomposition.csv) first.
|
59 |
+
- For a part is hard to describe, or you don't care, use '?' (full-width question mark, or does it matter?)
|
60 |
+
|
61 |
+
- What to do when the results are not as expected
|
62 |
+
- Pick a different 'sytle' which may have trained the model better
|
63 |
+
- Try again with a different random seed. This will change the overall structure in an unpredictable way
|
64 |
+
- Try again with a different 'step' number. This will change the local details in a continuous way
|
65 |
+
|
66 |
+
# Creative Uses
|
67 |
+
## Turning a bug into a feature
|
68 |
+
When you see a funny result you didn't expect (5 or 3 dots while it should be 4), don't throw it away immediately.
|
69 |
+
- Save the results to confuse/train OCR
|
70 |
+
- 3vade 3vil c3nsorship
|
71 |
+
- Share in discussion. The input text/seed/step will reliably reproduce the result.
|
72 |
+
|
73 |
+
# Future Features
|
74 |
+
- Typewriter keyboard for hard-to-input radicals, filtered by pinyin prefix
|
75 |
+
- Direct generation of a single char, auto decomposition%
|
app.py
CHANGED
@@ -1,36 +1,39 @@
|
|
1 |
-
import
|
2 |
-
import
|
3 |
import random
|
|
|
|
|
|
|
|
|
4 |
|
5 |
-
|
6 |
from diffusers import DiffusionPipeline
|
|
|
7 |
import torch
|
8 |
|
|
|
|
|
|
|
9 |
device = "cuda" if torch.cuda.is_available() else "cpu"
|
10 |
-
model_repo_id = "stabilityai/sdxl-turbo" # Replace to the model you would like to use
|
11 |
|
12 |
-
|
13 |
-
|
14 |
-
else:
|
15 |
-
torch_dtype = torch.float32
|
16 |
|
17 |
-
|
18 |
-
pipe = pipe.to(device)
|
19 |
|
20 |
-
|
21 |
-
MAX_IMAGE_SIZE = 1024
|
22 |
|
|
|
|
|
23 |
|
24 |
-
|
|
|
25 |
def infer(
|
26 |
prompt,
|
27 |
negative_prompt,
|
28 |
seed,
|
29 |
randomize_seed,
|
30 |
-
|
31 |
-
height,
|
32 |
-
guidance_scale,
|
33 |
-
num_inference_steps,
|
34 |
progress=gr.Progress(track_tqdm=True),
|
35 |
):
|
36 |
if randomize_seed:
|
@@ -39,34 +42,140 @@ def infer(
|
|
39 |
generator = torch.Generator().manual_seed(seed)
|
40 |
|
41 |
image = pipe(
|
42 |
-
prompt
|
43 |
-
|
44 |
-
guidance_scale=guidance_scale,
|
45 |
-
num_inference_steps=num_inference_steps,
|
46 |
-
width=width,
|
47 |
-
height=height,
|
48 |
generator=generator,
|
|
|
49 |
).images[0]
|
50 |
|
51 |
return image, seed
|
52 |
|
53 |
|
54 |
examples = [
|
55 |
-
"
|
56 |
-
"
|
57 |
-
"
|
|
|
|
|
|
|
|
|
|
|
|
|
58 |
]
|
59 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
60 |
css = """
|
61 |
#col-container {
|
62 |
margin: 0 auto;
|
63 |
-
max-width:
|
64 |
}
|
65 |
"""
|
66 |
|
67 |
with gr.Blocks(css=css) as demo:
|
68 |
with gr.Column(elem_id="col-container"):
|
69 |
-
gr.Markdown(" #
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
70 |
|
71 |
with gr.Row():
|
72 |
prompt = gr.Text(
|
@@ -76,9 +185,14 @@ with gr.Blocks(css=css) as demo:
|
|
76 |
placeholder="Enter your prompt",
|
77 |
container=False,
|
78 |
)
|
79 |
-
|
80 |
run_button = gr.Button("Run", scale=0, variant="primary")
|
81 |
|
|
|
|
|
|
|
|
|
|
|
|
|
82 |
result = gr.Image(label="Result", show_label=False)
|
83 |
|
84 |
with gr.Accordion("Advanced Settings", open=False):
|
@@ -100,40 +214,16 @@ with gr.Blocks(css=css) as demo:
|
|
100 |
randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
|
101 |
|
102 |
with gr.Row():
|
103 |
-
width = gr.Slider(
|
104 |
-
label="Width",
|
105 |
-
minimum=256,
|
106 |
-
maximum=MAX_IMAGE_SIZE,
|
107 |
-
step=32,
|
108 |
-
value=1024, # Replace with defaults that work for your model
|
109 |
-
)
|
110 |
-
|
111 |
-
height = gr.Slider(
|
112 |
-
label="Height",
|
113 |
-
minimum=256,
|
114 |
-
maximum=MAX_IMAGE_SIZE,
|
115 |
-
step=32,
|
116 |
-
value=1024, # Replace with defaults that work for your model
|
117 |
-
)
|
118 |
-
|
119 |
-
with gr.Row():
|
120 |
-
guidance_scale = gr.Slider(
|
121 |
-
label="Guidance scale",
|
122 |
-
minimum=0.0,
|
123 |
-
maximum=10.0,
|
124 |
-
step=0.1,
|
125 |
-
value=0.0, # Replace with defaults that work for your model
|
126 |
-
)
|
127 |
-
|
128 |
num_inference_steps = gr.Slider(
|
129 |
label="Number of inference steps",
|
130 |
minimum=1,
|
131 |
-
maximum=
|
132 |
step=1,
|
133 |
-
value=
|
134 |
)
|
135 |
|
136 |
gr.Examples(examples=examples, inputs=[prompt])
|
|
|
137 |
gr.on(
|
138 |
triggers=[run_button.click, prompt.submit],
|
139 |
fn=infer,
|
@@ -142,9 +232,6 @@ with gr.Blocks(css=css) as demo:
|
|
142 |
negative_prompt,
|
143 |
seed,
|
144 |
randomize_seed,
|
145 |
-
width,
|
146 |
-
height,
|
147 |
-
guidance_scale,
|
148 |
num_inference_steps,
|
149 |
],
|
150 |
outputs=[result, seed],
|
@@ -152,3 +239,4 @@ with gr.Blocks(css=css) as demo:
|
|
152 |
|
153 |
if __name__ == "__main__":
|
154 |
demo.launch()
|
|
|
|
1 |
+
import os
|
2 |
+
import re
|
3 |
import random
|
4 |
+
import numpy as np
|
5 |
+
|
6 |
+
# !!! spaces must be imported before torch/CUDA
|
7 |
+
import spaces
|
8 |
|
9 |
+
from huggingface_hub import login
|
10 |
from diffusers import DiffusionPipeline
|
11 |
+
import gradio as gr
|
12 |
import torch
|
13 |
|
14 |
+
from utils import QPipeline
|
15 |
+
|
16 |
+
|
17 |
device = "cuda" if torch.cuda.is_available() else "cpu"
|
|
|
18 |
|
19 |
+
login(token=os.environ["HF_TOKEN"])
|
20 |
+
model_repo_id = os.environ["MODEL_ID"]
|
|
|
|
|
21 |
|
22 |
+
torch_dtype = torch.float16 if torch.cuda.is_available() else torch.float32
|
|
|
23 |
|
24 |
+
pipe = QPipeline.from_pretrained(model_repo_id, torch_dtype=torch_dtype).to(device)
|
|
|
25 |
|
26 |
+
MAX_SEED = 65535
|
27 |
+
MAX_IMAGE_SIZE = 128
|
28 |
|
29 |
+
|
30 |
+
@spaces.GPU # Enable ZeroGPU if needed
|
31 |
def infer(
|
32 |
prompt,
|
33 |
negative_prompt,
|
34 |
seed,
|
35 |
randomize_seed,
|
36 |
+
num_inference_steps=10,
|
|
|
|
|
|
|
37 |
progress=gr.Progress(track_tqdm=True),
|
38 |
):
|
39 |
if randomize_seed:
|
|
|
42 |
generator = torch.Generator().manual_seed(seed)
|
43 |
|
44 |
image = pipe(
|
45 |
+
[prompt],
|
46 |
+
batch_size=1,
|
|
|
|
|
|
|
|
|
47 |
generator=generator,
|
48 |
+
num_inference_steps=num_inference_steps
|
49 |
).images[0]
|
50 |
|
51 |
return image, seed
|
52 |
|
53 |
|
54 |
examples = [
|
55 |
+
"Structure: (LR 文 英). Style: style001",
|
56 |
+
"Structure: (TL 广 東). Style: style028",
|
57 |
+
"Structure: (TB 艹 (LR 禾 魚)). Style: style015",
|
58 |
+
"Structure: (TB 敬 音). Style: style013",
|
59 |
+
"Structure: (LR 釒 馬). Style: style018",
|
60 |
+
"Structure: (BL 走 羽). Style: style022",
|
61 |
+
"Structure: (LR 羊 大). Style: style005",
|
62 |
+
"Structure: (LR 鹿 孚). Style: style017",
|
63 |
+
"Structure: (OI 口 也). Style: style002",
|
64 |
]
|
65 |
|
66 |
+
# Map style images to style names (use real image files later)
|
67 |
+
style_options = {
|
68 |
+
"images/style001.png": "style001",
|
69 |
+
"images/style002.png": "style002",
|
70 |
+
"images/style003.png": "style003",
|
71 |
+
"images/style004.png": "style004",
|
72 |
+
"images/style005.png": "style005",
|
73 |
+
"images/style006.png": "style006",
|
74 |
+
"images/style007.png": "style007",
|
75 |
+
"images/style008.png": "style008",
|
76 |
+
"images/style009.png": "style009",
|
77 |
+
"images/style010.png": "style010",
|
78 |
+
"images/style011.png": "style011",
|
79 |
+
"images/style012.png": "style012",
|
80 |
+
"images/style013.png": "style013",
|
81 |
+
"images/style014.png": "style014",
|
82 |
+
"images/style015.png": "style015",
|
83 |
+
# "images/style016.png": "style016", very similar to 002
|
84 |
+
"images/style017.png": "style017",
|
85 |
+
"images/style018.png": "style018",
|
86 |
+
"images/style019.png": "style019",
|
87 |
+
"images/style020.png": "style020",
|
88 |
+
"images/style021.png": "style021",
|
89 |
+
"images/style022.png": "style022",
|
90 |
+
"images/style023.png": "style023",
|
91 |
+
"images/style024.png": "style024",
|
92 |
+
"images/style025.png": "style025",
|
93 |
+
"images/style026.png": "style026",
|
94 |
+
"images/style027.png": "style027",
|
95 |
+
"images/style028.png": "style028",
|
96 |
+
"images/style029.png": "style029",
|
97 |
+
}
|
98 |
+
|
99 |
+
|
100 |
+
def apply_style_on_click(evt: gr.SelectData, prompt_text):
|
101 |
+
index = evt.index
|
102 |
+
style_label = list(style_options.values())[index]
|
103 |
+
|
104 |
+
if re.search(r"Style: [^\n]+", prompt_text):
|
105 |
+
return re.sub(r"Style: [^\n]+", f"Style: {style_label}", prompt_text)
|
106 |
+
else:
|
107 |
+
return prompt_text.strip() + f" Style: {style_label}"
|
108 |
+
|
109 |
+
|
110 |
+
# CSS for fixing Gallery layout
|
111 |
css = """
|
112 |
#col-container {
|
113 |
margin: 0 auto;
|
114 |
+
max-width: 800px;
|
115 |
}
|
116 |
"""
|
117 |
|
118 |
with gr.Blocks(css=css) as demo:
|
119 |
with gr.Column(elem_id="col-container"):
|
120 |
+
gr.Markdown(" # NeoChar ")
|
121 |
+
gr.Markdown(" ## What: Create New Characters. Write them in Style. ")
|
122 |
+
gr.Markdown(" * NO more missing glyphs - Make them when fonts don't support! ")
|
123 |
+
gr.Markdown(" * Create valid and new Hanzi/Kanji that never existed before ")
|
124 |
+
gr.Markdown(" * Calligraphy with diversity - powered by controlled chaos")
|
125 |
+
gr.Markdown(" ## How ")
|
126 |
+
gr.Markdown(" * Specify 'Structure' and 'Components'. LR: Left-Right, TB: Top-Bottom, etc. ")
|
127 |
+
gr.Markdown(" * Select a 'Style' ")
|
128 |
+
gr.Markdown(" * (examples below)")
|
129 |
+
gr.Markdown(" ## README for more ")
|
130 |
+
|
131 |
+
|
132 |
+
gr.HTML("""
|
133 |
+
<style>
|
134 |
+
.gallery-container .gallery-item {
|
135 |
+
width: 60px !important;
|
136 |
+
height: 60px !important;
|
137 |
+
padding: 0 !important;
|
138 |
+
margin: 4px !important;
|
139 |
+
border-radius: 4px;
|
140 |
+
overflow: hidden;
|
141 |
+
background: none !important;
|
142 |
+
box-shadow: none !important;
|
143 |
+
}
|
144 |
+
|
145 |
+
.gallery-container .gallery-item img {
|
146 |
+
width: 64px !important;
|
147 |
+
height: 64px !important;
|
148 |
+
object-fit: cover;
|
149 |
+
display: block;
|
150 |
+
margin: auto;
|
151 |
+
}
|
152 |
+
|
153 |
+
.gallery-container button {
|
154 |
+
all: unset !important;
|
155 |
+
padding: 0 !important;
|
156 |
+
margin: 0 !important;
|
157 |
+
border: none !important;
|
158 |
+
background: none !important;
|
159 |
+
box-shadow: none !important;
|
160 |
+
}
|
161 |
+
|
162 |
+
.gallery__modal,
|
163 |
+
.gallery-container .preview,
|
164 |
+
.gallery-container .gallery-item:focus-visible {
|
165 |
+
display: none !important;
|
166 |
+
pointer-events: none !important;
|
167 |
+
}
|
168 |
+
</style>
|
169 |
+
""")
|
170 |
+
|
171 |
+
gallery = gr.Gallery(
|
172 |
+
value=list(style_options.keys()),
|
173 |
+
label="Click any image",
|
174 |
+
columns=7,
|
175 |
+
allow_preview=False,
|
176 |
+
height=None,
|
177 |
+
elem_classes=["gallery-container"]
|
178 |
+
)
|
179 |
|
180 |
with gr.Row():
|
181 |
prompt = gr.Text(
|
|
|
185 |
placeholder="Enter your prompt",
|
186 |
container=False,
|
187 |
)
|
|
|
188 |
run_button = gr.Button("Run", scale=0, variant="primary")
|
189 |
|
190 |
+
gallery.select(
|
191 |
+
fn=apply_style_on_click,
|
192 |
+
inputs=[prompt],
|
193 |
+
outputs=prompt
|
194 |
+
)
|
195 |
+
|
196 |
result = gr.Image(label="Result", show_label=False)
|
197 |
|
198 |
with gr.Accordion("Advanced Settings", open=False):
|
|
|
214 |
randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
|
215 |
|
216 |
with gr.Row():
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
217 |
num_inference_steps = gr.Slider(
|
218 |
label="Number of inference steps",
|
219 |
minimum=1,
|
220 |
+
maximum=20,
|
221 |
step=1,
|
222 |
+
value=10,
|
223 |
)
|
224 |
|
225 |
gr.Examples(examples=examples, inputs=[prompt])
|
226 |
+
|
227 |
gr.on(
|
228 |
triggers=[run_button.click, prompt.submit],
|
229 |
fn=infer,
|
|
|
232 |
negative_prompt,
|
233 |
seed,
|
234 |
randomize_seed,
|
|
|
|
|
|
|
235 |
num_inference_steps,
|
236 |
],
|
237 |
outputs=[result, seed],
|
|
|
239 |
|
240 |
if __name__ == "__main__":
|
241 |
demo.launch()
|
242 |
+
|
images/style001.png
ADDED
![]() |
images/style002.png
ADDED
![]() |
images/style003.png
ADDED
![]() |
images/style004.png
ADDED
![]() |
images/style005.png
ADDED
![]() |
images/style006.png
ADDED
![]() |
images/style007.png
ADDED
![]() |
images/style008.png
ADDED
![]() |
images/style009.png
ADDED
![]() |
images/style010.png
ADDED
![]() |
images/style011.png
ADDED
![]() |
images/style012.png
ADDED
![]() |
images/style013.png
ADDED
![]() |
images/style014.png
ADDED
![]() |
images/style015.png
ADDED
![]() |
images/style016.png
ADDED
![]() |
images/style017.png
ADDED
![]() |
images/style018.png
ADDED
![]() |
images/style019.png
ADDED
![]() |
images/style020.png
ADDED
![]() |
images/style021.png
ADDED
![]() |
images/style022.png
ADDED
![]() |
images/style023.png
ADDED
![]() |
images/style024.png
ADDED
![]() |
images/style025.png
ADDED
![]() |
images/style026.png
ADDED
![]() |
images/style027.png
ADDED
![]() |
images/style028.png
ADDED
![]() |
images/style029.png
ADDED
![]() |
requirements.txt
CHANGED
@@ -1,6 +1,14 @@
|
|
|
|
1 |
accelerate
|
2 |
diffusers
|
3 |
invisible_watermark
|
4 |
torch
|
5 |
transformers
|
6 |
-
xformers
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
numpy
|
2 |
accelerate
|
3 |
diffusers
|
4 |
invisible_watermark
|
5 |
torch
|
6 |
transformers
|
7 |
+
xformers
|
8 |
+
sentencepiece
|
9 |
+
datasets
|
10 |
+
einops
|
11 |
+
Pillow
|
12 |
+
torchvision
|
13 |
+
tqdm
|
14 |
+
imwatermark
|
utils.py
ADDED
@@ -0,0 +1,108 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import os
|
2 |
+
import numpy as np
|
3 |
+
import torch
|
4 |
+
from torch import nn
|
5 |
+
from torch.utils.data import Dataset, DataLoader
|
6 |
+
from torchvision import transforms as T
|
7 |
+
from PIL import Image as PILImage, ImageDraw, ImageFont
|
8 |
+
from imwatermark import WatermarkEncoder
|
9 |
+
|
10 |
+
from diffusers.pipelines.pipeline_utils import DiffusionPipeline, ImagePipelineOutput
|
11 |
+
from diffusers.utils.torch_utils import randn_tensor
|
12 |
+
from transformers import MT5Tokenizer, MT5EncoderModel
|
13 |
+
from typing import List, Optional, Tuple, Union
|
14 |
+
|
15 |
+
# Determine device and torch dtype
|
16 |
+
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
17 |
+
torch_dtype = torch.float16 if torch.cuda.is_available() else torch.float32
|
18 |
+
|
19 |
+
# Load MT5 tokenizer and encoder (can be replaced with private model + token if needed)
|
20 |
+
tokenizer = MT5Tokenizer.from_pretrained("google/mt5-small", use_safetensors=True)
|
21 |
+
encoder_model = MT5EncoderModel.from_pretrained("google/mt5-small", use_safetensors=True).to(device=device, dtype=torch_dtype)
|
22 |
+
encoder_model.eval()
|
23 |
+
|
24 |
+
class QPipeline(DiffusionPipeline):
|
25 |
+
def __init__(self, unet, scheduler):
|
26 |
+
super().__init__()
|
27 |
+
self.register_modules(unet=unet, scheduler=scheduler)
|
28 |
+
|
29 |
+
def add_watermark(self, img: PILImage.Image) -> PILImage.Image:
|
30 |
+
# Resize image to 256, as 128 is too small for watermark
|
31 |
+
img = img.resize((256, 256), resample=PILImage.BICUBIC)
|
32 |
+
|
33 |
+
watermark_str = os.getenv("WATERMARK_URL", "hf.co/lqume/new-hanzi")
|
34 |
+
encoder = WatermarkEncoder()
|
35 |
+
encoder.set_watermark('bytes', watermark_str.encode('utf-8'))
|
36 |
+
|
37 |
+
# Convert PIL image to NumPy array
|
38 |
+
img_np = np.asarray(img.convert("RGB")) # ensure 3-channel RGB
|
39 |
+
watermarked_np = encoder.encode(img_np, 'dwtDct')
|
40 |
+
|
41 |
+
# Convert back to PIL
|
42 |
+
return PILImage.fromarray(watermarked_np)
|
43 |
+
|
44 |
+
@torch.no_grad()
|
45 |
+
def __call__(
|
46 |
+
self,
|
47 |
+
texts: List[str],
|
48 |
+
batch_size: int = 1,
|
49 |
+
generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
|
50 |
+
num_inference_steps: int = 20,
|
51 |
+
output_type: Optional[str] = "pil",
|
52 |
+
return_dict: bool = True,
|
53 |
+
) -> Union[ImagePipelineOutput, Tuple[List[PILImage.Image]]]:
|
54 |
+
|
55 |
+
batch_size = len(texts)
|
56 |
+
|
57 |
+
# Tokenize input text
|
58 |
+
tokenized = tokenizer(
|
59 |
+
texts,
|
60 |
+
return_tensors="pt",
|
61 |
+
padding="max_length",
|
62 |
+
truncation=True,
|
63 |
+
max_length=48
|
64 |
+
)
|
65 |
+
input_ids = tokenized["input_ids"].to(device=device, dtype=torch.long)
|
66 |
+
attention_mask = tokenized["attention_mask"].to(device=device, dtype=torch.long)
|
67 |
+
|
68 |
+
# Encode to latent space
|
69 |
+
encoded = encoder_model.encoder(input_ids=input_ids, attention_mask=attention_mask)
|
70 |
+
|
71 |
+
# Prepare noise tensor
|
72 |
+
if isinstance(self.unet.config.sample_size, int):
|
73 |
+
image_shape = (
|
74 |
+
batch_size,
|
75 |
+
self.unet.config.in_channels,
|
76 |
+
self.unet.config.sample_size,
|
77 |
+
self.unet.config.sample_size,
|
78 |
+
)
|
79 |
+
else:
|
80 |
+
image_shape = (batch_size, self.unet.config.in_channels, *self.unet.config.sample_size)
|
81 |
+
|
82 |
+
image = randn_tensor(image_shape, generator=generator, device=self.device, dtype=torch_dtype)
|
83 |
+
|
84 |
+
# Run denoising loop
|
85 |
+
self.scheduler.set_timesteps(num_inference_steps)
|
86 |
+
|
87 |
+
for timestep in self.progress_bar(self.scheduler.timesteps):
|
88 |
+
noise_pred = self.unet(
|
89 |
+
image,
|
90 |
+
timestep,
|
91 |
+
encoder_hidden_states=encoded.last_hidden_state,
|
92 |
+
encoder_attention_mask=attention_mask.bool(),
|
93 |
+
return_dict=False
|
94 |
+
)[0]
|
95 |
+
|
96 |
+
image = self.scheduler.step(noise_pred, timestep, image, generator=generator, return_dict=False)[0]
|
97 |
+
|
98 |
+
# Final image post-processing
|
99 |
+
image = image.clamp(0, 1).cpu().permute(0, 2, 3, 1).numpy()
|
100 |
+
if output_type == "pil":
|
101 |
+
image = self.numpy_to_pil(image)
|
102 |
+
image = [self.add_watermark(img) for img in image]
|
103 |
+
|
104 |
+
if not return_dict:
|
105 |
+
return (image,)
|
106 |
+
|
107 |
+
return ImagePipelineOutput(images=image)
|
108 |
+
|