Spaces:
Running on Zero
Running on Zero
File size: 6,364 Bytes
64db20c 287aa00 64db20c 287aa00 64db20c 287aa00 64db20c 287aa00 64db20c | 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 | """Per-request resource lifecycle.
Tracks every resource acquired during a generation (LoRAs including turbo,
VAE residency, keyframe images) and tears them all down in `_RequestContext.release()`,
even on failure. The authoritative record of what this task applied is
`_RequestContext.applied` / the parent `GenerationTask.applied_loras` — never
the global `*_merged` sets as request truth.
Intended to be used under `_MODEL_MUTATION_LOCK` for the mutate+generate
critical section when physical merge/unmerge is required.
"""
from __future__ import annotations
import time
from config import NSFW_LORA_CONFIGS, PLACEMENT
from lora.apply import _remove_custom, _remove_lora, _remove_nsfw, _remove_turbo
import state as _state
from state import (
_CUSTOM_LoRA,
_CUSTOM_merged,
_LORA_ACCESS,
_LORA_MAX_ENTRIES,
_NSFW_LoRA,
_NSFW_merged,
_log,
)
from task import _task_log
def _touch_lora(key):
"""Update the last-access timestamp for a LoRA entry."""
_LORA_ACCESS[key] = time.time()
def _evict_stale_loras():
"""Remove the oldest LoRA entries when either store exceeds _LORA_MAX_ENTRIES."""
for store, merged, _scale_fn in [
(_CUSTOM_LoRA, _CUSTOM_merged, None),
(_NSFW_LoRA, _NSFW_merged, None),
]:
if len(store) <= _LORA_MAX_ENTRIES:
continue
sorted_keys = sorted(store.keys(), key=lambda k: _LORA_ACCESS.get(k, 0))
n_evict = len(store) - _LORA_MAX_ENTRIES
for key in sorted_keys[:n_evict]:
try:
scale = NSFW_LORA_CONFIGS[key]["scale"] if key in NSFW_LORA_CONFIGS else 1.0
_remove_lora(key, scale, store, merged)
store.pop(key, None)
_LORA_ACCESS.pop(key, None)
_log(f"evicted stale LoRA {key}")
except Exception as _e:
_log(f"evict failed for {key}: {_e}", "ERROR")
class _RequestContext:
"""Per-request (task-scoped) resource lifecycle manager."""
def __init__(self, task=None):
self.task = task
self._custom: list[tuple[tuple[str, str], float]] = [] # (key, scale)
self._nsfw: list[tuple[str, float]] = [] # (key, scale)
self._turbo: bool = False
self._vae_on_gpu = False
self._pil_images: list = [] # keyframe PILs to del
self.applied: list[str] = [] # human-readable names of applied adapters
# ── Acquire helpers ────────────────────────────────────────────────────────
def acquire_vae(self):
"""Move the two VAEs to GPU for this request (pack mode)."""
if PLACEMENT == "pack" and _state.PIPE is not None:
_state.PIPE.vae.to("cuda")
_state.PIPE.audio_vae.to("cuda")
self._vae_on_gpu = True
def acquire_turbo(self):
"""Track that turbo LoRA was merged (will be unmerged on release)."""
self._turbo = True
self.applied.append("turbo")
if self.task is not None:
self.task.applied_loras.append("turbo")
def acquire_custom_lora(self, key, scale):
"""Track a custom LoRA that was merged (will be unmerged on release)."""
self._custom.append((key, scale))
label = f"custom:{key[0]}/{key[1]}"
self.applied.append(label)
if self.task is not None:
self.task.applied_loras.append(label)
_touch_lora(key)
def acquire_nsfw_lora(self, key, scale):
"""Track an NSFW LoRA that was merged (will be unmerged on release)."""
self._nsfw.append((key, scale))
label = f"nsfw:{key}"
self.applied.append(label)
if self.task is not None:
self.task.applied_loras.append(label)
_touch_lora(key)
def track_pil(self, img):
"""Register a PIL image for cleanup on release."""
if img is not None:
self._pil_images.append(img)
# ── Release ────────────────────────────────────────────────────────────────
def release(self):
"""Release ALL acquired resources. Safe to call multiple times.
Unmerge order is reverse of application so baseline is restored.
"""
# 1. Unmerge LoRAs in reverse application order — each independently.
for key, scale in reversed(self._custom):
try:
_remove_custom(key, scale)
except Exception as _e:
_task_log(self.task, f"custom LoRA unmerge failed for {key}: {_e}", "ERROR")
self._custom.clear()
for key, scale in reversed(self._nsfw):
try:
_remove_nsfw(key, scale)
except Exception as _e:
_task_log(self.task, f"NSFW LoRA unmerge failed for {key}: {_e}", "ERROR")
self._nsfw.clear()
if self._turbo:
try:
_remove_turbo()
except Exception as _e:
_task_log(self.task, f"turbo LoRA unmerge failed: {_e}", "ERROR")
self._turbo = False
# 2. Move VAEs back to CPU to free ~10 GB GPU between requests.
if self._vae_on_gpu and _state.PIPE is not None:
try:
_state.PIPE.vae.to("cpu")
_state.PIPE.audio_vae.to("cpu")
except Exception as _e:
_task_log(self.task, f"VAE CPU offload failed: {_e}", "ERROR")
self._vae_on_gpu = False
# 3. Drop PIL images.
for img in self._pil_images:
try:
img.close()
except Exception:
pass
self._pil_images.clear()
# 4. Evict stale LoRAs from the stores (LRU, bucket-cached so re-load is fast).
try:
_evict_stale_loras()
except Exception as _e:
_task_log(self.task, f"LoRA eviction error: {_e}", "ERROR")
# 5. Aggressive CUDA cleanup: del + gc + empty_cache.
import gc as _gc
_gc.collect()
try:
import torch as _torch
_torch.cuda.empty_cache()
except Exception:
pass
|