Skip to content

Commit 57243fd

Browse files
anton-lanton-l
committed
GLIDE integration test
1 parent d10441d commit 57243fd

3 files changed

Lines changed: 24 additions & 12 deletions

File tree

‎src/diffusers/configuration_utils.py‎

Lines changed: 0 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -225,9 +225,6 @@ def _dict_from_json_file(cls, json_file: Union[str, os.PathLike]):
225225
text = reader.read()
226226
return json.loads(text)
227227

228-
def __eq__(self, other):
229-
return self.__dict__ == other.__dict__
230-
231228
def __repr__(self):
232229
return f"{self.__class__.__name__} {self.to_json_string()}"
233230

‎src/diffusers/pipelines/pipeline_glide.py‎

Lines changed: 8 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -832,12 +832,12 @@ def text_model_fn(x_t, ts, transformer_out, **kwargs):
832832

833833
# 1. Sample gaussian noise
834834
batch_size = 2 # second image is empty for classifier-free guidance
835-
image = self.text_noise_scheduler.sample_noise(
836-
(batch_size, self.text_unet.in_channels, 64, 64), device=torch_device, generator=generator
837-
)
835+
image = torch.randn(
836+
(batch_size, self.text_unet.in_channels, 64, 64), generator=generator
837+
).to(torch_device)
838838

839839
# 2. Encode tokens
840-
# an empty input is needed to guide the model away from (
840+
# an empty input is needed to guide the model away from it
841841
inputs = self.tokenizer([prompt, ""], padding="max_length", max_length=128, return_tensors="pt")
842842
input_ids = inputs["input_ids"].to(torch_device)
843843
attention_mask = inputs["attention_mask"].to(torch_device)
@@ -850,7 +850,7 @@ def text_model_fn(x_t, ts, transformer_out, **kwargs):
850850
mean, variance, log_variance, pred_xstart = self.p_mean_variance(
851851
text_model_fn, self.text_noise_scheduler, image, t, transformer_out=transformer_out
852852
)
853-
noise = self.text_noise_scheduler.sample_noise(image.shape, device=torch_device, generator=generator)
853+
noise = torch.randn(image.shape, generator=generator).to(torch_device)
854854
nonzero_mask = (t != 0).float().view(-1, *([1] * (len(image.shape) - 1))) # no noise when t == 0
855855
image = mean + nonzero_mask * torch.exp(0.5 * log_variance) * noise
856856

@@ -873,8 +873,8 @@ def text_model_fn(x_t, ts, transformer_out, **kwargs):
873873
self.upscale_unet.resolution,
874874
),
875875
generator=generator,
876-
)
877-
image = image.to(torch_device) * upsample_temp
876+
).to(torch_device)
877+
image = image * upsample_temp
878878

879879
num_trained_timesteps = self.upscale_noise_scheduler.timesteps
880880
inference_step_times = range(0, num_trained_timesteps, num_trained_timesteps // num_inference_steps_upscale)
@@ -896,7 +896,7 @@ def text_model_fn(x_t, ts, transformer_out, **kwargs):
896896
# 3. optionally sample variance
897897
variance = 0
898898
if eta > 0:
899-
noise = torch.randn(image.shape, generator=generator).to(image.device)
899+
noise = torch.randn(image.shape, generator=generator).to(torch_device)
900900
variance = (
901901
self.upscale_noise_scheduler.get_variance(t, num_inference_steps_upscale).sqrt() * eta * noise
902902
)

‎tests/test_modeling_utils.py‎

Lines changed: 16 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -19,7 +19,7 @@
1919

2020
import torch
2121

22-
from diffusers import DDIM, DDPM, PNDM, DDIMScheduler, DDPMScheduler, LatentDiffusion, PNDMScheduler, UNetModel
22+
from diffusers import DDIM, DDPM, PNDM, GLIDE, DDIMScheduler, DDPMScheduler, LatentDiffusion, PNDMScheduler, UNetModel
2323
from diffusers.configuration_utils import ConfigMixin
2424
from diffusers.pipeline_utils import DiffusionPipeline
2525
from diffusers.testing_utils import floats_tensor, slow, torch_device
@@ -212,3 +212,18 @@ def test_ldm_text2img(self):
212212
assert image.shape == (1, 3, 256, 256)
213213
expected_slice = torch.tensor([0.7295, 0.7358, 0.7256, 0.7435, 0.7095, 0.6884, 0.7325, 0.6921, 0.6458])
214214
assert (image_slice.flatten() - expected_slice).abs().max() < 1e-2
215+
216+
@slow
217+
def test_glide_text2img(self):
218+
model_id = "fusing/glide-base"
219+
glide = GLIDE.from_pretrained(model_id)
220+
221+
prompt = "a pencil sketch of a corgi"
222+
generator = torch.manual_seed(0)
223+
image = glide(prompt, generator=generator, num_inference_steps_upscale=20)
224+
225+
image_slice = image[0, :3, :3, -1].cpu()
226+
227+
assert image.shape == (1, 256, 256, 3)
228+
expected_slice = torch.tensor([0.7119, 0.7073, 0.6460, 0.7780, 0.7423, 0.6926, 0.7378, 0.7189, 0.7784])
229+
assert (image_slice.flatten() - expected_slice).abs().max() < 1e-2

0 commit comments

Comments
 (0)