mirror of
https://github.com/huggingface/diffusers.git
synced 2026-01-29 07:22:12 +03:00
b80790f95a
* ⚙️chore(train_controlnet) fix typo in logger message * ⚙️chore(models) refactor modules order; make them the same as calling order When printing the BasicTransformerBlock to stdout, I think it's crucial that the attributes order are shown in proper order. And also previously the "3. Feed Forward" comment was not making sense. It should have been close to self.ff but it's instead next to self.norm3 * correct many tests * remove bogus file * make style * correct more tests * finish tests * fix one more * make style * make unclip deterministic * ⚙️chore(models/attention) reorganize comments in BasicTransformerBlock class --------- Co-authored-by: Patrick von Platen <patrick.v.platen@gmail.com>
247 lines
8.8 KiB
Python
247 lines
8.8 KiB
Python
# coding=utf-8
|
|
# Copyright 2023 HuggingFace Inc.
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
|
|
import gc
|
|
import unittest
|
|
|
|
import numpy as np
|
|
import torch
|
|
from transformers import CLIPTextConfig, CLIPTextModel, CLIPTokenizer
|
|
|
|
from diffusers import (
|
|
AutoencoderKL,
|
|
DDIMScheduler,
|
|
EulerAncestralDiscreteScheduler,
|
|
PNDMScheduler,
|
|
StableDiffusionModelEditingPipeline,
|
|
UNet2DConditionModel,
|
|
)
|
|
from diffusers.utils import slow, torch_device
|
|
from diffusers.utils.testing_utils import require_torch_gpu, skip_mps
|
|
|
|
from ...pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_PARAMS
|
|
from ...test_pipelines_common import PipelineTesterMixin
|
|
|
|
|
|
torch.backends.cuda.matmul.allow_tf32 = False
|
|
|
|
|
|
@skip_mps
|
|
class StableDiffusionModelEditingPipelineFastTests(PipelineTesterMixin, unittest.TestCase):
|
|
pipeline_class = StableDiffusionModelEditingPipeline
|
|
params = TEXT_TO_IMAGE_PARAMS
|
|
batch_params = TEXT_TO_IMAGE_BATCH_PARAMS
|
|
|
|
def get_dummy_components(self):
|
|
torch.manual_seed(0)
|
|
unet = UNet2DConditionModel(
|
|
block_out_channels=(32, 64),
|
|
layers_per_block=2,
|
|
sample_size=32,
|
|
in_channels=4,
|
|
out_channels=4,
|
|
down_block_types=("DownBlock2D", "CrossAttnDownBlock2D"),
|
|
up_block_types=("CrossAttnUpBlock2D", "UpBlock2D"),
|
|
cross_attention_dim=32,
|
|
)
|
|
scheduler = DDIMScheduler()
|
|
torch.manual_seed(0)
|
|
vae = AutoencoderKL(
|
|
block_out_channels=[32, 64],
|
|
in_channels=3,
|
|
out_channels=3,
|
|
down_block_types=["DownEncoderBlock2D", "DownEncoderBlock2D"],
|
|
up_block_types=["UpDecoderBlock2D", "UpDecoderBlock2D"],
|
|
latent_channels=4,
|
|
)
|
|
torch.manual_seed(0)
|
|
text_encoder_config = CLIPTextConfig(
|
|
bos_token_id=0,
|
|
eos_token_id=2,
|
|
hidden_size=32,
|
|
intermediate_size=37,
|
|
layer_norm_eps=1e-05,
|
|
num_attention_heads=4,
|
|
num_hidden_layers=5,
|
|
pad_token_id=1,
|
|
vocab_size=1000,
|
|
)
|
|
text_encoder = CLIPTextModel(text_encoder_config)
|
|
tokenizer = CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip")
|
|
|
|
components = {
|
|
"unet": unet,
|
|
"scheduler": scheduler,
|
|
"vae": vae,
|
|
"text_encoder": text_encoder,
|
|
"tokenizer": tokenizer,
|
|
"safety_checker": None,
|
|
"feature_extractor": None,
|
|
}
|
|
return components
|
|
|
|
def get_dummy_inputs(self, device, seed=0):
|
|
generator = torch.manual_seed(seed)
|
|
inputs = {
|
|
"prompt": "A field of roses",
|
|
"generator": generator,
|
|
# Setting height and width to None to prevent OOMs on CPU.
|
|
"height": None,
|
|
"width": None,
|
|
"num_inference_steps": 2,
|
|
"guidance_scale": 6.0,
|
|
"output_type": "numpy",
|
|
}
|
|
return inputs
|
|
|
|
def test_stable_diffusion_model_editing_default_case(self):
|
|
device = "cpu" # ensure determinism for the device-dependent torch.Generator
|
|
components = self.get_dummy_components()
|
|
sd_pipe = StableDiffusionModelEditingPipeline(**components)
|
|
sd_pipe = sd_pipe.to(device)
|
|
sd_pipe.set_progress_bar_config(disable=None)
|
|
|
|
inputs = self.get_dummy_inputs(device)
|
|
image = sd_pipe(**inputs).images
|
|
image_slice = image[0, -3:, -3:, -1]
|
|
assert image.shape == (1, 64, 64, 3)
|
|
|
|
expected_slice = np.array([0.4755, 0.5132, 0.4976, 0.3904, 0.3554, 0.4765, 0.5139, 0.5158, 0.4889])
|
|
|
|
assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-2
|
|
|
|
def test_stable_diffusion_model_editing_negative_prompt(self):
|
|
device = "cpu" # ensure determinism for the device-dependent torch.Generator
|
|
components = self.get_dummy_components()
|
|
sd_pipe = StableDiffusionModelEditingPipeline(**components)
|
|
sd_pipe = sd_pipe.to(device)
|
|
sd_pipe.set_progress_bar_config(disable=None)
|
|
|
|
inputs = self.get_dummy_inputs(device)
|
|
negative_prompt = "french fries"
|
|
output = sd_pipe(**inputs, negative_prompt=negative_prompt)
|
|
image = output.images
|
|
image_slice = image[0, -3:, -3:, -1]
|
|
|
|
assert image.shape == (1, 64, 64, 3)
|
|
|
|
expected_slice = np.array([0.4992, 0.5101, 0.5004, 0.3949, 0.3604, 0.4735, 0.5216, 0.5204, 0.4913])
|
|
|
|
assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-2
|
|
|
|
def test_stable_diffusion_model_editing_euler(self):
|
|
device = "cpu" # ensure determinism for the device-dependent torch.Generator
|
|
components = self.get_dummy_components()
|
|
components["scheduler"] = EulerAncestralDiscreteScheduler(
|
|
beta_start=0.00085, beta_end=0.012, beta_schedule="scaled_linear"
|
|
)
|
|
sd_pipe = StableDiffusionModelEditingPipeline(**components)
|
|
sd_pipe = sd_pipe.to(device)
|
|
sd_pipe.set_progress_bar_config(disable=None)
|
|
|
|
inputs = self.get_dummy_inputs(device)
|
|
image = sd_pipe(**inputs).images
|
|
image_slice = image[0, -3:, -3:, -1]
|
|
|
|
assert image.shape == (1, 64, 64, 3)
|
|
|
|
expected_slice = np.array([0.4747, 0.5372, 0.4779, 0.4982, 0.5543, 0.4816, 0.5238, 0.4904, 0.5027])
|
|
|
|
assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-2
|
|
|
|
def test_stable_diffusion_model_editing_pndm(self):
|
|
device = "cpu" # ensure determinism for the device-dependent torch.Generator
|
|
components = self.get_dummy_components()
|
|
components["scheduler"] = PNDMScheduler()
|
|
sd_pipe = StableDiffusionModelEditingPipeline(**components)
|
|
sd_pipe = sd_pipe.to(device)
|
|
sd_pipe.set_progress_bar_config(disable=None)
|
|
|
|
inputs = self.get_dummy_inputs(device)
|
|
# the pipeline does not expect pndm so test if it raises error.
|
|
with self.assertRaises(ValueError):
|
|
_ = sd_pipe(**inputs).images
|
|
|
|
|
|
@slow
|
|
@require_torch_gpu
|
|
class StableDiffusionModelEditingSlowTests(unittest.TestCase):
|
|
def tearDown(self):
|
|
super().tearDown()
|
|
gc.collect()
|
|
torch.cuda.empty_cache()
|
|
|
|
def get_inputs(self, seed=0):
|
|
generator = torch.manual_seed(seed)
|
|
inputs = {
|
|
"prompt": "A field of roses",
|
|
"generator": generator,
|
|
"num_inference_steps": 3,
|
|
"guidance_scale": 7.5,
|
|
"output_type": "numpy",
|
|
}
|
|
return inputs
|
|
|
|
def test_stable_diffusion_model_editing_default(self):
|
|
model_ckpt = "CompVis/stable-diffusion-v1-4"
|
|
pipe = StableDiffusionModelEditingPipeline.from_pretrained(model_ckpt, safety_checker=None)
|
|
pipe.to(torch_device)
|
|
pipe.set_progress_bar_config(disable=None)
|
|
pipe.enable_attention_slicing()
|
|
|
|
inputs = self.get_inputs()
|
|
image = pipe(**inputs).images
|
|
image_slice = image[0, -3:, -3:, -1].flatten()
|
|
|
|
assert image.shape == (1, 512, 512, 3)
|
|
|
|
expected_slice = np.array(
|
|
[0.6749496, 0.6386453, 0.51443267, 0.66094905, 0.61921215, 0.5491332, 0.5744417, 0.58075106, 0.5174658]
|
|
)
|
|
|
|
assert np.abs(expected_slice - image_slice).max() < 1e-2
|
|
|
|
# make sure image changes after editing
|
|
pipe.edit_model("A pack of roses", "A pack of blue roses")
|
|
|
|
image = pipe(**inputs).images
|
|
image_slice = image[0, -3:, -3:, -1].flatten()
|
|
|
|
assert image.shape == (1, 512, 512, 3)
|
|
|
|
assert np.abs(expected_slice - image_slice).max() > 1e-1
|
|
|
|
def test_stable_diffusion_model_editing_pipeline_with_sequential_cpu_offloading(self):
|
|
torch.cuda.empty_cache()
|
|
torch.cuda.reset_max_memory_allocated()
|
|
torch.cuda.reset_peak_memory_stats()
|
|
|
|
model_ckpt = "CompVis/stable-diffusion-v1-4"
|
|
scheduler = DDIMScheduler.from_pretrained(model_ckpt, subfolder="scheduler")
|
|
pipe = StableDiffusionModelEditingPipeline.from_pretrained(
|
|
model_ckpt, scheduler=scheduler, safety_checker=None
|
|
)
|
|
pipe = pipe.to(torch_device)
|
|
pipe.set_progress_bar_config(disable=None)
|
|
pipe.enable_attention_slicing(1)
|
|
pipe.enable_sequential_cpu_offload()
|
|
|
|
inputs = self.get_inputs()
|
|
_ = pipe(**inputs)
|
|
|
|
mem_bytes = torch.cuda.max_memory_allocated()
|
|
# make sure that less than 4.4 GB is allocated
|
|
assert mem_bytes < 4.4 * 10**9
|