mirror of
https://github.com/huggingface/diffusers.git
synced 2026-01-27 17:22:53 +03:00
85 lines
3.1 KiB
Python
85 lines
3.1 KiB
Python
import gc
|
|
|
|
import pytest
|
|
import torch
|
|
|
|
from diffusers import (
|
|
StableDiffusionUpscalePipeline,
|
|
)
|
|
from diffusers.utils import load_image
|
|
|
|
from ..testing_utils import (
|
|
backend_empty_cache,
|
|
enable_full_determinism,
|
|
numpy_cosine_similarity_distance,
|
|
require_torch_accelerator,
|
|
slow,
|
|
torch_device,
|
|
)
|
|
from .single_file_testing_utils import SDSingleFileTesterMixin
|
|
|
|
|
|
enable_full_determinism()
|
|
|
|
|
|
@slow
|
|
@require_torch_accelerator
|
|
class TestStableDiffusionUpscalePipelineSingleFileSlow(SDSingleFileTesterMixin):
|
|
pipeline_class = StableDiffusionUpscalePipeline
|
|
ckpt_path = "https://huggingface.co/stabilityai/stable-diffusion-x4-upscaler/blob/main/x4-upscaler-ema.safetensors"
|
|
original_config = "https://raw.githubusercontent.com/Stability-AI/stablediffusion/main/configs/stable-diffusion/x4-upscaling.yaml"
|
|
repo_id = "stabilityai/stable-diffusion-x4-upscaler"
|
|
|
|
def setup_method(self):
|
|
gc.collect()
|
|
backend_empty_cache(torch_device)
|
|
|
|
def teardown_method(self):
|
|
gc.collect()
|
|
backend_empty_cache(torch_device)
|
|
|
|
def test_single_file_format_inference_is_same_as_pretrained(self):
|
|
image = load_image(
|
|
"https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"
|
|
"/sd2-upscale/low_res_cat.png"
|
|
)
|
|
|
|
prompt = "a cat sitting on a park bench"
|
|
pipe = StableDiffusionUpscalePipeline.from_pretrained(self.repo_id)
|
|
pipe.enable_model_cpu_offload(device=torch_device)
|
|
|
|
generator = torch.Generator("cpu").manual_seed(0)
|
|
output = pipe(prompt=prompt, image=image, generator=generator, output_type="np", num_inference_steps=3)
|
|
image_from_pretrained = output.images[0]
|
|
|
|
pipe_from_single_file = StableDiffusionUpscalePipeline.from_single_file(self.ckpt_path)
|
|
pipe_from_single_file.enable_model_cpu_offload(device=torch_device)
|
|
|
|
generator = torch.Generator("cpu").manual_seed(0)
|
|
output_from_single_file = pipe_from_single_file(
|
|
prompt=prompt, image=image, generator=generator, output_type="np", num_inference_steps=3
|
|
)
|
|
image_from_single_file = output_from_single_file.images[0]
|
|
|
|
assert image_from_pretrained.shape == (512, 512, 3)
|
|
assert image_from_single_file.shape == (512, 512, 3)
|
|
assert (
|
|
numpy_cosine_similarity_distance(image_from_pretrained.flatten(), image_from_single_file.flatten()) < 1e-3
|
|
)
|
|
|
|
@pytest.mark.xfail(
|
|
condition=True,
|
|
reason="Test fails because of mismatches in the configs but it is very hard to properly fix this considering downstream usecase.",
|
|
strict=True,
|
|
)
|
|
def test_single_file_components_with_original_config(self):
|
|
super().test_single_file_components_with_original_config()
|
|
|
|
@pytest.mark.xfail(
|
|
condition=True,
|
|
reason="Test fails because of mismatches in the configs but it is very hard to properly fix this considering downstream usecase.",
|
|
strict=True,
|
|
)
|
|
def test_single_file_components_with_original_config_local_files_only(self):
|
|
super().test_single_file_components_with_original_config_local_files_only()
|