ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

[pre-commit.ci] auto fixes from pre-commit.com hooks

Browse source 
for more information, see https://pre-commit.ci
This commit is contained in:
pre-commit-ci[bot] 2026-07-25 08:59:55 +00:00
commit 7c9521810e
27 changed files with 298 additions and 196 deletions
Download patch file Download diff file Expand all files Collapse all files

34
studio/backend/core/inference/diffusion.py
View file

@ -506,7 +506,7 @@ def _install_gguf_prefix_strip(transformer_cls: Any, logger: Any) -> None:
def _stripped_mapping_fn(checkpoint = None, **kwargs):
checkpoint = {
(key[len(prefix):] if key.startswith(prefix) else key): value
(key[len(prefix) :] if key.startswith(prefix) else key): value
for key, value in (checkpoint or {}).items()
}
return original(checkpoint = checkpoint, **kwargs)
@ -1303,6 +1303,7 @@ class DiffusionBackend:
logger = logger,
)
if candidate is not None:
def _replan_candidate():
return self._plan_memory(
target,
@ -1982,8 +1983,16 @@ class DiffusionBackend:
)
if transformer is not None:
pipe = self._assemble_pipe(
pipeline_cls, base, transformer, dtype, hf_token, device, base_local_dir,
fam = fam, te_quant_mode = text_encoder_quant, target = target,
pipeline_cls,
base,
transformer,
dtype,
hf_token,
device,
base_local_dir,
fam = fam,
te_quant_mode = text_encoder_quant,
target = target,
)
return pipe, scheme
@ -1998,8 +2007,16 @@ class DiffusionBackend:
base, subfolder = "transformer", torch_dtype = dtype, token = hf_token
)
pipe = self._assemble_pipe(
pipeline_cls, base, transformer, dtype, hf_token, device, base_local_dir,
fam = fam, te_quant_mode = text_encoder_quant, target = target,
pipeline_cls,
base,
transformer,
dtype,
hf_token,
device,
base_local_dir,
fam = fam,
te_quant_mode = text_encoder_quant,
target = target,
)
if lora_specs:
# Bake the adapters BEFORE quantize_: peft injects its wrappers on the dense
@ -3032,8 +3049,7 @@ class DiffusionBackend:
chunk_kwargs = dict(kwargs)
shared = uniform_prompt(chunk)
generators = [
torch.Generator(device = state.device).manual_seed(s)
for _, s in chunk
torch.Generator(device = state.device).manual_seed(s) for _, s in chunk
]
if len(jobs) == 1:
# Single image: scalar prompt + generator, exactly the pre-batching
@ -3060,9 +3076,7 @@ class DiffusionBackend:
raise
# OOM backoff: halve the failed chunk and retry; finished chunks keep
# their images and per-image seeds keep every retry reproducible.
empty_cache = getattr(
getattr(torch, "cuda", None), "empty_cache", None
)
empty_cache = getattr(getattr(torch, "cuda", None), "empty_cache", None)
if callable(empty_cache):
empty_cache()
first_half, second_half = split_chunk(chunk)

4
studio/backend/core/inference/diffusion_batched.py
View file

