mirror of
https://github.com/huggingface/diffusers.git
synced 2026-01-27 17:22:53 +03:00
121 lines
4.4 KiB
Python
121 lines
4.4 KiB
Python
import gc
|
|
|
|
import pytest
|
|
import torch
|
|
|
|
from diffusers import (
|
|
StableDiffusionInpaintPipeline,
|
|
)
|
|
from diffusers.utils import load_image
|
|
|
|
from ..testing_utils import (
|
|
backend_empty_cache,
|
|
enable_full_determinism,
|
|
require_torch_accelerator,
|
|
slow,
|
|
torch_device,
|
|
)
|
|
from .single_file_testing_utils import SDSingleFileTesterMixin
|
|
|
|
|
|
enable_full_determinism()
|
|
|
|
|
|
@slow
|
|
@require_torch_accelerator
|
|
class TestStableDiffusionInpaintPipelineSingleFileSlow(SDSingleFileTesterMixin):
|
|
pipeline_class = StableDiffusionInpaintPipeline
|
|
ckpt_path = "https://huggingface.co/botp/stable-diffusion-v1-5-inpainting/blob/main/sd-v1-5-inpainting.ckpt"
|
|
original_config = "https://raw.githubusercontent.com/runwayml/stable-diffusion/main/configs/stable-diffusion/v1-inpainting-inference.yaml"
|
|
repo_id = "botp/stable-diffusion-v1-5-inpainting"
|
|
|
|
def setup_method(self):
|
|
gc.collect()
|
|
backend_empty_cache(torch_device)
|
|
|
|
def teardown_method(self):
|
|
gc.collect()
|
|
backend_empty_cache(torch_device)
|
|
|
|
def get_inputs(self, device, generator_device="cpu", dtype=torch.float32, seed=0):
|
|
generator = torch.Generator(device=generator_device).manual_seed(seed)
|
|
init_image = load_image(
|
|
"https://huggingface.co/datasets/diffusers/test-arrays/resolve/main"
|
|
"/stable_diffusion_inpaint/input_bench_image.png"
|
|
)
|
|
mask_image = load_image(
|
|
"https://huggingface.co/datasets/diffusers/test-arrays/resolve/main"
|
|
"/stable_diffusion_inpaint/input_bench_mask.png"
|
|
)
|
|
inputs = {
|
|
"prompt": "Face of a yellow cat, high resolution, sitting on a park bench",
|
|
"image": init_image,
|
|
"mask_image": mask_image,
|
|
"generator": generator,
|
|
"num_inference_steps": 3,
|
|
"guidance_scale": 7.5,
|
|
"output_type": "np",
|
|
}
|
|
return inputs
|
|
|
|
def test_single_file_format_inference_is_same_as_pretrained(self):
|
|
super().test_single_file_format_inference_is_same_as_pretrained(expected_max_diff=1e-3)
|
|
|
|
def test_single_file_loading_4_channel_unet(self):
|
|
# Test loading single file inpaint with a 4 channel UNet
|
|
ckpt_path = "https://huggingface.co/stable-diffusion-v1-5/stable-diffusion-v1-5/blob/main/v1-5-pruned-emaonly.safetensors"
|
|
pipe = self.pipeline_class.from_single_file(ckpt_path)
|
|
|
|
assert pipe.unet.config.in_channels == 4
|
|
|
|
@pytest.mark.skip(reason="runwayml original config has been removed")
|
|
def test_single_file_components_with_original_config(self):
|
|
return
|
|
|
|
@pytest.mark.skip(reason="runwayml original config has been removed")
|
|
def test_single_file_components_with_original_config_local_files_only(self):
|
|
return
|
|
|
|
|
|
@slow
|
|
@require_torch_accelerator
|
|
class TestStableDiffusion21InpaintPipelineSingleFileSlow(SDSingleFileTesterMixin):
|
|
pipeline_class = StableDiffusionInpaintPipeline
|
|
ckpt_path = (
|
|
"https://huggingface.co/stabilityai/stable-diffusion-2-inpainting/blob/main/512-inpainting-ema.safetensors"
|
|
)
|
|
original_config = "https://raw.githubusercontent.com/Stability-AI/stablediffusion/main/configs/stable-diffusion/v2-inpainting-inference.yaml"
|
|
repo_id = "stabilityai/stable-diffusion-2-inpainting"
|
|
|
|
def setup_method(self):
|
|
gc.collect()
|
|
backend_empty_cache(torch_device)
|
|
|
|
def teardown_method(self):
|
|
gc.collect()
|
|
backend_empty_cache(torch_device)
|
|
|
|
def get_inputs(self, device, generator_device="cpu", dtype=torch.float32, seed=0):
|
|
generator = torch.Generator(device=generator_device).manual_seed(seed)
|
|
init_image = load_image(
|
|
"https://huggingface.co/datasets/diffusers/test-arrays/resolve/main"
|
|
"/stable_diffusion_inpaint/input_bench_image.png"
|
|
)
|
|
mask_image = load_image(
|
|
"https://huggingface.co/datasets/diffusers/test-arrays/resolve/main"
|
|
"/stable_diffusion_inpaint/input_bench_mask.png"
|
|
)
|
|
inputs = {
|
|
"prompt": "Face of a yellow cat, high resolution, sitting on a park bench",
|
|
"image": init_image,
|
|
"mask_image": mask_image,
|
|
"generator": generator,
|
|
"num_inference_steps": 3,
|
|
"guidance_scale": 7.5,
|
|
"output_type": "np",
|
|
}
|
|
return inputs
|
|
|
|
def test_single_file_format_inference_is_same_as_pretrained(self):
|
|
super().test_single_file_format_inference_is_same_as_pretrained(expected_max_diff=1e-3)
|