diff --git a/scripts/distributed.py b/scripts/distributed.py index 83b3a36..b771250 100644 --- a/scripts/distributed.py +++ b/scripts/distributed.py @@ -235,6 +235,7 @@ class Script(scripts.Script): try: Script.world.initialize(batch_size) + Script.world.initial_payload = initial_payload logger.debug(f"World initialized!") except WorldAlreadyInitialized: Script.world.update_world(total_batch_size=batch_size) @@ -249,6 +250,9 @@ class Script(scripts.Script): current_thread().name = "distributed_main" Script.initialize(initial_payload=p) + # save original process_images_inner function for later if we monkeypatch it + Script.original_process_images_inner = processing.process_images_inner + # strip scripts that aren't yet supported and warn user packed_script_args: List[dict] = [] # list of api formatted per-script argument objects # { "script_name": { "args": ["value1", "value2", ...] } @@ -387,9 +391,12 @@ class Script(scripts.Script): if not Script.world.enabled: return - if len(processed.images) >= 1 and Script.master_start is not None: + if Script.master_start is not None: Script.add_to_gallery(p=p, processed=processed) + # restore process_images_inner if it was monkey-patched + processing.process_images_inner = Script.original_process_images_inner + @staticmethod def signal_handler(sig, frame): logger.debug("handling interrupt signal") diff --git a/scripts/spartan/ui.py b/scripts/spartan/ui.py index bf7f7e7..bbcdd74 100644 --- a/scripts/spartan/ui.py +++ b/scripts/spartan/ui.py @@ -354,8 +354,8 @@ class UI: with gradio.Tab('Settings'): thin_client_cbx = gradio.Checkbox( - label='Thin-client mode (experimental)', - info="(BROKEN) Only generate images using remote workers. There will be no previews when enabled.", + label='Thin-client mode', + info="Only generate images using remote workers. There will be no previews (yet) when enabled.", value=self.world.thin_client_mode ) job_timeout = gradio.Number( diff --git a/scripts/spartan/worker.py b/scripts/spartan/worker.py index cfb36bb..2dc8149 100644 --- a/scripts/spartan/worker.py +++ b/scripts/spartan/worker.py @@ -618,6 +618,10 @@ class Worker: except requests.exceptions.ConnectionError as e: logger.error(e) return False + except requests.ReadTimeout as e: + logger.critical(f"worker '{self.label}' is online but not responding (crashed?)") + logger.error(e) + return False def mark_unreachable(self): if self.state == State.DISABLED: diff --git a/scripts/spartan/world.py b/scripts/spartan/world.py index b412d74..494c6fe 100644 --- a/scripts/spartan/world.py +++ b/scripts/spartan/world.py @@ -18,8 +18,8 @@ from . import shared as sh from .pmodels import ConfigModel, Benchmark_Payload from .shared import logger, warmup_samples, extension_path from .worker import Worker, State -import asyncio from modules.call_queue import wrap_queued_call +from modules import processing class NotBenchmarked(Exception): @@ -565,6 +565,22 @@ class World: distro_summary += f"'{job.worker.label}' - {job.batch_size * iterations} image(s) @ {job.worker.avg_ipm:.2f} ipm\n" logger.info(distro_summary) + if self.thin_client_mode is True or self.master_job().batch_size == 0: + # save original process_images_inner for later so we can restore once we're done + logger.debug(f"bypassing local generation completely") + def process_images_inner_bypass(p) -> processing.Processed: + processed = processing.Processed(p, [], p.seed, info="") + processed.all_prompts = [] + processed.all_seeds = [] + processed.all_subseeds = [] + processed.all_negative_prompts = [] + processed.infotexts = [] + processed.prompt = None + + self.initial_payload.scripts.postprocess(p, processed) + return processed + processing.process_images_inner = process_images_inner_bypass + # delete any jobs that have no work last = len(self.jobs) - 1 while last > 0: