mirror of
https://github.com/huggingface/diffusers.git
synced 2026-01-27 17:22:53 +03:00
7aa6af1138
* update
* update
* update
* update
* update
* merge main
* Revert "merge main"
This reverts commit 65efbcead5.
273 lines
9.5 KiB
Python
273 lines
9.5 KiB
Python
import tempfile
|
|
|
|
import numpy as np
|
|
import torch
|
|
from transformers import AutoTokenizer, T5EncoderModel
|
|
|
|
from diffusers import DDPMScheduler, UNet2DConditionModel
|
|
from diffusers.models.attention_processor import AttnAddedKVProcessor
|
|
from diffusers.pipelines.deepfloyd_if import IFWatermarker
|
|
|
|
from ...testing_utils import torch_device
|
|
from ..test_pipelines_common import to_np
|
|
|
|
|
|
# WARN: the hf-internal-testing/tiny-random-t5 text encoder has some non-determinism in the `save_load` tests.
|
|
|
|
|
|
class IFPipelineTesterMixin:
|
|
def _get_dummy_components(self):
|
|
torch.manual_seed(0)
|
|
text_encoder = T5EncoderModel.from_pretrained("hf-internal-testing/tiny-random-t5")
|
|
|
|
torch.manual_seed(0)
|
|
tokenizer = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-t5")
|
|
|
|
torch.manual_seed(0)
|
|
unet = UNet2DConditionModel(
|
|
sample_size=32,
|
|
layers_per_block=1,
|
|
block_out_channels=[32, 64],
|
|
down_block_types=[
|
|
"ResnetDownsampleBlock2D",
|
|
"SimpleCrossAttnDownBlock2D",
|
|
],
|
|
mid_block_type="UNetMidBlock2DSimpleCrossAttn",
|
|
up_block_types=["SimpleCrossAttnUpBlock2D", "ResnetUpsampleBlock2D"],
|
|
in_channels=3,
|
|
out_channels=6,
|
|
cross_attention_dim=32,
|
|
encoder_hid_dim=32,
|
|
attention_head_dim=8,
|
|
addition_embed_type="text",
|
|
addition_embed_type_num_heads=2,
|
|
cross_attention_norm="group_norm",
|
|
resnet_time_scale_shift="scale_shift",
|
|
act_fn="gelu",
|
|
)
|
|
unet.set_attn_processor(AttnAddedKVProcessor()) # For reproducibility tests
|
|
|
|
torch.manual_seed(0)
|
|
scheduler = DDPMScheduler(
|
|
num_train_timesteps=1000,
|
|
beta_schedule="squaredcos_cap_v2",
|
|
beta_start=0.0001,
|
|
beta_end=0.02,
|
|
thresholding=True,
|
|
dynamic_thresholding_ratio=0.95,
|
|
sample_max_value=1.0,
|
|
prediction_type="epsilon",
|
|
variance_type="learned_range",
|
|
)
|
|
|
|
torch.manual_seed(0)
|
|
watermarker = IFWatermarker()
|
|
|
|
return {
|
|
"text_encoder": text_encoder,
|
|
"tokenizer": tokenizer,
|
|
"unet": unet,
|
|
"scheduler": scheduler,
|
|
"watermarker": watermarker,
|
|
"safety_checker": None,
|
|
"feature_extractor": None,
|
|
}
|
|
|
|
def _get_superresolution_dummy_components(self):
|
|
torch.manual_seed(0)
|
|
text_encoder = T5EncoderModel.from_pretrained("hf-internal-testing/tiny-random-t5")
|
|
|
|
torch.manual_seed(0)
|
|
tokenizer = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-t5")
|
|
|
|
torch.manual_seed(0)
|
|
unet = UNet2DConditionModel(
|
|
sample_size=32,
|
|
layers_per_block=[1, 2],
|
|
block_out_channels=[32, 64],
|
|
down_block_types=[
|
|
"ResnetDownsampleBlock2D",
|
|
"SimpleCrossAttnDownBlock2D",
|
|
],
|
|
mid_block_type="UNetMidBlock2DSimpleCrossAttn",
|
|
up_block_types=["SimpleCrossAttnUpBlock2D", "ResnetUpsampleBlock2D"],
|
|
in_channels=6,
|
|
out_channels=6,
|
|
cross_attention_dim=32,
|
|
encoder_hid_dim=32,
|
|
attention_head_dim=8,
|
|
addition_embed_type="text",
|
|
addition_embed_type_num_heads=2,
|
|
cross_attention_norm="group_norm",
|
|
resnet_time_scale_shift="scale_shift",
|
|
act_fn="gelu",
|
|
class_embed_type="timestep",
|
|
mid_block_scale_factor=1.414,
|
|
time_embedding_act_fn="gelu",
|
|
time_embedding_dim=32,
|
|
)
|
|
unet.set_attn_processor(AttnAddedKVProcessor()) # For reproducibility tests
|
|
|
|
torch.manual_seed(0)
|
|
scheduler = DDPMScheduler(
|
|
num_train_timesteps=1000,
|
|
beta_schedule="squaredcos_cap_v2",
|
|
beta_start=0.0001,
|
|
beta_end=0.02,
|
|
thresholding=True,
|
|
dynamic_thresholding_ratio=0.95,
|
|
sample_max_value=1.0,
|
|
prediction_type="epsilon",
|
|
variance_type="learned_range",
|
|
)
|
|
|
|
torch.manual_seed(0)
|
|
image_noising_scheduler = DDPMScheduler(
|
|
num_train_timesteps=1000,
|
|
beta_schedule="squaredcos_cap_v2",
|
|
beta_start=0.0001,
|
|
beta_end=0.02,
|
|
)
|
|
|
|
torch.manual_seed(0)
|
|
watermarker = IFWatermarker()
|
|
|
|
return {
|
|
"text_encoder": text_encoder,
|
|
"tokenizer": tokenizer,
|
|
"unet": unet,
|
|
"scheduler": scheduler,
|
|
"image_noising_scheduler": image_noising_scheduler,
|
|
"watermarker": watermarker,
|
|
"safety_checker": None,
|
|
"feature_extractor": None,
|
|
}
|
|
|
|
# this test is modified from the base class because if pipelines set the text encoder
|
|
# as optional with the intention that the user is allowed to encode the prompt once
|
|
# and then pass the embeddings directly to the pipeline. The base class test uses
|
|
# the unmodified arguments from `self.get_dummy_inputs` which will pass the unencoded
|
|
# prompt to the pipeline when the text encoder is set to None, throwing an error.
|
|
# So we make the test reflect the intended usage of setting the text encoder to None.
|
|
def _test_save_load_optional_components(self):
|
|
components = self.get_dummy_components()
|
|
pipe = self.pipeline_class(**components)
|
|
pipe.to(torch_device)
|
|
pipe.set_progress_bar_config(disable=None)
|
|
|
|
inputs = self.get_dummy_inputs(torch_device)
|
|
|
|
prompt = inputs["prompt"]
|
|
generator = inputs["generator"]
|
|
num_inference_steps = inputs["num_inference_steps"]
|
|
output_type = inputs["output_type"]
|
|
|
|
if "image" in inputs:
|
|
image = inputs["image"]
|
|
else:
|
|
image = None
|
|
|
|
if "mask_image" in inputs:
|
|
mask_image = inputs["mask_image"]
|
|
else:
|
|
mask_image = None
|
|
|
|
if "original_image" in inputs:
|
|
original_image = inputs["original_image"]
|
|
else:
|
|
original_image = None
|
|
|
|
prompt_embeds, negative_prompt_embeds = pipe.encode_prompt(prompt)
|
|
|
|
# inputs with prompt converted to embeddings
|
|
inputs = {
|
|
"prompt_embeds": prompt_embeds,
|
|
"negative_prompt_embeds": negative_prompt_embeds,
|
|
"generator": generator,
|
|
"num_inference_steps": num_inference_steps,
|
|
"output_type": output_type,
|
|
}
|
|
|
|
if image is not None:
|
|
inputs["image"] = image
|
|
|
|
if mask_image is not None:
|
|
inputs["mask_image"] = mask_image
|
|
|
|
if original_image is not None:
|
|
inputs["original_image"] = original_image
|
|
|
|
# set all optional components to None
|
|
for optional_component in pipe._optional_components:
|
|
setattr(pipe, optional_component, None)
|
|
|
|
output = pipe(**inputs)[0]
|
|
|
|
with tempfile.TemporaryDirectory() as tmpdir:
|
|
pipe.save_pretrained(tmpdir)
|
|
pipe_loaded = self.pipeline_class.from_pretrained(tmpdir)
|
|
pipe_loaded.to(torch_device)
|
|
pipe_loaded.set_progress_bar_config(disable=None)
|
|
|
|
pipe_loaded.unet.set_attn_processor(AttnAddedKVProcessor()) # For reproducibility tests
|
|
|
|
for optional_component in pipe._optional_components:
|
|
self.assertTrue(
|
|
getattr(pipe_loaded, optional_component) is None,
|
|
f"`{optional_component}` did not stay set to None after loading.",
|
|
)
|
|
|
|
inputs = self.get_dummy_inputs(torch_device)
|
|
|
|
generator = inputs["generator"]
|
|
num_inference_steps = inputs["num_inference_steps"]
|
|
output_type = inputs["output_type"]
|
|
|
|
# inputs with prompt converted to embeddings
|
|
inputs = {
|
|
"prompt_embeds": prompt_embeds,
|
|
"negative_prompt_embeds": negative_prompt_embeds,
|
|
"generator": generator,
|
|
"num_inference_steps": num_inference_steps,
|
|
"output_type": output_type,
|
|
}
|
|
|
|
if image is not None:
|
|
inputs["image"] = image
|
|
|
|
if mask_image is not None:
|
|
inputs["mask_image"] = mask_image
|
|
|
|
if original_image is not None:
|
|
inputs["original_image"] = original_image
|
|
|
|
output_loaded = pipe_loaded(**inputs)[0]
|
|
|
|
max_diff = np.abs(to_np(output) - to_np(output_loaded)).max()
|
|
self.assertLess(max_diff, 1e-4)
|
|
|
|
# Modified from `PipelineTesterMixin` to set the attn processor as it's not serialized.
|
|
# This should be handled in the base test and then this method can be removed.
|
|
def _test_save_load_local(self):
|
|
components = self.get_dummy_components()
|
|
pipe = self.pipeline_class(**components)
|
|
pipe.to(torch_device)
|
|
pipe.set_progress_bar_config(disable=None)
|
|
|
|
inputs = self.get_dummy_inputs(torch_device)
|
|
output = pipe(**inputs)[0]
|
|
|
|
with tempfile.TemporaryDirectory() as tmpdir:
|
|
pipe.save_pretrained(tmpdir)
|
|
pipe_loaded = self.pipeline_class.from_pretrained(tmpdir)
|
|
pipe_loaded.to(torch_device)
|
|
pipe_loaded.set_progress_bar_config(disable=None)
|
|
|
|
pipe_loaded.unet.set_attn_processor(AttnAddedKVProcessor()) # For reproducibility tests
|
|
|
|
inputs = self.get_dummy_inputs(torch_device)
|
|
output_loaded = pipe_loaded(**inputs)[0]
|
|
|
|
max_diff = np.abs(to_np(output) - to_np(output_loaded)).max()
|
|
self.assertLess(max_diff, 1e-4)
|