amildravid4292 commited on
Commit
31081a5
·
verified ·
1 Parent(s): ad09921

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +406 -442
app.py CHANGED
@@ -12,424 +12,391 @@ import warnings
12
  warnings.filterwarnings("ignore")
13
  from PIL import Image
14
  import numpy as np
 
15
  from editing import get_direction, debias
16
  from sampling import sample_weights
17
  from lora_w2w import LoRAw2w
18
- from transformers import CLIPTextModel
19
- from lora_w2w import LoRAw2w
20
- from diffusers import AutoencoderKL, DDPMScheduler, DiffusionPipeline, UNet2DConditionModel, LMSDiscreteScheduler
21
- from transformers import AutoTokenizer, PretrainedConfig
22
- from diffusers import (
23
- AutoencoderKL,
24
- DDPMScheduler,
25
- DiffusionPipeline,
26
- DPMSolverMultistepScheduler,
27
- UNet2DConditionModel,
28
- PNDMScheduler,
29
- StableDiffusionPipeline
30
- )
31
  from huggingface_hub import snapshot_download
32
  import spaces
33
 
 
 
 
 
 
 
 
 
 
 
 
34
  models_path = snapshot_download(repo_id="Snapchat/w2w")
35
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
36
  @spaces.GPU
37
- def load_models(device):
38
- pretrained_model_name_or_path = "stablediffusionapi/realistic-vision-v51"
39
-
40
- revision = None
41
- rank = 1
42
- weight_dtype = torch.bfloat16
43
-
44
- # Load scheduler, tokenizer and models.
45
- pipe = StableDiffusionPipeline.from_pretrained("stablediffusionapi/realistic-vision-v51",
46
- torch_dtype=torch.float16,safety_checker = None,
47
- requires_safety_checker = False).to(device)
48
- noise_scheduler = pipe.scheduler
49
- del pipe
50
- tokenizer = AutoTokenizer.from_pretrained(
51
- pretrained_model_name_or_path, subfolder="tokenizer", revision=revision
52
- )
53
- text_encoder = CLIPTextModel.from_pretrained(
54
- pretrained_model_name_or_path, subfolder="text_encoder", revision=revision
55
- )
56
- vae = AutoencoderKL.from_pretrained(pretrained_model_name_or_path, subfolder="vae", revision=revision)
57
- unet = UNet2DConditionModel.from_pretrained(
58
- pretrained_model_name_or_path, subfolder="unet", revision=revision
59
- )
60
-
61
- unet.requires_grad_(False)
62
- unet.to(device, dtype=weight_dtype)
63
- vae.requires_grad_(False)
64
-
65
- text_encoder.requires_grad_(False)
66
- vae.requires_grad_(False)
67
- vae.to(device, dtype=weight_dtype)
68
- text_encoder.to(device, dtype=weight_dtype)
69
- print("")
70
 
71
- return unet, vae, text_encoder, tokenizer, noise_scheduler
72
 
73
- class main():
74
- def __init__(self):
75
- super(main, self).__init__()
76
-
77
- device = "cuda"
78
- mean = torch.load(f"{models_path}/files/mean.pt", map_location=torch.device('cpu')).bfloat16().to(device)
79
- std = torch.load(f"{models_path}/files/std.pt", map_location=torch.device('cpu')).bfloat16().to(device)
80
- v = torch.load(f"{models_path}/files/V.pt", map_location=torch.device('cpu')).bfloat16().to(device)
81
- proj = torch.load(f"{models_path}/files/proj_1000pc.pt", map_location=torch.device('cpu')).bfloat16().to(device)
82
- df = torch.load(f"{models_path}/files/identity_df.pt")
83
- weight_dimensions = torch.load(f"{models_path}/files/weight_dimensions.pt")
84
- pinverse = torch.load(f"{models_path}/files/pinverse_1000pc.pt", map_location=torch.device('cpu')).bfloat16().to(device)
85
-
86
- self.device = device
87
- self.mean = mean
88
- self.std = std
89
- self.v = v
90
- self.proj = proj
91
- self.df = df
92
- self.weight_dimensions = weight_dimensions
93
- self.pinverse = pinverse
94
-
95
- self.unet, self.vae, self.text_encoder, self.tokenizer, self.noise_scheduler = load_models(self.device)
96
- print(self.text_encoder.device)
97
- self.network = None
98
-
99
- young = get_direction(df, "Young", pinverse, 1000, device)
100
- young = debias(young, "Male", df, pinverse, device)
101
- young = debias(young, "Pointy_Nose", df, pinverse, device)
102
- young = debias(young, "Wavy_Hair", df, pinverse, device)
103
- young = debias(young, "Chubby", df, pinverse, device)
104
- young = debias(young, "No_Beard", df, pinverse, device)
105
- young = debias(young, "Mustache", df, pinverse, device)
106
- self.young = young
107
-
108
- pointy = get_direction(df, "Pointy_Nose", pinverse, 1000, device)
109
- pointy = debias(pointy, "Young", df, pinverse, device)
110
- pointy = debias(pointy, "Male", df, pinverse, device)
111
- pointy = debias(pointy, "Wavy_Hair", df, pinverse, device)
112
- pointy = debias(pointy, "Chubby", df, pinverse, device)
113
- pointy = debias(pointy, "Heavy_Makeup", df, pinverse, device)
114
- self.pointy = pointy
115
-
116
- wavy = get_direction(df, "Wavy_Hair", pinverse, 1000, device)
117
- wavy = debias(wavy, "Young", df, pinverse, device)
118
- wavy = debias(wavy, "Male", df, pinverse, device)
119
- wavy = debias(wavy, "Pointy_Nose", df, pinverse, device)
120
- wavy = debias(wavy, "Chubby", df, pinverse, device)
121
- wavy = debias(wavy, "Heavy_Makeup", df, pinverse, device)
122
- self.wavy = wavy
 
 
 
 
 
 
 
123
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
124
 
125
- thick = get_direction(df, "Bushy_Eyebrows", pinverse, 1000, device)
126
- thick = debias(thick, "Male", df, pinverse, device)
127
- thick = debias(thick, "Young", df, pinverse, device)
128
- thick = debias(thick, "Pointy_Nose", df, pinverse, device)
129
- thick = debias(thick, "Wavy_Hair", df, pinverse, device)
130
- thick = debias(thick, "Mustache", df, pinverse, device)
131
- thick = debias(thick, "No_Beard", df, pinverse, device)
132
- thick = debias(thick, "Sideburns", df, pinverse, device)
133
- thick = debias(thick, "Big_Nose", df, pinverse, device)
134
- thick = debias(thick, "Big_Lips", df, pinverse, device)
135
- thick = debias(thick, "Black_Hair", df, pinverse, device)
136
- thick = debias(thick, "Brown_Hair", df, pinverse, device)
137
- thick = debias(thick, "Pale_Skin", df, pinverse, device)
138
- thick = debias(thick, "Heavy_Makeup", df, pinverse, device)
139
- self.thick = thick
140
-
141
-
142
-
143
- def sample_model(self):
144
- self.unet, _, _, _, _ = load_models(self.device)
145
- self.network = sample_weights(self.unet, self.proj, self.mean, self.std, self.v[:, :1000], self.device, factor = 1.00)
146
-
147
 
