Spaces:
Running on Zero
Running on Zero
File size: 8,276 Bytes
804ee23 | 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 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 | from __future__ import annotations
from collections.abc import Iterable
from copy import deepcopy
from dataclasses import dataclass
from dots_tts.data.pipelines.base import BaseSamplePipeline
from dots_tts.data.source_adapters.base_adapter import (
BaseSourceAdapter,
SourceContext,
)
@dataclass(frozen=True)
class SourceSpec:
name: str
weight: float
adapter: BaseSourceAdapter
pipeline: BaseSamplePipeline
_UINT64_MASK = 0xFFFFFFFFFFFFFFFF
def _mix_uint64(value: int) -> int:
value = (value ^ (value >> 30)) * 0xBF58476D1CE4E5B9
value &= _UINT64_MASK
value = (value ^ (value >> 27)) * 0x94D049BB133111EB
value &= _UINT64_MASK
return (value ^ (value >> 31)) & _UINT64_MASK
def _stable_seed(*parts: int) -> int:
value = 0x9E3779B97F4A7C15
for part in parts:
value = (value + int(part) + 0x9E3779B97F4A7C15) & _UINT64_MASK
value = _mix_uint64(value)
return value
class SequentialMultiSourceAdapter(BaseSourceAdapter):
"""Finite adapter that concatenates sources in the configured order."""
def __init__(self, *, sources: list[SourceSpec]):
if not sources:
raise ValueError(
"SequentialMultiSourceAdapter requires at least one source."
)
self.sources = list(sources)
def initial_state(self) -> dict:
return {
"source_index": 0,
"sources": {
source.name: source.adapter.initial_state() for source in self.sources
},
}
def is_cycle_start_state(self, state: dict | None) -> bool:
normalized = self.normalize_state(state)
if int(normalized["source_index"]) != 0:
return False
return all(
source.adapter.is_cycle_start_state(normalized["sources"][source.name])
for source in self.sources
)
def normalize_state(self, state: dict | None) -> dict:
normalized = super().normalize_state(state)
source_states = normalized.get("sources") or {}
normalized["sources"] = {
source.name: source.adapter.clone_state(source_states.get(source.name))
for source in self.sources
}
normalized["source_index"] = int(normalized.get("source_index", 0))
return normalized
def clone_state(self, state: dict | None) -> dict:
return deepcopy(self.normalize_state(state))
def iter_samples(
self,
context: SourceContext,
*,
state: dict | None = None,
) -> Iterable[dict]:
live_state = self.normalize_state(state)
start_index = int(live_state["source_index"])
for index in range(start_index, len(self.sources)):
source = self.sources[index]
child_state = live_state["sources"][source.name]
raw_iter = source.adapter.iter_samples(context, state=child_state)
for sample in source.pipeline(raw_iter):
item = dict(sample)
next_child_state = item.pop("_adapter_state", None)
if next_child_state is None:
raise RuntimeError(
f"{source.adapter.__class__.__name__} must attach '_adapter_state' to samples."
)
live_state["source_index"] = index
live_state["sources"][source.name] = source.adapter.clone_state(
next_child_state
)
item["source_name"] = source.name
item["_adapter_state"] = self.clone_state(live_state)
yield item
live_state["source_index"] = index + 1
class WeightedMultiSourceAdapter(BaseSourceAdapter):
"""Infinite weighted sampler that cycles each child source independently."""
def __init__(self, *, sources: list[SourceSpec]):
if not sources:
raise ValueError("WeightedMultiSourceAdapter requires at least one source.")
invalid = [source.name for source in sources if float(source.weight) <= 0.0]
if invalid:
raise ValueError(f"Source weights must be positive: {invalid}")
self.sources = list(sources)
self._cumulative_weights = []
total = 0.0
for source in self.sources:
total += float(source.weight)
self._cumulative_weights.append(total)
self._total_weight = total
def initial_state(self) -> dict:
return {
"draw_count": 0,
"sources": {
source.name: source.adapter.initial_state() for source in self.sources
},
}
def is_cycle_start_state(self, state: dict | None) -> bool:
normalized = self.normalize_state(state)
if int(normalized["draw_count"]) != 0:
return False
return all(
source.adapter.is_cycle_start_state(normalized["sources"][source.name])
for source in self.sources
)
def normalize_state(self, state: dict | None) -> dict:
normalized = super().normalize_state(state)
source_states = normalized.get("sources") or {}
normalized["sources"] = {
source.name: source.adapter.clone_state(source_states.get(source.name))
for source in self.sources
}
normalized["draw_count"] = int(normalized.get("draw_count", 0))
return normalized
def clone_state(self, state: dict | None) -> dict:
return deepcopy(self.normalize_state(state))
def _source_draw_value(self, context: SourceContext, draw_count: int) -> float:
raw = _stable_seed(
context.seed,
context.epoch,
context.rank,
context.worker_id,
draw_count,
)
return (raw / float(1 << 64)) * self._total_weight
def _pick_source(self, context: SourceContext, draw_count: int) -> SourceSpec:
draw_value = self._source_draw_value(context, draw_count)
for source, upper in zip(self.sources, self._cumulative_weights, strict=True):
if draw_value < upper:
return source
return self.sources[-1]
def iter_samples(
self,
context: SourceContext,
*,
state: dict | None = None,
) -> Iterable[dict]:
live_state = self.normalize_state(state)
iterators: dict[str, object] = {}
while True:
draw_count = int(live_state["draw_count"])
source = self._pick_source(context, draw_count)
while True:
child_state = live_state["sources"][source.name]
child_iter = iterators.get(source.name)
if child_iter is None:
raw_iter = source.adapter.iter_samples(context, state=child_state)
child_iter = iter(source.pipeline(raw_iter))
iterators[source.name] = child_iter
try:
sample = dict(next(child_iter))
except StopIteration:
if source.adapter.is_cycle_start_state(child_state):
raise RuntimeError(
"Weighted source yielded no samples for this worker. "
f"source={source.name!r}, worker={context.global_worker_id}, "
f"epoch={context.epoch}"
)
iterators.pop(source.name, None)
live_state["sources"][source.name] = source.adapter.advance_cycle(
child_state
)
continue
next_child_state = sample.pop("_adapter_state", None)
if next_child_state is None:
raise RuntimeError(
f"{source.adapter.__class__.__name__} must attach '_adapter_state' to samples."
)
live_state["sources"][source.name] = source.adapter.clone_state(
next_child_state
)
live_state["draw_count"] = draw_count + 1
sample["source_name"] = source.name
sample["_adapter_state"] = self.clone_state(live_state)
yield sample
break
|