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-04 14:45:54 +00:00
commit 02b0a082d5
9 changed files with 148 additions and 90 deletions
Download patch file Download diff file Expand all files Collapse all files

24
studio/backend/core/inference/diffusion_ideogram4.py
View file

@ -104,6 +104,7 @@ def _patch_create_causal_mask() -> None:
pipe_mod.create_causal_mask = create_causal_mask_compat
_CAUSAL_MASK_PATCHED = True
# The fp8 attention is stored as a single fused ``qkv`` matrix with the Q, K and V
# rows stacked in that order; each block is ``hidden_size`` rows tall. hidden_size =
# attention_head_dim * num_attention_heads, read from the transformer config so a
@ -207,7 +208,9 @@ def _text_encoder_shard_paths(repo_id: str, token: Optional[str]) -> list[str]:
raise FileNotFoundError(f"no text_encoder safetensors under {sub}")
try:
index_path = hf_hub_download(repo_id, "text_encoder/model.safetensors.index.json", token = token)
index_path = hf_hub_download(
repo_id, "text_encoder/model.safetensors.index.json", token = token
)
weight_map = json.loads(Path(index_path).read_text())["weight_map"]
shards = sorted(set(weight_map.values()))
except Exception: # noqa: BLE001 -- single-file text encoder has no index
@ -243,7 +246,11 @@ def _text_encoder_is_fp8(repo_id: str, token: Optional[str]) -> bool:
return False
def load_ideogram4_text_encoder(repo_id: str, dtype, hf_token: Optional[str] = None):
def load_ideogram4_text_encoder(
repo_id: str,
dtype,
hf_token: Optional[str] = None,
):
"""The Qwen3-VL text encoder for ``repo_id``.
The ``-fp8`` repo stores this encoder in the SAME float8-plus-per-channel-scale
@ -302,7 +309,12 @@ def load_ideogram4_text_encoder(repo_id: str, dtype, hf_token: Optional[str] = N
return model
def load_ideogram4_transformer(repo_id: str, subfolder: str, dtype, hf_token: Optional[str] = None):
def load_ideogram4_transformer(
repo_id: str,
subfolder: str,
dtype,
hf_token: Optional[str] = None,
):
"""An ``Ideogram4Transformer2DModel`` for ``repo_id/subfolder`` (still on CPU).
Reads the transformer config, and if the shards carry the vendor fp8 layout
@ -356,7 +368,11 @@ def load_ideogram4_transformer(repo_id: str, subfolder: str, dtype, hf_token: Op
return model
def load_ideogram4_pipeline(repo_id: str, dtype, hf_token: Optional[str] = None):
def load_ideogram4_pipeline(
repo_id: str,
dtype,
hf_token: Optional[str] = None,
):
"""Assemble Ideogram4Pipeline from ``repo_id`` per-component (see module doc)."""
import diffusers

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

@ -270,10 +270,7 @@ def estimate_image_runtime_mib(
def estimate_video_runtime_mib(
*,
width: Optional[int],
height: Optional[int],
num_frames: Optional[int],
*, width: Optional[int], height: Optional[int], num_frames: Optional[int]
) -> int:
"""Per-call activation / latent / decode headroom for a video generation.

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

@ -290,7 +290,6 @@ class VideoBackend:
# cache, so a cancelled pull costs nothing).
if kwargs.get("gguf_filename") and not Path(kwargs["repo_id"]).expanduser().exists():
from utils.hf_xet_fallback import hf_hub_download_with_xet_fallback
hf_hub_download_with_xet_fallback(
kwargs["repo_id"],
kwargs["gguf_filename"],
@ -354,7 +353,6 @@ class VideoBackend:
return 0
try:
from huggingface_hub import scan_cache_dir
rid = repo_id.strip()
for repo in scan_cache_dir().repos:
if repo.repo_id == rid:
@ -436,9 +434,7 @@ class VideoBackend:
components = fam.bf16_components_gb
mib_per_gb = 1000.0**3 / (1024.0 * 1024.0)
if kind == "pipeline":
model_dense_mib = (
int(sum(components) * mib_per_gb) if components is not None else None
)
model_dense_mib = int(sum(components) * mib_per_gb) if components is not None else None
companion_mib = None
else:
checkpoint_path = self._resolve_checkpoint_path(repo_id, gguf_filename, hf_token)
@ -502,12 +498,8 @@ class VideoBackend:
hf_token = hf_token,
)
else:
transformer = transformer_cls.from_single_file(
str(checkpoint_path), **sf_kwargs
)
pipe = pipeline_cls.from_pretrained(
base, transformer = transformer, **pipe_kwargs
)
transformer = transformer_cls.from_single_file(str(checkpoint_path), **sf_kwargs)
pipe = pipeline_cls.from_pretrained(base, transformer = transformer, **pipe_kwargs)
if _load_token is not None and _load_token != self._load_token:
del pipe
@ -534,9 +526,7 @@ class VideoBackend:
speed_mode = effective_speed,
offload_active = plan.offload_policy != "none",
)
offload_policy, vae_tiling = apply_memory_plan(
pipe, plan, device = device, logger = logger
)
offload_policy, vae_tiling = apply_memory_plan(pipe, plan, device = device, logger = logger)
if not vae_tiling:
# Decode of a whole clip is the video memory peak; tiling is near-free
# in quality and keeps the decode bounded, so it is always on.
@ -603,7 +593,11 @@ class VideoBackend:
)
logger.info(
"video.loaded: %s (%s, %s, offload=%s, speed=%s)",
repo_id, fam.name, kind, offload_policy, effective_speed,
repo_id,
fam.name,
kind,
offload_policy,
effective_speed,
)
return self.status()
@ -621,9 +615,7 @@ class VideoBackend:
return root
from utils.hf_xet_fallback import hf_hub_download_with_xet_fallback
return Path(
hf_hub_download_with_xet_fallback(repo_id, gguf_filename or "", hf_token)
)
return Path(hf_hub_download_with_xet_fallback(repo_id, gguf_filename or "", hf_token))
# ── generation ───────────────────────────────────────────────────────────
@ -641,7 +633,6 @@ class VideoBackend:
seed: Optional[int] = None,
) -> dict[str, Any]:
import torch
cancel = threading.Event()
with self._generate_lock:
with self._lock:
@ -652,7 +643,8 @@ class VideoBackend:
try:
fam = state.family
width, height = snap_video_size(
fam, width or fam.resolution_presets[0][0],
fam,
width or fam.resolution_presets[0][0],
height or fam.resolution_presets[0][1],
)
frames = snap_num_frames(fam, num_frames or fam.default_num_frames)
@ -688,8 +680,13 @@ class VideoBackend:
started = time.monotonic()
self._gen = {
"active": True, "phase": "denoise", "step": 0, "total": steps,
"started": started, "eta_seconds": None, "error": None,
"active": True,
"phase": "denoise",
"step": 0,
"total": steps,
"started": started,
"eta_seconds": None,
"error": None,
}
def _on_step(p, step_index, timestep, callback_kwargs):

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

@ -270,12 +270,10 @@ def read_checkpoint_header(checkpoint_path: Path | str) -> dict[str, tuple[int,
path = str(checkpoint_path)
if path.lower().endswith(".gguf"):
from gguf import GGUFReader
for tensor in GGUFReader(path).tensors:
names_shapes[str(tensor.name)] = tuple(int(x) for x in tensor.shape)
else:
from safetensors import safe_open
with safe_open(path, framework = "pt") as handle:
for name in handle.keys():
names_shapes[name] = tuple(handle.get_slice(name).get_shape())
@ -344,16 +342,20 @@ def _split_checkpoint(state: dict[str, Any]) -> dict[str, dict[str, Any]]:
connector keys, nothing else).
"""
groups: dict[str, dict[str, Any]] = {
"dit": {}, "connectors": {}, "vae": {}, "audio_vae": {}, "vocoder": {}
"dit": {},
"connectors": {},
"vae": {},
"audio_vae": {},
"vocoder": {},
}
for key, value in state.items():
bare = key[len(_DIT_PREFIX):] if key.startswith(_DIT_PREFIX) else key
bare = key[len(_DIT_PREFIX) :] if key.startswith(_DIT_PREFIX) else key
if bare.startswith("vae."):
groups["vae"][bare[len("vae."):]] = value
groups["vae"][bare[len("vae.") :]] = value
elif bare.startswith("audio_vae."):
groups["audio_vae"][bare[len("audio_vae."):]] = value
groups["audio_vae"][bare[len("audio_vae.") :]] = value
elif bare.startswith("vocoder."):
groups["vocoder"][bare[len("vocoder."):]] = value
groups["vocoder"][bare[len("vocoder.") :]] = value
elif bare.startswith(_CONNECTOR_KEY_PREFIXES):
groups["connectors"][bare] = value
else:
@ -382,9 +384,14 @@ def checkpoint_variant(checkpoint_path: Path | str) -> str:
# ── component builders ───────────────────────────────────────────────────────
def _build_from_config(model_cls: Any, config: dict[str, Any], state: dict[str, Any],
rename: dict[str, str], torch_dtype: Any,
remove_suffixes: tuple[str, ...] = ()) -> Any:
def _build_from_config(
model_cls: Any,
config: dict[str, Any],
state: dict[str, Any],
rename: dict[str, str],
torch_dtype: Any,
remove_suffixes: tuple[str, ...] = (),
) -> Any:
from accelerate import init_empty_weights
state = _apply_rename(_to_plain_dtype(state, torch_dtype), rename)
@ -396,8 +403,14 @@ def _build_from_config(model_cls: Any, config: dict[str, Any], state: dict[str,
return model.to(torch_dtype)
def load_ltx23_transformer(dit_state: dict[str, Any], *, base_repo: str, torch_dtype: Any,
is_gguf: bool, hf_token: Optional[str]) -> Any:
def load_ltx23_transformer(
dit_state: dict[str, Any],
*,
base_repo: str,
torch_dtype: Any,
is_gguf: bool,
hf_token: Optional[str],
) -> Any:
import diffusers
from diffusers import LTX2VideoTransformer3DModel
@ -406,7 +419,7 @@ def load_ltx23_transformer(dit_state: dict[str, Any], *, base_repo: str, torch_d
# base repo's 2.0 transformer config and runs the stock 2.0 key conversion.
for old, new in _TRANSFORMER_PRERENAME:
for key in [k for k in dit_state if k.startswith(old)]:
dit_state[new + key[len(old):]] = dit_state.pop(key)
dit_state[new + key[len(old) :]] = dit_state.pop(key)
kwargs: dict[str, Any] = {
"config": base_repo,
"subfolder": "transformer",
@ -415,44 +428,54 @@ def load_ltx23_transformer(dit_state: dict[str, Any], *, base_repo: str, torch_d
**LTX_2_3_TRANSFORMER_CONFIG_OVERRIDES,
}
if is_gguf:
kwargs["quantization_config"] = diffusers.GGUFQuantizationConfig(
compute_dtype = torch_dtype
)
kwargs["quantization_config"] = diffusers.GGUFQuantizationConfig(compute_dtype = torch_dtype)
return LTX2VideoTransformer3DModel.from_single_file(dit_state, **kwargs)
def load_ltx23_connectors(connector_state: dict[str, Any], *, variant: str,
torch_dtype: Any, hf_token: Optional[str]) -> Any:
def load_ltx23_connectors(
connector_state: dict[str, Any], *, variant: str, torch_dtype: Any, hf_token: Optional[str]
) -> Any:
from diffusers.pipelines.ltx2.connectors import LTX2TextConnectors
# Transformer-only checkpoints carry the connector stacks but not the huge
# per-modality text projections; fetch those from the companion file.
if not any(k.startswith("text_embedding_projection") for k in connector_state):
connector_state = dict(connector_state)
connector_state.update(_load_extras_file(
_EXTRAS_TEXT_PROJ.format(variant = variant), hf_token
))
connector_state.update(
_load_extras_file(_EXTRAS_TEXT_PROJ.format(variant = variant), hf_token)
)
return _build_from_config(
LTX2TextConnectors, _CONNECTORS_CONFIG, connector_state, _CONNECTORS_RENAME,
LTX2TextConnectors,
_CONNECTORS_CONFIG,
connector_state,
_CONNECTORS_RENAME,
torch_dtype,
)
def load_ltx23_vae(vae_state: dict[str, Any], *, variant: str, torch_dtype: Any,
hf_token: Optional[str]) -> Any:
def load_ltx23_vae(
vae_state: dict[str, Any], *, variant: str, torch_dtype: Any, hf_token: Optional[str]
) -> Any:
from diffusers import AutoencoderKLLTX2Video
if not vae_state:
vae_state = _load_extras_file(_EXTRAS_VIDEO_VAE.format(variant = variant), hf_token)
return _build_from_config(
AutoencoderKLLTX2Video, _VIDEO_VAE_CONFIG, vae_state, _VIDEO_VAE_RENAME,
torch_dtype, remove_suffixes = _VIDEO_VAE_REMOVE_SUFFIXES,
AutoencoderKLLTX2Video,
_VIDEO_VAE_CONFIG,
vae_state,
_VIDEO_VAE_RENAME,
torch_dtype,
remove_suffixes = _VIDEO_VAE_REMOVE_SUFFIXES,
)
def load_ltx23_audio_vae_and_vocoder(
audio_vae_state: dict[str, Any], vocoder_state: dict[str, Any], *, variant: str,
torch_dtype: Any, hf_token: Optional[str],
audio_vae_state: dict[str, Any],
vocoder_state: dict[str, Any],
*,
variant: str,
torch_dtype: Any,
hf_token: Optional[str],
) -> tuple[Any, Any]:
from diffusers import AutoencoderKLLTX2Audio
from diffusers.pipelines.ltx2.vocoder import LTX2VocoderWithBWE
@ -460,13 +483,16 @@ def load_ltx23_audio_vae_and_vocoder(
if not audio_vae_state or not vocoder_state:
combined = _load_extras_file(_EXTRAS_AUDIO_VAE.format(variant = variant), hf_token)
audio_vae_state = {
k[len("audio_vae."):]: v for k, v in combined.items() if k.startswith("audio_vae.")
k[len("audio_vae.") :]: v for k, v in combined.items() if k.startswith("audio_vae.")
}
vocoder_state = {
k[len("vocoder."):]: v for k, v in combined.items() if k.startswith("vocoder.")
k[len("vocoder.") :]: v for k, v in combined.items() if k.startswith("vocoder.")
}
audio_vae = _build_from_config(
AutoencoderKLLTX2Audio, _AUDIO_VAE_CONFIG, audio_vae_state, _AUDIO_VAE_RENAME,
AutoencoderKLLTX2Audio,
_AUDIO_VAE_CONFIG,
audio_vae_state,
_AUDIO_VAE_RENAME,
torch_dtype,
)
# The 2.3 vocoder is a composite (base vocoder + bandwidth-extension stack +
@ -485,8 +511,14 @@ def load_ltx23_audio_vae_and_vocoder(
# ── pipeline assembly ────────────────────────────────────────────────────────
def load_ltx23_pipeline(checkpoint_path: Path | str, *, base_repo: str, torch_dtype: Any,
is_gguf: bool, hf_token: Optional[str] = None) -> Any:
def load_ltx23_pipeline(
checkpoint_path: Path | str,
*,
base_repo: str,
torch_dtype: Any,
is_gguf: bool,
hf_token: Optional[str] = None,
) -> Any:
"""Full LTX-2.3 pipeline from a single-file/GGUF checkpoint.
Assembled per-component (constructor, not from_pretrained) because the base
@ -500,7 +532,9 @@ def load_ltx23_pipeline(checkpoint_path: Path | str, *, base_repo: str, torch_dt
variant = checkpoint_variant(checkpoint_path)
logger.info(
"video.ltx23_assembly: variant=%s gguf=%s extras=%s",
variant, is_gguf, LTX23_EXTRAS_REPO,
variant,
is_gguf,
LTX23_EXTRAS_REPO,
)
state = load_single_file_checkpoint(str(checkpoint_path))
groups = _split_checkpoint(state)
@ -518,19 +552,25 @@ def load_ltx23_pipeline(checkpoint_path: Path | str, *, base_repo: str, torch_dt
)
transformer = load_ltx23_transformer(
groups["dit"], base_repo = base_repo, torch_dtype = torch_dtype,
is_gguf = is_gguf, hf_token = hf_token,
)
connectors = load_ltx23_connectors(
groups["connectors"], variant = variant, torch_dtype = torch_dtype,
groups["dit"],
base_repo = base_repo,
torch_dtype = torch_dtype,
is_gguf = is_gguf,
hf_token = hf_token,
)
vae = load_ltx23_vae(
groups["vae"], variant = variant, torch_dtype = torch_dtype, hf_token = hf_token
connectors = load_ltx23_connectors(
groups["connectors"],
variant = variant,
torch_dtype = torch_dtype,
hf_token = hf_token,
)
vae = load_ltx23_vae(groups["vae"], variant = variant, torch_dtype = torch_dtype, hf_token = hf_token)
audio_vae, vocoder = load_ltx23_audio_vae_and_vocoder(
groups["audio_vae"], groups["vocoder"], variant = variant,
torch_dtype = torch_dtype, hf_token = hf_token,
groups["audio_vae"],
groups["vocoder"],
variant = variant,
torch_dtype = torch_dtype,
hf_token = hf_token,
)
# Shared 2.0/2.3 components from the base repo, resolved through model_index

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

@ -2439,7 +2439,9 @@ class VideoGenerationDefaults(BaseModel):
guidance: float = Field(..., description = "Default guidance scale")
num_frames: int = Field(..., description = "Default frame count")
fps: int = Field(..., description = "Default playback frame rate")
frame_step: int = Field(..., description = "Temporal lattice: valid counts are k * frame_step + 1")
frame_step: int = Field(
..., description = "Temporal lattice: valid counts are k * frame_step + 1"
)
resolution_multiple: int = Field(..., description = "Width/height must be divisible by this")
resolution_presets: list[list[int]] = Field(
default_factory = list, description = "(width, height) presets the UI offers, default first"

6
studio/backend/routes/video.py
View file

@ -55,7 +55,6 @@ def _guard_video_load_against_training() -> None:
diffusion_active = False
try:
from core.training.diffusion_training_service import get_diffusion_training_service
diffusion_active = get_diffusion_training_service().is_active()
except Exception: # noqa: BLE001
diffusion_active = False
@ -134,7 +133,6 @@ async def load_video_model(
@router.get("/video/load-progress", response_model = VideoLoadProgressResponse)
async def video_load_progress(current_subject: str = Depends(get_current_subject)):
from core.inference.video import get_video_backend
return VideoLoadProgressResponse(**get_video_backend().load_progress())
@ -212,14 +210,12 @@ async def generate_video(
@router.get("/video/generate-progress", response_model = VideoGenerateProgressResponse)
async def video_generate_progress(current_subject: str = Depends(get_current_subject)):
from core.inference.video import get_video_backend
return VideoGenerateProgressResponse(**get_video_backend().generate_progress())
@router.post("/video/generate/cancel")
async def cancel_video_generation(current_subject: str = Depends(get_current_subject)):
from core.inference.video import get_video_backend
cancelled = await asyncio.to_thread(get_video_backend().cancel_generate)
return {"cancelled": cancelled}
@ -227,7 +223,6 @@ async def cancel_video_generation(current_subject: str = Depends(get_current_sub
@router.get("/video/status", response_model = VideoStatusResponse)
async def video_status(current_subject: str = Depends(get_current_subject)):
from core.inference.video import get_video_backend
return VideoStatusResponse(**get_video_backend().status())
@ -298,6 +293,5 @@ async def delete_gallery_video(video_id: str, current_subject: str = Depends(get
@router.delete("/video/gallery")
async def clear_gallery_videos(current_subject: str = Depends(get_current_subject)):
from core.inference import video_gallery
removed = await asyncio.to_thread(video_gallery.clear)
return {"removed": removed}

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

@ -85,7 +85,6 @@ def test_excluded_model_reason_none_for_supported_and_unknown():
def test_validate_load_request_surfaces_exclusion_reason():
from core.inference.diffusion import DiffusionBackend
backend = DiffusionBackend()
with pytest.raises(ValueError, match = "trust_remote_code"):
backend.validate_load_request("tencent/HunyuanImage-3.0")

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

@ -240,17 +240,21 @@ def test_is_ltx23_checkpoint_gguf(monkeypatch, tmp_path):
gguf = types.ModuleType("gguf")
# GGUF headers store dims in GGML (reversed) order.
gguf.GGUFReader = _reader_for({
"model.diffusion_model.transformer_blocks.0.scale_shift_table": (4096, 9),
})
gguf.GGUFReader = _reader_for(
{
"model.diffusion_model.transformer_blocks.0.scale_shift_table": (4096, 9),
}
)
monkeypatch.setitem(sys.modules, "gguf", gguf)
path = tmp_path / "ltx23.gguf"
path.write_bytes(b"x")
assert is_ltx23_checkpoint(path) is True
gguf.GGUFReader = _reader_for({
"model.diffusion_model.transformer_blocks.0.scale_shift_table": (4096, 6),
})
gguf.GGUFReader = _reader_for(
{
"model.diffusion_model.transformer_blocks.0.scale_shift_table": (4096, 6),
}
)
assert is_ltx23_checkpoint(path) is False
def _boom(path):

15
studio/backend/tests/test_video_routes.py
View file

@ -99,7 +99,8 @@ class _FakeBackend:
"base_repo": kwargs.get("base_repo") or "Lightricks/LTX-2",
"device": "cpu",
"dtype": "float32",
"model_kind": kwargs.get("model_kind") or ("gguf" if kwargs.get("gguf_filename") else "pipeline"),
"model_kind": kwargs.get("model_kind")
or ("gguf" if kwargs.get("gguf_filename") else "pipeline"),
"memory_mode": kwargs.get("memory_mode") or "auto",
"has_audio": True,
"defaults": _defaults(),
@ -113,7 +114,13 @@ class _FakeBackend:
"error": None,
}
def generate(self, *, prompt, seed = None, **kwargs):
def generate(
self,
*,
prompt,
seed = None,
**kwargs,
):
if not self.loaded:
raise RuntimeError(VIDEO_NOT_LOADED_MSG)
return {
@ -260,7 +267,9 @@ def test_generate_happy_path_persists_and_returns_record(client):
"/api/inference/video/load",
json = {"model_path": "unsloth/LTX-2.3-GGUF", "gguf_filename": "q.gguf"},
)
gen = client.post("/api/inference/video/generate", json = {"prompt": "a sloth surfing", "seed": 7})
gen = client.post(
"/api/inference/video/generate", json = {"prompt": "a sloth surfing", "seed": 7}
)
assert gen.status_code == 200
video = gen.json()["video"]
assert video["seed"] == 7 and video["prompt"] == "a sloth surfing" and video["id"]