148
- @torch.no_grad()
149
- @spaces.GPU
150
- def inference(self, prompt, negative_prompt, guidance_scale, ddim_steps, seed):
151
- device = self.device
152
- generator = torch.Generator(device=device).manual_seed(seed)
153
- latents = torch.randn(
154
- (1, self.unet.in_channels, 512 // 8, 512 // 8),
155
- generator = generator,
156
- device = self.device
157
- ).bfloat16()
158
-
159
-
160
- text_input = self.tokenizer(prompt, padding="max_length", max_length=self.tokenizer.model_max_length, truncation=True, return_tensors="pt")
161
-
162
- text_embeddings = self.text_encoder(text_input.input_ids.to(device))[0]
163
-
164
- max_length = text_input.input_ids.shape[-1]
165
- uncond_input = self.tokenizer(
166
- [negative_prompt], padding="max_length", max_length=max_length, return_tensors="pt"
167
- )
168
- uncond_embeddings = self.text_encoder(uncond_input.input_ids.to(device))[0]
169
- text_embeddings = torch.cat([uncond_embeddings, text_embeddings])
170
- self.noise_scheduler.set_timesteps(ddim_steps)
171
- latents = latents * self.noise_scheduler.init_noise_sigma
172
-
173
- for i,t in enumerate(tqdm.tqdm(self.noise_scheduler.timesteps)):
174
- latent_model_input = torch.cat([latents] * 2)
175
- latent_model_input = self.noise_scheduler.scale_model_input(latent_model_input, timestep=t)
176
- with self.network:
177
- noise_pred = self.unet(latent_model_input, t, encoder_hidden_states=text_embeddings, timestep_cond= None).sample
178
- #guidance
179
- noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
180
- noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)
181
- latents = noise_scheduler.step(noise_pred, t, latents).prev_sample
182
 
183
- latents = 1 / 0.18215 * latents
184
- image = self.vae.decode(latents).sample
185
- image = (image / 2 + 0.5).clamp(0, 1)
186
- image = image.detach().cpu().float().permute(0, 2, 3, 1).numpy()[0]
187
-
188
- image = Image.fromarray((image * 255).round().astype("uint8"))
189
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
190
  return image
191
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
192
 
193
- @torch.no_grad()
194
- @spaces.GPU
195
- def edit_inference(self, prompt, negative_prompt, guidance_scale, ddim_steps, seed, start_noise, a1, a2, a3, a4):
196
- device = self.device
197
- original_weights = self,network.proj.clone()
198
-
199
- #pad to same number of PCs
200
- pcs_original = original_weights.shape[1]
201
- pcs_edits = self.young.shape[1]
202
- padding = torch.zeros((1,pcs_original-pcs_edits)).to(device)
203
- young_pad = torch.cat((self.young, padding), 1)
204
- pointy_pad = torch.cat((self.pointy, padding), 1)
205
- wavy_pad = torch.cat((self.wavy, padding), 1)
206
- thick_pad = torch.cat((self.thick, padding), 1)
207
-
208
-
209
- edited_weights = original_weights+a1*1e6*young_pad+a2*1e6*pointy_pad+a3*1e6*wavy_pad+a4*2e6*thick_pad
210
-
211
- generator = torch.Generator(device=device).manual_seed(seed)
212
- latents = torch.randn(
213
- (1, self.unet.in_channels, 512 // 8, 512 // 8),
214
- generator = generator,
215
- device = self.device
216
- ).bfloat16()
217
-
218
-
219
- text_input = self.tokenizer(prompt, padding="max_length", max_length=self.tokenizer.model_max_length, truncation=True, return_tensors="pt")
220
-
221
- text_embeddings = text_encoder(text_input.input_ids.to(device))[0]
222
-
223
- max_length = text_input.input_ids.shape[-1]
224
- uncond_input = tokenizer(
225
- [negative_prompt], padding="max_length", max_length=max_length, return_tensors="pt"
226
- )
227
- uncond_embeddings = text_encoder(uncond_input.input_ids.to(device))[0]
228
- text_embeddings = torch.cat([uncond_embeddings, text_embeddings])
229
- noise_scheduler.set_timesteps(ddim_steps)
230
- latents = latents * noise_scheduler.init_noise_sigma
231
-
232
-
233
 
234
- for i,t in enumerate(tqdm.tqdm(self.noise_scheduler.timesteps)):
235
- latent_model_input = torch.cat([latents] * 2)
236
- latent_model_input = self.noise_scheduler.scale_model_input(latent_model_input, timestep=t)
237
-
238
- if t>start_noise:
239
- pass
240
- elif t<=start_noise:
241
- self.network.proj = torch.nn.Parameter(edited_weights)
242
- self.network.reset()
243
-
244
-
245
- with self.network:
246
- noise_pred = self.unet(latent_model_input, t, encoder_hidden_states=text_embeddings, timestep_cond= None).sample
247
-
248
-
249
- #guidance
250
- noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
251
- noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)
252
- latents = noise_scheduler.step(noise_pred, t, latents).prev_sample
253
-
254
- latents = 1 / 0.18215 * latents
255
- image = self.vae.decode(latents).sample
256
- image = (image / 2 + 0.5).clamp(0, 1)
257
-
258
- image = image.detach().cpu().float().permute(0, 2, 3, 1).numpy()[0]
259
-
260
- image = Image.fromarray((image * 255).round().astype("uint8"))
261
-
262
- #reset weights back to original
263
- self.network.proj = torch.nn.Parameter(original_weights)
264
- self.network.reset()
265
-
266
- return image
267
 
268
- @spaces.GPU
269
- def sample_then_run(self):
270
- self.sample_model()
271
- prompt = "sks person"
272
- negative_prompt = "low quality, blurry, unfinished, nudity, weapon"
273
- seed = 5
274
- cfg = 3.0
275
- steps = 25
276
- image = self.inference( prompt, negative_prompt, cfg, steps, seed)
277
- torch.save(self.network.proj, "model.pt" )
278
- return image, "model.pt"
279
 
 
 
