diff --git a/.gitignore b/.gitignore index 46594ca..6fc50bf 100644 --- a/.gitignore +++ b/.gitignore @@ -151,6 +151,7 @@ hordelib/model_database/stable_diffusion.json hordelib/model_database/lora.json ComfyUI model.ckpt +models/ coverage.lcov profiles/ longprompts.zip diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 1ccb773..0964a3d 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -34,6 +34,6 @@ repos: types-tabulate, types-tqdm, types-urllib3, - horde_sdk==0.14.0, - horde_model_reference==0.8.1, + horde_sdk==0.14.3, + horde_model_reference==0.9.0, ] diff --git a/hordelib/horde.py b/hordelib/horde.py index bc81463..ce158d7 100644 --- a/hordelib/horde.py +++ b/hordelib/horde.py @@ -284,6 +284,12 @@ class HordeLib: "upscale_sampler.sampler_name": "sampler_name", "controlnet_apply.strength": "control_strength", "controlnet_model_loader.control_net_name": "control_type", + # Flux + "cfg_guider.cfg": "cfg_scale", + "random_noise.noise_seed": "seed", + "k_sampler_select.sampler_name": "sampler_name", + "basic_scheduler.denoise": "denoising_strength", + "basic_scheduler.steps": "ddim_steps", # Stable Cascade "stable_cascade_empty_latent_image.width": "width", "stable_cascade_empty_latent_image.height": "height", @@ -856,10 +862,15 @@ def _final_pipeline_adjustments(self, payload, pipeline_data) -> tuple[dict, lis ) # The last LORA always connects to the sampler and clip text encoders (via the clip_skip) - if lora_index == len(payload.get("loras")) - 1: - self.generator.reconnect_input(pipeline_data, "sampler.model", f"lora_{lora_index}") - self.generator.reconnect_input(pipeline_data, "upscale_sampler.model", f"lora_{lora_index}") - self.generator.reconnect_input(pipeline_data, "clip_skip.clip", f"lora_{lora_index}") + if lora_index == len(payload.get("loras")) - 1 and SharedModelManager.manager.compvis: + model_details = SharedModelManager.manager.compvis.get_model_reference_info(payload["model_name"]) + if model_details is not None and model_details["baseline"] == "flux_1": + self.generator.reconnect_input(pipeline_data, "cfg_guider.model", f"lora_{lora_index}") + self.generator.reconnect_input(pipeline_data, "basic_scheduler.model", f"lora_{lora_index}") + else: + self.generator.reconnect_input(pipeline_data, "sampler.model", f"lora_{lora_index}") + self.generator.reconnect_input(pipeline_data, "upscale_sampler.model", f"lora_{lora_index}") + self.generator.reconnect_input(pipeline_data, "clip_skip.clip", f"lora_{lora_index}") # Translate the payload parameters into pipeline parameters pipeline_params = {} @@ -885,7 +896,7 @@ def _final_pipeline_adjustments(self, payload, pipeline_data) -> tuple[dict, lis # 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 + # and don't want to have to rememeber to re-add those keys if "model_loader_stage_c.ckpt_name" in pipeline_params: pipeline_params["model_loader_stage_c.file_type"] = "stable_cascade_stage_c" if "model_loader_stage_b.ckpt_name" in pipeline_params: @@ -990,7 +1001,12 @@ def _final_pipeline_adjustments(self, payload, pipeline_data) -> tuple[dict, lis # We do this by reconnecting the nodes in the pipeline to make the input to the vae encoder # the source image instead of the latent noise generator if pipeline_params.get("image_loader.image"): - self.generator.reconnect_input(pipeline_data, "sampler.latent_image", "vae_encode") + if SharedModelManager.manager.compvis: + model_details = SharedModelManager.manager.compvis.get_model_reference_info(payload["model_name"]) + if isinstance(model_details, dict) and model_details.get("baseline") == "flux_1": + self.generator.reconnect_input(pipeline_data, "sampler_custom_advanced.latent_image", "vae_encode") + else: + self.generator.reconnect_input(pipeline_data, "sampler.latent_image", "vae_encode") if pipeline_params.get("sc_image_loader.image"): self.generator.reconnect_input( pipeline_data, @@ -1181,6 +1197,8 @@ def _get_appropriate_pipeline(self, params): if params.get("hires_fix", False): return "stable_cascade_2pass" return "stable_cascade" + if model_details.get("baseline") == "flux_1": + return "flux" if params.get("control_type"): if params.get("return_control_map", False): return "controlnet_annotator" diff --git a/hordelib/model_manager/lora.py b/hordelib/model_manager/lora.py index 04b3b9b..b3fdcc4 100644 --- a/hordelib/model_manager/lora.py +++ b/hordelib/model_manager/lora.py @@ -426,13 +426,13 @@ def _parse_civitai_lora_data(self, item, adhoc=False): logger.debug(f"Rejecting LoRa {lora.get('name')} because it doesn't have a url") return None # We don't want to start downloading GBs of a single LoRa. - # We just ignore anything over 150Mb. Them's the breaks... + # We just ignore anything over 400Mb. Them's the breaks... if ( lora["versions"][lora_version]["adhoc"] - and lora["versions"][lora_version]["size_mb"] > 220 + and lora["versions"][lora_version]["size_mb"] > 400 and lora["id"] not in self._default_lora_ids ): - logger.debug(f"Rejecting LoRa {lora.get('name')} version {lora_version} because its size is over 220Mb.") + logger.debug(f"Rejecting LoRa {lora.get('name')} version {lora_version} because its size is over 400Mb.") return None if lora["versions"][lora_version]["adhoc"] and lora["nsfw"] and not self.nsfw: logger.debug(f"Rejecting LoRa {lora.get('name')} because worker is SFW.") diff --git a/hordelib/pipeline_designs/pipeline_flux.json b/hordelib/pipeline_designs/pipeline_flux.json new file mode 100644 index 0000000..19b1a0a --- /dev/null +++ b/hordelib/pipeline_designs/pipeline_flux.json @@ -0,0 +1,733 @@ +{ + "last_node_id": 22, + "last_link_id": 51, + "nodes": [ + { + "id": 22, + "type": "CFGGuider", + "pos": { + "0": 659, + "1": 160 + }, + "size": { + "0": 315, + "1": 98 + }, + "flags": {}, + "order": 10, + "mode": 0, + "inputs": [ + { + "name": "model", + "type": "MODEL", + "link": 44 + }, + { + "name": "positive", + "type": "CONDITIONING", + "link": 46 + }, + { + "name": "negative", + "type": "CONDITIONING", + "link": 48 + } + ], + "outputs": [ + { + "name": "GUIDER", + "type": "GUIDER", + "links": [ + 39 + ], + "shape": 3, + "slot_index": 0 + } + ], + "title": "cfg_guider", + "properties": { + "Node name for S&R": "CFGGuider" + }, + "widgets_values": [ + 1 + ] + }, + { + "id": 6, + "type": "CLIPTextEncode", + "pos": { + "0": 389, + "1": 75 + }, + "size": { + "0": 236.33343505859375, + "1": 83.00006103515625 + }, + "flags": {}, + "order": 7, + "mode": 0, + "inputs": [ + { + "name": "clip", + "type": "CLIP", + "link": 37 + } + ], + "outputs": [ + { + "name": "CONDITIONING", + "type": "CONDITIONING", + "links": [ + 46 + ], + "slot_index": 0 + } + ], + "title": "prompt", + "properties": { + "Node name for S&R": "CLIPTextEncode" + }, + "widgets_values": [ + "a steampunk text that says \\\"Horde Engine\\\" floating" + ] + }, + { + "id": 7, + "type": "CLIPTextEncode", + "pos": { + "0": 378, + "1": 194 + }, + "size": { + "0": 243.0001220703125, + "1": 90.3333740234375 + }, + "flags": {}, + "order": 8, + "mode": 0, + "inputs": [ + { + "name": "clip", + "type": "CLIP", + "link": 49, + "slot_index": 0 + } + ], + "outputs": [ + { + "name": "CONDITIONING", + "type": "CONDITIONING", + "links": [ + 48 + ], + "slot_index": 0 + } + ], + "title": "negative_prompt", + "properties": { + "Node name for S&R": "CLIPTextEncode" + }, + "widgets_values": [ + "" + ] + }, + { + "id": 20, + "type": "SamplerCustomAdvanced", + "pos": { + "0": 1026, + "1": 288 + }, + "size": { + "0": 355.20001220703125, + "1": 106 + }, + "flags": {}, + "order": 11, + "mode": 0, + "inputs": [ + { + "name": "noise", + "type": "NOISE", + "link": 38 + }, + { + "name": "guider", + "type": "GUIDER", + "link": 39 + }, + { + "name": "sampler", + "type": "SAMPLER", + "link": 40 + }, + { + "name": "sigmas", + "type": "SIGMAS", + "link": 41 + }, + { + "name": "latent_image", + "type": "LATENT", + "link": 51 + } + ], + "outputs": [ + { + "name": "output", + "type": "LATENT", + "links": [ + 43 + ], + "shape": 3, + "slot_index": 0 + }, + { + "name": "denoised_output", + "type": "LATENT", + "links": null, + "shape": 3 + } + ], + "title": "sampler_custom_advanced", + "properties": { + "Node name for S&R": "SamplerCustomAdvanced" + } + }, + { + "id": 14, + "type": "VAEDecode", + "pos": { + "0": 1117, + "1": 473 + }, + "size": { + "0": 210, + "1": 46 + }, + "flags": {}, + "order": 12, + "mode": 0, + "inputs": [ + { + "name": "samples", + "type": "LATENT", + "link": 43 + }, + { + "name": "vae", + "type": "VAE", + "link": 19 + } + ], + "outputs": [ + { + "name": "IMAGE", + "type": "IMAGE", + "links": [ + 36 + ], + "slot_index": 0 + } + ], + "title": "vae_decode", + "properties": { + "Node name for S&R": "VAEDecode" + } + }, + { + "id": 9, + "type": "SaveImage", + "pos": { + "0": 1410, + "1": 166 + }, + "size": [ + 395.5999999999999, + 443.89999999999975 + ], + "flags": {}, + "order": 13, + "mode": 0, + "inputs": [ + { + "name": "images", + "type": "IMAGE", + "link": 36 + } + ], + "outputs": [], + "title": "output_image", + "properties": {}, + "widgets_values": [ + "ComfyUI" + ] + }, + { + "id": 11, + "type": "LoadImage", + "pos": { + "0": -14, + "1": 530 + }, + "size": { + "0": 315, + "1": 314 + }, + "flags": {}, + "order": 0, + "mode": 0, + "inputs": [], + "outputs": [ + { + "name": "IMAGE", + "type": "IMAGE", + "links": [ + 22 + ], + "slot_index": 0 + }, + { + "name": "MASK", + "type": "MASK", + "links": null + } + ], + "title": "image_loader", + "properties": { + "Node name for S&R": "LoadImage" + }, + "widgets_values": [ + "example.png", + "image" + ] + }, + { + "id": 15, + "type": "RepeatImageBatch", + "pos": { + "0": 320, + "1": 660 + }, + "size": { + "0": 315, + "1": 58 + }, + "flags": {}, + "order": 5, + "mode": 0, + "inputs": [ + { + "name": "image", + "type": "IMAGE", + "link": 22 + } + ], + "outputs": [ + { + "name": "IMAGE", + "type": "IMAGE", + "links": [ + 23 + ], + "slot_index": 0, + "shape": 3 + } + ], + "title": "repeat_image_batch", + "properties": { + "Node name for S&R": "RepeatImageBatch" + }, + "widgets_values": [ + 1 + ] + }, + { + "id": 12, + "type": "VAEEncode", + "pos": { + "0": 714, + "1": 883 + }, + "size": { + "0": 210, + "1": 46 + }, + "flags": {}, + "order": 9, + "mode": 0, + "inputs": [ + { + "name": "pixels", + "type": "IMAGE", + "link": 23 + }, + { + "name": "vae", + "type": "VAE", + "link": 14 + } + ], + "outputs": [ + { + "name": "LATENT", + "type": "LATENT", + "links": [], + "slot_index": 0 + } + ], + "title": "vae_encode", + "properties": { + "Node name for S&R": "VAEEncode" + } + }, + { + "id": 5, + "type": "EmptyLatentImage", + "pos": { + "0": 654, + "1": 704 + }, + "size": { + "0": 315, + "1": 106 + }, + "flags": {}, + "order": 1, + "mode": 0, + "inputs": [], + "outputs": [ + { + "name": "LATENT", + "type": "LATENT", + "links": [ + 51 + ], + "slot_index": 0 + } + ], + "title": "empty_latent_image", + "properties": { + "Node name for S&R": "EmptyLatentImage" + }, + "widgets_values": [ + 1024, + 1024, + 1 + ] + }, + { + "id": 18, + "type": "BasicScheduler", + "pos": { + "0": 653, + "1": 550 + }, + "size": { + "0": 315, + "1": 106 + }, + "flags": {}, + "order": 6, + "mode": 0, + "inputs": [ + { + "name": "model", + "type": "MODEL", + "link": 45 + } + ], + "outputs": [ + { + "name": "SIGMAS", + "type": "SIGMAS", + "links": [ + 41 + ], + "shape": 3, + "slot_index": 0 + } + ], + "title": "basic_scheduler", + "properties": { + "Node name for S&R": "BasicScheduler" + }, + "widgets_values": [ + "normal", + 4, + 1 + ] + }, + { + "id": 4, + "type": "CheckpointLoaderSimple", + "pos": { + "0": 23, + "1": 254 + }, + "size": { + "0": 315, + "1": 98 + }, + "flags": {}, + "order": 2, + "mode": 0, + "inputs": [], + "outputs": [ + { + "name": "MODEL", + "type": "MODEL", + "links": [ + 44, + 45 + ], + "slot_index": 0 + }, + { + "name": "CLIP", + "type": "CLIP", + "links": [ + 37, + 49 + ], + "slot_index": 1 + }, + { + "name": "VAE", + "type": "VAE", + "links": [ + 14, + 19 + ], + "slot_index": 2 + } + ], + "title": "model_loader", + "properties": { + "Node name for S&R": "CheckpointLoaderSimple" + }, + "widgets_values": [ + "flux1CompactCLIPAnd_Flux1SchnellFp8.safetensors" + ] + }, + { + "id": 21, + "type": "RandomNoise", + "pos": { + "0": 659, + "1": 313 + }, + "size": { + "0": 315, + "1": 82 + }, + "flags": {}, + "order": 3, + "mode": 0, + "inputs": [], + "outputs": [ + { + "name": "NOISE", + "type": "NOISE", + "links": [ + 38 + ], + "shape": 3, + "slot_index": 0 + } + ], + "title": "random_noise", + "properties": { + "Node name for S&R": "RandomNoise" + }, + "widgets_values": [ + 1312, + "fixed" + ] + }, + { + "id": 19, + "type": "KSamplerSelect", + "pos": { + "0": 651, + "1": 446 + }, + "size": { + "0": 315, + "1": 58 + }, + "flags": {}, + "order": 4, + "mode": 0, + "inputs": [], + "outputs": [ + { + "name": "SAMPLER", + "type": "SAMPLER", + "links": [ + 40 + ], + "shape": 3, + "slot_index": 0 + } + ], + "title": "k_sampler_select", + "properties": { + "Node name for S&R": "KSamplerSelect" + }, + "widgets_values": [ + "euler" + ] + } + ], + "links": [ + [ + 14, + 4, + 2, + 12, + 1, + "VAE" + ], + [ + 19, + 4, + 2, + 14, + 1, + "VAE" + ], + [ + 22, + 11, + 0, + 15, + 0, + "IMAGE" + ], + [ + 23, + 15, + 0, + 12, + 0, + "IMAGE" + ], + [ + 36, + 14, + 0, + 9, + 0, + "IMAGE" + ], + [ + 37, + 4, + 1, + 6, + 0, + "CLIP" + ], + [ + 38, + 21, + 0, + 20, + 0, + "NOISE" + ], + [ + 39, + 22, + 0, + 20, + 1, + "GUIDER" + ], + [ + 40, + 19, + 0, + 20, + 2, + "SAMPLER" + ], + [ + 41, + 18, + 0, + 20, + 3, + "SIGMAS" + ], + [ + 43, + 20, + 0, + 14, + 0, + "LATENT" + ], + [ + 44, + 4, + 0, + 22, + 0, + "MODEL" + ], + [ + 45, + 4, + 0, + 18, + 0, + "MODEL" + ], + [ + 46, + 6, + 0, + 22, + 1, + "CONDITIONING" + ], + [ + 48, + 7, + 0, + 22, + 2, + "CONDITIONING" + ], + [ + 49, + 4, + 1, + 7, + 0, + "CLIP" + ], + [ + 51, + 5, + 0, + 20, + 4, + "LATENT" + ] + ], + "groups": [], + "config": {}, + "extra": { + "ds": { + "scale": 1, + "offset": [ + -25.156096363635925, + 27.70414454545454 + ] + } + }, + "version": 0.4 +} diff --git a/hordelib/pipelines/pipeline_flux.json b/hordelib/pipelines/pipeline_flux.json new file mode 100644 index 0000000..f43bab3 --- /dev/null +++ b/hordelib/pipelines/pipeline_flux.json @@ -0,0 +1,198 @@ +{ + "4": { + "inputs": { + "ckpt_name": "flux1CompactCLIPAnd_Flux1SchnellFp8.safetensors" + }, + "class_type": "CheckpointLoaderSimple", + "_meta": { + "title": "model_loader" + } + }, + "5": { + "inputs": { + "width": 1024, + "height": 1024, + "batch_size": 1 + }, + "class_type": "EmptyLatentImage", + "_meta": { + "title": "empty_latent_image" + } + }, + "6": { + "inputs": { + "text": "a steampunk text that says \\\"Horde Engine\\\" floating", + "clip": [ + "4", + 1 + ] + }, + "class_type": "CLIPTextEncode", + "_meta": { + "title": "prompt" + } + }, + "7": { + "inputs": { + "text": "", + "clip": [ + "4", + 1 + ] + }, + "class_type": "CLIPTextEncode", + "_meta": { + "title": "negative_prompt" + } + }, + "9": { + "inputs": { + "filename_prefix": "ComfyUI", + "images": [ + "14", + 0 + ] + }, + "class_type": "SaveImage", + "_meta": { + "title": "output_image" + } + }, + "11": { + "inputs": { + "image": "example.png", + "upload": "image" + }, + "class_type": "LoadImage", + "_meta": { + "title": "image_loader" + } + }, + "12": { + "inputs": { + "pixels": [ + "15", + 0 + ], + "vae": [ + "4", + 2 + ] + }, + "class_type": "VAEEncode", + "_meta": { + "title": "vae_encode" + } + }, + "14": { + "inputs": { + "samples": [ + "20", + 0 + ], + "vae": [ + "4", + 2 + ] + }, + "class_type": "VAEDecode", + "_meta": { + "title": "vae_decode" + } + }, + "15": { + "inputs": { + "amount": 1, + "image": [ + "11", + 0 + ] + }, + "class_type": "RepeatImageBatch", + "_meta": { + "title": "repeat_image_batch" + } + }, + "18": { + "inputs": { + "scheduler": "normal", + "steps": 4, + "denoise": 1, + "model": [ + "4", + 0 + ] + }, + "class_type": "BasicScheduler", + "_meta": { + "title": "basic_scheduler" + } + }, + "19": { + "inputs": { + "sampler_name": "euler" + }, + "class_type": "KSamplerSelect", + "_meta": { + "title": "k_sampler_select" + } + }, + "20": { + "inputs": { + "noise": [ + "21", + 0 + ], + "guider": [ + "22", + 0 + ], + "sampler": [ + "19", + 0 + ], + "sigmas": [ + "18", + 0 + ], + "latent_image": [ + "5", + 0 + ] + }, + "class_type": "SamplerCustomAdvanced", + "_meta": { + "title": "sampler_custom_advanced" + } + }, + "21": { + "inputs": { + "noise_seed": 1312 + }, + "class_type": "RandomNoise", + "_meta": { + "title": "random_noise" + } + }, + "22": { + "inputs": { + "cfg": 1, + "model": [ + "4", + 0 + ], + "positive": [ + "6", + 0 + ], + "negative": [ + "7", + 0 + ] + }, + "class_type": "CFGGuider", + "_meta": { + "title": "cfg_guider" + } + } +} diff --git a/images_expected/flux_dev_fp8_image_to_image.png b/images_expected/flux_dev_fp8_image_to_image.png new file mode 100644 index 0000000..03b2d6d Binary files /dev/null and b/images_expected/flux_dev_fp8_image_to_image.png differ diff --git a/images_expected/flux_dev_fp8_text_to_image.png b/images_expected/flux_dev_fp8_text_to_image.png new file mode 100644 index 0000000..93cf9e8 Binary files /dev/null and b/images_expected/flux_dev_fp8_text_to_image.png differ diff --git a/images_expected/flux_schnell_fp8_image_to_image.png b/images_expected/flux_schnell_fp8_image_to_image.png new file mode 100644 index 0000000..84a6cd5 Binary files /dev/null and b/images_expected/flux_schnell_fp8_image_to_image.png differ diff --git a/images_expected/flux_schnell_fp8_text_to_image_lora.png b/images_expected/flux_schnell_fp8_text_to_image_lora.png new file mode 100644 index 0000000..0a16e91 Binary files /dev/null and b/images_expected/flux_schnell_fp8_text_to_image_lora.png differ diff --git a/images_expected/flux_schnell_fp8_text_to_image_n_iter_0.png b/images_expected/flux_schnell_fp8_text_to_image_n_iter_0.png new file mode 100644 index 0000000..ec7e93c Binary files /dev/null and b/images_expected/flux_schnell_fp8_text_to_image_n_iter_0.png differ diff --git a/images_expected/flux_schnell_fp8_text_to_image_n_iter_1.png b/images_expected/flux_schnell_fp8_text_to_image_n_iter_1.png new file mode 100644 index 0000000..d219ea8 Binary files /dev/null and b/images_expected/flux_schnell_fp8_text_to_image_n_iter_1.png differ diff --git a/requirements.txt b/requirements.txt index 5c6b137..53a4f54 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.14.0 -horde_model_reference>=0.8.1 +horde_sdk>=0.14.3 +horde_model_reference>=0.9.0 pydantic numpy==1.26.4 torch>=2.3.1 diff --git a/tests/conftest.py b/tests/conftest.py index 27c0d7d..f832c00 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -116,11 +116,66 @@ def isolated_comfy_horde_instance(init_horde) -> Comfy_Horde: return Comfy_Horde() +_testing_model_name = "Deliberate" +_sdxl_1_0_model_name = "SDXL 1.0" +_sdxl_refined_model_name = "AlbedoBase XL (SDXL)" +_stable_cascade_base_model_name = "Stable Cascade 1.0" +_flux1_schnell_fp8_base_model_name = "Flux.1-Schnell fp8 (Compact)" + +_all_model_names = [ + _testing_model_name, + _sdxl_1_0_model_name, + _sdxl_refined_model_name, + _stable_cascade_base_model_name, + _flux1_schnell_fp8_base_model_name, +] + +# !!!! +# If you're adding a model name, follow the pattern and **add it to `_all_model_names`** +# !!!! + + +@pytest.fixture(scope="session") +def stable_diffusion_model_name_for_testing(shared_model_manager: type[SharedModelManager]) -> str: + """The default stable diffusion 1.5 model name used for testing.""" + return _testing_model_name + + +@pytest.fixture(scope="session") +def sdxl_1_0_base_model_name(shared_model_manager: type[SharedModelManager]) -> str: + """The default SDXL 1.0 model name used for testing.""" + return _sdxl_1_0_model_name + + +@pytest.fixture(scope="session") +def sdxl_refined_model_name(shared_model_manager: type[SharedModelManager]) -> str: + """The default SDXL finetune model name used for testing.""" + return _sdxl_refined_model_name + + +@pytest.fixture(scope="session") +def stable_cascade_base_model_name(shared_model_manager: type[SharedModelManager]) -> str: + """The default stable cascade 1.0 model name used for testing.""" + return _stable_cascade_base_model_name + + +@pytest.fixture(scope="session") +def flux1_schnell_fp8_base_model_name(shared_model_manager: type[SharedModelManager]) -> str: + """The default flux1-schnell fp8 model name used for testing.""" + return _flux1_schnell_fp8_base_model_name + + +# !!!! +# If you're adding a model name, follow the pattern and **add it to `_all_model_names`** +# !!!! + + @pytest.fixture(scope="session") def shared_model_manager( custom_model_info_for_testing: tuple[str, str, str, str], hordelib_instance: HordeLib, ) -> Generator[type[SharedModelManager], None, None]: + assert hordelib_instance SharedModelManager(do_not_load_model_mangers=True) SharedModelManager.load_model_managers(ALL_MODEL_MANAGER_TYPES) @@ -132,17 +187,9 @@ def shared_model_manager( assert SharedModelManager.manager.miscellaneous.download_all_models() assert SharedModelManager.manager.compvis is not None - assert SharedModelManager.manager.download_model("Deliberate") - assert SharedModelManager.manager.validate_model("Deliberate") - assert SharedModelManager.manager.download_model("SDXL 1.0") - assert SharedModelManager.manager.validate_model("SDXL 1.0") - assert SharedModelManager.manager.download_model("AlbedoBase XL (SDXL)") - assert SharedModelManager.manager.validate_model("AlbedoBase XL (SDXL)") - assert SharedModelManager.manager.download_model("Rev Animated") - assert SharedModelManager.manager.validate_model("Rev Animated") - - assert SharedModelManager.manager.download_model("Stable Cascade 1.0") - assert SharedModelManager.manager.validate_model("Stable Cascade 1.0") + for model_name in _all_model_names: + assert SharedModelManager.manager.compvis.download_model(model_name) + assert SharedModelManager.manager.compvis.validate_model(model_name) custom_model_name, _, _, _ = custom_model_info_for_testing assert custom_model_name in SharedModelManager.manager.compvis.available_models @@ -165,29 +212,6 @@ def shared_model_manager( SharedModelManager.manager = None # type: ignore -_testing_model_name = "Deliberate" - - -@pytest.fixture(scope="session") -def stable_diffusion_model_name_for_testing(shared_model_manager: type[SharedModelManager]) -> str: - return _testing_model_name - - -@pytest.fixture(scope="session") -def sdxl_1_0_base_model_name(shared_model_manager: type[SharedModelManager]) -> str: - return "SDXL 1.0" - - -@pytest.fixture(scope="session") -def sdxl_refined_model_name(shared_model_manager: type[SharedModelManager]) -> str: - return "AlbedoBase XL (SDXL)" - - -@pytest.fixture(scope="session") -def stable_cascade_base_model_name(shared_model_manager: type[SharedModelManager]) -> str: - return "Stable Cascade 1.0" - - @pytest.fixture(scope="session") def custom_model_info_for_testing() -> tuple[str, str, str, str]: """Returns a tuple of the custom model name, its baseline, the on-disk file name and the download url.""" diff --git a/tests/test_horde_inference_flux.py b/tests/test_horde_inference_flux.py new file mode 100644 index 0000000..d7f2e6d --- /dev/null +++ b/tests/test_horde_inference_flux.py @@ -0,0 +1,238 @@ +# test_horde.py + +import pytest +from PIL import Image + +from hordelib.horde import HordeLib +from hordelib.shared_model_manager import SharedModelManager + +from .testing_shared_functions import check_list_inference_images_similarity, check_single_inference_image_similarity + + +class TestHordeInferenceFlux: + + # @pytest.mark.default_flux1_model + # def test_flux_dev_fp8_text_to_image( + # self, + # hordelib_instance: HordeLib, + # flux1_dev_fp8_base_model_name: str, + # ): + # data = { + # "sampler_name": "k_euler", + # "cfg_scale": 1, + # "denoising_strength": 1.0, + # "seed": 13122, + # "height": 1024, + # "width": 1024, + # "karras": False, + # "tiling": False, + # "hires_fix": False, + # "clip_skip": 1, + # "control_type": None, + # "image_is_control": False, + # "return_control_map": False, + # "prompt": 'a steampunk text that says "Horde Engine" floating', + # "ddim_steps": 20, + # "n_iter": 1, + # "model": flux1_dev_fp8_base_model_name, + # } + # pil_image = hordelib_instance.basic_inference_single_image(data).image + # assert pil_image is not None + # assert isinstance(pil_image, Image.Image) + + # img_filename = "flux_dev_fp8_text_to_image.png" + # pil_image.save(f"images/{img_filename}", quality=100) + + # assert check_single_inference_image_similarity( + # f"images_expected/{img_filename}", + # pil_image, + # ) + + @pytest.mark.default_flux1_model + def test_flux_schnell_fp8_text_to_image_n_iter( + self, + hordelib_instance: HordeLib, + flux1_schnell_fp8_base_model_name: str, + ): + data = { + "sampler_name": "k_euler", + "cfg_scale": 1, + "denoising_strength": 1.0, + "seed": 13122, + "height": 1024, + "width": 1024, + "karras": False, + "tiling": False, + "hires_fix": False, + "clip_skip": 1, + "control_type": None, + "image_is_control": False, + "return_control_map": False, + "prompt": 'a steampunk text that says "Horde Engine" floating', + "ddim_steps": 4, + "n_iter": 2, + "model": flux1_schnell_fp8_base_model_name, + } + image_results = hordelib_instance.basic_inference(data) + + assert len(image_results) == 2 + + img_pairs_to_check = [] + + img_filename_base = "flux_schnell_fp8_text_to_image_n_iter_{0}.png" + + for i, image_result in enumerate(image_results): + assert image_result.image is not None + assert isinstance(image_result.image, Image.Image) + + img_filename = img_filename_base.format(i) + + image_result.image.save(f"images/{img_filename}", quality=100) + img_pairs_to_check.append((f"images_expected/{img_filename}", image_result.image)) + + assert check_single_inference_image_similarity( + "images_expected/text_to_image.png", + "images/text_to_image_n_iter_0.png", + ) + + assert check_list_inference_images_similarity(img_pairs_to_check) + + @pytest.mark.default_flux1_model + def test_flux_schnell_fp8_image_to_image( + self, + hordelib_instance: HordeLib, + flux1_schnell_fp8_base_model_name: str, + ): + data = { + "sampler_name": "k_euler", + "cfg_scale": 1, + "denoising_strength": 0.8, + "seed": 13122, + "height": 1024, + "width": 1024, + "karras": False, + "tiling": False, + "hires_fix": False, + "clip_skip": 1, + "control_type": None, + "image_is_control": False, + "return_control_map": False, + "prompt": "a steampunk cowboy walking away from an explosion", + "ddim_steps": 4, + "n_iter": 1, + "source_image": Image.open("images/test_db0.jpg"), + "source_processing": "img2img", + "model": flux1_schnell_fp8_base_model_name, + } + pil_image = hordelib_instance.basic_inference_single_image(data).image + assert pil_image is not None + assert isinstance(pil_image, Image.Image) + + img_filename = "flux_schnell_fp8_image_to_image.png" + pil_image.save(f"images/{img_filename}", quality=100) + + assert check_single_inference_image_similarity( + f"images_expected/{img_filename}", + pil_image, + ) + + # @pytest.mark.default_flux1_model + # def test_flux_dev_fp8_image_to_image( + # self, + # hordelib_instance: HordeLib, + # flux1_dev_fp8_base_model_name: str, + # ): + # data = { + # "sampler_name": "k_euler", + # "cfg_scale": 2, + # "denoising_strength": 0.85, + # "seed": 13122, + # "height": 1024, + # "width": 1024, + # "karras": False, + # "tiling": False, + # "hires_fix": False, + # "clip_skip": 1, + # "control_type": None, + # "image_is_control": False, + # "return_control_map": False, + # "prompt": "a goth cowboy in black walking away from an explosion", + # "ddim_steps": 20, + # "n_iter": 1, + # "source_image": Image.open("images/test_db0.jpg"), + # "source_processing": "img2img", + # "model": flux1_dev_fp8_base_model_name, + # } + # pil_image = hordelib_instance.basic_inference_single_image(data).image + # assert pil_image is not None + # assert isinstance(pil_image, Image.Image) + + # img_filename = "flux_dev_fp8_image_to_image.png" + # pil_image.save(f"images/{img_filename}", quality=100) + + # assert check_single_inference_image_similarity( + # f"images_expected/{img_filename}", + # pil_image, + # ) + + @pytest.mark.default_flux1_model + def test_flux_dev_fp8_text_to_image_lora( + self, + shared_model_manager: type[SharedModelManager], + hordelib_instance: HordeLib, + flux1_schnell_fp8_base_model_name: str, + ): + + lora_schnell_version_id = "812384" + lora_dev_version_id = "735063" + + if shared_model_manager.manager.lora: + assert shared_model_manager.manager.lora.fetch_adhoc_lora( + lora_schnell_version_id, + timeout=45, + is_version=True, + ) + + assert shared_model_manager.manager.lora.fetch_adhoc_lora( + lora_dev_version_id, + timeout=45, + is_version=True, + ) + + data = { + "sampler_name": "k_euler", + "cfg_scale": 1, + "denoising_strength": 1.0, + "seed": 13122, + "height": 1024, + "width": 1024, + "karras": False, + "tiling": False, + "hires_fix": False, + "clip_skip": 1, + "control_type": None, + "image_is_control": False, + "return_control_map": False, + "prompt": ( + "A close-up of an magical and ominous obsidian monolith, with an glowing orange cat print inside. " + "detailmaximizer, dreamy art" + ), + "ddim_steps": 4, + "n_iter": 1, + "model": flux1_schnell_fp8_base_model_name, + "loras": [ + {"name": lora_schnell_version_id, "model": 1.0, "clip": 1.0, "is_version": True}, + {"name": lora_dev_version_id, "model": 0.5, "clip": 1.0, "is_version": True}, + ], + } + pil_image = hordelib_instance.basic_inference_single_image(data).image + assert pil_image is not None + assert isinstance(pil_image, Image.Image) + + img_filename = "flux_schnell_fp8_text_to_image_lora.png" + pil_image.save(f"images/{img_filename}", quality=100) + + assert check_single_inference_image_similarity( + f"images_expected/{img_filename}", + pil_image, + )