Jason HokuandClaude Opus 4.6
833158fab6
The LoRA field in job configs uses a combo string format: "name:model_str:clip_str + name2:model_str:clip_str" Parts can also be folder references like "FolderName/:0.90:0.90" or random selection syntax "Folder/[3,0.5]". _extract_required_models() was passing the entire combo string as a single filename, causing WinError 123 (invalid path) because colons are illegal in Windows filenames. Now splits on " + ", extracts the filename before ":", and handles: - Individual LoRA files: synced from master - Folder references (trailing /): expanded locally via folder_paths - Random syntax ([count,strength]): expanded locally via folder_paths Same fix applied to distribution_manager._get_upcoming_models() which builds the pre-fetch list for workers. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
803 lines
34 KiB
Python
803 lines
34 KiB
Python
"""
|
||
Distributed Worker Thread
|
||
Runs on remote ComfyUI instances. Polls the master for jobs,
|
||
processes them using the local GPU, and uploads results back.
|
||
"""
|
||
|
||
import threading
|
||
import time
|
||
import uuid
|
||
import json
|
||
import io
|
||
import os
|
||
import random
|
||
from urllib.parse import quote as url_quote
|
||
from .network_utils import (
|
||
distribution_get,
|
||
distribution_post_json,
|
||
distribution_post_multipart,
|
||
distribution_stream_download,
|
||
)
|
||
|
||
import torch
|
||
|
||
|
||
class _MasterFinishedError(Exception):
|
||
"""Raised when the master has deactivated (503 on submit).
|
||
Not a real error — means master already completed all jobs."""
|
||
pass
|
||
|
||
|
||
class WorkerThread(threading.Thread):
|
||
"""
|
||
Background thread that polls master for jobs and processes them.
|
||
|
||
Lifecycle:
|
||
1. Register with master
|
||
2. Claim job (GET /distribution/claim_job)
|
||
3. Process job locally (model load → encode → sample → VAE decode)
|
||
4. Submit result (POST /distribution/submit_result)
|
||
5. Repeat from 2 until no more jobs
|
||
6. Exit
|
||
"""
|
||
|
||
def __init__(self, master_url):
|
||
super().__init__(daemon=True)
|
||
self.master_url = master_url.rstrip("/")
|
||
self.worker_id = f"w_{uuid.uuid4().hex[:8]}"
|
||
self.jobs_processed = 0
|
||
self.current_model = None
|
||
self._stop_event = threading.Event()
|
||
self.poll_interval = 2 # seconds between polls when idle
|
||
self.heartbeat_interval = 30 # seconds between heartbeats
|
||
self.consecutive_empty = 0 # Track consecutive empty poll responses
|
||
self.max_empty_polls = 30 # Stop after this many consecutive empty polls
|
||
|
||
# Cached state for model reuse between jobs
|
||
self._loaded_model = None
|
||
self._loaded_clip = None
|
||
self._loaded_vae = None
|
||
self._patched_model = None
|
||
self._patched_clip = None
|
||
self._cached_model_key = None
|
||
self._cached_lora_key = None
|
||
self._incompatible_loras = {}
|
||
|
||
# Model sync: download missing models from master before processing
|
||
self._sync_models = False # Set by master when sync_models_to_workers enabled
|
||
self._prefetch_thread = None
|
||
self._prefetch_queue = [] # List of {category, filename} to pre-fetch
|
||
self._prefetch_lock = threading.Lock()
|
||
|
||
# =========================================================================
|
||
# MODEL SYNC HELPERS
|
||
# =========================================================================
|
||
|
||
def _download_model_from_master(self, category, filename):
|
||
"""
|
||
Download a single model file from the master instance.
|
||
Preserves full subdirectory structure.
|
||
|
||
Args:
|
||
category: Model category (checkpoints, loras, vae, etc.)
|
||
filename: Relative path including subdirs (e.g. SDXL/Illustrious/model.safetensors)
|
||
|
||
Returns:
|
||
str: Local path where file was saved, or None on failure
|
||
"""
|
||
import folder_paths
|
||
|
||
# Determine target directory
|
||
# folder_paths stores base directories per category
|
||
base_dirs = folder_paths.get_folder_paths(category)
|
||
if not base_dirs:
|
||
print(f"[Worker {self.worker_id}] ❌ Unknown model category: {category}")
|
||
return None
|
||
target_dir = base_dirs[0] # Use first (primary) directory
|
||
target_path = os.path.join(target_dir, filename.replace("/", os.sep))
|
||
|
||
# Skip if already exists
|
||
if os.path.exists(target_path):
|
||
return target_path
|
||
|
||
# Create subdirectories
|
||
os.makedirs(os.path.dirname(target_path), exist_ok=True)
|
||
|
||
# Download from master via network gateway
|
||
try:
|
||
print(f"[Worker {self.worker_id}] ⬇️ Downloading {category}/{filename}...")
|
||
last_progress = 0
|
||
|
||
def on_progress(downloaded, total_size):
|
||
nonlocal last_progress
|
||
if total_size > 0:
|
||
pct = int(downloaded * 100 / total_size)
|
||
if pct >= last_progress + 10:
|
||
last_progress = pct
|
||
size_str = f"{total_size / (1024**3):.2f} GB" if total_size > 1024**3 else f"{total_size / (1024**2):.0f} MB"
|
||
print(f"[Worker {self.worker_id}] ⬇️ {category}/{filename}: {pct}% of {size_str}")
|
||
|
||
distribution_stream_download(
|
||
url=f"{self.master_url}/distribution/download_model",
|
||
data_dict={
|
||
"worker_id": self.worker_id,
|
||
"category": category,
|
||
"filename": filename
|
||
},
|
||
target_path=target_path,
|
||
timeout=3600,
|
||
progress_callback=on_progress
|
||
)
|
||
print(f"[Worker {self.worker_id}] ✅ Downloaded {category}/{filename}")
|
||
return target_path
|
||
|
||
except Exception as e:
|
||
print(f"[Worker {self.worker_id}] ❌ Failed to download {category}/{filename}: {e}")
|
||
# Clean up partial download
|
||
temp_path = target_path + ".downloading"
|
||
if os.path.exists(temp_path):
|
||
os.remove(temp_path)
|
||
return None
|
||
|
||
def _extract_required_models(self, config):
|
||
"""
|
||
Extract all model file references from a job config.
|
||
|
||
Returns:
|
||
List of {category, filename} dicts
|
||
"""
|
||
import folder_paths
|
||
models = []
|
||
|
||
# Main model/checkpoint
|
||
model_name = config.get("model", "None")
|
||
if model_name and model_name != "None":
|
||
model_type = config.get("model_type", "checkpoint")
|
||
if model_type == "checkpoint":
|
||
category = "checkpoints"
|
||
elif model_type == "gguf":
|
||
category = "unet_gguf"
|
||
elif model_type == "diffusion_model":
|
||
category = "diffusion_models"
|
||
else:
|
||
category = "checkpoints"
|
||
models.append({"category": category, "filename": model_name})
|
||
|
||
# LoRA — combo string format: "name:model_str:clip_str + name2:model_str:clip_str"
|
||
# Parts can also be folder references like "FolderName/:0.90:0.90" which expand
|
||
# to all LoRA files in that folder, or random syntax "Folder/[3,0.5]"
|
||
lora = config.get("lora_expanded", config.get("lora", "None"))
|
||
if lora and lora != "None":
|
||
for lora_part in lora.split(" + "):
|
||
lora_part = lora_part.strip()
|
||
if not lora_part:
|
||
continue
|
||
# Extract just the filename (before first ":")
|
||
lora_filename = lora_part.split(":")[0].strip()
|
||
if not lora_filename or lora_filename == "None":
|
||
continue
|
||
|
||
# Check if it's a folder reference (ends with / or /*)
|
||
clean_name = lora_filename.rstrip("/*").rstrip("/")
|
||
is_folder = lora_filename.endswith("/") or lora_filename.endswith("/*")
|
||
# Also check for random selection syntax: folder/[count,strength]
|
||
has_bracket = "[" in lora_filename and "]" in lora_filename
|
||
if has_bracket:
|
||
clean_name = lora_filename.split("[")[0].rstrip("/")
|
||
is_folder = True
|
||
|
||
if is_folder:
|
||
# Expand folder to individual LoRA files
|
||
folder_prefix = clean_name.replace("\\", "/") + "/"
|
||
all_loras = folder_paths.get_filename_list("loras")
|
||
for lora_file in all_loras:
|
||
if lora_file.replace("\\", "/").startswith(folder_prefix):
|
||
models.append({"category": "loras", "filename": lora_file})
|
||
else:
|
||
models.append({"category": "loras", "filename": lora_filename})
|
||
|
||
# VAE (skip remote URLs and "Default")
|
||
vae = config.get("vae", "Default")
|
||
if vae and vae != "Default" and not vae.startswith("remote:"):
|
||
models.append({"category": "vae", "filename": vae})
|
||
|
||
# Text encoders
|
||
text_encoders = config.get("text_encoders", [])
|
||
for te in text_encoders:
|
||
if te and te != "None":
|
||
models.append({"category": "text_encoders", "filename": te})
|
||
|
||
return models
|
||
|
||
def _ensure_models_available(self, config):
|
||
"""
|
||
Check all required models exist locally, download missing ones from master.
|
||
Blocks until all models are available.
|
||
"""
|
||
import folder_paths
|
||
|
||
required = self._extract_required_models(config)
|
||
if not required:
|
||
return
|
||
|
||
missing = []
|
||
for m in required:
|
||
path = folder_paths.get_full_path(m["category"], m["filename"])
|
||
if path is None:
|
||
missing.append(m)
|
||
|
||
if not missing:
|
||
return
|
||
|
||
print(f"[Worker {self.worker_id}] 📦 Missing {len(missing)} model(s), downloading from master...")
|
||
for m in missing:
|
||
result = self._download_model_from_master(m["category"], m["filename"])
|
||
if result is None:
|
||
raise RuntimeError(
|
||
f"Failed to download required model: {m['category']}/{m['filename']}. "
|
||
f"Check that the file exists on the master instance."
|
||
)
|
||
|
||
def _start_prefetch(self, upcoming_models):
|
||
"""Start background thread to pre-fetch upcoming models."""
|
||
if not upcoming_models:
|
||
return
|
||
|
||
with self._prefetch_lock:
|
||
self._prefetch_queue = list(upcoming_models)
|
||
|
||
if self._prefetch_thread and self._prefetch_thread.is_alive():
|
||
return # Already running
|
||
|
||
def _prefetch_worker():
|
||
while True:
|
||
with self._prefetch_lock:
|
||
if not self._prefetch_queue:
|
||
break
|
||
model = self._prefetch_queue.pop(0)
|
||
try:
|
||
self._download_model_from_master(model["category"], model["filename"])
|
||
except Exception as e:
|
||
print(f"[Worker {self.worker_id}] ⚠️ Pre-fetch failed: {e}")
|
||
|
||
self._prefetch_thread = threading.Thread(target=_prefetch_worker, daemon=True)
|
||
self._prefetch_thread.start()
|
||
|
||
def run(self):
|
||
"""Main worker loop."""
|
||
print(f"[Worker {self.worker_id}] 🚀 Starting worker thread, master: {self.master_url}")
|
||
|
||
# Worker threads run outside ComfyUI's normal prompt execution context.
|
||
# The sampling callback (latent_preview.py → comfy/utils.py → main.py hook)
|
||
# expects PromptServer.instance.last_prompt_id to exist. Set a dummy value
|
||
# so the progress bar callback doesn't crash with AttributeError.
|
||
try:
|
||
from server import PromptServer
|
||
if PromptServer and hasattr(PromptServer, 'instance') and PromptServer.instance is not None:
|
||
if not hasattr(PromptServer.instance, 'last_prompt_id'):
|
||
PromptServer.instance.last_prompt_id = f"dist_worker_{self.worker_id}"
|
||
print(f"[Worker {self.worker_id}] 🔧 Set last_prompt_id for sampling callbacks")
|
||
except Exception:
|
||
pass
|
||
|
||
# Register with master
|
||
self._register()
|
||
|
||
# Start background heartbeat thread so heartbeats flow even during
|
||
# long-running job processing (model loading, generation, etc.).
|
||
# Without this, the main loop blocks on _process_job() and no heartbeats
|
||
# are sent, causing the master to wrongly reclaim jobs from slow workers.
|
||
self._heartbeat_thread = threading.Thread(
|
||
target=self._heartbeat_loop, daemon=True
|
||
)
|
||
self._heartbeat_thread.start()
|
||
|
||
while not self._stop_event.is_set():
|
||
# Claim a job
|
||
job = self._claim_job()
|
||
|
||
if job is None:
|
||
self.consecutive_empty += 1
|
||
if self.consecutive_empty >= self.max_empty_polls:
|
||
print(f"[Worker {self.worker_id}] 🏁 No more jobs after "
|
||
f"{self.max_empty_polls} polls, stopping")
|
||
break
|
||
self._stop_event.wait(self.poll_interval)
|
||
continue
|
||
|
||
self.consecutive_empty = 0
|
||
|
||
# Validate job has required fields
|
||
if not all(k in job for k in ("job_id", "config", "input_job")):
|
||
print(f"[Worker {self.worker_id}] ⚠️ Malformed job, skipping: {list(job.keys())}")
|
||
continue
|
||
|
||
# Sync missing models from master if enabled
|
||
try:
|
||
if self._sync_models:
|
||
self._ensure_models_available(job["config"])
|
||
|
||
# Start background pre-fetch for upcoming models
|
||
if job.get("upcoming_models"):
|
||
self._start_prefetch(job["upcoming_models"])
|
||
except Exception as e:
|
||
print(f"[Worker {self.worker_id}] ❌ Model sync failed for job {job['job_id']}: {e}")
|
||
self._report_failure(job["job_id"], f"Model sync failed: {e}")
|
||
continue
|
||
|
||
# Process the job
|
||
try:
|
||
image_bytes, meta = self._process_job(job)
|
||
|
||
# Submit result (with retry on network error)
|
||
self._submit_result(job["job_id"], image_bytes, meta)
|
||
self.jobs_processed += 1
|
||
|
||
print(f"[Worker {self.worker_id}] ✅ Job {job['job_id']} complete "
|
||
f"(total: {self.jobs_processed})")
|
||
|
||
# Periodic VRAM cleanup
|
||
if self.jobs_processed % 5 == 0:
|
||
import gc
|
||
gc.collect()
|
||
if torch.cuda.is_available():
|
||
torch.cuda.empty_cache()
|
||
|
||
except _MasterFinishedError:
|
||
# Master has shut down (503 on submit). This is NOT a failure —
|
||
# the master already finished all jobs (possibly reclaimed ours).
|
||
# Stop gracefully instead of reporting failure.
|
||
print(f"[Worker {self.worker_id}] 🏁 Master finished before we could submit "
|
||
f"job {job['job_id']} — stopping gracefully")
|
||
break
|
||
|
||
except Exception as e:
|
||
# Check if ComfyUI interrupted processing on this worker
|
||
try:
|
||
from comfy.model_management import InterruptProcessingException
|
||
if isinstance(e, InterruptProcessingException):
|
||
print(f"[Worker {self.worker_id}] 🛑 Processing interrupted, stopping worker")
|
||
self._report_failure(job["job_id"], "Worker interrupted")
|
||
break
|
||
except ImportError:
|
||
pass
|
||
|
||
print(f"[Worker {self.worker_id}] ❌ Job {job['job_id']} failed: {e}")
|
||
import traceback
|
||
traceback.print_exc()
|
||
self._report_failure(job["job_id"], str(e))
|
||
|
||
# Cleanup
|
||
self._cleanup_models()
|
||
print(f"[Worker {self.worker_id}] 🛑 Worker stopped. "
|
||
f"Processed {self.jobs_processed} jobs.")
|
||
|
||
def stop(self):
|
||
"""Signal the worker to stop after its current job."""
|
||
self._stop_event.set()
|
||
|
||
def _register(self):
|
||
"""Register with master. Retries on 503 (master still starting up)."""
|
||
max_retries = 5
|
||
retry_delay = 3 # seconds between retries
|
||
|
||
for attempt in range(1, max_retries + 1):
|
||
try:
|
||
status, resp_data = distribution_post_json(
|
||
f"{self.master_url}/distribution/register_worker",
|
||
{"worker_id": self.worker_id, "worker_url": ""},
|
||
timeout=10
|
||
)
|
||
result = json.loads(resp_data.decode("utf-8"))
|
||
print(f"[Worker {self.worker_id}] 🤝 Registered with master "
|
||
f"(session: {result.get('session_name', 'unknown')})")
|
||
return # Success
|
||
except Exception as e:
|
||
if hasattr(e, 'code') and e.code == 503 and attempt < max_retries:
|
||
print(f"[Worker {self.worker_id}] ⏳ Master not ready yet "
|
||
f"(attempt {attempt}/{max_retries}), retrying in {retry_delay}s...")
|
||
self._stop_event.wait(retry_delay)
|
||
if self._stop_event.is_set():
|
||
return
|
||
continue
|
||
print(f"[Worker {self.worker_id}] ⚠️ Registration failed: {e}")
|
||
return
|
||
|
||
def _claim_job(self):
|
||
"""Claim next job from master. Returns job dict or None."""
|
||
url = (f"{self.master_url}/distribution/claim_job"
|
||
f"?worker_id={url_quote(self.worker_id)}")
|
||
|
||
if self.consecutive_empty == 0 and self.jobs_processed == 0:
|
||
print(f"[Worker {self.worker_id}] 🔍 Claiming from: {url}")
|
||
|
||
try:
|
||
status, resp_data = distribution_get(url, timeout=10)
|
||
if status == 204:
|
||
return None
|
||
return json.loads(resp_data.decode("utf-8"))
|
||
except Exception as e:
|
||
if hasattr(e, 'code'):
|
||
if e.code == 204:
|
||
return None
|
||
if e.code == 503:
|
||
# Distribution not active yet — master may still be starting up
|
||
# (e.g. pre-encoding prompts). Don't immediately stop; let the
|
||
# normal empty-poll counter handle eventual timeout.
|
||
if self.jobs_processed == 0 and self.consecutive_empty < 5:
|
||
print(f"[Worker {self.worker_id}] ⏳ Master not active yet, "
|
||
f"waiting... ({self.consecutive_empty + 1}/5)")
|
||
elif self.jobs_processed > 0:
|
||
# Already processed jobs — master has finished and deactivated.
|
||
# Stop immediately to avoid wasting time.
|
||
print(f"[Worker {self.worker_id}] ℹ️ Distribution ended on master")
|
||
self.consecutive_empty = self.max_empty_polls
|
||
else:
|
||
print(f"[Worker {self.worker_id}] ℹ️ Distribution not active on master")
|
||
self.consecutive_empty = self.max_empty_polls
|
||
return None
|
||
print(f"[Worker {self.worker_id}] ⚠️ Claim failed: HTTP {e.code}")
|
||
return None
|
||
print(f"[Worker {self.worker_id}] ⚠️ Claim failed: {e}")
|
||
return None
|
||
|
||
@torch.inference_mode()
|
||
def _process_job(self, job):
|
||
"""
|
||
Process a single job using local ComfyUI infrastructure.
|
||
|
||
Steps:
|
||
1. Load model/LoRA (reuse if same as previous job)
|
||
2. Encode prompts with CLIP
|
||
3. Create latent noise tensor
|
||
4. Generate image (KSampler)
|
||
5. Decode with local VAE
|
||
6. Return image bytes + metadata
|
||
|
||
Args:
|
||
job: Job dict from master with config, input_job, gen_index
|
||
|
||
Returns:
|
||
tuple: (image_bytes, metadata_dict)
|
||
|
||
Note: @torch.inference_mode() is required because worker threads run
|
||
outside ComfyUI's execution engine which wraps everything in
|
||
inference_mode. Without it, Dynamic VRAM's lazy weight loading can
|
||
produce NaN values during VAE decode.
|
||
"""
|
||
from .model_loader import (
|
||
load_checkpoint, load_loras, load_vae_by_name,
|
||
load_diffusion_model_and_clip, get_latent_channels
|
||
)
|
||
from .image_generation import (
|
||
generate_image, decode_latent_with_vae, create_image_metadata
|
||
)
|
||
from .batch_encoding import encode_prompt_with_combinators
|
||
from .trigger_words import build_prompt_with_triggers
|
||
from .generation_orchestrator import get_model_cache_key
|
||
|
||
config = job["config"]
|
||
input_job = job["input_job"]
|
||
w = input_job["width"]
|
||
h = input_job["height"]
|
||
batch_idx = input_job.get("batch_idx", 0)
|
||
|
||
# Handle seed_behavior: "randomize" (same as main generation loop)
|
||
if config.get("seed_behavior") == "randomize":
|
||
config["seed"] = random.randint(0, 2**63 - 1)
|
||
|
||
# Full run seed behavior is handled by the orchestrator/master before dispatching jobs.
|
||
# Workers just use whatever seed is in the config they receive.
|
||
|
||
# ==== GPU COOLDOWN (if enabled in session settings) ====
|
||
cooldown_config = config.get("_session_cooldown", {})
|
||
if cooldown_config.get("enabled", False):
|
||
if not hasattr(self, '_gen_count'):
|
||
self._gen_count = 0
|
||
self._gen_count += 1
|
||
cooldown_every_n = int(cooldown_config.get("every_n", 1))
|
||
if self._gen_count % cooldown_every_n == 0:
|
||
cooldown_seconds = int(cooldown_config.get("seconds", 5))
|
||
clear_vram = cooldown_config.get("clear_vram", False)
|
||
print(f"[Worker {self.worker_id}] ❄️ GPU Cooldown: pausing {cooldown_seconds}s after {self._gen_count} generations")
|
||
if clear_vram:
|
||
import comfy.model_management as mm_cooldown
|
||
mm_cooldown.soft_empty_cache()
|
||
mm_cooldown.unload_all_models()
|
||
print(f"[Worker {self.worker_id}] ❄️ VRAM cleared")
|
||
self._cached_model_key = None # Force model reload on next iteration
|
||
import time as time_cooldown
|
||
time_cooldown.sleep(cooldown_seconds)
|
||
print(f"[Worker {self.worker_id}] ❄️ Cooldown complete, resuming")
|
||
|
||
# --- Model Loading (with caching between jobs) ---
|
||
target_model_key = get_model_cache_key(config)
|
||
self.current_model = config.get("model", "unknown")
|
||
|
||
if target_model_key != self._cached_model_key:
|
||
# Need to load a different model
|
||
self._cleanup_models()
|
||
|
||
model_type = config.get("model_type", "checkpoint")
|
||
if model_type == "checkpoint":
|
||
self._loaded_model, self._loaded_clip, self._loaded_vae = load_checkpoint(
|
||
config["model"], config["model"], False,
|
||
None, None, None, None, None, None, None,
|
||
model_cache=None
|
||
)
|
||
else:
|
||
self._loaded_model, self._loaded_clip, self._loaded_vae = load_diffusion_model_and_clip(
|
||
model_name=config["model"],
|
||
model_type=model_type,
|
||
text_encoder_paths=config.get("text_encoders", []),
|
||
clip_type_str=config.get("clip_type", "stable_diffusion"),
|
||
gguf_options=config.get("gguf_options"),
|
||
use_remote_vae=False,
|
||
optional_model=None,
|
||
optional_clip=None,
|
||
optional_vae=None,
|
||
model_cache=None
|
||
)
|
||
|
||
self._cached_model_key = target_model_key
|
||
self._cached_lora_key = None # Force LoRA reload on model switch
|
||
print(f"[Worker {self.worker_id}] 📦 Loaded model: {config['model']}")
|
||
|
||
# --- Per-config VAE loading ---
|
||
# Supports three modes:
|
||
# 1. "Default" → use model's bundled VAE
|
||
# 2. "remote:https://..." → remote HuggingFace VAE endpoint (decoded at decode step)
|
||
# 3. "ae.safetensors" → load local VAE file
|
||
config_vae = config.get("vae", "Default")
|
||
_remote_vae_url = None
|
||
if config_vae != "Default":
|
||
if config_vae.startswith("remote:"):
|
||
_remote_vae_url = config_vae[len("remote:"):]
|
||
vae = None # No local VAE needed — will use remote endpoint
|
||
print(f"[Worker {self.worker_id}] 🌐 Using remote VAE: {_remote_vae_url}")
|
||
else:
|
||
vae = load_vae_by_name(config_vae)
|
||
else:
|
||
vae = self._loaded_vae
|
||
|
||
# Validate VAE is available (non-checkpoint models like GGUF may not bundle a VAE)
|
||
# Skip this check when using remote VAE — no local VAE object needed.
|
||
if vae is None and _remote_vae_url is None:
|
||
raise RuntimeError(
|
||
f"No VAE available for model '{config.get('model', 'unknown')}'. "
|
||
f"Non-checkpoint models (GGUF/diffusion) require a VAE to be specified "
|
||
f"in the config builder, or connect one via optional_vae on the master."
|
||
)
|
||
|
||
# --- LoRA Loading (with caching) ---
|
||
lora_key = config.get("lora_expanded", config.get("lora", "None"))
|
||
if lora_key != self._cached_lora_key:
|
||
if lora_key != "None":
|
||
self._patched_model, self._patched_clip, should_skip = load_loras(
|
||
self._loaded_model, self._loaded_clip, lora_key,
|
||
config["model"], self._incompatible_loras,
|
||
model_cache=None
|
||
)
|
||
if should_skip:
|
||
raise RuntimeError(f"LoRA incompatible: {lora_key}")
|
||
else:
|
||
self._patched_model = self._loaded_model
|
||
self._patched_clip = self._loaded_clip
|
||
|
||
self._cached_lora_key = lora_key
|
||
print(f"[Worker {self.worker_id}] 🔗 Applied LoRA: {lora_key}")
|
||
|
||
# --- Prompt Encoding ---
|
||
# Check for master pre-encoded conditionings (skip CLIP encoding if available)
|
||
encoded_pos = job.get("encoded_positive")
|
||
encoded_neg = job.get("encoded_negative")
|
||
|
||
if encoded_pos is not None and encoded_neg is not None:
|
||
# Master already encoded these prompts — deserialize and use directly.
|
||
# Use the master's prompt text for metadata (avoids CivitAI trigger lookups
|
||
# on the worker since the master already applied trigger words during encoding).
|
||
from .generation_orchestrator import _serializable_to_conditioning
|
||
pos_cond = _serializable_to_conditioning(encoded_pos)
|
||
neg_cond = _serializable_to_conditioning(encoded_neg)
|
||
actual_positive = job.get("actual_positive_text", config.get("positive", ""))
|
||
actual_negative = job.get("actual_negative_text", config.get("negative", ""))
|
||
print(f"[Worker {self.worker_id}] 🧠 Using master pre-encoded conditionings")
|
||
else:
|
||
# Fallback: encode locally with CLIP (need to build prompt with trigger words)
|
||
lora_triggerwords_mode = config.get("_lora_triggerwords_mode", "None")
|
||
try:
|
||
actual_positive, _ = build_prompt_with_triggers(config, lora_triggerwords_mode)
|
||
except Exception:
|
||
actual_positive = config.get("positive", "")
|
||
actual_negative = config.get("negative", "")
|
||
clip_skip = config.get("clip_skip", 0)
|
||
pos_cond = encode_prompt_with_combinators(self._patched_clip, actual_positive, clip_skip)
|
||
neg_cond = encode_prompt_with_combinators(self._patched_clip, actual_negative, clip_skip)
|
||
|
||
# --- Create Latent ---
|
||
latent_channels = get_latent_channels(self._loaded_model, None)
|
||
latent_in = {"samples": torch.zeros([1, latent_channels, h // 8, w // 8])}
|
||
|
||
# --- Generate Image ---
|
||
attention_mode = config.get("attention_mode", "default")
|
||
result_latent, duration = generate_image(
|
||
self._patched_model, config["seed"], config["steps"], config["cfg"],
|
||
config["sampler"], config["scheduler"], pos_cond, neg_cond,
|
||
latent_in, config["denoise"],
|
||
attention_mode=attention_mode,
|
||
model_sampling_override=config.get("model_sampling_override", "none"),
|
||
model_sampling_shift=config.get("model_sampling_shift", 1.73),
|
||
model_sampling_flux_max_shift=config.get("model_sampling_flux_max_shift", 1.15),
|
||
model_sampling_flux_base_shift=config.get("model_sampling_flux_base_shift", 0.5),
|
||
use_advanced_sampling=config.get("use_advanced_sampling", False),
|
||
advanced_guider=config.get("advanced_guider", "cfg_guider"),
|
||
advanced_scheduler=config.get("advanced_scheduler", "basic"),
|
||
flux_guidance_value=config.get("flux_guidance_value", 0.0),
|
||
width=w,
|
||
height=h
|
||
)
|
||
|
||
# --- VAE Decode ---
|
||
# Clone latent samples and free UNet-related tensors BEFORE VAE decode.
|
||
# With Dynamic VRAM, the model manager needs to evict the UNet to load
|
||
# the VAE. Holding references to UNet output tensors on GPU can interfere
|
||
# with memory eviction, causing the VAE to run in a degraded state that
|
||
# produces NaN values.
|
||
latent_for_decode = result_latent["samples"].clone()
|
||
del result_latent, latent_in, pos_cond, neg_cond
|
||
if torch.cuda.is_available():
|
||
torch.cuda.empty_cache()
|
||
|
||
if _remote_vae_url:
|
||
# Remote VAE: send latents to HuggingFace endpoint for decoding
|
||
from .remote_vae import remote_decode_hf
|
||
from PIL import Image as PILImage
|
||
import numpy as np
|
||
if latent_for_decode.ndim == 3:
|
||
latent_for_decode = latent_for_decode.unsqueeze(0)
|
||
decoded = remote_decode_hf(_remote_vae_url, latent_for_decode, h, w)
|
||
del latent_for_decode
|
||
# Denormalize from [-1, 1] → [0, 1] (same as VaeImageProcessor.postprocess)
|
||
decoded = (decoded / 2 + 0.5).clamp(0, 1)
|
||
img_data = decoded[0].permute(1, 2, 0).cpu().numpy()
|
||
image = PILImage.fromarray((img_data * 255).round().astype(np.uint8))
|
||
else:
|
||
image = decode_latent_with_vae(vae, latent_for_decode)
|
||
del latent_for_decode
|
||
|
||
# --- Convert to bytes ---
|
||
buf = io.BytesIO()
|
||
image.save(buf, format="WEBP", quality=80)
|
||
image_bytes = buf.getvalue()
|
||
del image, buf # Free PIL image and buffer
|
||
|
||
# --- Create metadata ---
|
||
ts = int(time.time() * 100000) + random.randint(0, 1000)
|
||
meta = create_image_metadata(
|
||
config, w, h, duration, config["seed"],
|
||
batch_idx, actual_positive, actual_negative,
|
||
gen_index=job.get("gen_index")
|
||
)
|
||
meta["id"] = ts
|
||
|
||
return image_bytes, meta
|
||
|
||
def _submit_result(self, job_id, image_bytes, meta, max_retries=3):
|
||
"""
|
||
Upload image + metadata to master via multipart POST.
|
||
Uses network_utils.py gateway for all outbound requests.
|
||
Retries on transient network errors (connection reset, timeout).
|
||
"""
|
||
# JSON-safe meta (ensure no non-serializable types from config round-trip)
|
||
try:
|
||
safe_meta = json.loads(json.dumps(meta, default=str))
|
||
except Exception:
|
||
safe_meta = meta
|
||
|
||
boundary = f"----DistWorker{uuid.uuid4().hex[:16]}"
|
||
|
||
# Build multipart body manually
|
||
body = b""
|
||
|
||
# Metadata part
|
||
body += f"--{boundary}\r\n".encode()
|
||
body += b'Content-Disposition: form-data; name="metadata"\r\n'
|
||
body += b'Content-Type: application/json\r\n\r\n'
|
||
body += json.dumps({
|
||
"job_id": job_id,
|
||
"meta": safe_meta,
|
||
"worker_id": self.worker_id
|
||
}).encode("utf-8")
|
||
body += b"\r\n"
|
||
|
||
# Image part
|
||
body += f"--{boundary}\r\n".encode()
|
||
body += b'Content-Disposition: form-data; name="image"; filename="image.webp"\r\n'
|
||
body += b'Content-Type: image/webp\r\n\r\n'
|
||
body += image_bytes
|
||
body += b"\r\n"
|
||
|
||
body += f"--{boundary}--\r\n".encode()
|
||
|
||
last_error = None
|
||
for attempt in range(max_retries):
|
||
try:
|
||
status, _ = distribution_post_multipart(
|
||
f"{self.master_url}/distribution/submit_result",
|
||
f"multipart/form-data; boundary={boundary}",
|
||
body,
|
||
timeout=60
|
||
)
|
||
if status != 200:
|
||
raise RuntimeError(f"Submit failed: HTTP {status}")
|
||
return # Success
|
||
|
||
except Exception as e:
|
||
# 503 = master deactivated (all jobs done). Don't retry — it won't help.
|
||
# Raise a special exception so the main loop can stop gracefully.
|
||
if hasattr(e, 'code') and e.code == 503:
|
||
raise _MasterFinishedError(
|
||
f"Master distribution is no longer active (HTTP 503)"
|
||
)
|
||
last_error = e
|
||
if attempt < max_retries - 1:
|
||
wait = 2 ** attempt # Exponential backoff: 1s, 2s, 4s
|
||
print(f"[Worker {self.worker_id}] ⚠️ Submit retry {attempt + 1}/{max_retries} "
|
||
f"for job {job_id}: {e} (waiting {wait}s)")
|
||
time.sleep(wait)
|
||
|
||
raise RuntimeError(f"Submit failed after {max_retries} attempts: {last_error}")
|
||
|
||
def _send_heartbeat(self):
|
||
"""Send heartbeat to master."""
|
||
try:
|
||
distribution_post_json(
|
||
f"{self.master_url}/distribution/heartbeat",
|
||
{"worker_id": self.worker_id},
|
||
timeout=5
|
||
)
|
||
except Exception:
|
||
pass
|
||
|
||
def _heartbeat_loop(self):
|
||
"""
|
||
Background heartbeat loop. Runs in a separate daemon thread so that
|
||
heartbeats are sent continuously even while _process_job() blocks the
|
||
main worker loop. This prevents the master from incorrectly reclaiming
|
||
jobs from slow-but-alive workers.
|
||
"""
|
||
while not self._stop_event.is_set():
|
||
self._send_heartbeat()
|
||
self._stop_event.wait(self.heartbeat_interval)
|
||
|
||
def _report_failure(self, job_id, error):
|
||
"""Report job failure to master."""
|
||
try:
|
||
distribution_post_json(
|
||
f"{self.master_url}/distribution/fail_job",
|
||
{
|
||
"job_id": job_id,
|
||
"error": error,
|
||
"worker_id": self.worker_id
|
||
},
|
||
timeout=10
|
||
)
|
||
except Exception:
|
||
pass
|
||
|
||
def _cleanup_models(self):
|
||
"""Release cached model references."""
|
||
self._loaded_model = None
|
||
self._loaded_clip = None
|
||
self._loaded_vae = None
|
||
self._patched_model = None
|
||
self._patched_clip = None
|
||
self._cached_model_key = None
|
||
self._cached_lora_key = None
|
||
self.current_model = None
|
||
|
||
import gc
|
||
gc.collect()
|
||
if torch.cuda.is_available():
|
||
torch.cuda.empty_cache()
|