280
 
281
 
282
- class CustomImageDataset(Dataset):
283
- def __init__(self, images, transform=None):
284
- self.images = images
285
- self.transform = transform
286
-
287
- def __len__(self):
288
- return len(self.images)
289
-
290
- def __getitem__(self, idx):
291
- image = self.images[idx]
292
- if self.transform:
293
- image = self.transform(image)
294
- return image
295
-
296
- @spaces.GPU
297
- def invert(self, image, mask, pcs=10000, epochs=400, weight_decay = 1e-10, lr=1e-1):
298
-
299
- del unet
300
- del network
301
- unet, _, _, _, _ = load_models(device)
302
-
303
- proj = torch.zeros(1,pcs).bfloat16().to(device)
304
- network = LoRAw2w( proj, mean, std, v[:, :pcs],
305
- unet,
306
- rank=1,
307
- multiplier=1.0,
308
- alpha=27.0,
309
- train_method="xattn-strict"
310
- ).to(device, torch.bfloat16)
311
-
312
- ### load mask
313
- mask = transforms.Resize((64,64), interpolation=transforms.InterpolationMode.BILINEAR)(mask)
314
- mask = torchvision.transforms.functional.pil_to_tensor(mask).unsqueeze(0).to(device).bfloat16()[:,0,:,:].unsqueeze(1)
315
- ### check if an actual mask was draw, otherwise mask is just all ones
316
- if torch.sum(mask) == 0:
317
- mask = torch.ones((1,1,64,64)).to(device).bfloat16()
318
-
319
- ### single image dataset
320
- image_transforms = transforms.Compose([transforms.Resize(512, interpolation=transforms.InterpolationMode.BILINEAR),
321
- transforms.RandomCrop(512),
322
- transforms.ToTensor(),
323
- transforms.Normalize([0.5], [0.5])])
324
-
325
-
326
- train_dataset = CustomImageDataset(image, transform=image_transforms)
327
- train_dataloader = torch.utils.data.DataLoader(train_dataset, batch_size=1, shuffle=True)
328
-
329
- ### optimizer
330
- optim = torch.optim.Adam(network.parameters(), lr=lr, weight_decay=weight_decay)
331
-
332
- ### training loop
333
- unet.train()
334
- for epoch in tqdm.tqdm(range(epochs)):
335
- for batch in train_dataloader:
336
- ### prepare inputs
337
- batch = batch.to(device).bfloat16()
338
- latents = vae.encode(batch).latent_dist.sample()
339
- latents = latents*0.18215
340
- noise = torch.randn_like(latents)
341
- bsz = latents.shape[0]
342
-
343
- timesteps = torch.randint(0, noise_scheduler.config.num_train_timesteps, (bsz,), device=latents.device)
344
- timesteps = timesteps.long()
345
- noisy_latents = noise_scheduler.add_noise(latents, noise, timesteps)
346
- text_input = tokenizer("sks person", padding="max_length", max_length=tokenizer.model_max_length, truncation=True, return_tensors="pt")
347
- text_embeddings = text_encoder(text_input.input_ids.to(device))[0]
348
-
349
- ### loss + sgd step
350
- with network:
351
- model_pred = unet(noisy_latents, timesteps, text_embeddings).sample
352
- loss = torch.nn.functional.mse_loss(mask*model_pred.float(), mask*noise.float(), reduction="mean")
353
- optim.zero_grad()
354
- loss.backward()
355
- optim.step()
356
-
357
- ### return optimized network
358
- return network
359
 
360
 
361
- @spaces.GPU
362
- def run_inversion(self, dict, pcs, epochs, weight_decay,lr):
363
- init_image = dict["image"].convert("RGB").resize((512, 512))
364
- mask = dict["mask"].convert("RGB").resize((512, 512))
365
- network = invert([init_image], mask, pcs, epochs, weight_decay,lr)
366
-
367
-
368
- #sample an image
369
- prompt = "sks person"
370
- negative_prompt = "low quality, blurry, unfinished, nudity"
371
- seed = 5
372
- cfg = 3.0
373
- steps = 25
374
- image = inference( prompt, negative_prompt, cfg, steps, seed)
375
- torch.save(network.proj, "model.pt" )
376
- return image, "model.pt"
377
 
378
 
379
- @spaces.GPU
380
- def file_upload(self, file):
381
-
382
- proj = torch.load(file.name).to(device)
383
-
384
- #pad to 10000 Principal components to keep everything consistent
385
- pcs = proj.shape[1]
386
- padding = torch.zeros((1,10000-pcs)).to(device)
387
- proj = torch.cat((proj, padding), 1)
388
-
389
- unet, _, _, _, _ = load_models(device)
390
-
391
-
392
- network = LoRAw2w( proj, mean, std, v[:, :10000],
393
- unet,
394
- rank=1,
395
- multiplier=1.0,
396
- alpha=27.0,
397
- train_method="xattn-strict"
398
- ).to(device, torch.bfloat16)
399
-
400
-
401
- prompt = "sks person"
402
- negative_prompt = "low quality, blurry, unfinished, nudity"
403
- seed = 5
404
- cfg = 3.0
405
- steps = 25
406
- image = inference( prompt, negative_prompt, cfg, steps, seed)
407
- return image
408
 
409
 
 
 
 
 
 
 
410
 
 
 
 
 
 
 
 
 
 
 
 
411
 
412
- intro = """
413
- <div style="display: flex;align-items: center;justify-content: center">
414
- <h1 style="margin-left: 12px;text-align: center;margin-bottom: 7px;display: inline-block"><em>weights2weights</em> Demo</h1>
415
- <h3 style="display: inline-block;margin-left: 10px;margin-top: 6px;font-weight: 500">Interpreting the Weight Space of Customized Diffusion Models</h3>
416
- </div>
417
- <p style="font-size: 0.95rem;margin: 0rem;line-height: 1.2em;margin-top:1em;display: inline-block">
418
- <a href="https://snap-research.github.io/weights2weights/" target="_blank">Project Page</a> | <a href="https://arxiv.org/abs/2406.09413" target="_blank">Paper</a>
419
- | <a href="https://github.com/snap-research/weights2weights" target="_blank">Code</a> |
420
- <a href="https://huggingface.co/spaces/Snapchat/w2w-demo?duplicate=true" target="_blank" style="
421
- display: inline-block;
422
- ">
423
- <img style="margin-top: -1em;margin-bottom: 0em;position: absolute;" src="https://bit.ly/3CWLGkA" alt="Duplicate Space"></a>
424
- </p>
425
- """
426
 
427
 
 
428
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
429
  with gr.Blocks(css="style.css") as demo:
