diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index fe2d3b21..1ccb773c 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -34,5 +34,6 @@ repos: types-tabulate, types-tqdm, types-urllib3, - horde_sdk==0.12.0, + horde_sdk==0.14.0, + horde_model_reference==0.8.1, ] diff --git a/hordelib/comfy_horde.py b/hordelib/comfy_horde.py index 996132b8..72109f48 100644 --- a/hordelib/comfy_horde.py +++ b/hordelib/comfy_horde.py @@ -366,11 +366,8 @@ def __init__( # Load our pipelines self._load_pipelines() - stdio = OutputCollector() - with contextlib.redirect_stdout(stdio): - # Load our custom nodes - self._load_custom_nodes() - stdio.replay() + # Load our custom nodes + self._load_custom_nodes() self._comfyui_callback = comfyui_callback diff --git a/hordelib/horde.py b/hordelib/horde.py index 404dea7e..fe273add 100644 --- a/hordelib/horde.py +++ b/hordelib/horde.py @@ -13,6 +13,7 @@ from enum import Enum, auto from types import FunctionType +from horde_model_reference.meta_consts import STABLE_DIFFUSION_BASELINE_CATEGORY, get_baseline_native_resolution from horde_sdk.ai_horde_api.apimodels import ImageGenerateJobPopResponse from horde_sdk.ai_horde_api.apimodels.base import ( GenMetadataEntry, @@ -78,15 +79,55 @@ def __init__( self.faults = faults -def _calc_upscale_sampler_steps(payload): - """Calculates the amount of hires_fix upscaler steps based on the denoising used and the steps used for the - primary image""" - upscale_steps = round(payload["ddim_steps"] * (0.9 - payload["hires_fix_denoising_strength"])) - if upscale_steps < 3: - upscale_steps = 3 +def _calc_upscale_sampler_steps( + payload: dict, +) -> int: + """Use `ImageUtils.calc_upscale_sampler_steps(...)` to calculate the number of steps for the upscale sampler. - logger.debug(f"Upscale steps calculated as {upscale_steps}") - return upscale_steps + Args: + payload (dict): The payload to use for the calculation. + + Returns: + int: The number of steps to use. + """ + model_name = payload.get("model_name") + baseline = None + native_resolution = None + if model_name is not None: + baseline = SharedModelManager.model_reference_manager.stable_diffusion.get_model_baseline(model_name) + if baseline is not None: + try: + baseline = STABLE_DIFFUSION_BASELINE_CATEGORY(baseline) + except ValueError: + baseline = None + logger.warning( + f"Model {model_name} has an invalid baseline {baseline} so we cannot calculate " + "hires fix upscale steps.", + ) + if baseline is not None: + native_resolution = get_baseline_native_resolution(baseline) + + width: int | None = payload.get("width") + height: int | None = payload.get("height") + hires_fix_denoising_strength: float | None = payload.get("hires_fix_denoising_strength") + ddim_steps: int | None = payload.get("ddim_steps") + + if width is None or height is None: + raise ValueError("Width and height must be set to calculate upscale sampler steps") + + if hires_fix_denoising_strength is None: + raise ValueError("Hires fix denoising strength must be set to calculate upscale sampler steps") + + if ddim_steps is None: + raise ValueError("DDIM steps must be set to calculate upscale sampler steps") + + return ImageUtils.calc_upscale_sampler_steps( + model_native_resolution=native_resolution, + width=width, + height=height, + hires_fix_denoising_strength=hires_fix_denoising_strength, + ddim_steps=ddim_steps, + ) class HordeLib: @@ -825,13 +866,15 @@ def _final_pipeline_adjustments(self, payload, pipeline_data) -> tuple[dict, lis raise RuntimeError(f"Invalid key {key}") elif "*" in key: key, multiplier = key.split("*", 1) - elif key in payload: + + if key in payload: if multiplier: pipeline_params[newkey] = round(payload.get(key) * float(multiplier)) else: pipeline_params[newkey] = payload.get(key) - else: + elif not isinstance(key, FunctionType): logger.error(f"Parameter {key} not found") + # We inject these parameters to ensure the HordeCheckpointLoader knows what file to load, if necessary # We don't want to hardcode this into the pipeline.json as we export this directly from ComfyUI # and don't want to have to rememebr to re-add those keys @@ -874,16 +917,22 @@ def _final_pipeline_adjustments(self, payload, pipeline_data) -> tuple[dict, lis baseline = None if model_details: baseline = model_details.get("baseline") - if baseline and (baseline == "stable_cascade" or baseline == "stable_diffusion_xl"): - new_width, new_height = ImageUtils.get_first_pass_image_resolution_max( - original_width, - original_height, - ) - else: - new_width, new_height = ImageUtils.get_first_pass_image_resolution_min( - original_width, - original_height, - ) + if baseline: + if baseline == "stable_cascade": + new_width, new_height = ImageUtils.get_first_pass_image_resolution_max( + original_width, + original_height, + ) + elif baseline == "stable_diffusion_xl": + new_width, new_height = ImageUtils.get_first_pass_image_resolution_sdxl( + original_width, + original_height, + ) + else: # fall through case; only `stable diffusion 1`` at time of writing + new_width, new_height = ImageUtils.get_first_pass_image_resolution_min( + original_width, + original_height, + ) # This is the *target* resolution pipeline_params["latent_upscale.width"] = original_width diff --git a/hordelib/initialisation.py b/hordelib/initialisation.py index 0577637c..014d8e4c 100644 --- a/hordelib/initialisation.py +++ b/hordelib/initialisation.py @@ -27,6 +27,7 @@ def initialise( force_normal_vram_mode: bool = True, extra_comfyui_args: list[str] | None = None, disable_smart_memory: bool = False, + do_not_load_model_mangers: bool = False, ): """Initialise hordelib. This is required before using any other hordelib functions. @@ -96,7 +97,7 @@ def initialise( # Initialise model manager from hordelib.shared_model_manager import SharedModelManager - SharedModelManager() + SharedModelManager(do_not_load_model_mangers=do_not_load_model_mangers) sys.argv = sys_arg_bkp diff --git a/hordelib/model_manager/base.py b/hordelib/model_manager/base.py index eca5042c..56abfe71 100644 --- a/hordelib/model_manager/base.py +++ b/hordelib/model_manager/base.py @@ -138,22 +138,7 @@ def load_model_database(self) -> None: ) def download_model_reference(self) -> dict: - try: - logger.debug(f"Downloading Model Reference for {self.models_db_name}") - response = requests.get(self.remote_db) - logger.debug("Downloaded Model Reference successfully") - models = response.json() - logger.info("Updated Model Reference from remote.") - return models - except Exception as e: # XXX Double check and/or rework this - logger.error( - f"Download failed: {e}", - ) - logger.warning("Model Reference not downloaded, using local copy") - if self.models_db_path.exists(): - return json.loads(self.models_db_path.read_text()) - logger.error("No local copy of Model Reference found!") - return {} + raise NotImplementedError("Downloading model databases is no longer supported within hordelib.") def get_free_ram_mb(self) -> int: """Returns the amount of free RAM in MB rounded down to the nearest integer. diff --git a/hordelib/model_manager/lora.py b/hordelib/model_manager/lora.py index 104095b7..04b3b9b1 100644 --- a/hordelib/model_manager/lora.py +++ b/hordelib/model_manager/lora.py @@ -44,7 +44,7 @@ class LoraModelManager(BaseModelManager): ) LORA_API = "https://civitai.com/api/v1/models?types=LORA&sort=Highest%20Rated&primaryFileOnly=true" MAX_RETRIES = 10 if not TESTS_ONGOING else 3 - MAX_DOWNLOAD_THREADS = 5 if not TESTS_ONGOING else 15 + MAX_DOWNLOAD_THREADS = 5 if not TESTS_ONGOING else 75 RETRY_DELAY = 3 if not TESTS_ONGOING else 0.2 """The time to wait between retries in seconds""" REQUEST_METADATA_TIMEOUT = 20 # Longer because civitai performs poorly on metadata requests for more than 5 models diff --git a/hordelib/nodes/facerestore_cf/__init__.py b/hordelib/nodes/facerestore_cf/__init__.py index 81c6c1c6..bcad12b1 100644 --- a/hordelib/nodes/facerestore_cf/__init__.py +++ b/hordelib/nodes/facerestore_cf/__init__.py @@ -8,6 +8,7 @@ import model_management import numpy as np import torch + # from comfy_extras.chainner_models import model_loading from hordelib.nodes.facerestore_cf.r_chainner import model_loading from torchvision.transforms.functional import normalize diff --git a/hordelib/nodes/facerestore_cf/r_chainner/gfpganv1_clean_arch.py b/hordelib/nodes/facerestore_cf/r_chainner/gfpganv1_clean_arch.py index 7f2f0e75..455bdbb3 100644 --- a/hordelib/nodes/facerestore_cf/r_chainner/gfpganv1_clean_arch.py +++ b/hordelib/nodes/facerestore_cf/r_chainner/gfpganv1_clean_arch.py @@ -72,16 +72,12 @@ def forward( if randomize_noise: noise = [None] * self.num_layers # for each style conv layer else: # use the stored noise - noise = [ - getattr(self.noises, f"noise{i}") for i in range(self.num_layers) - ] + noise = [getattr(self.noises, f"noise{i}") for i in range(self.num_layers)] # style truncation if truncation < 1: style_truncation = [] for style in styles: - style_truncation.append( - truncation_latent + truncation * (style - truncation_latent) - ) + style_truncation.append(truncation_latent + truncation * (style - truncation_latent)) styles = style_truncation # get style latents with injection if len(styles) == 1: @@ -96,9 +92,7 @@ def forward( if inject_index is None: inject_index = random.randint(1, self.num_latent - 1) latent1 = styles[0].unsqueeze(1).repeat(1, inject_index, 1) - latent2 = ( - styles[1].unsqueeze(1).repeat(1, self.num_latent - inject_index, 1) - ) + latent2 = styles[1].unsqueeze(1).repeat(1, self.num_latent - inject_index, 1) latent = torch.cat([latent1, latent2], 1) # main generation @@ -160,14 +154,10 @@ def __init__(self, in_channels, out_channels, mode="down"): def forward(self, x): out = F.leaky_relu_(self.conv1(x), negative_slope=0.2) # upsample/downsample - out = F.interpolate( - out, scale_factor=self.scale_factor, mode="bilinear", align_corners=False - ) + out = F.interpolate(out, scale_factor=self.scale_factor, mode="bilinear", align_corners=False) out = F.leaky_relu_(self.conv2(out), negative_slope=0.2) # skip - x = F.interpolate( - x, scale_factor=self.scale_factor, mode="bilinear", align_corners=False - ) + x = F.interpolate(x, scale_factor=self.scale_factor, mode="bilinear", align_corners=False) skip = self.skip(x) out = out + skip return out @@ -283,9 +273,7 @@ def __init__( # load pre-trained stylegan2 model if necessary if decoder_load_path: self.stylegan_decoder.load_state_dict( - torch.load( - decoder_load_path, map_location=lambda storage, loc: storage - )["params_ema"] + torch.load(decoder_load_path, map_location=lambda storage, loc: storage)["params_ema"] ) # fix decoder without updating params if fix_decoder: @@ -317,9 +305,7 @@ def __init__( ) self.load_state_dict(state_dict) - def forward( - self, x, return_latents=False, return_rgb=True, randomize_noise=True, **kwargs - ): + def forward(self, x, return_latents=False, return_rgb=True, randomize_noise=True, **kwargs): """Forward function for GFPGANv1Clean. Args: x (Tensor): Input images. diff --git a/hordelib/nodes/facerestore_cf/r_chainner/model_loading.py b/hordelib/nodes/facerestore_cf/r_chainner/model_loading.py index 598e605c..7d10bb1c 100644 --- a/hordelib/nodes/facerestore_cf/r_chainner/model_loading.py +++ b/hordelib/nodes/facerestore_cf/r_chainner/model_loading.py @@ -1,4 +1,3 @@ - from hordelib.nodes.facerestore_cf.r_chainner.gfpganv1_clean_arch import GFPGANv1Clean from hordelib.nodes.facerestore_cf.r_chainner.types import PyTorchModel @@ -21,9 +20,6 @@ def load_state_dict(state_dict) -> PyTorchModel: state_dict_keys = list(state_dict.keys()) # GFPGAN - if ( - "toRGB.0.weight" in state_dict_keys - and "stylegan_decoder.style_mlp.1.weight" in state_dict_keys - ): + if "toRGB.0.weight" in state_dict_keys and "stylegan_decoder.style_mlp.1.weight" in state_dict_keys: model = GFPGANv1Clean(state_dict) return model diff --git a/hordelib/nodes/facerestore_cf/r_chainner/stylegan2_clean_arch.py b/hordelib/nodes/facerestore_cf/r_chainner/stylegan2_clean_arch.py index c48de9af..ddd013ba 100644 --- a/hordelib/nodes/facerestore_cf/r_chainner/stylegan2_clean_arch.py +++ b/hordelib/nodes/facerestore_cf/r_chainner/stylegan2_clean_arch.py @@ -117,9 +117,7 @@ def forward(self, x, style): demod = torch.rsqrt(weight.pow(2).sum([2, 3, 4]) + self.eps) weight = weight * demod.view(b, self.out_channels, 1, 1, 1) - weight = weight.view( - b * self.out_channels, c, self.kernel_size, self.kernel_size - ) + weight = weight.view(b * self.out_channels, c, self.kernel_size, self.kernel_size) # upsample or downsample if necessary if self.sample_mode == "upsample": @@ -224,9 +222,7 @@ def forward(self, x, style, skip=None): out = out + self.bias if skip is not None: if self.upsample: - skip = F.interpolate( - skip, scale_factor=2, mode="bilinear", align_corners=False - ) + skip = F.interpolate(skip, scale_factor=2, mode="bilinear", align_corners=False) out = out + skip return out @@ -257,9 +253,7 @@ class StyleGAN2GeneratorClean(nn.Module): narrow (float): Narrow ratio for channels. Default: 1.0. """ - def __init__( - self, out_size, num_style_feat=512, num_mlp=8, channel_multiplier=2, narrow=1 - ): + def __init__(self, out_size, num_style_feat=512, num_mlp=8, channel_multiplier=2, narrow=1): super(StyleGAN2GeneratorClean, self).__init__() # Style MLP layers self.num_style_feat = num_style_feat @@ -362,9 +356,7 @@ def get_latent(self, x): return self.style_mlp(x) def mean_latent(self, num_latent): - latent_in = torch.randn( - num_latent, self.num_style_feat, device=self.constant_input.weight.device - ) + latent_in = torch.randn(num_latent, self.num_style_feat, device=self.constant_input.weight.device) latent = self.style_mlp(latent_in).mean(0, keepdim=True) return latent @@ -398,16 +390,12 @@ def forward( if randomize_noise: noise = [None] * self.num_layers # for each style conv layer else: # use the stored noise - noise = [ - getattr(self.noises, f"noise{i}") for i in range(self.num_layers) - ] + noise = [getattr(self.noises, f"noise{i}") for i in range(self.num_layers)] # style truncation if truncation < 1: style_truncation = [] for style in styles: - style_truncation.append( - truncation_latent + truncation * (style - truncation_latent) - ) + style_truncation.append(truncation_latent + truncation * (style - truncation_latent)) styles = style_truncation # get style latents with injection if len(styles) == 1: @@ -422,9 +410,7 @@ def forward( if inject_index is None: inject_index = random.randint(1, self.num_latent - 1) latent1 = styles[0].unsqueeze(1).repeat(1, inject_index, 1) - latent2 = ( - styles[1].unsqueeze(1).repeat(1, self.num_latent - inject_index, 1) - ) + latent2 = styles[1].unsqueeze(1).repeat(1, self.num_latent - inject_index, 1) latent = torch.cat([latent1, latent2], 1) # main generation diff --git a/hordelib/nodes/facerestore_cf/r_chainner/types.py b/hordelib/nodes/facerestore_cf/r_chainner/types.py index 20e39f68..db215364 100644 --- a/hordelib/nodes/facerestore_cf/r_chainner/types.py +++ b/hordelib/nodes/facerestore_cf/r_chainner/types.py @@ -1,4 +1,3 @@ - from typing import Union from hordelib.nodes.facerestore_cf.r_chainner.gfpganv1_clean_arch import GFPGANv1Clean @@ -11,7 +10,8 @@ def is_pytorch_face_model(model: object): return isinstance(model, PyTorchFaceModels) -PyTorchModels = (*PyTorchFaceModels, ) + +PyTorchModels = (*PyTorchFaceModels,) PyTorchModel = Union[PyTorchFaceModel] diff --git a/hordelib/shared_model_manager.py b/hordelib/shared_model_manager.py index e5464237..6ac25d3e 100644 --- a/hordelib/shared_model_manager.py +++ b/hordelib/shared_model_manager.py @@ -6,7 +6,7 @@ import torch from horde_model_reference import get_model_reference_file_path -from horde_model_reference.legacy import LegacyReferenceDownloadManager +from horde_model_reference.model_reference_manager import ModelReferenceManager from loguru import logger from typing_extensions import Self @@ -59,6 +59,7 @@ def do_migrations(): class SharedModelManager: _instance: Self = None # type: ignore manager: ModelManager + model_reference_manager: ModelReferenceManager cuda_available: bool def __new__(cls, do_not_load_model_mangers: bool = False): @@ -101,16 +102,17 @@ def load_model_managers( args_passed = locals().copy() # XXX This is temporary args_passed.pop("cls") # XXX This is temporary - lrdm = LegacyReferenceDownloadManager() if download_legacy_references: try: - lrdm.download_all_legacy_model_references(overwrite_existing=overwrite_existing_references) + cls.model_reference_manager = ModelReferenceManager(download_and_convert_legacy_dbs=True) except Exception as e: logger.error(f"Failed to download legacy model references: {e}") logger.error( "If this continues to happen, " "github may be down or your internet connection may be having issues.", ) + else: + cls.model_reference_manager = ModelReferenceManager() references = {} for reference in managers_to_load: diff --git a/hordelib/utils/image_utils.py b/hordelib/utils/image_utils.py index 19c0dd23..6d75937c 100644 --- a/hordelib/utils/image_utils.py +++ b/hordelib/utils/image_utils.py @@ -10,6 +10,20 @@ DEFAULT_IMAGE_MIN_RESOLUTION = 512 DEFAULT_HIGHER_RES_MAX_RESOLUTION = 1024 +IDEAL_SDXL_RESOLUTIONS = [ + (1024, 1024), + (1152, 896), + (896, 1152), + (1216, 832), + (832, 1216), + (1344, 768), + (768, 1344), + (1536, 640), + (640, 1536), +] + +IDEAL_SDXL_RESOLUTIONS_ASPECT_RATIOS = [width / height for width, height in IDEAL_SDXL_RESOLUTIONS] + class ImageUtils: @classmethod @@ -110,6 +124,103 @@ def get_first_pass_image_resolution_max( ) return width, height + @classmethod + def get_first_pass_image_resolution_sdxl( + cls, + width: int, + height: int, + ): + """Resize the image dimensions fit in one of the pre-defined SDXL resolution buckets which most closely + matches the aspect ratio of the image. + """ + + aspect_ratio = width / height + closest_aspect_ratio = min( + IDEAL_SDXL_RESOLUTIONS_ASPECT_RATIOS, + key=lambda x: abs(aspect_ratio - x), + ) + + index = IDEAL_SDXL_RESOLUTIONS_ASPECT_RATIOS.index(closest_aspect_ratio) + return IDEAL_SDXL_RESOLUTIONS[index] + + @staticmethod + def calc_upscale_sampler_steps( + model_native_resolution: int | None, + width: int, + height: int, + hires_fix_denoising_strength: float, + ddim_steps: int, + ) -> int: + """Calculate the number of upscale steps to use for the upscale sampler based on the input parameters. + + Note: The resulting values are non-linear to the input values. The heuristic is based on the native resolution + of the model, the requested resolution, the denoising strength and the number of steps used for the ddim + sampler. + + Args: + model_name (str | None): The model name to use for the calculation. + width (int): The width of the image to generate. + height (int): The height of the image to generate. + hires_fix_denoising_strength (float): The denoising strength to use for the upscale sampler. + ddim_steps (int): The number of steps used for the sampler. + + Returns: + int: The number of upscale steps to use for the upscale sampler. + """ + MIN_DENOISING_STRENGTH = 0.01 + MAX_DENOISING_STRENGTH = 1.0 + DECAY_RATE = 2 + """The rate at which the upscale steps decay based on the denoising strength""" + MIN_STEPS = 3 + """The minimum number of steps to use for the upscaling sampler""" + UPSCALE_ADJUSTMENT_FACTOR = 0.5 + """The factor by which the upscale steps are adjusted based on the native resolution distance factor""" + UPSCALE_DIVISOR = 2.25 + """The divisor used to adjust the upscale steps based on the native resolution distance factor""" + + STANDARD_RESOLUTION = 512 + + native_resolution_distance_factor: float = 0 + + if model_native_resolution is not None: + native_resolution_pixels = model_native_resolution * model_native_resolution + + requested_pixels = width * height + native_resolution_distance_factor = requested_pixels / native_resolution_pixels + + resolution_penalty = 3 * (STANDARD_RESOLUTION / model_native_resolution) + native_resolution_distance_factor /= resolution_penalty + + hires_fix_denoising_strength = max( + MIN_DENOISING_STRENGTH, + min(MAX_DENOISING_STRENGTH, hires_fix_denoising_strength), + ) + + scale = ddim_steps - MIN_STEPS + upscale_steps = round(MIN_STEPS + scale * (hires_fix_denoising_strength**DECAY_RATE)) + + # if native_resolution_distance_factor > NATIVE_RESOLUTION_THRESHOLD: + upscale_steps = round( + upscale_steps * ((1 / (UPSCALE_ADJUSTMENT_FACTOR**native_resolution_distance_factor)) / UPSCALE_DIVISOR), + ) + + logger.debug(f"Upscale steps calculated as {upscale_steps}") + + if ddim_steps <= 18: + logger.debug(f"Upscale steps increased by {MIN_STEPS} due to low requested ddim steps") + upscale_steps += MIN_STEPS + + if upscale_steps > ddim_steps: + logger.debug(f"Upscale steps adjusted to {ddim_steps} from {upscale_steps}") + upscale_steps = ddim_steps + + step_floor = min(6, ddim_steps) + if step_floor > upscale_steps: + logger.debug(f"Upscale steps adjusted to {step_floor} from {upscale_steps}") + upscale_steps = step_floor + + return upscale_steps + @classmethod def add_image_alpha_channel(cls, source_image, alpha_image): # Convert images to RGBA mode diff --git a/images_expected/lora_blue_hires_fix.png b/images_expected/lora_blue_hires_fix.png index 23550bce..8c15b701 100644 Binary files a/images_expected/lora_blue_hires_fix.png and b/images_expected/lora_blue_hires_fix.png differ diff --git a/images_expected/lora_character_hires_fix.png b/images_expected/lora_character_hires_fix.png index c3a3156c..9bb2d46a 100644 Binary files a/images_expected/lora_character_hires_fix.png and b/images_expected/lora_character_hires_fix.png differ diff --git a/images_expected/lora_character_hires_fix_sdxl.png b/images_expected/lora_character_hires_fix_sdxl.png new file mode 100644 index 00000000..52c042d2 Binary files /dev/null and b/images_expected/lora_character_hires_fix_sdxl.png differ diff --git a/images_expected/sdxl_text_to_image_hires_fix.png b/images_expected/sdxl_text_to_image_hires_fix.png index fb254dc6..bfcd4073 100644 Binary files a/images_expected/sdxl_text_to_image_hires_fix.png and b/images_expected/sdxl_text_to_image_hires_fix.png differ diff --git a/pyproject.toml b/pyproject.toml index 7922ba9d..f3d7bca2 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -60,6 +60,12 @@ filterwarnings = [ testpaths = [ "tests" ] +markers = [ + "slow: marks tests as slow (deselect with '-m \"not slow\"')", + "default_sd15_model: marks tests as default_sd15_model (deselect with '-m \"not default_sd15_model\"')", + "default_sdxl_model: marks tests as default_sdxl_model (deselect with '-m \"not default_sdxl_model\"')", + "refined_sdxl_model: marks tests as refined_sdxl_model (deselect with '-m \"not refined_sdxl_model\"')", +] [tool.black] line-length = 119 diff --git a/requirements.txt b/requirements.txt index 90cec3c1..861f5d3c 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,7 +1,7 @@ # Add this in for tox, comment out for build --extra-index-url https://download.pytorch.org/whl/cu121 -horde_sdk>=0.13.0 -horde_model_reference>=0.7.0 +horde_sdk>=0.14.0 +horde_model_reference>=0.8.1 pydantic numpy==1.26.4 torch>=2.3.1 diff --git a/tests/conftest.py b/tests/conftest.py index 2b9bbab0..89a007d0 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -14,6 +14,8 @@ from hordelib.model_manager.hyper import ALL_MODEL_MANAGER_TYPES from hordelib.shared_model_manager import SharedModelManager +from .testing_shared_classes import ResolutionTestCase + @pytest.fixture(scope="function", autouse=True) def line_break(): @@ -87,7 +89,12 @@ def init_horde( import hordelib - hordelib.initialise(setup_logging=True, logging_verbosity=5, disable_smart_memory=True) + hordelib.initialise( + setup_logging=True, + logging_verbosity=5, + disable_smart_memory=False, + do_not_load_model_mangers=True, + ) from hordelib.settings import UserSettings UserSettings.set_ram_to_leave_free_mb("100%") @@ -109,7 +116,7 @@ def shared_model_manager( custom_model_info_for_testing: tuple[str, str, str, str], hordelib_instance: HordeLib, ) -> Generator[type[SharedModelManager], None, None]: - SharedModelManager() + SharedModelManager(do_not_load_model_mangers=True) SharedModelManager.load_model_managers(ALL_MODEL_MANAGER_TYPES) assert SharedModelManager()._instance is not None @@ -240,6 +247,180 @@ def lora_GlowingRunesAI(shared_model_manager: type[SharedModelManager]) -> str: return name +@pytest.fixture(scope="session") +def default_min_steps() -> int: + return 6 + + +@pytest.fixture(scope="session") +def default_first_pass_steps() -> int: + return 30 + + +@pytest.fixture(scope="session") +def default_hires_fix_denoise_strength() -> float: + return 0.65 + + +@pytest.fixture(scope="session") +def sdxl_hires_test_cases( + default_first_pass_steps: int, + default_hires_fix_denoise_strength: int, + default_min_steps: int, +) -> list[ResolutionTestCase]: + sdxl_model_native_resolution = 1024 + sdxl_default_second_pass_steps = int(default_first_pass_steps / 2.25) + sdxl_high_second_pass_steps = default_first_pass_steps * 0.6 + return [ + ResolutionTestCase( + width=1024, + height=1024, + ddim_steps=default_first_pass_steps, + hires_fix_denoise_strength=default_hires_fix_denoise_strength, + model_native_resolution=sdxl_model_native_resolution, + max_expected_steps=default_first_pass_steps / 2, + min_expected_steps=default_min_steps, + ), + ResolutionTestCase( + width=1280, + height=1280, + ddim_steps=default_first_pass_steps, + hires_fix_denoise_strength=default_hires_fix_denoise_strength, + model_native_resolution=sdxl_model_native_resolution, + max_expected_steps=sdxl_high_second_pass_steps + default_min_steps, + min_expected_steps=default_min_steps, + ), + ResolutionTestCase( + width=1156, + height=1480, + ddim_steps=default_first_pass_steps, + hires_fix_denoise_strength=default_hires_fix_denoise_strength, + model_native_resolution=sdxl_model_native_resolution, + max_expected_steps=sdxl_high_second_pass_steps + default_min_steps, + min_expected_steps=sdxl_default_second_pass_steps, + ), + ResolutionTestCase( + width=2048, + height=2048, + ddim_steps=default_first_pass_steps, + hires_fix_denoise_strength=default_hires_fix_denoise_strength, + model_native_resolution=sdxl_model_native_resolution, + max_expected_steps=default_first_pass_steps, + min_expected_steps=default_first_pass_steps - default_min_steps, + ), + ResolutionTestCase( + width=1600, + height=1600, + ddim_steps=default_first_pass_steps, + hires_fix_denoise_strength=default_hires_fix_denoise_strength, + model_native_resolution=sdxl_model_native_resolution, + max_expected_steps=default_first_pass_steps, + min_expected_steps=sdxl_high_second_pass_steps - default_min_steps, + ), + ResolutionTestCase( + width=1664, + height=1152, + ddim_steps=default_first_pass_steps, + hires_fix_denoise_strength=default_hires_fix_denoise_strength, + model_native_resolution=sdxl_model_native_resolution, + max_expected_steps=default_first_pass_steps, + min_expected_steps=sdxl_high_second_pass_steps - default_min_steps, + ), + ResolutionTestCase( + height=1664, + width=1152, + ddim_steps=12, + hires_fix_denoise_strength=0.65, + model_native_resolution=1024, + max_expected_steps=12, + min_expected_steps=9, + ), + ] + + +@pytest.fixture(scope="session") +def sd15_hires_test_cases( + default_first_pass_steps: int, + default_hires_fix_denoise_strength: int, + default_min_steps: int, +) -> list[ResolutionTestCase]: + sd15_model_native_resolution = 512 + sd15_high_second_pass_steps = default_first_pass_steps * 0.67 + return [ + ResolutionTestCase( + width=512, + height=512, + ddim_steps=default_first_pass_steps, + hires_fix_denoise_strength=default_hires_fix_denoise_strength, + model_native_resolution=sd15_model_native_resolution, + max_expected_steps=sd15_high_second_pass_steps, + min_expected_steps=default_min_steps, + ), + ResolutionTestCase( + width=640, + height=640, + ddim_steps=default_first_pass_steps, + hires_fix_denoise_strength=default_hires_fix_denoise_strength, + model_native_resolution=sd15_model_native_resolution, + max_expected_steps=sd15_high_second_pass_steps + default_min_steps, + min_expected_steps=default_min_steps, + ), + ResolutionTestCase( + width=578, + height=740, + ddim_steps=default_first_pass_steps, + hires_fix_denoise_strength=default_hires_fix_denoise_strength, + model_native_resolution=sd15_model_native_resolution, + max_expected_steps=default_min_steps * 2.25, + min_expected_steps=default_min_steps * 1.25, + ), + ResolutionTestCase( + width=1024, + height=1024, + ddim_steps=default_first_pass_steps, + hires_fix_denoise_strength=default_hires_fix_denoise_strength, + model_native_resolution=sd15_model_native_resolution, + max_expected_steps=18, + min_expected_steps=default_min_steps * 2, + ), + ResolutionTestCase( + width=1536, + height=1024, + ddim_steps=default_first_pass_steps, + hires_fix_denoise_strength=default_hires_fix_denoise_strength, + model_native_resolution=sd15_model_native_resolution, + max_expected_steps=default_first_pass_steps, + min_expected_steps=default_first_pass_steps - default_min_steps, + ), + ResolutionTestCase( + width=800, + height=800, + ddim_steps=default_first_pass_steps, + hires_fix_denoise_strength=default_hires_fix_denoise_strength, + model_native_resolution=sd15_model_native_resolution, + max_expected_steps=default_min_steps * 2.5, + min_expected_steps=default_min_steps * 1.5, + ), + ResolutionTestCase( + width=2048, + height=2048, + ddim_steps=default_first_pass_steps, + hires_fix_denoise_strength=default_hires_fix_denoise_strength, + model_native_resolution=sd15_model_native_resolution, + max_expected_steps=default_first_pass_steps, + min_expected_steps=default_first_pass_steps, + ), + ] + + +@pytest.fixture(scope="session") +def all_hires_test_cases( + sdxl_hires_test_cases: list[ResolutionTestCase], + sd15_hires_test_cases: list[ResolutionTestCase], +) -> list[ResolutionTestCase]: + return sdxl_hires_test_cases + sd15_hires_test_cases + + def pytest_collection_modifyitems(items): """Modifies test items to ensure test modules run in a given order.""" MODULES_TO_RUN_FIRST = [ @@ -260,26 +441,54 @@ def pytest_collection_modifyitems(items): "tests.test_horde_inference", "tests.test_horde_inference_img2img", "tests.test_horde_inference_qrcode", - "tests.test_horde_inference_cascade", "tests.test_horde_samplers", "tests.test_horde_ti", "tests.test_horde_lcm", "tests.test_horde_lora", "tests.test_horde_inference_controlnet", "tests.test_horde_inference_painting", + "tests.test_horde_inference_cascade", ] module_mapping = {item: item.module.__name__ for item in items} sorted_items = [] + PYTEST_MARK_LAST = [ + "default_sdxl_model", + "refined_sdxl_model", + ] + for module in MODULES_TO_RUN_FIRST: - sorted_items.extend([item for item in items if module_mapping[item] == module]) + sorted_module = [item for item in items if module_mapping[item] == module] + # Any items with a mark in PYTEST_MARK_LAST will be moved to the end (in the order of the list) + for mark in PYTEST_MARK_LAST: + marked_items = [ + item for item in sorted_module if any(own_marker.name == mark for own_marker in item.own_markers) + ] + sorted_module = [item for item in sorted_module if item not in marked_items] + marked_items + + sorted_items.extend(sorted_module) sorted_items.extend( [item for item in items if module_mapping[item] not in MODULES_TO_RUN_FIRST + MODULES_TO_RUN_LAST], ) for module in MODULES_TO_RUN_LAST: - sorted_items.extend([item for item in items if module_mapping[item] == module]) + sorted_module = [item for item in items if module_mapping[item] == module] + # Any items with a mark in PYTEST_MARK_LAST will be moved to the end (in the order of the list) + for mark in PYTEST_MARK_LAST: + marked_items = [ + item for item in sorted_module if any(own_marker.name == mark for own_marker in item.own_markers) + ] + sorted_module = [item for item in sorted_module if item not in marked_items] + marked_items + + sorted_items.extend(sorted_module) + + # Any items with a mark in PYTEST_MARK_LAST will be moved to the end (in the order of the list) + # for mark in PYTEST_MARK_LAST: + # marked_items = [ + # item for item in sorted_items if any(own_marker.name == mark for own_marker in item.own_markers) + # ] + # sorted_items = [item for item in sorted_items if item not in marked_items] + marked_items items[:] = sorted_items diff --git a/tests/model_managers/test_mm_lora.py b/tests/model_managers/test_mm_lora.py index e89b299b..fedf47c8 100644 --- a/tests/model_managers/test_mm_lora.py +++ b/tests/model_managers/test_mm_lora.py @@ -114,7 +114,7 @@ def test_fetch_adhoc_lora_conflicting_fuzz(self): lora_model_manager.wait_for_downloads(600) lora_model_manager.wait_for_adhoc_reset(15) - lora_model_manager.fetch_adhoc_lora("33970") + lora_model_manager.fetch_adhoc_lora("33970", timeout=300) lora_model_manager.ensure_lora_deleted("Eula Genshin Impact | Character Lora 1644") lora_model_manager.fetch_adhoc_lora("Eula Genshin Impact | Character Lora 1644") assert lora_model_manager.get_lora_name("33970") == str("Dehya Genshin Impact | Character Lora 809".lower()) diff --git a/tests/test_horde_inference.py b/tests/test_horde_inference.py index fae205a2..c0c58b0b 100644 --- a/tests/test_horde_inference.py +++ b/tests/test_horde_inference.py @@ -9,6 +9,8 @@ class TestHordeInference: + + @pytest.mark.default_sd15_model def test_text_to_image( self, hordelib_instance: HordeLib, @@ -45,6 +47,7 @@ def test_text_to_image( pil_image, ) + @pytest.mark.default_sd15_model def test_text_to_image_n_iter( self, hordelib_instance: HordeLib, @@ -93,6 +96,7 @@ def test_text_to_image_n_iter( assert check_list_inference_images_similarity(img_pairs_to_check) + @pytest.mark.default_sdxl_model def test_sdxl_text_to_image( self, hordelib_instance: HordeLib, @@ -130,6 +134,7 @@ def test_sdxl_text_to_image( pil_image, ) + @pytest.mark.refined_sdxl_model def test_sdxl_text_to_image_hires_fix( self, hordelib_instance: HordeLib, @@ -166,6 +171,7 @@ def test_sdxl_text_to_image_hires_fix( pil_image, ) + @pytest.mark.default_sdxl_model @pytest.mark.skip(reason="This test is too slow to run on every test run") def test_sdxl_text_to_image_recommended_resolutions( self, @@ -222,6 +228,7 @@ def test_sdxl_text_to_image_recommended_resolutions( pil_image, ) + @pytest.mark.default_sd15_model def test_text_to_image_small( self, hordelib_instance: HordeLib, @@ -258,6 +265,7 @@ def test_text_to_image_small( pil_image, ) + @pytest.mark.default_sd15_model def test_text_to_image_clip_skip_2( self, hordelib_instance: HordeLib, @@ -294,6 +302,7 @@ def test_text_to_image_clip_skip_2( pil_image, ) + @pytest.mark.default_sd15_model def test_text_to_image_hires_fix( self, hordelib_instance: HordeLib, @@ -330,6 +339,7 @@ def test_text_to_image_hires_fix( pil_image, ) + @pytest.mark.default_sd15_model def test_text_to_image_tiling( self, hordelib_instance: HordeLib, @@ -379,6 +389,7 @@ def test_text_to_image_tiling( pil_image_no_tiling, ) + @pytest.mark.default_sd15_model def test_text_to_image_hires_fix_n_iter( self, hordelib_instance: HordeLib, @@ -422,6 +433,7 @@ def test_text_to_image_hires_fix_n_iter( assert check_list_inference_images_similarity(img_pairs_to_check) + @pytest.mark.default_sd15_model def test_callback_with_post_processors( self, hordelib_instance: HordeLib, diff --git a/tests/test_horde_inference_cascade.py b/tests/test_horde_inference_cascade.py index 54efb8a1..dfe92f81 100644 --- a/tests/test_horde_inference_cascade.py +++ b/tests/test_horde_inference_cascade.py @@ -11,6 +11,8 @@ class TestHordeInferenceCascade: + + @pytest.mark.slow def test_cascade_text_to_image( self, hordelib_instance: HordeLib, @@ -51,6 +53,7 @@ def test_cascade_text_to_image( pil_image, ) + @pytest.mark.slow def test_cascade_text_to_image_n_iter( self, hordelib_instance: HordeLib, @@ -102,6 +105,7 @@ def test_cascade_text_to_image_n_iter( assert check_list_inference_images_similarity(img_pairs_to_check) + @pytest.mark.slow def test_cascade_image_to_image( self, stable_cascade_base_model_name: str, @@ -150,6 +154,7 @@ def test_cascade_image_to_image( assert check_list_inference_images_similarity(img_pairs_to_check) + @pytest.mark.slow def test_cascade_image_remix_single( self, stable_cascade_base_model_name: str, @@ -188,6 +193,7 @@ def test_cascade_image_remix_single( pil_image, ) + @pytest.mark.slow def test_cascade_image_remix_double( self, stable_cascade_base_model_name: str, @@ -231,6 +237,7 @@ def test_cascade_image_remix_double( pil_image, ) + @pytest.mark.slow def test_cascade_image_remix_double_weak( self, stable_cascade_base_model_name: str, @@ -275,6 +282,7 @@ def test_cascade_image_remix_double_weak( pil_image, ) + @pytest.mark.slow def test_cascade_image_remix_triple( self, stable_cascade_base_model_name: str, @@ -321,6 +329,7 @@ def test_cascade_image_remix_triple( pil_image, ) + @pytest.mark.slow def test_cascade_text_to_image_hires_2pass( self, hordelib_instance: HordeLib, diff --git a/tests/test_horde_inference_controlnet.py b/tests/test_horde_inference_controlnet.py index 843d380b..b8a6e93c 100644 --- a/tests/test_horde_inference_controlnet.py +++ b/tests/test_horde_inference_controlnet.py @@ -16,6 +16,7 @@ def setup_and_teardown(self, shared_model_manager: type[SharedModelManager]): for preproc in HordeLib.CONTROLNET_IMAGE_PREPROCESSOR_MAP.keys(): shared_model_manager.manager.controlnet.download_control_type(preproc, ["stable diffusion 1"]) + @pytest.mark.default_sd15_model def test_controlnet_sd1( self, shared_model_manager: type[SharedModelManager], @@ -68,6 +69,7 @@ def test_controlnet_sd1( pil_image, ) + @pytest.mark.default_sd15_model def test_controlnet_fake_cn( self, hordelib_instance: HordeLib, @@ -99,6 +101,8 @@ def test_controlnet_fake_cn( image = hordelib_instance.basic_inference_single_image(data).image assert image is not None + @pytest.mark.slow + @pytest.mark.default_sd15_model def test_controlnet_strength( self, hordelib_instance: HordeLib, @@ -145,6 +149,7 @@ def test_controlnet_strength( pil_image, ) + @pytest.mark.default_sd15_model def test_controlnet_hires_fix( self, hordelib_instance: HordeLib, @@ -186,6 +191,7 @@ def test_controlnet_hires_fix( pil_image.save(f"images/{img_filename}", quality=100) images_to_compare.append((f"images_expected/{img_filename}", pil_image)) + @pytest.mark.default_sd15_model def test_controlnet_image_is_control( self, hordelib_instance: HordeLib, @@ -230,6 +236,7 @@ def test_controlnet_image_is_control( pil_image, ) + @pytest.mark.default_sd15_model def test_controlnet_n_iter( self, hordelib_instance: HordeLib, diff --git a/tests/test_horde_inference_img2img.py b/tests/test_horde_inference_img2img.py index 3d150266..c34c9e4d 100644 --- a/tests/test_horde_inference_img2img.py +++ b/tests/test_horde_inference_img2img.py @@ -1,5 +1,6 @@ # test_horde.py +import pytest from PIL import Image from hordelib.horde import HordeLib @@ -8,6 +9,8 @@ class TestHordeInferenceImg2Img: + + @pytest.mark.default_sd15_model def test_image_to_image( self, stable_diffusion_model_name_for_testing: str, @@ -47,6 +50,7 @@ def test_image_to_image( pil_image, ) + @pytest.mark.default_sd15_model def test_image_to_image_tiling( self, stable_diffusion_model_name_for_testing: str, @@ -86,6 +90,7 @@ def test_image_to_image_tiling( pil_image, ) + @pytest.mark.default_sdxl_model def test_sdxl_image_to_image( self, hordelib_instance: HordeLib, @@ -126,6 +131,7 @@ def test_sdxl_image_to_image( pil_image, ) + @pytest.mark.default_sd15_model def test_image_to_image_hires_fix_small( self, stable_diffusion_model_name_for_testing: str, @@ -165,6 +171,7 @@ def test_image_to_image_hires_fix_small( pil_image, ) + @pytest.mark.default_sd15_model def test_image_to_image_hires_fix_large( self, stable_diffusion_model_name_for_testing: str, @@ -203,6 +210,7 @@ def test_image_to_image_hires_fix_large( pil_image, ) + @pytest.mark.default_sd15_model def test_img2img_masked_denoise_1( self, stable_diffusion_model_name_for_testing: str, @@ -239,6 +247,7 @@ def test_img2img_masked_denoise_1( pil_image, ) + @pytest.mark.default_sd15_model def test_img2img_masked_denoise_high( self, stable_diffusion_model_name_for_testing: str, @@ -273,6 +282,7 @@ def test_img2img_masked_denoise_high( pil_image, ) + @pytest.mark.default_sd15_model def test_img2img_masked_denoise_mid( self, stable_diffusion_model_name_for_testing: str, @@ -307,6 +317,7 @@ def test_img2img_masked_denoise_mid( pil_image, ) + @pytest.mark.default_sd15_model def test_img2img_masked_denoise_low( self, stable_diffusion_model_name_for_testing: str, @@ -341,6 +352,7 @@ def test_img2img_masked_denoise_low( pil_image, ) + @pytest.mark.default_sdxl_model def test_sdxl_img2img_masked_denoise_95( self, sdxl_1_0_base_model_name: str, @@ -375,6 +387,7 @@ def test_sdxl_img2img_masked_denoise_95( pil_image, ) + @pytest.mark.default_sd15_model def test_image_to_faulty_source_image( self, stable_diffusion_model_name_for_testing: str, @@ -413,6 +426,7 @@ def test_image_to_faulty_source_image( pil_image, ) + @pytest.mark.default_sd15_model def test_image_to_image_n_iter( self, hordelib_instance: HordeLib, @@ -458,6 +472,7 @@ def test_image_to_image_n_iter( assert check_list_inference_images_similarity(img_pairs_to_check) + @pytest.mark.default_sd15_model def test_img2img_masked_n_iter( self, stable_diffusion_model_name_for_testing: str, @@ -498,6 +513,7 @@ def test_img2img_masked_n_iter( assert check_list_inference_images_similarity(img_pairs_to_check) + @pytest.mark.default_sd15_model def test_image_to_image_hires_fix_n_iter( self, stable_diffusion_model_name_for_testing: str, diff --git a/tests/test_horde_inference_layerdiffusion.py b/tests/test_horde_inference_layerdiffusion.py index 8e95814a..8c54e263 100644 --- a/tests/test_horde_inference_layerdiffusion.py +++ b/tests/test_horde_inference_layerdiffusion.py @@ -1,5 +1,6 @@ # test_horde.py +import pytest from PIL import Image from hordelib.horde import HordeLib @@ -8,6 +9,8 @@ class TestHordeInferenceTransparent: + + @pytest.mark.default_sd15_model def test_layerdiffuse_sd15( self, hordelib_instance: HordeLib, @@ -50,6 +53,7 @@ def test_layerdiffuse_sd15( image_result.image.save(f"images/{img_filename}", quality=100) img_pairs_to_check.append((f"images_expected/{img_filename}", image_result.image)) + @pytest.mark.refined_sdxl_model def test_layerdiffuse_sdxl( self, hordelib_instance: HordeLib, @@ -93,6 +97,7 @@ def test_layerdiffuse_sdxl( image_result.image.save(f"images/{img_filename}", quality=100) img_pairs_to_check.append((f"images_expected/{img_filename}", image_result.image)) + @pytest.mark.default_sd15_model def test_layerdiffusion_hires_fix( self, hordelib_instance: HordeLib, diff --git a/tests/test_horde_inference_qr_code.py b/tests/test_horde_inference_qr_code.py index 50df141e..62f46aaa 100644 --- a/tests/test_horde_inference_qr_code.py +++ b/tests/test_horde_inference_qr_code.py @@ -15,6 +15,7 @@ def setup_and_teardown(self, shared_model_manager: type[SharedModelManager]): for preproc in HordeLib.CONTROLNET_IMAGE_PREPROCESSOR_MAP.keys(): shared_model_manager.manager.controlnet.download_control_type(preproc, ["stable diffusion 1"]) + @pytest.mark.default_sd15_model def test_qr_code_inference( self, shared_model_manager: type[SharedModelManager], @@ -70,6 +71,7 @@ def test_qr_code_inference( pil_image, ) + @pytest.mark.default_sd15_model def test_qr_code_inference_out_of_bounds( self, shared_model_manager: type[SharedModelManager], @@ -125,6 +127,7 @@ def test_qr_code_inference_out_of_bounds( pil_image, ) + @pytest.mark.refined_sdxl_model def test_qr_code_inference_xl( self, shared_model_manager: type[SharedModelManager], @@ -174,6 +177,7 @@ def test_qr_code_inference_xl( pil_image, ) + @pytest.mark.refined_sdxl_model def test_qr_code_inference_too_large_text( self, shared_model_manager: type[SharedModelManager], @@ -241,6 +245,7 @@ def test_qr_code_inference_too_large_text( pil_image, ) + @pytest.mark.refined_sdxl_model def test_qr_code_control_strength( self, shared_model_manager: type[SharedModelManager], @@ -290,6 +295,7 @@ def test_qr_code_control_strength( pil_image, ) + @pytest.mark.refined_sdxl_model def test_qr_code_control_non_square( self, shared_model_manager: type[SharedModelManager], @@ -339,6 +345,7 @@ def test_qr_code_control_non_square( pil_image, ) + @pytest.mark.refined_sdxl_model def test_qr_code_control_qr_texts( self, shared_model_manager: type[SharedModelManager], diff --git a/tests/test_horde_lcm.py b/tests/test_horde_lcm.py index 969f0147..6e12ce89 100644 --- a/tests/test_horde_lcm.py +++ b/tests/test_horde_lcm.py @@ -16,11 +16,12 @@ def setup_and_teardown(self, shared_model_manager: type[SharedModelManager]): shared_model_manager.manager.lora.wait_for_downloads() yield + @pytest.mark.refined_sdxl_model def test_use_lcm_turbomix_lora_euler_a( self, shared_model_manager: type[SharedModelManager], hordelib_instance: HordeLib, - stable_diffusion_model_name_for_testing: str, + sdxl_refined_model_name: str, ): assert shared_model_manager.manager.lora @@ -49,7 +50,7 @@ def test_use_lcm_turbomix_lora_euler_a( "loras": [{"name": lora_name, "model": 1.0, "clip": 1.0, "inject_trigger": "any", "is_version": True}], "ddim_steps": 5, "n_iter": 1, - "model": "AlbedoBase XL (SDXL)", + "model": sdxl_refined_model_name, } pil_image = hordelib_instance.basic_inference_single_image(data).image @@ -63,11 +64,12 @@ def test_use_lcm_turbomix_lora_euler_a( pil_image, ) + @pytest.mark.refined_sdxl_model def test_use_lcm_turbomix_lora_lcm( self, shared_model_manager: type[SharedModelManager], hordelib_instance: HordeLib, - stable_diffusion_model_name_for_testing: str, + sdxl_refined_model_name: str, ): assert shared_model_manager.manager.lora @@ -96,7 +98,7 @@ def test_use_lcm_turbomix_lora_lcm( "loras": [{"name": lora_name, "model": 1.0, "clip": 1.0, "inject_trigger": "any", "is_version": True}], "ddim_steps": 4, "n_iter": 1, - "model": "AlbedoBase XL (SDXL)", + "model": sdxl_refined_model_name, } pil_image = hordelib_instance.basic_inference_single_image(data).image @@ -110,11 +112,12 @@ def test_use_lcm_turbomix_lora_lcm( pil_image, ) + @pytest.mark.refined_sdxl_model def test_use_lcm_turbomix_lora_dpmpp_sde( self, shared_model_manager: type[SharedModelManager], hordelib_instance: HordeLib, - stable_diffusion_model_name_for_testing: str, + sdxl_refined_model_name: str, ): assert shared_model_manager.manager.lora @@ -143,7 +146,7 @@ def test_use_lcm_turbomix_lora_dpmpp_sde( "loras": [{"name": lora_name, "model": 1.0, "clip": 1.0, "inject_trigger": "any", "is_version": True}], "ddim_steps": 10, "n_iter": 1, - "model": "AlbedoBase XL (SDXL)", + "model": sdxl_refined_model_name, } pil_image = hordelib_instance.basic_inference_single_image(data).image @@ -157,6 +160,7 @@ def test_use_lcm_turbomix_lora_dpmpp_sde( pil_image, ) + @pytest.mark.default_sd15_model def test_use_sd1_5_lora_lcm( self, shared_model_manager: type[SharedModelManager], diff --git a/tests/test_horde_lora.py b/tests/test_horde_lora.py index f60817a3..3e0bd46d 100644 --- a/tests/test_horde_lora.py +++ b/tests/test_horde_lora.py @@ -18,6 +18,7 @@ def setup_and_teardown(self, shared_model_manager: type[SharedModelManager]): shared_model_manager.manager.lora.wait_for_downloads() yield + @pytest.mark.default_sd15_model def test_text_to_image_lora_red( self, shared_model_manager: type[SharedModelManager], @@ -69,6 +70,7 @@ def test_text_to_image_lora_red( if not (last_use > datetime.now() - timedelta(minutes=1)): raise Exception("Last use of lora was not updated") + @pytest.mark.default_sd15_model def test_text_to_image_lora_blue( self, shared_model_manager: type[SharedModelManager], @@ -113,6 +115,7 @@ def test_text_to_image_lora_blue( pil_image, ) + @pytest.mark.default_sd15_model def test_text_to_image_lora_blue_tiled( self, shared_model_manager: type[SharedModelManager], @@ -157,6 +160,7 @@ def test_text_to_image_lora_blue_tiled( pil_image, ) + @pytest.mark.default_sd15_model def test_text_to_image_lora_blue_weighted( self, shared_model_manager: type[SharedModelManager], @@ -201,6 +205,7 @@ def test_text_to_image_lora_blue_weighted( pil_image, ) + @pytest.mark.default_sd15_model def test_text_to_image_lora_blue_low_strength( self, shared_model_manager: type[SharedModelManager], @@ -268,6 +273,7 @@ def test_text_to_image_lora_blue_low_strength( pil_image, ) + @pytest.mark.default_sd15_model def test_text_to_image_lora_blue_negative_strength( self, shared_model_manager: type[SharedModelManager], @@ -335,6 +341,7 @@ def test_text_to_image_lora_blue_negative_strength( pil_image, ) + @pytest.mark.default_sd15_model def test_text_to_image_lora_blue_hires_fix( self, shared_model_manager: type[SharedModelManager], @@ -443,6 +450,60 @@ def test_text_to_image_lora_character_hires_fix( pil_image, ) + def test_text_to_image_lora_character_hires_fix_sdxl( + self, + shared_model_manager: type[SharedModelManager], + hordelib_instance: HordeLib, + ): + assert shared_model_manager.manager.lora + + shared_model_manager.manager.lora.fetch_adhoc_lora("247778", is_version=True) + lora_name_1 = shared_model_manager.manager.lora.get_lora_name("247778", is_version=True) + shared_model_manager.manager.lora.fetch_adhoc_lora("135867", is_version=True) + lora_name_2 = shared_model_manager.manager.lora.get_lora_name("135867", is_version=True) + + assert lora_name_1 + assert lora_name_2 + + assert shared_model_manager.manager.compvis + + data = { + "sampler_name": "k_dpmpp_sde", + "cfg_scale": 5, + "denoising_strength": 1.0, + "seed": 4061434610, + "height": 1664, + "width": 1152, + "karras": True, + "tiling": False, + "hires_fix": True, + "clip_skip": 2, + "control_type": None, + "image_is_control": False, + "return_control_map": False, + "prompt": "(\nVtuber, face closeup, green eyes, heart-shaped pupils, forward horns, white hair, long hair, fox ears, fix whiskers stickers, comical blush, earrings, wink, eyeshadow, makeup, long eyelashes, v sign,\n\n(Red background, abstract background, reflective surface:1.15),\n\n\nscore_9,\n\n:1.0), ### (\n\n(futanari, shemale, dickgirl, futa, trap, yaoi, black and white, b&w, monochrome, 2boys, multiple boys:1.2),\n\n(score_6, score_5, score_4, score_3, score_2, score_1, source_furry, source_pony, source_cartoon, source_anime, source_filmmaker:1.0),\n\n(tongue, licking, chubby, dehydrated, ribs, ribcage, teeth, fish eyes, dead eyes:1.0),\n\nugly, worst quality, low quality, normal quality, messy drawing, amateur drawing, lowres, low resolution, poor resolution, normal resolution, bad anatomy, bad hands, text, watermark, logo, people, plastic, figurine, semi-realistic, painting, surrealist, digital art, cgi, render, sketch, manga, visual novel, drawing, uncanny, 3D, daz3d, anime, cartoon, animation, comic, video game,\n\n:1.0),", # noqa + "loras": [ + {"name": lora_name_1, "model": 1, "clip": 1.0}, + {"name": lora_name_2, "model": 5, "clip": 1.0}, + ], + "ddim_steps": 12, + "n_iter": 1, + "model": "AMPonyXL", + } + + pil_image = hordelib_instance.basic_inference_single_image(data).image + assert pil_image is not None + assert isinstance(pil_image, Image.Image) + + img_filename = "lora_character_hires_fix_sdxl.png" + pil_image.save(f"images/{img_filename}", quality=100) + + # assert check_single_lora_image_similarity( + # f"images_expected/{img_filename}", + # pil_image, + # ) + + @pytest.mark.default_sd15_model def test_text_to_image_lora_chained( self, shared_model_manager: type[SharedModelManager], @@ -509,6 +570,7 @@ def test_text_to_image_lora_chained( pil_image, ) + @pytest.mark.default_sd15_model def test_text_to_image_lora_chained_bad( self, shared_model_manager: type[SharedModelManager], @@ -552,6 +614,7 @@ def test_text_to_image_lora_chained_bad( assert len(ret.faults) == 1 # Don't save this one, just testing we didn't crash and burn + @pytest.mark.default_sd15_model def test_lora_trigger_inject_red( self, shared_model_manager: type[SharedModelManager], @@ -595,6 +658,7 @@ def test_lora_trigger_inject_red( pil_image, ) + @pytest.mark.default_sd15_model def test_lora_trigger_inject_any( self, shared_model_manager: type[SharedModelManager], @@ -640,6 +704,7 @@ def test_lora_trigger_inject_any( img_filename = "lora_inject_any_2.png" pil_image_2.save(f"images/{img_filename}", quality=100) + @pytest.mark.default_sd15_model def test_download_and_use_adhoc_lora( self, shared_model_manager: type[SharedModelManager], @@ -685,6 +750,7 @@ def test_download_and_use_adhoc_lora( pil_image, ) + @pytest.mark.default_sd15_model def test_download_and_use_specific_version_lora( self, shared_model_manager: type[SharedModelManager], @@ -727,6 +793,7 @@ def test_download_and_use_specific_version_lora( pil_image, ) + @pytest.mark.default_sd15_model def test_for_probability_tensor_runtime_error( self, hordelib_instance: HordeLib, @@ -763,6 +830,7 @@ def test_for_probability_tensor_runtime_error( pil_image = hordelib_instance.basic_inference_single_image(data).image assert pil_image is not None + @pytest.mark.default_sd15_model def test_sd21_lora_against_sd15_model( self, hordelib_instance: HordeLib, @@ -806,6 +874,7 @@ def test_sd21_lora_against_sd15_model( pil_image, ) + @pytest.mark.default_sd15_model def test_stonepunk( self, hordelib_instance: HordeLib, @@ -851,6 +920,7 @@ def test_stonepunk( pil_image, ) + @pytest.mark.default_sd15_model def test_negative_model_power( self, hordelib_instance: HordeLib, diff --git a/tests/test_horde_samplers.py b/tests/test_horde_samplers.py index 0d56abf8..dd01e0c0 100644 --- a/tests/test_horde_samplers.py +++ b/tests/test_horde_samplers.py @@ -1,5 +1,6 @@ # test_horde.py +import pytest from loguru import logger from PIL import Image @@ -11,6 +12,8 @@ class TestHordeSamplers: + + @pytest.mark.default_sd15_model def test_ddim_sampler( self, stable_diffusion_model_name_for_testing: str, @@ -42,6 +45,7 @@ def test_ddim_sampler( pil_image = hordelib_instance.basic_inference_single_image(data).image assert pil_image is not None + @pytest.mark.default_sd15_model def test_k_dpmpp_sde_sampler( self, stable_diffusion_model_name_for_testing: str, @@ -81,6 +85,8 @@ def test_k_dpmpp_sde_sampler( pil_image, ) + @pytest.mark.default_sd15_model + @pytest.mark.slow def test_samplers( self, stable_diffusion_model_name_for_testing: str, @@ -132,6 +138,7 @@ def test_samplers( f"Skipping image similarity check for {img_filename} due to SDE samplers being non-deterministic.", ) + @pytest.mark.default_sd15_model def test_slow_samplers( self, stable_diffusion_model_name_for_testing: str, diff --git a/tests/test_horde_ti.py b/tests/test_horde_ti.py index 51d340a0..9d86d152 100644 --- a/tests/test_horde_ti.py +++ b/tests/test_horde_ti.py @@ -43,6 +43,7 @@ def basic_ti_payload_data( "model": stable_diffusion_model_name_for_testing, } + @pytest.mark.default_sd15_model def test_basic_ti( self, shared_model_manager: type[SharedModelManager], @@ -60,6 +61,7 @@ def test_basic_ti( img_filename = "ti_basic.png" pil_image.save(f"images/{img_filename}", quality=100) + @pytest.mark.default_sd15_model def test_inject_ti( self, hordelib_instance: HordeLib, @@ -117,6 +119,7 @@ def test_inject_ti( pil_image, ) + @pytest.mark.default_sd15_model def test_bad_inject_ti( self, hordelib_instance: HordeLib, diff --git a/tests/test_image_metadata.py b/tests/test_image_metadata.py index 58700079..9a2a18f9 100644 --- a/tests/test_image_metadata.py +++ b/tests/test_image_metadata.py @@ -1,12 +1,15 @@ # test_horde.py import json +import pytest from PIL import Image from hordelib.horde import HordeLib class TestHordeInferenceMetadata: + + @pytest.mark.default_sd15_model def test_text_to_image( self, hordelib_instance: HordeLib, diff --git a/tests/test_internal_comfyui_failures.py b/tests/test_internal_comfyui_failures.py index 77c899de..09cb8fe6 100644 --- a/tests/test_internal_comfyui_failures.py +++ b/tests/test_internal_comfyui_failures.py @@ -4,6 +4,7 @@ from hordelib.shared_model_manager import SharedModelManager +@pytest.mark.default_sd15_model def test_lora_failure( shared_model_manager: type[SharedModelManager], hordelib_instance: HordeLib, diff --git a/tests/test_utils.py b/tests/test_utils.py index 1bb0cabe..594dbc80 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -1,4 +1,5 @@ import pytest +from loguru import logger from hordelib.settings import UserSettings from hordelib.utils.distance import ( @@ -8,6 +9,8 @@ from hordelib.utils.gpuinfo import GPUInfo from hordelib.utils.image_utils import ImageUtils +from .testing_shared_classes import ResolutionTestCase + def test_worker_settings_singleton(): a = UserSettings._instance @@ -147,3 +150,59 @@ def test_both_dimensions_oversized_unevenly_max(self): assert calculated_cascade != calculated_default assert calculated_cascade == expected + + def test_get_first_pass_image_resolution_sdxl(self): + expected = (1024, 1024) + calculated = ImageUtils.get_first_pass_image_resolution_sdxl(1024, 1024) + + assert calculated == expected + + def test_known_buckets_first_pass_sdxl(self): + from hordelib.utils.image_utils import IDEAL_SDXL_RESOLUTIONS + + for resolution in IDEAL_SDXL_RESOLUTIONS: + calculated = ImageUtils.get_first_pass_image_resolution_sdxl(resolution[0], resolution[1]) + assert calculated == resolution + + def test_not_known_buckets_first_pass_sdxl(self): + calculated = ImageUtils.get_first_pass_image_resolution_sdxl(1024, 1024 + 64) + assert calculated == (1024, 1024) + + def test_calc_upscale_sampler_steps(self, all_hires_test_cases: list[ResolutionTestCase]): + for hires_test_case in all_hires_test_cases: + logger.debug( + f"Testing upscale steps for hires denoise {hires_test_case.hires_fix_denoise_strength} " + f"{hires_test_case.width}x{hires_test_case.height} for {hires_test_case.ddim_steps} @ " + f"{hires_test_case.model_native_resolution} model resolution", + ) + result_1 = ImageUtils.calc_upscale_sampler_steps( + hires_test_case.model_native_resolution, + hires_test_case.width, + hires_test_case.height, + hires_test_case.hires_fix_denoise_strength, + hires_test_case.ddim_steps, + ) + + if hires_test_case.min_expected_steps is not None: + logger.debug(f"Expecting at least {hires_test_case.min_expected_steps} upscale steps") + assert result_1 >= hires_test_case.min_expected_steps, ( + f"Expected at least {hires_test_case.min_expected_steps} upscale steps for " + f"{hires_test_case.width}x{hires_test_case.height} resolution" + ) + + if hires_test_case.max_expected_steps is not None: + logger.debug(f"Expecting at most {hires_test_case.max_expected_steps} upscale steps") + assert result_1 <= hires_test_case.max_expected_steps, ( + f"Expected at most {hires_test_case.max_expected_steps} upscale steps for " + f"{hires_test_case.width}x{hires_test_case.height} resolution" + ) + + result_2 = ImageUtils.calc_upscale_sampler_steps( + hires_test_case.model_native_resolution, + hires_test_case.height, + hires_test_case.width, + hires_test_case.hires_fix_denoise_strength, + hires_test_case.ddim_steps, + ) + + assert result_1 == result_2, "Upscale steps should be the same with reversed width and height" diff --git a/tests/testing_shared_classes.py b/tests/testing_shared_classes.py new file mode 100644 index 00000000..fa7d555e --- /dev/null +++ b/tests/testing_shared_classes.py @@ -0,0 +1,20 @@ +class ResolutionTestCase: + + def __init__( + self, + *, + width: int | float, + height: int | float, + ddim_steps: int | float, + hires_fix_denoise_strength: float, + model_native_resolution: int, + max_expected_steps: int | float | None, + min_expected_steps: int | float | None, + ): + self.width = int(width) + self.height = int(height) + self.ddim_steps = int(ddim_steps) + self.hires_fix_denoise_strength = hires_fix_denoise_strength + self.model_native_resolution = int(model_native_resolution) + self.max_expected_steps = int(max_expected_steps) if max_expected_steps is not None else None + self.min_expected_steps = int(min_expected_steps) if min_expected_steps is not None else None