Spaces:
Runtime error
Runtime error
Upload 2 files
Browse files- app_flux.py +327 -0
- app_v1_1.py +252 -0
app_flux.py
ADDED
@@ -0,0 +1,327 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import time
|
2 |
+
|
3 |
+
import gradio as gr
|
4 |
+
import torch
|
5 |
+
from einops import rearrange
|
6 |
+
from PIL import Image
|
7 |
+
|
8 |
+
from flux.sampling import denoise, get_noise, get_schedule, prepare, unpack
|
9 |
+
from flux.util import (
|
10 |
+
SamplingOptions,
|
11 |
+
load_ae,
|
12 |
+
load_clip,
|
13 |
+
load_flow_model,
|
14 |
+
load_flow_model_quintized,
|
15 |
+
load_t5,
|
16 |
+
)
|
17 |
+
from pulid.pipeline_flux import PuLIDPipeline
|
18 |
+
from pulid.utils import resize_numpy_image_long
|
19 |
+
|
20 |
+
|
21 |
+
def get_models(name: str, device: torch.device, offload: bool, fp8: bool):
|
22 |
+
t5 = load_t5(device, max_length=128)
|
23 |
+
clip = load_clip(device)
|
24 |
+
if fp8:
|
25 |
+
model = load_flow_model_quintized(name, device="cpu" if offload else device)
|
26 |
+
else:
|
27 |
+
model = load_flow_model(name, device="cpu" if offload else device)
|
28 |
+
model.eval()
|
29 |
+
ae = load_ae(name, device="cpu" if offload else device)
|
30 |
+
return model, ae, t5, clip
|
31 |
+
|
32 |
+
|
33 |
+
class FluxGenerator:
|
34 |
+
def __init__(self, model_name: str, device: str, offload: bool, aggressive_offload: bool, args):
|
35 |
+
self.device = torch.device(device)
|
36 |
+
self.offload = offload
|
37 |
+
self.aggressive_offload = aggressive_offload
|
38 |
+
self.model_name = model_name
|
39 |
+
self.model, self.ae, self.t5, self.clip = get_models(
|
40 |
+
model_name,
|
41 |
+
device=self.device,
|
42 |
+
offload=self.offload,
|
43 |
+
fp8=args.fp8,
|
44 |
+
)
|
45 |
+
self.pulid_model = PuLIDPipeline(self.model, device="cpu" if offload else device, weight_dtype=torch.bfloat16,
|
46 |
+
onnx_provider=args.onnx_provider)
|
47 |
+
if offload:
|
48 |
+
self.pulid_model.face_helper.face_det.mean_tensor = self.pulid_model.face_helper.face_det.mean_tensor.to(torch.device("cuda"))
|
49 |
+
self.pulid_model.face_helper.face_det.device = torch.device("cuda")
|
50 |
+
self.pulid_model.face_helper.device = torch.device("cuda")
|
51 |
+
self.pulid_model.device = torch.device("cuda")
|
52 |
+
self.pulid_model.load_pretrain(args.pretrained_model, version=args.version)
|
53 |
+
|
54 |
+
@torch.inference_mode()
|
55 |
+
def generate_image(
|
56 |
+
self,
|
57 |
+
width,
|
58 |
+
height,
|
59 |
+
num_steps,
|
60 |
+
start_step,
|
61 |
+
guidance,
|
62 |
+
seed,
|
63 |
+
prompt,
|
64 |
+
id_image=None,
|
65 |
+
id_weight=1.0,
|
66 |
+
neg_prompt="",
|
67 |
+
true_cfg=1.0,
|
68 |
+
timestep_to_start_cfg=1,
|
69 |
+
max_sequence_length=128,
|
70 |
+
):
|
71 |
+
self.t5.max_length = max_sequence_length
|
72 |
+
|
73 |
+
seed = int(seed)
|
74 |
+
if seed == -1:
|
75 |
+
seed = None
|
76 |
+
|
77 |
+
opts = SamplingOptions(
|
78 |
+
prompt=prompt,
|
79 |
+
width=width,
|
80 |
+
height=height,
|
81 |
+
num_steps=num_steps,
|
82 |
+
guidance=guidance,
|
83 |
+
seed=seed,
|
84 |
+
)
|
85 |
+
|
86 |
+
if opts.seed is None:
|
87 |
+
opts.seed = torch.Generator(device="cpu").seed()
|
88 |
+
print(f"Generating '{opts.prompt}' with seed {opts.seed}")
|
89 |
+
t0 = time.perf_counter()
|
90 |
+
|
91 |
+
use_true_cfg = abs(true_cfg - 1.0) > 1e-2
|
92 |
+
|
93 |
+
# prepare input
|
94 |
+
x = get_noise(
|
95 |
+
1,
|
96 |
+
opts.height,
|
97 |
+
opts.width,
|
98 |
+
device=self.device,
|
99 |
+
dtype=torch.bfloat16,
|
100 |
+
seed=opts.seed,
|
101 |
+
)
|
102 |
+
timesteps = get_schedule(
|
103 |
+
opts.num_steps,
|
104 |
+
x.shape[-1] * x.shape[-2] // 4,
|
105 |
+
shift=True,
|
106 |
+
)
|
107 |
+
|
108 |
+
if self.offload:
|
109 |
+
self.t5, self.clip = self.t5.to(self.device), self.clip.to(self.device)
|
110 |
+
inp = prepare(t5=self.t5, clip=self.clip, img=x, prompt=opts.prompt)
|
111 |
+
inp_neg = prepare(t5=self.t5, clip=self.clip, img=x, prompt=neg_prompt) if use_true_cfg else None
|
112 |
+
|
113 |
+
# offload TEs to CPU, load processor models and id encoder to gpu
|
114 |
+
if self.offload:
|
115 |
+
self.t5, self.clip = self.t5.cpu(), self.clip.cpu()
|
116 |
+
torch.cuda.empty_cache()
|
117 |
+
self.pulid_model.components_to_device(torch.device("cuda"))
|
118 |
+
|
119 |
+
if id_image is not None:
|
120 |
+
id_image = resize_numpy_image_long(id_image, 1024)
|
121 |
+
id_embeddings, uncond_id_embeddings = self.pulid_model.get_id_embedding(id_image, cal_uncond=use_true_cfg)
|
122 |
+
else:
|
123 |
+
id_embeddings = None
|
124 |
+
uncond_id_embeddings = None
|
125 |
+
|
126 |
+
# offload processor models and id encoder to CPU, load dit model to gpu
|
127 |
+
if self.offload:
|
128 |
+
self.pulid_model.components_to_device(torch.device("cpu"))
|
129 |
+
torch.cuda.empty_cache()
|
130 |
+
if self.aggressive_offload:
|
131 |
+
self.model.components_to_gpu()
|
132 |
+
else:
|
133 |
+
self.model = self.model.to(self.device)
|
134 |
+
|
135 |
+
# denoise initial noise
|
136 |
+
x = denoise(
|
137 |
+
self.model, **inp, timesteps=timesteps, guidance=opts.guidance, id=id_embeddings, id_weight=id_weight,
|
138 |
+
start_step=start_step, uncond_id=uncond_id_embeddings, true_cfg=true_cfg,
|
139 |
+
timestep_to_start_cfg=timestep_to_start_cfg,
|
140 |
+
neg_txt=inp_neg["txt"] if use_true_cfg else None,
|
141 |
+
neg_txt_ids=inp_neg["txt_ids"] if use_true_cfg else None,
|
142 |
+
neg_vec=inp_neg["vec"] if use_true_cfg else None,
|
143 |
+
aggressive_offload=self.aggressive_offload,
|
144 |
+
)
|
145 |
+
|
146 |
+
# offload model, load autoencoder to gpu
|
147 |
+
if self.offload:
|
148 |
+
self.model.cpu()
|
149 |
+
torch.cuda.empty_cache()
|
150 |
+
self.ae.decoder.to(x.device)
|
151 |
+
|
152 |
+
# decode latents to pixel space
|
153 |
+
x = unpack(x.float(), opts.height, opts.width)
|
154 |
+
with torch.autocast(device_type=self.device.type, dtype=torch.bfloat16):
|
155 |
+
x = self.ae.decode(x)
|
156 |
+
|
157 |
+
if self.offload:
|
158 |
+
self.ae.decoder.cpu()
|
159 |
+
torch.cuda.empty_cache()
|
160 |
+
|
161 |
+
t1 = time.perf_counter()
|
162 |
+
|
163 |
+
print(f"Done in {t1 - t0:.1f}s.")
|
164 |
+
# bring into PIL format
|
165 |
+
x = x.clamp(-1, 1)
|
166 |
+
# x = embed_watermark(x.float())
|
167 |
+
x = rearrange(x[0], "c h w -> h w c")
|
168 |
+
|
169 |
+
img = Image.fromarray((127.5 * (x + 1.0)).cpu().byte().numpy())
|
170 |
+
return img, str(opts.seed), self.pulid_model.debug_img_list
|
171 |
+
|
172 |
+
_HEADER_ = '''
|
173 |
+
<div style="text-align: center; max-width: 650px; margin: 0 auto;">
|
174 |
+
<h1 style="font-size: 2.5rem; font-weight: 700; margin-bottom: 1rem; display: contents;">PuLID for FLUX</h1>
|
175 |
+
<p style="font-size: 1rem; margin-bottom: 1.5rem;">Paper: <a href='https://arxiv.org/abs/2404.16022' target='_blank'>PuLID: Pure and Lightning ID Customization via Contrastive Alignment</a> | Codes: <a href='https://github.com/ToTheBeginning/PuLID' target='_blank'>GitHub</a></p>
|
176 |
+
</div>
|
177 |
+
|
178 |
+
❗️❗️❗️**Tips:**
|
179 |
+
- `timestep to start inserting ID:` The smaller the value, the higher the fidelity, but the lower the editability; the higher the value, the lower the fidelity, but the higher the editability. **The recommended range for this value is between 0 and 4**. For photorealistic scenes, we recommend using 4; for stylized scenes, we recommend using 0-1. If you are not satisfied with the similarity, you can lower this value; conversely, if you are not satisfied with the editability, you can increase this value.
|
180 |
+
- `true CFG scale:` In most scenarios, it is recommended to use a fake CFG, i.e., setting the true CFG scale to 1, and just adjusting the guidance scale. This is also more efficiency. However, in a few cases, utilizing a true CFG can yield better results. For more detaileds, please refer to the [doc](https://github.com/ToTheBeginning/PuLID/blob/main/docs/pulid_for_flux.md#useful-tips).
|
181 |
+
- please refer to the <a href='https://github.com/ToTheBeginning/PuLID/blob/main/docs/pulid_for_flux.md' target='_blank'>github doc</a> for more details and info about the model, we provide the detail explanation about the above two parameters in the doc.
|
182 |
+
- we provide some examples in the bottom, you can try these example prompts first
|
183 |
+
|
184 |
+
''' # noqa E501
|
185 |
+
|
186 |
+
_CITE_ = r"""
|
187 |
+
If PuLID is helpful, please help to ⭐ the <a href='https://github.com/ToTheBeginning/PuLID' target='_blank'> Github Repo</a>. Thanks!
|
188 |
+
---
|
189 |
+
|
190 |
+
📧 **Contact**
|
191 |
+
If you have any questions or feedbacks, feel free to open a discussion or contact <b>[email protected]</b>.
|
192 |
+
""" # noqa E501
|
193 |
+
|
194 |
+
|
195 |
+
def create_demo(args, model_name: str, device: str = "cuda" if torch.cuda.is_available() else "cpu",
|
196 |
+
offload: bool = False, aggressive_offload: bool = False):
|
197 |
+
generator = FluxGenerator(model_name, device, offload, aggressive_offload, args)
|
198 |
+
|
199 |
+
with gr.Blocks() as demo:
|
200 |
+
gr.Markdown(_HEADER_)
|
201 |
+
|
202 |
+
with gr.Row():
|
203 |
+
with gr.Column():
|
204 |
+
prompt = gr.Textbox(label="Prompt", value="portrait, color, cinematic")
|
205 |
+
id_image = gr.Image(label="ID Image")
|
206 |
+
id_weight = gr.Slider(0.0, 3.0, 1, step=0.05, label="id weight")
|
207 |
+
|
208 |
+
width = gr.Slider(256, 1536, 896, step=16, label="Width")
|
209 |
+
height = gr.Slider(256, 1536, 1152, step=16, label="Height")
|
210 |
+
num_steps = gr.Slider(1, 20, 20, step=1, label="Number of steps")
|
211 |
+
start_step = gr.Slider(0, 10, 0, step=1, label="timestep to start inserting ID")
|
212 |
+
guidance = gr.Slider(1.0, 10.0, 4, step=0.1, label="Guidance")
|
213 |
+
seed = gr.Textbox(-1, label="Seed (-1 for random)")
|
214 |
+
max_sequence_length = gr.Slider(128, 512, 128, step=128,
|
215 |
+
label="max_sequence_length for prompt (T5), small will be faster")
|
216 |
+
|
217 |
+
with gr.Accordion("Advanced Options (True CFG, true_cfg_scale=1 means use fake CFG, >1 means use true CFG, if using true CFG, we recommend set the guidance scale to 1)", open=False): # noqa E501
|
218 |
+
neg_prompt = gr.Textbox(
|
219 |
+
label="Negative Prompt",
|
220 |
+
value="bad quality, worst quality, text, signature, watermark, extra limbs")
|
221 |
+
true_cfg = gr.Slider(1.0, 10.0, 1, step=0.1, label="true CFG scale")
|
222 |
+
timestep_to_start_cfg = gr.Slider(0, 20, 1, step=1, label="timestep to start cfg", visible=args.dev)
|
223 |
+
|
224 |
+
generate_btn = gr.Button("Generate")
|
225 |
+
|
226 |
+
with gr.Column():
|
227 |
+
output_image = gr.Image(label="Generated Image")
|
228 |
+
seed_output = gr.Textbox(label="Used Seed")
|
229 |
+
intermediate_output = gr.Gallery(label='Output', elem_id="gallery", visible=args.dev)
|
230 |
+
gr.Markdown(_CITE_)
|
231 |
+
|
232 |
+
with gr.Row(), gr.Column():
|
233 |
+
gr.Markdown("## Examples")
|
234 |
+
example_inps = [
|
235 |
+
[
|
236 |
+
'a woman holding sign with glowing green text \"PuLID for FLUX\"',
|
237 |
+
'example_inputs/liuyifei.png',
|
238 |
+
4, 4, 2680261499100305976, 1
|
239 |
+
],
|
240 |
+
[
|
241 |
+
'portrait, side view',
|
242 |
+
'example_inputs/liuyifei.png',
|
243 |
+
4, 4, 1205240166692517553, 1
|
244 |
+
],
|
245 |
+
[
|
246 |
+
'white-haired woman with vr technology atmosphere, revolutionary exceptional magnum with remarkable details', # noqa E501
|
247 |
+
'example_inputs/liuyifei.png',
|
248 |
+
4, 4, 6349424134217931066, 1
|
249 |
+
],
|
250 |
+
[
|
251 |
+
'a young child is eating Icecream',
|
252 |
+
'example_inputs/liuyifei.png',
|
253 |
+
4, 4, 10606046113565776207, 1
|
254 |
+
],
|
255 |
+
[
|
256 |
+
'a man is holding a sign with text \"PuLID for FLUX\", winter, snowing, top of the mountain',
|
257 |
+
'example_inputs/pengwei.jpg',
|
258 |
+
4, 4, 2410129802683836089, 1
|
259 |
+
],
|
260 |
+
[
|
261 |
+
'portrait, candle light',
|
262 |
+
'example_inputs/pengwei.jpg',
|
263 |
+
4, 4, 17522759474323955700, 1
|
264 |
+
],
|
265 |
+
[
|
266 |
+
'profile shot dark photo of a 25-year-old male with smoke escaping from his mouth, the backlit smoke gives the image an ephemeral quality, natural face, natural eyebrows, natural skin texture, award winning photo, highly detailed face, atmospheric lighting, film grain, monochrome', # noqa E501
|
267 |
+
'example_inputs/pengwei.jpg',
|
268 |
+
4, 4, 17733156847328193625, 1
|
269 |
+
],
|
270 |
+
[
|
271 |
+
'American Comics, 1boy',
|
272 |
+
'example_inputs/pengwei.jpg',
|
273 |
+
1, 4, 13223174453874179686, 1
|
274 |
+
],
|
275 |
+
[
|
276 |
+
'portrait, pixar',
|
277 |
+
'example_inputs/pengwei.jpg',
|
278 |
+
1, 4, 9445036702517583939, 1
|
279 |
+
],
|
280 |
+
]
|
281 |
+
gr.Examples(examples=example_inps, inputs=[prompt, id_image, start_step, guidance, seed, true_cfg],
|
282 |
+
label='fake CFG')
|
283 |
+
|
284 |
+
example_inps = [
|
285 |
+
[
|
286 |
+
'portrait, made of ice sculpture',
|
287 |
+
'example_inputs/lecun.jpg',
|
288 |
+
1, 1, 3811899118709451814, 5
|
289 |
+
],
|
290 |
+
]
|
291 |
+
gr.Examples(examples=example_inps, inputs=[prompt, id_image, start_step, guidance, seed, true_cfg],
|
292 |
+
label='true CFG')
|
293 |
+
|
294 |
+
generate_btn.click(
|
295 |
+
fn=generator.generate_image,
|
296 |
+
inputs=[width, height, num_steps, start_step, guidance, seed, prompt, id_image, id_weight, neg_prompt,
|
297 |
+
true_cfg, timestep_to_start_cfg, max_sequence_length],
|
298 |
+
outputs=[output_image, seed_output, intermediate_output],
|
299 |
+
)
|
300 |
+
|
301 |
+
return demo
|
302 |
+
|
303 |
+
|
304 |
+
if __name__ == "__main__":
|
305 |
+
import argparse
|
306 |
+
|
307 |
+
parser = argparse.ArgumentParser(description="PuLID for FLUX.1-dev")
|
308 |
+
parser.add_argument('--version', type=str, default='v0.9.1', help='version of the model', choices=['v0.9.0', 'v0.9.1'])
|
309 |
+
parser.add_argument("--name", type=str, default="flux-dev", choices=list('flux-dev'),
|
310 |
+
help="currently only support flux-dev")
|
311 |
+
parser.add_argument("--device", type=str, default="cuda", help="Device to use")
|
312 |
+
parser.add_argument("--offload", action="store_true", help="Offload model to CPU when not in use")
|
313 |
+
parser.add_argument("--aggressive_offload", action="store_true", help="Offload model more aggressively to CPU when not in use, for 24G GPUs")
|
314 |
+
parser.add_argument("--fp8", action="store_true", help="use flux-dev-fp8 model")
|
315 |
+
parser.add_argument("--onnx_provider", type=str, default="gpu", choices=["gpu", "cpu"],
|
316 |
+
help="set onnx_provider to cpu (default gpu) can help reduce RAM usage, and when combined with"
|
317 |
+
"fp8 option, the peak RAM is under 15GB")
|
318 |
+
parser.add_argument("--port", type=int, default=8080, help="Port to use")
|
319 |
+
parser.add_argument("--dev", action='store_true', help="Development mode")
|
320 |
+
parser.add_argument("--pretrained_model", type=str, help='for development')
|
321 |
+
args = parser.parse_args()
|
322 |
+
|
323 |
+
if args.aggressive_offload:
|
324 |
+
args.offload = True
|
325 |
+
|
326 |
+
demo = create_demo(args, args.name, args.device, args.offload, args.aggressive_offload)
|
327 |
+
demo.launch(server_name='0.0.0.0', server_port=args.port)
|
app_v1_1.py
ADDED
@@ -0,0 +1,252 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import argparse
|
2 |
+
|
3 |
+
import gradio as gr
|
4 |
+
import numpy as np
|
5 |
+
import torch
|
6 |
+
|
7 |
+
from pulid import attention_processor as attention
|
8 |
+
from pulid.pipeline_v1_1 import PuLIDPipeline
|
9 |
+
from pulid.utils import resize_numpy_image_long
|
10 |
+
|
11 |
+
torch.set_grad_enabled(False)
|
12 |
+
|
13 |
+
parser = argparse.ArgumentParser()
|
14 |
+
parser.add_argument(
|
15 |
+
'--base',
|
16 |
+
type=str,
|
17 |
+
default='RunDiffusion/Juggernaut-XL-v9',
|
18 |
+
choices=[
|
19 |
+
'Lykon/dreamshaper-xl-lightning',
|
20 |
+
# 'SG161222/RealVisXL_V4.0', will add it later
|
21 |
+
'RunDiffusion/Juggernaut-XL-v9',
|
22 |
+
],
|
23 |
+
)
|
24 |
+
# parser.add_argument('--sampler', type=str, default='dpmpp_2m', choices=['dpmpp_sde', 'dpmpp_2m'])
|
25 |
+
parser.add_argument('--port', type=int, default=7860)
|
26 |
+
args = parser.parse_args()
|
27 |
+
|
28 |
+
use_lightning_model = 'lightning' in args.base.lower()
|
29 |
+
# currently we only support two commonly used sampler
|
30 |
+
args.sampler = 'dpmpp_sde' if use_lightning_model else 'dpmpp_2m'
|
31 |
+
if use_lightning_model:
|
32 |
+
default_cfg = 2.0
|
33 |
+
default_steps = 5
|
34 |
+
else:
|
35 |
+
default_cfg = 7.0
|
36 |
+
default_steps = 25
|
37 |
+
|
38 |
+
pipeline = PuLIDPipeline(sdxl_repo=args.base, sampler=args.sampler)
|
39 |
+
|
40 |
+
# other params
|
41 |
+
DEFAULT_NEGATIVE_PROMPT = (
|
42 |
+
'flaws in the eyes, flaws in the face, flaws, lowres, non-HDRi, low quality, worst quality,'
|
43 |
+
'artifacts noise, text, watermark, glitch, deformed, mutated, ugly, disfigured, hands, '
|
44 |
+
'low resolution, partially rendered objects, deformed or partially rendered eyes, '
|
45 |
+
'deformed, deformed eyeballs, cross-eyed,blurry'
|
46 |
+
)
|
47 |
+
|
48 |
+
dreamshaper_example_inps = [
|
49 |
+
['portrait, blacklight', 'example_inputs/liuyifei.png', 42, 0.8, 10],
|
50 |
+
['pixel art, 1boy', 'example_inputs/lecun.jpg', 42, 0.8, 10],
|
51 |
+
[
|
52 |
+
'cinematic film still, close up, photo of redheaded girl near grasses, fictional landscapes, (intense sunlight:1.4), realist detail, brooding mood, ue5, detailed character expressions, light amber and red, amazing quality, wallpaper, analog film grain',
|
53 |
+
'example_inputs/liuyifei.png',
|
54 |
+
42,
|
55 |
+
0.8,
|
56 |
+
10,
|
57 |
+
],
|
58 |
+
[
|
59 |
+
'A minimalist line art depiction of an Artificial Intelligence being\'s thought process, lines and nodes forming intricate patterns.',
|
60 |
+
'example_inputs/hinton.jpeg',
|
61 |
+
42,
|
62 |
+
0.8,
|
63 |
+
10,
|
64 |
+
],
|
65 |
+
[
|
66 |
+
'instagram photo, photo of 23 y.o man in black sweater, pale skin, (smile:0.4), hard shadows',
|
67 |
+
'example_inputs/pengwei.jpg',
|
68 |
+
42,
|
69 |
+
0.8,
|
70 |
+
10,
|
71 |
+
],
|
72 |
+
[
|
73 |
+
'by Tsutomu Nihei,(strange but extremely beautiful:1.4),(masterpiece, best quality:1.4),in the style of nicola samori,The Joker,',
|
74 |
+
'example_inputs/lecun.jpg',
|
75 |
+
1675432759740519133,
|
76 |
+
0.8,
|
77 |
+
10,
|
78 |
+
],
|
79 |
+
]
|
80 |
+
|
81 |
+
jugger_example_inps = [
|
82 |
+
[
|
83 |
+
'robot,simple robot,robot with glass face,ellipse head robot,(made partially out of glass),hexagonal shapes,ferns growing inside head,butterflies on head,butterflies flying around',
|
84 |
+
'example_inputs/hinton.jpeg',
|
85 |
+
15022214902832471291,
|
86 |
+
0.8,
|
87 |
+
20,
|
88 |
+
],
|
89 |
+
['sticker art, 1girl', 'example_inputs/liuyifei.png', 42, 0.8, 20],
|
90 |
+
[
|
91 |
+
'1girl, cute model, Long thick Maxi Skirt, Knit sweater, swept back hair, alluring smile, working at a clothing store, perfect eyes, highly detailed beautiful expressive eyes, detailed eyes, 35mm photograph, film, bokeh, professional, 4k, highly detailed dynamic lighting, photorealistic, 8k, raw, rich, intricate details,',
|
92 |
+
'example_inputs/liuyifei.png',
|
93 |
+
42,
|
94 |
+
0.8,
|
95 |
+
20,
|
96 |
+
],
|
97 |
+
['Chinese paper-cut, 1girl', 'example_inputs/liuyifei.png', 42, 0.8, 20],
|
98 |
+
['Studio Ghibli, 1boy', 'example_inputs/hinton.jpeg', 42, 0.8, 20],
|
99 |
+
['1man made of ice sculpture', 'example_inputs/lecun.jpg', 42, 0.8, 20],
|
100 |
+
['portrait of green-skinned shrek, wearing lacoste purple sweater', 'example_inputs/lecun.jpg', 42, 0.8, 20],
|
101 |
+
['1990s Japanese anime, 1girl', 'example_inputs/liuyifei.png', 42, 0.8, 20],
|
102 |
+
['made of little stones, portrait', 'example_inputs/hinton.jpeg', 42, 0.8, 20],
|
103 |
+
]
|
104 |
+
|
105 |
+
|
106 |
+
@torch.inference_mode()
|
107 |
+
def run(*args):
|
108 |
+
id_image = args[0]
|
109 |
+
supp_images = args[1:4]
|
110 |
+
prompt, neg_prompt, scale, seed, steps, H, W, id_scale, num_zero, ortho = args[4:]
|
111 |
+
seed = int(seed)
|
112 |
+
if seed == -1:
|
113 |
+
seed = torch.Generator(device="cpu").seed()
|
114 |
+
|
115 |
+
pipeline.debug_img_list = []
|
116 |
+
|
117 |
+
attention.NUM_ZERO = num_zero
|
118 |
+
if ortho == 'v2':
|
119 |
+
attention.ORTHO = False
|
120 |
+
attention.ORTHO_v2 = True
|
121 |
+
elif ortho == 'v1':
|
122 |
+
attention.ORTHO = True
|
123 |
+
attention.ORTHO_v2 = False
|
124 |
+
else:
|
125 |
+
attention.ORTHO = False
|
126 |
+
attention.ORTHO_v2 = False
|
127 |
+
|
128 |
+
if id_image is not None:
|
129 |
+
id_image = resize_numpy_image_long(id_image, 1024)
|
130 |
+
supp_id_image_list = [
|
131 |
+
resize_numpy_image_long(supp_id_image, 1024) for supp_id_image in supp_images if supp_id_image is not None
|
132 |
+
]
|
133 |
+
id_image_list = [id_image] + supp_id_image_list
|
134 |
+
uncond_id_embedding, id_embedding = pipeline.get_id_embedding(id_image_list)
|
135 |
+
else:
|
136 |
+
uncond_id_embedding = None
|
137 |
+
id_embedding = None
|
138 |
+
|
139 |
+
img = pipeline.inference(
|
140 |
+
prompt, (1, H, W), neg_prompt, id_embedding, uncond_id_embedding, id_scale, scale, steps, seed
|
141 |
+
)[0]
|
142 |
+
|
143 |
+
return np.array(img), str(seed), pipeline.debug_img_list
|
144 |
+
|
145 |
+
|
146 |
+
_HEADER_ = '''
|
147 |
+
<h2><b>Official Gradio Demo</b></h2><h2><a href='https://github.com/ToTheBeginning/PuLID' target='_blank'><b>PuLID: Pure and Lightning ID Customization via Contrastive Alignment</b></a></h2>
|
148 |
+
|
149 |
+
**PuLID** is a tuning-free ID customization approach. PuLID maintains high ID fidelity while effectively reducing interference with the original model’s behavior.
|
150 |
+
|
151 |
+
Code: <a href='https://github.com/ToTheBeginning/PuLID' target='_blank'>GitHub</a>. Paper: <a href='https://arxiv.org/abs/2404.16022' target='_blank'>ArXiv</a>.
|
152 |
+
|
153 |
+
❗️❗️❗️**Tips:**
|
154 |
+
- we provide some examples in the bottom, you can try these example prompts first
|
155 |
+
- a single ID image is usually sufficient, you can also supplement with additional auxiliary images
|
156 |
+
- You can adjust the trade-off between ID fidelity and editability in the advanced options, but generally, the default settings are good enough.
|
157 |
+
|
158 |
+
''' # noqa E501
|
159 |
+
|
160 |
+
_CITE_ = r"""
|
161 |
+
If PuLID is helpful, please help to ⭐ the <a href='https://github.com/ToTheBeginning/PuLID' target='_blank'>Github Repo</a>. Thanks! [![GitHub Stars](https://img.shields.io/github/stars/ToTheBeginning/PuLID?style=social)](https://github.com/ToTheBeginning/PuLID)
|
162 |
+
---
|
163 |
+
📧 **Contact**
|
164 |
+
If you have any questions, feel free to open a discussion or contact us at <b>[email protected]</b> or <b>[email protected]</b>.
|
165 |
+
""" # noqa E501
|
166 |
+
|
167 |
+
|
168 |
+
with gr.Blocks(title="PuLID", css=".gr-box {border-color: #8136e2}") as demo:
|
169 |
+
gr.Markdown(_HEADER_)
|
170 |
+
with gr.Row():
|
171 |
+
with gr.Column():
|
172 |
+
with gr.Row():
|
173 |
+
face_image = gr.Image(label="ID image (main)", height=256)
|
174 |
+
supp_image1 = gr.Image(label="Additional ID image (auxiliary)", height=256)
|
175 |
+
supp_image2 = gr.Image(label="Additional ID image (auxiliary)", height=256)
|
176 |
+
supp_image3 = gr.Image(label="Additional ID image (auxiliary)", height=256)
|
177 |
+
prompt = gr.Textbox(label="Prompt", value='portrait,color,cinematic,in garden,soft light,detailed face')
|
178 |
+
submit = gr.Button("Generate")
|
179 |
+
neg_prompt = gr.Textbox(label="Negative Prompt", value=DEFAULT_NEGATIVE_PROMPT)
|
180 |
+
scale = gr.Slider(
|
181 |
+
label="CFG (recommend 2 for lightning model and 7 for non-accelerated model)",
|
182 |
+
value=default_cfg,
|
183 |
+
minimum=1,
|
184 |
+
maximum=10,
|
185 |
+
step=0.1,
|
186 |
+
)
|
187 |
+
seed = gr.Textbox(-1, label="Seed (-1 for random)")
|
188 |
+
steps = gr.Slider(label="Steps", value=default_steps, minimum=1, maximum=30, step=1)
|
189 |
+
with gr.Row():
|
190 |
+
H = gr.Slider(label="Height", value=1152, minimum=512, maximum=2024, step=64)
|
191 |
+
W = gr.Slider(label="Width", value=896, minimum=512, maximum=2024, step=64)
|
192 |
+
with gr.Row(), gr.Accordion(
|
193 |
+
"Advanced Options (adjust the trade-off between ID fidelity and editability)", open=False
|
194 |
+
):
|
195 |
+
id_scale = gr.Slider(
|
196 |
+
label="ID scale (Increasing it enhances ID similarity but reduces editability)",
|
197 |
+
minimum=0,
|
198 |
+
maximum=5,
|
199 |
+
step=0.05,
|
200 |
+
value=0.8,
|
201 |
+
interactive=True,
|
202 |
+
)
|
203 |
+
num_zero = gr.Slider(
|
204 |
+
label="num zero (Increasing it enhances ID editability but reduces similarity)",
|
205 |
+
minimum=0,
|
206 |
+
maximum=80,
|
207 |
+
step=1,
|
208 |
+
value=20,
|
209 |
+
interactive=True,
|
210 |
+
)
|
211 |
+
ortho = gr.Dropdown(label="ortho", choices=['off', 'v1', 'v2'], value='v2', visible=False)
|
212 |
+
|
213 |
+
with gr.Column():
|
214 |
+
output = gr.Image(label="Generated Image")
|
215 |
+
seed_output = gr.Textbox(label="Used Seed")
|
216 |
+
intermediate_output = gr.Gallery(label='DebugImage', elem_id="gallery", visible=False)
|
217 |
+
gr.Markdown(_CITE_)
|
218 |
+
|
219 |
+
with gr.Row(), gr.Column():
|
220 |
+
gr.Markdown("## Examples")
|
221 |
+
if args.base == 'Lykon/dreamshaper-xl-lightning':
|
222 |
+
gr.Examples(
|
223 |
+
examples=dreamshaper_example_inps,
|
224 |
+
inputs=[prompt, face_image, seed, id_scale, num_zero],
|
225 |
+
label='dreamshaper-xl-lightning examples',
|
226 |
+
)
|
227 |
+
elif args.base == 'RunDiffusion/Juggernaut-XL-v9':
|
228 |
+
gr.Examples(
|
229 |
+
examples=jugger_example_inps,
|
230 |
+
inputs=[prompt, face_image, seed, id_scale, num_zero],
|
231 |
+
label='Juggernaut-XL-v9 examples',
|
232 |
+
)
|
233 |
+
|
234 |
+
inps = [
|
235 |
+
face_image,
|
236 |
+
supp_image1,
|
237 |
+
supp_image2,
|
238 |
+
supp_image3,
|
239 |
+
prompt,
|
240 |
+
neg_prompt,
|
241 |
+
scale,
|
242 |
+
seed,
|
243 |
+
steps,
|
244 |
+
H,
|
245 |
+
W,
|
246 |
+
id_scale,
|
247 |
+
num_zero,
|
248 |
+
ortho,
|
249 |
+
]
|
250 |
+
submit.click(fn=run, inputs=inps, outputs=[output, seed_output, intermediate_output])
|
251 |
+
|
252 |
+
demo.launch(server_name='0.0.0.0', server_port=args.port)
|