430
- model = main()
431
  gr.HTML(intro)
432
-
433
  gr.Markdown("""<div style="text-align: justify;"> In this demo, you can get an identity-encoding model by sampling or inverting. To use a model previously downloaded from this demo see \"Uploading a model\" in the Advanced Options. Next, you can generate new images from it, or edit the identity encoded in the model and generate images from the edited model. We provide detailed instructions and tips at the bottom of the page.""")
434
  with gr.Column():
435
  with gr.Row():
@@ -437,26 +404,26 @@ with gr.Blocks(css="style.css") as demo:
437
  gr.Markdown("""1) Either sample a new model, or upload an image (optionally draw a mask over the head) and click `invert`.""")
438
  sample = gr.Button("🎲 Sample New Model")
439
  input_image = gr.ImageEditor(elem_id="image_upload", type='pil', label="Reference Identity",
440
- width=512, height=512)
441
-
442
  with gr.Row():
443
  invert_button = gr.Button("⬆️ Invert")
444
-
445
-
446
-
447
  with gr.Column():
448
  gr.Markdown("""2) Generate images of the sampled/inverted identity or edit the identity with the sliders and generate new images with various prompts and seeds.""")
449
  gallery = gr.Image(label="Generated Image",height=512, width=512, interactive=False)
450
  submit = gr.Button("Generate")
451
-
452
-
453
  prompt = gr.Textbox(label="Prompt",
454
- info="Make sure to include 'sks person'" ,
455
- placeholder="sks person",
456
- value="sks person")
457
-
458
  seed = gr.Number(value=5, label="Seed", precision=0, interactive=True)
459
-
460
  # Editing
461
  with gr.Column():
462
  with gr.Row():
@@ -465,8 +432,8 @@ with gr.Blocks(css="style.css") as demo:
465
  with gr.Row():
466
  a3 = gr.Slider(label="- Curly Hair +", value=0, step=0.001, minimum=-1, maximum=1, interactive=True)
467
  a4 = gr.Slider(label="- Thick Eyebrows +", value=0, step=0.001, minimum=-1, maximum=1, interactive=True)
468
-
469
-
470
  with gr.Accordion("Advanced Options", open=False):
471
  with gr.Tab("Inversion"):
472
  with gr.Row():
@@ -477,65 +444,62 @@ with gr.Blocks(css="style.css") as demo:
477
  weight_decay = gr.Number(value=1e-10, label="Weight Decay", interactive=True)
478
  with gr.Tab("Sampling"):
479
  with gr.Row():
480
- cfg= gr.Slider(label="CFG", value=3.0, step=0.1, minimum=0, maximum=10, interactive=True)
481
- steps = gr.Slider(label="Inference Steps", value=25, step=1, minimum=0, maximum=100, interactive=True)
482
  with gr.Row():
483
- negative_prompt = gr.Textbox(label="Negative Prompt", placeholder="low quality, blurry, unfinished, nudity, weapon", value="low quality, blurry, unfinished, nudity, weapon")
484
- injection_step = gr.Slider(label="Injection Step", value=800, step=1, minimum=0, maximum=1000, interactive=True)
485
-
486
  with gr.Tab("Uploading a model"):
487
  gr.Markdown("""<div style="text-align: justify;">Upload a model below downloaded from this demo.""")
488
-
489
  file_input = gr.File(label="Upload Model", container=True)
490
-
491
-
492
-
493
-
494
  gr.Markdown("""<div style="text-align: justify;"> After sampling a new model or inverting, you can download the model below.""")
495
-
496
  with gr.Row():
497
  file_output = gr.File(label="Download Sampled/Inverted Model", container=True, interactive=False)
 
 
 
 
 
 
 
 
498
 
499
-
500
-
501
-
502
- invert_button.click(fn=model.run_inversion,
503
- inputs=[input_image, pcs, epochs, weight_decay,lr],
504
- outputs = [input_image, file_output])
505
-
506
-
507
- sample.click(fn=model.sample_then_run, outputs=[input_image, file_output])
508
-
509
  submit.click(
510
- fn=model.edit_inference, inputs=[prompt, negative_prompt, cfg, steps, seed, injection_step, a1, a2, a3, a4], outputs=[gallery]
511
- )
512
- file_input.change(fn=model.file_upload, inputs=file_input, outputs = gallery)
513
-
514
-
515
-
516
  help_text1 = """
517
- <b>Instructions</b>:
518
- 1. To get results faster without waiting in queue, you can duplicate into a private space with an A100 GPU.
519
- 2. To begin, you will have to get an identity-encoding model. You can either sample one from *weights2weights* space by clicking `Sample New Model` or by uploading an image and clicking `invert` to invert the identity into a model. You can optionally draw over the head to define a mask in the image for better results. Sampling a model takes around 10 seconds and inversion takes around 2 minutes. After this is done, you can optionally download this model for later use. A model can be uploaded in the \"Uploading a model\" tab in the `Advanced Options`.
520
- 3. After getting a model, an image of the identity will be displayed on the right. You can sample from the model by changing seeds as well as prompts and then clicking `Generate`. Make sure to include \"sks person\" in your prompt to keep the same identity.
521
- 4. The identity in the model can be edited by changing the sliders for various attributes. After clicking `Generate`, you can see how the identity has changed and the effects are maintained across different seeds and prompts.
522
- """
523
  help_text2 = """<b>Tips</b>:
524
- 1. Editing and Identity Generation
525
- * If you are interested in preserving more of the image during identity-editing (i.e., where the same seed and prompt results in the same image with only the identity changed), you can play with the "Injection Step" parameter in the \"Sampling\" tab in the `Advanced Options`. During the first *n* timesteps, the original model's weights will be used, and then the edited weights will be set during the remaining steps. Values closer to 1000 will set the edited weights early, having a more pronounced effect, which may disrupt some semantics and structure of the generated image. Lower values will set the edited weights later, better preserving image context. We notice that around 600-800 tends to produce the best results. Larger values in the range (700-1000) are helpful for more global attribute changes, while smaller (400-700) can be used for more finegrained edits. Although it is not always needed.
526
- * You can play around with negative prompts, number of inference steps, and CFG in the \"Sampling\" tab in the `Advanced Options` to affect the ultimate image quality.
527
- * Sometimes the identity will not be perfectly consistent (e.g., there might be small variations of the face) when you use some seeds or prompts. This is a limitation of our method as well as an open-problem in personalized models.
528
- 2. Inversion
529
- * To obtain the best results for inversion, upload a high resolution photo of the face with minimal occlusion. It is recommended to draw over the face and hair to define a mask. But inversion should still work generally for non-closeup face shots.
530
- * For inverting a realistic photo of an identity, typically 800 epochs with lr=1e-1 and 10,000 principal components (PCs) works well. If the resulting generations have artifacted and unrealstic textures, there is probably overfitting and you may want to reduce the number of epochs or learning rate, or play with weight decay. If the generations do not look like the input photo, then you may want to increase the number of epochs.
531
- * For inverting out-of-distribution identities, such as artistic renditions of people or non-humans (e.g. the ones shown in the paper), it is recommended to use 1000 PCs, lr=1, and train for 800 epochs.
532
- * Note that if you change the number of PCs, you will probably need to change the learning rate. For less PCs, higher learning rates are typically required."""
 
