Spaces:
Running on Zero
Running on Zero
File size: 8,060 Bytes
cb617a4 5fff6c4 cb617a4 5fff6c4 cb617a4 5c49101 5fff6c4 cb617a4 5fff6c4 b5e5a50 5fff6c4 cb617a4 03cd0e3 cb617a4 9bdeeef 43b47b6 cb617a4 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 | """Pure runtime tests; no ComfyUI clone or model download is required."""
from __future__ import annotations
import os
import sys
import tempfile
import unittest
from pathlib import Path
from unittest.mock import patch
os.environ["KREA_SKIP_STARTUP"] = "1"
sys.path.insert(0, os.path.dirname(__file__))
import app # noqa: E402
class WorkflowTests(unittest.TestCase):
def test_text_to_image_loader_and_prompt_injection(self):
workflow = app._t2i_workflow()
app._inject_t2i(
workflow,
prompt="a blue glass sculpture",
width=1024,
height=768,
steps=8,
cfg=1.0,
sampler="euler",
scheduler="beta",
seed=42,
)
self.assertEqual(workflow["1"]["inputs"]["unet_name"], app.CUSTOM_KREA_FILE)
self.assertEqual(app.BASE_MODELS[app.CUSTOM_KREA_FILE]["repo"], "mpasila/Krea-2-Models")
self.assertIn(app.MUSE_KREA_FILE, app.BASE_MODELS)
self.assertEqual(workflow["2"]["inputs"]["clip_name"], "qwen3vl_4b_fp8_scaled.safetensors")
self.assertEqual(workflow["5"]["inputs"]["text"], "a blue glass sculpture")
self.assertEqual(workflow["7"]["inputs"]["height"], 768)
self.assertEqual(workflow["8"]["inputs"]["seed"], 42)
def test_edit_workflow_has_mandatory_adapter_and_optional_second_reference(self):
workflow = app._edit_workflow(True)
self.assertEqual(workflow["6"]["class_type"], "LoraLoaderModelOnly")
self.assertTrue(workflow["6"]["inputs"]["lora_name"].endswith(app.IDENTITY_FILE))
self.assertIn("source_latent_b", workflow["9"]["inputs"])
self.assertIn("image_b", workflow["10"]["inputs"])
def test_edit_workflow_without_second_reference_omits_b_inputs(self):
workflow = app._edit_workflow(False)
self.assertNotIn("2", workflow)
self.assertNotIn("source_latent_b", workflow["9"]["inputs"])
self.assertNotIn("source_image_b", workflow["9"]["inputs"])
self.assertNotIn("image_b", workflow["10"]["inputs"])
def test_text_to_image_supports_selected_model_and_ordered_loras(self):
workflow = app._t2i_workflow(app.MUSE_KREA_FILE)
app._inject_t2i(
workflow,
prompt="a portrait",
width=1024,
height=1024,
steps=8,
cfg=1.0,
sampler="euler",
scheduler="beta",
seed=7,
enabled_loras=[("krea2/first.safetensors", 0.5), ("huggingface/second.safetensors", -0.25)],
)
self.assertEqual(workflow["1"]["inputs"]["unet_name"], app.MUSE_KREA_FILE)
self.assertEqual(workflow["user_lora_0"]["inputs"]["model"], ["1", 0])
self.assertEqual(workflow["user_lora_1"]["inputs"]["model"], ["user_lora_0", 0])
self.assertEqual(workflow["4"]["inputs"]["model"], ["user_lora_1", 0])
self.assertEqual(workflow["5"]["inputs"]["clip"], ["user_lora_1", 1])
def test_custom_base_model_download_uses_managed_namespace(self):
original_models = app.MODELS
with tempfile.TemporaryDirectory() as directory:
app.MODELS = Path(directory) / "models"
def fake_download(**kwargs):
destination = Path(kwargs["local_dir"]) / kwargs["filename"]
destination.parent.mkdir(parents=True, exist_ok=True)
destination.write_bytes(b"test")
return str(destination)
try:
with patch.object(app, "hf_hub_download", side_effect=fake_download):
relative = app._ensure_custom_base_model({
"repo_id": "org/repo",
"filename": "sub/base.safetensors",
"revision": "main",
})
finally:
app.MODELS = original_models
self.assertTrue(relative.startswith("huggingface/"))
self.assertTrue(relative.endswith("/sub/base.safetensors"))
workflow = app._t2i_workflow(relative)
self.assertEqual(workflow["1"]["inputs"]["unet_name"], relative)
def test_edit_user_loras_follow_identity_adapter(self):
workflow = app._edit_workflow(False, app.MUSE_KREA_FILE)
app._inject_edit(
workflow,
primary_name="source.png",
second_name=None,
width=1024,
height=1024,
edit_prompt="change the shirt",
grounding_px=768,
ref_boost=1.0,
ref_boost_a=1.0,
steps=8,
cfg=1.0,
sampler="euler",
scheduler="beta",
seed=7,
enabled_loras=[("krea2/style.safetensors", 1.0)],
)
self.assertEqual(workflow["5"]["inputs"]["unet_name"], app.MUSE_KREA_FILE)
self.assertEqual(workflow["user_lora_0"]["inputs"]["model"], ["6", 0])
self.assertEqual(workflow["9"]["inputs"]["model"], ["user_lora_0", 0])
self.assertEqual(workflow["10"]["inputs"]["clip"], ["user_lora_0", 1])
self.assertEqual(workflow["6"]["class_type"], "LoraLoaderModelOnly")
def test_custom_lora_download_uses_managed_namespace(self):
original_root = app.LORA_ROOT
original_custom_root = app.CUSTOM_LORA_DEST_DIR
with tempfile.TemporaryDirectory() as directory:
app.LORA_ROOT = Path(directory) / "loras"
app.CUSTOM_LORA_DEST_DIR = app.LORA_ROOT / "huggingface"
def fake_download(**kwargs):
destination = Path(kwargs["local_dir"]) / kwargs["filename"]
destination.parent.mkdir(parents=True, exist_ok=True)
destination.write_bytes(b"test")
return str(destination)
try:
with patch.object(app, "hf_hub_download", side_effect=fake_download):
relative = app._ensure_custom_lora({
"repo_id": "org/repo",
"filename": "sub/style.safetensors",
"revision": "main",
"weight": 0.8,
})
finally:
app.LORA_ROOT = original_root
app.CUSTOM_LORA_DEST_DIR = original_custom_root
self.assertTrue(relative.startswith("huggingface/"))
self.assertTrue(relative.endswith("/sub/style.safetensors"))
def test_mode_validation(self):
with self.assertRaises(ValueError):
app._validate_request("text2image", "", "", None)
with self.assertRaises(ValueError):
app._validate_request("edit", "", "change the color", None)
with self.assertRaises(ValueError):
app._validate_request("edit", "", "", "source.png")
app._validate_request("edit", "", "change the color", "source.png")
app._validate_request("edit", "change the color", "", "source.png")
def test_gpu_duration_accepts_spaces_positional_callback(self):
values = [
"text2image", "prompt", "", None, None,
1024, 1024, 1.4, 768, 1.0, 1.0,
8, 1.0, "euler", "beta", 2, False, 0, None,
]
duration = app.get_gpu_duration(*values)
self.assertGreaterEqual(duration, app.MIN_GPU_SECONDS)
def test_gpu_duration_accepts_keyword_callback(self):
duration = app.get_gpu_duration(steps=8, width=1024, height=1024, gen_budget=120)
self.assertEqual(duration, 120)
def test_runtime_uses_ram_pressure_executor_configuration(self):
with open(app.__file__, encoding="utf-8") as source_file:
source = source_file.read()
self.assertIn("cache_type=execution.CacheType.RAM_PRESSURE", source)
self.assertIn('"ram": 2.0', source)
self.assertIn('"ram_inactive": 8.0', source)
if __name__ == "__main__":
unittest.main()
|