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-06-25 03:21:52 +00:00
commit 0a2a423b08
8 changed files with 129 additions and 69 deletions
Download patch file Download diff file Expand all files Collapse all files

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

@ -134,7 +134,9 @@ class DiffusionBackend:
) -> dict[str, Any]:
"""Validate, then run the (slow) load on a daemon thread. Returns at once."""
if not gguf_filename:
raise ValueError("gguf_filename is required: this backend loads single-file GGUF checkpoints only.")
raise ValueError(
"gguf_filename is required: this backend loads single-file GGUF checkpoints only."
)
fam = detect_family(repo_id, family_override)
if fam is None:
raise ValueError(
@ -173,7 +175,9 @@ class DiffusionBackend:
# calls) so begin_load returns instantly; the bar shows raw bytes until
# the total lands. This is the only writer of _loading's fields here.
fam = detect_family(kwargs["repo_id"], kwargs.get("family_override"))
base = _resolve_base_repo(kwargs["repo_id"], kwargs.get("base_repo"), fam, kwargs.get("hf_token"))
base = _resolve_base_repo(
kwargs["repo_id"], kwargs.get("base_repo"), fam, kwargs.get("hf_token")
)
kwargs["base_repo"] = base
loading = self._loading
if loading is not None:
@ -270,7 +274,9 @@ class DiffusionBackend:
import diffusers
if not gguf_filename:
raise ValueError("gguf_filename is required: this backend loads single-file GGUF checkpoints only.")
raise ValueError(
"gguf_filename is required: this backend loads single-file GGUF checkpoints only."
)
fam = detect_family(repo_id, family_override)
if fam is None:
raise ValueError(
@ -343,7 +349,6 @@ class DiffusionBackend:
batch_size: int = 1,
) -> dict[str, Any]:
import torch
with self._lock:
state = self._state
if state is None:
@ -406,7 +411,13 @@ class DiffusionBackend:
"""Live per-step progress for an in-flight generation (lock-free read)."""
gen = self._gen
if gen is None or gen.total_steps <= 0:
return {"active": False, "step": 0, "total_steps": 0, "fraction": 0.0, "eta_seconds": None}
return {
"active": False,
"step": 0,
"total_steps": 0,
"fraction": 0.0,
"eta_seconds": None,
}
return {
"active": True,
"step": gen.step,
@ -474,7 +485,6 @@ def _hf_base_model(repo_id: str, hf_token: Optional[str]) -> Optional[str]:
return None
try:
from huggingface_hub import HfApi
meta = HfApi().model_info(repo_id, token = hf_token).cardData or {}
except Exception: # noqa: BLE001 — best-effort; fall back to the family default
return None

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

@ -52,7 +52,6 @@ def _evict_chat() -> None:
def _evict_diffusion() -> None:
from core.inference.diffusion import get_diffusion_backend
get_diffusion_backend().unload()

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

@ -143,7 +143,7 @@ def list_images(limit: Optional[int] = None, offset: int = 0) -> list[dict[str,
except OSError:
return []
paths.sort(key = _mtime, reverse = True)
window = paths[offset:] if limit is None else paths[offset:offset + limit]
window = paths[offset:] if limit is None else paths[offset : offset + limit]
records = []
for path in window:
meta = _read_meta(path)

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

@ -1690,7 +1690,9 @@ class DiffusionLoadRequest(BaseModel):
"""Request to load a local diffusion (text-to-image) checkpoint."""
model_path: str = Field(..., description = "Diffusion GGUF repo id or local path")
gguf_filename: str = Field(..., description = "The chosen single-file GGUF quant inside model_path")
gguf_filename: str = Field(
..., description = "The chosen single-file GGUF quant inside model_path"
)
base_repo: Optional[str] = Field(
None, description = "Companion diffusers repo for VAE/text-encoders (default: family base)"
)
@ -1705,9 +1707,13 @@ class DiffusionGenerateRequest(BaseModel):
"""Request to generate one image from the loaded diffusion model."""
prompt: str = Field(..., min_length = 1, description = "Text prompt")
negative_prompt: Optional[str] = Field(None, description = "What to avoid (if the model supports it)")
negative_prompt: Optional[str] = Field(
None, description = "What to avoid (if the model supports it)"
)
width: int = Field(1024, ge = 256, le = 2048, description = "Image width in pixels (multiple of 16)")
height: int = Field(1024, ge = 256, le = 2048, description = "Image height in pixels (multiple of 16)")
height: int = Field(
1024, ge = 256, le = 2048, description = "Image height in pixels (multiple of 16)"
)
steps: int = Field(9, ge = 1, le = 100, description = "Number of denoising steps")
guidance: float = Field(0.0, ge = 0.0, le = 20.0, description = "Classifier-free guidance scale")
seed: Optional[int] = Field(
@ -1748,9 +1754,7 @@ class GalleryImage(BaseModel):
class DiffusionGenerateResponse(BaseModel):
"""The persisted gallery records for one generation call (a batch)."""
images: list[GalleryImage] = Field(
..., description = "Saved records, one per image in the batch"
)
images: list[GalleryImage] = Field(..., description = "Saved records, one per image in the batch")
class GalleryListResponse(BaseModel):

52
studio/backend/routes/inference.py
View file

@ -2524,6 +2524,7 @@ async def load_model(
# ownership is asserted without depending on chat/diffusion exclusivity
# holding. No-op when diffusion isn't loaded.
from core.inference.gpu_arbiter import acquire_for, CHAT
await asyncio.to_thread(acquire_for, CHAT)
# ── Already-loaded check: skip reload if the exact model is active ──
@ -10054,8 +10055,7 @@ async def _openai_passthrough_non_streaming(
@studio_router.post("/images/load", response_model = DiffusionStatusResponse)
async def load_diffusion_model(
request: DiffusionLoadRequest,
current_subject: str = Depends(get_current_subject),
request: DiffusionLoadRequest, current_subject: str = Depends(get_current_subject)
):
from core.inference.diffusion import get_diffusion_backend
from core.inference.gpu_arbiter import acquire_for, DIFFUSION
@ -10085,8 +10085,7 @@ async def load_diffusion_model(
@studio_router.post("/images/generate", response_model = DiffusionGenerateResponse)
async def generate_diffusion_image(
request: DiffusionGenerateRequest,
current_subject: str = Depends(get_current_subject),
request: DiffusionGenerateRequest, current_subject: str = Depends(get_current_subject)
):
from core.inference import image_gallery
from core.inference.diffusion import get_diffusion_backend
@ -10119,20 +10118,25 @@ async def generate_diffusion_image(
def _persist() -> list[dict]:
records = []
for index, image in enumerate(result["images"]):
records.append(image_gallery.save(image, {
"prompt": request.prompt,
"negative_prompt": request.negative_prompt,
"width": request.width,
"height": request.height,
"steps": request.steps,
"guidance": request.guidance,
"seed": result["seed"],
# Position within the batch: images here share a seed + timestamp,
# so the export filename needs this to stay unique.
"batch_index": index,
"model": result.get("repo_id"),
"created_at": created_at,
}))
records.append(
image_gallery.save(
image,
{
"prompt": request.prompt,
"negative_prompt": request.negative_prompt,
"width": request.width,
"height": request.height,
"steps": request.steps,
"guidance": request.guidance,
"seed": result["seed"],
# Position within the batch: images here share a seed + timestamp,
# so the export filename needs this to stay unique.
"batch_index": index,
"model": result.get("repo_id"),
"created_at": created_at,
},
)
)
return records
try:
@ -10165,8 +10169,7 @@ async def list_gallery_images(
@studio_router.get("/images/gallery/{image_id}/file")
async def get_gallery_image_file(
image_id: str,
current_subject: str = Depends(get_current_subject),
image_id: str, current_subject: str = Depends(get_current_subject)
):
from core.inference import image_gallery
@ -10183,10 +10186,7 @@ async def get_gallery_image_file(
@studio_router.delete("/images/gallery/{image_id}")
async def delete_gallery_image(
image_id: str,
current_subject: str = Depends(get_current_subject),
):
async def delete_gallery_image(image_id: str, current_subject: str = Depends(get_current_subject)):
from core.inference import image_gallery
deleted = await asyncio.to_thread(image_gallery.delete, image_id)
@ -10198,7 +10198,6 @@ async def delete_gallery_image(
@studio_router.delete("/images/gallery")
async def clear_gallery_images(current_subject: str = Depends(get_current_subject)):
from core.inference import image_gallery
removed = await asyncio.to_thread(image_gallery.clear)
return {"removed": removed}
@ -10216,19 +10215,16 @@ async def unload_diffusion_model(current_subject: str = Depends(get_current_subj
@studio_router.get("/images/status", response_model = DiffusionStatusResponse)
async def diffusion_status(current_subject: str = Depends(get_current_subject)):
from core.inference.diffusion import get_diffusion_backend
return DiffusionStatusResponse(**get_diffusion_backend().status())
@studio_router.get("/images/load-progress", response_model = DiffusionLoadProgressResponse)
async def diffusion_load_progress(current_subject: str = Depends(get_current_subject)):
from core.inference.diffusion import get_diffusion_backend
return DiffusionLoadProgressResponse(**get_diffusion_backend().load_progress())
@studio_router.get("/images/generate-progress", response_model = DiffusionGenerateProgressResponse)
async def diffusion_generate_progress(current_subject: str = Depends(get_current_subject)):
from core.inference.diffusion import get_diffusion_backend
return DiffusionGenerateProgressResponse(**get_diffusion_backend().generate_progress())

26
studio/backend/routes/models.py
View file

@ -844,8 +844,7 @@ async def list_local_models(
models = [m for m in models if not _is_hidden_model(m.id, m.path)]
# Tag each GGUF with its task so the Images picker can filter to diffusion.
models = [
m.model_copy(update = {"task": _local_model_task(m.path, m.model_format)})
for m in models
m.model_copy(update = {"task": _local_model_task(m.path, m.model_format)}) for m in models
]
return LocalModelListResponse(
@ -3063,11 +3062,24 @@ def _repo_gguf_last_modified(repo_info) -> float:
# GGUF general.architecture values that denote a diffusion (image) model;
# everything else is treated as a text model. Lets the Images picker show only
# image GGUFs in its On Device list.
_DIFFUSION_GGUF_ARCHS = frozenset({
"flux", "flux2", "sd1", "sd2", "sd3", "sdxl", "stable_diffusion",
"lumina2", "qwen_image", "qwenimage", "auraflow", "pixart",
"hunyuan_video", "wan",
})
_DIFFUSION_GGUF_ARCHS = frozenset(
{
"flux",
"flux2",
"sd1",
"sd2",
"sd3",
"sdxl",
"stable_diffusion",
"lumina2",
"qwen_image",
"qwenimage",
"auraflow",
"pixart",
"hunyuan_video",
"wan",
}
)
def _gguf_architecture(path: str) -> Optional[str]:

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

@ -280,12 +280,21 @@ def test_resolve_base_repo_prefers_caller_then_hf_tag_then_fallback(monkeypatch)
fam = detect_family("unsloth/Qwen-Image-2512-GGUF")
monkeypatch.setattr(diffusion, "_hf_base_model", lambda repo, tok: "Qwen/Qwen-Image-2512")
# Caller's explicit base wins and the HF tag is not consulted.
assert diffusion._resolve_base_repo("unsloth/Qwen-Image-2512-GGUF", "my/base", fam, None) == "my/base"
assert (
diffusion._resolve_base_repo("unsloth/Qwen-Image-2512-GGUF", "my/base", fam, None)
== "my/base"
)
# No caller base: the repo's base_model tag (the variant base) is used.
assert diffusion._resolve_base_repo("unsloth/Qwen-Image-2512-GGUF", None, fam, None) == "Qwen/Qwen-Image-2512"
assert (
diffusion._resolve_base_repo("unsloth/Qwen-Image-2512-GGUF", None, fam, None)
== "Qwen/Qwen-Image-2512"
)
# No caller base and no tag: the family fallback.
monkeypatch.setattr(diffusion, "_hf_base_model", lambda repo, tok: None)
assert diffusion._resolve_base_repo("unsloth/Qwen-Image-2512-GGUF", " ", fam, None) == fam.base_repo
assert (
diffusion._resolve_base_repo("unsloth/Qwen-Image-2512-GGUF", " ", fam, None)
== fam.base_repo
)
def test_load_without_gguf_raises():
@ -340,7 +349,9 @@ def test_base_file_downloaded_excludes_undownloaded():
assert _base_file_downloaded("vae/diffusion_pytorch_model.safetensors")
# Excluded: the GGUF supplies the transformer; docs/assets and top-level files
# are never downloaded, so counting them would peg the bar short of 100%.
assert not _base_file_downloaded("transformer/diffusion_pytorch_model-00001-of-00003.safetensors")
assert not _base_file_downloaded(
"transformer/diffusion_pytorch_model-00001-of-00003.safetensors"
)
assert not _base_file_downloaded("assets/Z-Image-Gallery.pdf")
assert not _base_file_downloaded("README.md")
assert not _base_file_downloaded(".gitattributes")
@ -374,7 +385,10 @@ def test_generate_qwen_uses_true_cfg_scale(fake_runtime, tmp_path):
(tmp_path / "model.gguf").write_bytes(b"weights")
backend = DiffusionBackend()
backend.load_pipeline(
str(tmp_path), gguf_filename = "model.gguf", base_repo = "Qwen/Qwen-Image", family_override = "qwen-image"
str(tmp_path),
gguf_filename = "model.gguf",
base_repo = "Qwen/Qwen-Image",
family_override = "qwen-image",
)
backend.generate(prompt = "a sloth", guidance = 4.0)
# Qwen-Image's distilled guidance is off; the real CFG must land on true_cfg_scale.
@ -386,9 +400,13 @@ def test_begin_load_rejects_concurrent(monkeypatch):
backend = DiffusionBackend()
# The worker resolves the base via a network lookup; stub it so the test is offline.
monkeypatch.setattr("core.inference.diffusion._hf_base_model", lambda *a, **k: None)
monkeypatch.setattr(DiffusionBackend, "_estimate_download_bytes", staticmethod(lambda *a, **k: 0))
monkeypatch.setattr(
DiffusionBackend, "_estimate_download_bytes", staticmethod(lambda *a, **k: 0)
)
# Block the spawned worker so the load stays "in progress".
monkeypatch.setattr(DiffusionBackend, "load_pipeline", lambda self, **k: __import__("time").sleep(0.2))
monkeypatch.setattr(
DiffusionBackend, "load_pipeline", lambda self, **k: __import__("time").sleep(0.2)
)
backend.begin_load("unsloth/Z-Image-Turbo-GGUF", gguf_filename = "z-image-turbo-Q4_K_S.gguf")
with pytest.raises(RuntimeError):
backend.begin_load("unsloth/Z-Image-Turbo-GGUF", gguf_filename = "z-image-turbo-Q4_K_S.gguf")

47
studio/backend/tests/test_diffusion_routes.py
View file

@ -51,7 +51,13 @@ class _FakeBackend:
"error": None,
}
def generate(self, *, seed = None, batch_size = 1, **kwargs):
def generate(
self,
*,
seed = None,
batch_size = 1,
**kwargs,
):
if not self.loaded:
raise RuntimeError("No diffusion model is loaded.")
# The real backend returns the PIL images; the route persists them. The
@ -110,13 +116,15 @@ def client(monkeypatch, tmp_path):
monkeypatch.setattr(gallery_module, "save", _save)
monkeypatch.setattr(gallery_module, "image_b64", lambda i: "QUJD" if i in store else None)
def _list_images(limit = None, offset = 0):
ordered = sorted(store.values(), key = lambda r: r.get("created_at", 0.0), reverse = True)
return ordered[offset:] if limit is None else ordered[offset:offset + limit]
return ordered[offset:] if limit is None else ordered[offset : offset + limit]
monkeypatch.setattr(gallery_module, "list_images", _list_images)
monkeypatch.setattr(
gallery_module, "image_path",
gallery_module,
"image_path",
lambda i: (tmp_path / f"{i}.png") if i in store else None,
)
monkeypatch.setattr(gallery_module, "delete", lambda i: store.pop(i, None) is not None)
@ -129,11 +137,14 @@ def client(monkeypatch, tmp_path):
def test_load_generate_status_unload_roundtrip(client):
loaded = client.post("/api/inference/images/load", json = {
"model_path": "unsloth/Z-Image-Turbo-GGUF",
"gguf_filename": "z-image-turbo-Q4_K_S.gguf",
"base_repo": "base/repo",
})
loaded = client.post(
"/api/inference/images/load",
json = {
"model_path": "unsloth/Z-Image-Turbo-GGUF",
"gguf_filename": "z-image-turbo-Q4_K_S.gguf",
"base_repo": "base/repo",
},
)
assert loaded.status_code == 200
body = loaded.json()
assert body["loaded"] is True and body["family"] == "z-image"
@ -161,7 +172,9 @@ def test_load_generate_status_unload_roundtrip(client):
def test_generate_batch_size_persists_each_image(client):
client.post("/api/inference/images/load", json = {"model_path": "x/z-image", "gguf_filename": "q.gguf"})
client.post(
"/api/inference/images/load", json = {"model_path": "x/z-image", "gguf_filename": "q.gguf"}
)
resp = client.post(
"/api/inference/images/generate",
json = {"prompt": "p", "batch_size": 3, "seed": 5},
@ -175,7 +188,9 @@ def test_generate_batch_size_persists_each_image(client):
def test_gallery_pagination(client):
client.post("/api/inference/images/load", json = {"model_path": "x/z-image", "gguf_filename": "q.gguf"})
client.post(
"/api/inference/images/load", json = {"model_path": "x/z-image", "gguf_filename": "q.gguf"}
)
client.post("/api/inference/images/generate", json = {"prompt": "p", "batch_size": 5, "seed": 1})
page1 = client.get("/api/inference/images/gallery?limit=2&offset=0").json()
assert len(page1["images"]) == 2 and page1["has_more"] is True
@ -184,7 +199,9 @@ def test_gallery_pagination(client):
def test_generate_rejects_non_multiple_of_16(client):
client.post("/api/inference/images/load", json = {"model_path": "x/z-image", "gguf_filename": "q.gguf"})
client.post(
"/api/inference/images/load", json = {"model_path": "x/z-image", "gguf_filename": "q.gguf"}
)
# Odd, and a multiple of 8 that isn't a multiple of 16: both rejected, since
# Z-Image requires dimensions divisible by 16.
for bad in (1001, 1000):
@ -213,7 +230,9 @@ def test_load_unknown_family_returns_400(client, monkeypatch):
backend = _FakeBackend()
backend.begin_load = _raise
monkeypatch.setattr(diffusion_module, "get_diffusion_backend", lambda: backend)
resp = client.post("/api/inference/images/load", json = {"model_path": "x/y", "gguf_filename": "q.gguf"})
resp = client.post(
"/api/inference/images/load", json = {"model_path": "x/y", "gguf_filename": "q.gguf"}
)
assert resp.status_code == 400
assert "family" in resp.json()["detail"]
@ -223,7 +242,9 @@ def test_load_progress_route(client):
idle = client.get("/api/inference/images/load-progress")
assert idle.status_code == 200 and idle.json()["phase"] is None
# After load: the fake reports ready.
client.post("/api/inference/images/load", json = {"model_path": "x/z-image", "gguf_filename": "q.gguf"})
client.post(
"/api/inference/images/load", json = {"model_path": "x/z-image", "gguf_filename": "q.gguf"}
)
ready = client.get("/api/inference/images/load-progress")
assert ready.json()["phase"] == "ready"