533
 
534
-
535
  gr.Markdown(help_text1)
536
  gr.Markdown(help_text2)
537
-
538
  demo.queue().launch()
539
-
540
-
541
-
 
12
  warnings.filterwarnings("ignore")
13
  from PIL import Image
14
  import numpy as np
15
+ from utils import load_models
16
  from editing import get_direction, debias
17
  from sampling import sample_weights
18
  from lora_w2w import LoRAw2w
 
 
 
 
 
 
 
 
 
 
 
 
 
19
  from huggingface_hub import snapshot_download
20
  import spaces
21
 
22
+ device = gr.State()
23
+ generator = gr.State()
24
+ unet = gr.State()
25
+ vae = gr.State()
26
+ text_encoder = gr.State()
27
+ tokenizer = gr.State()
28
+ noise_scheduler = gr.State()
29
+ network = gr.State()
30
+ device = gr.State("cuda")
31
+
32
+
33
  models_path = snapshot_download(repo_id="Snapchat/w2w")
34
 
35
+ mean = torch.load(f"{models_path}/files/mean.pt", map_location=torch.device('cpu')).bfloat16().to(device)
36
+ std = torch.load(f"{models_path}/files/std.pt", map_location=torch.device('cpu')).bfloat16().to(device)
37
+ v = torch.load(f"{models_path}/files/V.pt", map_location=torch.device('cpu')).bfloat16().to(device)
38
+ proj = torch.load(f"{models_path}/files/proj_1000pc.pt", map_location=torch.device('cpu')).bfloat16().to(device)
39
+ df = torch.load(f"{models_path}/files/identity_df.pt")
40
+ weight_dimensions = torch.load(f"{models_path}/files/weight_dimensions.pt")
41
+ pinverse = torch.load(f"{models_path}/files/pinverse_1000pc.pt", map_location=torch.device('cpu')).bfloat16().to(device)
42
+
43
+ unet, vae, text_encoder, tokenizer, noise_scheduler = load_models(device)
44
+
45
+ def sample_model():
46
+ unet.value, _, _, _, _ = load_models(device)
47
+ network.value = sample_weights(unet, proj, mean, std, v[:, :1000], device, factor = 1.00)
48
+
49
+ @torch.no_grad()
50
  @spaces.GPU
51
+ def inference( prompt, negative_prompt, guidance_scale, ddim_steps, seed):
52
+ global device
53
+ #global generator
54
+ global unet
55
+ global vae
56
+ global text_encoder
57
+ global tokenizer
58
+ global noise_scheduler
59
+ generator = torch.Generator(device=device).manual_seed(seed)
60
+ latents = torch.randn(
61
+ (1, unet.in_channels, 512 // 8, 512 // 8),
62
+ generator = generator,
63
+ device = device
64
+ ).bfloat16()
65
+
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
66
 
67
+ text_input = tokenizer(prompt, padding="max_length", max_length=tokenizer.model_max_length, truncation=True, return_tensors="pt")
68
 
69
+ text_embeddings = text_encoder(text_input.input_ids.to(device))[0]
70
+
71
+ max_length = text_input.input_ids.shape[-1]
72
+ uncond_input = tokenizer(
73
+ [negative_prompt], padding="max_length", max_length=max_length, return_tensors="pt"
74
+ )
75
+ uncond_embeddings = text_encoder(uncond_input.input_ids.to(device))[0]
76
+ text_embeddings = torch.cat([uncond_embeddings, text_embeddings])
77
+ noise_scheduler.set_timesteps(ddim_steps)
78
+ latents = latents * noise_scheduler.init_noise_sigma
79
+
80
+ for i,t in enumerate(tqdm.tqdm(noise_scheduler.timesteps)):
81
+ latent_model_input = torch.cat([latents] * 2)
82
+ latent_model_input = noise_scheduler.scale_model_input(latent_model_input, timestep=t)
83
+ with network:
84
+ noise_pred = unet(latent_model_input, t, encoder_hidden_states=text_embeddings, timestep_cond= None).sample
85
+ #guidance
86
+ noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
87
+ noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)
88
+ latents = noise_scheduler.step(noise_pred, t, latents).prev_sample
89
+
90
+ latents = 1 / 0.18215 * latents
91
+ image = vae.decode(latents).sample
92
+ image = (image / 2 + 0.5).clamp(0, 1)
93
+ image = image.detach().cpu().float().permute(0, 2, 3, 1).numpy()[0]
94
+
95
+ image = Image.fromarray((image * 255).round().astype("uint8"))
96
+
97
+ return image
98
+
99
+
100
+ @torch.no_grad()
101
+ @spaces.GPU
102
+ def edit_inference(prompt, negative_prompt, guidance_scale, ddim_steps, seed, start_noise, a1, a2, a3, a4):
103
+ start_items()
104
+ global device
105
+ #global generator
106
+ global unet
107
+ global vae
108
+ global text_encoder
109
+ global tokenizer
110
+ global noise_scheduler
111
+ global young
112
+ global pointy
113
+ global wavy
114
+ global thick
115
+
116
+ original_weights = network.proj.clone()
117
+
118
+ #pad to same number of PCs
119
+ pcs_original = original_weights.shape[1]
120
+ pcs_edits = young.shape[1]
121
+ padding = torch.zeros((1,pcs_original-pcs_edits)).to(device)
122
+ young_pad = torch.cat((young, padding), 1)
123
+ pointy_pad = torch.cat((pointy, padding), 1)
124
+ wavy_pad = torch.cat((wavy, padding), 1)
125
+ thick_pad = torch.cat((thick, padding), 1)
126
 
