ChenWu98 commited on
Commit
3e4185b
β€’
1 Parent(s): b84d23c

Update cac support

Browse files
Files changed (2) hide show
  1. app.py +2 -1
  2. ptp_utils.py +1 -156
app.py CHANGED
@@ -282,7 +282,8 @@ with gr.Blocks(css=css) as demo:
282
  </div>
283
  <p>
284
  Demo for CycleDiffusion with Stable Diffusion. <br>
285
- <a href="https://huggingface.co/docs/diffusers/main/en/api/pipelines/cycle_diffusion">🧨 Pipeline doc</a> | <a href="https://arxiv.org/abs/2210.05559">πŸ“„ Paper link</a>
 
286
  </p>
287
  <p>You can skip the queue in the colab: <a href="https://colab.research.google.com/gist/ChenWu98/0aa4fe7be80f6b45d3d055df9f14353a/copy-of-fine-tuned-diffusion-gradio.ipynb"><img data-canonical-src="https://colab.research.google.com/assets/colab-badge.svg" alt="Open In Colab" src="https://colab.research.google.com/assets/colab-badge.svg"></a></p>
288
  Running on <b>{device_print}</b>{(" in a <b>Google Colab</b>." if is_colab else "")}
 
282
  </div>
283
  <p>
284
  Demo for CycleDiffusion with Stable Diffusion. <br>
285
+ CycleDiffusion (<a href="https://github.com/ChenWu98/cycle-diffusion">Github</a> | <a href="https://arxiv.org/abs/2210.05559">πŸ“„ Paper link</a> | <a href="https://huggingface.co/docs/diffusers/main/en/api/pipelines/cycle_diffusion">🧨 Pipeline doc</a>) is an image-to-image translation method that supports stochastic samplers for diffusion models. <br>
286
+ It also supports Cross Attention Control (<a href="https://github.com/google/prompt-to-prompt">Github</a> | <a href="https://arxiv.org/abs/2208.01626">πŸ“„ Paper link</a>), which is a technique to transfer the attention map from the source prompt to the target prompt. <br>
287
  </p>
288
  <p>You can skip the queue in the colab: <a href="https://colab.research.google.com/gist/ChenWu98/0aa4fe7be80f6b45d3d055df9f14353a/copy-of-fine-tuned-diffusion-gradio.ipynb"><img data-canonical-src="https://colab.research.google.com/assets/colab-badge.svg" alt="Open In Colab" src="https://colab.research.google.com/assets/colab-badge.svg"></a></p>
289
  Running on <b>{device_print}</b>{(" in a <b>Google Colab</b>." if is_colab else "")}
ptp_utils.py CHANGED
@@ -14,162 +14,7 @@
14
 
15
  import numpy as np
16
  import torch
