261 lines
8.6 KiB
Python
261 lines
8.6 KiB
Python
import asyncio
|
|
import importlib.util
|
|
import sys
|
|
import types
|
|
import unittest
|
|
from pathlib import Path
|
|
|
|
import torch
|
|
|
|
|
|
def _load_static_mode_module():
|
|
module_path = Path(__file__).resolve().parents[1] / "upscale" / "modes" / "static.py"
|
|
package_name = "dist_static_mode_testpkg"
|
|
|
|
for mod_name in list(sys.modules):
|
|
if mod_name == package_name or mod_name.startswith(f"{package_name}."):
|
|
del sys.modules[mod_name]
|
|
|
|
root_pkg = types.ModuleType(package_name)
|
|
root_pkg.__path__ = []
|
|
sys.modules[package_name] = root_pkg
|
|
|
|
upscale_pkg = types.ModuleType(f"{package_name}.upscale")
|
|
upscale_pkg.__path__ = []
|
|
sys.modules[f"{package_name}.upscale"] = upscale_pkg
|
|
|
|
modes_pkg = types.ModuleType(f"{package_name}.upscale.modes")
|
|
modes_pkg.__path__ = []
|
|
sys.modules[f"{package_name}.upscale.modes"] = modes_pkg
|
|
|
|
utils_pkg = types.ModuleType(f"{package_name}.utils")
|
|
utils_pkg.__path__ = []
|
|
sys.modules[f"{package_name}.utils"] = utils_pkg
|
|
|
|
created_comfy_stub = False
|
|
if "comfy" not in sys.modules:
|
|
created_comfy_stub = True
|
|
comfy_module = types.ModuleType("comfy")
|
|
model_mgmt = types.ModuleType("comfy.model_management")
|
|
|
|
class _InterruptProcessingException(Exception):
|
|
pass
|
|
|
|
model_mgmt.processing_interrupted = lambda: False
|
|
model_mgmt.throw_exception_if_processing_interrupted = lambda: None
|
|
model_mgmt.InterruptProcessingException = _InterruptProcessingException
|
|
|
|
comfy_module.model_management = model_mgmt
|
|
sys.modules["comfy"] = comfy_module
|
|
sys.modules["comfy.model_management"] = model_mgmt
|
|
|
|
logging_module = types.ModuleType(f"{package_name}.utils.logging")
|
|
logging_module.debug_log = lambda *_args, **_kwargs: None
|
|
logging_module.log = lambda *_args, **_kwargs: None
|
|
sys.modules[f"{package_name}.utils.logging"] = logging_module
|
|
|
|
image_module = types.ModuleType(f"{package_name}.utils.image")
|
|
from PIL import Image as PILImage
|
|
import numpy as np
|
|
|
|
def _tensor_to_pil(img_tensor, batch_index=0):
|
|
return PILImage.fromarray((255 * img_tensor[batch_index].cpu().numpy()).astype(np.uint8))
|
|
|
|
def _pil_to_tensor(image):
|
|
arr = np.array(image).astype(np.float32) / 255.0
|
|
return torch.from_numpy(arr).unsqueeze(0)
|
|
|
|
image_module.tensor_to_pil = _tensor_to_pil
|
|
image_module.pil_to_tensor = _pil_to_tensor
|
|
sys.modules[f"{package_name}.utils.image"] = image_module
|
|
|
|
async_helpers_module = types.ModuleType(f"{package_name}.utils.async_helpers")
|
|
|
|
def _run_async_in_server_loop(coro, timeout=None):
|
|
if timeout is not None:
|
|
return asyncio.run(asyncio.wait_for(coro, timeout=timeout))
|
|
return asyncio.run(coro)
|
|
|
|
async_helpers_module.run_async_in_server_loop = _run_async_in_server_loop
|
|
sys.modules[f"{package_name}.utils.async_helpers"] = async_helpers_module
|
|
|
|
config_module = types.ModuleType(f"{package_name}.utils.config")
|
|
config_module.get_worker_timeout_seconds = lambda: 60
|
|
sys.modules[f"{package_name}.utils.config"] = config_module
|
|
|
|
constants_module = types.ModuleType(f"{package_name}.utils.constants")
|
|
constants_module.HEARTBEAT_INTERVAL = 10.0
|
|
constants_module.JOB_POLL_INTERVAL = 0.0
|
|
constants_module.JOB_POLL_MAX_ATTEMPTS = 3
|
|
constants_module.MAX_BATCH = 20
|
|
constants_module.TILE_SEND_TIMEOUT = 1.0
|
|
constants_module.TILE_WAIT_TIMEOUT = 1.0
|
|
sys.modules[f"{package_name}.utils.constants"] = constants_module
|
|
|
|
job_store_module = types.ModuleType(f"{package_name}.upscale.job_store")
|
|
|
|
async def _noop(*_args, **_kwargs):
|
|
return None
|
|
|
|
job_store_module.ensure_tile_jobs_initialized = lambda: types.SimpleNamespace(
|
|
distributed_tile_jobs_lock=asyncio.Lock(),
|
|
distributed_pending_tile_jobs={},
|
|
)
|
|
job_store_module.init_static_job_batched = _noop
|
|
job_store_module._mark_task_completed = _noop
|
|
job_store_module._cleanup_job = _noop
|
|
job_store_module._drain_results_queue = _noop
|
|
job_store_module._get_completed_count = _noop
|
|
sys.modules[f"{package_name}.upscale.job_store"] = job_store_module
|
|
|
|
job_models_module = types.ModuleType(f"{package_name}.upscale.job_models")
|
|
|
|
class _TileJobState:
|
|
pass
|
|
|
|
job_models_module.TileJobState = _TileJobState
|
|
sys.modules[f"{package_name}.upscale.job_models"] = job_models_module
|
|
|
|
spec = importlib.util.spec_from_file_location(f"{package_name}.upscale.modes.static", module_path)
|
|
module = importlib.util.module_from_spec(spec)
|
|
assert spec is not None and spec.loader is not None
|
|
spec.loader.exec_module(module)
|
|
|
|
if created_comfy_stub:
|
|
sys.modules.pop("comfy.model_management", None)
|
|
sys.modules.pop("comfy", None)
|
|
|
|
return module
|
|
|
|
|
|
static_mode = _load_static_mode_module()
|
|
|
|
|
|
class _FakeStaticWorker(static_mode.StaticModeMixin):
|
|
def __init__(self):
|
|
self.sent_batches = []
|
|
self.request_calls = 0
|
|
self.heartbeat_calls = 0
|
|
self.job_ready = True
|
|
self.tile_sequence = [(0, 0, True), (None, 0, True)]
|
|
|
|
def round_to_multiple(self, value):
|
|
return value
|
|
|
|
def calculate_tiles(self, _width, _height, _tile_width, _tile_height, _force_uniform_tiles):
|
|
return [(0, 0)]
|
|
|
|
def _poll_job_ready(self, *_args, **_kwargs):
|
|
return self.job_ready
|
|
|
|
async def _request_tile_from_master(self, *_args, **_kwargs):
|
|
self.request_calls += 1
|
|
return self.tile_sequence.pop(0)
|
|
|
|
async def _send_heartbeat_to_master(self, *_args, **_kwargs):
|
|
self.heartbeat_calls += 1
|
|
|
|
async def send_tiles_batch_to_master(
|
|
self,
|
|
processed_tiles,
|
|
_multi_job_id,
|
|
_master_url,
|
|
_padding,
|
|
_worker_id,
|
|
is_final_flush=False,
|
|
):
|
|
self.sent_batches.append(
|
|
{
|
|
"tiles": list(processed_tiles),
|
|
"is_final_flush": bool(is_final_flush),
|
|
}
|
|
)
|
|
|
|
def _extract_and_process_tile(self, upscaled_image, *_args, **_kwargs):
|
|
batch_size = upscaled_image.shape[0]
|
|
processed_batch = torch.zeros((batch_size, 2, 2, 3), dtype=torch.float32)
|
|
return processed_batch, 0, 0, 2, 2
|
|
|
|
def create_tile_mask(self, *_args, **_kwargs):
|
|
from PIL import Image
|
|
return Image.new("L", (4, 4), 255)
|
|
|
|
def blend_tile(self, base_image, *_args, **_kwargs):
|
|
return base_image
|
|
|
|
|
|
def _call_worker_static(fake_worker):
|
|
image = torch.zeros((1, 4, 4, 3), dtype=torch.float32)
|
|
return fake_worker._process_worker_static_sync(
|
|
image,
|
|
model=None,
|
|
positive=None,
|
|
negative=None,
|
|
vae=None,
|
|
seed=1,
|
|
steps=1,
|
|
cfg=1.0,
|
|
sampler_name="euler",
|
|
scheduler="normal",
|
|
denoise=0.5,
|
|
tile_width=4,
|
|
tile_height=4,
|
|
padding=8,
|
|
mask_blur=4,
|
|
force_uniform_tiles=True,
|
|
tiled_decode=False,
|
|
multi_job_id="job-1",
|
|
master_url="http://master:8188",
|
|
worker_id="worker-1",
|
|
enabled_workers=["worker-1"],
|
|
)
|
|
|
|
|
|
class StaticModeWorkerFlowTests(unittest.TestCase):
|
|
def test_worker_static_aborts_when_job_not_ready(self):
|
|
worker = _FakeStaticWorker()
|
|
worker.job_ready = False
|
|
|
|
result = _call_worker_static(worker)
|
|
|
|
self.assertEqual(result[0].shape[0], 1)
|
|
self.assertEqual(worker.request_calls, 0)
|
|
self.assertEqual(worker.heartbeat_calls, 0)
|
|
self.assertEqual(worker.sent_batches, [])
|
|
|
|
def test_worker_static_requests_tiles_and_flushes_final_batch(self):
|
|
worker = _FakeStaticWorker()
|
|
|
|
_call_worker_static(worker)
|
|
|
|
self.assertEqual(worker.request_calls, 2) # one tile, then sentinel
|
|
self.assertEqual(worker.heartbeat_calls, 1)
|
|
self.assertEqual(len(worker.sent_batches), 1)
|
|
self.assertTrue(worker.sent_batches[0]["is_final_flush"])
|
|
tiles = worker.sent_batches[0]["tiles"]
|
|
self.assertEqual(len(tiles), 1)
|
|
self.assertEqual(tiles[0]["tile_idx"], 0)
|
|
self.assertEqual(tiles[0]["global_idx"], 0)
|
|
self.assertEqual(tiles[0]["batch_idx"], 0)
|
|
|
|
def test_flush_empty_final_still_sends_completion_signal(self):
|
|
worker = _FakeStaticWorker()
|
|
|
|
returned = worker._flush_tiles_to_master(
|
|
[],
|
|
"job-1",
|
|
"http://master:8188",
|
|
8,
|
|
"worker-1",
|
|
is_final_flush=True,
|
|
)
|
|
|
|
self.assertEqual(returned, [])
|
|
self.assertEqual(len(worker.sent_batches), 1)
|
|
self.assertEqual(worker.sent_batches[0]["tiles"], [])
|
|
self.assertTrue(worker.sent_batches[0]["is_final_flush"])
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|