127
+
128
+ edited_weights = original_weights+a1*1e6*young_pad+a2*1e6*pointy_pad+a3*1e6*wavy_pad+a4*2e6*thick_pad
129
+
130
+ generator = torch.Generator(device=device).manual_seed(seed)
131
+ latents = torch.randn(
132
+ (1, unet.in_channels, 512 // 8, 512 // 8),
133
+ generator = generator,
134
+ device = device
135
+ ).bfloat16()
136
+
137
+
138
+ text_input = tokenizer(prompt, padding="max_length", max_length=tokenizer.model_max_length, truncation=True, return_tensors="pt")
139
+
140
+ text_embeddings = text_encoder(text_input.input_ids.to(device))[0]
141
+
142
+ max_length = text_input.input_ids.shape[-1]
143
+ uncond_input = tokenizer(
144
+ [negative_prompt], padding="max_length", max_length=max_length, return_tensors="pt"
145
+ )
146
+ uncond_embeddings = text_encoder(uncond_input.input_ids.to(device))[0]
147
+ text_embeddings = torch.cat([uncond_embeddings, text_embeddings])
148
+ noise_scheduler.set_timesteps(ddim_steps)
149
+ latents = latents * noise_scheduler.init_noise_sigma
150
 
151
+
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
152
 
153
+ for i,t in enumerate(tqdm.tqdm(noise_scheduler.timesteps)):
154
+ latent_model_input = torch.cat([latents] * 2)
155
+ latent_model_input = noise_scheduler.scale_model_input(latent_model_input, timestep=t)
156
+
157
+ if t>start_noise:
158
+ pass
159
+ elif t<=start_noise:
160
+ network.proj = torch.nn.Parameter(edited_weights)
161
+ network.reset()
162
+
163
+
164
+ with network:
165
+ noise_pred = unet(latent_model_input, t, encoder_hidden_states=text_embeddings, timestep_cond= None).sample
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
166
 
 
 
 
 
 
 
167
 
168
+ #guidance
169
+ noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
170
+ noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)
171
+ latents = noise_scheduler.step(noise_pred, t, latents).prev_sample
172
+
173
+ latents = 1 / 0.18215 * latents
174
+ image = vae.decode(latents).sample
175
+ image = (image / 2 + 0.5).clamp(0, 1)
176
+
177
+ image = image.detach().cpu().float().permute(0, 2, 3, 1).numpy()[0]
178
+
179
+ image = Image.fromarray((image * 255).round().astype("uint8"))
180
+
181
+ #reset weights back to original
182
+ network.proj = torch.nn.Parameter(original_weights)
183
+ network.reset()
184
+
185
+ return image
186
+
187
+ @spaces.GPU
188
+ def sample_then_run():
189
+ sample_model()
190
+ prompt = "sks person"
191
+ negative_prompt = "low quality, blurry, unfinished, nudity, weapon"
192
+ seed = 5
193
+ cfg = 3.0
194
+ steps = 25
195
+ image = inference( prompt, negative_prompt, cfg, steps, seed)
196
+ torch.save(network.proj, "model.pt" )
197
+ return image, "model.pt"
198
+
199
+ #@spaces.GPU
200
+ def start_items():
201
+ print("Starting items")
202
+ global young
203
+ global pointy
204
+ global wavy
205
+ global thick
206
+ young = get_direction(df, "Young", pinverse, 1000, device)
207
+ young = debias(young, "Male", df, pinverse, device)
208
+ young = debias(young, "Pointy_Nose", df, pinverse, device)
209
+ young = debias(young, "Wavy_Hair", df, pinverse, device)
210
+ young = debias(young, "Chubby", df, pinverse, device)
211
+ young = debias(young, "No_Beard", df, pinverse, device)
212
+ young = debias(young, "Mustache", df, pinverse, device)
213
+
214
+ pointy = get_direction(df, "Pointy_Nose", pinverse, 1000, device)
215
+ pointy = debias(pointy, "Young", df, pinverse, device)
216
+ pointy = debias(pointy, "Male", df, pinverse, device)
217
+ pointy = debias(pointy, "Wavy_Hair", df, pinverse, device)
218
+ pointy = debias(pointy, "Chubby", df, pinverse, device)
219
+ pointy = debias(pointy, "Heavy_Makeup", df, pinverse, device)
220
+
221
+ wavy = get_direction(df, "Wavy_Hair", pinverse, 1000, device)
222
+ wavy = debias(wavy, "Young", df, pinverse, device)
223
+ wavy = debias(wavy, "Male", df, pinverse, device)
224
+ wavy = debias(wavy, "Pointy_Nose", df, pinverse, device)
225
+ wavy = debias(wavy, "Chubby", df, pinverse, device)
226
+ wavy = debias(wavy, "Heavy_Makeup", df, pinverse, device)
227
+
228
+ thick = get_direction(df, "Bushy_Eyebrows", pinverse, 1000, device)
229
+ thick = debias(thick, "Male", df, pinverse, device)
230
+ thick = debias(thick, "Young", df, pinverse, device)
231
+ thick = debias(thick, "Pointy_Nose", df, pinverse, device)
232
+ thick = debias(thick, "Wavy_Hair", df, pinverse, device)
233
+ thick = debias(thick, "Mustache", df, pinverse, device)
234
+ thick = debias(thick, "No_Beard", df, pinverse, device)
235
+ thick = debias(thick, "Sideburns", df, pinverse, device)
236
+ thick = debias(thick, "Big_Nose", df, pinverse, device)
237
+ thick = debias(thick, "Big_Lips", df, pinverse, device)
238
+ thick = debias(thick, "Black_Hair", df, pinverse, device)
239
+ thick = debias(thick, "Brown_Hair", df, pinverse, device)
240
+ thick = debias(thick, "Pale_Skin", df, pinverse, device)
241
+ thick = debias(thick, "Heavy_Makeup", df, pinverse, device)
242
+
243
+ class CustomImageDataset(Dataset):
244
+ def __init__(self, images, transform=None):
245
+ self.images = images
246
+ self.transform = transform
247
+
248
+ def __len__(self):
249
+ return len(self.images)
250
+
251
+ def __getitem__(self, idx):
252
+ image = self.images[idx]
253
+ if self.transform:
254
+ image = self.transform(image)
255
  return image
256
 
257
+ @spaces.GPU
258
+ def invert(image, mask, pcs=10000, epochs=400, weight_decay = 1e-10, lr=1e-1):
259
+ global unet
260
+ del unet
261
+ global network
262
+ unet, _, _, _, _ = load_models(device)
263
+
264
+ proj = torch.zeros(1,pcs).bfloat16().to(device)
265
+ network = LoRAw2w( proj, mean, std, v[:, :pcs],
266
+ unet,
267
+ rank=1,
268
+ multiplier=1.0,
269
+ alpha=27.0,
270
+ train_method="xattn-strict"
271
+ ).to(device, torch.bfloat16)
272
 