17
- from PIL import Image, ImageDraw, ImageFont
18
- import cv2
19
- from typing import Optional, Union, Tuple, List, Callable, Dict
20
- from IPython.display import display
21
- from tqdm.notebook import tqdm
22
-
23
-
24
- def text_under_image(image: np.ndarray, text: str, text_color: Tuple[int, int, int] = (0, 0, 0)):
25
- h, w, c = image.shape
26
- offset = int(h * .2)
27
- img = np.ones((h + offset, w, c), dtype=np.uint8) * 255
28
- font = cv2.FONT_HERSHEY_SIMPLEX
29
- # font = ImageFont.truetype("/usr/share/fonts/truetype/noto/NotoMono-Regular.ttf", font_size)
30
- img[:h] = image
31
- textsize = cv2.getTextSize(text, font, 1, 2)[0]
32
- text_x, text_y = (w - textsize[0]) // 2, h + offset - textsize[1] // 2
33
- cv2.putText(img, text, (text_x, text_y ), font, 1, text_color, 2)
34
- return img
35
-
36
-
37
- def view_images(images, num_rows=1, offset_ratio=0.02):
38
- if type(images) is list:
39
- num_empty = len(images) % num_rows
40
- elif images.ndim == 4:
41
- num_empty = images.shape[0] % num_rows
42
- else:
43
- images = [images]
44
- num_empty = 0
45
-
46
- empty_images = np.ones(images[0].shape, dtype=np.uint8) * 255
47
- images = [image.astype(np.uint8) for image in images] + [empty_images] * num_empty
48
- num_items = len(images)
49
-
50
- h, w, c = images[0].shape
51
- offset = int(h * offset_ratio)
52
- num_cols = num_items // num_rows
53
- image_ = np.ones((h * num_rows + offset * (num_rows - 1),
54
- w * num_cols + offset * (num_cols - 1), 3), dtype=np.uint8) * 255
55
- for i in range(num_rows):
56
- for j in range(num_cols):
57
- image_[i * (h + offset): i * (h + offset) + h:, j * (w + offset): j * (w + offset) + w] = images[
58
- i * num_cols + j]
59
-
60
- pil_img = Image.fromarray(image_)
61
- display(pil_img)
62
-
63
-
64
-
65
- def diffusion_step(model, controller, latents, context, t, guidance_scale, low_resource=False):
66
- if low_resource:
67
- noise_pred_uncond = model.unet(latents, t, encoder_hidden_states=context[0])["sample"]
68
- noise_prediction_text = model.unet(latents, t, encoder_hidden_states=context[1])["sample"]
69
- else:
70
- latents_input = torch.cat([latents] * 2)
71
- noise_pred = model.unet(latents_input, t, encoder_hidden_states=context)["sample"]
72
- noise_pred_uncond, noise_prediction_text = noise_pred.chunk(2)
73
- noise_pred = noise_pred_uncond + guidance_scale * (noise_prediction_text - noise_pred_uncond)
74
- latents = model.scheduler.step(noise_pred, t, latents)["prev_sample"]
75
- latents = controller.step_callback(latents)
76
- return latents
77
-
78
-
79
- def latent2image(vae, latents):
80
- latents = 1 / 0.18215 * latents
81
- image = vae.decode(latents)['sample']
82
- image = (image / 2 + 0.5).clamp(0, 1)
83
- image = image.cpu().permute(0, 2, 3, 1).numpy()
84
- image = (image * 255).astype(np.uint8)
85
- return image
86
-
87
-
88
- def init_latent(latent, model, height, width, generator, batch_size):
89
- if latent is None:
90
- latent = torch.randn(
91
- (1, model.unet.in_channels, height // 8, width // 8),
92
- generator=generator,
93
- )
94
- latents = latent.expand(batch_size, model.unet.in_channels, height // 8, width // 8).to(model.device)
95
- return latent, latents
96
-
97
-
98
- @torch.no_grad()
99
- def text2image_ldm(
100
- model,
101
- prompt: List[str],
102
- controller,
103
- num_inference_steps: int = 50,
104
- guidance_scale: Optional[float] = 7.,
105
- generator: Optional[torch.Generator] = None,
106
- latent: Optional[torch.FloatTensor] = None,
107
- ):
108
- register_attention_control(model, controller)
109
- height = width = 256
110
- batch_size = len(prompt)
111
-
112
- uncond_input = model.tokenizer([""] * batch_size, padding="max_length", max_length=77, return_tensors="pt")
113
- uncond_embeddings = model.bert(uncond_input.input_ids.to(model.device))[0]
114
-
115
- text_input = model.tokenizer(prompt, padding="max_length", max_length=77, return_tensors="pt")
116
- text_embeddings = model.bert(text_input.input_ids.to(model.device))[0]
117
- latent, latents = init_latent(latent, model, height, width, generator, batch_size)
118
- context = torch.cat([uncond_embeddings, text_embeddings])
119
-
120
- model.scheduler.set_timesteps(num_inference_steps)
121
- for t in tqdm(model.scheduler.timesteps):
122
- latents = diffusion_step(model, controller, latents, context, t, guidance_scale)
123
-
124
- image = latent2image(model.vqvae, latents)
125
-
126
- return image, latent
127
-
128
-
129
-
130
- @torch.no_grad()
131
- def text2image_ldm_stable(
132
- model,
133
- prompt: List[str],
134
- controller,
135
- num_inference_steps: int = 50,
136
- guidance_scale: float = 7.5,
137
- generator: Optional[torch.Generator] = None,
138
- latent: Optional[torch.FloatTensor] = None,
139
- low_resource: bool = False,
140
- ):
141
- register_attention_control(model, controller)
142
- height = width = 512
143
- batch_size = len(prompt)
144
-
145
- text_input = model.tokenizer(
146
- prompt,
147
- padding="max_length",
148
- max_length=model.tokenizer.model_max_length,
149
- truncation=True,
150
- return_tensors="pt",
151
- )
152
- text_embeddings = model.text_encoder(text_input.input_ids.to(model.device))[0]
153
- max_length = text_input.input_ids.shape[-1]
154
- uncond_input = model.tokenizer(
155
- [""] * batch_size, padding="max_length", max_length=max_length, return_tensors="pt"
156
- )
157
- uncond_embeddings = model.text_encoder(uncond_input.input_ids.to(model.device))[0]
158
-
159
- context = [uncond_embeddings, text_embeddings]
160
- if not low_resource:
161
- context = torch.cat(context)
162
- latent, latents = init_latent(latent, model, height, width, generator, batch_size)
163
-
164
- # set timesteps
165
- extra_set_kwargs = {"offset": 1}
166
- model.scheduler.set_timesteps(num_inference_steps, **extra_set_kwargs)
167
- for t in tqdm(model.scheduler.timesteps):
168
- latents = diffusion_step(model, controller, latents, context, t, guidance_scale, low_resource)
169
-
170
- image = latent2image(model.vae, latents)
171
-
172
- return image, latent
173
 
174
 
175
  def register_attention_control(model, controller):
 
14
 
15
  import numpy as np
16
  import torch
17
+ from typing import Optional, Union, Tuple, Dict
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
18
 
19
 
20
  def register_attention_control(model, controller):