diff --git a/studio/backend/core/inference/diffusion.py b/studio/backend/core/inference/diffusion.py index 4cfe310f7a..dfa0f1f114 100644 --- a/studio/backend/core/inference/diffusion.py +++ b/studio/backend/core/inference/diffusion.py @@ -35,6 +35,7 @@ from .diffusion_families import ( IDEOGRAM4_FAMILY_NAME, LUMINA2_FAMILY_NAME, DiffusionFamily, + assert_pipeline_class_available, default_generation_params, detect_family_for_pick, excluded_model_reason, @@ -727,6 +728,9 @@ class DiffusionBackend: f"pass family_override with that family name. (Video models and image models " f"whose diffusers transformer has no single-file loader are not supported.)" ) + # Refuse a too-old diffusers here rather than deep in the load, after the checkpoint has + # already been downloaded (Flux2Klein / Z-Image / Krea 2 / HunyuanImage are 0.39-only). + assert_pipeline_class_available(fam.pipeline_class, fam.name) # Families whose single file IS the whole pipeline have no GGUF path; reject before eviction. if kind == "gguf" and fam.single_file_is_pipeline: raise ValueError( @@ -3161,6 +3165,11 @@ class DiffusionBackend: # compile and the generate request then carries none, so a recipe built from # the request alone claimed no LoRA for an image that plainly used one. "active_loras": _active_lora_pairs(state.pipe), + # The workflow this generation ACTUALLY ran, as resolved above from which + # conditioning inputs were present. The recipe records it so a conditioned + # image is not presented as a plain Create recipe that would replay as + # something unrelated. + "workflow": workflow, } finally: # Deregister so a later unload/load can't poke a finished generation (if still ours). diff --git a/studio/backend/core/inference/diffusion_families.py b/studio/backend/core/inference/diffusion_families.py index 70f5247fbd..4212ea13cd 100644 --- a/studio/backend/core/inference/diffusion_families.py +++ b/studio/backend/core/inference/diffusion_families.py @@ -620,6 +620,27 @@ def family_prequant_repo( return None +def assert_pipeline_class_available(pipeline_class: str, family_name: str) -> None: + """Raise before any download when the installed diffusers has no ``pipeline_class``. + + The newer families (Flux2Klein, Z-Image, Krea 2, LTX-2, HunyuanImage) only exist from diffusers + 0.39, and the packaging leaves an older diffusers installable on Python 3.9 -- diffusers dropped + 3.9 in 0.38 and this project still supports it, so the 0.39 floor has to be conditional or the + whole extra becomes unresolvable. Without this check the getattr chain died with a bare + AttributeError deep in the load, after the checkpoint had already been fetched, which is an + expensive way to learn the environment is too old. Krea 2 already guarded itself this way; this + is the same check for every family, run from validation.""" + import diffusers + + if hasattr(diffusers, pipeline_class): + return + raise RuntimeError( + f"'{family_name}' needs diffusers >= 0.39.0 ({pipeline_class}); this environment has " + f"diffusers {getattr(diffusers, '__version__', 'unknown')}. Upgrade with: " + f"pip install -U diffusers (which needs Python >= 3.10; diffusers dropped 3.9 in 0.38)." + ) + + def family_gguf_loadable(fam: DiffusionFamily) -> bool: """True when a GGUF transformer can be assembled for this family. diff --git a/studio/backend/core/inference/video.py b/studio/backend/core/inference/video.py index 8475ced363..ee6cd6f7c2 100644 --- a/studio/backend/core/inference/video.py +++ b/studio/backend/core/inference/video.py @@ -394,6 +394,11 @@ class VideoBackend: f"{', '.join(supported_video_family_names())}. If this is a variant of one " f"of them, pass family_override with that family name." ) + # Refuse a too-old diffusers here rather than deep in the load, after the checkpoint has + # already been downloaded (LTX2Pipeline / HunyuanVideo15Pipeline are 0.39-only). + from .diffusion_families import assert_pipeline_class_available + + assert_pipeline_class_available(fam.pipeline_class, fam.name) if kind != "gguf" and not _is_trusted_video_repo(repo_id): raise ValueError( f"Non-GGUF video loads are limited to unsloth/* repos, the official " diff --git a/studio/backend/core/inference/video_gallery.py b/studio/backend/core/inference/video_gallery.py index de2be1ba90..6a57bc37f8 100644 --- a/studio/backend/core/inference/video_gallery.py +++ b/studio/backend/core/inference/video_gallery.py @@ -78,33 +78,61 @@ def video_path(video_id: str) -> Optional[Path]: return path if path.is_file() else None -def transcode(video_id: str, fmt: str) -> Optional[bytes]: - """Re-encode a stored MP4 for the Download menu: "webm" (VP9) or "gif". Returns the bytes, or - None when the id doesn't resolve. Raises RuntimeError on missing codec/deps (route 501s). MP4 - downloads stream the original via /file, not here.""" +def transcode_to_file(video_id: str, fmt: str) -> Optional[Path]: + """Re-encode a stored MP4 for the Download menu into a TEMP FILE and return its path, or None + when the id doesn't resolve. Raises RuntimeError on missing codec/deps (route 501s). The caller + owns the file and must delete it after serving. + + A file rather than a buffer because the request caps allow 2048x2048 x 1024 frames: a VP9 + export of a clip that size runs to hundreds of MB, and holding it as one ``bytes`` (then again + in the response) let a couple of concurrent export clicks exhaust the process. The MP4 route + already streams from disk; this makes the transcodes behave the same way.""" # Ownership-gate like /file: only transcode a Studio-owned clip (readable sidecar), so a guessed # stem for a foreign/orphan MP4 the gallery hides can't be re-encoded out either. path = owned_video_path(video_id) if path is None: return None normalized = fmt.strip().lower() - if normalized == "webm": - return _transcode_webm(path) - if normalized == "gif": - return _transcode_gif(path) - raise ValueError(f"Unsupported export format '{fmt}'. Use webm or gif.") + if normalized not in ("webm", "gif"): + raise ValueError(f"Unsupported export format '{fmt}'. Use webm or gif.") + import tempfile + + fd, tmp_name = tempfile.mkstemp(prefix = f"unsloth-export-{video_id}-", suffix = f".{normalized}") + os.close(fd) + dest = Path(tmp_name) + try: + if normalized == "webm": + _transcode_webm(path, dest) + else: + # GIF is already bounded by _GIF_MAX_FRAMES / _GIF_MAX_EDGE, so it is built in memory + # and written out; the cap is what keeps that safe. + dest.write_bytes(_transcode_gif(path)) + except BaseException: + dest.unlink(missing_ok = True) + raise + return dest -def _transcode_webm(path: Path) -> bytes: - import io +def transcode(video_id: str, fmt: str) -> Optional[bytes]: + """``transcode_to_file`` read back into memory. Kept for callers that want the bytes; the route + uses the file form so a large export is never fully resident.""" + dest = transcode_to_file(video_id, fmt) + if dest is None: + return None + try: + return dest.read_bytes() + finally: + dest.unlink(missing_ok = True) + +def _transcode_webm(path: Path, dest: Path) -> None: + """Transcode ``path`` to VP9 (+ Opus when the clip has audio) at ``dest``.""" try: import av except Exception as exc: # noqa: BLE001 -- no PyAV -> no transcode raise RuntimeError("WebM export needs the 'av' package (PyAV).") from exc - buf = io.BytesIO() try: - with av.open(str(path)) as src, av.open(buf, "w", format = "webm") as dst: + with av.open(str(path)) as src, av.open(str(dest), "w", format = "webm") as dst: if not src.streams.video: raise RuntimeError("WebM export failed: the clip has no video stream.") in_v = src.streams.video[0] @@ -171,7 +199,6 @@ def _transcode_webm(path: Path) -> bytes: raise except Exception as exc: # noqa: BLE001 -- surface as "encoder unavailable" raise RuntimeError(f"WebM export failed (libvpx-vp9 unavailable?): {exc}") from exc - return buf.getvalue() # Ceilings for a GIF export, which must hold every kept frame in memory before encoding. 720 px diff --git a/studio/backend/models/inference.py b/studio/backend/models/inference.py index 2ca11df59f..fe7aaa120e 100644 --- a/studio/backend/models/inference.py +++ b/studio/backend/models/inference.py @@ -2587,6 +2587,24 @@ class GalleryImage(BaseModel): controlnet: Optional[str] = Field( None, description = "ControlNet applied, formatted as 'id:control_type:strength'" ) + # Conditioned-workflow settings. The images themselves are NOT persisted (user uploads with + # their own lifetime), so these say what ran and let the client tell the user which inputs it + # needs back instead of silently restoring a conditioned image as a plain Create. + workflow: Optional[str] = Field( + None, + description = "Workflow that produced it: txt2img, img2img, inpaint, upscale, edit, " + "reference or controlnet. Absent on records written before this was recorded.", + ) + strength: Optional[float] = Field( + None, description = "img2img/inpaint denoise strength, when the workflow used one" + ) + upscale: Optional[float] = Field(None, description = "Upscale factor, for the upscale workflow") + controlnet_guidance: Optional[str] = Field( + None, description = "ControlNet guidance interval, formatted as 'start:end'" + ) + reference_image_count: Optional[int] = Field( + None, description = "How many reference images the reference workflow used" + ) created_at: float = Field(..., description = "Creation time (epoch seconds)") diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index fcd65cef61..16fe47b115 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -16390,6 +16390,22 @@ async def generate_diffusion_image( if request.controlnet and request.controlnet.strength > 0 else None ), + # The conditioned workflows (Transform/Inpaint/Extend/Upscale/Edit/reference/ + # ControlNet) keep their scalar settings here. The source, mask, reference and + # control IMAGES are deliberately not persisted -- they are user uploads with + # their own lifetime, and copying them into every recipe would grow the gallery + # without bound -- so a recipe records what it ran, and the client says plainly + # that the images have to be supplied again rather than silently replaying as + # a plain Create. + "workflow": result.get("workflow"), + "strength": request.strength, + "upscale": request.upscale, + "controlnet_guidance": ( + f"{request.controlnet.guidance_start:g}:{request.controlnet.guidance_end:g}" + if request.controlnet and request.controlnet.strength > 0 + else None + ), + "reference_image_count": len(request.reference_images or []) or None, "created_at": created_at, }, ) diff --git a/studio/backend/routes/video.py b/studio/backend/routes/video.py index 99ed2652ef..07777d7cee 100644 --- a/studio/backend/routes/video.py +++ b/studio/backend/routes/video.py @@ -356,18 +356,29 @@ async def export_gallery_video( if fmt not in ("webm", "gif"): raise HTTPException(status_code = 400, detail = "Unsupported format. Use webm or gif.") try: - data = await asyncio.to_thread(video_gallery.transcode, video_id, fmt) + path = await asyncio.to_thread(video_gallery.transcode_to_file, video_id, fmt) except RuntimeError as exc: raise HTTPException(status_code = 501, detail = str(exc)) from exc - if data is None: + if path is None: raise HTTPException(status_code = 404, detail = "Video not found.") - from fastapi.responses import Response + from fastapi.responses import FileResponse + from starlette.background import BackgroundTask - return Response( - content = data, + def _cleanup() -> None: + try: + path.unlink(missing_ok = True) + except OSError as e: # noqa: BLE001 -- a leaked temp file must not fail the download + logger.debug(f"Could not remove the export temp file {path}: {e}") + + # FileResponse streams from disk, so a large VP9 export is never fully resident (the caps allow + # 2048x2048 x 1024 frames). The temp file is deleted once the response has been sent. + return FileResponse( + path, media_type = "video/webm" if fmt == "webm" else "image/gif", + filename = f"{video_id}.{fmt}", # Transcodes are deterministic per id+format; let the browser cache them. headers = {"Cache-Control": "private, max-age=31536000, immutable"}, + background = BackgroundTask(_cleanup), ) diff --git a/studio/backend/tests/test_cached_gguf_routes.py b/studio/backend/tests/test_cached_gguf_routes.py index 0c0951ec79..57f8246a42 100644 --- a/studio/backend/tests/test_cached_gguf_routes.py +++ b/studio/backend/tests/test_cached_gguf_routes.py @@ -1778,3 +1778,36 @@ def test_hub_local_rows_are_tagged_with_their_task(): src = inspect.getsource(local_inventory.list_local_models_response) assert "_local_model_task" in src assert 'model_copy(update = {"task"' in src + + +def test_pipeline_class_guard_fires_before_any_download(): + # The 0.39-only families (Flux2Klein, Z-Image, Krea 2, LTX-2, HunyuanImage) used to die with a + # bare AttributeError deep in the load, after the checkpoint had been fetched, on the older + # diffusers that packaging still allows on Python 3.9. Validation refuses first, with the + # version and the fix in the message. + import pytest + + from core.inference.diffusion_families import _FAMILIES, assert_pipeline_class_available + + # Present -> no raise (every shipped family resolves on a current diffusers). + import diffusers + + for fam in _FAMILIES: + assert_pipeline_class_available(fam.pipeline_class, fam.name) + + stub = types.SimpleNamespace(__version__ = "0.37.0") + real = sys.modules.get("diffusers") + sys.modules["diffusers"] = stub + try: + with pytest.raises(RuntimeError) as excinfo: + assert_pipeline_class_available("ZImagePipeline", "z-image") + finally: + if real is not None: + sys.modules["diffusers"] = real + else: # pragma: no cover + del sys.modules["diffusers"] + msg = str(excinfo.value) + assert "z-image" in msg and "ZImagePipeline" in msg + assert "0.39" in msg and "0.37.0" in msg + assert "3.10" in msg # names the Python floor that carries a new enough diffusers + assert diffusers is not None diff --git a/studio/backend/tests/test_diffusion_routes.py b/studio/backend/tests/test_diffusion_routes.py index 863a978214..855a42496d 100644 --- a/studio/backend/tests/test_diffusion_routes.py +++ b/studio/backend/tests/test_diffusion_routes.py @@ -123,6 +123,12 @@ class _FakeBackend: "images": [object() for _ in range(batch_size)], "seed": seed if seed is not None else 4242, "repo_id": "x/z-image", + # The real backend reports the workflow it resolved; the recipe records it. + "workflow": ( + "inpaint" + if kwargs.get("mask_image") + else ("img2img" if kwargs.get("init_image") else "txt2img") + ), } def generate_progress(self): @@ -1093,3 +1099,41 @@ def test_load_refused_when_only_the_diffusion_probe_can_be_read(client, monkeypa json = {"model_path": "x/z-image", "gguf_filename": "q.gguf"}, ) assert resp.status_code == 200 + + +def test_recipe_records_the_conditioned_workflow_settings(client, monkeypatch): + # A conditioned generation's recipe used to carry only the txt2img fields, so the gallery + # presented an inpaint result as a complete Create recipe and restoring it replayed an + # unrelated text-to-image request. The images are still not persisted (user uploads with their + # own lifetime), but what ran IS, so the client can say which inputs to supply again. + import base64 + import io + + from PIL import Image + + client.post( + "/api/inference/images/load", json = {"model_path": "x/z-image", "gguf_filename": "q.gguf"} + ) + buf = io.BytesIO() + Image.new("RGB", (8, 8), (120, 30, 90)).save(buf, format = "PNG") + px = base64.b64encode(buf.getvalue()).decode() + gen = client.post( + "/api/inference/images/generate", + json = { + "prompt": "a sloth", + "seed": 7, + "init_image": px, + "mask_image": px, + "strength": 0.42, + }, + ) + assert gen.status_code == 200 + img = gen.json()["images"][0] + assert img["workflow"] == "inpaint" + assert img["strength"] == 0.42 + # A plain txt2img still records its own workflow and leaves the conditioning fields empty. + plain = client.post("/api/inference/images/generate", json = {"prompt": "a sloth", "seed": 7}) + assert plain.status_code == 200 + plain_img = plain.json()["images"][0] + assert plain_img["workflow"] == "txt2img" + assert plain_img["strength"] is None and plain_img["upscale"] is None diff --git a/studio/backend/tests/test_video_gallery.py b/studio/backend/tests/test_video_gallery.py index f18db6d9a8..85bbdc389b 100644 --- a/studio/backend/tests/test_video_gallery.py +++ b/studio/backend/tests/test_video_gallery.py @@ -470,6 +470,41 @@ def test_transcode_unknown_id_and_bad_format(): gallery.transcode(record["id"], "avi") +def test_transcode_to_file_writes_a_temp_file_the_caller_owns(): + # The route streams the export from disk instead of materialising it: the request caps allow + # 2048x2048 x 1024 frames, and a VP9 export of that size held as bytes (then again in the + # response) let concurrent exports exhaust the process. + record = gallery.save(_real_mp4_bytes(), _meta()) + for fmt, magic in (("webm", b"\x1a\x45\xdf\xa3"), ("gif", b"GIF8")): + path = gallery.transcode_to_file(record["id"], fmt) + assert path is not None and path.is_file(), fmt + assert path.suffix == f".{fmt}" + assert path.read_bytes()[: len(magic)] == magic + # It is a temp file, NOT something inside the gallery: deleting it must not touch the clip. + path.unlink() + assert gallery.video_path(record["id"]) is not None + assert gallery.transcode_to_file("does-not-exist", "webm") is None + + +def test_transcode_to_file_leaves_no_temp_file_when_the_encode_fails(monkeypatch): + # A half-written export must not accumulate in the temp dir on a host with no VP9 encoder. + import tempfile + + from core.inference import video_gallery as vg + + record = gallery.save(_real_mp4_bytes(), _meta()) + + def _boom(src, dest): + dest.write_bytes(b"partial") + raise RuntimeError("WebM export failed (libvpx-vp9 unavailable?)") + + monkeypatch.setattr(vg, "_transcode_webm", _boom) + before = set(Path(tempfile.gettempdir()).glob("unsloth-export-*")) + with pytest.raises(RuntimeError): + vg.transcode_to_file(record["id"], "webm") + assert set(Path(tempfile.gettempdir()).glob("unsloth-export-*")) == before + + def test_gif_export_bounds_frames_and_edge(monkeypatch): """Every kept frame is held as a paletted image before the encoder runs, so an unbounded walk is a memory bomb: the generate request allows 2048x2048 for 1024 frames, and at the 12 fps diff --git a/studio/backend/tests/test_video_routes.py b/studio/backend/tests/test_video_routes.py index 82cdea2a21..5bbd33c80c 100644 --- a/studio/backend/tests/test_video_routes.py +++ b/studio/backend/tests/test_video_routes.py @@ -774,7 +774,7 @@ def test_export_endpoint_validation(client, monkeypatch): def _boom(video_id, fmt): raise RuntimeError("WebM export needs the 'av' package (PyAV).") - monkeypatch.setattr(gallery_module, "transcode", _boom) + monkeypatch.setattr(gallery_module, "transcode_to_file", _boom) client.post( "/api/inference/video/load", json = {"model_path": "unsloth/LTX-2.3-GGUF", "gguf_filename": "q.gguf"}, diff --git a/studio/frontend/src/features/images/api.ts b/studio/frontend/src/features/images/api.ts index d4f414eb60..c25f8d89ff 100644 --- a/studio/frontend/src/features/images/api.ts +++ b/studio/frontend/src/features/images/api.ts @@ -175,6 +175,14 @@ export interface GalleryImage { model: string | null; loras?: string[]; controlnet?: string | null; + // Conditioned-workflow settings. The source/mask/reference/control images are not persisted, so + // these say what ran and let restore name the inputs the user has to supply again. Absent on + // records written before they were recorded. + workflow?: string | null; + strength?: number | null; + upscale?: number | null; + controlnet_guidance?: string | null; + reference_image_count?: number | null; created_at: number; } diff --git a/studio/frontend/src/features/images/images-page.tsx b/studio/frontend/src/features/images/images-page.tsx index dbf49dae22..79ff340a39 100644 --- a/studio/frontend/src/features/images/images-page.tsx +++ b/studio/frontend/src/features/images/images-page.tsx @@ -181,6 +181,20 @@ const WORKFLOW_TABS: Array<{ // Per-model generation defaults (steps + guidance), matched by repo-id substring, // most specific first. Distilled "turbo/schnell" models want few steps and little // guidance; the full "dev" models want more steps and real CFG. +// The images each conditioned workflow consumed, named for the restore toast. A recipe keeps the +// scalar settings but not the uploads themselves, so restoring one of these lands on Create and the +// user needs to know which input to add back rather than pressing Generate on a silently different +// request. Keys are the backend's own workflow strings (diffusion.py); txt2img is absent because it +// restores completely. +const CONDITIONED_WORKFLOW_INPUTS: Record = { + img2img: "the source image", + inpaint: "the source image and mask", + upscale: "the source image", + edit: "the source image", + reference: "the source and reference images", + controlnet: "the control image", +}; + // Generation defaults when the model is unrecognised: the distilled few-step / // no-CFG shape. Also seeds the sliders' initial state. const DEFAULT_GEN = { steps: 9, guidance: 0 }; @@ -1550,11 +1564,18 @@ export function ImagesPage({ active = true }: { active?: boolean }) { if (id && Number.isFinite(weight)) restoredLoras.push({ id, weight }); } setLoras(restoredLoras); + // The conditioned workflows' scalar settings ARE persisted, so restore them even though the + // form returns to Create: they are what the user re-applies after adding the image back. + if (typeof image.strength === "number") { + if (image.workflow === "upscale") setUpscaleStrength(image.strength); + else setStrength(image.strength); + } + if (typeof image.upscale === "number") setUpscaleFactor(image.upscale); // None of the conditioning images are persisted in a recipe, so a restore must not leave the // form pointing at whatever is currently loaded in the Transform / Inpaint / Edit tabs: the // next Generate would condition on an unrelated image and reproduce something else entirely. - // Clear them all and return to Create, the workflow the restored recipe actually describes - // (the workflow-validity effect snaps this back if the loaded model is edit-only). + // Clear them all and return to Create (the workflow-validity effect snaps this back if the + // loaded model is edit-only). setWorkflow("create"); setInitImage(null); setMaskImage(null); @@ -1562,7 +1583,15 @@ export function ImagesPage({ active = true }: { active?: boolean }) { // The control image isn't persisted, so clear any stale ControlNet selection. setControlnetId(""); setControlImage(null); - toast.success("Settings restored to inputs"); + // Say so, rather than letting a conditioned image restore as a plain Create that quietly + // generates something unrelated. The scalar settings ARE restored above; only the images the + // workflow consumed have to be supplied again. + const conditioned = CONDITIONED_WORKFLOW_INPUTS[image.workflow ?? ""]; + if (conditioned) { + toast.success(`Settings restored. Add ${conditioned} again to reproduce this image.`); + } else { + toast.success("Settings restored to inputs"); + } }, []); // A locked ratio keeps the paired dimension in step while dragging; "custom"