273
+ ### load mask
274
+ mask = transforms.Resize((64,64), interpolation=transforms.InterpolationMode.BILINEAR)(mask)
275
+ mask = torchvision.transforms.functional.pil_to_tensor(mask).unsqueeze(0).to(device).bfloat16()[:,0,:,:].unsqueeze(1)
276
+ ### check if an actual mask was draw, otherwise mask is just all ones
277
+ if torch.sum(mask) == 0:
278
+ mask = torch.ones((1,1,64,64)).to(device).bfloat16()
279
+
280
+ ### single image dataset
281
+ image_transforms = transforms.Compose([transforms.Resize(512, interpolation=transforms.InterpolationMode.BILINEAR),
282
+ transforms.RandomCrop(512),
283
+ transforms.ToTensor(),
284
+ transforms.Normalize([0.5], [0.5])])
285
+
286
+
287
+ train_dataset = CustomImageDataset(image, transform=image_transforms)
288
+ train_dataloader = torch.utils.data.DataLoader(train_dataset, batch_size=1, shuffle=True)
289
+
290
+ ### optimizer
291
+ optim = torch.optim.Adam(network.parameters(), lr=lr, weight_decay=weight_decay)
292
+
293
+ ### training loop
294
+ unet.train()
295
+ for epoch in tqdm.tqdm(range(epochs)):
296
+ for batch in train_dataloader:
297
+ ### prepare inputs
298
+ batch = batch.to(device).bfloat16()
299
+ latents = vae.encode(batch).latent_dist.sample()
300
+ latents = latents*0.18215
301
+ noise = torch.randn_like(latents)
302
+ bsz = latents.shape[0]
 
 
 
 
 
 
 
 
 
 
303
 
304
+ timesteps = torch.randint(0, noise_scheduler.config.num_train_timesteps, (bsz,), device=latents.device)
305
+ timesteps = timesteps.long()
306
+ noisy_latents = noise_scheduler.add_noise(latents, noise, timesteps)
307
+ text_input = tokenizer("sks person", padding="max_length", max_length=tokenizer.model_max_length, truncation=True, return_tensors="pt")
308
+ text_embeddings = text_encoder(text_input.input_ids.to(device))[0]
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
309
 
310
+ ### loss + sgd step
311
+ with network:
312
+ model_pred = unet(noisy_latents, timesteps, text_embeddings).sample
313
+ loss = torch.nn.functional.mse_loss(mask*model_pred.float(), mask*noise.float(), reduction="mean")
314
+ optim.zero_grad()
315
+ loss.backward()
316
+ optim.step()
 
 
 
 
317
 
318
+ ### return optimized network
319
+ return network
320
 
321
 
322
+ @spaces.GPU
323
+ def run_inversion(dict, pcs, epochs, weight_decay,lr):
324
+ global network
325
+ init_image = dict["image"].convert("RGB").resize((512, 512))
326
+ mask = dict["mask"].convert("RGB").resize((512, 512))
327
+ network = invert([init_image], mask, pcs, epochs, weight_decay,lr)
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
328
 
329
 
330
+ #sample an image
331
+ prompt = "sks person"
332
+ negative_prompt = "low quality, blurry, unfinished, nudity"
333
+ seed = 5
334
+ cfg = 3.0
335
+ steps = 25
336
+ image = inference( prompt, negative_prompt, cfg, steps, seed)
337
+ torch.save(network.proj, "model.pt" )
338
+ return image, "model.pt"
 
 
 
 
 
 
 
339
 
340
 
341
+ @spaces.GPU
342
+ def file_upload(file):
343
+ global unet
344
+ del unet
345
+ global network
346
+ global device
347
+
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
348
 
349
 
350
+ proj = torch.load(file.name).to(device)
351
+
352
+ #pad to 10000 Principal components to keep everything consistent
353
+ pcs = proj.shape[1]
354
+ padding = torch.zeros((1,10000-pcs)).to(device)
355
+ proj = torch.cat((proj, padding), 1)
356
 
357
+ unet, _, _, _, _ = load_models(device)
358
+
359
+
360
+ network = LoRAw2w( proj, mean, std, v[:, :10000],
361
+ unet,
362
+ rank=1,
363
+ multiplier=1.0,
364
+ alpha=27.0,
365
+ train_method="xattn-strict"
366
+ ).to(device, torch.bfloat16)
367
+
368
 
369
+ prompt = "sks person"
370
+ negative_prompt = "low quality, blurry, unfinished, nudity"
371
+ seed = 5
372
+ cfg = 3.0
373
+ steps = 25
374
+ image = inference( prompt, negative_prompt, cfg, steps, seed)
375
+ return image
 
 
 
 
 
 
 
376
 
377
 
378
+
379
 
380
+ intro = """
381
+ <div style="display: flex;align-items: center;justify-content: center">
382
+ <h1 style="margin-left: 12px;text-align: center;margin-bottom: 7px;display: inline-block"><em>weights2weights</em> Demo</h1>
383
+ <h3 style="display: inline-block;margin-left: 10px;margin-top: 6px;font-weight: 500">Interpreting the Weight Space of Customized Diffusion Models</h3>
384
+ </div>
385
+ <p style="font-size: 0.95rem;margin: 0rem;line-height: 1.2em;margin-top:1em;display: inline-block">
386
+ <a href="https://snap-research.github.io/weights2weights/" target="_blank">Project Page</a> | <a href="https://arxiv.org/abs/2406.09413" target="_blank">Paper</a>
387
+ | <a href="https://github.com/snap-research/weights2weights" target="_blank">Code</a> |
388
+ <a href="https://huggingface.co/spaces/Snapchat/w2w-demo?duplicate=true" target="_blank" style="
389
+ display: inline-block;
390
+ ">
391
+ <img style="margin-top: -1em;margin-bottom: 0em;position: absolute;" src="https://bit.ly/3CWLGkA" alt="Duplicate Space"></a>
392
+ </p>
393
+ """
394
+
395
+
396
+
397
  with gr.Blocks(css="style.css") as demo:
 
398
  gr.HTML(intro)
399
+
400
  gr.Markdown("""<div style="text-align: justify;"> In this demo, you can get an identity-encoding model by sampling or inverting. To use a model previously downloaded from this demo see \"Uploading a model\" in the Advanced Options. Next, you can generate new images from it, or edit the identity encoded in the model and generate images from the edited model. We provide detailed instructions and tips at the bottom of the page.""")
401
  with gr.Column():
402
  with gr.Row():
 
404
  gr.Markdown("""1) Either sample a new model, or upload an image (optionally draw a mask over the head) and click `invert`.""")
405
  sample = gr.Button("🎲 Sample New Model")