@ -96,9 +96,7 @@ def resolve_batch_jobs(
return list(zip(job_prompts, job_seeds)), base_seed
def chunk_jobs(
jobs: list[tuple[str, int]], batch_size: int
) -> list[list[tuple[str, int]]]:
def chunk_jobs(jobs: list[tuple[str, int]], batch_size: int) -> list[list[tuple[str, int]]]:
"""Split the jobs into per-forward chunks.
``batch_size`` doubles as the per-forward cap when a prompt/seed list drives

6
studio/backend/core/inference/diffusion_cond_cache.py
View file

@ -135,7 +135,6 @@ def install(
# Lazy: core.training imports parts of core.inference, so the module-level
# import would be circular; the extras module itself is stdlib-only.
from core.training.diffusion_train_extras import PersistentConditioningCache
signature = inspect.signature(encode)
cache = PersistentConditioningCache(root, family, 0)
except Exception as exc: # noqa: BLE001 — cache is best-effort
@ -167,9 +166,7 @@ def install(
if not all(_json_safe(v) for v in keyed.values()):
# Tensor/object arguments (pre-supplied embeds, images) are not keyable.
return encode(*args, **kwargs)
payload = json.dumps(
{"load": load_fp, "args": keyed}, sort_keys = True, default = str
)
payload = json.dumps({"load": load_fp, "args": keyed}, sort_keys = True, default = str)
key = cache.text_key(
f"inference::{hashlib.sha256(payload.encode('utf-8')).hexdigest()}"
)
@ -217,7 +214,6 @@ def _target_device(pipe: Any, bound: inspect.BoundArguments) -> Any:
def _diffusers_version() -> Optional[str]:
try:
import diffusers # noqa: PLC0415
return str(getattr(diffusers, "__version__", None))
except Exception: # noqa: BLE001
return None

4
studio/backend/core/inference/diffusion_families.py
View file

@ -623,7 +623,9 @@ def default_generation_params(*identifiers: Optional[str]) -> tuple[int, float]:
def family_prequant_repo(
fam: DiffusionFamily, scheme: str, base_repo: Optional[str] = None
fam: DiffusionFamily,
scheme: str,
base_repo: Optional[str] = None,
) -> Optional[str]:
"""The hosted pre-quantized transformer repo for ``scheme`` in this family, or None.

1
studio/backend/core/inference/diffusion_hidream.py
View file

@ -75,7 +75,6 @@ def hidream_te4_kwargs(
load_prequant_text_encoder,
resolve_te_prequant_source,
)
source = resolve_te_prequant_source(fam, "text_encoder_4", "fp8")
if source is not None:
encoder = load_prequant_text_encoder(

6
studio/backend/core/inference/diffusion_memory.py
View file

@ -149,7 +149,9 @@ def snapshot_device_memory(target: Any) -> DeviceMemory:
def settled_snapshot_device_memory(
target: Any, attempts: int = 3, delay_s: float = 1.0
target: Any,
attempts: int = 3,
delay_s: float = 1.0,
) -> DeviceMemory:
"""``snapshot_device_memory`` hardened against TRANSIENT free-VRAM undercounts on cuda.
@ -165,7 +167,6 @@ def settled_snapshot_device_memory(
return snapshot_device_memory(target)
try:
import torch
torch.cuda.synchronize()
torch.cuda.empty_cache()
except Exception: # noqa: BLE001 — settle is best-effort; the snapshot below still runs
@ -178,7 +179,6 @@ def settled_snapshot_device_memory(
break
try:
import time
time.sleep(delay_s)
except Exception: # noqa: BLE001
break

10
studio/backend/core/inference/diffusion_prequant.py
View file

@ -159,9 +159,7 @@ def usable_prequant_source(
memory planning falls back to dense-fit checks up front, instead of the loader refusing
the path only after the resident pipeline was evicted and dense bf16 materialises under
a plan that never budgeted for it (evict-then-OOM). Hosted-repo sources are unaffected."""
src = resolve_prequant_source(
fam, scheme, path_override = path_override, base_repo = base_repo
)
src = resolve_prequant_source(fam, scheme, path_override = path_override, base_repo = base_repo)
if src is not None and src.kind == "path" and not local_prequant_path_ready(src.location):
return None
return src
@ -271,10 +269,14 @@ def _resolve_checkpoint_path(source: PrequantSource, hf_token: Optional[str]) ->
try:
from huggingface_hub.errors import EntryNotFoundError
except Exception: # noqa: BLE001 — older hub layouts; fall back to a private marker
class EntryNotFoundError(Exception): # type: ignore[no-redef]
pass
try:
return hf_hub_download(repo_id = source.location, filename = source.filename, token = hf_token)
return hf_hub_download(
repo_id = source.location, filename = source.filename, token = hf_token
)
except EntryNotFoundError:
if not source.fallback_filename or source.fallback_filename == source.filename:
raise

7
studio/backend/core/inference/diffusion_te_prequant.py
View file

@ -334,15 +334,11 @@ def _resolve_checkpoint_path(source: TePrequantSource, hf_token: Optional[str])
"""The local file path for ``source``, downloading from the Hub if needed; None if absent."""
if source.kind == "path":
import os
expanded = os.path.expanduser(source.location)
return expanded if os.path.isfile(expanded) else None
if source.kind == "repo":
from huggingface_hub import hf_hub_download
return hf_hub_download(
repo_id = source.location, filename = source.filename, token = hf_token
)
return hf_hub_download(repo_id = source.location, filename = source.filename, token = hf_token)
return None
@ -391,7 +387,6 @@ def _validate_checkpoint(ckpt: Any, scheme: str, component: str, base: str, logg
def _has_meta_tensors(module: Any) -> bool:
"""True if any parameter or buffer is still on the meta device after loading."""
from itertools import chain
try:
return any(
getattr(t, "is_meta", False) for t in chain(module.parameters(), module.buffers())

2
studio/backend/core/inference/diffusion_transformer_quant.py
View file

@ -363,6 +363,7 @@ def _make_quant_config(scheme: str, fast_accum: Optional[bool] = None) -> Any:
# (torchao >= 0.13); older versions keep today's behaviour via the signature check.
import inspect
from torchao.quantization import PerRow
fp8_kwargs: dict = {"granularity": PerRow()}
config_params = inspect.signature(Float8DynamicActivationFloat8WeightConfig).parameters
if "activation_value_lb" in config_params:
@ -378,7 +379,6 @@ def _make_quant_config(scheme: str, fast_accum: Optional[bool] = None) -> Any:
from torchao.quantization.quantize_.common.kernel_preference import (
KernelPreference,
)
fp8_kwargs["kernel_preference"] = KernelPreference.TORCH
except Exception: # noqa: BLE001 — enum moved: keep the library default
pass

1
studio/backend/core/inference/video.py
View file

@ -1607,7 +1607,6 @@ class VideoBackend:
ltx2_distilled_ids,
ltx23_verbatim_sigmas,
)
if steps == len(LTX23_DISTILLED_SIGMAS) and ltx2_distilled_ids(
state.gguf_filename, state.repo_id, state.base_repo
):

11
studio/backend/core/inference/video_ltx2.py
View file

@ -355,7 +355,14 @@ def checkpoint_variant(checkpoint_path: Path | str) -> str:
# (measured second sigma 0.945-0.981 vs 0.99375, and a 0.37-0.61 -> 0.1 tail vs 0.725 -> 0.42),
# so the distilled default of 8 steps must pass this list verbatim.
LTX23_DISTILLED_SIGMAS: tuple[float, ...] = (
1.0, 0.99375, 0.9875, 0.98125, 0.975, 0.909375, 0.725, 0.421875,
1.0,
0.99375,
0.9875,
0.98125,
0.975,
0.909375,
0.725,
0.421875,
)
@ -395,8 +402,6 @@ def ltx23_verbatim_sigmas(pipe: Any) -> Any:
return _ctx()
# ── component builders ───────────────────────────────────────────────────────

43
studio/backend/core/training/diffusion_dit_trainer.py
View file

@ -201,6 +201,7 @@ def _bell_loss_weights(num_train_timesteps):
centered mid-schedule, floored at 0 and normalized to mean 1 (so the expected loss
scale is unchanged). Indexed by round(sigma * num_train_timesteps)."""
import torch
steps = num_train_timesteps
t = torch.arange(steps, dtype = torch.float32)
w = torch.exp(-2.0 * ((t - steps / 2) / steps) ** 2)
@ -370,7 +371,6 @@ def _fp8_training_config():
dim never aborts the scaled_mm (the rowwise recipe manages its own padding rules and
rejects the knob, hence the split)."""
from torchao.float8 import Float8LinearConfig
try:
return Float8LinearConfig.from_recipe_name("rowwise")
except Exception: # noqa: BLE001 -- older torchao without the rowwise recipe
@ -1218,9 +1218,7 @@ _SPECS: dict[str, _FamilySpec] = {
# HF repos that gate access behind a license acceptance: training needs a token whose account
# accepted the license. Checked by name (no network) so a missing token fails fast with an
# actionable message instead of a confusing 401 mid-load.
_GATED_TRAIN_REPOS = frozenset(
{"black-forest-labs/flux.1-dev", "black-forest-labs/flux.2-dev"}
)
_GATED_TRAIN_REPOS = frozenset({"black-forest-labs/flux.1-dev", "black-forest-labs/flux.2-dev"})
def _assert_gated_access(base_model: str, hf_token: Optional[str]) -> None:
@ -1290,8 +1288,16 @@ def _load_pixel_tensor_planned(path, resolution, center_crop, u_left, u_top, fli
def _build_latent_cache(
spec, vae, image_paths, cfg, device, weight_dtype, on_event, check_stop,
pcache = None, plan = None,
spec,
vae,
image_paths,
cfg,
device,
weight_dtype,
on_event,
check_stop,
pcache = None,
plan = None,
):
"""Precompute the per-image latent posterior cache: for each planned crop/flip variant,
encode once and store the affine (A, B) pair on CPU (pinned when possible) in fp32. The
@ -1688,8 +1694,13 @@ def _train_dit(cfg, spec, pairs, rng, device, weight_dtype, on_event, _check_sto
pcache, image_paths, plan, to_encode, device
)
if caption_embeds is not None:
_emit(on_event, "preparing", stage = "cache_latents",
done = len(image_paths), total = len(image_paths))
_emit(
on_event,
"preparing",
stage = "cache_latents",
done = len(image_paths),
total = len(image_paths),
)
if caption_embeds is None:
# Phase 1 (cold): conditioning only. The pipeline loads WITHOUT its transformer, so
# the text encoders + VAE never share VRAM with the multi-GB denoiser. Caption
@ -1709,8 +1720,16 @@ def _train_dit(cfg, spec, pairs, rng, device, weight_dtype, on_event, _check_sto
# preserved).
if use_cache:
latent_cache = _build_latent_cache(
spec, vae, image_paths, cfg, device, weight_dtype, on_event, _check_stop,
pcache = pcache, plan = plan,
spec,
vae,
image_paths,
cfg,
device,
weight_dtype,
on_event,
_check_stop,
pcache = pcache,
plan = plan,
)
if latent_cache is LATENT_CACHE_OVER_BUDGET:
# The estimated cache exceeded the host-memory budget; keep the VAE resident
@ -1900,9 +1919,7 @@ def _train_dit(cfg, spec, pairs, rng, device, weight_dtype, on_event, _check_sto
else:
per = F.mse_loss(model_pred.float(), target.float(), reduction = "none")
w_idx = (
(sigmas.flatten().float() * num_train_ts)
.long()
.clamp(0, num_train_ts - 1)
(sigmas.flatten().float() * num_train_ts).long().clamp(0, num_train_ts - 1)
)
w = bell_weights[w_idx].view(-1, *([1] * (per.ndim - 1)))
loss = (per * w).mean()

15
studio/backend/core/training/diffusion_train_common.py
View file

@ -252,7 +252,10 @@ FAMILY_TRAIN_DEFAULTS: dict[str, dict[str, Any]] = {
# mid-schedule); the small warmups below are scaled for Studio's short-run step budgets.
"flux.1": {"lora_rank": 16, "learning_rate": 1e-4, "resolution": 512, "lr_warmup_steps": 20},
"qwen-image": {
"lora_rank": 16, "learning_rate": 5e-5, "resolution": 512, "lr_warmup_steps": 20,
"lora_rank": 16,
"learning_rate": 5e-5,
"resolution": 512,
"lr_warmup_steps": 20,
},
"z-image": {"lora_rank": 16, "learning_rate": 1e-4, "resolution": 768},
# The Krea 2 authors' recommended starting point (their DreamBooth script defaults):
@ -261,10 +264,16 @@ FAMILY_TRAIN_DEFAULTS: dict[str, dict[str, Any]] = {
# The upstream FLUX.2 DreamBooth references default to rank 16 / lr 1e-4; FLUX.2's
# uniform timestep draw benefits most from a warmup ramp.
"flux.2-klein": {
"lora_rank": 16, "learning_rate": 1e-4, "resolution": 512, "lr_warmup_steps": 20,
"lora_rank": 16,
"learning_rate": 1e-4,
"resolution": 512,
"lr_warmup_steps": 20,
},
"flux.2-dev": {
"lora_rank": 16, "learning_rate": 1e-4, "resolution": 512, "lr_warmup_steps": 20,
"lora_rank": 16,
"learning_rate": 1e-4,
"resolution": 512,
"lr_warmup_steps": 20,
},
}

12
studio/backend/core/training/diffusion_train_extras.py
View file

@ -61,7 +61,12 @@ class LoRAEMA:
shadow = decay * shadow + (1 - decay) * param
"""
def __init__(self, model: Any, decay: float = 0.99, warmup: bool = True):
def __init__(
self,
model: Any,
decay: float = 0.99,
warmup: bool = True,
):
if not 0.0 <= float(decay) < 1.0:
raise ValueError(f"ema decay must be in [0, 1), got {decay}")
self.decay = float(decay)
@ -115,7 +120,6 @@ class LoRAEMA:
def restore(self, model: Any, backup: dict[str, Any]) -> None:
import torch
with torch.no_grad():
for name, p in model.named_parameters():
if name in backup:
@ -194,8 +198,7 @@ class PersistentConditioningCache:
geom = f"{shape[0]}x{shape[1]}" if shape else str(self.resolution)
var = f"{u_left:.6f}_{u_top:.6f}_{int(bool(flip))}"
return (
f"lat_v{_CACHE_VERSION}_{self.family}_{geom}_"
f"{_file_content_hash(image_path)}_{var}"
f"lat_v{_CACHE_VERSION}_{self.family}_{geom}_" f"{_file_content_hash(image_path)}_{var}"
)
def text_key(self, caption: str) -> str:
@ -237,7 +240,6 @@ class PersistentConditioningCache:
return None
try:
from safetensors import safe_open
with safe_open(str(path), framework = "pt", device = "cpu") as f:
meta = f.metadata() or {}
count = int(meta.get("count", "0"))

7
studio/backend/models/inference.py
View file

@ -2427,14 +2427,17 @@ class DiffusionGenerateRequest(BaseModel):
@model_validator(mode = "after")
def _prompts_seeds_lengths_match(self) -> "DiffusionGenerateRequest":
if self.prompts is not None and self.seeds is not None and len(self.prompts) != len(
self.seeds
if (
self.prompts is not None
and self.seeds is not None
and len(self.prompts) != len(self.seeds)
):
raise ValueError(
f"prompts and seeds must have the same length (got {len(self.prompts)} "
f"prompts, {len(self.seeds)} seeds)"
)
return self
# Image-conditioned workflows (base64 or data-URL): init_image alone runs img2img,
# init_image + mask_image runs inpaint. Both require a family with the matching pipeline or
# the load is rejected. Cap each base64 string so one request can't buffer a multi-GB payload

2
studio/backend/tests/test_cached_gguf_routes.py
View file

@ -1490,4 +1490,4 @@ def test_list_cached_models_flags_single_file_diffusion_repos(monkeypatch, tmp_p
rows = {r["repo_id"]: r for r in result["cached"]}
assert rows["unsloth/Qwen-Image-fp8-single"].get("single_file") is True
assert "single_file" not in rows["unsloth/Qwen-Image-pipeline"]
assert "single_file" not in rows["Org/ChatRepo"]
assert "single_file" not in rows["Org/ChatRepo"]

83
studio/backend/tests/test_diffusion_backend.py
View file

@ -2673,7 +2673,9 @@ def test_dense_quant_skipped_when_dense_transformer_does_not_fit(
assert _FakeTransformer.last["path"] # GGUF path used
def test_dense_quant_prequant_proceeds_but_forbids_dense_fallback(fake_runtime, tmp_path, monkeypatch):
def test_dense_quant_prequant_proceeds_but_forbids_dense_fallback(
fake_runtime, tmp_path, monkeypatch
):
# With a prequant checkpoint, the fast path loads the small quantized file, so a dense
# misfit must NOT decline the fast path -- but the dense re-check still runs to gate the
# in-loader fallback: if the prequant later fails, the loader must raise to GGUF instead of
@ -2756,7 +2758,12 @@ def test_dense_quant_replan_retries_once_on_transient_free_undercount(
replan_calls = []
orig_plan = DiffusionBackend._plan_memory
def spy_plan(self, *a, transformer_resident_override_mib = None, **k):
def spy_plan(
self,
*a,
transformer_resident_override_mib = None,
**k,
):
real = orig_plan(
self, *a, transformer_resident_override_mib = transformer_resident_override_mib, **k
)
@ -2797,9 +2804,7 @@ def test_dense_quant_replan_retries_once_on_transient_free_undercount(
assert attempted == [False] # fast path attempted; prequant-sized plan forbids dense fallback
def test_dense_quant_replan_no_retry_when_capacity_truly_short(
fake_runtime, tmp_path, monkeypatch
):
def test_dense_quant_replan_no_retry_when_capacity_truly_short(fake_runtime, tmp_path, monkeypatch):
# When the candidate does NOT fit total capacity, the decline is real: no retry.
import dataclasses
@ -2821,7 +2826,12 @@ def test_dense_quant_replan_no_retry_when_capacity_truly_short(
replan_calls = []
orig_plan = DiffusionBackend._plan_memory
def spy_plan(self, *a, transformer_resident_override_mib = None, **k):
def spy_plan(
self,
*a,
transformer_resident_override_mib = None,
**k,
):
real = orig_plan(
self, *a, transformer_resident_override_mib = transformer_resident_override_mib, **k
)
@ -2869,7 +2879,12 @@ def _decline_dense_quant(backend, monkeypatch, tmp_path):
)
orig_plan = DiffusionBackend._plan_memory
def spy_plan(self, *a, transformer_resident_override_mib = None, **k):
def spy_plan(
self,
*a,
transformer_resident_override_mib = None,
**k,
):
real = orig_plan(
self, *a, transformer_resident_override_mib = transformer_resident_override_mib, **k
)
@ -2906,9 +2921,7 @@ def test_declined_dense_with_baked_loras_fails_instead_of_silent_drop(
)
def test_declined_dense_without_loras_still_falls_back_to_gguf(
fake_runtime, tmp_path, monkeypatch
):
def test_declined_dense_without_loras_still_falls_back_to_gguf(fake_runtime, tmp_path, monkeypatch):
# The plain decline (no adapters requested) keeps the silent GGUF fallback: weight-0
# adapters count as "none" (an explicit disable is not a bake request).
backend = DiffusionBackend()
@ -2928,10 +2941,18 @@ class _BakePipe:
def __init__(self):
self.calls: list = []
def load_lora_weights(self, path, adapter_name = None):
def load_lora_weights(
self,
path,
adapter_name = None,
):
self.calls.append(("load", path, adapter_name))
def set_adapters(self, names, adapter_weights = None):
def set_adapters(
self,
names,
adapter_weights = None,
):
self.calls.append(("set", tuple(names), tuple(adapter_weights)))
@ -2960,9 +2981,7 @@ def test_dense_quant_lora_bake_attaches_before_quantize(fake_runtime, monkeypatc
return object()
pipe = _BakePipe()
monkeypatch.setattr(
DiffusionBackend, "_assemble_pipe", staticmethod(lambda *a, **k: pipe)
)
monkeypatch.setattr(DiffusionBackend, "_assemble_pipe", staticmethod(lambda *a, **k: pipe))
monkeypatch.setattr(
DiffusionBackend,
"_resolve_lora_set",
@ -3013,9 +3032,7 @@ def test_apply_loras_quant_unbaked_requires_reload(monkeypatch):
backend = DiffusionBackend()
pipe = _BakePipe()
with pytest.raises(ValueError, match = "Reload the model with the adapter selection"):
backend._apply_loras(
_quant_lora_state(pipe), [("sloth", 1.0)], threading.Event()
)
backend._apply_loras(_quant_lora_state(pipe), [("sloth", 1.0)], threading.Event())
# ...but a no-adapter generation on the same pipe stays a plain no-op.
backend._apply_loras(_quant_lora_state(pipe), [], threading.Event())
assert pipe.calls == []
@ -3030,9 +3047,7 @@ def test_apply_loras_quant_baked_matrix(monkeypatch):
DiffusionBackend,
"_resolve_lora_set",
staticmethod(
lambda specs, **k: tuple(
(i, f"/adapters/{i}.safetensors", w) for (i, w) in specs
)
lambda specs, **k: tuple((i, f"/adapters/{i}.safetensors", w) for (i, w) in specs)
),
)
@ -3079,7 +3094,13 @@ def test_assemble_pipe_routes_krea2_per_component(monkeypatch):
calls["device"] = device
return self
def fake_loader(base, dtype, hf_token = None, transformer = None, text_encoder = None):
def fake_loader(
base,
dtype,
hf_token = None,
transformer = None,
text_encoder = None,
):
calls["base"] = base
calls["transformer"] = transformer
return Pipe()
@ -3632,9 +3653,7 @@ def test_generate_batched_seed_matches_solo_replay(fake_runtime, tmp_path):
def test_generate_prompt_list_rejected_off_txt2img(fake_runtime, tmp_path):
backend = _load_zimage_backend(tmp_path)
with pytest.raises(ValueError, match = "text-to-image only"):
backend.generate(
prompt = "x", prompts = ["a", "b"], init_image = _tiny_png_b64()
)
backend.generate(prompt = "x", prompts = ["a", "b"], init_image = _tiny_png_b64())
class _CountingPipe(_FakePipe):
@ -3645,7 +3664,12 @@ class _CountingPipe(_FakePipe):
self.batch_attempts = []
self.max_images = max_images
def __call__(self, *, prompt = None, **kwargs):
def __call__(
self,
*,
prompt = None,
**kwargs,
):
n = kwargs.get("num_images_per_prompt", 1)
if isinstance(prompt, list):
n *= len(prompt)
@ -3683,7 +3707,12 @@ def test_generate_oom_backoff_halves_the_batch(fake_runtime, tmp_path):
class _BoomPipe(_CountingPipe):
"""Fails every forward with a NON-OOM error (must not trigger backoff)."""
def __call__(self, *, prompt = None, **kwargs):
def __call__(
self,
*,
prompt = None,
**kwargs,
):
self.batch_attempts.append(kwargs.get("num_images_per_prompt", 1))
raise RuntimeError("shape mismatch")

7
studio/backend/tests/test_diffusion_cond_cache.py
View file

@ -135,7 +135,12 @@ def test_lora_attached_bypasses_the_cache(cache_env):
class _ListEncodePipe(_EncodePipe):
"""Returns per-prompt embedding LISTS like Z-Image's ``encode_prompt``."""
def encode_prompt(self, prompt, device = None, do_classifier_free_guidance = True):
def encode_prompt(
self,
prompt,
device = None,
do_classifier_free_guidance = True,
):
self.calls += 1
prompts = prompt if isinstance(prompt, list) else [prompt]
embeds = [torch.full((1, 4), float(sum(map(ord, p)))) for p in prompts]

11
studio/backend/tests/test_diffusion_dit_trainer.py
View file

@ -40,7 +40,12 @@ from core.training.diffusion_train_common import (
def test_specs_cover_the_dit_families():
assert set(_SPECS) == {
"flux.1", "qwen-image", "z-image", "krea-2", "flux.2-klein", "flux.2-dev"
"flux.1",
"qwen-image",
"z-image",
"krea-2",
"flux.2-klein",
"flux.2-dev",
}
# FLUX / Qwen share the added-kv attention target set; Z-Image and Krea 2 are
# single-stream.
@ -132,9 +137,7 @@ def test_flux2_rejects_fp16_before_loading():
).normalized()
assert ok.resolved_family == "flux.2-klein"
assert (
DiffusionLoraConfig(
base_model = "black-forest-labs/FLUX.2-dev", data_dir = "d", output_dir = "o"
)
DiffusionLoraConfig(base_model = "black-forest-labs/FLUX.2-dev", data_dir = "d", output_dir = "o")
.normalized()
.resolved_family
== "flux.2-dev"

10
studio/backend/tests/test_diffusion_dit_trainer_flow_shift.py
View file

@ -30,7 +30,6 @@ def _qwen_scheduler():
# use_dynamic_shifting is true, base_shift = max_shift = log 3 (constant inference mu),
# exponential time shift, terminal stretch to 0.02.
from diffusers import FlowMatchEulerDiscreteScheduler
return FlowMatchEulerDiscreteScheduler(
num_train_timesteps = 1000,
shift = 1.0,
@ -45,7 +44,6 @@ def _qwen_scheduler():
def _flux_static_scheduler():
# A static-shift scheduler (shift baked into sigmas at init, no dynamic shifting).
from diffusers import FlowMatchEulerDiscreteScheduler
return FlowMatchEulerDiscreteScheduler(num_train_timesteps = 1000, shift = 3.0)
@ -73,17 +71,13 @@ def test_flow_shift_explicit_values_and_validation():
assert cfg.flow_shift == 2.2
# String numerics from the Studio config path coerce; "auto" passes through.
assert (
DiffusionLoraConfig(
base_model = "b", data_dir = "d", output_dir = "o", flow_shift = "3.0"
)
DiffusionLoraConfig(base_model = "b", data_dir = "d", output_dir = "o", flow_shift = "3.0")
.normalized()
.flow_shift
== 3.0
)
assert (
DiffusionLoraConfig(
base_model = "b", data_dir = "d", output_dir = "o", flow_shift = "AUTO"
)
DiffusionLoraConfig(base_model = "b", data_dir = "d", output_dir = "o", flow_shift = "AUTO")
.normalized()
.flow_shift
== "auto"

1
studio/backend/tests/test_diffusion_lora.py
View file

@ -414,6 +414,7 @@ def test_diffusers_apply_rejects_unsupported_quant():
# selection. nvfp4/mxfp8 are never baked, so the same reload error is unreachable there
# via the API (supports_lora blocks the load), but the backend path is shared.
import threading
pipe = _FakePipe()
with pytest.raises(ValueError, match = "Reload the model with the adapter selection"):
_backend()._apply_loras(

10
studio/backend/tests/test_diffusion_memory.py
View file

@ -588,9 +588,7 @@ def test_settled_snapshot_stops_early_when_device_already_idle(monkeypatch):
def fake_snapshot(target):
calls.append(1)
return DeviceMemory(
"cuda", "cuda", "discrete_vram", free_mib = 170_000, total_mib = 183_359
)
return DeviceMemory("cuda", "cuda", "discrete_vram", free_mib = 170_000, total_mib = 183_359)
monkeypatch.setattr(dm, "snapshot_device_memory", fake_snapshot)
snap = dm.settled_snapshot_device_memory(_target(device = "cuda"), attempts = 3, delay_s = 0)
@ -619,7 +617,11 @@ def test_plan_fits_total_capacity():
# stem from the instantaneous free reading, so a settled retry is worthwhile.
from core.inference.diffusion_memory import plan_fits_total_capacity
def plan(required, total, kind = "discrete_vram"):
def plan(
required,
total,
kind = "discrete_vram",
):
return types.SimpleNamespace(
estimates = {"resident_required_mib": required},
device_memory = DeviceMemory("cuda", "cuda", kind, free_mib = 1, total_mib = total),

10
studio/backend/tests/test_diffusion_more_families.py
View file

@ -127,7 +127,6 @@ def test_lumina2_generation_defaults():
def test_lumina2_prequant_wiring():
# Hosted int8/fp8 checkpoints (gate-validated) serve the family default base.
from core.inference.diffusion_families import family_prequant_repo
fam = detect_family("Alpha-VLLM/Lumina-Image-2.0")
for scheme in ("int8", "fp8"):
assert family_prequant_repo(fam, scheme) == "unsloth/Lumina-Image-2.0-FP8"
@ -186,15 +185,15 @@ def test_hunyuanimage21_is_trusted_non_gguf():
def test_hunyuanimage21_generation_defaults():
# Card recipe: 50 steps; guidance feeds the call's distilled_guidance_scale (3.25
# default), while classifier-free guidance runs inside the repo's guider components.
assert default_generation_params(
"hunyuanvideo-community/HunyuanImage-2.1-Diffusers"
) == (50, 3.25)
assert default_generation_params("hunyuanvideo-community/HunyuanImage-2.1-Diffusers") == (
50,
3.25,
)
def test_hunyuanimage21_prequant_wiring():
# Hosted int8/fp8 checkpoints, verified bit-identical to on-the-fly quantize.
from core.inference.diffusion_families import family_prequant_repo
fam = detect_family("hunyuanvideo-community/HunyuanImage-2.1-Diffusers")
for scheme in ("int8", "fp8"):
assert family_prequant_repo(fam, scheme) == "unsloth/HunyuanImage-2.1-FP8"
@ -300,7 +299,6 @@ def test_hidream_prequant_wiring():
# Hosted int8/fp8 checkpoints (28/28 per-case gate pairs per scheme; int8 verified
# bit-identical to on-the-fly quantize) serve the family default base.
from core.inference.diffusion_families import family_prequant_repo
fam = detect_family("HiDream-ai/HiDream-I1-Full")
for scheme in ("int8", "fp8"):
assert family_prequant_repo(fam, scheme) == "unsloth/HiDream-I1-Full-FP8"

31
studio/backend/tests/test_diffusion_prequant.py
View file

@ -54,9 +54,12 @@ def test_resolve_family_repo_by_scheme():
def test_prequant_repo_filename_convention():
from core.inference.diffusion_prequant import prequant_repo_filename
assert prequant_repo_filename("unsloth/Z-Image-Turbo-FP8", "int8") == "Z-Image-Turbo-INT8.pt"
assert prequant_repo_filename("unsloth/Z-Image-Turbo-FP8", "fp8") == "Z-Image-Turbo-FP8.pt"
assert prequant_repo_filename("unsloth/Qwen-Image-2512-INT8", "int8") == "Qwen-Image-2512-INT8.pt"
assert (
prequant_repo_filename("unsloth/Qwen-Image-2512-INT8", "int8") == "Qwen-Image-2512-INT8.pt"
)
assert prequant_repo_filename("org/Some-Model-quantized", "fp8") == "Some-Model-FP8.pt"
assert prequant_repo_filename("org/PlainRepo", "int8") == "PlainRepo-INT8.pt"
@ -85,29 +88,21 @@ def test_resolve_variant_base_falls_back_to_default():
== "org/default-fp8"
)
# Scheme still has to match within the variant table.
assert (
resolve_prequant_source(fam, "int8", base_repo = "org/model-dev").location
== "org/dev-fp8"
)
assert resolve_prequant_source(fam, "int8", base_repo = "org/model-dev").location == "org/dev-fp8"
def test_flux1_variant_prequant_wiring():
# The real flux.1 entry serves schnell by default and dev / Krea-dev via variants.
from core.inference.diffusion_families import detect_family, family_prequant_repo
fam = detect_family("black-forest-labs/FLUX.1-schnell")
for scheme in ("int8", "fp8"):
assert family_prequant_repo(fam, scheme) == "unsloth/FLUX.1-schnell-FP8"
assert (
family_prequant_repo(
fam, scheme, base_repo = "black-forest-labs/FLUX.1-dev"
)
family_prequant_repo(fam, scheme, base_repo = "black-forest-labs/FLUX.1-dev")
== "unsloth/FLUX.1-dev-FP8"
)
assert (
family_prequant_repo(
fam, scheme, base_repo = "black-forest-labs/FLUX.1-Krea-dev"
)
family_prequant_repo(fam, scheme, base_repo = "black-forest-labs/FLUX.1-Krea-dev")
== "unsloth/FLUX.1-Krea-dev-FP8"
)
@ -540,7 +535,11 @@ def test_load_repo_source_falls_back_to_legacy_filename(monkeypatch, tmp_path):
errors.EntryNotFoundError = _NotFound
requested = []
def _dl(repo_id, filename, token = None):
def _dl(
repo_id,
filename,
token = None,
):
requested.append(filename)
if filename != "transformer_fp8.pt":
raise _NotFound(filename)
@ -553,8 +552,10 @@ def test_load_repo_source_falls_back_to_legacy_filename(monkeypatch, tmp_path):
monkeypatch.setitem(sys.modules, "huggingface_hub.errors", errors)
source = PrequantSource(
kind = "repo", location = "org/Z-Image-Turbo-FP8",
filename = "Z-Image-Turbo-FP8.pt", fallback_filename = "transformer_fp8.pt",
kind = "repo",
location = "org/Z-Image-Turbo-FP8",
filename = "Z-Image-Turbo-FP8.pt",
fallback_filename = "transformer_fp8.pt",
)
result = load_prequantized_transformer(
_FakeTransformer,

150
studio/backend/tests/test_diffusion_te_prequant.py
View file

@ -57,7 +57,9 @@ def test_family_repo_by_scheme_and_component():
assert family_te_prequant_repo(fam, "fp8", "text_encoder_3") is None
assert family_te_prequant_repo(fam, "int8", "text_encoder") is None
# A malformed entry is skipped, not fatal.
assert family_te_prequant_repo(_fam(te_prequant_repos = (("bad",),)), "fp8", "text_encoder") is None
assert (
family_te_prequant_repo(_fam(te_prequant_repos = (("bad",),)), "fp8", "text_encoder") is None
)
# Families without the field resolve to None (both dataclasses default it, but a fake
# or an older family object must not break).
assert family_te_prequant_repo(types.SimpleNamespace(name = "x"), "fp8", "text_encoder") is None
@ -80,7 +82,11 @@ def test_resolve_priority_and_scheme_gate():
# ── checkpoint validation ────────────────────────────────────────────────────
def _good_ckpt(scheme = "fp8", component = "text_encoder", base = "Lightricks/LTX-2"):
def _good_ckpt(
scheme = "fp8",
component = "text_encoder",
base = "Lightricks/LTX-2",
):
return {
"format": TE_PREQUANT_FORMAT,
"metadata": {
@ -155,18 +161,23 @@ def _target():
def test_pipe_kwargs_empty_when_mode_not_fp8(monkeypatch):
fam = _fam(te_prequant_repos = (("fp8", "text_encoder", "org/hosted"),))
for mode in (None, "", "off", "int8", "fp8_dynamic"):
assert te_prequant_pipe_kwargs(
fam, "Lightricks/LTX-2", te_quant_mode = mode, target = _target(), dtype = None
) == {}
assert (
te_prequant_pipe_kwargs(
fam, "Lightricks/LTX-2", te_quant_mode = mode, target = _target(), dtype = None
)
== {}
)
def test_pipe_kwargs_empty_without_hosted_entry(monkeypatch):
import core.inference.diffusion_precision as precision
monkeypatch.setattr(precision, "te_quant_supported", lambda target, mode: True)
assert te_prequant_pipe_kwargs(
_fam(), "Lightricks/LTX-2", te_quant_mode = "fp8", target = _target(), dtype = None
) == {}
assert (
te_prequant_pipe_kwargs(
_fam(), "Lightricks/LTX-2", te_quant_mode = "fp8", target = _target(), dtype = None
)
== {}
)
def test_pipe_kwargs_empty_when_device_unsupported(monkeypatch):
@ -174,9 +185,12 @@ def test_pipe_kwargs_empty_when_device_unsupported(monkeypatch):
fam = _fam(te_prequant_repos = (("fp8", "text_encoder", "org/hosted"),))
monkeypatch.setattr(precision, "te_quant_supported", lambda target, mode: False)
assert te_prequant_pipe_kwargs(
fam, "Lightricks/LTX-2", te_quant_mode = "fp8", target = _target(), dtype = None
) == {}
assert (
te_prequant_pipe_kwargs(
fam, "Lightricks/LTX-2", te_quant_mode = "fp8", target = _target(), dtype = None
)
== {}
)
def test_pipe_kwargs_respects_family_deny(monkeypatch):
@ -188,9 +202,12 @@ def test_pipe_kwargs_respects_family_deny(monkeypatch):
monkeypatch.setattr(
precision, "_te_family_denied", lambda family, mode: family == "ltx-2", raising = False
)
assert te_prequant_pipe_kwargs(
fam, "Lightricks/LTX-2", te_quant_mode = "fp8", target = _target(), dtype = None
) == {}
assert (
te_prequant_pipe_kwargs(
fam, "Lightricks/LTX-2", te_quant_mode = "fp8", target = _target(), dtype = None
)
== {}
)
def test_pipe_kwargs_injects_loaded_encoder(monkeypatch):
@ -220,9 +237,12 @@ def test_pipe_kwargs_empty_when_load_fails(monkeypatch):
fam = _fam(te_prequant_repos = (("fp8", "text_encoder", "org/hosted"),))
monkeypatch.setattr(precision, "te_quant_supported", lambda target, mode: True)
monkeypatch.setattr(tpq, "load_prequant_text_encoder", lambda *a, **k: None)
assert te_prequant_pipe_kwargs(
fam, "Lightricks/LTX-2", te_quant_mode = "fp8", target = _target(), dtype = None
) == {}
assert (
te_prequant_pipe_kwargs(
fam, "Lightricks/LTX-2", te_quant_mode = "fp8", target = _target(), dtype = None
)
== {}
)
def test_pipe_kwargs_injects_every_hosted_component(monkeypatch):
@ -259,9 +279,7 @@ def test_te_base_equivalent_groups():
assert te_base_equivalent(
"Qwen/Qwen-Image", "hunyuanvideo-community/HunyuanImage-2.1-Diffusers"
)
assert te_base_equivalent(
"black-forest-labs/FLUX.1-schnell", "black-forest-labs/FLUX.1-dev"
)
assert te_base_equivalent("black-forest-labs/FLUX.1-schnell", "black-forest-labs/FLUX.1-dev")
assert te_base_equivalent(
"black-forest-labs/FLUX.1-Krea-dev", "black-forest-labs/FLUX.1-schnell"
)
@ -281,8 +299,11 @@ def test_validate_accepts_equivalent_base():
},
}
assert tpq._validate_checkpoint(
ckpt, "fp8", "text_encoder",
"hunyuanvideo-community/HunyuanImage-2.1-Diffusers", None,
ckpt,
"fp8",
"text_encoder",
"hunyuanvideo-community/HunyuanImage-2.1-Diffusers",
None,
)
assert not tpq._validate_checkpoint(
ckpt, "fp8", "text_encoder", "black-forest-labs/FLUX.1-schnell", None
@ -317,31 +338,36 @@ def test_hosted_te_prequant_entries():
("fp8", "text_encoder", "unsloth/LTX-2-FP8"),
)
# The hosted filenames follow the repo naming convention the resolver derives.
assert te_prequant_repo_filename(
"unsloth/Qwen-Image-FP8", "text_encoder", "fp8"
) == "Qwen-Image-text_encoder-FP8.pt"
assert te_prequant_repo_filename(
"unsloth/FLUX.2-dev-FP8", "text_encoder", "fp8"
) == "FLUX.2-dev-text_encoder-FP8.pt"
assert te_prequant_repo_filename(
"unsloth/LTX-2-FP8", "text_encoder", "fp8"
) == "LTX-2-text_encoder-FP8.pt"
assert (
te_prequant_repo_filename("unsloth/Qwen-Image-FP8", "text_encoder", "fp8")
== "Qwen-Image-text_encoder-FP8.pt"
)
assert (
te_prequant_repo_filename("unsloth/FLUX.2-dev-FP8", "text_encoder", "fp8")
== "FLUX.2-dev-text_encoder-FP8.pt"
)
assert (
te_prequant_repo_filename("unsloth/LTX-2-FP8", "text_encoder", "fp8")
== "LTX-2-text_encoder-FP8.pt"
)
# HiDream's heavyweight is TE4 (Llama-3.1-8B), engaged via hidream_te4_kwargs because
# the generic quantize_text_encoders pass only covers text_encoder.._3.
assert detect_family("HiDream-ai/HiDream-I1-Full").te_prequant_repos == (
("fp8", "text_encoder_4", "unsloth/HiDream-I1-Full-FP8"),
)
assert te_prequant_repo_filename(
"unsloth/HiDream-I1-Full-FP8", "text_encoder_4", "fp8"
) == "HiDream-I1-Full-text_encoder_4-FP8.pt"
assert (
te_prequant_repo_filename("unsloth/HiDream-I1-Full-FP8", "text_encoder_4", "fp8")
== "HiDream-I1-Full-text_encoder_4-FP8.pt"
)
# Round 2: T5-XXL for every flux.1 base (byte-identical weights, one artifact),
# Gemma2-2B, Qwen3-4B, Qwen3-VL-4B, and hunyuanimage reusing the Qwen-Image artifact.
assert detect_family("black-forest-labs/FLUX.1-schnell").te_prequant_repos == (
("fp8", "text_encoder_2", "unsloth/FLUX.1-schnell-FP8"),
)
assert te_prequant_repo_filename(
"unsloth/FLUX.1-schnell-FP8", "text_encoder_2", "fp8"
) == "FLUX.1-schnell-text_encoder_2-FP8.pt"
assert (
te_prequant_repo_filename("unsloth/FLUX.1-schnell-FP8", "text_encoder_2", "fp8")
== "FLUX.1-schnell-text_encoder_2-FP8.pt"
)
assert detect_family("Alpha-VLLM/Lumina-Image-2.0").te_prequant_repos == (
("fp8", "text_encoder", "unsloth/Lumina-Image-2.0-FP8"),
)
@ -351,9 +377,9 @@ def test_hosted_te_prequant_entries():
assert detect_family("krea/Krea-2-Turbo").te_prequant_repos == (
("fp8", "text_encoder", "unsloth/Krea-2-Turbo-FP8"),
)
assert detect_family(
"hunyuanvideo-community/HunyuanImage-2.1-Diffusers"
).te_prequant_repos == (("fp8", "text_encoder", "unsloth/Qwen-Image-FP8"),)
assert detect_family("hunyuanvideo-community/HunyuanImage-2.1-Diffusers").te_prequant_repos == (
("fp8", "text_encoder", "unsloth/Qwen-Image-FP8"),
)
# flux.2-klein-4B hosts NO TE entry: its Qwen3-4B retrained layer 35's MLP, so the
# z-image artifact must not serve it (verified tensor diff, maxdiff 0.86).
assert detect_family("black-forest-labs/FLUX.2-klein-4B").te_prequant_repos == ()
@ -422,18 +448,13 @@ def test_hidream_te4_prefers_precast_checkpoint(monkeypatch):
te_prequant_repos = (("fp8", "text_encoder_4", "unsloth/HiDream-I1-Full-FP8"),),
name = "hidream-i1",
)
out = dh.hidream_te4_kwargs(
None, None, fam = fam, te_quant_mode = "fp8", target = _target()
)
out = dh.hidream_te4_kwargs(None, None, fam = fam, te_quant_mode = "fp8", target = _target())
assert out["text_encoder_4"] is precast
assert calls["base"] == "unsloth/Meta-Llama-3.1-8B-Instruct"
assert calls["component"] == "text_encoder_4"
# Standalone repo: config at the root, forward flags the pipeline needs applied.
assert calls["config_subfolder"] == ""
assert calls["config_overrides"] == {
"output_hidden_states": True,
"output_attentions": True,
}
assert calls["config_overrides"] == {"output_hidden_states": True, "output_attentions": True}
# The dense Llama download never ran.
assert ("llama_from_pretrained", "unsloth/Meta-Llama-3.1-8B-Instruct") not in recorder
@ -452,9 +473,7 @@ def test_hidream_te4_falls_back_to_dense_cast(monkeypatch):
te_prequant_repos = (("fp8", "text_encoder_4", "unsloth/HiDream-I1-Full-FP8"),),
name = "hidream-i1",
)
out = dh.hidream_te4_kwargs(
None, None, fam = fam, te_quant_mode = "fp8", target = _target()
)
out = dh.hidream_te4_kwargs(None, None, fam = fam, te_quant_mode = "fp8", target = _target())
assert cast == [out["text_encoder_4"]]
assert ("llama_from_pretrained", "unsloth/Meta-Llama-3.1-8B-Instruct") in recorder
@ -473,9 +492,7 @@ def test_hidream_te4_partial_cast_reloads_dense(monkeypatch):
monkeypatch.setattr(precision, "_cast_fp8", _boom)
fam = _fam(name = "hidream-i1") # no hosted entry -> dense + cast path
out = dh.hidream_te4_kwargs(
None, None, fam = fam, te_quant_mode = "fp8", target = _target()
)
out = dh.hidream_te4_kwargs(None, None, fam = fam, te_quant_mode = "fp8", target = _target())
dense_loads = [r for r in recorder if r[0] == "llama_from_pretrained"]
assert len(dense_loads) == 2 # initial load + the fail-safe reload
assert getattr(out["text_encoder_4"], "tag", "").startswith("dense")
@ -498,18 +515,31 @@ def test_assemble_pipe_injects_precast_te(monkeypatch):
seen.update(kw)
return FakePipe()
monkeypatch.setattr(
dif, "te_prequant_pipe_kwargs", lambda *a, **k: {"text_encoder": "PRECAST"}
)
monkeypatch.setattr(dif, "te_prequant_pipe_kwargs", lambda *a, **k: {"text_encoder": "PRECAST"})
dif.DiffusionBackend._assemble_pipe(
FakePipelineCls, "org/base", "TR", None, None, "cpu", None,
fam = None, te_quant_mode = "fp8", target = object(),
FakePipelineCls,
"org/base",
"TR",
None,
None,
"cpu",
None,
fam = None,
te_quant_mode = "fp8",
target = object(),
)
assert seen["text_encoder"] == "PRECAST"
seen.clear()
# No target (defensive default) keeps the assembly unchanged.
dif.DiffusionBackend._assemble_pipe(
FakePipelineCls, "org/base", "TR", None, None, "cpu", None, fam = None,
FakePipelineCls,
"org/base",
"TR",
None,
None,
"cpu",
None,
fam = None,
)
assert "text_encoder" not in seen

6
studio/backend/tests/test_diffusion_train_extras.py
View file

@ -95,7 +95,11 @@ def test_ema_rejects_bad_decay():
# ── persistent conditioning cache ─────────────────────────────────────────────
def _make_image(tmp_path, name = "a.png", color = (255, 0, 0)):
def _make_image(
tmp_path,
name = "a.png",
color = (255, 0, 0),
):
from PIL import Image
p = tmp_path / name

10
studio/backend/tests/test_video_backend.py
View file

@ -897,9 +897,7 @@ def test_ltx23_verbatim_sigmas_restores_scheduler_config():
class _Sched:
def __init__(self):
self.config = _Cfg(
use_dynamic_shifting = True, shift = 1.0, shift_terminal = 0.1
)
self.config = _Cfg(use_dynamic_shifting = True, shift = 1.0, shift_terminal = 0.1)
def register_to_config(self, **kw):
self.config.update(kw)
@ -1094,11 +1092,7 @@ def test_generate_progress_and_cancel_idle(fake_runtime):
backend = VideoBackend()
# Idle shape carries the image-endpoint-compatible aliases (total_steps / fraction)
# so one poller works against both generate-progress APIs.
assert backend.generate_progress() == {
"active": False,
"total_steps": 0,
"fraction": 0.0,
}
assert backend.generate_progress() == {"active": False, "total_steps": 0, "fraction": 0.0}
assert backend.cancel_generate() is False