mirror of
https://github.com/huggingface/diffusers.git
synced 2026-01-27 17:22:53 +03:00
Remove unnecessary single file tests for SD Cascade UNet (#7996)
update
This commit is contained in:
Dhruv Nair
GitHub
@@ -1,191 +0,0 @@
|
||||
# coding=utf-8
|
||||
# Copyright 2024 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 torch
|
||||
|
||||
from diffusers import StableCascadeUNet
|
||||
from diffusers.utils import logging
|
||||
from diffusers.utils.testing_utils import (
|
||||
enable_full_determinism,
|
||||
numpy_cosine_similarity_distance,
|
||||
require_torch_gpu,
|
||||
slow,
|
||||
)
|
||||
from diffusers.utils.torch_utils import randn_tensor
|
||||
|
||||
|
||||
logger = logging.get_logger(__name__)
|
||||
|
||||
enable_full_determinism()
|
||||
|
||||
|
||||
@slow
|
||||
class StableCascadeUNetModelSlowTests(unittest.TestCase):
|
||||
def tearDown(self) -> None:
|
||||
super().tearDown()
|
||||
gc.collect()
|
||||
torch.cuda.empty_cache()
|
||||
|
||||
def test_stable_cascade_unet_prior_single_file_components(self):
|
||||
single_file_url = "https://huggingface.co/stabilityai/stable-cascade/blob/main/stage_c_bf16.safetensors"
|
||||
single_file_unet = StableCascadeUNet.from_single_file(single_file_url)
|
||||
|
||||
single_file_unet_config = single_file_unet.config
|
||||
del single_file_unet
|
||||
gc.collect()
|
||||
torch.cuda.empty_cache()
|
||||
|
||||
unet = StableCascadeUNet.from_pretrained("stabilityai/stable-cascade-prior", subfolder="prior", variant="bf16")
|
||||
unet_config = unet.config
|
||||
del unet
|
||||
gc.collect()
|
||||
torch.cuda.empty_cache()
|
||||
|
||||
PARAMS_TO_IGNORE = ["torch_dtype", "_name_or_path", "_use_default_values", "_diffusers_version"]
|
||||
for param_name, param_value in single_file_unet_config.items():
|
||||
if param_name in PARAMS_TO_IGNORE:
|
||||
continue
|
||||
|
||||
assert unet_config[param_name] == param_value
|
||||
|
||||
def test_stable_cascade_unet_decoder_single_file_components(self):
|
||||
single_file_url = "https://huggingface.co/stabilityai/stable-cascade/blob/main/stage_b_bf16.safetensors"
|
||||
single_file_unet = StableCascadeUNet.from_single_file(single_file_url)
|
||||
|
||||
single_file_unet_config = single_file_unet.config
|
||||
del single_file_unet
|
||||
gc.collect()
|
||||
torch.cuda.empty_cache()
|
||||
|
||||
unet = StableCascadeUNet.from_pretrained("stabilityai/stable-cascade", subfolder="decoder", variant="bf16")
|
||||
unet_config = unet.config
|
||||
del unet
|
||||
gc.collect()
|
||||
torch.cuda.empty_cache()
|
||||
|
||||
PARAMS_TO_IGNORE = ["torch_dtype", "_name_or_path", "_use_default_values", "_diffusers_version"]
|
||||
for param_name, param_value in single_file_unet_config.items():
|
||||
if param_name in PARAMS_TO_IGNORE:
|
||||
continue
|
||||
|
||||
assert unet_config[param_name] == param_value
|
||||
|
||||
def test_stable_cascade_unet_config_loading(self):
|
||||
config = StableCascadeUNet.load_config(
|
||||
pretrained_model_name_or_path="diffusers/stable-cascade-configs", subfolder="prior"
|
||||
)
|
||||
single_file_url = "https://huggingface.co/stabilityai/stable-cascade/blob/main/stage_c_bf16.safetensors"
|
||||
|
||||
single_file_unet = StableCascadeUNet.from_single_file(single_file_url, config=config)
|
||||
single_file_unet_config = single_file_unet.config
|
||||
del single_file_unet
|
||||
gc.collect()
|
||||
torch.cuda.empty_cache()
|
||||
|
||||
PARAMS_TO_IGNORE = ["torch_dtype", "_name_or_path", "_use_default_values", "_diffusers_version"]
|
||||
for param_name, param_value in config.items():
|
||||
if param_name in PARAMS_TO_IGNORE:
|
||||
continue
|
||||
|
||||
assert single_file_unet_config[param_name] == param_value
|
||||
|
||||
@require_torch_gpu
|
||||
def test_stable_cascade_unet_single_file_prior_forward_pass(self):
|
||||
dtype = torch.bfloat16
|
||||
generator = torch.Generator("cpu")
|
||||
|
||||
model_inputs = {
|
||||
"sample": randn_tensor((1, 16, 24, 24), generator=generator.manual_seed(0)).to("cuda", dtype),
|
||||
"timestep_ratio": torch.tensor([1]).to("cuda", dtype),
|
||||
"clip_text_pooled": randn_tensor((1, 1, 1280), generator=generator.manual_seed(0)).to("cuda", dtype),
|
||||
"clip_text": randn_tensor((1, 77, 1280), generator=generator.manual_seed(0)).to("cuda", dtype),
|
||||
"clip_img": randn_tensor((1, 1, 768), generator=generator.manual_seed(0)).to("cuda", dtype),
|
||||
"pixels": randn_tensor((1, 3, 8, 8), generator=generator.manual_seed(0)).to("cuda", dtype),
|
||||
}
|
||||
|
||||
unet = StableCascadeUNet.from_pretrained(
|
||||
"stabilityai/stable-cascade-prior",
|
||||
subfolder="prior",
|
||||
revision="refs/pr/2",
|
||||
variant="bf16",
|
||||
torch_dtype=dtype,
|
||||
)
|
||||
unet.to("cuda")
|
||||
with torch.no_grad():
|
||||
prior_output = unet(**model_inputs).sample.float().cpu().numpy()
|
||||
|
||||
# Remove UNet from GPU memory before loading the single file UNet model
|
||||
del unet
|
||||
gc.collect()
|
||||
torch.cuda.empty_cache()
|
||||
|
||||
single_file_url = "https://huggingface.co/stabilityai/stable-cascade/blob/main/stage_c_bf16.safetensors"
|
||||
single_file_unet = StableCascadeUNet.from_single_file(single_file_url, torch_dtype=dtype)
|
||||
single_file_unet.to("cuda")
|
||||
with torch.no_grad():
|
||||
prior_single_file_output = single_file_unet(**model_inputs).sample.float().cpu().numpy()
|
||||
|
||||
# Remove UNet from GPU memory before loading the single file UNet model
|
||||
del single_file_unet
|
||||
gc.collect()
|
||||
torch.cuda.empty_cache()
|
||||
|
||||
max_diff = numpy_cosine_similarity_distance(prior_output.flatten(), prior_single_file_output.flatten())
|
||||
assert max_diff < 8e-3
|
||||
|
||||
@require_torch_gpu
|
||||
def test_stable_cascade_unet_single_file_decoder_forward_pass(self):
|
||||
dtype = torch.float32
|
||||
generator = torch.Generator("cpu")
|
||||
|
||||
model_inputs = {
|
||||
"sample": randn_tensor((1, 4, 256, 256), generator=generator.manual_seed(0)).to("cuda", dtype),
|
||||
"timestep_ratio": torch.tensor([1]).to("cuda", dtype),
|
||||
"clip_text": randn_tensor((1, 77, 1280), generator=generator.manual_seed(0)).to("cuda", dtype),
|
||||
"clip_text_pooled": randn_tensor((1, 1, 1280), generator=generator.manual_seed(0)).to("cuda", dtype),
|
||||
"pixels": randn_tensor((1, 3, 8, 8), generator=generator.manual_seed(0)).to("cuda", dtype),
|
||||
}
|
||||
|
||||
unet = StableCascadeUNet.from_pretrained(
|
||||
"stabilityai/stable-cascade",
|
||||
subfolder="decoder",
|
||||
revision="refs/pr/44",
|
||||
torch_dtype=dtype,
|
||||
)
|
||||
unet.to("cuda")
|
||||
with torch.no_grad():
|
||||
prior_output = unet(**model_inputs).sample.float().cpu().numpy()
|
||||
|
||||
# Remove UNet from GPU memory before loading the single file UNet model
|
||||
del unet
|
||||
gc.collect()
|
||||
torch.cuda.empty_cache()
|
||||
|
||||
single_file_url = "https://huggingface.co/stabilityai/stable-cascade/blob/main/stage_b.safetensors"
|
||||
single_file_unet = StableCascadeUNet.from_single_file(single_file_url, torch_dtype=dtype)
|
||||
single_file_unet.to("cuda")
|
||||
with torch.no_grad():
|
||||
prior_single_file_output = single_file_unet(**model_inputs).sample.float().cpu().numpy()
|
||||
|
||||
# Remove UNet from GPU memory before loading the single file UNet model
|
||||
del single_file_unet
|
||||
gc.collect()
|
||||
torch.cuda.empty_cache()
|
||||
|
||||
max_diff = numpy_cosine_similarity_distance(prior_output.flatten(), prior_single_file_output.flatten())
|
||||
assert max_diff < 1e-4
|
||||
Reference in New Issue
Block a user