406
  input_image = gr.ImageEditor(elem_id="image_upload", type='pil', label="Reference Identity",
407
+ width=512, height=512)
408
+
409
  with gr.Row():
410
  invert_button = gr.Button("⬆️ Invert")
411
+
412
+
413
+
414
  with gr.Column():
415
  gr.Markdown("""2) Generate images of the sampled/inverted identity or edit the identity with the sliders and generate new images with various prompts and seeds.""")
416
  gallery = gr.Image(label="Generated Image",height=512, width=512, interactive=False)
417
  submit = gr.Button("Generate")
418
+
419
+
420
  prompt = gr.Textbox(label="Prompt",
421
+ info="Make sure to include 'sks person'" ,
422
+ placeholder="sks person",
423
+ value="sks person")
424
+
425
  seed = gr.Number(value=5, label="Seed", precision=0, interactive=True)
426
+
427
  # Editing
428
  with gr.Column():
429
  with gr.Row():
 
432
  with gr.Row():
433
  a3 = gr.Slider(label="- Curly Hair +", value=0, step=0.001, minimum=-1, maximum=1, interactive=True)
434
  a4 = gr.Slider(label="- Thick Eyebrows +", value=0, step=0.001, minimum=-1, maximum=1, interactive=True)
435
+
436
+
437
  with gr.Accordion("Advanced Options", open=False):
438
  with gr.Tab("Inversion"):
439
  with gr.Row():
 
444
  weight_decay = gr.Number(value=1e-10, label="Weight Decay", interactive=True)
445
  with gr.Tab("Sampling"):
446
  with gr.Row():
447
+ cfg= gr.Slider(label="CFG", value=3.0, step=0.1, minimum=0, maximum=10, interactive=True)
448
+ steps = gr.Slider(label="Inference Steps", value=25, step=1, minimum=0, maximum=100, interactive=True)
449
  with gr.Row():
450
+ negative_prompt = gr.Textbox(label="Negative Prompt", placeholder="low quality, blurry, unfinished, nudity, weapon", value="low quality, blurry, unfinished, nudity, weapon")
451
+ injection_step = gr.Slider(label="Injection Step", value=800, step=1, minimum=0, maximum=1000, interactive=True)
452
+
453
  with gr.Tab("Uploading a model"):
454
  gr.Markdown("""<div style="text-align: justify;">Upload a model below downloaded from this demo.""")
455
+
456
  file_input = gr.File(label="Upload Model", container=True)
457
+
458
+
459
+
460
+
461
  gr.Markdown("""<div style="text-align: justify;"> After sampling a new model or inverting, you can download the model below.""")
462
+
463
  with gr.Row():
464
  file_output = gr.File(label="Download Sampled/Inverted Model", container=True, interactive=False)
465
+
466
+
467
+
468
+
469
+ invert_button.click(fn=run_inversion,
470
+ inputs=[input_image, pcs, epochs, weight_decay,lr],
471
+ outputs = [input_image, file_output])
472
+
473
 
474
+ sample.click(fn=sample_then_run, outputs=[input_image, file_output])
475
+
 
 
 
 
 
 
 
 
476
  submit.click(
477
+ fn=edit_inference, inputs=[prompt, negative_prompt, cfg, steps, seed, injection_step, a1, a2, a3, a4], outputs=[gallery]
478
+ )
479
+ file_input.change(fn=file_upload, inputs=file_input, outputs = gallery)
480
+
481
+
482
+
483
  help_text1 = """
484
+ <b>Instructions</b>:
485
+ 1. To get results faster without waiting in queue, you can duplicate into a private space with an A100 GPU.
486
+ 2. To begin, you will have to get an identity-encoding model. You can either sample one from *weights2weights* space by clicking `Sample New Model` or by uploading an image and clicking `invert` to invert the identity into a model. You can optionally draw over the head to define a mask in the image for better results. Sampling a model takes around 10 seconds and inversion takes around 2 minutes. After this is done, you can optionally download this model for later use. A model can be uploaded in the \"Uploading a model\" tab in the `Advanced Options`.
487
+ 3. After getting a model, an image of the identity will be displayed on the right. You can sample from the model by changing seeds as well as prompts and then clicking `Generate`. Make sure to include \"sks person\" in your prompt to keep the same identity.
488
+ 4. The identity in the model can be edited by changing the sliders for various attributes. After clicking `Generate`, you can see how the identity has changed and the effects are maintained across different seeds and prompts.
489
+ """
490
  help_text2 = """<b>Tips</b>:
491
+ 1. Editing and Identity Generation
492
+ * If you are interested in preserving more of the image during identity-editing (i.e., where the same seed and prompt results in the same image with only the identity changed), you can play with the "Injection Step" parameter in the \"Sampling\" tab in the `Advanced Options`. During the first *n* timesteps, the original model's weights will be used, and then the edited weights will be set during the remaining steps. Values closer to 1000 will set the edited weights early, having a more pronounced effect, which may disrupt some semantics and structure of the generated image. Lower values will set the edited weights later, better preserving image context. We notice that around 600-800 tends to produce the best results. Larger values in the range (700-1000) are helpful for more global attribute changes, while smaller (400-700) can be used for more finegrained edits. Although it is not always needed.
493
+ * You can play around with negative prompts, number of inference steps, and CFG in the \"Sampling\" tab in the `Advanced Options` to affect the ultimate image quality.
494
+ * Sometimes the identity will not be perfectly consistent (e.g., there might be small variations of the face) when you use some seeds or prompts. This is a limitation of our method as well as an open-problem in personalized models.
495
+ 2. Inversion
496
+ * To obtain the best results for inversion, upload a high resolution photo of the face with minimal occlusion. It is recommended to draw over the face and hair to define a mask. But inversion should still work generally for non-closeup face shots.
497
+ * For inverting a realistic photo of an identity, typically 800 epochs with lr=1e-1 and 10,000 principal components (PCs) works well. If the resulting generations have artifacted and unrealstic textures, there is probably overfitting and you may want to reduce the number of epochs or learning rate, or play with weight decay. If the generations do not look like the input photo, then you may want to increase the number of epochs.
498
+ * For inverting out-of-distribution identities, such as artistic renditions of people or non-humans (e.g. the ones shown in the paper), it is recommended to use 1000 PCs, lr=1, and train for 800 epochs.
499
+ * Note that if you change the number of PCs, you will probably need to change the learning rate. For less PCs, higher learning rates are typically required."""
500
+
501
 
 
502
  gr.Markdown(help_text1)
503
  gr.Markdown(help_text2)
504
+ #demo.load(fn=start_items)
505
  demo.queue().launch()