From be5c91ebf2ef9d224615603e398d614b38970233 Mon Sep 17 00:00:00 2001 From: Irozuku Date: Wed, 1 Jul 2026 11:51:56 -0400 Subject: [PATCH 01/89] feat: add job progress reporting to queue infrastructure Add progress and progress_message columns to task_copy (with a migration for existing databases), a report_progress method on the queue that refreshes last_update so changes surface through polling, and a safe no-op report_progress helper on BaseJob for jobs to opt in. --- .../dependencies/job_queues/base_job_queue.py | 22 +++++++ .../dependencies/job_queues/huey_job_queue.py | 64 +++++++++++++++++-- DashAI/back/job/base_job.py | 34 +++++++++- 3 files changed, 114 insertions(+), 6 deletions(-) diff --git a/DashAI/back/dependencies/job_queues/base_job_queue.py b/DashAI/back/dependencies/job_queues/base_job_queue.py index 17c579ca4..b965affca 100644 --- a/DashAI/back/dependencies/job_queues/base_job_queue.py +++ b/DashAI/back/dependencies/job_queues/base_job_queue.py @@ -107,6 +107,28 @@ def to_list(self) -> List[BaseJob]: """ raise NotImplementedError + def report_progress( + self, + job_id: str, + progress: Optional[float], + message: Optional[str] = None, + ) -> None: + """Update the progress of a running job. + + The default implementation is a no-op so queues that do not track + progress remain valid. Concrete queues may override it. + + Parameters + ---------- + job_id: str + Identifier of the job to update. + progress: Optional float + Completion percentage in the range 0-100, or None when unknown. + message: Optional str + Short description of the current phase. + """ + return None + class JobQueueError(Exception): """Exception raised when a method of the job queue fails.""" diff --git a/DashAI/back/dependencies/job_queues/huey_job_queue.py b/DashAI/back/dependencies/job_queues/huey_job_queue.py index bfff16821..e4b18b96a 100644 --- a/DashAI/back/dependencies/job_queues/huey_job_queue.py +++ b/DashAI/back/dependencies/job_queues/huey_job_queue.py @@ -78,6 +78,7 @@ def __init__(self, queue_name: str, path_db: str): ) self._enable_wal() self._ensure_task_copy_table() + self._ensure_progress_columns() self._register_signals() @self.huey.task(context=True, priority=0) @@ -190,7 +191,7 @@ def on_start(signal, task): def on_success(signal, task, *args): exec_sql( ( - "UPDATE task_copy SET status = ?, " + "UPDATE task_copy SET status = ?, progress = 100, " f"last_update = {NOW_MICRO} " "WHERE id = ?" ), @@ -229,6 +230,8 @@ def _ensure_task_copy_table(self): 'deleted', 'error' - last_update (DATETIME NOT NULL): defaults to CURRENT_TIMESTAMP (UTC) - error_msg (TEXT): optional error message when a task fails + - progress (REAL): optional completion percentage in the range 0-100 + - progress_message (TEXT): optional short description of the current phase """ with sqlite3.connect(self.db_path) as conn: conn.execute( @@ -240,7 +243,9 @@ def _ensure_task_copy_table(self): enqueued_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP, status TEXT NOT NULL, last_update DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP, - error_msg TEXT + error_msg TEXT, + progress REAL, + progress_message TEXT ) """ ) @@ -251,12 +256,29 @@ def _ensure_task_copy_table(self): ) ) + def _ensure_progress_columns(self): + """Add the progress columns to an existing 'task_copy' table. + + Installs created before progress tracking existed have a 'task_copy' + table without the 'progress' and 'progress_message' columns. SQLite has + no 'ADD COLUMN IF NOT EXISTS', so inspect the current columns via + PRAGMA table_info and add only the ones that are missing. + """ + with sqlite3.connect(self.db_path) as conn: + cur = conn.execute("PRAGMA table_info(task_copy)") + existing = {row[1] for row in cur.fetchall()} + if "progress" not in existing: + conn.execute("ALTER TABLE task_copy ADD COLUMN progress REAL") + if "progress_message" not in existing: + conn.execute("ALTER TABLE task_copy ADD COLUMN progress_message TEXT") + def status(self, job_id: str) -> dict: conn = sqlite3.connect(self.db_path) cur = conn.cursor() cur.execute( """ - SELECT status, last_update, error_msg, job_name + SELECT status, last_update, error_msg, job_name, progress, + progress_message FROM task_copy WHERE id = ? """, (str(job_id),), @@ -270,8 +292,40 @@ def status(self, job_id: str) -> dict: "updated": row[1], "error": row[2], "job_name": row[3], + "progress": row[4], + "progress_message": row[5], } + def report_progress( + self, job_id: str, progress: float | None, message: str | None = None + ) -> None: + """Update the progress of a running job. + + Parameters + ---------- + job_id : str + The UUID of the job (its Huey task id). + progress : float or None + Completion percentage in the range 0-100, or None for jobs whose + total work is unknown (the frontend renders an indeterminate bar). + message : str or None + Optional short description of the current phase. + + Notes + ----- + This also refreshes 'last_update' so the change surfaces through + 'changes_since' and the frontend polling channel. + """ + with sqlite3.connect(self.db_path) as conn: + conn.execute( + ( + "UPDATE task_copy SET progress = ?, progress_message = ?, " + "last_update = STRFTIME('%Y-%m-%d %H:%M:%f','now') " + "WHERE id = ?" + ), + (progress, message, str(job_id)), + ) + def put(self, job: BaseJob) -> int: result = self._execute(job) @@ -284,7 +338,7 @@ def to_list(self) -> list[dict]: cur.execute( """ SELECT id, task_type, job_name, enqueued_at, status, last_update, - error_msg + error_msg, progress, progress_message FROM task_copy ORDER BY last_update DESC """ @@ -304,7 +358,7 @@ def changes_since(self, since: str) -> list[dict]: cur.execute( """ SELECT id, task_type, job_name, enqueued_at, status, last_update, - error_msg + error_msg, progress, progress_message FROM task_copy WHERE last_update >= ? ORDER BY last_update DESC diff --git a/DashAI/back/job/base_job.py b/DashAI/back/job/base_job.py index f1fcaaf52..53a521db4 100644 --- a/DashAI/back/job/base_job.py +++ b/DashAI/back/job/base_job.py @@ -1,7 +1,10 @@ """Base Job abstract class.""" +import logging from abc import ABCMeta, abstractmethod -from typing import Final +from typing import Final, Optional + +logger = logging.getLogger(__name__) class BaseJob(metaclass=ABCMeta): @@ -20,6 +23,35 @@ def __init__(self, **kwargs): job_kwargs = kwargs.pop("kwargs", {}) self.kwargs = {**kwargs, **job_kwargs} + def report_progress( + self, fraction: Optional[float], message: Optional[str] = None + ) -> None: + """Report the job's progress to the job queue. + + Jobs opt in by calling this at meaningful checkpoints. It is safe to + call from any job: it never raises and does nothing when the job has no + Huey id (e.g. immediate mode used in tests) or the queue is unavailable. + + Parameters + ---------- + fraction: Optional float + Completion in the range 0-1, or None when the total work is unknown + (the frontend then shows an indeterminate bar). + message: Optional str + Short description of the current phase. + """ + try: + from kink import di + + huey_id = self.kwargs.get("huey_id") + if not huey_id: + return + + progress = None if fraction is None else max(0.0, min(1.0, fraction)) * 100 + di["job_queue"].report_progress(huey_id, progress, message) + except Exception as e: # pragma: no cover - progress must never break a job + logger.debug(f"Could not report job progress: {e}") + @abstractmethod def set_status_as_delivered(self) -> None: """Set the status of the job as delivered.""" From 33bb5b7fad567f98fdfb554c5f1cd7f6b2dcec94 Mon Sep 17 00:00:00 2001 From: Irozuku Date: Wed, 1 Jul 2026 11:52:22 -0400 Subject: [PATCH 02/89] feat: report progress from training, dataset, prediction and converter jobs Emit phase checkpoints from ModelJob, DatasetJob, PredictJob and ConverterJob so the queue exposes a real completion signal. The converter loop reports determinate progress across its converters. --- DashAI/back/job/converter_job.py | 11 ++++++++++- DashAI/back/job/dataset_job.py | 6 ++++++ DashAI/back/job/model_job.py | 4 ++++ DashAI/back/job/predict_job.py | 6 ++++++ 4 files changed, 26 insertions(+), 1 deletion(-) diff --git a/DashAI/back/job/converter_job.py b/DashAI/back/job/converter_job.py index 97cb18c06..2d85d5423 100644 --- a/DashAI/back/job/converter_job.py +++ b/DashAI/back/job/converter_job.py @@ -233,6 +233,8 @@ def instantiate_converters( log.exception(e) raise JobError("Error loading converter info") from e + self.report_progress(0.1, "Loading dataset") + # Get dataset try: dataset_id = converter.notebook.dataset_id @@ -309,11 +311,17 @@ def instantiate_converters( i += 1 # Apply each converter in sequence - for converter_info in converter_instances: + total_converters = len(converter_instances) + for converter_index, converter_info in enumerate(converter_instances): converter_instance = converter_info["instance"] converter_name = converter_info["name"] converter_scope = converter_info["scope"] + # Map converter progress onto the 0.2-0.9 band. + self.report_progress( + 0.2 + 0.7 * (converter_index / max(total_converters, 1)), + f"Applying {converter_name}", + ) log.info(f"Applying converter: {converter_name}") columns_scope = [ @@ -394,6 +402,7 @@ def instantiate_converters( dataset_original_columns = loaded_dataset.column_names + self.report_progress(0.95, "Saving dataset") save_dataset(loaded_dataset, f"{dataset_path}") converter.set_status_as_finished() db.commit() diff --git a/DashAI/back/job/dataset_job.py b/DashAI/back/job/dataset_job.py index 6d9fc925e..4193dc3a3 100644 --- a/DashAI/back/job/dataset_job.py +++ b/DashAI/back/job/dataset_job.py @@ -126,6 +126,8 @@ def run( db.commit() db.refresh(dataset) + self.report_progress(0.1, "Loading data") + if n_sample and dataset.file_path != "": folder_path = Path(dataset.file_path) else: @@ -290,6 +292,8 @@ def run( new_dataset = transform_dataset_with_schema(new_dataset, schema) + self.report_progress(0.5, "Computing metadata") + compute_meta = params.get("compute_metadata", True) extended_keys = ( "general_info", @@ -327,6 +331,8 @@ def run( new_dataset.splits.pop(stale_key, None) gc.collect() + self.report_progress(0.8, "Saving dataset") + dataset_save_path = folder_path / "dataset" log.debug("Saving dataset in %s", str(dataset_save_path)) save_dataset(new_dataset, dataset_save_path) diff --git a/DashAI/back/job/model_job.py b/DashAI/back/job/model_job.py index 583c6b5fb..1109696c9 100644 --- a/DashAI/back/job/model_job.py +++ b/DashAI/back/job/model_job.py @@ -115,6 +115,7 @@ def run( run: Run = db.get(Run, run_id) run.huey_id = self.kwargs.get("huey_id", None) db.commit() + self.report_progress(0.05, "Preparing data") try: # Get the model session, dataset, task, metrics and splits model_session: ModelSession = db.get(ModelSession, run.model_session_id) @@ -271,6 +272,7 @@ def run( raise JobError( "Connection with the database failed", ) from e + self.report_progress(0.2, "Training") try: # Hyperparameter Tunning plot_paths = [] @@ -332,6 +334,7 @@ def run( f"Hyperparameter plot path saving failed {e}", ) from e + self.report_progress(0.85, "Computing metrics") # Calculate metrics at the end of training if not done already try: last_train_metric = ( @@ -370,6 +373,7 @@ def run( f"Metric calculation failed {e}", ) from e + self.report_progress(0.95, "Saving model") try: run_path = os.path.join(config["RUNS_PATH"], str(run.id)) model.save(run_path) diff --git a/DashAI/back/job/predict_job.py b/DashAI/back/job/predict_job.py index efbdee0ca..64683090e 100644 --- a/DashAI/back/job/predict_job.py +++ b/DashAI/back/job/predict_job.py @@ -305,6 +305,8 @@ def run( prediction.set_status_as_started() db.commit() + self.report_progress(0.1, "Loading model") + dataset_id = prediction.dataset_id # Validate input data @@ -422,6 +424,8 @@ def run( manual_input_data, dataset_trained_path ) + self.report_progress(0.4, "Running prediction") + prepared_dataset, y_pred = _run_prediction_pipeline( task=task, trained_model=trained_model, @@ -452,6 +456,8 @@ def run( "Model prediction failed", ) from e + self.report_progress(0.9, "Saving predictions") + # Save Predictions to Arrow file try: # Create unique folder for predictions From 0774f8fa8d0165efd1f79176568a7d277cd64782 Mon Sep 17 00:00:00 2001 From: Irozuku Date: Wed, 1 Jul 2026 11:52:43 -0400 Subject: [PATCH 03/89] feat: show per-job progress bar and phase in the job queue widget Render a LinearProgress bar (determinate when a percent is known, indeterminate otherwise) plus the current phase message for running jobs in the queue widget and the job details dialog. Add the progress label to the five locale files. --- .../src/components/jobs/JobDetailsDialog.jsx | 39 +++++++++++ .../src/components/jobs/JobQueueWidget.jsx | 66 ++++++++++++++----- .../src/utils/i18n/locales/de/common.json | 1 + .../src/utils/i18n/locales/en/common.json | 1 + .../src/utils/i18n/locales/es/common.json | 1 + .../src/utils/i18n/locales/pt/common.json | 1 + .../src/utils/i18n/locales/zh/common.json | 1 + 7 files changed, 92 insertions(+), 18 deletions(-) diff --git a/DashAI/front/src/components/jobs/JobDetailsDialog.jsx b/DashAI/front/src/components/jobs/JobDetailsDialog.jsx index 747af922a..bcecc0588 100644 --- a/DashAI/front/src/components/jobs/JobDetailsDialog.jsx +++ b/DashAI/front/src/components/jobs/JobDetailsDialog.jsx @@ -12,6 +12,7 @@ import { Box, Divider, CircularProgress, + LinearProgress, } from "@mui/material"; import { useTranslation } from "react-i18next"; import { getJobDetails } from "../../api/job"; @@ -189,6 +190,44 @@ const JobDetailsDialog = ({ job, open, onClose }) => { + {displayJob.status === "started" && ( + + + {t("common:jobQueue.details.progress")} + + {displayJob.progress_message && ( + + {displayJob.progress_message} + + )} + + + + + {displayJob.progress != null && ( + + {Math.round(displayJob.progress)}% + + )} + + + )} + {displayJob.error_msg && ( diff --git a/DashAI/front/src/components/jobs/JobQueueWidget.jsx b/DashAI/front/src/components/jobs/JobQueueWidget.jsx index caf892e30..ab9c03f2e 100644 --- a/DashAI/front/src/components/jobs/JobQueueWidget.jsx +++ b/DashAI/front/src/components/jobs/JobQueueWidget.jsx @@ -12,6 +12,7 @@ import { Badge, Collapse, List, + LinearProgress, CircularProgress, ListItem, ListItemButton, @@ -517,6 +518,7 @@ const JobQueueWidget = () => { @@ -531,26 +533,54 @@ const JobQueueWidget = () => { } secondary={ - - {getStatusText(job.status, t)} - {job.status === "error" && ( - + + + {job.status === "started" && + job.progress_message + ? job.progress_message + : getStatusText(job.status, t)} + + {job.status === "error" && ( + + + + )} + + {job.status === "started" && ( + - - + sx={{ + mt: 0.5, + height: 4, + borderRadius: 1, + }} + /> )} - + } /> diff --git a/DashAI/front/src/utils/i18n/locales/de/common.json b/DashAI/front/src/utils/i18n/locales/de/common.json index f11216b06..04d5433f4 100644 --- a/DashAI/front/src/utils/i18n/locales/de/common.json +++ b/DashAI/front/src/utils/i18n/locales/de/common.json @@ -186,6 +186,7 @@ "jobType": "Aufgabentyp", "lastUpdated": "Zuletzt aktualisiert", "errorMessage": "Fehlermeldung", + "progress": "Fortschritt", "unnamedJob": "Unbenannte Aufgabe" }, "status": { diff --git a/DashAI/front/src/utils/i18n/locales/en/common.json b/DashAI/front/src/utils/i18n/locales/en/common.json index dc3d26b7f..e16ba9fdd 100644 --- a/DashAI/front/src/utils/i18n/locales/en/common.json +++ b/DashAI/front/src/utils/i18n/locales/en/common.json @@ -186,6 +186,7 @@ "jobType": "Job Type", "lastUpdated": "Last Updated", "errorMessage": "Error Message", + "progress": "Progress", "unnamedJob": "Unnamed Job" }, "status": { diff --git a/DashAI/front/src/utils/i18n/locales/es/common.json b/DashAI/front/src/utils/i18n/locales/es/common.json index acf0f6980..f01424a64 100644 --- a/DashAI/front/src/utils/i18n/locales/es/common.json +++ b/DashAI/front/src/utils/i18n/locales/es/common.json @@ -188,6 +188,7 @@ "jobType": "Tipo de trabajo", "lastUpdated": "Última actualización", "errorMessage": "Mensaje de error", + "progress": "Progreso", "unnamedJob": "Trabajo sin nombre" }, "status": { diff --git a/DashAI/front/src/utils/i18n/locales/pt/common.json b/DashAI/front/src/utils/i18n/locales/pt/common.json index 201094865..22c81d1a9 100644 --- a/DashAI/front/src/utils/i18n/locales/pt/common.json +++ b/DashAI/front/src/utils/i18n/locales/pt/common.json @@ -188,6 +188,7 @@ "jobType": "Tipo de trabalho", "lastUpdated": "Última atualização", "errorMessage": "Mensagem de erro", + "progress": "Progresso", "unnamedJob": "Trabalho sem nome" }, "status": { diff --git a/DashAI/front/src/utils/i18n/locales/zh/common.json b/DashAI/front/src/utils/i18n/locales/zh/common.json index 529d2c2de..1e6e7085e 100644 --- a/DashAI/front/src/utils/i18n/locales/zh/common.json +++ b/DashAI/front/src/utils/i18n/locales/zh/common.json @@ -186,6 +186,7 @@ "jobType": "任务类型", "lastUpdated": "最后更新", "errorMessage": "错误信息", + "progress": "进度", "unnamedJob": "未命名任务" }, "status": { From 5fa8db5a283926b019b88668fc49da3ff1d25b99 Mon Sep 17 00:00:00 2001 From: Irozuku Date: Wed, 1 Jul 2026 11:52:56 -0400 Subject: [PATCH 04/89] test: cover job progress reporting and widget progress bar Add queue tests for progress writes, completion setting 100, changes payload, the old-table migration and the BaseJob no-op, plus widget tests for the determinate and indeterminate progress bars. --- .../components/jobs/JobQueueWidget.test.jsx | 48 ++++++++++- tests/back/job_queue/test_huey_job_queue.py | 79 +++++++++++++++++++ 2 files changed, 126 insertions(+), 1 deletion(-) diff --git a/DashAI/front/src/components/jobs/JobQueueWidget.test.jsx b/DashAI/front/src/components/jobs/JobQueueWidget.test.jsx index ea66f8d53..e8b32c303 100644 --- a/DashAI/front/src/components/jobs/JobQueueWidget.test.jsx +++ b/DashAI/front/src/components/jobs/JobQueueWidget.test.jsx @@ -2,6 +2,9 @@ import React from "react"; import { screen } from "@testing-library/react"; import { renderWithProviders } from "../../test-utils/renderWithProviders"; +// Mutable job list the mocked useJobManager returns; set per test. +let mockJobs = []; + // Mock API modules before importing component jest.mock("../../api/job", () => ({ deleteJob: jest.fn(), @@ -11,7 +14,7 @@ jest.mock("../../api/job", () => ({ jest.mock("../../hooks/useJobPolling", () => ({ useJobManager: () => ({ - jobs: [], + jobs: mockJobs, loading: false, error: null, refresh: jest.fn(), @@ -20,6 +23,12 @@ jest.mock("../../hooks/useJobPolling", () => ({ import JobQueueWidget from "./JobQueueWidget"; +beforeEach(() => { + mockJobs = []; + // Ensure the job list is expanded so rows (and their bars) render. + localStorage.setItem("jobQueueWidgetExpanded", "true"); +}); + describe("JobQueueWidget", () => { it("renders without crashing", () => { renderWithProviders(); @@ -29,4 +38,41 @@ describe("JobQueueWidget", () => { renderWithProviders(); expect(screen.getByText("Job Queue")).toBeInTheDocument(); }); + + it("shows a determinate bar and phase message for a running job", () => { + mockJobs = [ + { + id: "job-1", + task_type: "ModelJob", + job_name: "Train model", + status: "started", + last_update: "2026-07-01 00:00:00.000", + progress: 42, + progress_message: "Training", + }, + ]; + renderWithProviders(); + + expect(screen.getByText("Training")).toBeInTheDocument(); + const bar = screen.getByRole("progressbar"); + expect(bar).toHaveAttribute("aria-valuenow", "42"); + }); + + it("shows an indeterminate bar when a running job has no progress", () => { + mockJobs = [ + { + id: "job-2", + task_type: "ModelJob", + job_name: "Train model", + status: "started", + last_update: "2026-07-01 00:00:00.000", + progress: null, + progress_message: null, + }, + ]; + renderWithProviders(); + + const bar = screen.getByRole("progressbar"); + expect(bar).not.toHaveAttribute("aria-valuenow"); + }); }); diff --git a/tests/back/job_queue/test_huey_job_queue.py b/tests/back/job_queue/test_huey_job_queue.py index 4879b1b97..67a5c8956 100644 --- a/tests/back/job_queue/test_huey_job_queue.py +++ b/tests/back/job_queue/test_huey_job_queue.py @@ -1,3 +1,4 @@ +import sqlite3 import time # noqa: F401 import pytest @@ -110,3 +111,81 @@ def test_peek_and_get_nonexistent(test_job_queue: HueyJobQueue): with pytest.raises(JobQueueError): test_job_queue.get(job_id) + + +def test_completion_sets_progress_to_100(test_job_queue: HueyJobQueue): + job = DummyJob() + job_id = test_job_queue.put(job).id + + status = test_job_queue.status(job_id) + assert status["status"] == "finished" + assert status["progress"] == 100 + + +def test_report_progress_updates_status(test_job_queue: HueyJobQueue): + job = DummyJob() + job_id = test_job_queue.put(job).id + + test_job_queue.report_progress(job_id, 42.0, "Halfway there") + + status = test_job_queue.status(job_id) + assert status["progress"] == 42.0 + assert status["progress_message"] == "Halfway there" + + +def test_report_progress_none_is_indeterminate(test_job_queue: HueyJobQueue): + job = DummyJob() + job_id = test_job_queue.put(job).id + + test_job_queue.report_progress(job_id, None, "Working") + + status = test_job_queue.status(job_id) + assert status["progress"] is None + assert status["progress_message"] == "Working" + + +def test_report_progress_surfaces_in_changes(test_job_queue: HueyJobQueue): + job = DummyJob() + job_id = test_job_queue.put(job).id + + test_job_queue.report_progress(job_id, 25.0, "Quarter") + + changes = test_job_queue.changes_since("1970-01-01 00:00:00.000000") + changed = next(j for j in changes if j["id"] == job_id) + assert changed["progress"] == 25.0 + assert changed["progress_message"] == "Quarter" + + +def test_ensure_progress_columns_migrates_old_table(tmp_path): + # Simulate an install created before progress tracking: a 'task_copy' + # table without the progress columns. + db_path = tmp_path / "legacy.db" + with sqlite3.connect(db_path) as conn: + conn.execute( + """ + CREATE TABLE task_copy ( + id TEXT PRIMARY KEY, + task_type TEXT NOT NULL, + job_name TEXT, + enqueued_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP, + status TEXT NOT NULL, + last_update DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP, + error_msg TEXT + ) + """ + ) + + # Constructing the queue against the same directory must add the columns. + HueyJobQueue("legacy", path_db=str(tmp_path)) + + with sqlite3.connect(db_path) as conn: + cols = {row[1] for row in conn.execute("PRAGMA table_info(task_copy)")} + assert "progress" in cols + assert "progress_message" in cols + + +def test_base_job_report_progress_noop_without_huey_id(): + # A job with no huey_id (e.g. immediate mode) must not raise. + job = DummyJob() + assert job.kwargs.get("huey_id") is None + job.report_progress(0.5, "Should be ignored") From 109370911b91c1bee7cafa1b394ab4d5d8afc2ce Mon Sep 17 00:00:00 2001 From: Irozuku Date: Wed, 1 Jul 2026 12:54:54 -0400 Subject: [PATCH 05/89] feat: add COMPONENT_PATH config for downloadable components --- .gitignore | 2 ++ DashAI/back/config.py | 1 + DashAI/back/dependencies/config_builder.py | 1 + tests/back/downloads/__init__.py | 0 tests/back/downloads/test_component_path.py | 11 +++++++++++ 5 files changed, 15 insertions(+) create mode 100644 tests/back/downloads/__init__.py create mode 100644 tests/back/downloads/test_component_path.py diff --git a/.gitignore b/.gitignore index 54c488e68..dbcdb085e 100644 --- a/.gitignore +++ b/.gitignore @@ -11,6 +11,8 @@ build/ develop-eggs/ dist/ downloads/ +!DashAI/back/dependencies/downloads/ +!tests/back/downloads/ eggs/ .eggs/ lib/ diff --git a/DashAI/back/config.py b/DashAI/back/config.py index 728c8a760..6f3a0d507 100644 --- a/DashAI/back/config.py +++ b/DashAI/back/config.py @@ -22,3 +22,4 @@ class DefaultSettings(BaseSettings): EXPLANATIONS_PATH: str = "explanations" NOTEBOOK_PATH: str = "notebook" DATAFILE_PATH: str = "datafiles" + COMPONENT_PATH: str = "components" diff --git a/DashAI/back/dependencies/config_builder.py b/DashAI/back/dependencies/config_builder.py index 3323340b9..2966bc203 100644 --- a/DashAI/back/dependencies/config_builder.py +++ b/DashAI/back/dependencies/config_builder.py @@ -59,6 +59,7 @@ def build_config_dict( config["RUNS_PATH"] = local_path / config["RUNS_PATH"] config["IMAGES_PATH"] = local_path / config["IMAGES_PATH"] config["DATAFILE_PATH"] = local_path / config["DATAFILE_PATH"] + config["COMPONENT_PATH"] = local_path / config["COMPONENT_PATH"] config["FRONT_BUILD_PATH"] = pathlib.Path(config["FRONT_BUILD_PATH"]).absolute() config["BACK_PATH"] = pathlib.Path(config["BACK_PATH"]).absolute() config["LOGGING_LEVEL"] = getattr(logging, logging_level) diff --git a/tests/back/downloads/__init__.py b/tests/back/downloads/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/back/downloads/test_component_path.py b/tests/back/downloads/test_component_path.py new file mode 100644 index 000000000..765dc9f35 --- /dev/null +++ b/tests/back/downloads/test_component_path.py @@ -0,0 +1,11 @@ +from DashAI.back.config import DefaultSettings +from DashAI.back.dependencies.config_builder import build_config_dict + + +def test_default_settings_has_component_path(): + assert DefaultSettings().COMPONENT_PATH == "components" + + +def test_component_path_resolved_under_local_path(tmp_path): + config = build_config_dict(local_path=tmp_path, logging_level="INFO") + assert config["COMPONENT_PATH"] == tmp_path / "components" From 54c66ba762123f8bad008d935c8b6a2692cf392d Mon Sep 17 00:00:00 2001 From: Irozuku Date: Wed, 1 Jul 2026 12:59:07 -0400 Subject: [PATCH 06/89] feat: add DownloadableMixin and HuggingFace download helpers --- .../back/dependencies/downloads/__init__.py | 0 .../dependencies/downloads/downloadable.py | 102 ++++++++++++++++++ tests/back/downloads/test_downloadable.py | 72 +++++++++++++ 3 files changed, 174 insertions(+) create mode 100644 DashAI/back/dependencies/downloads/__init__.py create mode 100644 DashAI/back/dependencies/downloads/downloadable.py create mode 100644 tests/back/downloads/test_downloadable.py diff --git a/DashAI/back/dependencies/downloads/__init__.py b/DashAI/back/dependencies/downloads/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/DashAI/back/dependencies/downloads/downloadable.py b/DashAI/back/dependencies/downloads/downloadable.py new file mode 100644 index 000000000..3f148d92d --- /dev/null +++ b/DashAI/back/dependencies/downloads/downloadable.py @@ -0,0 +1,102 @@ +"""Source-agnostic mixin for components that download external artifacts. + +Each downloadable component owns a directory ``/`` and +is responsible for downloading its artifacts into that directory and removing +them from it. There is no shared cache; ``HFDownloadableMixin`` covers the common +HuggingFace case by downloading each repo into the component's own folder. +""" + +import logging +import pathlib +import shutil +from typing import Callable, List, Optional, Tuple + +from huggingface_hub import snapshot_download +from kink import di + +logger = logging.getLogger(__name__) + +# report(fraction, message): fraction in [0, 1] or None for indeterminate. +ProgressReporter = Callable[[Optional[float], Optional[str]], None] + + +def _components_root() -> pathlib.Path: + """Return the base directory that holds one folder per component.""" + return pathlib.Path(di["config"]["COMPONENT_PATH"]) + + +class DownloadableMixin: + """Marks a component as requiring an explicit download. + + Subclasses implement ``is_downloaded`` and ``download`` for their own source + and store artifacts under ``component_dir()``. ``delete`` defaults to removing + that directory. See ``HFDownloadableMixin`` for the HuggingFace case. + """ + + REQUIRES_DOWNLOAD: bool = True + DOWNLOAD_SIZE_BYTES: Optional[int] = None + + @classmethod + def component_dir(cls) -> pathlib.Path: + """Return this component's own storage directory. + + Returns + ------- + pathlib.Path + ``/``. + """ + return _components_root() / cls.__name__ + + @classmethod + def is_downloaded(cls) -> bool: + """Return whether the component's artifacts are present locally.""" + raise NotImplementedError + + @classmethod + def download(cls, report: Optional[ProgressReporter] = None) -> None: + """Fetch the component's artifacts into ``component_dir()``.""" + raise NotImplementedError + + @classmethod + def delete(cls) -> None: + """Remove the component's downloaded artifacts.""" + shutil.rmtree(cls.component_dir(), ignore_errors=True) + + +class HFDownloadableMixin(DownloadableMixin): + """Downloadable component whose artifacts are HuggingFace repos. + + Each repo is downloaded into ``component_dir()/``. Subclasses set + ``HF_REPOS`` or, for a dynamic repo (e.g. derived from a per-subclass + ``MODEL_NAME``), override ``hf_repos``. + """ + + HF_REPOS: List[Tuple[str, str]] = [] + + @classmethod + def hf_repos(cls) -> List[Tuple[str, str]]: + """Return the ``(repo_id, repo_type)`` pairs this component needs.""" + return list(cls.HF_REPOS) + + @classmethod + def _repo_dir(cls, repo_id: str) -> pathlib.Path: + return cls.component_dir() / repo_id.split("/")[-1] + + @classmethod + def is_downloaded(cls) -> bool: + repos = cls.hf_repos() + return bool(repos) and all( + cls._repo_dir(rid).is_dir() and any(cls._repo_dir(rid).iterdir()) + for rid, _ in repos + ) + + @classmethod + def download(cls, report: Optional[ProgressReporter] = None) -> None: + for rid, rtype in cls.hf_repos(): + target = cls._repo_dir(rid) + target.mkdir(parents=True, exist_ok=True) + if report is not None: + # snapshot_download exposes no aggregate byte count, so progress + # is reported as indeterminate (None) with a phase message. + report(None, f"Downloading {rid}") + snapshot_download(repo_id=rid, repo_type=rtype, local_dir=str(target)) diff --git a/tests/back/downloads/test_downloadable.py b/tests/back/downloads/test_downloadable.py new file mode 100644 index 000000000..337db44a1 --- /dev/null +++ b/tests/back/downloads/test_downloadable.py @@ -0,0 +1,72 @@ +import pathlib +from unittest import mock + +import pytest +from kink import di + +from DashAI.back.dependencies.downloads import downloadable as dl + +_SENTINEL = object() + + +@pytest.fixture +def components_root(tmp_path): + # kink Container has no .get(), so monkeypatch.setitem is not usable. + # Save and restore the "config" key manually. + try: + old = di["config"] + except Exception: + old = _SENTINEL + di["config"] = {"COMPONENT_PATH": str(tmp_path)} + yield pathlib.Path(tmp_path) + if old is _SENTINEL: + del di["config"] + else: + di["config"] = old + + +class _Dummy(dl.HFDownloadableMixin): + HF_REPOS = [("owner/model-a", "model")] + + +def _populate(root: pathlib.Path, cls, repo_leaf: str): + d = root / cls.__name__ / repo_leaf + d.mkdir(parents=True) + (d / "config.json").write_text("{}") + + +def test_component_dir_under_root(components_root): + assert _Dummy.component_dir() == components_root / "_Dummy" + + +def test_is_downloaded_false_when_absent(components_root): + assert _Dummy.is_downloaded() is False + + +def test_is_downloaded_true_when_present(components_root): + _populate(components_root, _Dummy, "model-a") + assert _Dummy.is_downloaded() is True + + +def test_is_downloaded_false_when_no_repos(components_root): + class Empty(dl.HFDownloadableMixin): + HF_REPOS = [] + + assert Empty.is_downloaded() is False + + +def test_download_fetches_into_component_dir(components_root): + calls = [] + with mock.patch.object(dl, "snapshot_download") as snap: + _Dummy.download(lambda frac, msg: calls.append((frac, msg))) + snap.assert_called_once() + kwargs = snap.call_args.kwargs + assert kwargs["repo_id"] == "owner/model-a" + assert kwargs["local_dir"] == str(components_root / "_Dummy" / "model-a") + assert calls # progress was reported + + +def test_delete_removes_component_dir(components_root): + _populate(components_root, _Dummy, "model-a") + _Dummy.delete() + assert not (components_root / "_Dummy").exists() From c0c8f374e00cd4773eaed7eebb726cd7441298df Mon Sep 17 00:00:00 2001 From: Irozuku Date: Wed, 1 Jul 2026 13:02:36 -0400 Subject: [PATCH 07/89] test: harden downloadable mixin tests and add _repo_dir docstring --- DashAI/back/dependencies/downloads/downloadable.py | 1 + tests/back/downloads/test_downloadable.py | 13 +++++++------ 2 files changed, 8 insertions(+), 6 deletions(-) diff --git a/DashAI/back/dependencies/downloads/downloadable.py b/DashAI/back/dependencies/downloads/downloadable.py index 3f148d92d..366d1c814 100644 --- a/DashAI/back/dependencies/downloads/downloadable.py +++ b/DashAI/back/dependencies/downloads/downloadable.py @@ -80,6 +80,7 @@ def hf_repos(cls) -> List[Tuple[str, str]]: @classmethod def _repo_dir(cls, repo_id: str) -> pathlib.Path: + """Return the local directory for a single repo under component_dir().""" return cls.component_dir() / repo_id.split("/")[-1] @classmethod diff --git a/tests/back/downloads/test_downloadable.py b/tests/back/downloads/test_downloadable.py index 337db44a1..847c74666 100644 --- a/tests/back/downloads/test_downloadable.py +++ b/tests/back/downloads/test_downloadable.py @@ -15,7 +15,7 @@ def components_root(tmp_path): # Save and restore the "config" key manually. try: old = di["config"] - except Exception: + except KeyError: old = _SENTINEL di["config"] = {"COMPONENT_PATH": str(tmp_path)} yield pathlib.Path(tmp_path) @@ -59,11 +59,12 @@ def test_download_fetches_into_component_dir(components_root): calls = [] with mock.patch.object(dl, "snapshot_download") as snap: _Dummy.download(lambda frac, msg: calls.append((frac, msg))) - snap.assert_called_once() - kwargs = snap.call_args.kwargs - assert kwargs["repo_id"] == "owner/model-a" - assert kwargs["local_dir"] == str(components_root / "_Dummy" / "model-a") - assert calls # progress was reported + snap.assert_called_once_with( + repo_id="owner/model-a", + repo_type="model", + local_dir=str(components_root / "_Dummy" / "model-a"), + ) + assert calls[0] == (None, "Downloading owner/model-a") def test_delete_removes_component_dir(components_root): From 21d38646f98cc5ebcc1771973160234924424ebf Mon Sep 17 00:00:00 2001 From: Irozuku Date: Wed, 1 Jul 2026 13:04:26 -0400 Subject: [PATCH 08/89] feat: expose requires_download and download_size_bytes in model metadata --- DashAI/back/models/base_model.py | 3 +- tests/back/models/test_download_metadata.py | 31 +++++++++++++++++++++ 2 files changed, 33 insertions(+), 1 deletion(-) create mode 100644 tests/back/models/test_download_metadata.py diff --git a/DashAI/back/models/base_model.py b/DashAI/back/models/base_model.py index 6cebb7818..36a08a962 100644 --- a/DashAI/back/models/base_model.py +++ b/DashAI/back/models/base_model.py @@ -43,7 +43,8 @@ def get_metadata(cls) -> Dict[str, Any]: """ metadata: Dict[str, Any] = {} metadata["icon"] = cls.ICON if cls.ICON else "Science" - + metadata["requires_download"] = bool(getattr(cls, "REQUIRES_DOWNLOAD", False)) + metadata["download_size_bytes"] = getattr(cls, "DOWNLOAD_SIZE_BYTES", None) return metadata @abstractmethod diff --git a/tests/back/models/test_download_metadata.py b/tests/back/models/test_download_metadata.py new file mode 100644 index 000000000..916df75ad --- /dev/null +++ b/tests/back/models/test_download_metadata.py @@ -0,0 +1,31 @@ +from DashAI.back.dependencies.downloads.downloadable import HFDownloadableMixin +from DashAI.back.models.base_model import BaseModel + + +class _PlainModel(BaseModel): + def save(self, filename): ... + @classmethod + def load(cls, filename): ... + def train(self, x, y, xv, yv): ... + + +class _DownloadableModel(HFDownloadableMixin, BaseModel): + HF_REPOS = [("owner/model", "model")] + DOWNLOAD_SIZE_BYTES = 1234 + + def save(self, filename): ... + @classmethod + def load(cls, filename): ... + def train(self, x, y, xv, yv): ... + + +def test_plain_model_not_downloadable(): + meta = _PlainModel.get_metadata() + assert meta["requires_download"] is False + assert meta["download_size_bytes"] is None + + +def test_downloadable_model_metadata(): + meta = _DownloadableModel.get_metadata() + assert meta["requires_download"] is True + assert meta["download_size_bytes"] == 1234 From 06d3c5827c5d573a9d6350665e5504d413f576b6 Mon Sep 17 00:00:00 2001 From: Irozuku Date: Wed, 1 Jul 2026 14:29:22 -0400 Subject: [PATCH 09/89] feat: track component download status in the registry --- .../registry/component_registry.py | 55 +++++++++++++++++++ tests/back/registries/test_download_status.py | 37 +++++++++++++ tests/back/registries/test_registry.py | 6 ++ 3 files changed, 98 insertions(+) create mode 100644 tests/back/registries/test_download_status.py diff --git a/DashAI/back/dependencies/registry/component_registry.py b/DashAI/back/dependencies/registry/component_registry.py index d99145aaf..959a69e4a 100644 --- a/DashAI/back/dependencies/registry/component_registry.py +++ b/DashAI/back/dependencies/registry/component_registry.py @@ -64,6 +64,8 @@ def __init__( for component in initial_components: self.register_component(component) + self.seed_download_status() + @property @beartype def registry(self) -> Dict[str, Dict[str, Any]]: @@ -225,6 +227,8 @@ def register_component(self, new_component: Type) -> None: else: self._registry[base_type][new_component.__name__] = new_register_component + self._set_download_status(new_register_component) + if hasattr(new_component, "COMPATIBLE_COMPONENTS"): for compatible_component in new_component.COMPATIBLE_COMPONENTS: self._relationship_manager.add_relationship( @@ -232,6 +236,57 @@ def register_component(self, new_component: Type) -> None: compatible_component, ) + def seed_download_status(self) -> None: + """Populate the ``downloaded`` flag for every registered component. + + Downloadable components are checked via their ``is_downloaded`` method; + all other components are always considered available. + """ + for type_dict in self._registry.values(): + for _name, component_dict in type_dict.items(): + self._set_download_status(component_dict) + + def refresh_download_status(self, name: str) -> bool: + """Recheck a single component's download status and update the registry. + + Parameters + ---------- + name : str + The component class name. + + Returns + ------- + bool + The reconciled ``downloaded`` value. + """ + component_dict = self[name] + return self._set_download_status(component_dict) + + def _set_download_status(self, component_dict: dict) -> bool: + """Set the ``downloaded`` key on a component dict and return the value. + + Parameters + ---------- + component_dict : dict + A registry component dict to update in place. + + Returns + ------- + bool + The resolved download status for the component. + """ + component_class = component_dict["class"] + requires = bool(getattr(component_class, "REQUIRES_DOWNLOAD", False)) + if requires: + try: + downloaded = bool(component_class.is_downloaded()) + except Exception: + downloaded = False + else: + downloaded = True + component_dict["downloaded"] = downloaded + return downloaded + @beartype def unregister_component(self, component: Type) -> None: """Remove a component from the registry. diff --git a/tests/back/registries/test_download_status.py b/tests/back/registries/test_download_status.py new file mode 100644 index 000000000..ff86ac8bb --- /dev/null +++ b/tests/back/registries/test_download_status.py @@ -0,0 +1,37 @@ +from DashAI.back.config_object import ConfigObject +from DashAI.back.dependencies.downloads.downloadable import DownloadableMixin +from DashAI.back.dependencies.registry import ComponentRegistry + + +class _FakeBase(ConfigObject): + TYPE = "Model" + + +_STATE = {"downloaded": False} + + +class _FakeDownloadable(DownloadableMixin, _FakeBase): + DOWNLOAD_SIZE_BYTES = 10 + + @classmethod + def is_downloaded(cls): + return _STATE["downloaded"] + + @classmethod + def download(cls, report=None): ... + @classmethod + def delete(cls): ... + + +def test_seed_sets_downloaded_flag(): + _STATE["downloaded"] = False + registry = ComponentRegistry(initial_components=[_FakeDownloadable]) + assert registry["_FakeDownloadable"]["downloaded"] is False + + +def test_refresh_reconciles_single_component(): + _STATE["downloaded"] = False + registry = ComponentRegistry(initial_components=[_FakeDownloadable]) + _STATE["downloaded"] = True + assert registry.refresh_download_status("_FakeDownloadable") is True + assert registry["_FakeDownloadable"]["downloaded"] is True diff --git a/tests/back/registries/test_registry.py b/tests/back/registries/test_registry.py index c15b17238..eabd08f44 100644 --- a/tests/back/registries/test_registry.py +++ b/tests/back/registries/test_registry.py @@ -80,6 +80,7 @@ class NoComponent: ... "description": None, "display_name": None, "color": None, + "downloaded": True, } COMPONENT2_DICT = { "name": "Component2", @@ -91,6 +92,7 @@ class NoComponent: ... "description": None, "display_name": None, "color": None, + "downloaded": True, } SUBCOMPONENT1_DICT = { "name": "SubComponent1", @@ -102,6 +104,7 @@ class NoComponent: ... "description": None, "display_name": None, "color": None, + "downloaded": True, } COMPONENT3_DICT = { "name": "Component3", @@ -113,6 +116,7 @@ class NoComponent: ... "description": "Some static component", "display_name": None, "color": None, + "downloaded": True, } COMPONENT3_DICT_MS = COMPONENT3_DICT.copy() COMPONENT3_DICT_MS["description"] = MultilingualString(en="Some static component") @@ -127,6 +131,7 @@ class NoComponent: ... "description": None, "display_name": None, "color": None, + "downloaded": True, } RELATED_COMPONENT2_DICT = { "name": "RelatedComponent2", @@ -138,6 +143,7 @@ class NoComponent: ... "description": None, "display_name": None, "color": None, + "downloaded": True, } From 139e0961f41994d6f96aaaab0192bfb488369383 Mon Sep 17 00:00:00 2001 From: Irozuku Date: Wed, 1 Jul 2026 14:32:12 -0400 Subject: [PATCH 10/89] docs: document downloaded flag in component registry --- DashAI/back/dependencies/registry/component_registry.py | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/DashAI/back/dependencies/registry/component_registry.py b/DashAI/back/dependencies/registry/component_registry.py index 959a69e4a..c6f55603d 100644 --- a/DashAI/back/dependencies/registry/component_registry.py +++ b/DashAI/back/dependencies/registry/component_registry.py @@ -30,6 +30,7 @@ class ComponentRegistry: "description": "...", # An object description. "display_name": "...", # A readable label. "color": "...", # A color associated to the component. + "downloaded": True, # False until a download-required component is fetched. } ``` @@ -64,6 +65,7 @@ def __init__( for component in initial_components: self.register_component(component) + # Ensure every component carries the downloaded flag. self.seed_download_status() @property @@ -274,6 +276,11 @@ def _set_download_status(self, component_dict: dict) -> bool: ------- bool The resolved download status for the component. + + Notes + ----- + Exceptions from ``is_downloaded()`` are suppressed and treated as + not-downloaded, so a component may appear not-downloaded without raising. """ component_class = component_dict["class"] requires = bool(getattr(component_class, "REQUIRES_DOWNLOAD", False)) From 45a007d23d805b8bcddfa21e81a1981c507e601e Mon Sep 17 00:00:00 2001 From: Irozuku Date: Wed, 1 Jul 2026 14:33:47 -0400 Subject: [PATCH 11/89] test: verify components API exposes download fields --- .../api/test_components_download_fields.py | 43 +++++++++++++++++++ 1 file changed, 43 insertions(+) create mode 100644 tests/back/api/test_components_download_fields.py diff --git a/tests/back/api/test_components_download_fields.py b/tests/back/api/test_components_download_fields.py new file mode 100644 index 000000000..5c5c00e8e --- /dev/null +++ b/tests/back/api/test_components_download_fields.py @@ -0,0 +1,43 @@ +"""Test that components API exposes download-related fields.""" + + +def test_components_list_includes_download_fields(client): + """Verify that the components list endpoint includes download fields. + + Parameters + ---------- + client : TestClient + The FastAPI test client fixture. + """ + response = client.get("/api/v1/component/", params={"select_types": ["Model"]}) + assert response.status_code == 200 + components = response.json() + assert components, "expected at least one Model component" + for component in components: + assert "downloaded" in component + assert "requires_download" in component["metadata"] + assert "download_size_bytes" in component["metadata"] + + +def test_component_by_id_includes_download_fields(client): + """Verify that the single component endpoint includes download fields. + + Parameters + ---------- + client : TestClient + The FastAPI test client fixture. + """ + # Get a Model component to test + response = client.get("/api/v1/component/", params={"select_types": ["Model"]}) + assert response.status_code == 200 + components = response.json() + assert components, "expected at least one Model component" + + # Test the first component via direct ID endpoint + component_id = components[0]["name"] + response = client.get(f"/api/v1/component/{component_id}/") + assert response.status_code == 200 + component = response.json() + assert "downloaded" in component + assert "requires_download" in component["metadata"] + assert "download_size_bytes" in component["metadata"] From c2a8d3ac0ec391a1eb3fff0afdf258c4422f398d Mon Sep 17 00:00:00 2001 From: Irozuku Date: Wed, 1 Jul 2026 14:36:57 -0400 Subject: [PATCH 12/89] feat: add ComponentDownloadJob --- DashAI/back/initial_components.py | 2 + DashAI/back/job/component_download_job.py | 58 +++++++++++++++++++ .../job_queue/test_component_download_job.py | 35 +++++++++++ 3 files changed, 95 insertions(+) create mode 100644 DashAI/back/job/component_download_job.py create mode 100644 tests/back/job_queue/test_component_download_job.py diff --git a/DashAI/back/initial_components.py b/DashAI/back/initial_components.py index 49ea258f1..d12cbafa9 100644 --- a/DashAI/back/initial_components.py +++ b/DashAI/back/initial_components.py @@ -103,6 +103,7 @@ from DashAI.back.exploration.explorers.wordcloud import WordcloudExplorer # Jobs +from DashAI.back.job.component_download_job import ComponentDownloadJob from DashAI.back.job.converter_job import ConverterJob from DashAI.back.job.datafile_job import DatafileJob from DashAI.back.job.dataset_job import DatasetJob @@ -424,6 +425,7 @@ def get_initial_components(): OptunaOptimizer, HyperOptOptimizer, # Jobs + ComponentDownloadJob, DatafileJob, ExplainerJob, ModelJob, diff --git a/DashAI/back/job/component_download_job.py b/DashAI/back/job/component_download_job.py new file mode 100644 index 000000000..09bc647ab --- /dev/null +++ b/DashAI/back/job/component_download_job.py @@ -0,0 +1,58 @@ +"""Job that downloads a component's external artifacts.""" + +import logging + +from kink import di + +from DashAI.back.job.base_job import BaseJob, JobError + +log = logging.getLogger(__name__) + + +class ComponentDownloadJob(BaseJob): + """Download the artifacts required by a downloadable component. + + Parameters + ---------- + kwargs : dict + Must contain ``component_name`` (the component class name). + """ + + def set_status_as_delivered(self) -> None: + """No dedicated DB entity; nothing to mark as delivered.""" + + def set_status_as_error(self) -> None: + """No dedicated DB entity; nothing to mark as error.""" + + def get_job_name(self) -> str: + """Return a descriptive name for the job. + + Returns + ------- + str + A human-readable name including the component name. + """ + return f"Download: {self.kwargs.get('component_name', 'component')}" + + def run(self) -> None: + """Resolve the component and download its artifacts, reporting progress. + + Raises + ------ + JobError + If the component is not registered or does not require a download. + """ + component_registry = di["component_registry"] + name = self.kwargs["component_name"] + + try: + component_class = component_registry[name]["class"] + except Exception as e: + raise JobError(f"Component {name} is not registered") from e + + if not getattr(component_class, "REQUIRES_DOWNLOAD", False): + raise JobError(f"Component {name} does not require a download") + + self.report_progress(0.0, "Starting download") + component_class.download(self.report_progress) + self.report_progress(1.0, "Download complete") diff --git a/tests/back/job_queue/test_component_download_job.py b/tests/back/job_queue/test_component_download_job.py new file mode 100644 index 000000000..026c0fa91 --- /dev/null +++ b/tests/back/job_queue/test_component_download_job.py @@ -0,0 +1,35 @@ +import pytest +from kink import di + +from DashAI.back.job.component_download_job import ComponentDownloadJob + + +@pytest.fixture(autouse=False) +def fake_registry(): + """Inject a minimal component_registry into the kink container.""" + + class FakeComponent: + REQUIRES_DOWNLOAD = True + + calls = {"download": 0} + + @classmethod + def download(cls, report=None): + cls.calls["download"] += 1 + report(None, "Downloading") + + registry = {"FakeComponent": {"class": FakeComponent}} + di["component_registry"] = registry + yield FakeComponent + del di["component_registry"] + + +def test_run_downloads_component(fake_registry): + job = ComponentDownloadJob(component_name="FakeComponent") + job.run() + assert fake_registry.calls["download"] == 1 + + +def test_get_job_name_uses_component_name(): + job = ComponentDownloadJob(component_name="FakeComponent") + assert "FakeComponent" in job.get_job_name() From e077f08d8b31c495a9d0a59da35bfd87d67b1f45 Mon Sep 17 00:00:00 2001 From: Irozuku Date: Wed, 1 Jul 2026 14:39:35 -0400 Subject: [PATCH 13/89] fix: narrow component download job registry lookup exception --- DashAI/back/job/component_download_job.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/DashAI/back/job/component_download_job.py b/DashAI/back/job/component_download_job.py index 09bc647ab..191e96b63 100644 --- a/DashAI/back/job/component_download_job.py +++ b/DashAI/back/job/component_download_job.py @@ -47,7 +47,7 @@ def run(self) -> None: try: component_class = component_registry[name]["class"] - except Exception as e: + except KeyError as e: raise JobError(f"Component {name} is not registered") from e if not getattr(component_class, "REQUIRES_DOWNLOAD", False): From 831cf2e4214831bcd0da6fb2d23599f886d89038 Mon Sep 17 00:00:00 2001 From: Irozuku Date: Wed, 1 Jul 2026 14:43:21 -0400 Subject: [PATCH 14/89] feat: add component download, delete and status endpoints --- .../back/api/api_v1/endpoints/components.py | 98 +++++++++++++++++++ .../api/test_component_download_endpoints.py | 54 ++++++++++ 2 files changed, 152 insertions(+) create mode 100644 tests/back/api/test_component_download_endpoints.py diff --git a/DashAI/back/api/api_v1/endpoints/components.py b/DashAI/back/api/api_v1/endpoints/components.py index 3c2d5cc53..9eec92d1e 100644 --- a/DashAI/back/api/api_v1/endpoints/components.py +++ b/DashAI/back/api/api_v1/endpoints/components.py @@ -236,6 +236,104 @@ async def get_components( ] +@router.get("/{name}/download") +@inject +async def get_component_download_status( + name: str, + component_registry: "ComponentRegistry" = Depends(lambda: di["component_registry"]), +): + """Return the reconciled download status of a component. + + Parameters + ---------- + name : str + The component class name. + + Returns + ------- + dict + ``{"downloaded": bool, "requires_download": bool}``. + """ + try: + component_class = component_registry[name]["class"] + except Exception as e: + raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail=str(e)) from e + requires = bool(getattr(component_class, "REQUIRES_DOWNLOAD", False)) + downloaded = component_registry.refresh_download_status(name) + return {"downloaded": downloaded, "requires_download": requires} + + +@router.post("/{name}/download", status_code=status.HTTP_201_CREATED) +@inject +async def download_component( + name: str, + component_registry: "ComponentRegistry" = Depends(lambda: di["component_registry"]), + job_queue=Depends(lambda: di["job_queue"]), +): + """Enqueue a job to download the component's artifacts. + + Parameters + ---------- + name : str + The component class name. + + Returns + ------- + dict + ``{"id": job_id}`` of the enqueued download job. + + Raises + ------ + HTTPException + 404 if unknown; 409 if not downloadable or already downloaded. + """ + from DashAI.back.job.component_download_job import ComponentDownloadJob + + try: + component_class = component_registry[name]["class"] + except Exception as e: + raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail=str(e)) from e + if not getattr(component_class, "REQUIRES_DOWNLOAD", False): + raise HTTPException( + status_code=status.HTTP_409_CONFLICT, + detail=f"Component {name} does not require a download", + ) + if component_registry.refresh_download_status(name): + raise HTTPException( + status_code=status.HTTP_409_CONFLICT, + detail=f"Component {name} is already downloaded", + ) + job = ComponentDownloadJob(component_name=name) + job.set_status_as_delivered() + job_id = job_queue.put(job).id + return {"id": job_id} + + +@router.delete("/{name}/download", status_code=status.HTTP_204_NO_CONTENT) +@inject +async def delete_component_download( + name: str, + component_registry: "ComponentRegistry" = Depends(lambda: di["component_registry"]), +): + """Delete a component's downloaded artifacts and reconcile its status. + + Parameters + ---------- + name : str + The component class name. + """ + from fastapi import Response + + try: + component_class = component_registry[name]["class"] + except Exception as e: + raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail=str(e)) from e + if getattr(component_class, "REQUIRES_DOWNLOAD", False): + component_class.delete() + component_registry.refresh_download_status(name) + return Response(status_code=status.HTTP_204_NO_CONTENT) + + @router.get("/{id}/") @inject def get_component_by_id( diff --git a/tests/back/api/test_component_download_endpoints.py b/tests/back/api/test_component_download_endpoints.py new file mode 100644 index 000000000..66e768d8c --- /dev/null +++ b/tests/back/api/test_component_download_endpoints.py @@ -0,0 +1,54 @@ +"""Tests for the component download/delete/status endpoints.""" + + +def test_get_download_status_nondownloadable(client): + # Any component that does NOT require download should return the status dict. + models = client.get("/api/v1/component/", params={"select_types": ["Model"]}).json() + non_downloadable = [m for m in models if not m["metadata"]["requires_download"]] + if not non_downloadable: + return # skip if all models require download (unlikely) + name = non_downloadable[0]["name"] + resp = client.get(f"/api/v1/component/{name}/download") + assert resp.status_code == 200 + body = resp.json() + assert "downloaded" in body + assert "requires_download" in body + assert body["requires_download"] is False + + +def test_get_download_status(client): + # Pick any downloadable Model from the list. + models = client.get("/api/v1/component/", params={"select_types": ["Model"]}).json() + downloadable = [m for m in models if m["metadata"]["requires_download"]] + if not downloadable: + return # no downloadable component registered in this environment + name = downloadable[0]["name"] + resp = client.get(f"/api/v1/component/{name}/download") + assert resp.status_code == 200 + assert "downloaded" in resp.json() + + +def test_get_download_status_unknown_404(client): + resp = client.get("/api/v1/component/NopeComponent/download") + assert resp.status_code == 404 + + +def test_download_nonexistent_component_404(client): + resp = client.post("/api/v1/component/NopeComponent/download") + assert resp.status_code == 404 + + +def test_post_download_non_downloadable_409(client): + # A component that does not require download should yield 409. + models = client.get("/api/v1/component/", params={"select_types": ["Model"]}).json() + non_downloadable = [m for m in models if not m["metadata"]["requires_download"]] + if not non_downloadable: + return + name = non_downloadable[0]["name"] + resp = client.post(f"/api/v1/component/{name}/download") + assert resp.status_code == 409 + + +def test_delete_download_unknown_404(client): + resp = client.delete("/api/v1/component/NopeComponent/download") + assert resp.status_code == 404 From 5df76f65200fb7c78c615d7ab8560595a8e573b8 Mon Sep 17 00:00:00 2001 From: Irozuku Date: Wed, 1 Jul 2026 14:47:33 -0400 Subject: [PATCH 15/89] fix: return 409 on delete of non-downloadable component and cover download enqueue --- .../back/api/api_v1/endpoints/components.py | 14 ++-- .../api/test_component_download_endpoints.py | 73 +++++++++++++++++++ 2 files changed, 81 insertions(+), 6 deletions(-) diff --git a/DashAI/back/api/api_v1/endpoints/components.py b/DashAI/back/api/api_v1/endpoints/components.py index 9eec92d1e..873d2f6df 100644 --- a/DashAI/back/api/api_v1/endpoints/components.py +++ b/DashAI/back/api/api_v1/endpoints/components.py @@ -3,7 +3,7 @@ import logging from typing import TYPE_CHECKING, Any, Dict, List, Union -from fastapi import APIRouter, Depends, Header, Query, status +from fastapi import APIRouter, Depends, Header, Query, Response, status from fastapi.exceptions import HTTPException from fastapi.responses import StreamingResponse from kink import di, inject @@ -322,15 +322,17 @@ async def delete_component_download( name : str The component class name. """ - from fastapi import Response - try: component_class = component_registry[name]["class"] except Exception as e: raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail=str(e)) from e - if getattr(component_class, "REQUIRES_DOWNLOAD", False): - component_class.delete() - component_registry.refresh_download_status(name) + if not getattr(component_class, "REQUIRES_DOWNLOAD", False): + raise HTTPException( + status_code=status.HTTP_409_CONFLICT, + detail=f"Component {name} does not support download and cannot be deleted", + ) + component_class.delete() + component_registry.refresh_download_status(name) return Response(status_code=status.HTTP_204_NO_CONTENT) diff --git a/tests/back/api/test_component_download_endpoints.py b/tests/back/api/test_component_download_endpoints.py index 66e768d8c..9405f8806 100644 --- a/tests/back/api/test_component_download_endpoints.py +++ b/tests/back/api/test_component_download_endpoints.py @@ -1,5 +1,78 @@ """Tests for the component download/delete/status endpoints.""" +import pytest +from kink import di + + +@pytest.fixture +def fake_downloadable_registry(): + """Swap di with a minimal registry containing a downloadable component.""" + + class FakeDownloadable: + REQUIRES_DOWNLOAD = True + _delete_called = False + + @classmethod + def delete(cls): + cls._delete_called = True + + class FakeRegistry: + def __getitem__(self, name): + if name == "FakeComponent": + return {"class": FakeDownloadable} + raise KeyError(f"Component {name!r} not found") + + def refresh_download_status(self, name): + return False + + class FakeJob: + id = "job-xyz" + + class FakeJobQueue: + def put(self, job): + return FakeJob() + + old_registry = di["component_registry"] + old_queue = di["job_queue"] + di["component_registry"] = FakeRegistry() + di["job_queue"] = FakeJobQueue() + yield + di["component_registry"] = old_registry + di["job_queue"] = old_queue + + +@pytest.fixture +def fake_non_downloadable_registry(): + """Swap di with a minimal registry containing a non-downloadable component.""" + + class FakeNonDownloadable: + REQUIRES_DOWNLOAD = False + + class FakeRegistry: + def __getitem__(self, name): + if name == "FakeComponent": + return {"class": FakeNonDownloadable} + raise KeyError(f"Component {name!r} not found") + + def refresh_download_status(self, name): + return False + + old_registry = di["component_registry"] + di["component_registry"] = FakeRegistry() + yield + di["component_registry"] = old_registry + + +def test_post_download_enqueue_201(client, fake_downloadable_registry): + resp = client.post("/api/v1/component/FakeComponent/download") + assert resp.status_code == 201 + assert resp.json() == {"id": "job-xyz"} + + +def test_delete_non_downloadable_409(client, fake_non_downloadable_registry): + resp = client.delete("/api/v1/component/FakeComponent/download") + assert resp.status_code == 409 + def test_get_download_status_nondownloadable(client): # Any component that does NOT require download should return the status dict. From e8916880da0907262075402f6d84e5a54c1a372b Mon Sep 17 00:00:00 2001 From: Irozuku Date: Wed, 1 Jul 2026 14:56:09 -0400 Subject: [PATCH 16/89] feat: block use of undownloaded components at run creation and training --- DashAI/back/api/api_v1/endpoints/runs.py | 18 +++- DashAI/back/job/model_job.py | 7 ++ tests/back/api/test_run_download_gate.py | 110 +++++++++++++++++++++++ 3 files changed, 134 insertions(+), 1 deletion(-) create mode 100644 tests/back/api/test_run_download_gate.py diff --git a/DashAI/back/api/api_v1/endpoints/runs.py b/DashAI/back/api/api_v1/endpoints/runs.py index 04959dc63..70becdc5f 100644 --- a/DashAI/back/api/api_v1/endpoints/runs.py +++ b/DashAI/back/api/api_v1/endpoints/runs.py @@ -289,17 +289,21 @@ async def get_hyperparameter_optimization_plot( async def upload_run( params: RunParams, session_factory: "sessionmaker" = Depends(lambda: di["session_factory"]), + component_registry=Depends(lambda: di["component_registry"]), ): """Create a new run. Parameters ---------- - params : int + params : RunParams The parameters of the new run, which includes the model session, model name, run name and description, among others. session_factory : Callable[..., ContextManager[Session]] A factory that creates a context manager that handles a SQLAlchemy session. The generated session can be used to access and query the database. + component_registry : ComponentRegistry + The application component registry, used to check whether the requested + model has been downloaded. Returns ------- @@ -310,6 +314,8 @@ async def upload_run( ------ HTTPException If the model session with id model_session_id is not registered in the DB. + HTTPException + If the model requires a download but has not been downloaded yet (HTTP 409). """ with session_factory() as db: try: @@ -319,6 +325,16 @@ async def upload_run( status_code=status.HTTP_404_NOT_FOUND, detail="Model session not found", ) + entry = component_registry[params.model_name] + if getattr(entry["class"], "REQUIRES_DOWNLOAD", False) and not entry.get( + "downloaded", False + ): + raise HTTPException( + status_code=status.HTTP_409_CONFLICT, + detail=( + f"Model {params.model_name} must be downloaded before use." + ), + ) run = Run( model_session_id=params.model_session_id, model_name=params.model_name, diff --git a/DashAI/back/job/model_job.py b/DashAI/back/job/model_job.py index 1109696c9..b3829f54f 100644 --- a/DashAI/back/job/model_job.py +++ b/DashAI/back/job/model_job.py @@ -222,6 +222,13 @@ def run( raise JobError( f"Unable to find Model with name {run.model_name} in registry.", ) from e + if getattr(run_model_class, "REQUIRES_DOWNLOAD", False) and not ( + run_model_class.is_downloaded() + ): + raise JobError( + f"Model {run.model_name} is not downloaded. " + "Download it before training." + ) try: factory = ModelFactory( run_model_class, diff --git a/tests/back/api/test_run_download_gate.py b/tests/back/api/test_run_download_gate.py new file mode 100644 index 000000000..b54a37ddd --- /dev/null +++ b/tests/back/api/test_run_download_gate.py @@ -0,0 +1,110 @@ +"""Tests for the run-creation download gate (Task 8).""" + +from kink import di + +_RUN_PAYLOAD_BASE = { + "parameters": {}, + "optimizer_name": "OptunaOptimizer", + "optimizer_parameters": {}, + "goal_metric": "Accuracy", + "plot_history_path": "", + "plot_slice_path": "", + "plot_contour_path": "", + "plot_importance_path": "", + "name": "gate-test-run", + "description": "", +} + + +class _FakeDownloadableModel: + """Minimal stub that looks like an undownloaded, download-required model.""" + + REQUIRES_DOWNLOAD = True + + @classmethod + def is_downloaded(cls): + return False + + +class _FakeRegistry: + """Registry wrapper that exposes FakeDownloadableModel on top of the real one.""" + + def __init__(self, real): + self._real = real + + def __getitem__(self, name): + if name == "FakeDownloadableModel": + return {"class": _FakeDownloadableModel, "downloaded": False} + return self._real[name] + + def get_components_by_types(self, select=None, ignore=None): + return self._real.get_components_by_types(select=select, ignore=ignore) + + def __contains__(self, name): + return name == "FakeDownloadableModel" or name in self._real + + +def _make_model_session(client): + """Insert a Dataset + ModelSession row and return the ModelSession id.""" + from DashAI.back.dependencies.database.models import Dataset, ModelSession + + sf = client.app.container["session_factory"] + with sf() as db: + ds = Dataset(name="__gate_test_ds__", file_path="") + db.add(ds) + db.flush() + ms = ModelSession( + dataset_id=ds.id, + name="__gate_test_ms__", + task_name="TabularClassificationTask", + input_columns=[], + output_columns=[], + train_metrics=[], + validation_metrics=[], + test_metrics=[], + splits={}, + ) + db.add(ms) + db.commit() + db.refresh(ms) + return ms.id + + +def test_upload_run_rejects_undownloaded_model_synthetic(client): + """Creating a run for a not-yet-downloaded model must return HTTP 409.""" + ms_id = _make_model_session(client) + old = di["component_registry"] + di["component_registry"] = _FakeRegistry(old) + try: + resp = client.post( + "/api/v1/run/", + json={ + "model_session_id": ms_id, + "model_name": "FakeDownloadableModel", + **_RUN_PAYLOAD_BASE, + }, + ) + finally: + di["component_registry"] = old + + assert resp.status_code == 409 + assert "download" in resp.json()["detail"].lower() + + +def test_upload_run_rejects_undownloaded_model(client, monkeypatch): + """Force a real downloadable model undownloaded; run creation must return 409.""" + registry = di["component_registry"] + models = client.get("/api/v1/component/", params={"select_types": ["Model"]}).json() + downloadable = [m for m in models if m["metadata"]["requires_download"]] + if not downloadable: + return # no downloadable model registered in this environment; skip + name = downloadable[0]["name"] + registry[name]["downloaded"] = False + + resp = client.post( + "/api/v1/run/", + json={"model_session_id": 1, "model_name": name, **_RUN_PAYLOAD_BASE}, + ) + assert resp.status_code in (404, 409) + if resp.status_code == 409: + assert "download" in resp.json()["detail"].lower() From 46f90a9541afeb78171005374be473ff3a4afbf7 Mon Sep 17 00:00:00 2001 From: Irozuku Date: Wed, 1 Jul 2026 15:00:57 -0400 Subject: [PATCH 17/89] fix: return 422 for unknown model in run gate and restore registry state in tests --- DashAI/back/api/api_v1/endpoints/runs.py | 7 +++++++ tests/back/api/test_run_download_gate.py | 25 ++++++++++++++++++++---- 2 files changed, 28 insertions(+), 4 deletions(-) diff --git a/DashAI/back/api/api_v1/endpoints/runs.py b/DashAI/back/api/api_v1/endpoints/runs.py index 70becdc5f..0c4dd2119 100644 --- a/DashAI/back/api/api_v1/endpoints/runs.py +++ b/DashAI/back/api/api_v1/endpoints/runs.py @@ -325,6 +325,13 @@ async def upload_run( status_code=status.HTTP_404_NOT_FOUND, detail="Model session not found", ) + # REQUIRES_DOWNLOAD is read from the class (static contract); + # "downloaded" is read from the registry dict (runtime state, Task 4). + if params.model_name not in component_registry: + raise HTTPException( + status_code=status.HTTP_422_UNPROCESSABLE_ENTITY, + detail=f"Unknown model '{params.model_name}'", + ) entry = component_registry[params.model_name] if getattr(entry["class"], "REQUIRES_DOWNLOAD", False) and not entry.get( "downloaded", False diff --git a/tests/back/api/test_run_download_gate.py b/tests/back/api/test_run_download_gate.py index b54a37ddd..e8e89ab9b 100644 --- a/tests/back/api/test_run_download_gate.py +++ b/tests/back/api/test_run_download_gate.py @@ -44,18 +44,21 @@ def __contains__(self, name): return name == "FakeDownloadableModel" or name in self._real -def _make_model_session(client): +def _make_model_session(client, suffix=""): """Insert a Dataset + ModelSession row and return the ModelSession id.""" + import uuid + from DashAI.back.dependencies.database.models import Dataset, ModelSession + unique = suffix or uuid.uuid4().hex[:8] sf = client.app.container["session_factory"] with sf() as db: - ds = Dataset(name="__gate_test_ds__", file_path="") + ds = Dataset(name=f"__gate_test_ds_{unique}__", file_path="") db.add(ds) db.flush() ms = ModelSession( dataset_id=ds.id, - name="__gate_test_ms__", + name=f"__gate_test_ms_{unique}__", task_name="TabularClassificationTask", input_columns=[], output_columns=[], @@ -99,7 +102,7 @@ def test_upload_run_rejects_undownloaded_model(client, monkeypatch): if not downloadable: return # no downloadable model registered in this environment; skip name = downloadable[0]["name"] - registry[name]["downloaded"] = False + monkeypatch.setitem(registry[name], "downloaded", False) resp = client.post( "/api/v1/run/", @@ -108,3 +111,17 @@ def test_upload_run_rejects_undownloaded_model(client, monkeypatch): assert resp.status_code in (404, 409) if resp.status_code == 409: assert "download" in resp.json()["detail"].lower() + + +def test_upload_run_unknown_model_422(client): + """POSTing a run with an unregistered model_name must return HTTP 422.""" + ms_id = _make_model_session(client) + resp = client.post( + "/api/v1/run/", + json={ + "model_session_id": ms_id, + "model_name": "__totally_bogus_model_xyz__", + **_RUN_PAYLOAD_BASE, + }, + ) + assert resp.status_code == 422 From 82c9270b6e125f7104e9854b01bf055f45fcae5c Mon Sep 17 00:00:00 2001 From: Irozuku Date: Wed, 1 Jul 2026 15:05:39 -0400 Subject: [PATCH 18/89] feat: make Opus-MT translation models downloadable --- .../hugging_face/base_opus_mt_transformer.py | 22 +++++++-- .../back/models/test_opus_mt_downloadable.py | 47 +++++++++++++++++++ 2 files changed, 66 insertions(+), 3 deletions(-) create mode 100644 tests/back/models/test_opus_mt_downloadable.py diff --git a/DashAI/back/models/hugging_face/base_opus_mt_transformer.py b/DashAI/back/models/hugging_face/base_opus_mt_transformer.py index 1b80e1269..a49b53228 100644 --- a/DashAI/back/models/hugging_face/base_opus_mt_transformer.py +++ b/DashAI/back/models/hugging_face/base_opus_mt_transformer.py @@ -7,6 +7,7 @@ from sklearn.exceptions import NotFittedError +from DashAI.back.dependencies.downloads.downloadable import HFDownloadableMixin from DashAI.back.models.translation_model import TranslationModel from DashAI.back.models.utils import ( GPU_OR_CPU_PLACEHOLDER, @@ -17,7 +18,7 @@ from DashAI.back.dataloaders.classes.dashai_dataset import DashAIDataset -class OpusMtTransformerMixin(TranslationModel): +class OpusMtTransformerMixin(HFDownloadableMixin, TranslationModel): """Shared implementation for Helsinki-NLP Opus-MT translation wrappers. Subclasses must define ``MODEL_NAME`` (the HuggingFace checkpoint ID) and @@ -35,6 +36,20 @@ class OpusMtTransformerMixin(TranslationModel): MODEL_NAME: str = "" TEMP_CHECKPOINT_DIR: str = "DashAI/back/user_models/temp_checkpoints_opus_mt" + # Marian Opus-MT checkpoints are ~300 MB; declared statically for the UI. + DOWNLOAD_SIZE_BYTES: int = 300_000_000 + + @classmethod + def hf_repos(cls): + """Derive the single HuggingFace repo from the subclass MODEL_NAME. + + Returns + ------- + list of tuple of (str, str) + A single ``(repo_id, repo_type)`` pair derived from ``MODEL_NAME``, + or an empty list when ``MODEL_NAME`` is not set. + """ + return [(cls.MODEL_NAME, "model")] if cls.MODEL_NAME else [] def __init__(self, model=None, **kwargs): """Initialize tokenizer and seq2seq model. @@ -56,7 +71,8 @@ def __init__(self, model=None, **kwargs): ) self.model_name = self.MODEL_NAME - self.tokenizer = AutoTokenizer.from_pretrained(self.model_name) + local_dir = str(self._repo_dir(self.MODEL_NAME)) + self.tokenizer = AutoTokenizer.from_pretrained(local_dir) self.training_args = { "num_train_epochs": kwargs.get("num_train_epochs", 2), @@ -77,7 +93,7 @@ def __init__(self, model=None, **kwargs): if model is None: from transformers import AutoModelForSeq2SeqLM - self.model = AutoModelForSeq2SeqLM.from_pretrained(self.model_name) + self.model = AutoModelForSeq2SeqLM.from_pretrained(local_dir) else: self.model = model diff --git a/tests/back/models/test_opus_mt_downloadable.py b/tests/back/models/test_opus_mt_downloadable.py new file mode 100644 index 000000000..fc91adf81 --- /dev/null +++ b/tests/back/models/test_opus_mt_downloadable.py @@ -0,0 +1,47 @@ +"""Tests that OpusMtTransformerMixin subclasses expose download metadata.""" + +import pytest +from kink import di + +from DashAI.back.dependencies.downloads.downloadable import HFDownloadableMixin +from DashAI.back.models.hugging_face.opus_mt_en_es_transformer import ( + OpusMtEnESTransformer, +) + + +@pytest.fixture +def component_root(tmp_path): + """Inject a temporary COMPONENT_PATH into the kink DI container.""" + sentinel = object() + old = di["config"] if "config" in di else sentinel # noqa: SIM401 + di["config"] = {"COMPONENT_PATH": str(tmp_path)} + yield tmp_path + if old is sentinel: + del di["config"] + else: + di["config"] = old + + +def test_opus_mt_is_downloadable(): + assert issubclass(OpusMtEnESTransformer, HFDownloadableMixin) + assert OpusMtEnESTransformer.REQUIRES_DOWNLOAD is True + assert OpusMtEnESTransformer.DOWNLOAD_SIZE_BYTES is not None + + +def test_opus_mt_hf_repos_derived_from_model_name(): + assert OpusMtEnESTransformer.hf_repos() == [("Helsinki-NLP/opus-mt-en-es", "model")] + + +def test_opus_mt_metadata_flags_download(): + meta = OpusMtEnESTransformer.get_metadata() + assert meta["requires_download"] is True + assert meta["download_size_bytes"] == OpusMtEnESTransformer.DOWNLOAD_SIZE_BYTES + + +def test_opus_mt_is_downloaded_uses_component_dir(component_root): + assert OpusMtEnESTransformer.is_downloaded() is False + + repo_dir = component_root / "OpusMtEnESTransformer" / "opus-mt-en-es" + repo_dir.mkdir(parents=True) + (repo_dir / "config.json").write_text("{}") + assert OpusMtEnESTransformer.is_downloaded() is True From 5c2354ca312fc2f66e33a4d0d182554a7216f439 Mon Sep 17 00:00:00 2001 From: Irozuku Date: Wed, 1 Jul 2026 15:08:58 -0400 Subject: [PATCH 19/89] docs: correct Opus-MT loading note for local downloads --- DashAI/back/models/hugging_face/base_opus_mt_transformer.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/DashAI/back/models/hugging_face/base_opus_mt_transformer.py b/DashAI/back/models/hugging_face/base_opus_mt_transformer.py index a49b53228..38eeecf04 100644 --- a/DashAI/back/models/hugging_face/base_opus_mt_transformer.py +++ b/DashAI/back/models/hugging_face/base_opus_mt_transformer.py @@ -30,8 +30,10 @@ class OpusMtTransformerMixin(HFDownloadableMixin, TranslationModel): here so each language-pair subclass only needs to set class attributes. .. note:: - Requires internet access on first use to download pretrained weights - from the Hugging Face Hub. + The pretrained weights must be downloaded first via ``download()`` + (this component requires a download). ``__init__`` loads the tokenizer + and model from the component's local folder; it does not fetch from + the Hugging Face Hub. """ MODEL_NAME: str = "" From 3878e3c55bb7400c3b4ed218c0f3693778654dec Mon Sep 17 00:00:00 2001 From: Irozuku Date: Wed, 1 Jul 2026 15:13:22 -0400 Subject: [PATCH 20/89] feat: add frontend component download API and types --- DashAI/front/src/api/__mocks__/api.ts | 5 +++++ DashAI/front/src/api/component.test.ts | 20 ++++++++++++++++++++ DashAI/front/src/api/component.ts | 23 +++++++++++++++++++++++ DashAI/front/src/types/component.ts | 1 + 4 files changed, 49 insertions(+) create mode 100644 DashAI/front/src/api/__mocks__/api.ts create mode 100644 DashAI/front/src/api/component.test.ts diff --git a/DashAI/front/src/api/__mocks__/api.ts b/DashAI/front/src/api/__mocks__/api.ts new file mode 100644 index 000000000..22bcb5e06 --- /dev/null +++ b/DashAI/front/src/api/__mocks__/api.ts @@ -0,0 +1,5 @@ +export default { + post: jest.fn(), + get: jest.fn(), + delete: jest.fn(), +}; diff --git a/DashAI/front/src/api/component.test.ts b/DashAI/front/src/api/component.test.ts new file mode 100644 index 000000000..56a5f8581 --- /dev/null +++ b/DashAI/front/src/api/component.test.ts @@ -0,0 +1,20 @@ +jest.mock("./api"); + +import api from "./api"; +import { downloadComponent } from "./component"; + +describe("component download api", () => { + beforeEach(() => { + jest.clearAllMocks(); + }); + + it("posts to the download endpoint and returns the job id", async () => { + (api.post as jest.Mock).mockResolvedValue({ data: { id: "job-1" } }); + + const result = await downloadComponent("OpusMtEnRoaTransformer"); + expect(api.post).toHaveBeenCalledWith( + "/v1/component/OpusMtEnRoaTransformer/download", + ); + expect(result).toEqual({ id: "job-1" }); + }); +}); diff --git a/DashAI/front/src/api/component.ts b/DashAI/front/src/api/component.ts index cc66da278..a4008831c 100644 --- a/DashAI/front/src/api/component.ts +++ b/DashAI/front/src/api/component.ts @@ -54,3 +54,26 @@ export const getComponentById = async (id: string): Promise => { const response = await api.get(`/v1/component/${id}/`); return response.data; }; + +export const downloadComponent = async ( + name: string, +): Promise<{ id: string }> => { + const response = await api.post<{ id: string }>( + `/v1/component/${name}/download`, + ); + return response.data; +}; + +export const deleteComponentDownload = async (name: string): Promise => { + await api.delete(`/v1/component/${name}/download`); +}; + +export const getComponentDownloadStatus = async ( + name: string, +): Promise<{ downloaded: boolean; requires_download: boolean }> => { + const response = await api.get<{ + downloaded: boolean; + requires_download: boolean; + }>(`/v1/component/${name}/download`); + return response.data; +}; diff --git a/DashAI/front/src/types/component.ts b/DashAI/front/src/types/component.ts index 9bedf444b..f7fad0a32 100644 --- a/DashAI/front/src/types/component.ts +++ b/DashAI/front/src/types/component.ts @@ -10,4 +10,5 @@ export interface IComponent { description: string; display_name?: string; color?: string; + downloaded?: boolean; } From 7ca82c5fd93a1c0b207090c8a873b5e4a9de6e21 Mon Sep 17 00:00:00 2001 From: Irozuku Date: Wed, 1 Jul 2026 15:15:02 -0400 Subject: [PATCH 21/89] feat: add download metadata fields to component metadata type --- DashAI/front/src/types/task.ts | 2 ++ 1 file changed, 2 insertions(+) diff --git a/DashAI/front/src/types/task.ts b/DashAI/front/src/types/task.ts index 3bca2b158..33e17cfe2 100644 --- a/DashAI/front/src/types/task.ts +++ b/DashAI/front/src/types/task.ts @@ -11,4 +11,6 @@ export interface ITaskMetadataParameters { outputs_columns: string[]; inputs_cardinality: "n" | number; outputs_cardinality: "n" | number; + requires_download?: boolean; + download_size_bytes?: number | null; } From 95d2f7ce38fe18f7b147246f3c3d95241cc1f0eb Mon Sep 17 00:00:00 2001 From: Irozuku Date: Wed, 1 Jul 2026 15:21:47 -0400 Subject: [PATCH 22/89] feat: add component download control with progress and usage gating --- .../src/components/models/AddModelDialog.jsx | 59 ++++++++-- .../src/components/models/ModelsRightBar.jsx | 1 + .../models/model/ComponentDownloadControl.jsx | 111 ++++++++++++++++++ .../model/ComponentDownloadControl.test.jsx | 41 +++++++ .../src/utils/i18n/locales/de/common.json | 8 ++ .../src/utils/i18n/locales/en/common.json | 8 ++ .../src/utils/i18n/locales/es/common.json | 8 ++ .../src/utils/i18n/locales/pt/common.json | 8 ++ .../src/utils/i18n/locales/zh/common.json | 8 ++ 9 files changed, 240 insertions(+), 12 deletions(-) create mode 100644 DashAI/front/src/components/models/model/ComponentDownloadControl.jsx create mode 100644 DashAI/front/src/components/models/model/ComponentDownloadControl.test.jsx diff --git a/DashAI/front/src/components/models/AddModelDialog.jsx b/DashAI/front/src/components/models/AddModelDialog.jsx index 61cde7eb3..be97edb44 100644 --- a/DashAI/front/src/components/models/AddModelDialog.jsx +++ b/DashAI/front/src/components/models/AddModelDialog.jsx @@ -12,6 +12,7 @@ import { TextField, Box, IconButton, + Tooltip, Typography, } from "@mui/material"; import { Close as CloseIcon } from "@mui/icons-material"; @@ -26,6 +27,7 @@ import { createRun } from "../../api/run"; import { useTranslation } from "react-i18next"; import { useTourContext } from "../tour/TourProvider"; import { checkIfHaveOptimazers } from "../../utils/schema"; +import ComponentDownloadControl from "./model/ComponentDownloadControl"; /** * Dialog for adding a new model run to a session @@ -37,6 +39,7 @@ function AddModelDialog({ onClose, session, preselectedModel, + preselectedModelObject, existingRuns = [], onRunCreated, }) { @@ -51,6 +54,7 @@ function AddModelDialog({ const [hasUserTouchedName, setHasUserTouchedName] = useState(false); const [goalMetric, setGoalMetric] = useState(""); const [hasLoadedInitialParams, setHasLoadedInitialParams] = useState(false); + const [modelDownloaded, setModelDownloaded] = useState(true); const { t } = useTranslation(["models", "common"]); const { defaultValues: defaultModelParams } = useSchema({ @@ -96,6 +100,13 @@ function AddModelDialog({ } }, [preselectedModel, selectedModel]); + useEffect(() => { + const comp = preselectedModelObject; + const requiresDownload = Boolean(comp?.metadata?.requires_download); + const isDownloaded = Boolean(comp?.downloaded); + setModelDownloaded(!requiresDownload || isDownloaded); + }, [preselectedModelObject]); + useEffect(() => { if ( selectedModel && @@ -314,6 +325,12 @@ function AddModelDialog({ )} + {preselectedModelObject?.metadata?.requires_download && ( + + )} )} @@ -371,20 +388,37 @@ function AddModelDialog({ {t("common:back")} )} - + + + + ); @@ -399,6 +433,7 @@ AddModelDialog.propTypes = { task_name: PropTypes.string, }), preselectedModel: PropTypes.string, + preselectedModelObject: PropTypes.object, existingRuns: PropTypes.array, onRunCreated: PropTypes.func, }; diff --git a/DashAI/front/src/components/models/ModelsRightBar.jsx b/DashAI/front/src/components/models/ModelsRightBar.jsx index 1d87f889e..f77ab0a0d 100644 --- a/DashAI/front/src/components/models/ModelsRightBar.jsx +++ b/DashAI/front/src/components/models/ModelsRightBar.jsx @@ -255,6 +255,7 @@ export default function ModelsRightBar({ onToggle }) { open={configOpen} onClose={closeConfig} preselectedModel={selectedModel?.name} + preselectedModelObject={selectedModel} session={session} existingRuns={existingRuns} onRunCreated={onRunCreated} diff --git a/DashAI/front/src/components/models/model/ComponentDownloadControl.jsx b/DashAI/front/src/components/models/model/ComponentDownloadControl.jsx new file mode 100644 index 000000000..4fd5694ac --- /dev/null +++ b/DashAI/front/src/components/models/model/ComponentDownloadControl.jsx @@ -0,0 +1,111 @@ +import React, { useState, useEffect } from "react"; +import { Box, Button, LinearProgress, Typography } from "@mui/material"; +import DownloadIcon from "@mui/icons-material/Download"; +import DeleteIcon from "@mui/icons-material/Delete"; +import { useTranslation } from "react-i18next"; +import { useSnackbar } from "notistack"; +import { + downloadComponent, + deleteComponentDownload, + getComponentDownloadStatus, +} from "../../../api/component"; +import { startJobPolling } from "../../../utils/jobPoller"; + +const formatSize = (bytes) => { + if (bytes == null) return ""; + const mb = bytes / 1024 / 1024; + if (mb >= 1024) return `${(mb / 1024).toFixed(1)} GB`; + return `${Math.round(mb)} MB`; +}; + +const ComponentDownloadControl = ({ component, onStatusChange }) => { + const { t } = useTranslation(["common"]); + const { enqueueSnackbar } = useSnackbar(); + const meta = component.metadata || {}; + const [downloaded, setDownloaded] = useState(Boolean(component.downloaded)); + const [downloading, setDownloading] = useState(false); + + useEffect(() => { + setDownloaded(Boolean(component.downloaded)); + }, [component.name, component.downloaded]); + + if (!meta.requires_download) return null; + + const finish = (isDownloaded) => { + setDownloading(false); + setDownloaded(isDownloaded); + if (onStatusChange) onStatusChange(isDownloaded); + }; + + const handleDownload = async () => { + setDownloading(true); + try { + const { id } = await downloadComponent(component.name); + startJobPolling( + id, + async () => { + const status = await getComponentDownloadStatus(component.name); + finish(status.downloaded); + enqueueSnackbar(t("common:componentDownload.done"), { + variant: "success", + }); + }, + () => { + finish(false); + enqueueSnackbar(t("common:componentDownload.failed"), { + variant: "error", + }); + }, + ); + } catch (e) { + finish(false); + enqueueSnackbar(t("common:componentDownload.failed"), { + variant: "error", + }); + } + }; + + const handleDelete = async () => { + await deleteComponentDownload(component.name); + finish(false); + }; + + if (downloading) { + return ( + + + {t("common:componentDownload.downloading")} + + + + ); + } + + if (downloaded) { + return ( + + ); + } + + return ( + + ); +}; + +export default ComponentDownloadControl; diff --git a/DashAI/front/src/components/models/model/ComponentDownloadControl.test.jsx b/DashAI/front/src/components/models/model/ComponentDownloadControl.test.jsx new file mode 100644 index 000000000..6c90798e3 --- /dev/null +++ b/DashAI/front/src/components/models/model/ComponentDownloadControl.test.jsx @@ -0,0 +1,41 @@ +import React from "react"; +import { screen, fireEvent, waitFor } from "@testing-library/react"; +import { renderWithProviders } from "../../../test-utils/renderWithProviders"; + +jest.mock("../../../api/component", () => ({ + downloadComponent: jest.fn(() => Promise.resolve({ id: "job-1" })), + deleteComponentDownload: jest.fn(() => Promise.resolve()), + getComponentDownloadStatus: jest.fn(() => + Promise.resolve({ downloaded: false, requires_download: true }), + ), +})); +jest.mock("../../../utils/jobPoller", () => ({ + startJobPolling: jest.fn(), + stopJobPolling: jest.fn(), + subscribeJobs: jest.fn(() => () => {}), +})); + +import ComponentDownloadControl from "./ComponentDownloadControl"; +import { downloadComponent } from "../../../api/component"; + +const component = { + name: "OpusMtEnRoaTransformer", + downloaded: false, + metadata: { requires_download: true, download_size_bytes: 310000000 }, +}; + +describe("ComponentDownloadControl", () => { + it("shows a download button with the size and triggers download", async () => { + renderWithProviders( + {}} + />, + ); + const button = await screen.findByRole("button", { name: /download/i }); + fireEvent.click(button); + await waitFor(() => + expect(downloadComponent).toHaveBeenCalledWith("OpusMtEnRoaTransformer"), + ); + }); +}); diff --git a/DashAI/front/src/utils/i18n/locales/de/common.json b/DashAI/front/src/utils/i18n/locales/de/common.json index 04d5433f4..2f777cdbb 100644 --- a/DashAI/front/src/utils/i18n/locales/de/common.json +++ b/DashAI/front/src/utils/i18n/locales/de/common.json @@ -160,6 +160,14 @@ "columnNameInvalidCharacters": "Spaltenname darf nur Buchstaben, Zahlen und Unterstriche enthalten", "columnNameAlreadyExists": "Eine Spalte mit diesem Namen existiert bereits", "errorRenamingColumn": "Fehler beim Umbenennen der Spalte", + "componentDownload": { + "download": "Herunterladen ({{size}})", + "delete": "Download loeschen", + "downloading": "Wird heruntergeladen...", + "done": "Komponente heruntergeladen", + "failed": "Download der Komponente fehlgeschlagen", + "mustDownload": "Dieses Modell muss vor der Nutzung heruntergeladen werden" + }, "jobQueue": { "title": "Aufgabenwarteschlange", "refresh": "Aktualisieren", diff --git a/DashAI/front/src/utils/i18n/locales/en/common.json b/DashAI/front/src/utils/i18n/locales/en/common.json index e16ba9fdd..0220e806f 100644 --- a/DashAI/front/src/utils/i18n/locales/en/common.json +++ b/DashAI/front/src/utils/i18n/locales/en/common.json @@ -160,6 +160,14 @@ "columnNameInvalidCharacters": "Column name can only contain letters, numbers and underscores", "columnNameAlreadyExists": "A column with this name already exists", "errorRenamingColumn": "Error renaming column", + "componentDownload": { + "download": "Download ({{size}})", + "delete": "Delete download", + "downloading": "Downloading...", + "done": "Component downloaded", + "failed": "Component download failed", + "mustDownload": "This model must be downloaded before use" + }, "jobQueue": { "title": "Job Queue", "refresh": "Refresh", diff --git a/DashAI/front/src/utils/i18n/locales/es/common.json b/DashAI/front/src/utils/i18n/locales/es/common.json index f01424a64..dbe4bbf16 100644 --- a/DashAI/front/src/utils/i18n/locales/es/common.json +++ b/DashAI/front/src/utils/i18n/locales/es/common.json @@ -160,6 +160,14 @@ "columnNameInvalidCharacters": "El nombre de la columna solo puede contener letras, números y guiones bajos", "columnNameAlreadyExists": "Ya existe una columna con este nombre", "errorRenamingColumn": "Error al renombrar la columna", + "componentDownload": { + "download": "Descargar ({{size}})", + "delete": "Eliminar descarga", + "downloading": "Descargando...", + "done": "Componente descargado", + "failed": "La descarga del componente ha fallado", + "mustDownload": "Este modelo debe descargarse antes de usarlo" + }, "jobQueue": { "title": "Cola de trabajos", "refresh": "Actualizar", diff --git a/DashAI/front/src/utils/i18n/locales/pt/common.json b/DashAI/front/src/utils/i18n/locales/pt/common.json index 22c81d1a9..e52110881 100644 --- a/DashAI/front/src/utils/i18n/locales/pt/common.json +++ b/DashAI/front/src/utils/i18n/locales/pt/common.json @@ -160,6 +160,14 @@ "columnNameInvalidCharacters": "O nome da coluna só pode conter letras, números e underscores", "columnNameAlreadyExists": "Já existe uma coluna com este nome", "errorRenamingColumn": "Erro ao renomear a coluna", + "componentDownload": { + "download": "Baixar ({{size}})", + "delete": "Remover download", + "downloading": "Baixando...", + "done": "Componente baixado", + "failed": "Falha ao baixar o componente", + "mustDownload": "Este modelo precisa ser baixado antes de usar" + }, "jobQueue": { "title": "Fila de trabalhos", "refresh": "Atualizar", diff --git a/DashAI/front/src/utils/i18n/locales/zh/common.json b/DashAI/front/src/utils/i18n/locales/zh/common.json index 1e6e7085e..d0623fb42 100644 --- a/DashAI/front/src/utils/i18n/locales/zh/common.json +++ b/DashAI/front/src/utils/i18n/locales/zh/common.json @@ -160,6 +160,14 @@ "columnNameInvalidCharacters": "列名只能包含字母、数字和下划线", "columnNameAlreadyExists": "已存在同名列", "errorRenamingColumn": "重命名列时出错", + "componentDownload": { + "download": "下载 ({{size}})", + "delete": "删除下载", + "downloading": "下载中...", + "done": "组件已下载", + "failed": "组件下载失败", + "mustDownload": "使用前必须先下载此模型" + }, "jobQueue": { "title": "任务队列", "refresh": "刷新", From 513ae7a32b87d2cf3696d0cfbde761d7e721d514 Mon Sep 17 00:00:00 2001 From: Irozuku Date: Wed, 1 Jul 2026 15:26:12 -0400 Subject: [PATCH 23/89] fix: stop component download poller on unmount and handle delete errors --- .../src/components/models/AddModelDialog.jsx | 6 ++++- .../models/model/ComponentDownloadControl.jsx | 24 +++++++++++++++---- 2 files changed, 25 insertions(+), 5 deletions(-) diff --git a/DashAI/front/src/components/models/AddModelDialog.jsx b/DashAI/front/src/components/models/AddModelDialog.jsx index be97edb44..80604a474 100644 --- a/DashAI/front/src/components/models/AddModelDialog.jsx +++ b/DashAI/front/src/components/models/AddModelDialog.jsx @@ -433,7 +433,11 @@ AddModelDialog.propTypes = { task_name: PropTypes.string, }), preselectedModel: PropTypes.string, - preselectedModelObject: PropTypes.object, + preselectedModelObject: PropTypes.shape({ + name: PropTypes.string, + downloaded: PropTypes.bool, + metadata: PropTypes.object, + }), existingRuns: PropTypes.array, onRunCreated: PropTypes.func, }; diff --git a/DashAI/front/src/components/models/model/ComponentDownloadControl.jsx b/DashAI/front/src/components/models/model/ComponentDownloadControl.jsx index 4fd5694ac..bcc6e207c 100644 --- a/DashAI/front/src/components/models/model/ComponentDownloadControl.jsx +++ b/DashAI/front/src/components/models/model/ComponentDownloadControl.jsx @@ -1,4 +1,4 @@ -import React, { useState, useEffect } from "react"; +import React, { useState, useEffect, useRef } from "react"; import { Box, Button, LinearProgress, Typography } from "@mui/material"; import DownloadIcon from "@mui/icons-material/Download"; import DeleteIcon from "@mui/icons-material/Delete"; @@ -9,7 +9,7 @@ import { deleteComponentDownload, getComponentDownloadStatus, } from "../../../api/component"; -import { startJobPolling } from "../../../utils/jobPoller"; +import { startJobPolling, stopJobPolling } from "../../../utils/jobPoller"; const formatSize = (bytes) => { if (bytes == null) return ""; @@ -24,11 +24,18 @@ const ComponentDownloadControl = ({ component, onStatusChange }) => { const meta = component.metadata || {}; const [downloaded, setDownloaded] = useState(Boolean(component.downloaded)); const [downloading, setDownloading] = useState(false); + const pollerIdRef = useRef(null); useEffect(() => { setDownloaded(Boolean(component.downloaded)); }, [component.name, component.downloaded]); + useEffect(() => { + return () => { + if (pollerIdRef.current != null) stopJobPolling(pollerIdRef.current); + }; + }, []); + if (!meta.requires_download) return null; const finish = (isDownloaded) => { @@ -41,9 +48,11 @@ const ComponentDownloadControl = ({ component, onStatusChange }) => { setDownloading(true); try { const { id } = await downloadComponent(component.name); + pollerIdRef.current = id; startJobPolling( id, async () => { + pollerIdRef.current = null; const status = await getComponentDownloadStatus(component.name); finish(status.downloaded); enqueueSnackbar(t("common:componentDownload.done"), { @@ -51,6 +60,7 @@ const ComponentDownloadControl = ({ component, onStatusChange }) => { }); }, () => { + pollerIdRef.current = null; finish(false); enqueueSnackbar(t("common:componentDownload.failed"), { variant: "error", @@ -66,8 +76,14 @@ const ComponentDownloadControl = ({ component, onStatusChange }) => { }; const handleDelete = async () => { - await deleteComponentDownload(component.name); - finish(false); + try { + await deleteComponentDownload(component.name); + finish(false); + } catch { + enqueueSnackbar(t("common:componentDownload.failed"), { + variant: "error", + }); + } }; if (downloading) { From 5c9caac164ce481eb0892a02b279f203156796e6 Mon Sep 17 00:00:00 2001 From: Irozuku Date: Wed, 1 Jul 2026 15:38:15 -0400 Subject: [PATCH 24/89] fix: persist Opus-MT tokenizer with runs so prediction survives download deletion --- .../hugging_face/base_opus_mt_transformer.py | 27 +++-- .../back/models/test_opus_mt_downloadable.py | 105 ++++++++++++++++++ 2 files changed, 124 insertions(+), 8 deletions(-) diff --git a/DashAI/back/models/hugging_face/base_opus_mt_transformer.py b/DashAI/back/models/hugging_face/base_opus_mt_transformer.py index 38eeecf04..a755cf80d 100644 --- a/DashAI/back/models/hugging_face/base_opus_mt_transformer.py +++ b/DashAI/back/models/hugging_face/base_opus_mt_transformer.py @@ -30,10 +30,13 @@ class OpusMtTransformerMixin(HFDownloadableMixin, TranslationModel): here so each language-pair subclass only needs to set class attributes. .. note:: - The pretrained weights must be downloaded first via ``download()`` - (this component requires a download). ``__init__`` loads the tokenizer - and model from the component's local folder; it does not fetch from - the Hugging Face Hub. + For fresh training the pretrained weights must be downloaded first via + ``download()`` (this component requires a download). ``__init__`` with + no ``pretrained_dir`` loads the tokenizer and model from the + component's local download folder; it does not fetch from the Hugging + Face Hub. When loading a previously saved run, ``pretrained_dir`` is + set to the run directory so the tokenizer is read from there, making + trained runs self-contained. """ MODEL_NAME: str = "" @@ -53,13 +56,19 @@ def hf_repos(cls): """ return [(cls.MODEL_NAME, "model")] if cls.MODEL_NAME else [] - def __init__(self, model=None, **kwargs): + def __init__(self, model=None, pretrained_dir: Optional[str] = None, **kwargs): """Initialize tokenizer and seq2seq model. Parameters ---------- model : transformers.PreTrainedModel or None Preloaded model to reuse instead of downloading weights. + pretrained_dir : str or None + Directory from which to load the tokenizer (and model weights when + ``model`` is ``None``). When ``None`` the component's download + folder is used, which is the correct path for fresh training. Pass + the run directory when restoring a saved run so the trained run + becomes self-contained and independent of the download folder. **kwargs Training hyperparameters forwarded to ``validate_and_transform``. """ @@ -73,8 +82,8 @@ def __init__(self, model=None, **kwargs): ) self.model_name = self.MODEL_NAME - local_dir = str(self._repo_dir(self.MODEL_NAME)) - self.tokenizer = AutoTokenizer.from_pretrained(local_dir) + source = pretrained_dir or str(self._repo_dir(self.MODEL_NAME)) + self.tokenizer = AutoTokenizer.from_pretrained(source) self.training_args = { "num_train_epochs": kwargs.get("num_train_epochs", 2), @@ -95,7 +104,7 @@ def __init__(self, model=None, **kwargs): if model is None: from transformers import AutoModelForSeq2SeqLM - self.model = AutoModelForSeq2SeqLM.from_pretrained(local_dir) + self.model = AutoModelForSeq2SeqLM.from_pretrained(source) else: self.model = model @@ -252,6 +261,7 @@ def save(self, filename: Union[str, "Path"]) -> None: save_dir.mkdir(parents=True, exist_ok=True) self.model.save_pretrained(save_dir) + self.tokenizer.save_pretrained(save_dir) config = AutoConfig.from_pretrained(save_dir) config.custom_params = { "num_train_epochs": self.training_args.get("num_train_epochs"), @@ -274,6 +284,7 @@ def load(cls, filename: Union[str, "Path"]): loaded_model = cls( model=model, + pretrained_dir=str(filename), num_train_epochs=custom_params.get("num_train_epochs"), batch_size=custom_params.get("batch_size"), learning_rate=custom_params.get("learning_rate"), diff --git a/tests/back/models/test_opus_mt_downloadable.py b/tests/back/models/test_opus_mt_downloadable.py index fc91adf81..f674b40d8 100644 --- a/tests/back/models/test_opus_mt_downloadable.py +++ b/tests/back/models/test_opus_mt_downloadable.py @@ -1,5 +1,7 @@ """Tests that OpusMtTransformerMixin subclasses expose download metadata.""" +from unittest.mock import MagicMock, patch + import pytest from kink import di @@ -45,3 +47,106 @@ def test_opus_mt_is_downloaded_uses_component_dir(component_root): repo_dir.mkdir(parents=True) (repo_dir / "config.json").write_text("{}") assert OpusMtEnESTransformer.is_downloaded() is True + + +# --------------------------------------------------------------------------- +# Self-contained run tests (no real weights required) +# --------------------------------------------------------------------------- + + +@pytest.fixture +def component_root_for_save_load(tmp_path): + """Inject a temporary COMPONENT_PATH so _repo_dir resolves without touching + the real filesystem or requiring a real download. + """ + sentinel = object() + old = di["config"] if "config" in di else sentinel # noqa: SIM401 + di["config"] = {"COMPONENT_PATH": str(tmp_path)} + yield tmp_path + if old is sentinel: + del di["config"] + else: + di["config"] = old + + +def test_save_persists_tokenizer(tmp_path, component_root_for_save_load): + """save() must call both model.save_pretrained and tokenizer.save_pretrained.""" + mock_tokenizer = MagicMock() + mock_model = MagicMock() + + with ( + patch("transformers.AutoTokenizer") as tok_cls, + patch("transformers.AutoModelForSeq2SeqLM"), + patch("transformers.AutoConfig") as cfg_cls, + ): + tok_cls.from_pretrained.return_value = mock_tokenizer + + instance = OpusMtEnESTransformer( + model=mock_model, + pretrained_dir=str(tmp_path), + num_train_epochs=1, + batch_size=2, + learning_rate=2e-5, + device="CPU", + weight_decay=0.01, + log_train_every_n_epochs=None, + log_train_every_n_steps=None, + log_validation_every_n_epochs=None, + log_validation_every_n_steps=None, + ) + instance.fitted = True + + save_dir = tmp_path / "run" + cfg_cls.from_pretrained.return_value = MagicMock() + + instance.save(save_dir) + + mock_model.save_pretrained.assert_called_once_with(save_dir) + mock_tokenizer.save_pretrained.assert_called_once_with(save_dir) + + +def test_load_tokenizer_from_run_dir_not_download_folder( + tmp_path, component_root_for_save_load +): + """load() must load the tokenizer from the run dir, not the component download + folder. This verifies that trained runs are self-contained. + """ + run_dir = tmp_path / "my_run" + run_dir.mkdir() + + mock_model = MagicMock() + mock_tokenizer = MagicMock() + mock_config = MagicMock() + mock_config.custom_params = { + "num_train_epochs": 1, + "batch_size": 2, + "learning_rate": 2e-5, + "device": "CPU", + "weight_decay": 0.01, + "fitted": True, + } + + with ( + patch("transformers.AutoTokenizer") as tok_cls, + patch("transformers.AutoModelForSeq2SeqLM") as model_cls, + patch("transformers.AutoConfig") as cfg_cls, + ): + tok_cls.from_pretrained.return_value = mock_tokenizer + model_cls.from_pretrained.return_value = mock_model + cfg_cls.from_pretrained.return_value = mock_config + + loaded = OpusMtEnESTransformer.load(run_dir) + + # The tokenizer must be loaded from the run dir. + tok_cls.from_pretrained.assert_called_once_with(str(run_dir)) + + # The component download folder (under COMPONENT_PATH) must NOT be used. + component_download_dir = str( + component_root_for_save_load / "OpusMtEnESTransformer" / "opus-mt-en-es" + ) + for call in tok_cls.from_pretrained.call_args_list: + assert call.args[0] != component_download_dir, ( + "tokenizer was loaded from the component download folder, not the run dir" + ) + + assert loaded.fitted is True From 2b2ffaea2a616dab7c047ab83502ff79782b6cb5 Mon Sep 17 00:00:00 2001 From: Irozuku Date: Wed, 1 Jul 2026 23:59:05 -0400 Subject: [PATCH 25/89] feat: block undownloaded models in the models list with inline download An undownloaded, download-required model now renders disabled in the ModelsRightBar list and cannot open the config dialog; a download control is shown beneath it and the list refreshes on completion. The in-dialog download control is removed (the list is the download surface); the disabled-create-button guard remains as a safety net. --- .../src/components/models/AddModelDialog.jsx | 7 --- .../src/components/models/ModelsRightBar.jsx | 59 ++++++++++++++++--- 2 files changed, 50 insertions(+), 16 deletions(-) diff --git a/DashAI/front/src/components/models/AddModelDialog.jsx b/DashAI/front/src/components/models/AddModelDialog.jsx index 80604a474..2bd46d5c1 100644 --- a/DashAI/front/src/components/models/AddModelDialog.jsx +++ b/DashAI/front/src/components/models/AddModelDialog.jsx @@ -27,7 +27,6 @@ import { createRun } from "../../api/run"; import { useTranslation } from "react-i18next"; import { useTourContext } from "../tour/TourProvider"; import { checkIfHaveOptimazers } from "../../utils/schema"; -import ComponentDownloadControl from "./model/ComponentDownloadControl"; /** * Dialog for adding a new model run to a session @@ -325,12 +324,6 @@ function AddModelDialog({ )} - {preselectedModelObject?.metadata?.requires_download && ( - - )} )} diff --git a/DashAI/front/src/components/models/ModelsRightBar.jsx b/DashAI/front/src/components/models/ModelsRightBar.jsx index f77ab0a0d..b4864cd07 100644 --- a/DashAI/front/src/components/models/ModelsRightBar.jsx +++ b/DashAI/front/src/components/models/ModelsRightBar.jsx @@ -7,6 +7,7 @@ import { useSnackbar } from "notistack"; import SideBar from "../threeSectionLayout/panelContainers/SideBar"; import { getComponents } from "../../api/component"; import ModelListItem from "./model/ModelListItem"; +import ComponentDownloadControl from "./model/ComponentDownloadControl"; import { useTranslation } from "react-i18next"; import { useTourContext } from "../tour/TourProvider"; import { useModels } from "./ModelsContext"; @@ -20,7 +21,7 @@ export default function ModelsRightBar({ onToggle }) { const [searchQuery, setSearchQuery] = useState(""); const [loading, setLoading] = useState(false); const { enqueueSnackbar } = useSnackbar(); - const { t } = useTranslation(["models"]); + const { t } = useTranslation(["models", "common"]); const { selectedSession: session, @@ -89,6 +90,11 @@ export default function ModelsRightBar({ onToggle }) { }); return; } + // A download-required model that has not been downloaded cannot be + // configured; it is blocked in the list with an inline download control. + if (model.metadata?.requires_download && !model.downloaded) { + return; + } selectModel(model); if (tourContext?.run && tourContext?.stepIndex === 2) { const waitForElement = () => { @@ -236,14 +242,49 @@ export default function ModelsRightBar({ onToggle }) { ) : ( - {filteredModels.map((model, index) => ( - handleModelClick(model)} - data-tour={index === 0 ? "first-model" : undefined} - /> - ))} + {filteredModels.map((model, index) => { + const needsDownload = + Boolean(model.metadata?.requires_download) && + !model.downloaded; + return ( + + handleModelClick(model) + } + data-tour={index === 0 ? "first-model" : undefined} + /> + {needsDownload && ( + { + if (isDownloaded) fetchModels(); + }} + /> + )} + + ); + })} )} From 06ec7edf2fa3ce93c88b8216e472d42b9b306460 Mon Sep 17 00:00:00 2001 From: Irozuku Date: Thu, 2 Jul 2026 00:00:46 -0400 Subject: [PATCH 26/89] feat: block undownloaded components in ComponentSelector cards with inline download Cards for a download-required, not-yet-downloaded component are disabled (unclickable, dimmed) and render an inline download control; an onDownloadChange callback lets parents refresh. Non-downloadable components are unchanged. --- .../components/custom/ComponentSelector.jsx | 107 +++++++++++------- 1 file changed, 64 insertions(+), 43 deletions(-) diff --git a/DashAI/front/src/components/custom/ComponentSelector.jsx b/DashAI/front/src/components/custom/ComponentSelector.jsx index b12423a2c..97be244b8 100644 --- a/DashAI/front/src/components/custom/ComponentSelector.jsx +++ b/DashAI/front/src/components/custom/ComponentSelector.jsx @@ -18,6 +18,7 @@ import { Check as CheckIcon, } from "@mui/icons-material"; import { useTranslation } from "react-i18next"; +import ComponentDownloadControl from "../models/model/ComponentDownloadControl"; const ALL_CATEGORY = "All"; const SEARCH_THRESHOLD = 10; @@ -41,8 +42,9 @@ function ComponentSelector({ flat = false, tourDataFor = null, tourDataMatchFn = null, + onDownloadChange = null, }) { - const { t } = useTranslation("custom"); + const { t } = useTranslation(["custom", "common"]); const [search, setSearch] = useState(""); const [activeCategory, setActiveCategory] = useState(ALL_CATEGORY); @@ -107,6 +109,8 @@ function ComponentSelector({ const renderCard = (component) => { const isSelected = selected?.name === component.name; const icon = getIcon?.(component); + const needsDownload = + Boolean(component.metadata?.requires_download) && !component.downloaded; const isCsvComponent = tourDataFor && (tourDataMatchFn @@ -116,61 +120,76 @@ function ComponentSelector({ handleSelect(component)} + onClick={needsDownload ? undefined : () => handleSelect(component)} data-tour={isCsvComponent ? tourDataFor : undefined} sx={{ p: 3, display: "flex", + flexDirection: "column", gap: 3, - alignItems: "flex-start", - cursor: "pointer", + cursor: needsDownload ? "not-allowed" : "pointer", border: 1, borderColor: isSelected ? "primary.main" : "divider", bgcolor: isSelected ? "action.selected" : "background.paper", + opacity: needsDownload ? 0.6 : 1, transition: "border-color 0.15s, background 0.15s", - "&:hover": { borderColor: "secondary.main" }, + "&:hover": { + borderColor: needsDownload ? "divider" : "secondary.main", + }, }} > - {icon && ( - - {icon} + + {icon && ( + + {icon} + + )} + + + {getLabel(component)} + + + {getDescription(component, t("noDescriptionAvailable"))} + - )} - - - {getLabel(component)} - - - {getDescription(component, t("noDescriptionAvailable"))} - + {isSelected && ( + + )} - {isSelected && ( - + {needsDownload && ( + e.stopPropagation()}> + { + if (isDownloaded) onDownloadChange?.(component); + }} + /> + )} ); @@ -358,7 +377,9 @@ ComponentSelector.propTypes = { emptyText: PropTypes.string, getIcon: PropTypes.func, flat: PropTypes.bool, + tourDataFor: PropTypes.string, tourDataMatchFn: PropTypes.func, + onDownloadChange: PropTypes.func, }; export default ComponentSelector; From 4da81af7f4e1aa27144a5a87fc6be20ed9a8ffa5 Mon Sep 17 00:00:00 2001 From: Irozuku Date: Thu, 2 Jul 2026 00:11:19 -0400 Subject: [PATCH 27/89] feat: expose download metadata and gate generative models Add BaseGenerativeModel.get_metadata exposing requires_download and download_size_bytes, a 409/422 download gate in upload_generative_session, and a defensive JobError in GenerativeJob.run for undownloaded models. --- .../api_v1/endpoints/generative_session.py | 19 +++++ DashAI/back/job/generative_job.py | 12 +++ DashAI/back/models/base_generative_model.py | 17 ++++- .../test_generative_session_download_gate.py | 73 +++++++++++++++++++ .../test_generative_download_metadata.py | 35 +++++++++ 5 files changed, 155 insertions(+), 1 deletion(-) create mode 100644 tests/back/api/test_generative_session_download_gate.py create mode 100644 tests/back/models/test_generative_download_metadata.py diff --git a/DashAI/back/api/api_v1/endpoints/generative_session.py b/DashAI/back/api/api_v1/endpoints/generative_session.py index 8211dbcf9..982b2b7c3 100644 --- a/DashAI/back/api/api_v1/endpoints/generative_session.py +++ b/DashAI/back/api/api_v1/endpoints/generative_session.py @@ -38,6 +38,13 @@ async def upload_generative_session( with session_factory() as db: try: + # Guard: unknown model name -> 422 + if params.model_name not in component_registry: + raise HTTPException( + status_code=status.HTTP_422_UNPROCESSABLE_ENTITY, + detail=f"Unknown model '{params.model_name}'", + ) + # Check if the model is registered try: model_class = component_registry[params.model_name]["class"] @@ -47,6 +54,18 @@ async def upload_generative_session( detail=f"Model {params.model_name} is not registered.", ) from e + # Guard: model requires download but has not been downloaded -> 409 + entry = component_registry[params.model_name] + if getattr(entry["class"], "REQUIRES_DOWNLOAD", False) and not entry.get( + "downloaded", False + ): + raise HTTPException( + status_code=status.HTTP_409_CONFLICT, + detail=( + f"Model {params.model_name} must be downloaded before use." + ), + ) + # Check if the model is a subclass of GenerativeModel if not issubclass(model_class, BaseGenerativeModel): raise HTTPException( diff --git a/DashAI/back/job/generative_job.py b/DashAI/back/job/generative_job.py index e52c58a72..bbf5de47e 100644 --- a/DashAI/back/job/generative_job.py +++ b/DashAI/back/job/generative_job.py @@ -151,8 +151,20 @@ def run( model_class = component_registry[generative_session.model_name][ "class" ] + if ( + getattr(model_class, "REQUIRES_DOWNLOAD", False) + and not model_class.is_downloaded() + ): + raise JobError( + f"Model {generative_session.model_name} is not downloaded." + " Download it before use." + ) params = generative_session.parameters model: BaseGenerativeModel = model_class(**params) + except JobError: + generative_process.set_status_as_error() + db.commit() + raise except Exception as e: log.exception(e) generative_process.set_status_as_error() diff --git a/DashAI/back/models/base_generative_model.py b/DashAI/back/models/base_generative_model.py index cff0e18e1..62dbcb9af 100644 --- a/DashAI/back/models/base_generative_model.py +++ b/DashAI/back/models/base_generative_model.py @@ -1,5 +1,5 @@ from abc import ABCMeta, abstractmethod -from typing import Any, Final, List, Tuple, Union +from typing import Any, Dict, Final, List, Tuple, Union from DashAI.back.config_object import ConfigObject @@ -15,6 +15,21 @@ class BaseGenerativeModel(ConfigObject, metaclass=ABCMeta): TYPE: Final[str] = "GenerativeModel" + @classmethod + def get_metadata(cls) -> Dict[str, Any]: + """Get metadata values for the current generative model. + + Returns + ------- + Dict[str, Any] + Dictionary indicating whether the model requires a download + before use and the expected download size in bytes. + """ + metadata: Dict[str, Any] = {} + metadata["requires_download"] = bool(getattr(cls, "REQUIRES_DOWNLOAD", False)) + metadata["download_size_bytes"] = getattr(cls, "DOWNLOAD_SIZE_BYTES", None) + return metadata + @abstractmethod def __init__(self, **kwargs): """Initialize the generative model with configuration parameters. diff --git a/tests/back/api/test_generative_session_download_gate.py b/tests/back/api/test_generative_session_download_gate.py new file mode 100644 index 000000000..b284200a2 --- /dev/null +++ b/tests/back/api/test_generative_session_download_gate.py @@ -0,0 +1,73 @@ +"""Tests for the generative-session creation download gate.""" + +from kink import di + +_SESSION_PAYLOAD_BASE = { + "parameters": {}, + "name": "gen-gate-test-session", + "description": None, +} + + +class _FakeDownloadableGenerativeModel: + """Minimal stub: download-required generative model that is not downloaded.""" + + REQUIRES_DOWNLOAD = True + + @classmethod + def is_downloaded(cls): + return False + + +class _FakeGenerativeRegistry: + """Registry wrapper that injects FakeDownloadableGenerativeModel.""" + + def __init__(self, real): + self._real = real + + def __getitem__(self, name): + if name == "FakeDownloadableGenerativeModel": + return { + "class": _FakeDownloadableGenerativeModel, + "downloaded": False, + } + return self._real[name] + + def get_components_by_types(self, select=None, ignore=None): + return self._real.get_components_by_types(select=select, ignore=ignore) + + def __contains__(self, name): + return name == "FakeDownloadableGenerativeModel" or name in self._real + + +def test_upload_generative_session_rejects_undownloaded_model(client): + """Creating a session for a not-yet-downloaded model must return HTTP 409.""" + old = di["component_registry"] + di["component_registry"] = _FakeGenerativeRegistry(old) + try: + resp = client.post( + "/api/v1/generative-session/", + json={ + "model_name": "FakeDownloadableGenerativeModel", + "task_name": "TextToTextGenerationTask", + **_SESSION_PAYLOAD_BASE, + }, + ) + finally: + di["component_registry"] = old + + assert resp.status_code == 409 + assert "download" in resp.json()["detail"].lower() + + +def test_upload_generative_session_unknown_model_422(client): + """POSTing a session with an unregistered model_name must return HTTP 422.""" + resp = client.post( + "/api/v1/generative-session/", + json={ + "model_name": "__totally_bogus_generative_model_xyz__", + "task_name": "TextToTextGenerationTask", + **_SESSION_PAYLOAD_BASE, + }, + ) + assert resp.status_code == 422 diff --git a/tests/back/models/test_generative_download_metadata.py b/tests/back/models/test_generative_download_metadata.py new file mode 100644 index 000000000..09ff25d01 --- /dev/null +++ b/tests/back/models/test_generative_download_metadata.py @@ -0,0 +1,35 @@ +"""Tests for BaseGenerativeModel.get_metadata download fields.""" + +from DashAI.back.dependencies.downloads.downloadable import HFDownloadableMixin +from DashAI.back.models.base_generative_model import BaseGenerativeModel + + +class _PlainGenerativeModel(BaseGenerativeModel): + def __init__(self, **kwargs): + pass + + def generate(self, input): + return [] + + +class _DownloadableGenerativeModel(HFDownloadableMixin, BaseGenerativeModel): + HF_REPOS = [("owner/x", "model")] + DOWNLOAD_SIZE_BYTES = 1234 + + def __init__(self, **kwargs): + pass + + def generate(self, input): + return [] + + +def test_plain_generative_model_not_downloadable(): + meta = _PlainGenerativeModel.get_metadata() + assert meta["requires_download"] is False + assert meta["download_size_bytes"] is None + + +def test_downloadable_generative_model_metadata(): + meta = _DownloadableGenerativeModel.get_metadata() + assert meta["requires_download"] is True + assert meta["download_size_bytes"] == 1234 From f284a7e97e34bc71512cdfaedc3848cdef1ddb85 Mon Sep 17 00:00:00 2001 From: Irozuku Date: Thu, 2 Jul 2026 00:23:20 -0400 Subject: [PATCH 28/89] feat: support single-file (allow_patterns) downloads in HFDownloadableMixin --- .../dependencies/downloads/downloadable.py | 98 +++++++++++++++++-- tests/back/downloads/test_downloadable.py | 42 ++++++++ 2 files changed, 132 insertions(+), 8 deletions(-) diff --git a/DashAI/back/dependencies/downloads/downloadable.py b/DashAI/back/dependencies/downloads/downloadable.py index 366d1c814..33bef1a25 100644 --- a/DashAI/back/dependencies/downloads/downloadable.py +++ b/DashAI/back/dependencies/downloads/downloadable.py @@ -9,7 +9,7 @@ import logging import pathlib import shutil -from typing import Callable, List, Optional, Tuple +from typing import Callable, List, Optional, Tuple, Union from huggingface_hub import snapshot_download from kink import di @@ -69,35 +69,117 @@ class HFDownloadableMixin(DownloadableMixin): Each repo is downloaded into ``component_dir()/``. Subclasses set ``HF_REPOS`` or, for a dynamic repo (e.g. derived from a per-subclass ``MODEL_NAME``), override ``hf_repos``. + + ``HF_REPOS`` entries accept two shapes: + + * ``(repo_id, repo_type)`` -- full snapshot download (original behavior). + * ``(repo_id, repo_type, allow_patterns)`` -- partial download; only files + matching the glob patterns in ``allow_patterns`` are fetched. """ - HF_REPOS: List[Tuple[str, str]] = [] + HF_REPOS: List[Union[Tuple[str, str], Tuple[str, str, List[str]]]] = [] @classmethod - def hf_repos(cls) -> List[Tuple[str, str]]: - """Return the ``(repo_id, repo_type)`` pairs this component needs.""" + def hf_repos(cls) -> List[Union[Tuple[str, str], Tuple[str, str, List[str]]]]: + """Return the repo entries this component needs. + + Returns + ------- + list of tuple + Each entry is either ``(repo_id, repo_type)`` or + ``(repo_id, repo_type, allow_patterns)``. + """ return list(cls.HF_REPOS) + @classmethod + def _unpack_entry( + cls, + entry: Union[Tuple[str, str], Tuple[str, str, List[str]]], + ) -> Tuple[str, str, Optional[List[str]]]: + """Normalise a repo entry into ``(repo_id, repo_type, allow_patterns)``. + Parameters + ---------- + entry : tuple + Either a 2-tuple ``(repo_id, repo_type)`` or a 3-tuple + ``(repo_id, repo_type, allow_patterns)``. + + Returns + ------- + tuple of (str, str, list[str] or None) + ``repo_id``, ``repo_type``, and ``allow_patterns`` (``None`` when + the entry was a 2-tuple, meaning a full snapshot download). + + Raises + ------ + ValueError + If ``entry`` has a length other than 2 or 3. + """ + if len(entry) == 2: + rid, rtype = entry + return rid, rtype, None + if len(entry) == 3: + rid, rtype, patterns = entry + return rid, rtype, patterns + raise ValueError( + f"HF_REPOS entries must be 2- or 3-tuples; " + f"got length {len(entry)}: {entry!r}" + ) + @classmethod def _repo_dir(cls, repo_id: str) -> pathlib.Path: - """Return the local directory for a single repo under component_dir().""" + """Return the local directory for a single repo under component_dir(). + + Parameters + ---------- + repo_id : str + HuggingFace repo identifier, e.g. ``"owner/model-name"``. + + Returns + ------- + pathlib.Path + ``component_dir()/``. + """ return cls.component_dir() / repo_id.split("/")[-1] @classmethod def is_downloaded(cls) -> bool: + """Return whether all repo directories exist and are non-empty. + + Returns + ------- + bool + ``True`` when every repo listed in ``hf_repos()`` has a non-empty + local directory; ``False`` otherwise (including when the list is + empty). + """ repos = cls.hf_repos() return bool(repos) and all( cls._repo_dir(rid).is_dir() and any(cls._repo_dir(rid).iterdir()) - for rid, _ in repos + for rid, *_ in repos ) @classmethod def download(cls, report: Optional[ProgressReporter] = None) -> None: - for rid, rtype in cls.hf_repos(): + """Download all repos listed in ``hf_repos()`` into ``component_dir()``. + + Parameters + ---------- + report : ProgressReporter, optional + Callback invoked before each repo download with + ``report(None, "Downloading ")``. ``None`` means no + progress reporting. + """ + for entry in cls.hf_repos(): + rid, rtype, allow_patterns = cls._unpack_entry(entry) target = cls._repo_dir(rid) target.mkdir(parents=True, exist_ok=True) if report is not None: # snapshot_download exposes no aggregate byte count, so progress # is reported as indeterminate (None) with a phase message. report(None, f"Downloading {rid}") - snapshot_download(repo_id=rid, repo_type=rtype, local_dir=str(target)) + kwargs = {} + if allow_patterns is not None: + kwargs["allow_patterns"] = allow_patterns + snapshot_download( + repo_id=rid, repo_type=rtype, local_dir=str(target), **kwargs + ) diff --git a/tests/back/downloads/test_downloadable.py b/tests/back/downloads/test_downloadable.py index 847c74666..6013a276d 100644 --- a/tests/back/downloads/test_downloadable.py +++ b/tests/back/downloads/test_downloadable.py @@ -71,3 +71,45 @@ def test_delete_removes_component_dir(components_root): _populate(components_root, _Dummy, "model-a") _Dummy.delete() assert not (components_root / "_Dummy").exists() + + +# --------------------------------------------------------------------------- +# 3-tuple (allow_patterns) support +# --------------------------------------------------------------------------- + + +class _DummyPartial(dl.HFDownloadableMixin): + HF_REPOS = [("owner/model-a", "model", ["*8_0.gguf"])] + + +def test_download_3tuple_passes_allow_patterns(components_root): + with mock.patch.object(dl, "snapshot_download") as snap: + _DummyPartial.download(lambda frac, msg: None) + snap.assert_called_once_with( + repo_id="owner/model-a", + repo_type="model", + local_dir=str(components_root / "_DummyPartial" / "model-a"), + allow_patterns=["*8_0.gguf"], + ) + + +def test_download_2tuple_no_allow_patterns(components_root): + with mock.patch.object(dl, "snapshot_download") as snap: + _Dummy.download(lambda frac, msg: None) + _call_kwargs = snap.call_args.kwargs + assert "allow_patterns" not in _call_kwargs + + +def test_is_downloaded_3tuple_true_when_present(components_root): + _populate(components_root, _DummyPartial, "model-a") + assert _DummyPartial.is_downloaded() is True + + +def test_is_downloaded_3tuple_false_when_absent(components_root): + assert _DummyPartial.is_downloaded() is False + + +def test_is_downloaded_3tuple_false_when_empty_dir(components_root): + d = components_root / "_DummyPartial" / "model-a" + d.mkdir(parents=True) + assert _DummyPartial.is_downloaded() is False From a07686c1d6ec107902f2c2298cd34204c6414584 Mon Sep 17 00:00:00 2001 From: Irozuku Date: Thu, 2 Jul 2026 00:37:57 -0400 Subject: [PATCH 29/89] feat: split GGUF text models into per-checkpoint downloadable components Replace the four enum-based model classes (QwenModel, SmolLMModel, LlamaModel, MistralModel) with nine thin subclasses of GGUFTextGenerationModel (HFDownloadableMixin). Each subclass represents one checkpoint and carries REPO_ID, GGUF_PATTERN, and DOWNLOAD_SIZE_BYTES so the download machinery can manage it independently. The shared schema and __init__/generate logic live in gguf_text_generation_base.py. --- DashAI/back/initial_components.py | 34 +- .../hugging_face/gguf_text_generation_base.py | 389 ++++++++++++ .../back/models/hugging_face/llama_model.py | 587 +++++------------- .../back/models/hugging_face/mistral_model.py | 497 +++------------ DashAI/back/models/hugging_face/qwen_model.py | 534 +++------------- .../back/models/hugging_face/smol_lm_model.py | 507 +++------------ .../back/models/test_gguf_text_generation.py | 89 +++ 7 files changed, 927 insertions(+), 1710 deletions(-) create mode 100644 DashAI/back/models/hugging_face/gguf_text_generation_base.py create mode 100644 tests/back/models/test_gguf_text_generation.py diff --git a/DashAI/back/initial_components.py b/DashAI/back/initial_components.py index d12cbafa9..54e166f7e 100644 --- a/DashAI/back/initial_components.py +++ b/DashAI/back/initial_components.py @@ -145,10 +145,17 @@ from DashAI.back.models.hugging_face.deberta_v3_transformer import DebertaV3Transformer from DashAI.back.models.hugging_face.distilbert_transformer import DistilBertTransformer from DashAI.back.models.hugging_face.electra_transformer import ElectraTransformer -from DashAI.back.models.hugging_face.llama_model import LlamaModel +from DashAI.back.models.hugging_face.llama_model import ( + Llama31_8BInstruct, + Llama32_1BInstruct, + Llama32_3BInstruct, +) from DashAI.back.models.hugging_face.m2m100_transformer import M2M100Transformer from DashAI.back.models.hugging_face.minilm_transformer import MiniLMTransformer -from DashAI.back.models.hugging_face.mistral_model import MistralModel +from DashAI.back.models.hugging_face.mistral_model import ( + Mistral7BInstructV03, + MistralNemoInstruct2407, +) from DashAI.back.models.hugging_face.mixtral_model import MixtralModel from DashAI.back.models.hugging_face.modernbert_transformer import ModernBertTransformer from DashAI.back.models.hugging_face.multilingual_bert_transformer import ( @@ -174,7 +181,10 @@ OpusMtFrEnTransformer, ) from DashAI.back.models.hugging_face.pixart_sigma_model import PixArtSigmaModel -from DashAI.back.models.hugging_face.qwen_model import QwenModel +from DashAI.back.models.hugging_face.qwen_model import ( + Qwen25_05BInstruct, + Qwen25_15BInstruct, +) from DashAI.back.models.hugging_face.roberta_transformer import RobertaTransformer from DashAI.back.models.hugging_face.sd15_depth_controlnet_model import ( SD15DepthControlNetModel, @@ -189,7 +199,10 @@ SDXLCannyControlNetModel, ) from DashAI.back.models.hugging_face.sdxl_turbo_model import SDXLTurboModel -from DashAI.back.models.hugging_face.smol_lm_model import SmolLMModel +from DashAI.back.models.hugging_face.smol_lm_model import ( + SmolLM2_17BInstruct, + SmolLM2_360MInstruct, +) from DashAI.back.models.hugging_face.stable_diffusion_v1_depth_controlnet import ( StableDiffusionXLV1ControlNet, ) @@ -346,11 +359,14 @@ def get_initial_components(): LinearRegression, LinearSVCClassifier, LinearSVR, - LlamaModel, + Llama31_8BInstruct, + Llama32_1BInstruct, + Llama32_3BInstruct, LogisticRegression, M2M100Transformer, MiniLMTransformer, - MistralModel, + Mistral7BInstructV03, + MistralNemoInstruct2407, MixtralModel, MultilingualBertTransformer, MLPClassifier, @@ -364,7 +380,8 @@ def get_initial_components(): OpusMtEsENTransformer, OpusMtFrEnTransformer, PixArtSigmaModel, - QwenModel, + Qwen25_05BInstruct, + Qwen25_15BInstruct, RandomForestClassifier, RobertaTransformer, RandomForestRegression, @@ -375,7 +392,8 @@ def get_initial_components(): SDXLCannyControlNetModel, SDXLTurboModel, SGDClassifier, - SmolLMModel, + SmolLM2_360MInstruct, + SmolLM2_17BInstruct, StableDiffusionV2Model, StableDiffusionV3Model, StableDiffusionXLModel, diff --git a/DashAI/back/models/hugging_face/gguf_text_generation_base.py b/DashAI/back/models/hugging_face/gguf_text_generation_base.py new file mode 100644 index 000000000..189f15629 --- /dev/null +++ b/DashAI/back/models/hugging_face/gguf_text_generation_base.py @@ -0,0 +1,389 @@ +"""Shared base for GGUF-backed text-generation models loaded via llama.cpp.""" + +from typing import List, Optional, Union + +from DashAI.back.core.schema_fields import ( + BaseSchema, + enum_field, + float_field, + int_field, + schema_field, +) +from DashAI.back.core.utils import MultilingualString +from DashAI.back.dependencies.downloads.downloadable import HFDownloadableMixin +from DashAI.back.models.text_to_text_generation_model import ( + TextToTextGenerationTaskModel, +) +from DashAI.back.models.utils import ( + LLAMA_DEVICE_ENUM, + LLAMA_DEVICE_PLACEHOLDER, + LLAMA_DEVICE_TO_IDX, +) + + +class GGUFTextGenerationSchema(BaseSchema): + """Schema for GGUF-based text-generation model hyperparameters. + + All GGUF checkpoint subclasses share this schema. The schema controls + generation length, sampling randomness, repetition penalty, context budget, + and the hardware device used by llama.cpp. + """ + + max_tokens: schema_field( + int_field(ge=1), + placeholder=100, + description=MultilingualString( + en=( + "Maximum number of new tokens the model will generate per response. " + "Roughly 1 token ≈ 0.75 English words. Set to 100-200 for short " + "answers, 500-1000 for detailed explanations or code. Must not " + "exceed the context window minus the prompt length." + ), + es=( + "Número máximo de tokens nuevos que el modelo generará por respuesta. " + "Aproximadamente 1 token ≈ 0.75 palabras en español. Use 100-200 " + "para respuestas cortas, 500-1000 para explicaciones detalladas o " + "código. No debe superar la ventana de contexto menos la longitud " + "del prompt." + ), + pt=( + "Número máximo de tokens novos que o modelo gerará por resposta. " + "Aproximadamente 1 token ≈ 0.75 palavras em português. Use 100-200 " + "para respostas curtas, 500-1000 para explicações detalhadas ou " + "código. Não deve exceder a janela de contexto menos o comprimento " + "do prompt." + ), + de=( + "Maximale Anzahl neuer Token, die das Modell pro Antwort erzeugt. " + "Ungefähr 1 Token ≈ 0,75 englische Wörter. 100-200 für kurze " + "Antworten, 500-1000 für ausführliche Erklärungen oder Code. " + "Darf die Kontextfenstergröße abzüglich der Prompt-Länge nicht " + "überschreiten." + ), + zh=( + "模型每次响应生成的最大新 token 数量。" + "大约 1 token 约等于 0.75 个英文单词。短答案设为 100-200，" + "详细说明或代码设为 500-1000。不得超过上下文窗口减去提示词长度的值。" + ), + ), + alias=MultilingualString( + en="Max tokens", + es="Tokens máximos", + pt="Tokens máximos", + de="Maximale neue Token", + zh="最大 token 数", + ), + ) # type: ignore + + temperature: schema_field( + float_field(ge=0.0, le=1.0), + placeholder=0.7, + description=MultilingualString( + en=( + "Sampling temperature controlling output randomness (range 0.0-1.0). " + "At 0.0 the model always picks the most likely token (greedy, fully " + "deterministic). Around 0.7 is a good balance for conversational " + "tasks. At 1.0 outputs are maximally varied and unpredictable." + ), + es=( + "Temperatura de muestreo que controla la aleatoriedad de la salida " + "(rango 0.0-1.0). En 0.0 el modelo siempre elige el token más " + "probable (greedy, totalmente determinista). Alrededor de 0.7 es " + "un buen equilibrio para tareas conversacionales. En 1.0 las " + "salidas son máximamente variadas e impredecibles." + ), + pt=( + "Temperatura de amostragem que controla a aleatoriedade da saída " + "(intervalo 0.0-1.0). Em 0.0 o modelo sempre escolhe o token mais " + "provável (greedy, totalmente determinístico). Em torno de 0.7 é " + "um bom equilíbrio para tarefas conversacionais. Em 1.0 as " + "saídas são maximamente variadas e imprevisíveis." + ), + de=( + "Stichprobentemperatur zur Steuerung der Ausgabezufälligkeit (0.0-1.0)." + "Bei 0.0 wählt das Modell stets den wahrscheinlichsten Token (greedy, " + "vollständig deterministisch). Um 0.7 ist ein gutes Gleichgewicht für " + "Konversationsaufgaben. Bei 1.0 sind Ausgaben maximal variiert und " + "unvorhersehbar." + ), + zh=( + "控制输出随机性的采样温度（范围 0.0-1.0）。" + "0.0 时模型始终选择最可能的 token（贪心，完全确定性）。" + "0.7 左右是对话任务的良好平衡点。1.0 时输出变化最大，不可预测。" + ), + ), + alias=MultilingualString( + en="Temperature", + es="Temperatura", + pt="Temperatura", + de="Temperatur", + zh="温度", + ), + ) # type: ignore + + frequency_penalty: schema_field( + float_field(ge=0.0, le=2.0), + placeholder=0.1, + description=MultilingualString( + en=( + "Penalizes tokens that have already appeared in the output based on " + "how often they occur (range 0.0-2.0). At 0.0 there is no penalty " + "and the model may repeat itself. Values around 0.1-0.3 gently " + "discourage repetition. High values (1.5+) strongly prevent reuse " + "of any word, which may produce less coherent text." + ), + es=( + "Penaliza los tokens que ya aparecieron en la salida según su " + "frecuencia (rango 0.0-2.0). En 0.0 no hay penalización y el modelo " + "puede repetirse. Valores en torno a 0.1-0.3 desincentivan " + "suavemente la repetición. Valores altos (1.5+) previenen " + "fuertemente la reutilización de palabras, lo que puede producir " + "texto menos coherente." + ), + pt=( + "Penaliza os tokens que já apareceram na saída com base em " + "sua frequência (intervalo 0.0-2.0). Em 0.0 não há penalização e o " + "modelo pode se repetir. Valores em torno de 0.1-0.3 desestimulam " + "suavemente a repetição. Valores altos (1.5+) impedem fortemente a " + "reutilização de palavras, o que pode produzir texto menos coerente." + ), + de=( + "Bestraft Token, die bereits in der Ausgabe erschienen sind, " + "basierend auf ihrer Häufigkeit (0.0-2.0). Bei 0.0 gibt es keine " + "Strafe und das Modell kann sich wiederholen. Werte um 0.1-0.3 " + "hemmen Wiederholungen sanft. Hohe Werte (1.5+) verhindern die " + "Wiederverwendung von Wörtern stark, was zu weniger kohärentem Text " + "führen kann." + ), + zh=( + "根据 token 在输出中出现的频率对其进行惩罚（范围 0.0-2.0）。" + "0.0 时无惩罚，模型可能重复输出。0.1-0.3 左右可轻微抑制重复。" + "高值（1.5+）会强烈阻止任何词的复用，可能导致文本连贯性下降。" + ), + ), + alias=MultilingualString( + en="Frequency penalty", + es="Penalización de frecuencia", + pt="Penalização de frequência", + de="Häufigkeitsstrafe", + zh="频率惩罚", + ), + ) # type: ignore + + context_window: schema_field( + int_field(ge=1, le=131072), + placeholder=512, + description=MultilingualString( + en=( + "Total token budget for a single forward pass, including both the " + "input prompt and the generated response. Larger values allow longer " + "conversations but consume more RAM/VRAM. Llama 3.1 supports up to " + "128K tokens natively; Llama 3.2 models support up to 128K tokens." + ), + es=( + "Presupuesto total de tokens para una sola pasada, incluyendo tanto " + "el prompt de entrada como la respuesta generada. Valores más altos " + "permiten conversaciones más largas pero consumen más RAM/VRAM. " + "Llama 3.1 soporta hasta 128K tokens de forma nativa; los modelos " + "Llama 3.2 soportan hasta 128K tokens." + ), + pt=( + "Orçamento total de tokens para uma única passagem, incluindo tanto " + "o prompt de entrada quanto a resposta gerada. Valores maiores " + "permitem conversas mais longas mas consomem mais RAM/VRAM. " + "Llama 3.1 suporta até 128K tokens nativamente; os modelos " + "Llama 3.2 suportam até 128K tokens." + ), + de=( + "Gesamtes Token-Budget für einen einzelnen Vorwärtsdurchlauf, " + "einschließlich Eingabe-Prompt und generierter Antwort. Größere Werte " + "ermöglichen längere Gespräche, verbrauchen jedoch mehr RAM/VRAM. " + "Llama 3.1 unterstützt nativ bis zu 128K Token; Llama-3.2-Modelle " + "unterstützen ebenfalls bis zu 128K Token." + ), + zh=( + "单次前向传播的总 token 预算，包含输入提示和生成响应。" + "较大的值允许更长的对话，但会消耗更多 RAM/VRAM。" + "Llama 3.1 原生支持最多 128K token；" + "Llama 3.2 模型同样支持最多 128K token。" + ), + ), + alias=MultilingualString( + en="Context window", + es="Ventana de contexto", + pt="Janela de contexto", + de="Kontextfenster", + zh="上下文窗口", + ), + ) # type: ignore + + device: schema_field( + enum_field(enum=LLAMA_DEVICE_ENUM), + placeholder=LLAMA_DEVICE_PLACEHOLDER, + description=MultilingualString( + en=( + "Hardware device for llama.cpp inference. 'CPU' runs the model " + "fully in RAM with no GPU requirement. Selecting a GPU option " + "offloads all layers for faster inference, setting n_gpu_layers=-1 " + "so every transformer layer is GPU-accelerated." + ), + es=( + "Dispositivo de hardware para la inferencia con llama.cpp. 'CPU' " + "ejecuta el modelo completamente en RAM sin requisito de GPU. " + "Seleccionar una opción de GPU descarga todas las capas para " + "inferencia más rápida, estableciendo n_gpu_layers=-1 para que " + "cada capa del transformer sea acelerada por GPU." + ), + pt=( + "Dispositivo de hardware para inferência com llama.cpp. 'CPU' " + "executa o modelo completamente na RAM sem requisito de GPU. " + "Selecionar uma opção de GPU descarrega todas as camadas para " + "inferência mais rápida, definindo n_gpu_layers=-1 para que " + "cada camada do transformer seja acelerada por GPU." + ), + de=( + "Hardware-Gerät für die llama.cpp-Inferenz. 'CPU' führt das Modell " + "vollständig im RAM ohne GPU-Anforderung aus. Eine GPU-Option " + "lagert alle Schichten für schnellere Inferenz aus und setzt " + "n_gpu_layers=-1, damit jede Transformer-Schicht GPU-beschleunigt wird." + ), + zh=( + "llama.cpp 推理所使用的硬件设备。'CPU' 完全在内存中运行模型，无需 GPU。" + "选择 GPU 选项会将所有层卸载以加快推理速度，" + "并设置 n_gpu_layers=-1 使每个 Transformer 层均由 GPU 加速。" + ), + ), + alias=MultilingualString( + en="Device", + es="Dispositivo", + pt="Dispositivo", + de="Gerät", + zh="设备", + ), + ) # type: ignore + + +class GGUFTextGenerationModel(HFDownloadableMixin, TextToTextGenerationTaskModel): + """Base class for GGUF quantized text-generation models loaded via llama.cpp. + + Each concrete subclass represents one specific checkpoint and sets the class + attributes ``REPO_ID``, ``GGUF_PATTERN``, and ``DOWNLOAD_SIZE_BYTES``. The + base class provides the shared ``hf_repos`` classmethod, a helper that locates + the downloaded GGUF file on disk, the ``__init__`` that loads the model, and + the ``generate`` method. + + Subclasses must NOT override ``hf_repos`` or ``_local_gguf_path`` unless they + need non-standard repo layout. + """ + + REPO_ID: str = "" + GGUF_PATTERN: str = "" + DOWNLOAD_SIZE_BYTES: Optional[int] = None + COMPATIBLE_COMPONENTS = ["TextToTextGenerationTask"] + SCHEMA = GGUFTextGenerationSchema + + @classmethod + def hf_repos( + cls, + ) -> List[Union[tuple, tuple]]: + """Return the single HuggingFace repo entry for this checkpoint. + + Returns + ------- + list of tuple + A list containing one 3-tuple ``(repo_id, "model", [gguf_pattern])`` + when ``REPO_ID`` is set, or an empty list otherwise. + """ + if cls.REPO_ID: + return [(cls.REPO_ID, "model", [cls.GGUF_PATTERN])] + return [] + + @classmethod + def _local_gguf_path(cls): + """Locate the downloaded GGUF file within this component's repo directory. + + Returns + ------- + pathlib.Path + Absolute path to the first ``*.gguf`` file found under the repo + directory for ``REPO_ID``. + + Raises + ------ + StopIteration + If no ``*.gguf`` file exists inside the repo directory. + """ + return next(iter(cls._repo_dir(cls.REPO_ID).glob("*.gguf"))) + + def __init__(self, **kwargs): + """Load a GGUF checkpoint from disk and initialise the llama.cpp model. + + Parameters + ---------- + **kwargs : dict + max_tokens : int, optional + Maximum number of new tokens to generate per response. Default 100. + temperature : float, optional + Sampling temperature in [0.0, 1.0]. Default 0.7. + frequency_penalty : float, optional + Token-frequency penalty in [0.0, 2.0]. Default 0.1. + context_window : int, optional + Total token budget for a single forward pass. Default 512. + device : str, optional + Target device from ``LLAMA_DEVICE_ENUM``. CPU runs in RAM only; + a GPU label enables full GPU offload via ``n_gpu_layers=-1``. + + Raises + ------ + RuntimeError + If ``llama-cpp-python`` is not installed. + """ + try: + from llama_cpp import Llama + except ImportError as e: + raise RuntimeError( + "llama-cpp-python is not installed. " + "Please install it to use this model." + ) from e + + kwargs = self.validate_and_transform(kwargs) + self.max_tokens = kwargs.pop("max_tokens", 100) + self.temperature = kwargs.pop("temperature", 0.7) + self.frequency_penalty = kwargs.pop("frequency_penalty", 0.1) + self.n_ctx = kwargs.pop("context_window", 512) + + device_val = kwargs.get("device") + use_gpu = LLAMA_DEVICE_TO_IDX.get(device_val, -1) >= 0 + main_gpu = LLAMA_DEVICE_TO_IDX.get(device_val, 0) if use_gpu else 0 + + self.model = Llama( + model_path=str(self._local_gguf_path()), + verbose=True, + n_ctx=self.n_ctx, + n_gpu_layers=-1 if use_gpu else 0, + main_gpu=main_gpu, + ) + + def generate(self, prompt: list) -> List[str]: + """Generate a reply for the given chat prompt. + + Parameters + ---------- + prompt : list of dict + Conversation history in OpenAI chat format. Each dict must contain + at least ``"role"`` (``"system"``, ``"user"``, or ``"assistant"``) + and ``"content"`` (the message text). + + Returns + ------- + list of str + A single-element list containing the model's reply text, extracted + from ``choices[0]["message"]["content"]``. + """ + output = self.model.create_chat_completion( + messages=prompt, + max_tokens=self.max_tokens, + temperature=self.temperature, + frequency_penalty=self.frequency_penalty, + ) + return [output["choices"][0]["message"]["content"]] diff --git a/DashAI/back/models/hugging_face/llama_model.py b/DashAI/back/models/hugging_face/llama_model.py index 18f36239b..170817f39 100644 --- a/DashAI/back/models/hugging_face/llama_model.py +++ b/DashAI/back/models/hugging_face/llama_model.py @@ -1,485 +1,186 @@ -from typing import List +"""Llama 3.x Instruct GGUF checkpoint subclasses for DashAI.""" -from DashAI.back.core.schema_fields import ( - BaseSchema, - enum_field, - float_field, - int_field, - schema_field, -) from DashAI.back.core.utils import MultilingualString -from DashAI.back.models.text_to_text_generation_model import ( - TextToTextGenerationTaskModel, -) -from DashAI.back.models.utils import ( - LLAMA_DEVICE_ENUM, - LLAMA_DEVICE_PLACEHOLDER, - LLAMA_DEVICE_TO_IDX, +from DashAI.back.models.hugging_face.gguf_text_generation_base import ( + GGUFTextGenerationModel, + GGUFTextGenerationSchema, ) -LLAMA_FILENAME_MAP = { - "bartowski/Meta-Llama-3.1-8B-Instruct-GGUF": "*Q4_K_M.gguf", - "bartowski/Llama-3.2-1B-Instruct-GGUF": "*Q4_K_M.gguf", - "bartowski/Llama-3.2-3B-Instruct-GGUF": "*Q4_K_M.gguf", -} +class Llama31_8BInstruct(GGUFTextGenerationModel): # noqa: N801 + """Meta Llama 3.1 8B Instruct GGUF checkpoint (Q4_K_M quantization). -class LlamaSchema(BaseSchema): - """Configuration schema for Meta Llama 3.x text generation. + An 8B-parameter instruction-tuned model from Meta with strong general + reasoning and multilingual ability. Weights are stored locally after a + one-time download from HuggingFace. - Configures the GGUF checkpoint variant (``model_name``), generation - behaviour (``max_tokens``, ``temperature``, ``frequency_penalty``), - context length (``context_window``), device target (``device``), and - system prompt (``system_prompt``) for ``LlamaModel``. + References + ---------- + - https://huggingface.co/bartowski/Meta-Llama-3.1-8B-Instruct-GGUF """ - model_name: schema_field( - enum_field( - enum=[ - "bartowski/Meta-Llama-3.1-8B-Instruct-GGUF", - "bartowski/Llama-3.2-1B-Instruct-GGUF", - "bartowski/Llama-3.2-3B-Instruct-GGUF", - ] + REPO_ID = "bartowski/Meta-Llama-3.1-8B-Instruct-GGUF" + GGUF_PATTERN = "*Q4_K_M.gguf" + DOWNLOAD_SIZE_BYTES = 4_900_000_000 + SCHEMA = GGUFTextGenerationSchema + COLOR: str = "#1a237e" + DISPLAY_NAME = MultilingualString( + en="Llama 3.1 8B Instruct", + es="Llama 3.1 8B Instruct", + pt="Llama 3.1 8B Instruct", + de="Llama 3.1 8B Instruct", + zh="Llama 3.1 8B Instruct", + ) + DESCRIPTION = MultilingualString( + en=( + "Meta Llama 3.1 8B Instruct is an 8B-parameter instruction-tuned language " + "model, loaded as a Q4_K_M GGUF for efficient CPU or GPU inference. It " + "offers strong reasoning, coding, and multilingual capabilities. This is " + "the largest text-generation model in DashAI and benefits from a GPU. " + "Model available at " + "https://huggingface.co/bartowski/Meta-Llama-3.1-8B-Instruct-GGUF." ), - placeholder="bartowski/Llama-3.2-3B-Instruct-GGUF", - description=MultilingualString( - en=( - "The Meta Llama 3.x Instruct checkpoint to load in GGUF format via " - "bartowski's community quantizations. 'Llama-3.2-1B' (~1B parameters) " - "is the smallest and fastest, ideal for CPU-only systems. " - "'Llama-3.2-3B' (~3B parameters) offers a good speed/quality " - "trade-off. " - "'Meta-Llama-3.1-8B' (~8B parameters) delivers the highest quality " - "at the cost of more RAM and slower inference." - ), - es=( - "El checkpoint Meta Llama 3.x Instruct a cargar en formato GGUF " - "mediante las cuantizaciones comunitarias de bartowski. " - "'Llama-3.2-1B' (~1B parámetros) es el más pequeño y rápido, " - "ideal para sistemas solo con CPU. " - "'Llama-3.2-3B' (~3B parámetros) ofrece un buen equilibrio entre " - "velocidad y calidad. 'Meta-Llama-3.1-8B' (~8B parámetros) entrega " - "la mayor calidad a costa de más RAM e inferencia más lenta." - ), - pt=( - "O checkpoint Meta Llama 3.x Instruct para carregar em formato GGUF " - "via quantizações comunitárias de bartowski. " - "'Llama-3.2-1B' (~1B parâmetros) é o menor e mais rápido, " - "ideal para sistemas apenas com CPU. " - "'Llama-3.2-3B' (~3B parâmetros) oferece um bom equilíbrio entre " - "velocidade e qualidade. 'Meta-Llama-3.1-8B' (~8B parâmetros) " - "entrega a maior qualidade ao custo de mais RAM e " - "inferência mais lenta." - ), - de=( - "Der im GGUF-Format zu ladende Meta Llama 3.x Instruct-Checkpoint " - "über bartowskis Community-Quantisierungen. " - "'Llama-3.2-1B' (~1B Parameter) ist der kleinste und schnellste, " - "ideal für reine CPU-Systeme. " - "'Llama-3.2-3B' (~3B Parameter) bietet ein gutes Geschwindigkeit-" - "Qualitäts-Verhältnis. 'Meta-Llama-3.1-8B' (~8B Parameter) liefert " - "die höchste Qualität auf Kosten von mehr RAM und langsamerer Inferenz." - ), - zh=( - "通过 bartowski 社区量化加载的 Meta Llama 3.x Instruct GGUF 检查点。" - "'Llama-3.2-1B'（约 1B 参数）是最小最快的版本，适合仅使用 CPU 的系统。" - "'Llama-3.2-3B'（约 3B 参数）在速度与质量之间取得良好平衡。" - "'Meta-Llama-3.1-8B'（约 8B 参数）质量最高，" - "但需要更多内存且推理速度较慢。" - ), + es=( + "Meta Llama 3.1 8B Instruct es un modelo de lenguaje de 8B parametros " + "ajustado para instrucciones, cargado como GGUF Q4_K_M para inferencia " + "eficiente en CPU o GPU. Ofrece solida capacidad de razonamiento, " + "programacion y multilingue. Es el modelo de generacion de texto mas " + "grande de DashAI y se beneficia de una GPU." ), - alias=MultilingualString( - en="Model name", - es="Nombre del modelo", - pt="Nome do modelo", - de="Modellname", - zh="模型名称", + pt=( + "Meta Llama 3.1 8B Instruct e um modelo de linguagem de 8B parametros " + "ajustado para instrucoes, carregado como GGUF Q4_K_M para inferencia " + "eficiente em CPU ou GPU. Oferece solida capacidade de raciocinio, " + "programacao e multilingue. E o maior modelo de geracao de texto do " + "DashAI e se beneficia de uma GPU." ), - ) # type: ignore - - max_tokens: schema_field( - int_field(ge=1), - placeholder=100, - description=MultilingualString( - en=( - "Maximum number of new tokens the model will generate per response. " - "Roughly 1 token ≈ 0.75 English words. Set to 100-200 for short " - "answers, 500-1000 for detailed explanations or code. Must not " - "exceed the context window minus the prompt length." - ), - es=( - "Número máximo de tokens nuevos que el modelo generará por respuesta. " - "Aproximadamente 1 token ≈ 0.75 palabras en español. Use 100-200 " - "para respuestas cortas, 500-1000 para explicaciones detalladas o " - "código. No debe superar la ventana de contexto menos la longitud " - "del prompt." - ), - pt=( - "Número máximo de tokens novos que o modelo gerará por resposta. " - "Aproximadamente 1 token ≈ 0.75 palavras em português. Use 100-200 " - "para respostas curtas, 500-1000 para explicações detalhadas ou " - "código. Não deve exceder a janela de contexto menos o comprimento " - "do prompt." - ), - de=( - "Maximale Anzahl neuer Token, die das Modell pro Antwort erzeugt. " - "Ungefähr 1 Token ≈ 0,75 englische Wörter. 100-200 für kurze " - "Antworten, 500-1000 für ausführliche Erklärungen oder Code. " - "Darf die Kontextfenstergröße abzüglich der Prompt-Länge nicht " - "überschreiten." - ), - zh=( - "模型每次响应生成的最大新 token 数。" - "约 1 token ≈ 0.75 个英文单词。简短回答设置 100-200，" - "详细说明或代码设置 500-1000。不得超过上下文窗口减去提示词长度的值。" - ), + de=( + "Meta Llama 3.1 8B Instruct ist ein 8B-Parameter-Instruktionsmodell, " + "als Q4_K_M-GGUF fuer effiziente CPU- oder GPU-Inferenz geladen. Es " + "bietet starke Faehigkeiten in Schlussfolgern, Programmierung und " + "Mehrsprachigkeit. Es ist das groesste Textgenerierungsmodell in DashAI " + "und profitiert von einer GPU." ), - alias=MultilingualString( - en="Max tokens", - es="Tokens máximos", - pt="Tokens máximos", - de="Maximale neue Token", - zh="最大 token 数", + zh=( + "Meta Llama 3.1 8B Instruct 是 80 亿参数的指令微调语言模型，" + "以 Q4_K_M GGUF 格式加载，支持高效的 CPU 或 GPU 推理。" + "它具备强大的推理、编程和多语言能力。" + "这是 DashAI 中最大的文本生成模型，使用 GPU 效果更佳。" ), - ) # type: ignore + ) - temperature: schema_field( - float_field(ge=0.0, le=1.0), - placeholder=0.7, - description=MultilingualString( - en=( - "Sampling temperature controlling output randomness (range 0.0-1.0). " - "At 0.0 the model always picks the most likely token (greedy, fully " - "deterministic). Around 0.7 is a good balance for conversational " - "tasks. At 1.0 outputs are maximally varied and unpredictable." - ), - es=( - "Temperatura de muestreo que controla la aleatoriedad de la salida " - "(rango 0.0-1.0). En 0.0 el modelo siempre elige el token más " - "probable (greedy, totalmente determinista). Alrededor de 0.7 es " - "un buen equilibrio para tareas conversacionales. En 1.0 las " - "salidas son máximamente variadas e impredecibles." - ), - pt=( - "Temperatura de amostragem que controla a aleatoriedade da saída " - "(intervalo 0.0-1.0). Em 0.0 o modelo sempre escolhe o token mais " - "provável (greedy, totalmente determinístico). Em torno de 0.7 é " - "um bom equilíbrio para tarefas conversacionais. Em 1.0 as " - "saídas são maximamente variadas e imprevisíveis." - ), - de=( - "Stichprobentemperatur zur Steuerung der Ausgabezufälligkeit (0.0-1.0)." - "Bei 0.0 wählt das Modell stets den wahrscheinlichsten Token (greedy, " - "vollständig deterministisch). Um 0.7 ist ein gutes Gleichgewicht für " - "Konversationsaufgaben. Bei 1.0 sind Ausgaben maximal variiert und " - "unvorhersehbar." - ), - zh=( - "控制输出随机性的采样温度（范围 0.0-1.0）。" - "0.0 时模型始终选择最可能的 token（贪心，完全确定性）。" - "0.7 左右是对话任务的良好平衡点。1.0 时输出变化最大，不可预测。" - ), - ), - alias=MultilingualString( - en="Temperature", - es="Temperatura", - pt="Temperatura", - de="Temperatur", - zh="温度", - ), - ) # type: ignore - frequency_penalty: schema_field( - float_field(ge=0.0, le=2.0), - placeholder=0.1, - description=MultilingualString( - en=( - "Penalizes tokens that have already appeared in the output based on " - "how often they occur (range 0.0-2.0). At 0.0 there is no penalty " - "and the model may repeat itself. Values around 0.1-0.3 gently " - "discourage repetition. High values (1.5+) strongly prevent reuse " - "of any word, which may produce less coherent text." - ), - es=( - "Penaliza los tokens que ya aparecieron en la salida según su " - "frecuencia (rango 0.0-2.0). En 0.0 no hay penalización y el modelo " - "puede repetirse. Valores en torno a 0.1-0.3 desincentivan " - "suavemente la repetición. Valores altos (1.5+) previenen " - "fuertemente la reutilización de palabras, lo que puede producir " - "texto menos coherente." - ), - pt=( - "Penaliza tokens que já apareceram na saída com base em sua " - "frequência (intervalo 0.0-2.0). Em 0.0 não há penalização e o modelo " - "pode se repetir. Valores em torno de 0.1-0.3 desencorajam " - "suavemente a repetição. Valores altos (1.5+) impedem fortemente " - "o reuso de palavras, o que pode produzir texto menos coerente." - ), - de=( - "Bestraft Token, die bereits in der Ausgabe erschienen sind, " - "basierend auf ihrer Häufigkeit (0.0-2.0). Bei 0.0 gibt es keine " - "Strafe und das Modell kann sich wiederholen. Werte um 0.1-0.3 " - "hemmen Wiederholungen sanft. Hohe Werte (1.5+) verhindern die " - "Wiederverwendung von Wörtern stark, was zu weniger kohärentem Text " - "führen kann." - ), - zh=( - "根据 token 在输出中出现的频率对其进行惩罚（范围 0.0-2.0）。" - "0.0 时无惩罚，模型可能重复自身。0.1-0.3 左右的值可温和抑制重复。" - "高值（1.5+）会强烈阻止任何词的重复使用，可能导致文本连贯性下降。" - ), - ), - alias=MultilingualString( - en="Frequency penalty", - es="Penalización de frecuencia", - pt="Penalidade de frequência", - de="Häufigkeitsstrafe", - zh="频率惩罚", - ), - ) # type: ignore +class Llama32_1BInstruct(GGUFTextGenerationModel): # noqa: N801 + """Meta Llama 3.2 1B Instruct GGUF checkpoint (Q4_K_M quantization). + + A lightweight 1B-parameter instruction-tuned model from Meta suitable for + CPU inference. Weights are stored locally after a one-time download from + HuggingFace. + + References + ---------- + - https://huggingface.co/bartowski/Llama-3.2-1B-Instruct-GGUF + """ - context_window: schema_field( - int_field(ge=1, le=131072), - placeholder=512, - description=MultilingualString( - en=( - "Total token budget for a single forward pass, including both the " - "input prompt and the generated response. Larger values allow longer " - "conversations but consume more RAM/VRAM. Llama 3.1 supports up to " - "128K tokens natively; Llama 3.2 models support up to 128K tokens." - ), - es=( - "Presupuesto total de tokens para una sola pasada, incluyendo tanto " - "el prompt de entrada como la respuesta generada. Valores más altos " - "permiten conversaciones más largas pero consumen más RAM/VRAM. " - "Llama 3.1 soporta hasta 128K tokens de forma nativa; los modelos " - "Llama 3.2 soportan hasta 128K tokens." - ), - pt=( - "Orçamento total de tokens para uma única passagem, incluindo tanto " - "o prompt de entrada quanto a resposta gerada. Valores maiores " - "permitem conversas mais longas mas consomem mais RAM/VRAM. " - "Llama 3.1 suporta até 128K tokens nativamente; os modelos " - "Llama 3.2 suportam até 128K tokens." - ), - de=( - "Gesamtes Token-Budget für einen einzelnen Vorwärtsdurchlauf, " - "einschließlich Eingabe-Prompt und generierter Antwort. Größere Werte " - "ermöglichen längere Gespräche, verbrauchen jedoch mehr RAM/VRAM. " - "Llama 3.1 unterstützt nativ bis zu 128K Token; Llama-3.2-Modelle " - "unterstützen ebenfalls bis zu 128K Token." - ), - zh=( - "单次前向传播的总 token 预算，包含输入提示和生成响应。" - "较大的值允许更长的对话，但会消耗更多 RAM/VRAM。" - "Llama 3.1 原生支持最多 128K token；" - "Llama 3.2 模型同样支持最多 128K token。" - ), + REPO_ID = "bartowski/Llama-3.2-1B-Instruct-GGUF" + GGUF_PATTERN = "*Q4_K_M.gguf" + DOWNLOAD_SIZE_BYTES = 800_000_000 + SCHEMA = GGUFTextGenerationSchema + COLOR: str = "#1a237e" + DISPLAY_NAME = MultilingualString( + en="Llama 3.2 1B Instruct", + es="Llama 3.2 1B Instruct", + pt="Llama 3.2 1B Instruct", + de="Llama 3.2 1B Instruct", + zh="Llama 3.2 1B Instruct", + ) + DESCRIPTION = MultilingualString( + en=( + "Meta Llama 3.2 1B Instruct is a lightweight 1B-parameter " + "instruction-tuned language model, loaded as a Q4_K_M GGUF for fast CPU " + "inference. It is a good balance of speed and quality for everyday tasks. " + "Model available at " + "https://huggingface.co/bartowski/Llama-3.2-1B-Instruct-GGUF." ), - alias=MultilingualString( - en="Context window", - es="Ventana de contexto", - pt="Janela de contexto", - de="Kontextfenster", - zh="上下文窗口", + es=( + "Meta Llama 3.2 1B Instruct es un modelo de lenguaje ligero de 1B " + "parametros ajustado para instrucciones, cargado como GGUF Q4_K_M para " + "inferencia rapida en CPU. Ofrece un buen equilibrio entre velocidad y " + "calidad para tareas cotidianas." ), - ) # type: ignore - - device: schema_field( - enum_field(enum=LLAMA_DEVICE_ENUM), - placeholder=LLAMA_DEVICE_PLACEHOLDER, - description=MultilingualString( - en=( - "Hardware device for llama.cpp inference. 'CPU' runs the model " - "fully in RAM with no GPU requirement. Selecting a GPU option " - "offloads all layers for faster inference, setting n_gpu_layers=-1 " - "so every transformer layer is GPU-accelerated." - ), - es=( - "Dispositivo de hardware para la inferencia con llama.cpp. 'CPU' " - "ejecuta el modelo completamente en RAM sin requisito de GPU. " - "Seleccionar una opción de GPU descarga todas las capas para " - "inferencia más rápida, estableciendo n_gpu_layers=-1 para que " - "cada capa del transformer sea acelerada por GPU." - ), - pt=( - "Dispositivo de hardware para inferência com llama.cpp. 'CPU' " - "executa o modelo completamente em RAM sem requisito de GPU. " - "Selecionar uma opção de GPU descarrega todas as camadas para " - "inferência mais rápida, definindo n_gpu_layers=-1 para que " - "cada camada do transformer seja acelerada por GPU." - ), - de=( - "Hardware-Gerät für die llama.cpp-Inferenz. 'CPU' führt das Modell " - "vollständig im RAM ohne GPU-Anforderung aus. Eine GPU-Option " - "lagert alle Schichten für schnellere Inferenz aus und setzt " - "n_gpu_layers=-1, damit jede Transformer-Schicht GPU-beschleunigt wird." - ), - zh=( - "llama.cpp 推理的硬件设备。'CPU' 完全在内存中运行模型，无需 GPU。" - "选择 GPU 选项可卸载所有层以加快推理速度，" - "设置 n_gpu_layers=-1 使每个 Transformer 层均由 GPU 加速。" - ), + pt=( + "Meta Llama 3.2 1B Instruct e um modelo de linguagem leve de 1B " + "parametros ajustado para instrucoes, carregado como GGUF Q4_K_M para " + "inferencia rapida em CPU. Oferece um bom equilibrio entre velocidade e " + "qualidade para tarefas cotidianas." ), - alias=MultilingualString( - en="Device", - es="Dispositivo", - pt="Dispositivo", - de="Gerät", - zh="设备", + de=( + "Meta Llama 3.2 1B Instruct ist ein leichtes 1B-Parameter-" + "Instruktionsmodell, als Q4_K_M-GGUF fuer schnelle CPU-Inferenz geladen. " + "Es bietet ein gutes Gleichgewicht aus Geschwindigkeit und Qualitaet fuer " + "alltaegliche Aufgaben." ), - ) # type: ignore - + zh=( + "Meta Llama 3.2 1B Instruct 是轻量级的 10 亿参数指令微调语言模型，" + "以 Q4_K_M GGUF 格式加载，支持快速 CPU 推理。" + "在日常任务中兼顾速度与质量。" + ), + ) -class LlamaModel(TextToTextGenerationTaskModel): - """Meta Llama 3.x instruction-tuned model for text generation via llama.cpp. - Wraps the Meta Llama 3.x family of open-weight instruction-tuned LLMs - loaded in Q4_K_M GGUF format using the ``llama-cpp-python`` library. - GGUF quantization enables efficient CPU and GPU inference without requiring - full-precision weights, making the models practical on consumer hardware. +class Llama32_3BInstruct(GGUFTextGenerationModel): # noqa: N801 + """Meta Llama 3.2 3B Instruct GGUF checkpoint (Q4_K_M quantization). - Three sizes are available via bartowski's community quantizations: - 1B (fastest, CPU-friendly), 3B (balanced), and 8B (highest quality). + A 3B-parameter instruction-tuned model from Meta offering higher quality + than the 1B variant while remaining CPU-friendly. Weights are stored locally + after a one-time download from HuggingFace. References ---------- - - [1] Meta AI, "Llama 3", 2024. https://ai.meta.com/blog/meta-llama-3/ - - [2] https://huggingface.co/bartowski + - https://huggingface.co/bartowski/Llama-3.2-3B-Instruct-GGUF """ - SCHEMA = LlamaSchema + REPO_ID = "bartowski/Llama-3.2-3B-Instruct-GGUF" + GGUF_PATTERN = "*Q4_K_M.gguf" + DOWNLOAD_SIZE_BYTES = 2_000_000_000 + SCHEMA = GGUFTextGenerationSchema COLOR: str = "#1a237e" - DISPLAY_NAME: str = MultilingualString( - en="Llama Model", - es="Modelo Llama", - pt="Modelo Llama", - de="Llama-Modell", - zh="Llama 模型", + DISPLAY_NAME = MultilingualString( + en="Llama 3.2 3B Instruct", + es="Llama 3.2 3B Instruct", + pt="Llama 3.2 3B Instruct", + de="Llama 3.2 3B Instruct", + zh="Llama 3.2 3B Instruct", ) - DESCRIPTION: str = MultilingualString( + DESCRIPTION = MultilingualString( en=( - "Meta Llama 3.x is a family of open instruction-tuned large language " - "models developed by Meta AI, loaded in GGUF format for efficient CPU " - "and GPU inference via the llama.cpp library. It supports multi-turn " - "conversation, reasoning, coding, and general text generation. Available " - "in 1B, 3B, and 8B parameter sizes. Models are hosted at " - "https://huggingface.co/bartowski." + "Meta Llama 3.2 3B Instruct is a 3B-parameter instruction-tuned language " + "model, loaded as a Q4_K_M GGUF for efficient CPU or GPU inference. It " + "offers stronger reasoning and generation quality than the 1B variant. " + "Model available at " + "https://huggingface.co/bartowski/Llama-3.2-3B-Instruct-GGUF." ), es=( - "Meta Llama 3.x es una familia de modelos de lenguaje grande de código " - "abierto ajustados para instrucciones, desarrollados por Meta AI, cargados " - "en formato GGUF para inferencia eficiente en CPU y GPU mediante la " - "librería llama.cpp. Soporta conversación multi-turno, razonamiento, " - "programación y generación de texto en general. Disponible en tamaños de " - "1B, 3B y 8B parámetros. Los modelos están en " - "https://huggingface.co/bartowski." + "Meta Llama 3.2 3B Instruct es un modelo de lenguaje de 3B parametros " + "ajustado para instrucciones, cargado como GGUF Q4_K_M para inferencia " + "eficiente en CPU o GPU. Ofrece mejor razonamiento y calidad de " + "generacion que la variante de 1B." ), pt=( - "Meta Llama 3.x é uma família de modelos de linguagem grande de código " - "aberto ajustados para instruções, desenvolvidos pela Meta AI, carregados " - "em formato GGUF para inferência eficiente em CPU e GPU via a biblioteca " - "llama.cpp. Suporta conversação multi-turno, raciocínio, programação e " - "geração de texto em geral. Disponível nos tamanhos de parâmetros 1B, 3B " - "e 8B. Os modelos estão em https://huggingface.co/bartowski." + "Meta Llama 3.2 3B Instruct e um modelo de linguagem de 3B parametros " + "ajustado para instrucoes, carregado como GGUF Q4_K_M para inferencia " + "eficiente em CPU ou GPU. Oferece melhor raciocinio e qualidade de " + "geracao do que a variante de 1B." ), de=( - "Meta Llama 3.x ist eine Familie offener instruktionsoptimierter großer " - "Sprachmodelle von Meta AI, im GGUF-Format für effiziente CPU- und " - "GPU-Inferenz über die llama.cpp-Bibliothek geladen. Unterstützt " - "Mehrfachdialog, Schlussfolgerung, Programmierung und allgemeine " - "Textgenerierung. Verfügbar in den Parametergrößen 1B, 3B und 8B. " - "Modelle unter https://huggingface.co/bartowski." + "Meta Llama 3.2 3B Instruct ist ein 3B-Parameter-Instruktionsmodell, " + "als Q4_K_M-GGUF fuer effiziente CPU- oder GPU-Inferenz geladen. Es " + "bietet besseres Schlussfolgern und Generierungsqualitaet als die " + "1B-Variante." ), zh=( - "Meta Llama 3.x 是 Meta AI 开发的开放指令微调大语言模型系列，" - "以 GGUF 格式加载，通过 llama.cpp 库实现高效的 CPU 和 GPU 推理。" - "支持多轮对话、推理、编程和通用文本生成。提供 1B、3B 和 8B 参数规格。" - "模型托管于 https://huggingface.co/bartowski。" + "Meta Llama 3.2 3B Instruct 是 30 亿参数的指令微调语言模型，" + "以 Q4_K_M GGUF 格式加载，支持高效的 CPU 或 GPU 推理。" + "与 1B 变体相比，它具有更强的推理能力和生成质量。" ), ) - - def __init__(self, **kwargs): - """Download and initialise a Llama 3.x GGUF model via llama.cpp. - - The model weights are fetched from HuggingFace Hub using - ``Llama.from_pretrained`` and kept in memory for repeated calls to - ``generate``. - - Parameters - ---------- - **kwargs : dict - model_name : str, optional - HuggingFace repo ID for the GGUF checkpoint. - Defaults to ``"bartowski/Llama-3.2-3B-Instruct-GGUF"``. - max_tokens : int, optional - Maximum number of new tokens to generate per call. Default 100. - temperature : float, optional - Sampling temperature in [0.0, 1.0]. Default 0.7. - frequency_penalty : float, optional - Token-frequency penalty in [0.0, 2.0]. Default 0.1. - context_window : int, optional - Total token budget (prompt + response) for a single forward - pass. Default 512. - device : str, optional - Target device from ``LLAMA_DEVICE_ENUM``. Any value whose - index is >= 0 enables full GPU offload (``n_gpu_layers=-1``); - ``"CPU"`` runs fully in RAM. - - Raises - ------ - RuntimeError - If ``llama-cpp-python`` is not installed. - """ - try: - from llama_cpp import Llama - except ImportError as e: - raise RuntimeError( - "llama-cpp-python is not installed. " - "Please install it to use this model." - ) from e - - kwargs = self.validate_and_transform(kwargs) - self.model_name = kwargs.get( - "model_name", "bartowski/Llama-3.2-3B-Instruct-GGUF" - ) - self.max_tokens = kwargs.pop("max_tokens", 100) - self.temperature = kwargs.pop("temperature", 0.7) - self.frequency_penalty = kwargs.pop("frequency_penalty", 0.1) - self.n_ctx = kwargs.pop("context_window", 512) - - self.filename = LLAMA_FILENAME_MAP.get(self.model_name, "*Q4_K_M.gguf") - use_gpu = LLAMA_DEVICE_TO_IDX.get(kwargs.get("device")) >= 0 - - self.model = Llama.from_pretrained( - repo_id=self.model_name, - filename=self.filename, - verbose=True, - n_ctx=self.n_ctx, - n_gpu_layers=-1 if use_gpu else 0, - main_gpu=(LLAMA_DEVICE_TO_IDX.get(kwargs.get("device")) if use_gpu else 0), - ) - - def generate(self, prompt: list[dict[str, str]]) -> List[str]: - """Generate a reply for the given chat prompt. - - Parameters - ---------- - prompt : list of dict - Conversation history in OpenAI chat format. Each dict must contain - at least ``"role"`` (``"system"``, ``"user"``, or ``"assistant"``) - and ``"content"`` (the message text). - - Returns - ------- - list of str - A single-element list containing the model's reply text, extracted - from ``choices[0]["message"]["content"]``. - """ - output = self.model.create_chat_completion( - messages=prompt, - max_tokens=self.max_tokens, - temperature=self.temperature, - frequency_penalty=self.frequency_penalty, - ) - return [output["choices"][0]["message"]["content"]] diff --git a/DashAI/back/models/hugging_face/mistral_model.py b/DashAI/back/models/hugging_face/mistral_model.py index ab0f450bf..5b07d1996 100644 --- a/DashAI/back/models/hugging_face/mistral_model.py +++ b/DashAI/back/models/hugging_face/mistral_model.py @@ -1,441 +1,124 @@ -from typing import List +"""Mistral Instruct GGUF checkpoint subclasses for DashAI.""" -from DashAI.back.core.schema_fields import ( - BaseSchema, - enum_field, - float_field, - int_field, - schema_field, -) from DashAI.back.core.utils import MultilingualString -from DashAI.back.models.text_to_text_generation_model import ( - TextToTextGenerationTaskModel, -) -from DashAI.back.models.utils import ( - LLAMA_DEVICE_ENUM, - LLAMA_DEVICE_PLACEHOLDER, - LLAMA_DEVICE_TO_IDX, +from DashAI.back.models.hugging_face.gguf_text_generation_base import ( + GGUFTextGenerationModel, + GGUFTextGenerationSchema, ) -class MistralSchema(BaseSchema): - """Schema for MistralModel hyperparameters. - - Configures the checkpoint variant, generation length, sampling temperature, - frequency penalty, context window, and target device for Mistral Instruct - models loaded via ``llama-cpp-python`` in GGUF format. - """ - - model_name: schema_field( - enum_field( - enum=[ - "bartowski/Mistral-7B-Instruct-v0.3-GGUF", - "bartowski/Mistral-Nemo-Instruct-2407-GGUF", - ] - ), - placeholder="bartowski/Mistral-7B-Instruct-v0.3-GGUF", - description=MultilingualString( - en=( - "The Mistral Instruct checkpoint to load in GGUF format. " - "'Mistral-7B-Instruct-v0.3' is a 7B-parameter instruction model " - "that delivers strong performance for its size. " - "'Mistral-Nemo-Instruct-2407' is a 12B-parameter model jointly " - "developed with NVIDIA, featuring a 128K context window and " - "improved multilingual capabilities." - ), - es=( - "El checkpoint Mistral Instruct a cargar en formato GGUF. " - "'Mistral-7B-Instruct-v0.3' es un modelo de instrucción de 7B " - "parámetros con fuerte rendimiento para su tamaño. " - "'Mistral-Nemo-Instruct-2407' es un modelo de 12B parámetros " - "desarrollado conjuntamente con NVIDIA, con una ventana de contexto " - "de 128K y mejores capacidades multilingües." - ), - pt=( - "O checkpoint Mistral Instruct para carregar em formato GGUF. " - "'Mistral-7B-Instruct-v0.3' é um modelo de instrução de 7B " - "parâmetros com forte desempenho para seu tamanho. " - "'Mistral-Nemo-Instruct-2407' é um modelo de 12B parâmetros " - "desenvolvido conjuntamente com a NVIDIA, com uma janela de contexto " - "de 128K e melhores capacidades multilíngues." - ), - de=( - "Der im GGUF-Format zu ladende Mistral Instruct-Checkpoint. " - "'Mistral-7B-Instruct-v0.3' ist ein 7B-Parameter-Instruktionsmodell " - "mit starker Leistung für seine Größe. " - "'Mistral-Nemo-Instruct-2407' ist ein 12B-Parameter-Modell, gemeinsam " - "mit NVIDIA entwickelt, mit einem 128K-Kontextfenster und verbesserten " - "mehrsprachigen Fähigkeiten." - ), - zh=( - "要加载的 Mistral Instruct 检查点（GGUF 格式）。" - "'Mistral-7B-Instruct-v0.3' 是 7B 参数指令模型，性能出色。" - "'Mistral-Nemo-Instruct-2407' 是与 NVIDIA 联合开发的 12B 参数模型，" - "支持 128K 上下文窗口，多语言能力更强。" - ), - ), - alias=MultilingualString( - en="Model name", - es="Nombre del modelo", - pt="Nome do modelo", - de="Modellname", - zh="模型名称", - ), - ) # type: ignore +class Mistral7BInstructV03(GGUFTextGenerationModel): + """Mistral 7B Instruct v0.3 GGUF checkpoint (Q4_K_M quantization). - max_tokens: schema_field( - int_field(ge=1), - placeholder=100, - description=MultilingualString( - en=( - "Maximum number of new tokens the model will generate per response. " - "Roughly 1 token ≈ 0.75 English words. Set to 100-200 for short " - "answers, 500-1000 for detailed explanations or code." - ), - es=( - "Número máximo de tokens nuevos que el modelo generará por respuesta. " - "Aproximadamente 1 token ≈ 0.75 palabras en español. Use 100-200 " - "para respuestas cortas, 500-1000 para explicaciones detalladas " - "o código." - ), - pt=( - "Número máximo de tokens novos que o modelo gerará por resposta. " - "Aproximadamente 1 token ≈ 0.75 palavras em português. Use 100-200 " - "para respostas curtas, 500-1000 para explicações detalhadas " - "ou código." - ), - de=( - "Maximale Anzahl neuer Token, die das Modell pro Antwort erzeugt. " - "Ungefähr 1 Token ≈ 0,75 englische Wörter. 100-200 für kurze " - "Antworten, 500-1000 für ausführliche Erklärungen oder Code." - ), - zh=( - "模型每次响应生成的最大新 token 数。" - "大约 1 token 约等于 0.75 个英文单词。" - "短回答设为 100-200，详细解释或代码设为 500-1000。" - ), - ), - alias=MultilingualString( - en="Max tokens", - es="Tokens máximos", - pt="Tokens máximos", - de="Maximale neue Token", - zh="最大 token 数", - ), - ) # type: ignore + A 7B-parameter instruction-tuned model from Mistral AI with strong general + performance. Weights are stored locally after a one-time download from + HuggingFace. - temperature: schema_field( - float_field(ge=0.0, le=1.0), - placeholder=0.7, - description=MultilingualString( - en=( - "Sampling temperature controlling output randomness (range 0.0-1.0). " - "At 0.0 the model picks the most likely token (deterministic). " - "Around 0.7 balances quality and creativity. At 1.0 outputs are " - "maximally varied." - ), - es=( - "Temperatura de muestreo que controla la aleatoriedad (rango 0.0-1.0). " - "En 0.0 el modelo elige el token más probable (determinista). " - "Alrededor de 0.7 equilibra calidad y creatividad. En 1.0 las salidas " - "son máximamente variadas." - ), - pt=( - "Temperatura de amostragem que controla a aleatoriedade " - "(intervalo 0.0-1.0). " - "Em 0.0 o modelo escolhe o token mais provável (determinístico). " - "Em torno de 0.7 equilibra qualidade e criatividade. Em 1.0 as saídas " - "são maximamente variadas." - ), - de=( - "Stichprobentemperatur zur Steuerung der Ausgabezufälligkeit (0.0-1.0)." - "Bei 0.0 wählt das Modell den wahrscheinlichsten Token " - "(deterministisch). " - "Ca. 0.7 balanciert Qualität und Kreativität. Bei 1.0 sind Ausgaben " - "maximal variiert." - ), - zh=( - "控制输出随机性的采样温度（范围 0.0-1.0）。" - "0.0 时模型选择最可能的 token（确定性）。" - "0.7 左右在质量与创造性之间取得平衡。1.0 时输出变化最大。" - ), - ), - alias=MultilingualString( - en="Temperature", - es="Temperatura", - pt="Temperatura", - de="Temperatur", - zh="温度", - ), - ) # type: ignore + References + ---------- + - https://huggingface.co/bartowski/Mistral-7B-Instruct-v0.3-GGUF + """ - frequency_penalty: schema_field( - float_field(ge=0.0, le=2.0), - placeholder=0.1, - description=MultilingualString( - en=( - "Penalizes tokens that have already appeared in the output based on " - "frequency (range 0.0-2.0). Higher values discourage repetition." - ), - es=( - "Penaliza los tokens que ya aparecieron en la salida según su " - "frecuencia (rango 0.0-2.0). Valores más altos desincentivan " - "la repetición." - ), - pt=( - "Penaliza tokens que já apareceram na saída com base na " - "frequência (intervalo 0.0-2.0). Valores mais altos desencorajam " - "a repetição." - ), - de=( - "Bestraft Token, die bereits in der Ausgabe erschienen sind, " - "basierend auf ihrer Häufigkeit (0.0-2.0). Höhere Werte reduzieren " - "Wiederholungen." - ), - zh=( - "根据频率对已出现在输出中的 token 施加惩罚（范围 0.0-2.0）。" - "较高值可抑制重复。" - ), - ), - alias=MultilingualString( - en="Frequency penalty", - es="Penalización de frecuencia", - pt="Penalidade de frequência", - de="Häufigkeitsstrafe", - zh="频率惩罚", + REPO_ID = "bartowski/Mistral-7B-Instruct-v0.3-GGUF" + GGUF_PATTERN = "*Q4_K_M.gguf" + DOWNLOAD_SIZE_BYTES = 4_400_000_000 + SCHEMA = GGUFTextGenerationSchema + COLOR: str = "#ff6f00" + DISPLAY_NAME = MultilingualString( + en="Mistral 7B Instruct v0.3", + es="Mistral 7B Instruct v0.3", + pt="Mistral 7B Instruct v0.3", + de="Mistral 7B Instruct v0.3", + zh="Mistral 7B Instruct v0.3", + ) + DESCRIPTION = MultilingualString( + en=( + "Mistral 7B Instruct v0.3 is a 7B-parameter instruction-tuned language " + "model from Mistral AI, loaded as a Q4_K_M GGUF for efficient CPU or GPU " + "inference. It offers strong general reasoning and generation quality. " + "Model available at " + "https://huggingface.co/bartowski/Mistral-7B-Instruct-v0.3-GGUF." ), - ) # type: ignore - - context_window: schema_field( - int_field(ge=1, le=131072), - placeholder=512, - description=MultilingualString( - en=( - "Total token budget for a single forward pass, including prompt and " - "response. Mistral-7B supports up to 32K tokens; Mistral-Nemo " - "supports up to 128K tokens." - ), - es=( - "Presupuesto total de tokens por pasada, incluyendo prompt y " - "respuesta. Mistral-7B soporta hasta 32K tokens; Mistral-Nemo " - "soporta hasta 128K tokens." - ), - pt=( - "Orçamento total de tokens por passagem, incluindo prompt e " - "resposta. Mistral-7B suporta até 32K tokens; Mistral-Nemo " - "suporta até 128K tokens." - ), - de=( - "Gesamtes Token-Budget für einen einzelnen Vorwärtsdurchlauf, " - "einschließlich Eingabeaufforderung und Antwort. Mistral-7B unterstützt" - "bis zu 32K Token; Mistral-Nemo bis zu 128K Token." - ), - zh=( - "单次前向传播的总 token 预算，包含提示词和回复。" - "Mistral-7B 支持最多 32K token；Mistral-Nemo 支持最多 128K token。" - ), + es=( + "Mistral 7B Instruct v0.3 es un modelo de lenguaje de 7B parametros " + "ajustado para instrucciones por Mistral AI, cargado como GGUF Q4_K_M " + "para inferencia eficiente en CPU o GPU. Ofrece solida capacidad de " + "razonamiento y calidad de generacion." ), - alias=MultilingualString( - en="Context window", - es="Ventana de contexto", - pt="Janela de contexto", - de="Kontextfenster", - zh="上下文窗口", + pt=( + "Mistral 7B Instruct v0.3 e um modelo de linguagem de 7B parametros " + "ajustado para instrucoes pela Mistral AI, carregado como GGUF Q4_K_M " + "para inferencia eficiente em CPU ou GPU. Oferece solida capacidade de " + "raciocinio e qualidade de geracao." ), - ) # type: ignore - - device: schema_field( - enum_field(enum=LLAMA_DEVICE_ENUM), - placeholder=LLAMA_DEVICE_PLACEHOLDER, - description=MultilingualString( - en=( - "Hardware device for llama.cpp inference. 'CPU' runs the model " - "fully in RAM. A GPU option offloads all layers for faster inference." - ), - es=( - "Dispositivo de hardware para inferencia con llama.cpp. 'CPU' ejecuta " - "el modelo en RAM. Una opción de GPU descarga todas las capas para " - "inferencia más rápida." - ), - pt=( - "Dispositivo de hardware para inferência com llama.cpp. 'CPU' executa " - "o modelo em RAM. Uma opção de GPU descarrega todas as camadas para " - "inferência mais rápida." - ), - de=( - "Hardware-Gerät für die llama.cpp-Inferenz. 'CPU' führt das Modell " - "vollständig im RAM aus. Eine GPU-Option lagert alle Schichten für " - "schnellere Inferenz aus." - ), - zh=( - "llama.cpp 推理所用的硬件设备。'CPU' 完全在内存中运行模型。" - "选择 GPU 选项可卸载所有层以加快推理速度。" - ), + de=( + "Mistral 7B Instruct v0.3 ist ein 7B-Parameter-Instruktionsmodell von " + "Mistral AI, als Q4_K_M-GGUF fuer effiziente CPU- oder GPU-Inferenz " + "geladen. Es bietet starke allgemeine Schlussfolgerungs- und " + "Generierungsqualitaet." ), - alias=MultilingualString( - en="Device", - es="Dispositivo", - pt="Dispositivo", - de="Gerät", - zh="设备", + zh=( + "Mistral 7B Instruct v0.3 是 Mistral AI 推出的 70 亿参数指令微调语言模型，" + "以 Q4_K_M GGUF 格式加载，支持高效的 CPU 或 GPU 推理。" + "它具备强大的通用推理和生成质量。" ), - ) # type: ignore - + ) -class MistralModel(TextToTextGenerationTaskModel): - """Mistral Instruct model for open-ended text generation via llama.cpp. - Mistral is a 7B-parameter transformer language model developed by Mistral AI, - designed to deliver high performance with efficient inference. It uses grouped- - query attention (GQA) for faster decoding and sliding-window attention (SWA) to - handle long contexts efficiently. The 12B Mistral-Nemo variant, developed jointly - with NVIDIA, extends the context window to 128 K tokens and improves multilingual - capability. +class MistralNemoInstruct2407(GGUFTextGenerationModel): + """Mistral Nemo Instruct 2407 GGUF checkpoint (Q4_K_M quantization). - Models are loaded as GGUF quantized checkpoints via ``llama-cpp-python``, - allowing CPU and GPU inference without requiring a full PyTorch stack. + A 12B-parameter instruction-tuned model from Mistral AI and NVIDIA with a + large context window. This is a heavy model that benefits from a GPU. + Weights are stored locally after a one-time download from HuggingFace. References ---------- - - [1] Jiang et al. (2023) "Mistral 7B" https://arxiv.org/abs/2310.06825 - - [2] https://huggingface.co/mistralai + - https://huggingface.co/bartowski/Mistral-Nemo-Instruct-2407-GGUF """ - SCHEMA = MistralSchema + REPO_ID = "bartowski/Mistral-Nemo-Instruct-2407-GGUF" + GGUF_PATTERN = "*Q4_K_M.gguf" + DOWNLOAD_SIZE_BYTES = 7_100_000_000 + SCHEMA = GGUFTextGenerationSchema COLOR: str = "#ff6f00" - DISPLAY_NAME: str = MultilingualString( - en="Mistral Model", - es="Modelo Mistral", - pt="Modelo Mistral", - de="Mistral-Modell", - zh="Mistral 模型", + DISPLAY_NAME = MultilingualString( + en="Mistral Nemo Instruct 2407", + es="Mistral Nemo Instruct 2407", + pt="Mistral Nemo Instruct 2407", + de="Mistral Nemo Instruct 2407", + zh="Mistral Nemo Instruct 2407", ) - DESCRIPTION: str = MultilingualString( + DESCRIPTION = MultilingualString( en=( - "Mistral instruction-tuned models by Mistral AI, loaded in GGUF format " - "for efficient CPU and GPU inference via the llama.cpp library. Mistral " - "models are known for strong performance relative to their parameter count " - "and efficient inference. Supports multi-turn conversation, reasoning, " - "and general text generation. Available in 7B (Mistral-7B-v0.3) and 12B " - "(Mistral-Nemo-2407) variants. Models hosted at " - "https://huggingface.co/bartowski." + "Mistral Nemo Instruct 2407 is a 12B-parameter instruction-tuned language " + "model built by Mistral AI and NVIDIA, loaded as a Q4_K_M GGUF. It offers " + "high generation quality and a large context window, and is the heaviest " + "text-generation model in DashAI; a GPU is recommended. Model available at " + "https://huggingface.co/bartowski/Mistral-Nemo-Instruct-2407-GGUF." ), es=( - "Modelos ajustados para instrucciones de Mistral AI, cargados en formato " - "GGUF para inferencia eficiente en CPU y GPU mediante llama.cpp. " - "Los modelos " - "Mistral son conocidos por su fuerte rendimiento relativo a su cantidad de " - "parámetros e inferencia eficiente. Soporta conversación multi-turno, " - "razonamiento y generación de texto en general. Disponible en variantes de " - "7B (Mistral-7B-v0.3) y 12B (Mistral-Nemo-2407). Modelos en " - "https://huggingface.co/bartowski." + "Mistral Nemo Instruct 2407 es un modelo de lenguaje de 12B parametros " + "ajustado para instrucciones, creado por Mistral AI y NVIDIA, cargado como " + "GGUF Q4_K_M. Ofrece alta calidad de generacion y una gran ventana de " + "contexto; es el modelo mas pesado de DashAI y se recomienda una GPU." ), pt=( - "Modelos ajustados para instruções da Mistral AI, carregados em formato " - "GGUF para inferência eficiente em CPU e GPU via llama.cpp. Os modelos " - "Mistral são conhecidos pelo forte desempenho em relação à sua quantidade " - "de parâmetros e inferência eficiente. Suporta conversação multi-turno, " - "raciocínio e geração de texto em geral. Disponível nas variantes de " - "7B (Mistral-7B-v0.3) e 12B (Mistral-Nemo-2407). Modelos em " - "https://huggingface.co/bartowski." + "Mistral Nemo Instruct 2407 e um modelo de linguagem de 12B parametros " + "ajustado para instrucoes, criado pela Mistral AI e NVIDIA, carregado como " + "GGUF Q4_K_M. Oferece alta qualidade de geracao e uma grande janela de " + "contexto; e o modelo mais pesado do DashAI e uma GPU e recomendada." ), de=( - "Instruktionsoptimierte Mistral-Modelle von Mistral AI, im GGUF-Format " - "für effiziente CPU- und GPU-Inferenz über die llama.cpp-Bibliothek. " - "Mistral-Modelle sind bekannt für starke Leistung relativ zu ihrer " - "Parameteranzahl und effizienter Inferenz. Unterstützt Mehrfachdialog, " - "Schlussfolgerung und allgemeine Textgenerierung. Verfügbar in 7B " - "(Mistral-7B-v0.3) und 12B (Mistral-Nemo-2407) Varianten. Modelle unter " - "https://huggingface.co/bartowski." + "Mistral Nemo Instruct 2407 ist ein 12B-Parameter-Instruktionsmodell von " + "Mistral AI und NVIDIA, als Q4_K_M-GGUF geladen. Es bietet hohe " + "Generierungsqualitaet und ein grosses Kontextfenster und ist das " + "schwerste Textgenerierungsmodell in DashAI; eine GPU wird empfohlen." ), zh=( - "Mistral AI 的指令微调模型，以 GGUF 格式加载，" - "通过 llama.cpp 库实现高效的 CPU 和 GPU 推理。" - "支持多轮对话、推理和通用文本生成。提供 7B 和 12B 两种规格。" + "Mistral Nemo Instruct 2407 是 Mistral AI 与 NVIDIA 共同打造的 " + "120 亿参数指令微调语言模型，以 Q4_K_M GGUF 格式加载。" + "它具有高生成质量和大上下文窗口，是 DashAI 中最重的文本生成模型，" + "建议使用 GPU。" ), ) - - def __init__(self, **kwargs): - """Download and initialise a Mistral Instruct GGUF model via llama.cpp. - - The model weights are fetched from HuggingFace Hub using - ``Llama.from_pretrained`` and kept in memory for repeated calls to - ``generate``. - - Parameters - ---------- - **kwargs : dict - model_name : str, optional - HuggingFace repo ID for the GGUF checkpoint. - Defaults to ``"bartowski/Mistral-7B-Instruct-v0.3-GGUF"``. - max_tokens : int, optional - Maximum number of new tokens to generate per call. Default 100. - temperature : float, optional - Sampling temperature in [0.0, 1.0]. Default 0.7. - frequency_penalty : float, optional - Token-frequency penalty in [0.0, 2.0]. Default 0.1. - context_window : int, optional - Total token budget (prompt + response) for a single forward - pass. Default 512. - device : str, optional - Target device from ``LLAMA_DEVICE_ENUM``. Any value whose - index is >= 0 enables full GPU offload (``n_gpu_layers=-1``); - ``"CPU"`` runs fully in RAM. - - Raises - ------ - RuntimeError - If ``llama-cpp-python`` is not installed. - """ - try: - from llama_cpp import Llama - except ImportError as e: - raise RuntimeError( - "llama-cpp-python is not installed. " - "Please install it to use this model." - ) from e - - kwargs = self.validate_and_transform(kwargs) - self.model_name = kwargs.get( - "model_name", "bartowski/Mistral-7B-Instruct-v0.3-GGUF" - ) - self.max_tokens = kwargs.pop("max_tokens", 100) - self.temperature = kwargs.pop("temperature", 0.7) - self.frequency_penalty = kwargs.pop("frequency_penalty", 0.1) - self.n_ctx = kwargs.pop("context_window", 512) - - self.filename = "*Q4_K_M.gguf" - use_gpu = LLAMA_DEVICE_TO_IDX.get(kwargs.get("device")) >= 0 - - self.model = Llama.from_pretrained( - repo_id=self.model_name, - filename=self.filename, - verbose=True, - n_ctx=self.n_ctx, - n_gpu_layers=-1 if use_gpu else 0, - main_gpu=(LLAMA_DEVICE_TO_IDX.get(kwargs.get("device")) if use_gpu else 0), - ) - - def generate(self, prompt: list[dict[str, str]]) -> List[str]: - """Generate a reply for the given chat prompt. - - Parameters - ---------- - prompt : list of dict - Conversation history in OpenAI chat format. Each dict must contain - at least ``"role"`` (``"system"``, ``"user"``, or ``"assistant"``) - and ``"content"`` (the message text). - - Returns - ------- - list of str - A single-element list containing the model's reply text, extracted - from ``choices[0]["message"]["content"]``. - """ - output = self.model.create_chat_completion( - messages=prompt, - max_tokens=self.max_tokens, - temperature=self.temperature, - frequency_penalty=self.frequency_penalty, - ) - return [output["choices"][0]["message"]["content"]] diff --git a/DashAI/back/models/hugging_face/qwen_model.py b/DashAI/back/models/hugging_face/qwen_model.py index 475c536bc..56afdcb6c 100644 --- a/DashAI/back/models/hugging_face/qwen_model.py +++ b/DashAI/back/models/hugging_face/qwen_model.py @@ -1,468 +1,136 @@ -from typing import List +"""Qwen 2.5 Instruct GGUF checkpoint subclasses for DashAI.""" -from DashAI.back.core.schema_fields import ( - BaseSchema, - enum_field, - float_field, - int_field, - schema_field, -) from DashAI.back.core.utils import MultilingualString -from DashAI.back.models.text_to_text_generation_model import ( - TextToTextGenerationTaskModel, -) -from DashAI.back.models.utils import ( - LLAMA_DEVICE_ENUM, - LLAMA_DEVICE_PLACEHOLDER, - LLAMA_DEVICE_TO_IDX, +from DashAI.back.models.hugging_face.gguf_text_generation_base import ( + GGUFTextGenerationModel, + GGUFTextGenerationSchema, ) -class QwenSchema(BaseSchema): - """Schema for QwenModel hyperparameters. - - Configures the Qwen 2.5 Instruct checkpoint variant (0.5B or 1.5B), generation - length, sampling temperature, frequency penalty, context window, and target - device. The GGUF filename is selected automatically using a Q8_0 quantization - pattern; no manual filename override is exposed. - """ - - model_name: schema_field( - enum_field( - enum=[ - "Qwen/Qwen2.5-0.5B-Instruct-GGUF", - "Qwen/Qwen2.5-1.5B-Instruct-GGUF", - ] - ), - placeholder="Qwen/Qwen2.5-1.5B-Instruct-GGUF", - description=MultilingualString( - en=( - "The Qwen 2.5 Instruct checkpoint to load in GGUF format. " - "'0.5B' (500M parameters) is faster and uses less memory, suitable " - "for lightweight tasks on CPU. '1.5B' (1.5B parameters) is more " - "capable and produces higher-quality responses at the cost of " - "more memory and slightly slower inference." - ), - es=( - "El checkpoint Qwen 2.5 Instruct a cargar en formato GGUF. " - "'0.5B' (500M parámetros) es más rápido y usa menos memoria, " - "adecuado para tareas ligeras en CPU. '1.5B' (1.5B parámetros) " - "es más capaz y produce respuestas de mayor calidad a costa de " - "más memoria e inferencia levemente más lenta." - ), - pt=( - "O checkpoint Qwen 2.5 Instruct a carregar em formato GGUF. " - "'0.5B' (500M parâmetros) é mais rápido e usa menos memória, " - "adequado para tarefas leves em CPU. '1.5B' (1.5B parâmetros) " - "é mais capaz e produz respostas de maior qualidade ao custo de " - "mais memória e inferência levemente mais lenta." - ), - de=( - "Der im GGUF-Format zu ladende Qwen 2.5 Instruct-Checkpoint. " - "'0.5B' (500M Parameter) ist schneller und verbraucht weniger Speicher," - "geeignet für leichte CPU-Aufgaben. '1.5B' (1,5B Parameter) ist " - "leistungsfähiger und liefert qualitativ hochwertigere Antworten " - "auf Kosten von mehr Speicher und etwas langsamerer Inferenz." - ), - zh=( - "要加载的 Qwen 2.5 Instruct GGUF 格式检查点。" - "'0.5B'（5亿参数）速度更快、内存占用更少，适合 CPU 上的轻量级任务。" - "'1.5B'（15亿参数）能力更强，生成质量更高，但需要更多内存且推理速度略慢。" - ), - ), - alias=MultilingualString( - en="Model name", - es="Nombre del modelo", - pt="Nome do modelo", - de="Modellname", - zh="模型名称", - ), - ) # type: ignore +class Qwen25_05BInstruct(GGUFTextGenerationModel): # noqa: N801 + """Qwen 2.5 0.5B Instruct GGUF checkpoint (Q8_0 quantization). - max_tokens: schema_field( - int_field(ge=1), - placeholder=100, - description=MultilingualString( - en=( - "Maximum number of new tokens the model will generate per response. " - "Roughly 1 token ≈ 0.75 English words. Set to 100-200 for short " - "answers, 500-1000 for detailed explanations or code. Must not " - "exceed the context window minus the prompt length." - ), - es=( - "Número máximo de tokens nuevos que el modelo generará por respuesta. " - "Aproximadamente 1 token ≈ 0.75 palabras en español. Use 100-200 " - "para respuestas cortas, 500-1000 para explicaciones detalladas o " - "código. No debe superar la ventana de contexto menos la longitud " - "del prompt." - ), - pt=( - "Número máximo de tokens novos que o modelo gerará por resposta. " - "Aproximadamente 1 token ≈ 0.75 palavras em português. Use 100-200 " - "para respostas curtas, 500-1000 para explicações detalhadas ou " - "código. Não deve exceder a janela de contexto menos o comprimento " - "do prompt." - ), - de=( - "Maximale Anzahl neuer Token, die das Modell pro Antwort erzeugt. " - "Ungefähr 1 Token ≈ 0,75 englische Wörter. 100-200 für kurze " - "Antworten, 500-1000 für ausführliche Erklärungen oder Code. " - "Darf die Kontextfenstergröße abzüglich der Prompt-Länge nicht " - "überschreiten." - ), - zh=( - "模型每次响应生成的最大新 token 数量。" - "大约 1 token 约等于 0.75 个英文单词。短答案设为 100-200，" - "详细说明或代码设为 500-1000。不得超过上下文窗口减去提示词长度的值。" - ), - ), - alias=MultilingualString( - en="Max tokens", - es="Tokens máximos", - pt="Tokens máximos", - de="Maximale neue Token", - zh="最大 token 数", - ), - ) # type: ignore + A compact 500M-parameter instruction-tuned model from Alibaba Cloud, + well suited for lightweight CPU inference. Weights are stored locally + after a one-time download from HuggingFace. - temperature: schema_field( - float_field(ge=0.0, le=1.0), - placeholder=0.7, - description=MultilingualString( - en=( - "Sampling temperature controlling output randomness (range 0.0-1.0). " - "At 0.0 the model always picks the most likely token (greedy, fully " - "deterministic). Around 0.7 is a good balance for conversational " - "tasks. At 1.0 outputs are maximally varied and unpredictable." - ), - es=( - "Temperatura de muestreo que controla la aleatoriedad de la salida " - "(rango 0.0-1.0). En 0.0 el modelo siempre elige el token más " - "probable (greedy, totalmente determinista). Alrededor de 0.7 es " - "un buen equilibrio para tareas conversacionales. En 1.0 las " - "salidas son máximamente variadas e impredecibles." - ), - pt=( - "Temperatura de amostragem que controla a aleatoriedade da saída " - "(intervalo 0.0-1.0). Em 0.0 o modelo sempre escolhe o token mais " - "provável (greedy, totalmente determinístico). Em torno de 0.7 é " - "um bom equilíbrio para tarefas conversacionais. Em 1.0 as " - "saídas são maximamente variadas e imprevisíveis." - ), - de=( - "Stichprobentemperatur zur Steuerung der Ausgabezufälligkeit (0.0-1.0)." - "Bei 0.0 wählt das Modell stets den wahrscheinlichsten Token (greedy, " - "vollständig deterministisch). Um 0.7 ist ein gutes Gleichgewicht für " - "Konversationsaufgaben. Bei 1.0 sind Ausgaben maximal variiert und " - "unvorhersehbar." - ), - zh=( - "控制输出随机性的采样温度（范围 0.0-1.0）。" - "0.0 时模型始终选择最可能的 token（贪心，完全确定性）。" - "0.7 左右是对话任务的良好平衡点。1.0 时输出变化最大，不可预测。" - ), - ), - alias=MultilingualString( - en="Temperature", - es="Temperatura", - pt="Temperatura", - de="Temperatur", - zh="温度", - ), - ) # type: ignore + References + ---------- + - Qwen Team (2024). "Qwen2.5 Technical Report." + https://arxiv.org/abs/2412.15115 + - https://huggingface.co/Qwen/Qwen2.5-0.5B-Instruct-GGUF + """ - frequency_penalty: schema_field( - float_field(ge=0.0, le=2.0), - placeholder=0.1, - description=MultilingualString( - en=( - "Penalizes tokens that have already appeared in the output based on " - "how often they occur (range 0.0-2.0). At 0.0 there is no penalty " - "and the model may repeat itself. Values around 0.1-0.3 gently " - "discourage repetition. High values (1.5+) strongly prevent reuse " - "of any word, which may produce less coherent text." - ), - es=( - "Penaliza los tokens que ya aparecieron en la salida según su " - "frecuencia (rango 0.0-2.0). En 0.0 no hay penalización y el modelo " - "puede repetirse. Valores en torno a 0.1-0.3 desincentivan " - "suavemente la repetición. Valores altos (1.5+) previenen " - "fuertemente la reutilización de palabras, lo que puede producir " - "texto menos coherente." - ), - pt=( - "Penaliza os tokens que já apareceram na saída com base em " - "sua frequência (intervalo 0.0-2.0). Em 0.0 não há penalização e o " - "modelo pode se repetir. Valores em torno de 0.1-0.3 desestimulam " - "suavemente a repetição. Valores altos (1.5+) impedem fortemente a " - "reutilização de palavras, o que pode produzir texto menos coerente." - ), - de=( - "Bestraft Token, die bereits in der Ausgabe erschienen sind, " - "basierend auf ihrer Häufigkeit (0.0-2.0). Bei 0.0 gibt es keine " - "Strafe und das Modell kann sich wiederholen. Werte um 0.1-0.3 " - "hemmen Wiederholungen sanft. Hohe Werte (1.5+) verhindern die " - "Wiederverwendung von Wörtern stark, was zu weniger kohärentem Text " - "führen kann." - ), - zh=( - "根据 token 在输出中出现的频率对其进行惩罚（范围 0.0-2.0）。" - "0.0 时无惩罚，模型可能重复输出。0.1-0.3 左右可轻微抑制重复。" - "高值（1.5+）会强烈阻止任何词的复用，可能导致文本连贯性下降。" - ), - ), - alias=MultilingualString( - en="Frequency penalty", - es="Penalización de frecuencia", - pt="Penalização de frequência", - de="Häufigkeitsstrafe", - zh="频率惩罚", + REPO_ID = "Qwen/Qwen2.5-0.5B-Instruct-GGUF" + GGUF_PATTERN = "*8_0.gguf" + DOWNLOAD_SIZE_BYTES = 700_000_000 + SCHEMA = GGUFTextGenerationSchema + COLOR: str = "#2e7d32" + DISPLAY_NAME = MultilingualString( + en="Qwen2.5 0.5B Instruct", + es="Qwen2.5 0.5B Instruct", + pt="Qwen2.5 0.5B Instruct", + de="Qwen2.5 0.5B Instruct", + zh="Qwen2.5 0.5B Instruct", + ) + DESCRIPTION = MultilingualString( + en=( + "Qwen 2.5 0.5B Instruct is a 500M-parameter instruction-tuned " + "language model by Alibaba Cloud, loaded as a Q8_0 GGUF for " + "efficient CPU inference. It is the fastest and most " + "memory-efficient Qwen 2.5 variant in DashAI, ideal for rapid " + "prototyping or devices with limited RAM. Model available at " + "https://huggingface.co/Qwen/Qwen2.5-0.5B-Instruct-GGUF." ), - ) # type: ignore - - context_window: schema_field( - int_field(ge=1, le=32768), - placeholder=512, - description=MultilingualString( - en=( - "Total token budget for a single forward pass, including both the " - "input prompt and the generated response. Larger values allow longer " - "conversations but consume more RAM/VRAM. Qwen 2.5 supports up to " - "32768 tokens natively; keep this at or below that limit." - ), - es=( - "Presupuesto total de tokens para una sola pasada, incluyendo tanto " - "el prompt de entrada como la respuesta generada. Valores más altos " - "permiten conversaciones más largas pero consumen más RAM/VRAM. " - "Qwen 2.5 soporta hasta 32768 tokens de forma nativa; mantenga " - "este valor igual o por debajo de ese límite." - ), - pt=( - "Orçamento total de tokens para uma única passagem, incluindo tanto " - "o prompt de entrada quanto a resposta gerada. Valores maiores " - "permitem conversas mais longas mas consomem mais RAM/VRAM. " - "Qwen 2.5 suporta até 32768 tokens nativamente; mantenha " - "este valor igual ou abaixo desse limite." - ), - de=( - "Gesamtes Token-Budget für einen einzelnen Vorwärtsdurchlauf, " - "einschließlich Eingabe-Prompt und generierter Antwort. Größere Werte " - "ermöglichen längere Gespräche, verbrauchen jedoch mehr RAM/VRAM. " - "Qwen 2.5 unterstützt nativ bis zu 32768 Token; halten Sie " - "diesen Wert gleich oder unter diesem Limit." - ), - zh=( - "单次前向传播的总 token 预算，包含输入提示词和生成的响应。" - "较大的值允许更长的对话，但会消耗更多 RAM/VRAM。" - "Qwen 2.5 原生支持最多 32768 个 token，请保持此值不超过该限制。" - ), + es=( + "Qwen 2.5 0.5B Instruct es un modelo de 500M parámetros ajustado " + "para instrucciones por Alibaba Cloud, cargado como GGUF Q8_0 " + "para inferencia eficiente en CPU. Es la variante Qwen 2.5 más " + "rápida y con menor uso de memoria en DashAI, ideal para " + "prototipado rápido o dispositivos con RAM limitada." ), - alias=MultilingualString( - en="Context window", - es="Ventana de contexto", - pt="Janela de contexto", - de="Kontextfenster", - zh="上下文窗口", + pt=( + "Qwen 2.5 0.5B Instruct é um modelo de 500M parâmetros ajustado " + "para instruções pela Alibaba Cloud, carregado como GGUF Q8_0 " + "para inferência eficiente em CPU. É a variante Qwen 2.5 mais " + "rápida e com menor uso de memória no DashAI, ideal para " + "prototipagem rápida ou dispositivos com RAM limitada." ), - ) # type: ignore - - device: schema_field( - enum_field(enum=LLAMA_DEVICE_ENUM), - placeholder=LLAMA_DEVICE_PLACEHOLDER, - description=MultilingualString( - en=( - "Hardware device for llama.cpp inference. 'CPU' runs the model " - "fully in RAM with no GPU requirement. Selecting a GPU option " - "offloads all layers for faster inference, setting n_gpu_layers=-1 " - "so every transformer layer is GPU-accelerated." - ), - es=( - "Dispositivo de hardware para la inferencia con llama.cpp. 'CPU' " - "ejecuta el modelo completamente en RAM sin requisito de GPU. " - "Seleccionar una opción de GPU descarga todas las capas para " - "inferencia más rápida, estableciendo n_gpu_layers=-1 para que " - "cada capa del transformer sea acelerada por GPU." - ), - pt=( - "Dispositivo de hardware para inferência com llama.cpp. 'CPU' " - "executa o modelo completamente na RAM sem requisito de GPU. " - "Selecionar uma opção de GPU descarrega todas as camadas para " - "inferência mais rápida, definindo n_gpu_layers=-1 para que " - "cada camada do transformer seja acelerada por GPU." - ), - de=( - "Hardware-Gerät für die llama.cpp-Inferenz. 'CPU' führt das Modell " - "vollständig im RAM ohne GPU-Anforderung aus. Eine GPU-Option " - "lagert alle Schichten für schnellere Inferenz aus und setzt " - "n_gpu_layers=-1, damit jede Transformer-Schicht GPU-beschleunigt wird." - ), - zh=( - "llama.cpp 推理所使用的硬件设备。'CPU' 完全在内存中运行模型，无需 GPU。" - "选择 GPU 选项会将所有层卸载以加快推理速度，" - "并设置 n_gpu_layers=-1 使每个 Transformer 层均由 GPU 加速。" - ), + de=( + "Qwen 2.5 0.5B Instruct ist ein 500M-Parameter-Instruktionsmodell " + "von Alibaba Cloud, als Q8_0-GGUF für effiziente CPU-Inferenz " + "geladen. Es ist die schnellste und speichereffizienteste " + "Qwen-2.5-Variante in DashAI, ideal für schnelles Prototyping " + "oder Geräte mit begrenztem RAM." ), - alias=MultilingualString( - en="Device", - es="Dispositivo", - pt="Dispositivo", - de="Gerät", - zh="设备", + zh=( + "Qwen 2.5 0.5B Instruct 是阿里云推出的 5 亿参数指令微调语言模型，" + "以 Q8_0 GGUF 格式加载，支持高效 CPU 推理。" + "这是 DashAI 中速度最快、内存占用最低的 Qwen 2.5 变体，" + "非常适合快速原型开发或内存受限的设备。" ), - ) # type: ignore - + ) -class QwenModel(TextToTextGenerationTaskModel): - """Qwen 2.5 Instruct model for efficient text generation via llama.cpp. - Qwen 2.5 is a series of dense transformer language models from Alibaba Cloud, - spanning 0.5B to 72B parameters. The DashAI integration exposes the 0.5B and - 1.5B Instruct variants, which run comfortably on CPU. Both are trained on 18 - trillion tokens with improved coding, mathematics, and multilingual capability - over Qwen 2. +class Qwen25_15BInstruct(GGUFTextGenerationModel): # noqa: N801 + """Qwen 2.5 1.5B Instruct GGUF checkpoint (Q8_0 quantization). - Models are loaded as GGUF Q8_0 quantized checkpoints via ``llama-cpp-python``; - the quantization file is selected automatically from the HuggingFace repo. + A 1.5B-parameter instruction-tuned model from Alibaba Cloud that offers + higher response quality than the 0.5B variant while still running on CPU. + Weights are stored locally after a one-time download from HuggingFace. References ---------- - - [1] Qwen Team (2024). "Qwen2.5 Technical Report." https://arxiv.org/abs/2412.15115 - - [2] https://huggingface.co/Qwen/Qwen2.5-1.5B-Instruct-GGUF + - Qwen Team (2024). "Qwen2.5 Technical Report." + https://arxiv.org/abs/2412.15115 + - https://huggingface.co/Qwen/Qwen2.5-1.5B-Instruct-GGUF """ - SCHEMA = QwenSchema + REPO_ID = "Qwen/Qwen2.5-1.5B-Instruct-GGUF" + GGUF_PATTERN = "*8_0.gguf" + DOWNLOAD_SIZE_BYTES = 1_900_000_000 + SCHEMA = GGUFTextGenerationSchema COLOR: str = "#2e7d32" - DISPLAY_NAME: str = MultilingualString( - en="Qwen Model", - es="Modelo Qwen", - pt="Modelo Qwen", - de="Qwen-Modell", - zh="Qwen 模型", + DISPLAY_NAME = MultilingualString( + en="Qwen2.5 1.5B Instruct", + es="Qwen2.5 1.5B Instruct", + pt="Qwen2.5 1.5B Instruct", + de="Qwen2.5 1.5B Instruct", + zh="Qwen2.5 1.5B Instruct", ) - DESCRIPTION: str = MultilingualString( + DESCRIPTION = MultilingualString( en=( - "Qwen 2.5 is an instruction-tuned large language model by Alibaba Cloud, " - "loaded in GGUF format for efficient CPU and GPU inference via the " - "llama.cpp library. It supports multi-turn conversation, reasoning, " - "coding, and general text generation. Available in 0.5B and 1.5B " - "parameter sizes. Models are available at " - "https://huggingface.co/Qwen/Qwen2.5-0.5B-Instruct-GGUF and " + "Qwen 2.5 1.5B Instruct is a 1.5B-parameter instruction-tuned " + "language model by Alibaba Cloud, loaded as a Q8_0 GGUF for " + "efficient CPU inference. It provides stronger reasoning and " + "generation quality than the 0.5B variant at the cost of slightly " + "more memory. Model available at " "https://huggingface.co/Qwen/Qwen2.5-1.5B-Instruct-GGUF." ), es=( - "Qwen 2.5 es un modelo de lenguaje grande ajustado para instrucciones " - "por Alibaba Cloud, cargado en formato GGUF para inferencia eficiente en " - "CPU y GPU mediante la librería llama.cpp. Soporta conversación " - "multi-turno, razonamiento, programación y generación de texto en " - "general. Disponible en tamaños de 0.5B y 1.5B parámetros. Los modelos " - "están disponibles en " - "https://huggingface.co/Qwen/Qwen2.5-0.5B-Instruct-GGUF y " - "https://huggingface.co/Qwen/Qwen2.5-1.5B-Instruct-GGUF." + "Qwen 2.5 1.5B Instruct es un modelo de 1.5B parámetros ajustado " + "para instrucciones por Alibaba Cloud, cargado como GGUF Q8_0 " + "para inferencia eficiente en CPU. Ofrece mayor capacidad de " + "razonamiento y calidad de generación que la variante de 0.5B a " + "costa de un poco más de memoria." ), pt=( - "Qwen 2.5 é um modelo de linguagem grande ajustado para instruções " - "pela Alibaba Cloud, carregado em formato GGUF para inferência eficiente " - "em CPU e GPU via biblioteca llama.cpp. Suporta conversa multi-turno, " - "raciocínio, programação e geração de texto em geral. Disponível nos " - "tamanhos 0.5B e 1.5B parâmetros. Os modelos estão disponíveis em " - "https://huggingface.co/Qwen/Qwen2.5-0.5B-Instruct-GGUF e " - "https://huggingface.co/Qwen/Qwen2.5-1.5B-Instruct-GGUF." + "Qwen 2.5 1.5B Instruct é um modelo de 1.5B parâmetros ajustado " + "para instruções pela Alibaba Cloud, carregado como GGUF Q8_0 " + "para inferência eficiente em CPU. Oferece melhor raciocínio e " + "qualidade de geração que a variante de 0.5B ao custo de um pouco " + "mais de memória." ), de=( - "Qwen 2.5 ist ein instruktionsoptimiertes großes Sprachmodell von " - "Alibaba Cloud, im GGUF-Format für effiziente CPU- und GPU-Inferenz über " - "die llama.cpp-Bibliothek geladen. Unterstützt Mehrfachdialog, " - "Schlussfolgerung, Programmierung und allgemeine Textgenerierung. " - "Verfügbar in den Parametergrößen 0,5B und 1,5B. Modelle verfügbar unter " - "https://huggingface.co/Qwen/Qwen2.5-0.5B-Instruct-GGUF und " - "https://huggingface.co/Qwen/Qwen2.5-1.5B-Instruct-GGUF." + "Qwen 2.5 1.5B Instruct ist ein 1,5B-Parameter-Instruktionsmodell " + "von Alibaba Cloud, als Q8_0-GGUF für effiziente CPU-Inferenz " + "geladen. Es bietet besseres Schlussfolgern und " + "Generierungsqualität als die 0,5B-Variante, benötigt jedoch " + "etwas mehr Speicher." ), zh=( - "Qwen 2.5 是阿里云开发的指令微调大语言模型，" - "以 GGUF 格式加载，通过 llama.cpp 库实现高效的 CPU 和 GPU 推理。" - "支持多轮对话、推理、编程和通用文本生成。提供 0.5B 和 1.5B 参数规格。" + "Qwen 2.5 1.5B Instruct 是阿里云推出的 15 亿参数指令微调语言模型，" + "以 Q8_0 GGUF 格式加载，支持高效 CPU 推理。" + "与 0.5B 变体相比，它具有更强的推理能力和生成质量，但需要稍多的内存。" ), ) - - def __init__(self, **kwargs): - """Download and initialise a Qwen 2.5 Instruct GGUF model via llama.cpp. - - The model weights are fetched from HuggingFace Hub using - ``Llama.from_pretrained`` and kept in memory for repeated calls to - ``generate``. The Q8_0 quantization file is always selected regardless - of the chosen model variant. - - Parameters - ---------- - **kwargs : dict - model_name : str, optional - HuggingFace repo ID for the GGUF checkpoint. - Defaults to ``"Qwen/Qwen2.5-1.5B-Instruct-GGUF"``. - max_tokens : int, optional - Maximum number of new tokens to generate per call. Default 100. - temperature : float, optional - Sampling temperature in [0.0, 1.0]. Default 0.7. - frequency_penalty : float, optional - Token-frequency penalty in [0.0, 2.0]. Default 0.1. - context_window : int, optional - Total token budget (prompt + response) for a single forward - pass. Default 512. - device : str, optional - Target device from ``LLAMA_DEVICE_ENUM``. Any value whose - index is >= 0 enables full GPU offload (``n_gpu_layers=-1``); - ``"CPU"`` runs fully in RAM. - - Raises - ------ - RuntimeError - If ``llama-cpp-python`` is not installed. - """ - try: - from llama_cpp import Llama - except ImportError as e: - raise RuntimeError( - "llama-cpp-python is not installed. Please install it to use QwenModel." - ) from e - - kwargs = self.validate_and_transform(kwargs) - self.model_name = kwargs.get("model_name", "Qwen/Qwen2.5-1.5B-Instruct-GGUF") - self.max_tokens = kwargs.pop("max_tokens", 100) - self.temperature = kwargs.pop("temperature", 0.7) - self.frequency_penalty = kwargs.pop("frequency_penalty", 0.1) - self.n_ctx = kwargs.pop("context_window", 512) - - self.filename = "*8_0.gguf" - use_gpu = LLAMA_DEVICE_TO_IDX.get(kwargs.get("device")) >= 0 - - self.model = Llama.from_pretrained( - repo_id=self.model_name, - filename=self.filename, - verbose=True, - n_ctx=self.n_ctx, - n_gpu_layers=-1 if use_gpu else 0, - main_gpu=LLAMA_DEVICE_TO_IDX.get(kwargs.get("device")) if use_gpu else 0, - ) - - def generate(self, prompt: list[dict[str, str]]) -> List[str]: - """Generate a reply for the given chat prompt. - - Parameters - ---------- - prompt : list of dict - Conversation history in OpenAI chat format. Each dict must contain - at least ``"role"`` (``"system"``, ``"user"``, or ``"assistant"``) - and ``"content"`` (the message text). - - Returns - ------- - list of str - A single-element list containing the model's reply text, extracted - from ``choices[0]["message"]["content"]``. - """ - output = self.model.create_chat_completion( - messages=prompt, - max_tokens=self.max_tokens, - temperature=self.temperature, - frequency_penalty=self.frequency_penalty, - ) - return [output["choices"][0]["message"]["content"]] diff --git a/DashAI/back/models/hugging_face/smol_lm_model.py b/DashAI/back/models/hugging_face/smol_lm_model.py index a5f56bdd9..b2999369d 100644 --- a/DashAI/back/models/hugging_face/smol_lm_model.py +++ b/DashAI/back/models/hugging_face/smol_lm_model.py @@ -1,454 +1,123 @@ -from typing import List +"""SmolLM2 Instruct GGUF checkpoint subclasses for DashAI.""" -from DashAI.back.core.schema_fields import ( - BaseSchema, - enum_field, - float_field, - int_field, - schema_field, -) from DashAI.back.core.utils import MultilingualString -from DashAI.back.models.text_to_text_generation_model import ( - TextToTextGenerationTaskModel, -) -from DashAI.back.models.utils import ( - LLAMA_DEVICE_ENUM, - LLAMA_DEVICE_PLACEHOLDER, - LLAMA_DEVICE_TO_IDX, +from DashAI.back.models.hugging_face.gguf_text_generation_base import ( + GGUFTextGenerationModel, + GGUFTextGenerationSchema, ) -SMOLLM_FILENAME_MAP = { - "HuggingFaceTB/SmolLM2-1.7B-Instruct-GGUF": "*q4_k_m.gguf", - "HuggingFaceTB/SmolLM2-360M-Instruct-GGUF": "*q8_0.gguf", -} +class SmolLM2_360MInstruct(GGUFTextGenerationModel): # noqa: N801 + """SmolLM2 360M Instruct GGUF checkpoint (Q8_0 quantization). -class SmolLMSchema(BaseSchema): - """Schema for SmolLM2 model hyperparameters. + A very small 360M-parameter instruction-tuned model from HuggingFace, + designed for fast on-device inference. Weights are stored locally after a + one-time download from HuggingFace. - Configures the SmolLM2 Instruct checkpoint variant (360M or 1.7B), generation - length, sampling temperature, frequency penalty, context window, and target - device. The GGUF filename is resolved automatically from ``SMOLLM_FILENAME_MAP``: - Q4_K_M quantization for 1.7B and Q8_0 quantization for 360M. + References + ---------- + - https://huggingface.co/HuggingFaceTB/SmolLM2-360M-Instruct-GGUF """ - model_name: schema_field( - enum_field( - enum=[ - "HuggingFaceTB/SmolLM2-1.7B-Instruct-GGUF", - "HuggingFaceTB/SmolLM2-360M-Instruct-GGUF", - ] - ), - placeholder="HuggingFaceTB/SmolLM2-1.7B-Instruct-GGUF", - description=MultilingualString( - en=( - "The SmolLM2 Instruct checkpoint to load in GGUF format. " - "'SmolLM2-1.7B' is a 1.7B-parameter instruction model with strong " - "performance for on-device and edge inference. " - "'SmolLM2-360M' is an ultra-compact 360M-parameter model for " - "extremely fast CPU inference with minimal memory usage (~300 MB). " - "Both models are trained on diverse synthetic datasets by Hugging Face." - ), - es=( - "El checkpoint SmolLM2 Instruct a cargar en formato GGUF. " - "'SmolLM2-1.7B' es un modelo de instrucción de 1.7B parámetros con " - "fuerte rendimiento para inferencia en dispositivos y en el borde. " - "'SmolLM2-360M' es un modelo ultra-compacto de 360M parámetros para " - "inferencia CPU extremadamente rápida con uso mínimo de memoria " - "(~300 MB). " - "Ambos modelos son entrenados en datasets sintéticos diversos por " - "Hugging Face." - ), - pt=( - "O checkpoint SmolLM2 Instruct a carregar em formato GGUF. " - "'SmolLM2-1.7B' é um modelo de instrução de 1.7B parâmetros com " - "forte desempenho para inferência em dispositivos e na borda. " - "'SmolLM2-360M' é um modelo ultra-compacto de 360M parâmetros para " - "inferência CPU extremamente rápida com uso mínimo de memória " - "(~300 MB). " - "Ambos os modelos são treinados em conjuntos de dados sintéticos " - "diversos pelo Hugging Face." - ), - de=( - "Der im GGUF-Format zu ladende SmolLM2 Instruct-Checkpoint. " - "'SmolLM2-1.7B' ist ein 1,7B-Parameter-Instruktionsmodell mit starker " - "Leistung für Inferenz auf Endgeräten und Edge-Systemen. " - "'SmolLM2-360M' ist ein ultra-kompaktes 360M-Parameter-Modell für " - "extrem schnelle CPU-Inferenz mit minimalem Speicherbedarf (~300 MB). " - "Beide Modelle werden von Hugging Face auf diversen synthetischen " - "Datensätzen trainiert." - ), - zh=( - "以 GGUF 格式加载的 SmolLM2 Instruct 检查点。" - "'SmolLM2-1.7B' 是 17 亿参数指令模型，适用于端侧和边缘推理。" - "'SmolLM2-360M' 是 3.6 亿参数超紧凑模型，CPU 推理极快，" - "内存占用极低（约 300 MB）。" - "两款模型均由 Hugging Face 在多样化合成数据集上训练。" - ), - ), - alias=MultilingualString( - en="Model name", - es="Nombre del modelo", - pt="Nome do modelo", - de="Modellname", - zh="模型名称", - ), - ) # type: ignore - - max_tokens: schema_field( - int_field(ge=1), - placeholder=100, - description=MultilingualString( - en=( - "Maximum number of new tokens the model will generate per response. " - "Roughly 1 token ≈ 0.75 English words. SmolLM2 models are optimized " - "for short to medium-length responses." - ), - es=( - "Número máximo de tokens nuevos que el modelo generará por respuesta. " - "Aproximadamente 1 token ≈ 0.75 palabras en español. Los modelos " - "SmolLM2 están optimizados para respuestas cortas a medianas." - ), - pt=( - "Número máximo de tokens novos que o modelo gerará por resposta. " - "Aproximadamente 1 token ≈ 0.75 palavras em português. Os modelos " - "SmolLM2 são otimizados para respostas curtas a médias." - ), - de=( - "Maximale Anzahl neuer Token, die das Modell pro Antwort erzeugt. " - "Ungefähr 1 Token ≈ 0,75 englische Wörter. SmolLM2-Modelle sind " - "für kurze bis mittellange Antworten optimiert." - ), - zh=( - "模型每次响应生成的最大新词元数。" - "约 1 词元 ≈ 0.75 个英文单词。" - "SmolLM2 模型针对短至中等长度的响应进行了优化。" - ), - ), - alias=MultilingualString( - en="Max tokens", - es="Tokens máximos", - pt="Tokens máximos", - de="Maximale neue Token", - zh="最大词元数", - ), - ) # type: ignore - - temperature: schema_field( - float_field(ge=0.0, le=1.0), - placeholder=0.7, - description=MultilingualString( - en=( - "Sampling temperature controlling output randomness (range 0.0-1.0). " - "At 0.0 outputs are deterministic. Around 0.7 balances quality and " - "creativity." - ), - es=( - "Temperatura de muestreo que controla la aleatoriedad (rango 0.0-1.0). " - "En 0.0 las salidas son deterministas. Alrededor de 0.7 equilibra " - "calidad y creatividad." - ), - pt=( - "Temperatura de amostragem que controla a aleatoriedade da saída " - "(intervalo 0.0-1.0). Em 0.0 as saídas são determinísticas. " - "Em torno de 0.7 equilibra qualidade e criatividade." - ), - de=( - "Stichprobentemperatur zur Steuerung der Ausgabezufälligkeit (0.0-1.0)." - "Bei 0.0 sind die Ausgaben deterministisch. Um 0.7 balanciert " - "Qualität und Kreativität." - ), - zh=( - "控制输出随机性的采样温度（范围 0.0-1.0）。" - "0.0 时输出为确定性结果，0.7 左右可平衡质量与创造力。" - ), - ), - alias=MultilingualString( - en="Temperature", - es="Temperatura", - pt="Temperatura", - de="Temperatur", - zh="温度", - ), - ) # type: ignore - - frequency_penalty: schema_field( - float_field(ge=0.0, le=2.0), - placeholder=0.1, - description=MultilingualString( - en=( - "Penalizes tokens that have already appeared in the output based on " - "frequency (range 0.0-2.0). Higher values discourage repetition." - ), - es=( - "Penaliza los tokens que ya aparecieron en la salida según su " - "frecuencia (rango 0.0-2.0). Valores más altos desincentivan " - "la repetición." - ), - pt=( - "Penaliza os tokens que já apareceram na saída com base na " - "frequência (intervalo 0.0-2.0). Valores mais altos desestimulam " - "a repetição." - ), - de=( - "Bestraft Token, die bereits in der Ausgabe erschienen sind, " - "basierend auf ihrer Häufigkeit (0.0-2.0). Höhere Werte reduzieren " - "Wiederholungen." - ), - zh=( - "根据词元在输出中出现的频率对其进行惩罚（范围 0.0-2.0）。" - "较高的值可抑制重复内容。" - ), - ), - alias=MultilingualString( - en="Frequency penalty", - es="Penalización de frecuencia", - pt="Penalização de frequência", - de="Häufigkeitsstrafe", - zh="频率惩罚", + REPO_ID = "HuggingFaceTB/SmolLM2-360M-Instruct-GGUF" + GGUF_PATTERN = "*q8_0.gguf" + DOWNLOAD_SIZE_BYTES = 400_000_000 + SCHEMA = GGUFTextGenerationSchema + COLOR: str = "#00695c" + DISPLAY_NAME = MultilingualString( + en="SmolLM2 360M Instruct", + es="SmolLM2 360M Instruct", + pt="SmolLM2 360M Instruct", + de="SmolLM2 360M Instruct", + zh="SmolLM2 360M Instruct", + ) + DESCRIPTION = MultilingualString( + en=( + "SmolLM2 360M Instruct is a compact 360M-parameter instruction-tuned " + "language model from HuggingFace, loaded as a Q8_0 GGUF for very fast " + "CPU inference. It is the lightest text-generation model in DashAI, " + "ideal for constrained devices. Model available at " + "https://huggingface.co/HuggingFaceTB/SmolLM2-360M-Instruct-GGUF." ), - ) # type: ignore - - context_window: schema_field( - int_field(ge=1, le=8192), - placeholder=512, - description=MultilingualString( - en=( - "Total token budget for a single forward pass, including both the " - "input prompt and the generated response. SmolLM2 models support " - "up to 8K tokens natively." - ), - es=( - "Presupuesto total de tokens por pasada, incluyendo prompt y " - "respuesta. Los modelos SmolLM2 soportan hasta 8K tokens de " - "forma nativa." - ), - pt=( - "Orçamento total de tokens por passagem, incluindo prompt e " - "resposta. Os modelos SmolLM2 suportam até 8K tokens nativamente." - ), - de=( - "Gesamtes Token-Budget für einen einzelnen Vorwärtsdurchlauf, " - "einschließlich Eingabe-Prompt und Antwort. " - "SmolLM2-Modelle unterstützen nativ bis zu 8K Token." - ), - zh=( - "单次前向传播的词元总预算，包含输入提示和生成响应。" - "SmolLM2 模型原生支持最多 8K 词元。" - ), + es=( + "SmolLM2 360M Instruct es un modelo de lenguaje compacto de 360M " + "parametros ajustado para instrucciones por HuggingFace, cargado como " + "GGUF Q8_0 para inferencia muy rapida en CPU. Es el modelo de generacion " + "de texto mas ligero de DashAI, ideal para dispositivos limitados." ), - alias=MultilingualString( - en="Context window", - es="Ventana de contexto", - pt="Janela de contexto", - de="Kontextfenster", - zh="上下文窗口", + pt=( + "SmolLM2 360M Instruct e um modelo de linguagem compacto de 360M " + "parametros ajustado para instrucoes pela HuggingFace, carregado como " + "GGUF Q8_0 para inferencia muito rapida em CPU. E o modelo de geracao de " + "texto mais leve do DashAI, ideal para dispositivos limitados." ), - ) # type: ignore - - device: schema_field( - enum_field(enum=LLAMA_DEVICE_ENUM), - placeholder=LLAMA_DEVICE_PLACEHOLDER, - description=MultilingualString( - en=( - "Hardware device for llama.cpp inference. 'CPU' runs the model " - "fully in RAM with no GPU requirement. SmolLM2 models are small " - "enough to run efficiently on CPU even on modest hardware." - ), - es=( - "Dispositivo de hardware para inferencia con llama.cpp. 'CPU' ejecuta " - "el modelo en RAM sin requisito de GPU. Los modelos SmolLM2 son lo " - "suficientemente pequeños para ejecutarse eficientemente en CPU " - "incluso " - "en hardware modesto." - ), - pt=( - "Dispositivo de hardware para inferência com llama.cpp. 'CPU' executa " - "o modelo na RAM sem requisito de GPU. Os modelos SmolLM2 são " - "pequenos o suficiente para rodar eficientemente em CPU " - "mesmo em hardware modesto." - ), - de=( - "Hardware-Gerät für die llama.cpp-Inferenz. 'CPU' führt das Modell " - "im RAM ohne GPU-Anforderung aus. SmolLM2-Modelle sind klein genug, " - "um auch auf bescheidener Hardware effizient auf der CPU zu laufen." - ), - zh=( - "llama.cpp 推理所用的硬件设备。'CPU' 将模型完全加载至内存运行，" - "无需 GPU。" - "SmolLM2 模型体积小巧，即使在普通硬件上也能高效地在 CPU 上运行。" - ), + de=( + "SmolLM2 360M Instruct ist ein kompaktes 360M-Parameter-Instruktionsmodell " + "von HuggingFace, als Q8_0-GGUF fuer sehr schnelle CPU-Inferenz geladen. " + "Es ist das leichteste Textgenerierungsmodell in DashAI, ideal fuer " + "eingeschraenkte Geraete." ), - alias=MultilingualString( - en="Device", es="Dispositivo", pt="Dispositivo", de="Gerät", zh="设备" + zh=( + "SmolLM2 360M Instruct 是 HuggingFace 推出的 3.6 亿参数指令微调语言模型，" + "以 Q8_0 GGUF 格式加载，支持极快的 CPU 推理。" + "这是 DashAI 中最轻量的文本生成模型，非常适合资源受限的设备。" ), - ) # type: ignore - - -class SmolLMModel(TextToTextGenerationTaskModel): - """SmolLM2 Instruct model for on-device text generation via llama.cpp. + ) - SmolLM2 is a family of compact, instruction-tuned language models developed by - Hugging Face TB, designed for efficient on-device and edge deployment. Unlike - larger language models, SmolLM2 achieves competitive benchmark results at very - small parameter counts by training on high-quality synthetic datasets including - cosmopedia-v2, FineWeb-Edu, and StackEdu. - The DashAI integration exposes the 360M and 1.7B Instruct variants. The 360M - model requires under 300 MB of RAM and runs comfortably on modest CPU hardware; - the 1.7B model delivers higher-quality responses while remaining deployable - without a GPU. +class SmolLM2_17BInstruct(GGUFTextGenerationModel): # noqa: N801 + """SmolLM2 1.7B Instruct GGUF checkpoint (Q4_K_M quantization). - Models are loaded as GGUF quantized checkpoints via ``llama-cpp-python``. The - quantization level is variant-dependent: Q8_0 for 360M (higher fidelity at small - size) and Q4_K_M for 1.7B (balanced quality/size trade-off). The filename is - resolved automatically from ``SMOLLM_FILENAME_MAP``. + A 1.7B-parameter instruction-tuned model from HuggingFace offering stronger + generation quality than the 360M variant. Weights are stored locally after a + one-time download from HuggingFace. References ---------- - - [1] Allal, L.B. et al. (2024). "SmolLM2: with great data, comes great - performance." Hugging Face Blog. - https://huggingface.co/blog/smollm2 - - [2] https://huggingface.co/HuggingFaceTB/SmolLM2-1.7B-Instruct-GGUF - - [3] https://huggingface.co/HuggingFaceTB/SmolLM2-360M-Instruct-GGUF + - https://huggingface.co/HuggingFaceTB/SmolLM2-1.7B-Instruct-GGUF """ - SCHEMA = SmolLMSchema + REPO_ID = "HuggingFaceTB/SmolLM2-1.7B-Instruct-GGUF" + GGUF_PATTERN = "*q4_k_m.gguf" + DOWNLOAD_SIZE_BYTES = 1_100_000_000 + SCHEMA = GGUFTextGenerationSchema COLOR: str = "#00695c" - DISPLAY_NAME: str = MultilingualString( - en="SmolLM Model", - es="Modelo SmolLM", - pt="Modelo SmolLM", - de="SmolLM-Modell", - zh="SmolLM 模型", + DISPLAY_NAME = MultilingualString( + en="SmolLM2 1.7B Instruct", + es="SmolLM2 1.7B Instruct", + pt="SmolLM2 1.7B Instruct", + de="SmolLM2 1.7B Instruct", + zh="SmolLM2 1.7B Instruct", ) - DESCRIPTION: str = MultilingualString( + DESCRIPTION = MultilingualString( en=( - "SmolLM2 is a family of compact instruction-tuned language models by " - "Hugging Face, loaded in GGUF format for efficient CPU and GPU inference " - "via the llama.cpp library. Designed for on-device and edge deployment, " - "SmolLM2 achieves strong benchmark results at very small parameter counts. " - "The 360M variant requires less than 300 MB of RAM, making it ideal for " - "resource-constrained environments. Available in 360M and 1.7B variants. " - "Models available at https://huggingface.co/HuggingFaceTB." + "SmolLM2 1.7B Instruct is a 1.7B-parameter instruction-tuned language " + "model from HuggingFace, loaded as a Q4_K_M GGUF for efficient CPU " + "inference. It provides better reasoning and generation quality than the " + "360M variant. Model available at " + "https://huggingface.co/HuggingFaceTB/SmolLM2-1.7B-Instruct-GGUF." ), es=( - "SmolLM2 es una familia de modelos de lenguaje compactos ajustados para " - "instrucciones por Hugging Face, cargados en formato GGUF para inferencia " - "eficiente en CPU y GPU mediante llama.cpp. Diseñados para despliegue en " - "dispositivo y en el borde, SmolLM2 logra fuertes resultados de benchmark " - "con muy pocos parámetros. La variante de 360M requiere menos de 300 MB de " - "RAM, ideal para entornos con recursos limitados. Disponible en variantes " - "de 360M y 1.7B. Modelos en https://huggingface.co/HuggingFaceTB." + "SmolLM2 1.7B Instruct es un modelo de lenguaje de 1.7B parametros " + "ajustado para instrucciones por HuggingFace, cargado como GGUF Q4_K_M " + "para inferencia eficiente en CPU. Ofrece mejor razonamiento y calidad de " + "generacion que la variante de 360M." ), pt=( - "SmolLM2 é uma família de modelos de linguagem compactos ajustados para " - "instruções pelo Hugging Face, carregados em formato GGUF para inferência " - "eficiente em CPU e GPU via llama.cpp. Projetados para implantação em " - "dispositivos e na borda, SmolLM2 alcança fortes resultados de benchmark " - "com pouquíssimos parâmetros. A variante de 360M requer menos de 300 MB de " - "RAM, ideal para ambientes com recursos limitados. " - "Disponível nas variantes " - "360M e 1.7B. Modelos disponíveis em https://huggingface.co/HuggingFaceTB." + "SmolLM2 1.7B Instruct e um modelo de linguagem de 1.7B parametros " + "ajustado para instrucoes pela HuggingFace, carregado como GGUF Q4_K_M " + "para inferencia eficiente em CPU. Oferece melhor raciocinio e qualidade " + "de geracao do que a variante de 360M." ), de=( - "SmolLM2 ist eine Familie kompakter instruktionsoptimierter Sprachmodelle " - "von Hugging Face, im GGUF-Format für effiziente CPU- und GPU-Inferenz " - "über llama.cpp geladen. Für Deployment auf Endgeräten und Edge-Systemen " - "konzipiert, erzielt SmolLM2 starke Benchmark-Ergebnisse mit sehr wenigen " - "Parametern. Die 360M-Variante benötigt weniger als 300 MB RAM und ist " - "ideal für ressourcenbeschränkte Umgebungen. Verfügbar in den Varianten " - "360M und 1,7B. Modelle unter https://huggingface.co/HuggingFaceTB." + "SmolLM2 1.7B Instruct ist ein 1,7B-Parameter-Instruktionsmodell von " + "HuggingFace, als Q4_K_M-GGUF fuer effiziente CPU-Inferenz geladen. " + "Es bietet besseres Schlussfolgern und Generierungsqualitaet als die " + "360M-Variante." ), zh=( - "SmolLM2 是 Hugging Face 推出的紧凑型指令微调语言模型系列，" - "以 GGUF 格式加载，通过 llama.cpp 库高效推理。" - "专为端侧和边缘部署设计，提供 360M 和 1.7B 两种规格。" + "SmolLM2 1.7B Instruct 是 HuggingFace 推出的 17 亿参数指令微调语言模型，" + "以 Q4_K_M GGUF 格式加载，支持高效 CPU 推理。" + "与 360M 变体相比，它具有更强的推理能力和生成质量。" ), ) - - def __init__(self, **kwargs): - """Download and initialise a SmolLM2 Instruct GGUF model via llama.cpp. - - The model weights are fetched from HuggingFace Hub using - ``Llama.from_pretrained`` and kept in memory for repeated calls to - ``generate``. The GGUF filename is resolved from ``SMOLLM_FILENAME_MAP`` - (Q4_K_M for 1.7B, Q8_0 for 360M). - - Parameters - ---------- - **kwargs : dict - model_name : str, optional - HuggingFace repo ID for the GGUF checkpoint. - Defaults to ``"HuggingFaceTB/SmolLM2-1.7B-Instruct-GGUF"``. - max_tokens : int, optional - Maximum number of new tokens to generate per call. Default 100. - temperature : float, optional - Sampling temperature in [0.0, 1.0]. Default 0.7. - frequency_penalty : float, optional - Token-frequency penalty in [0.0, 2.0]. Default 0.1. - context_window : int, optional - Total token budget (prompt + response) for a single forward - pass. Default 512. - device : str, optional - Target device from ``LLAMA_DEVICE_ENUM``. Any value whose - index is >= 0 enables full GPU offload (``n_gpu_layers=-1``); - ``"CPU"`` runs fully in RAM. - - Raises - ------ - RuntimeError - If ``llama-cpp-python`` is not installed. - """ - try: - from llama_cpp import Llama - except ImportError as e: - raise RuntimeError( - "llama-cpp-python is not installed. " - "Please install it to use this model." - ) from e - - kwargs = self.validate_and_transform(kwargs) - self.model_name = kwargs.get( - "model_name", "HuggingFaceTB/SmolLM2-1.7B-Instruct-GGUF" - ) - self.max_tokens = kwargs.pop("max_tokens", 100) - self.temperature = kwargs.pop("temperature", 0.7) - self.frequency_penalty = kwargs.pop("frequency_penalty", 0.1) - self.n_ctx = kwargs.pop("context_window", 512) - - self.filename = SMOLLM_FILENAME_MAP.get(self.model_name, "*q4_k_m.gguf") - use_gpu = LLAMA_DEVICE_TO_IDX.get(kwargs.get("device")) >= 0 - - self.model = Llama.from_pretrained( - repo_id=self.model_name, - filename=self.filename, - verbose=True, - n_ctx=self.n_ctx, - n_gpu_layers=-1 if use_gpu else 0, - main_gpu=(LLAMA_DEVICE_TO_IDX.get(kwargs.get("device")) if use_gpu else 0), - ) - - def generate(self, prompt: list[dict[str, str]]) -> List[str]: - """Generate a reply for the given chat prompt. - - Parameters - ---------- - prompt : list of dict - Conversation history in OpenAI chat format. Each dict must contain - at least ``"role"`` (``"system"``, ``"user"``, or ``"assistant"``) - and ``"content"`` (the message text). - - Returns - ------- - list of str - A single-element list containing the model's reply text, extracted - from ``choices[0]["message"]["content"]``. - """ - output = self.model.create_chat_completion( - messages=prompt, - max_tokens=self.max_tokens, - temperature=self.temperature, - frequency_penalty=self.frequency_penalty, - ) - return [output["choices"][0]["message"]["content"]] diff --git a/tests/back/models/test_gguf_text_generation.py b/tests/back/models/test_gguf_text_generation.py new file mode 100644 index 000000000..5c40da2c2 --- /dev/null +++ b/tests/back/models/test_gguf_text_generation.py @@ -0,0 +1,89 @@ +"""Tests for the per-checkpoint GGUF text-generation components.""" + +import pathlib + +import pytest +from kink import di + +from DashAI.back.dependencies.downloads.downloadable import HFDownloadableMixin +from DashAI.back.initial_components import get_initial_components +from DashAI.back.models.hugging_face.llama_model import ( + Llama31_8BInstruct, + Llama32_1BInstruct, + Llama32_3BInstruct, +) +from DashAI.back.models.hugging_face.mistral_model import ( + Mistral7BInstructV03, + MistralNemoInstruct2407, +) +from DashAI.back.models.hugging_face.qwen_model import ( + Qwen25_05BInstruct, + Qwen25_15BInstruct, +) +from DashAI.back.models.hugging_face.smol_lm_model import ( + SmolLM2_17BInstruct, + SmolLM2_360MInstruct, +) + +ALL_CHECKPOINTS = [ + Qwen25_05BInstruct, + Qwen25_15BInstruct, + SmolLM2_360MInstruct, + SmolLM2_17BInstruct, + Llama31_8BInstruct, + Llama32_1BInstruct, + Llama32_3BInstruct, + Mistral7BInstructV03, + MistralNemoInstruct2407, +] + + +@pytest.fixture +def component_root(tmp_path): + sentinel = object() + old = di["config"] if "config" in di else sentinel # noqa: SIM401 + di["config"] = {"COMPONENT_PATH": str(tmp_path)} + yield pathlib.Path(tmp_path) + if old is sentinel: + del di["config"] + else: + di["config"] = old + + +@pytest.mark.parametrize("cls", ALL_CHECKPOINTS) +def test_checkpoint_is_downloadable(cls): + assert issubclass(cls, HFDownloadableMixin) + assert cls.REQUIRES_DOWNLOAD is True + assert cls.DOWNLOAD_SIZE_BYTES is not None + assert cls.REPO_ID + assert cls.GGUF_PATTERN + + +@pytest.mark.parametrize("cls", ALL_CHECKPOINTS) +def test_hf_repos_single_file_entry(cls): + assert cls.hf_repos() == [(cls.REPO_ID, "model", [cls.GGUF_PATTERN])] + + +@pytest.mark.parametrize("cls", ALL_CHECKPOINTS) +def test_metadata_flags_download(cls): + meta = cls.get_metadata() + assert meta["requires_download"] is True + assert meta["download_size_bytes"] == cls.DOWNLOAD_SIZE_BYTES + + +def test_is_downloaded_reflects_component_dir(component_root): + cls = Qwen25_05BInstruct + assert cls.is_downloaded() is False + + repo_dir = component_root / cls.__name__ / cls.REPO_ID.split("/")[-1] + repo_dir.mkdir(parents=True) + (repo_dir / "model-q8_0.gguf").write_text("weights") + assert cls.is_downloaded() is True + + +def test_new_classes_registered_old_ones_gone(): + registered = {c.__name__ for c in get_initial_components()} + for cls in ALL_CHECKPOINTS: + assert cls.__name__ in registered + for old in {"QwenModel", "SmolLMModel", "LlamaModel", "MistralModel"}: + assert old not in registered From 6108b45cbe60595db424103eb11cc954aa115522 Mon Sep 17 00:00:00 2001 From: Irozuku Date: Thu, 2 Jul 2026 00:43:03 -0400 Subject: [PATCH 30/89] fix: keep generative session unknown-model response at 400 and update GGUF-split tests Reverts the generative session-creation guard to the endpoint's existing 400 "not registered" behavior for unknown models (dropping the added 422 that broke existing tests); the 409 download gate is unchanged. Updates the session fixtures to use a non-download model after the GGUF split removed the old QwenModel component. --- .../api_v1/endpoints/generative_session.py | 7 ---- .../test_generative_session_download_gate.py | 7 ++-- tests/back/api/test_session_api.py | 36 +++++++++++-------- 3 files changed, 25 insertions(+), 25 deletions(-) diff --git a/DashAI/back/api/api_v1/endpoints/generative_session.py b/DashAI/back/api/api_v1/endpoints/generative_session.py index 982b2b7c3..92c1b105e 100644 --- a/DashAI/back/api/api_v1/endpoints/generative_session.py +++ b/DashAI/back/api/api_v1/endpoints/generative_session.py @@ -38,13 +38,6 @@ async def upload_generative_session( with session_factory() as db: try: - # Guard: unknown model name -> 422 - if params.model_name not in component_registry: - raise HTTPException( - status_code=status.HTTP_422_UNPROCESSABLE_ENTITY, - detail=f"Unknown model '{params.model_name}'", - ) - # Check if the model is registered try: model_class = component_registry[params.model_name]["class"] diff --git a/tests/back/api/test_generative_session_download_gate.py b/tests/back/api/test_generative_session_download_gate.py index b284200a2..42c77a969 100644 --- a/tests/back/api/test_generative_session_download_gate.py +++ b/tests/back/api/test_generative_session_download_gate.py @@ -60,8 +60,8 @@ def test_upload_generative_session_rejects_undownloaded_model(client): assert "download" in resp.json()["detail"].lower() -def test_upload_generative_session_unknown_model_422(client): - """POSTing a session with an unregistered model_name must return HTTP 422.""" +def test_upload_generative_session_unknown_model_400(client): + """POSTing a session with an unregistered model_name must return HTTP 400.""" resp = client.post( "/api/v1/generative-session/", json={ @@ -70,4 +70,5 @@ def test_upload_generative_session_unknown_model_422(client): **_SESSION_PAYLOAD_BASE, }, ) - assert resp.status_code == 422 + assert resp.status_code == 400 + assert "is not registered" in resp.json()["detail"] diff --git a/tests/back/api/test_session_api.py b/tests/back/api/test_session_api.py index 1572a28eb..ece8d95cc 100644 --- a/tests/back/api/test_session_api.py +++ b/tests/back/api/test_session_api.py @@ -54,17 +54,20 @@ def create_session_2(client: TestClient): @pytest.fixture(scope="module", name="response_3") def create_session_3(client: TestClient): - """Create testing session 3 using job system.""" + """Create testing session 3 using a non-download-required model.""" params = { - "model_name": "QwenModel", - "task_name": "TextToTextGenerationTask", + "model_name": "StableDiffusionV2Model", + "task_name": "TextToImageGenerationTask", "parameters": { - "model_name": "Qwen/Qwen2.5-1.5B-Instruct-GGUF", - "max_tokens": 100, - "temperature": 0.9, - "frequency_penalty": 0.1, - "context_window": 512, + "num_inference_steps": 1, + "model_name": "sd2-community/stable-diffusion-2", + "guidance_scale": 6.0, "device": "CPU", + "negative_prompt": "", + "seed": 42, + "width": 256, + "height": 256, + "num_images_per_prompt": 1, }, "name": "session_3", "description": None, @@ -80,17 +83,20 @@ def create_session_3(client: TestClient): @pytest.fixture(scope="module", name="response_4") def create_session_4(client: TestClient): - """Create testing session 4 using job system.""" + """Create testing session 4 with an invalid task (valid model).""" params = { - "model_name": "QwenModel", + "model_name": "StableDiffusionV2Model", "task_name": "SomeTask", "parameters": { - "model_name": "Qwen/Qwen2.5-1.5B-Instruct-GGUF", - "max_tokens": 100, - "temperature": 0.9, - "frequency_penalty": 0.1, - "context_window": 512, + "num_inference_steps": 1, + "model_name": "sd2-community/stable-diffusion-2", + "guidance_scale": 6.0, "device": "CPU", + "negative_prompt": "", + "seed": 42, + "width": 256, + "height": 256, + "num_images_per_prompt": 1, }, "name": "session_4", "description": None, From a59d251e40e641b08e4836ec5957191ad67581de Mon Sep 17 00:00:00 2001 From: Irozuku Date: Thu, 2 Jul 2026 00:45:11 -0400 Subject: [PATCH 31/89] feat: allow changing a generative session's model with a download gate Extend PATCH /generative-session/{id} to accept model_name, validating it is a registered generative model and rejecting undownloaded download-required models with 409. --- .../api_v1/endpoints/generative_session.py | 38 +++++++++- .../test_generative_session_download_gate.py | 74 +++++++++++++++++++ 2 files changed, 110 insertions(+), 2 deletions(-) diff --git a/DashAI/back/api/api_v1/endpoints/generative_session.py b/DashAI/back/api/api_v1/endpoints/generative_session.py index 92c1b105e..6c59fb3d6 100644 --- a/DashAI/back/api/api_v1/endpoints/generative_session.py +++ b/DashAI/back/api/api_v1/endpoints/generative_session.py @@ -321,7 +321,9 @@ async def update_generative_session( session_id: int, name: Union[str, None] = None, description: Union[str, None] = None, + model_name: Union[str, None] = None, session_factory: "sessionmaker" = Depends(lambda: di["session_factory"]), + component_registry: "ComponentRegistry" = Depends(lambda: di["component_registry"]), ): """Update the generative session associated with the provided ID. @@ -333,9 +335,15 @@ async def update_generative_session( New name for the session. description : Union[str, None], optional New description for the session. + model_name : Union[str, None], optional + New model (component name) for the session. Must be a registered + generative model; if it requires a download it must already be + downloaded. session_factory : Callable[..., ContextManager[Session]] A factory that creates a context manager that handles a SQLAlchemy session. The generated session can be used to access and query the database. + component_registry : ComponentRegistry + The DashAI component registry, used to validate the new model. Returns ------- @@ -345,8 +353,10 @@ async def update_generative_session( Raises ------ HTTPException - If the session does not exist, name is invalid, or name already exists. + If the session does not exist, the name is invalid or taken, or the new + model is unknown, not a generative model, or not yet downloaded. """ + from DashAI.back.models.base_generative_model import BaseGenerativeModel with session_factory() as db: try: session = db.get(GenerativeSession, session_id) @@ -385,7 +395,31 @@ async def update_generative_session( if description is not None: setattr(session, "description", description) - if name is not None or description is not None: + # Validate and apply a model change if provided + if model_name is not None: + try: + model_class = component_registry[model_name]["class"] + except KeyError as e: + raise HTTPException( + status_code=status.HTTP_400_BAD_REQUEST, + detail=f"Model {model_name} is not registered.", + ) from e + if not issubclass(model_class, BaseGenerativeModel): + raise HTTPException( + status_code=status.HTTP_400_BAD_REQUEST, + detail=f"Model {model_name} is not a valid generative model.", + ) + entry = component_registry[model_name] + if getattr( + entry["class"], "REQUIRES_DOWNLOAD", False + ) and not entry.get("downloaded", False): + raise HTTPException( + status_code=status.HTTP_409_CONFLICT, + detail=f"Model {model_name} must be downloaded before use.", + ) + setattr(session, "model_name", model_name) + + if name is not None or description is not None or model_name is not None: session.last_modified = datetime.now() db.commit() db.refresh(session) diff --git a/tests/back/api/test_generative_session_download_gate.py b/tests/back/api/test_generative_session_download_gate.py index 42c77a969..57de0d995 100644 --- a/tests/back/api/test_generative_session_download_gate.py +++ b/tests/back/api/test_generative_session_download_gate.py @@ -72,3 +72,77 @@ def test_upload_generative_session_unknown_model_400(client): ) assert resp.status_code == 400 assert "is not registered" in resp.json()["detail"] + + +_SD_PARAMS = { + "num_inference_steps": 1, + "model_name": "sd2-community/stable-diffusion-2", + "guidance_scale": 6.0, + "device": "CPU", + "negative_prompt": "", + "seed": 42, + "width": 256, + "height": 256, + "num_images_per_prompt": 1, +} + + +def _create_sd_session(client, name): + return client.post( + "/api/v1/generative-session/", + json={ + "model_name": "StableDiffusionV2Model", + "task_name": "TextToImageGenerationTask", + "parameters": _SD_PARAMS, + "name": name, + "description": None, + }, + ) + + +def test_change_session_model_to_undownloaded_returns_409(client): + """Switching a session to a not-downloaded model must return HTTP 409.""" + created = _create_sd_session(client, "gen-switch-409") + assert created.status_code == 201 + session_id = created.json()["id"] + + resp = client.patch( + f"/api/v1/generative-session/{session_id}", + params={"model_name": "Qwen25_15BInstruct"}, + ) + assert resp.status_code == 409 + assert "download" in resp.json()["detail"].lower() + + client.delete(f"/api/v1/generative-session/{session_id}") + + +def test_change_session_model_unknown_returns_400(client): + """Switching a session to an unregistered model must return HTTP 400.""" + created = _create_sd_session(client, "gen-switch-400") + assert created.status_code == 201 + session_id = created.json()["id"] + + resp = client.patch( + f"/api/v1/generative-session/{session_id}", + params={"model_name": "__totally_bogus_generative_model_xyz__"}, + ) + assert resp.status_code == 400 + assert "is not registered" in resp.json()["detail"] + + client.delete(f"/api/v1/generative-session/{session_id}") + + +def test_change_session_model_valid_returns_200(client): + """Switching a session to a valid (non-download) model must succeed.""" + created = _create_sd_session(client, "gen-switch-200") + assert created.status_code == 201 + session_id = created.json()["id"] + + resp = client.patch( + f"/api/v1/generative-session/{session_id}", + params={"model_name": "StableDiffusionV2Model"}, + ) + assert resp.status_code == 200 + assert resp.json()["model_name"] == "StableDiffusionV2Model" + + client.delete(f"/api/v1/generative-session/{session_id}") From adb678121158fb79128e28a0b9a7fd69dfc1bb7d Mon Sep 17 00:00:00 2001 From: Irozuku Date: Thu, 2 Jul 2026 09:04:04 -0400 Subject: [PATCH 32/89] feat: split Mixtral into per-quantization downloadable components Replace the enum-based MixtralModel with two per-quantization components (Q4_K_M ~26GB and Q2_K ~16GB) on the shared GGUFTextGenerationModel base, each downloading only its one GGUF file. --- DashAI/back/initial_components.py | 8 +- .../back/models/hugging_face/mixtral_model.py | 584 +++--------------- .../back/models/test_gguf_text_generation.py | 14 +- 3 files changed, 115 insertions(+), 491 deletions(-) diff --git a/DashAI/back/initial_components.py b/DashAI/back/initial_components.py index 54e166f7e..056769619 100644 --- a/DashAI/back/initial_components.py +++ b/DashAI/back/initial_components.py @@ -156,7 +156,10 @@ Mistral7BInstructV03, MistralNemoInstruct2407, ) -from DashAI.back.models.hugging_face.mixtral_model import MixtralModel +from DashAI.back.models.hugging_face.mixtral_model import ( + Mixtral8x7BInstructQ2K, + Mixtral8x7BInstructQ4KM, +) from DashAI.back.models.hugging_face.modernbert_transformer import ModernBertTransformer from DashAI.back.models.hugging_face.multilingual_bert_transformer import ( MultilingualBertTransformer, @@ -367,7 +370,8 @@ def get_initial_components(): MiniLMTransformer, Mistral7BInstructV03, MistralNemoInstruct2407, - MixtralModel, + Mixtral8x7BInstructQ2K, + Mixtral8x7BInstructQ4KM, MultilingualBertTransformer, MLPClassifier, MLPRegression, diff --git a/DashAI/back/models/hugging_face/mixtral_model.py b/DashAI/back/models/hugging_face/mixtral_model.py index 402eb281a..978818d87 100644 --- a/DashAI/back/models/hugging_face/mixtral_model.py +++ b/DashAI/back/models/hugging_face/mixtral_model.py @@ -1,524 +1,132 @@ -from typing import List +"""Mixtral 8x7B Instruct GGUF checkpoint subclasses for DashAI. + +Mixtral 8x7B is a single Sparse Mixture-of-Experts repo published in several +quantizations. Each quantization is exposed as its own downloadable component so +the user fetches only the one GGUF file they intend to run. +""" -from DashAI.back.core.schema_fields import ( - BaseSchema, - enum_field, - float_field, - int_field, - schema_field, -) from DashAI.back.core.utils import MultilingualString -from DashAI.back.models.text_to_text_generation_model import ( - TextToTextGenerationTaskModel, -) -from DashAI.back.models.utils import ( - LLAMA_DEVICE_ENUM, - LLAMA_DEVICE_PLACEHOLDER, - LLAMA_DEVICE_TO_IDX, +from DashAI.back.models.hugging_face.gguf_text_generation_base import ( + GGUFTextGenerationModel, + GGUFTextGenerationSchema, ) +_MIXTRAL_REPO = "mradermacher/Mixtral-8x7B-Instruct-v0.1-GGUF" -class MixtralSchema(BaseSchema): - """Schema for MixtralModel hyperparameters. - Configures the checkpoint variant (with optional GGUF filename override), - generation length, sampling temperature, frequency penalty, context window, - and target device for Mixtral Sparse-MoE models loaded via - ``llama-cpp-python``. - """ - - model_name: schema_field( - enum_field( - enum=[ - "mradermacher/Mixtral-8x7B-Instruct-v0.1-GGUF", - ] - ), - placeholder="mradermacher/Mixtral-8x7B-Instruct-v0.1-GGUF", - description=MultilingualString( - en=( - "The Mixtral Instruct checkpoint to load in GGUF format. " - "'Mixtral-8x7B-Instruct-v0.1' is a Sparse Mixture-of-Experts (SMoE) " - "model with 8 expert networks of 7B parameters each, activating 2 " - "experts per token. It achieves quality comparable to larger dense " - "models while being more efficient at inference. " - "Warning: this model requires ~26 GB of RAM for the Q4_K_M " - "quantization." - ), - es=( - "El checkpoint Mixtral Instruct a cargar en formato GGUF. " - "'Mixtral-8x7B-Instruct-v0.1' es un modelo de Mezcla Dispersa de " - "Expertos (SMoE) con 8 redes expertas de 7B parámetros cada una, " - "activando 2 expertos por token. Logra calidad comparable a modelos " - "densos más grandes siendo más eficiente en inferencia. " - "Advertencia: este modelo requiere ~26 GB de RAM para la " - "cuantización Q4_K_M." - ), - pt=( - "O checkpoint Mixtral Instruct a carregar em formato GGUF. " - "'Mixtral-8x7B-Instruct-v0.1' é um modelo de Mistura Esparsa de " - "Especialistas (SMoE) com 8 redes especialistas de 7B parâmetros cada, " - "ativando 2 especialistas por token. Alcança qualidade comparável a " - "modelos densos maiores sendo mais eficiente na inferência. " - "Aviso: este modelo requer ~26 GB de RAM para a " - "quantização Q4_K_M." - ), - de=( - "Der im GGUF-Format zu ladende Mixtral Instruct-Checkpoint. " - "'Mixtral-8x7B-Instruct-v0.1' ist ein Sparse Mixture-of-Experts " - "(SMoE)-Modell mit 8 Expertennetzwerken à 7B Parameter, das 2 " - "Experten pro Token aktiviert. Es erreicht eine mit größeren dichten " - "Modellen vergleichbare Qualität bei effizienterer Inferenz. " - "Warnung: dieses Modell benötigt ~26 GB RAM für die " - "Q4_K_M-Quantisierung." - ), - zh=( - "以 GGUF 格式加载的 Mixtral Instruct 检查点。" - "'Mixtral-8x7B-Instruct-v0.1' 是一个稀疏混合专家（SMoE）模型，" - "包含 8 个各 70 亿参数的专家网络，每个 token 激活 2 个专家。" - "推理效率高于同等质量的稠密模型。" - "警告：Q4_K_M 量化需要约 26 GB 内存。" - ), - ), - alias=MultilingualString( - en="Model name", - es="Nombre del modelo", - pt="Nome do modelo", - de="Modellname", - zh="模型名称", - ), - ) # type: ignore +class Mixtral8x7BInstructQ4KM(GGUFTextGenerationModel): + """Mixtral 8x7B Instruct GGUF checkpoint (Q4_K_M quantization). - filename: schema_field( - enum_field( - enum=[ - "Mixtral-8x7B-Instruct-v0.1.Q2_K.gguf", - "Mixtral-8x7B-Instruct-v0.1.Q3_K_M.gguf", - "Mixtral-8x7B-Instruct-v0.1.Q4_0.gguf", - "Mixtral-8x7B-Instruct-v0.1.Q4_K_M.gguf", - "Mixtral-8x7B-Instruct-v0.1.Q5_0.gguf", - "Mixtral-8x7B-Instruct-v0.1.Q5_K_M.gguf", - "Mixtral-8x7B-Instruct-v0.1.Q6_K.gguf", - "Mixtral-8x7B-Instruct-v0.1.Q8_0.gguf", - ] - ), - placeholder="Mixtral-8x7B-Instruct-v0.1.Q2_K.gguf", - description=MultilingualString( - en=( - "The specific GGUF file to load for the Mixtral model. The different " - "quantization levels (Q2_K, Q3_K_M, Q4_0, Q4_K_M, Q5_0, Q5_K_M, " - "Q6_K, Q8_0) represent various trade-offs between model size, " - "inference speed, and output quality. Q4_K_M is a popular choice " - "for balancing performance and resource requirements." - ), - es=( - "El archivo GGUF específico a cargar para el modelo Mixtral. Los " - "diferentes niveles de cuantización (Q2_K, Q3_K_M, Q4_0, Q4_K_M, " - "Q5_0, Q5_K_M, Q6_K, Q8_0) representan varios compromisos entre " - "tamaño del modelo, velocidad de inferencia y calidad de salida. " - "Q4_K_M es una opción popular para equilibrar rendimiento y " - "requisitos de recursos." - ), - pt=( - "O arquivo GGUF específico a carregar para o modelo Mixtral. Os " - "diferentes níveis de quantização (Q2_K, Q3_K_M, Q4_0, Q4_K_M, " - "Q5_0, Q5_K_M, Q6_K, Q8_0) representam vários compromissos entre " - "tamanho do modelo, velocidade de inferência e qualidade de saída. " - "Q4_K_M é uma escolha popular para equilibrar desempenho e " - "requisitos de recursos." - ), - de=( - "Die zu ladende spezifische GGUF-Datei für das Mixtral-Modell. " - "Die verschiedenen Quantisierungsstufen (Q2_K, Q3_K_M, Q4_0, Q4_K_M, " - "Q5_0, Q5_K_M, Q6_K, Q8_0) stellen verschiedene Kompromisse zwischen " - "Modellgröße, Inferenzgeschwindigkeit und Ausgabequalität dar. " - "Q4_K_M ist eine beliebte Wahl für ausgewogene Leistung und " - "Ressourcenbedarf." - ), - zh=( - "为 Mixtral 模型加载的具体 GGUF 文件。" - "不同量化级别（Q2_K、Q3_K_M、Q4_0、Q4_K_M、Q5_0、Q5_K_M、Q6_K、Q8_0）" - "在模型大小、推理速度和输出质量之间存在不同权衡。" - "Q4_K_M 是兼顾性能与资源需求的常用选择。" - ), - ), - alias=MultilingualString( - en="Filename", - es="Nombre del archivo", - pt="Nome do archivo", - de="Dateiname", - zh="文件名", - ), - ) # type: ignore + A Sparse Mixture-of-Experts model (8 experts of 7B parameters, 2 active per + token) from Mistral AI. The Q4_K_M quantization balances quality and size + and requires roughly 26 GB of RAM. Weights are stored locally after a + one-time download from HuggingFace. - max_tokens: schema_field( - int_field(ge=1), - placeholder=100, - description=MultilingualString( - en=( - "Maximum number of new tokens the model will generate per response. " - "Roughly 1 token ≈ 0.75 English words. Set to 100-200 for short " - "answers, 500-1000 for detailed explanations or code." - ), - es=( - "Número máximo de tokens nuevos que el modelo generará por respuesta. " - "Aproximadamente 1 token ≈ 0.75 palabras en español. Use 100-200 " - "para respuestas cortas, 500-1000 para explicaciones detalladas " - "o código." - ), - pt=( - "Número máximo de tokens novos que o modelo gerará por resposta. " - "Aproximadamente 1 token ≈ 0.75 palavras em português. Use 100-200 " - "para respostas curtas, 500-1000 para explicações detalhadas " - "ou código." - ), - de=( - "Maximale Anzahl neuer Token, die das Modell pro Antwort erzeugt. " - "Ungefähr 1 Token ≈ 0,75 englische Wörter. 100-200 für kurze " - "Antworten, 500-1000 für ausführliche Erklärungen oder Code." - ), - zh=( - "模型每次响应生成的最大新 token 数。" - "约 1 token ≈ 0.75 个英文单词。" - "短回答设为 100-200，详细说明或代码设为 500-1000。" - ), - ), - alias=MultilingualString( - en="Max tokens", - es="Tokens máximos", - pt="Tokens máximos", - de="Maximale neue Token", - zh="最大 token 数", - ), - ) # type: ignore - - temperature: schema_field( - float_field(ge=0.0, le=1.0), - placeholder=0.7, - description=MultilingualString( - en=( - "Sampling temperature controlling output randomness (range 0.0-1.0). " - "At 0.0 outputs are deterministic. Around 0.7 balances quality and " - "creativity." - ), - es=( - "Temperatura de muestreo que controla la aleatoriedad (rango 0.0-1.0). " - "En 0.0 las salidas son deterministas. Alrededor de 0.7 equilibra " - "calidad y creatividad." - ), - pt=( - "Temperatura de amostragem que controla a aleatoriedade " - "(intervalo 0.0-1.0). " - "Em 0.0 as saídas são determinísticas. Em torno de 0.7 equilibra " - "qualidade e criatividade." - ), - de=( - "Stichprobentemperatur zur Steuerung der Ausgabezufälligkeit (0.0-1.0)." - "Bei 0.0 sind die Ausgaben deterministisch. Um 0.7 balanciert " - "Qualität und Kreativität." - ), - zh=( - "控制输出随机性的采样温度（范围 0.0-1.0）。" - "0.0 时输出确定性最强，0.7 左右可平衡质量与创造性。" - ), - ), - alias=MultilingualString( - en="Temperature", - es="Temperatura", - pt="Temperatura", - de="Temperatur", - zh="温度", - ), - ) # type: ignore + References + ---------- + - Jiang et al. (2024) "Mixtral of Experts" https://arxiv.org/abs/2401.04088 + - https://huggingface.co/mradermacher/Mixtral-8x7B-Instruct-v0.1-GGUF + """ - frequency_penalty: schema_field( - float_field(ge=0.0, le=2.0), - placeholder=0.1, - description=MultilingualString( - en=( - "Penalizes tokens that have already appeared in the output based on " - "frequency (range 0.0-2.0). Higher values discourage repetition." - ), - es=( - "Penaliza los tokens que ya aparecieron en la salida según su " - "frecuencia (rango 0.0-2.0). Valores más altos desincentivan " - "la repetición." - ), - pt=( - "Penaliza os tokens que já apareceram na saída com base na " - "frequência (intervalo 0.0-2.0). Valores mais altos desestimulam " - "a repetição." - ), - de=( - "Bestraft Token, die bereits in der Ausgabe erschienen sind, " - "basierend auf ihrer Häufigkeit (0.0-2.0). Höhere Werte reduzieren " - "Wiederholungen." - ), - zh=( - "根据频率对已出现在输出中的 token 施加惩罚（范围 0.0-2.0）。" - "较高的值可抑制重复内容。" - ), - ), - alias=MultilingualString( - en="Frequency penalty", - es="Penalización de frecuencia", - pt="Penalização de frequência", - de="Häufigkeitsstrafe", - zh="频率惩罚", + REPO_ID = _MIXTRAL_REPO + GGUF_PATTERN = "*Q4_K_M.gguf" + DOWNLOAD_SIZE_BYTES = 26_000_000_000 + SCHEMA = GGUFTextGenerationSchema + COLOR: str = "#4a148c" + DISPLAY_NAME = MultilingualString( + en="Mixtral 8x7B Instruct (Q4_K_M)", + es="Mixtral 8x7B Instruct (Q4_K_M)", + pt="Mixtral 8x7B Instruct (Q4_K_M)", + de="Mixtral 8x7B Instruct (Q4_K_M)", + zh="Mixtral 8x7B Instruct (Q4_K_M)", + ) + DESCRIPTION = MultilingualString( + en=( + "Mixtral 8x7B Instruct is a Sparse Mixture-of-Experts model by Mistral " + "AI (8 experts of 7B parameters, 2 active per token), loaded as a Q4_K_M " + "GGUF for a balance of quality and size. It matches larger dense models " + "on many tasks. Warning: requires ~26 GB of RAM; a GPU with >= 24 GB " + "VRAM is recommended. Model available at " + "https://huggingface.co/mradermacher/Mixtral-8x7B-Instruct-v0.1-GGUF." ), - ) # type: ignore - - context_window: schema_field( - int_field(ge=1, le=32768), - placeholder=512, - description=MultilingualString( - en=( - "Total token budget for a single forward pass, including both the " - "input prompt and the generated response. Mixtral 8x7B supports " - "up to 32K tokens natively." - ), - es=( - "Presupuesto total de tokens por pasada, incluyendo prompt y " - "respuesta. Mixtral 8x7B soporta hasta 32K tokens de forma nativa." - ), - pt=( - "Orçamento total de tokens por passagem, incluindo prompt e " - "resposta. Mixtral 8x7B suporta até 32K tokens nativamente." - ), - de=( - "Gesamtes Token-Budget für einen einzelnen Vorwärtsdurchlauf, " - "einschließlich Eingabe-Prompt und generierter Antwort. " - "Mixtral 8x7B unterstützt nativ bis zu 32K Token." - ), - zh=( - "单次前向传播的总 token 预算，包含输入提示和生成响应。" - "Mixtral 8x7B 原生支持最多 32K token。" - ), + es=( + "Mixtral 8x7B Instruct es un modelo de Mezcla Dispersa de Expertos de " + "Mistral AI (8 expertos de 7B parametros, 2 activos por token), cargado " + "como GGUF Q4_K_M para equilibrar calidad y tamano. Advertencia: " + "requiere ~26 GB de RAM; se recomienda una GPU con >= 24 GB de VRAM." ), - alias=MultilingualString( - en="Context window", - es="Ventana de contexto", - pt="Janela de contexto", - de="Kontextfenster", - zh="上下文窗口", + pt=( + "Mixtral 8x7B Instruct e um modelo de Mistura Esparsa de Especialistas " + "da Mistral AI (8 especialistas de 7B parametros, 2 ativos por token), " + "carregado como GGUF Q4_K_M para equilibrar qualidade e tamanho. Aviso: " + "requer ~26 GB de RAM; recomenda-se uma GPU com >= 24 GB de VRAM." ), - ) # type: ignore - - device: schema_field( - enum_field(enum=LLAMA_DEVICE_ENUM), - placeholder=LLAMA_DEVICE_PLACEHOLDER, - description=MultilingualString( - en=( - "Hardware device for llama.cpp inference. 'CPU' runs the model " - "fully in RAM. A GPU option offloads all layers for faster inference. " - "Due to the large size of Mixtral, a GPU with at least 24 GB VRAM " - "is recommended for full GPU offloading." - ), - es=( - "Dispositivo de hardware para inferencia con llama.cpp. 'CPU' ejecuta " - "el modelo en RAM. Una opción de GPU descarga todas las capas para " - "inferencia más rápida. Debido al gran tamaño de Mixtral, " - "se recomienda " - "una GPU con al menos 24 GB de VRAM para descarga completa." - ), - pt=( - "Dispositivo de hardware para inferência com llama.cpp. 'CPU' executa " - "o modelo na RAM. Uma opção de GPU descarrega todas as camadas para " - "inferência mais rápida. Devido ao grande tamanho do Mixtral, " - "recomenda-se " - "uma GPU com pelo menos 24 GB de VRAM para descarregamento completo." - ), - de=( - "Hardware-Gerät für die llama.cpp-Inferenz. 'CPU' führt das Modell " - "im RAM aus. Eine GPU-Option lagert alle Schichten für schnellere " - "Inferenz aus. Aufgrund der Größe von Mixtral wird eine GPU mit " - "mindestens 24 GB VRAM für vollständiges GPU-Offloading empfohlen." - ), - zh=( - "llama.cpp 推理所用的硬件设备。'CPU' 在内存中运行模型。" - "GPU 选项可将所有层卸载以加快推理速度。" - "由于 Mixtral 体量较大，完整 GPU 卸载建议使用至少 24 GB 显存的 GPU。" - ), + de=( + "Mixtral 8x7B Instruct ist ein Sparse-Mixture-of-Experts-Modell von " + "Mistral AI (8 Experten mit 7B Parametern, 2 pro Token aktiv), als " + "Q4_K_M-GGUF fuer ein Gleichgewicht aus Qualitaet und Groesse geladen. " + "Warnung: benoetigt ~26 GB RAM; eine GPU mit >= 24 GB VRAM wird empfohlen." ), - alias=MultilingualString( - en="Device", - es="Dispositivo", - pt="Dispositivo", - de="Gerät", - zh="设备", + zh=( + "Mixtral 8x7B Instruct 是 Mistral AI 的稀疏混合专家模型" + "（8 个 70 亿参数专家，每 token 激活 2 个），" + "以 Q4_K_M GGUF 格式加载，兼顾质量与体积。" + "警告：需要约 26 GB 内存；建议使用显存不低于 24 GB 的 GPU。" ), - ) # type: ignore - + ) -class MixtralModel(TextToTextGenerationTaskModel): - """Mixtral Sparse Mixture-of-Experts (SMoE) model for text generation via llama.cpp. - Mixtral 8x7B is a transformer language model with 8 expert feed-forward - networks per layer; only 2 experts are activated per token, giving it the - computational cost of a 12B-parameter dense model while retaining capacity - equivalent to a 47B model. It matches or surpasses Llama 2 70B and GPT-3.5 - on most benchmarks. +class Mixtral8x7BInstructQ2K(GGUFTextGenerationModel): + """Mixtral 8x7B Instruct GGUF checkpoint (Q2_K quantization). - Models are loaded as GGUF quantized checkpoints via ``llama-cpp-python``. - The Q4_K_M quantization requires approximately 26 GB of RAM. + The smallest quantization of the Mixtral 8x7B Sparse Mixture-of-Experts + model, trading some quality for a lower memory footprint (roughly 16 GB). + Weights are stored locally after a one-time download from HuggingFace. References ---------- - - [1] Jiang et al. (2024) "Mixtral of Experts" https://arxiv.org/abs/2401.04088 - - [2] https://huggingface.co/mistralai/Mixtral-8x7B-Instruct-v0.1 + - Jiang et al. (2024) "Mixtral of Experts" https://arxiv.org/abs/2401.04088 + - https://huggingface.co/mradermacher/Mixtral-8x7B-Instruct-v0.1-GGUF """ - SCHEMA = MixtralSchema + REPO_ID = _MIXTRAL_REPO + GGUF_PATTERN = "*Q2_K.gguf" + DOWNLOAD_SIZE_BYTES = 16_000_000_000 + SCHEMA = GGUFTextGenerationSchema COLOR: str = "#4a148c" - DISPLAY_NAME: str = MultilingualString( - en="Mixtral Model", - es="Modelo Mixtral", - pt="Modelo Mixtral", - de="Mixtral-Modell", - zh="Mixtral 模型", + DISPLAY_NAME = MultilingualString( + en="Mixtral 8x7B Instruct (Q2_K)", + es="Mixtral 8x7B Instruct (Q2_K)", + pt="Mixtral 8x7B Instruct (Q2_K)", + de="Mixtral 8x7B Instruct (Q2_K)", + zh="Mixtral 8x7B Instruct (Q2_K)", ) - DESCRIPTION: str = MultilingualString( + DESCRIPTION = MultilingualString( en=( - "Mixtral 8x7B Instruct, a Sparse Mixture-of-Experts (SMoE) model by " - "Mistral AI, loaded in GGUF format for efficient CPU and GPU inference " - "via the llama.cpp library. The model uses 8 expert networks of 7B " - "parameters each, activating only 2 experts per token, achieving " - "performance comparable to larger dense models while being more efficient " - "at inference. Supports multi-turn conversation, reasoning, coding, and " - "general text generation. Warning: requires ~26 GB of RAM for the Q4_K_M " - "quantization. Model hosted at " + "Mixtral 8x7B Instruct is a Sparse Mixture-of-Experts model by Mistral " + "AI (8 experts of 7B parameters, 2 active per token). This Q2_K " + "quantization is the smallest variant (~16 GB), trading some output " + "quality for lower memory use. Model available at " "https://huggingface.co/mradermacher/Mixtral-8x7B-Instruct-v0.1-GGUF." ), es=( - "Mixtral 8x7B Instruct, un modelo de Mezcla Dispersa de Expertos (SMoE) " - "de Mistral AI, cargado en formato GGUF para inferencia eficiente en CPU " - "y GPU mediante llama.cpp. El modelo usa 8 redes expertas de 7B parámetros " - "cada una, activando solo 2 expertos por token, logrando un rendimiento " - "comparable a modelos densos más grandes " - "siendo más eficiente en inferencia. " - "Soporta conversación multi-turno, razonamiento, programación y generación " - "de texto en general. Advertencia: requiere ~26 GB de RAM para la " - "cuantización Q4_K_M. Modelo en " - "https://huggingface.co/mradermacher/Mixtral-8x7B-Instruct-v0.1-GGUF." + "Mixtral 8x7B Instruct es un modelo de Mezcla Dispersa de Expertos de " + "Mistral AI. Esta cuantizacion Q2_K es la variante mas pequena (~16 GB), " + "sacrificando algo de calidad por menor uso de memoria." ), pt=( - "Mixtral 8x7B Instruct, um modelo de Mistura Esparsa de Especialistas " - "(SMoE) da Mistral AI, carregado em formato GGUF para inferência eficiente " - "em CPU e GPU via biblioteca llama.cpp. O modelo usa 8 redes especialistas " - "de 7B parâmetros cada, ativando apenas 2 especialistas por token, " - "alcançando desempenho comparável a modelos densos maiores sendo mais " - "eficiente na inferência. Suporta conversa multi-turno, raciocínio, " - "programação e geração de texto em geral. Aviso: requer ~26 GB de RAM para " - "a quantização Q4_K_M. Modelo disponível em " - "https://huggingface.co/mradermacher/Mixtral-8x7B-Instruct-v0.1-GGUF." + "Mixtral 8x7B Instruct e um modelo de Mistura Esparsa de Especialistas " + "da Mistral AI. Esta quantizacao Q2_K e a variante menor (~16 GB), " + "trocando alguma qualidade por menor uso de memoria." ), de=( - "Mixtral 8x7B Instruct, ein Sparse Mixture-of-Experts (SMoE)-Modell von " - "Mistral AI, im GGUF-Format für effiziente CPU- und GPU-Inferenz über die " - "llama.cpp-Bibliothek geladen. Das Modell nutzt 8 Expertennetzwerke à 7B " - "Parameter und aktiviert nur 2 Experten pro Token, was mit größeren dichten" - "Modellen vergleichbare Leistung bei effizienterer Inferenz ermöglicht. " - "Unterstützt Mehrfachdialog, Schlussfolgerung, Programmierung und " - "allgemeine " - "Textgenerierung. Warnung: erfordert ~26 GB RAM für die " - "Q4_K_M-Quantisierung. " - "Modell unter " - "https://huggingface.co/mradermacher/Mixtral-8x7B-Instruct-v0.1-GGUF." + "Mixtral 8x7B Instruct ist ein Sparse-Mixture-of-Experts-Modell von " + "Mistral AI. Diese Q2_K-Quantisierung ist die kleinste Variante " + "(~16 GB) und tauscht etwas Qualitaet gegen geringeren Speicherbedarf." ), zh=( - "Mixtral 8x7B Instruct 是 Mistral AI 的稀疏混合专家（SMoE）模型，" - "以 GGUF 格式加载，通过 llama.cpp 库高效推理。" - "支持多轮对话、推理、编程和通用文本生成。" - "注意：Q4_K_M 量化需要约 26 GB 内存。" + "Mixtral 8x7B Instruct 是 Mistral AI 的稀疏混合专家模型。" + "此 Q2_K 量化是最小的变体（约 16 GB），" + "以部分输出质量换取更低的内存占用。" ), ) - - def __init__(self, **kwargs): - """Download and initialise a Mixtral 8x7B Instruct GGUF model via llama.cpp. - - The model weights are fetched from HuggingFace Hub using - ``Llama.from_pretrained`` and kept in memory for repeated calls to - ``generate``. - - Parameters - ---------- - **kwargs : dict - model_name : str, optional - HuggingFace repo ID for the GGUF checkpoint. - Defaults to - ``"mradermacher/Mixtral-8x7B-Instruct-v0.1-GGUF"``. - filename : str, optional - Specific GGUF quantization file to load (e.g. - ``"mixtral-8x7b-instruct-v0.1.Q4_K_M.gguf"``). Defaults to - the Q2_K variant. Higher quantizations use more RAM but - produce better output quality. - max_tokens : int, optional - Maximum number of new tokens to generate per call. Default 100. - temperature : float, optional - Sampling temperature in [0.0, 1.0]. Default 0.7. - frequency_penalty : float, optional - Token-frequency penalty in [0.0, 2.0]. Default 0.1. - context_window : int, optional - Total token budget (prompt + response) for a single forward - pass. Default 512. - device : str, optional - Target device from ``LLAMA_DEVICE_ENUM``. Any value whose - index is >= 0 enables full GPU offload (``n_gpu_layers=-1``); - ``"CPU"`` runs fully in RAM. - - Raises - ------ - RuntimeError - If ``llama-cpp-python`` is not installed. - """ - try: - from llama_cpp import Llama - except ImportError as e: - raise RuntimeError( - "llama-cpp-python is not installed. " - "Please install it to use this model." - ) from e - - kwargs = self.validate_and_transform(kwargs) - self.model_name = kwargs.get( - "model_name", "mradermacher/Mixtral-8x7B-Instruct-v0.1-GGUF" - ) - self.max_tokens = kwargs.pop("max_tokens", 100) - self.temperature = kwargs.pop("temperature", 0.7) - self.frequency_penalty = kwargs.pop("frequency_penalty", 0.1) - self.n_ctx = kwargs.pop("context_window", 512) - - self.filename = kwargs.get("filename", "Mixtral-8x7B-Instruct-v0.1.Q2_K.gguf") - use_gpu = LLAMA_DEVICE_TO_IDX.get(kwargs.get("device")) >= 0 - - self.model = Llama.from_pretrained( - repo_id=self.model_name, - filename=self.filename, - verbose=True, - n_ctx=self.n_ctx, - n_gpu_layers=-1 if use_gpu else 0, - main_gpu=(LLAMA_DEVICE_TO_IDX.get(kwargs.get("device")) if use_gpu else 0), - ) - - def generate(self, prompt: list[dict[str, str]]) -> List[str]: - """Generate a reply for the given chat prompt. - - Parameters - ---------- - prompt : list of dict - Conversation history in OpenAI chat format. Each dict must contain - at least ``"role"`` (``"system"``, ``"user"``, or ``"assistant"``) - and ``"content"`` (the message text). - - Returns - ------- - list of str - A single-element list containing the model's reply text, extracted - from ``choices[0]["message"]["content"]``. - """ - output = self.model.create_chat_completion( - messages=prompt, - max_tokens=self.max_tokens, - temperature=self.temperature, - frequency_penalty=self.frequency_penalty, - ) - return [output["choices"][0]["message"]["content"]] diff --git a/tests/back/models/test_gguf_text_generation.py b/tests/back/models/test_gguf_text_generation.py index 5c40da2c2..714800608 100644 --- a/tests/back/models/test_gguf_text_generation.py +++ b/tests/back/models/test_gguf_text_generation.py @@ -16,6 +16,10 @@ Mistral7BInstructV03, MistralNemoInstruct2407, ) +from DashAI.back.models.hugging_face.mixtral_model import ( + Mixtral8x7BInstructQ2K, + Mixtral8x7BInstructQ4KM, +) from DashAI.back.models.hugging_face.qwen_model import ( Qwen25_05BInstruct, Qwen25_15BInstruct, @@ -35,6 +39,8 @@ Llama32_3BInstruct, Mistral7BInstructV03, MistralNemoInstruct2407, + Mixtral8x7BInstructQ4KM, + Mixtral8x7BInstructQ2K, ] @@ -85,5 +91,11 @@ def test_new_classes_registered_old_ones_gone(): registered = {c.__name__ for c in get_initial_components()} for cls in ALL_CHECKPOINTS: assert cls.__name__ in registered - for old in {"QwenModel", "SmolLMModel", "LlamaModel", "MistralModel"}: + for old in { + "QwenModel", + "SmolLMModel", + "LlamaModel", + "MistralModel", + "MixtralModel", + }: assert old not in registered From e90d4ca9fc8b817ad046f20396b85da619c9ec1e Mon Sep 17 00:00:00 2001 From: Irozuku Date: Thu, 2 Jul 2026 09:12:50 -0400 Subject: [PATCH 33/89] feat: add session-level model switcher for generative sessions Add a model switcher in the generative chat header that lists same-task models (undownloaded ones disabled), persists the choice via the gated PATCH endpoint, and reloads the params panel. Extend updateGenerativeSession with model_name and refetch the create-session picker when a model download completes. --- DashAI/front/src/api/session.ts | 2 +- .../generative/CreateSessionCenter.jsx | 2 + .../generative/CreateSessionContext.jsx | 26 ++--- .../components/generative/GenerativeChat.jsx | 15 ++- .../components/generative/ModelSwitcher.jsx | 104 ++++++++++++++++++ .../src/components/generative/ParamsBar.jsx | 3 +- .../src/utils/i18n/locales/de/generative.json | 7 +- .../src/utils/i18n/locales/en/generative.json | 7 +- .../src/utils/i18n/locales/es/generative.json | 7 +- .../src/utils/i18n/locales/pt/generative.json | 7 +- .../src/utils/i18n/locales/zh/generative.json | 7 +- 11 files changed, 161 insertions(+), 26 deletions(-) create mode 100644 DashAI/front/src/components/generative/ModelSwitcher.jsx diff --git a/DashAI/front/src/api/session.ts b/DashAI/front/src/api/session.ts index a08e4d534..4f1ed1602 100644 --- a/DashAI/front/src/api/session.ts +++ b/DashAI/front/src/api/session.ts @@ -31,7 +31,7 @@ export const updateGenerativeSession = async ({ formData, }: { id: string; - formData: { name?: string; task_name?: string }; + formData: { name?: string; task_name?: string; model_name?: string }; }): Promise => { const response = await api.patch(`/v1/generative-session/${id}`, null, { params: formData, diff --git a/DashAI/front/src/components/generative/CreateSessionCenter.jsx b/DashAI/front/src/components/generative/CreateSessionCenter.jsx index 2480df2b4..f15ce47ea 100644 --- a/DashAI/front/src/components/generative/CreateSessionCenter.jsx +++ b/DashAI/front/src/components/generative/CreateSessionCenter.jsx @@ -29,6 +29,7 @@ export default function CreateSessionCenter() { step, models, loadingModels, + refetchModels, selectedModel, handleSelectModel, formik, @@ -126,6 +127,7 @@ export default function CreateSessionCenter() { components={models} selected={selectedModel} onSelect={handleSelectModelWithTour} + onDownloadChange={() => refetchModels()} categoryKey="task_display_name" searchPlaceholder={t("generative:label.searchModels")} tourDataFor={tourContext?.run ? "model-card-qwen" : null} diff --git a/DashAI/front/src/components/generative/CreateSessionContext.jsx b/DashAI/front/src/components/generative/CreateSessionContext.jsx index 43a639088..34bdd62ce 100644 --- a/DashAI/front/src/components/generative/CreateSessionContext.jsx +++ b/DashAI/front/src/components/generative/CreateSessionContext.jsx @@ -42,14 +42,14 @@ export function CreateSessionProvider({ children }) { const [selectedModel, setSelectedModel] = useState(null); const [submitting, setSubmitting] = useState(false); - // Load all generative models grouped by their compatible task. - // Re-fetches when language changes so display_name/description are translated. - useEffect(() => { - if (!tasks || tasks.length === 0) return; - let cancelled = false; + // Load all generative models grouped by their compatible task. Exposed as + // refetchModels so callers (e.g. an inline download control) can refresh the + // list when a model's downloaded state changes. Re-fetches when language + // changes so display_name/description are translated. + const loadModels = useCallback(() => { + if (!tasks || tasks.length === 0) return Promise.resolve(); setLoadingModels(true); - - Promise.all( + return Promise.all( tasks.map((task) => getRelatedComponents(task.name).then((components) => components.map((c) => ({ @@ -61,7 +61,6 @@ export function CreateSessionProvider({ children }) { ), ) .then((perTaskLists) => { - if (cancelled) return; // Deduplicate by model name (a model may appear under several tasks). const seen = new Set(); const flat = []; @@ -79,14 +78,14 @@ export function CreateSessionProvider({ children }) { }); }) .finally(() => { - if (!cancelled) setLoadingModels(false); + setLoadingModels(false); }); - - return () => { - cancelled = true; - }; }, [tasks, enqueueSnackbar, t]); + useEffect(() => { + loadModels(); + }, [loadModels]); + const processedProperties = useMemo( () => selectedModel?.schema?.properties @@ -219,6 +218,7 @@ export function CreateSessionProvider({ children }) { step, models, loadingModels, + refetchModels: loadModels, selectedModel, handleSelectModel, formik, diff --git a/DashAI/front/src/components/generative/GenerativeChat.jsx b/DashAI/front/src/components/generative/GenerativeChat.jsx index 5f12cd2d3..6609cb543 100644 --- a/DashAI/front/src/components/generative/GenerativeChat.jsx +++ b/DashAI/front/src/components/generative/GenerativeChat.jsx @@ -15,6 +15,7 @@ import { enqueueGenerativeProcessJob } from "../../api/job"; import { startJobQueue } from "../../api/job"; import { getHistoryBySessionId, getSessionById } from "../../api/session"; import InfoSessionModal from "./InfoSessionModal"; +import ModelSwitcher from "./ModelSwitcher"; import { useSnackbar } from "notistack"; import { MediaInput } from "./MediaInput"; import { Trans, useTranslation } from "react-i18next"; @@ -29,6 +30,8 @@ export default function GenerativeChat() { selectedTaskName: taskName, tasks, paramsVersion, + setParamsVersion, + fetchSessions, } = useGenerative(); const inputsCardinality = useMemo(() => { @@ -254,7 +257,17 @@ export default function GenerativeChat() { {sessionInfo?.description ? ":" : null} {sessionInfo?.description} - + + { + getSessionInfo(); + fetchSessions(); + setParamsVersion((v) => v + 1); + }} + /> setSessionInfoVisible(true)}> { + if (!taskName) return; + getRelatedComponents(taskName) + .then((components) => setModels(components || [])) + .catch(() => setModels([])); + }, [taskName]); + + const handleChange = async (event) => { + const newModel = event.target.value; + if (!newModel || newModel === currentModelName) return; + setSaving(true); + try { + await updateGenerativeSession({ + id: sessionId, + formData: { model_name: newModel }, + }); + if (onChanged) onChanged(newModel); + } catch (error) { + const status = error?.response?.status; + enqueueSnackbar( + status === 409 + ? t("common:componentDownload.mustDownload") + : t("generative:error.modelSwitchFailed"), + { variant: "error" }, + ); + } finally { + setSaving(false); + } + }; + + // Ensure the current model is always a selectable value even if the list + // has not loaded yet (avoids an out-of-range MUI Select warning). + const hasCurrent = models.some((m) => m.name === currentModelName); + const options = hasCurrent + ? models + : [{ name: currentModelName, display_name: currentModelName }, ...models]; + + if (!currentModelName) return null; + + return ( + + + {t("generative:label.sessionModel")} + + + + ); +} + +ModelSwitcher.propTypes = { + sessionId: PropTypes.oneOfType([PropTypes.string, PropTypes.number]), + taskName: PropTypes.string, + currentModelName: PropTypes.string, + onChanged: PropTypes.func, +}; diff --git a/DashAI/front/src/components/generative/ParamsBar.jsx b/DashAI/front/src/components/generative/ParamsBar.jsx index a6aa35ca7..3b85c76dd 100644 --- a/DashAI/front/src/components/generative/ParamsBar.jsx +++ b/DashAI/front/src/components/generative/ParamsBar.jsx @@ -21,6 +21,7 @@ export default function ParamsBar({ onToggle }) { const { selectedSessionId, selectedTaskName: taskName, + paramsVersion, setParamsVersion, } = useGenerative(); const [parameters, setParameters] = useState({}); @@ -60,7 +61,7 @@ export default function ParamsBar({ onToggle }) { } }); }); - }, [selectedSessionId, t]); + }, [selectedSessionId, t, paramsVersion]); useEffect(() => { if (selectedModel?.schema?.properties) { diff --git a/DashAI/front/src/utils/i18n/locales/de/generative.json b/DashAI/front/src/utils/i18n/locales/de/generative.json index 2446ff33d..701767f7c 100644 --- a/DashAI/front/src/utils/i18n/locales/de/generative.json +++ b/DashAI/front/src/utils/i18n/locales/de/generative.json @@ -17,7 +17,8 @@ "nameRequired": "Name ist erforderlich", "processError": "Der Prozess ist fehlgeschlagen. Wird gelöscht... {{error}}", "sessionNameEmpty": "Sitzungsname darf nicht leer sein", - "sessionNameExists": "Eine Sitzung mit diesem Namen existiert bereits" + "sessionNameExists": "Eine Sitzung mit diesem Namen existiert bereits", + "modelSwitchFailed": "Sitzungsmodell konnte nicht geändert werden" }, "label": { "confirmDeleteSession": "Sind Sie sicher, dass Sie die Sitzung \"{{name}}\" löschen möchten? Diese Aktion kann nicht rückgängig gemacht werden.", @@ -57,7 +58,9 @@ "attachMedia": "Medien anhängen", "attachMediaToContinue": "Medien anhängen, um fortzufahren", "noInputAvailable": "Keine Eingabe für diese Aufgabe verfügbar", - "selectModelCrumb": "Modell auswählen" + "selectModelCrumb": "Modell auswählen", + "sessionModel": "Modell", + "downloadRequired": "Download erforderlich" }, "message": { "sessionCreatedSuccess": "Sitzung erfolgreich erstellt.", diff --git a/DashAI/front/src/utils/i18n/locales/en/generative.json b/DashAI/front/src/utils/i18n/locales/en/generative.json index 654597724..9bd5ddb0c 100644 --- a/DashAI/front/src/utils/i18n/locales/en/generative.json +++ b/DashAI/front/src/utils/i18n/locales/en/generative.json @@ -17,7 +17,8 @@ "nameRequired": "Name is required", "processError": "The process has failed. Deleting it... {{error}}", "sessionNameEmpty": "Session name cannot be empty", - "sessionNameExists": "A session with this name already exists" + "sessionNameExists": "A session with this name already exists", + "modelSwitchFailed": "Failed to change the session model" }, "label": { "confirmDeleteSession": "Are you sure you want to delete the session \"{{name}}\"? This action cannot be undone.", @@ -57,7 +58,9 @@ "attachMedia": "Attach media", "attachMediaToContinue": "Attach media to continue", "noInputAvailable": "No input available for this task", - "selectModelCrumb": "Select Model" + "selectModelCrumb": "Select Model", + "sessionModel": "Model", + "downloadRequired": "download required" }, "message": { "sessionCreatedSuccess": "Session successfully created.", diff --git a/DashAI/front/src/utils/i18n/locales/es/generative.json b/DashAI/front/src/utils/i18n/locales/es/generative.json index 97f059735..08f4bdefe 100644 --- a/DashAI/front/src/utils/i18n/locales/es/generative.json +++ b/DashAI/front/src/utils/i18n/locales/es/generative.json @@ -17,7 +17,8 @@ "nameRequired": "Se requiere un nombre", "processError": "El proceso ha fallado. Eliminándolo... {{error}}", "sessionNameEmpty": "El nombre de la sesión no puede estar vacío", - "sessionNameExists": "Ya existe una sesión con este nombre" + "sessionNameExists": "Ya existe una sesión con este nombre", + "modelSwitchFailed": "No se pudo cambiar el modelo de la sesión" }, "label": { "confirmDeleteSession": "¿Seguro que quieres eliminar la sesión \"{{name}}\"? Esta acción no se puede deshacer.", @@ -57,7 +58,9 @@ "attachMedia": "Adjuntar medio", "attachMediaToContinue": "Adjunta medios para continuar", "noInputAvailable": "No hay entrada disponible para esta tarea", - "selectModelCrumb": "Select Model" + "selectModelCrumb": "Select Model", + "sessionModel": "Modelo", + "downloadRequired": "descarga requerida" }, "message": { "sessionCreatedSuccess": "Sesión creada exitosamente.", diff --git a/DashAI/front/src/utils/i18n/locales/pt/generative.json b/DashAI/front/src/utils/i18n/locales/pt/generative.json index c2b513a51..69f867d4b 100644 --- a/DashAI/front/src/utils/i18n/locales/pt/generative.json +++ b/DashAI/front/src/utils/i18n/locales/pt/generative.json @@ -17,7 +17,8 @@ "nameRequired": "É necessário um nome", "processError": "O processo falhou. Excluindo... {{error}}", "sessionNameEmpty": "O nome da sessão não pode estar vazio", - "sessionNameExists": "Já existe uma sessão com este nome" + "sessionNameExists": "Já existe uma sessão com este nome", + "modelSwitchFailed": "Falha ao alterar o modelo da sessão" }, "label": { "confirmDeleteSession": "Tem certeza de que deseja excluir a sessão \"{{name}}\"? Esta ação não pode ser desfeita.", @@ -57,7 +58,9 @@ "attachMedia": "Anexar mídia", "attachMediaToContinue": "Anexe uma mídia para continuar", "noInputAvailable": "Não há entrada disponível para esta tarefa", - "selectModelCrumb": "Selecionar Modelo" + "selectModelCrumb": "Selecionar Modelo", + "sessionModel": "Modelo", + "downloadRequired": "download necessário" }, "message": { "sessionCreatedSuccess": "Sessão criada com sucesso.", diff --git a/DashAI/front/src/utils/i18n/locales/zh/generative.json b/DashAI/front/src/utils/i18n/locales/zh/generative.json index bd86dc015..6f3b9870c 100644 --- a/DashAI/front/src/utils/i18n/locales/zh/generative.json +++ b/DashAI/front/src/utils/i18n/locales/zh/generative.json @@ -17,7 +17,8 @@ "nameRequired": "名称为必填项", "processError": "处理失败，正在删除...{{error}}", "sessionNameEmpty": "会话名称不能为空", - "sessionNameExists": "已存在同名会话" + "sessionNameExists": "已存在同名会话", + "modelSwitchFailed": "更改会话模型失败" }, "label": { "confirmDeleteSession": "确定要删除会话 \"{{name}}\" 吗？此操作无法撤销。", @@ -57,7 +58,9 @@ "attachMedia": "附加媒体", "attachMediaToContinue": "附加媒体以继续", "noInputAvailable": "此任务无可用输入", - "selectModelCrumb": "选择模型" + "selectModelCrumb": "选择模型", + "sessionModel": "模型", + "downloadRequired": "需下载" }, "message": { "sessionCreatedSuccess": "会话创建成功。", From ddc682d1745368207a3859d336f82fb10feaf0e4 Mon Sep 17 00:00:00 2001 From: Irozuku Date: Thu, 2 Jul 2026 09:36:59 -0400 Subject: [PATCH 34/89] fix: reconcile component download status against filesystem on read --- .../back/api/api_v1/endpoints/components.py | 8 ++++++ tests/back/api/test_components_api.py | 26 +++++++++++++++++++ 2 files changed, 34 insertions(+) diff --git a/DashAI/back/api/api_v1/endpoints/components.py b/DashAI/back/api/api_v1/endpoints/components.py index 873d2f6df..64b59ec84 100644 --- a/DashAI/back/api/api_v1/endpoints/components.py +++ b/DashAI/back/api/api_v1/endpoints/components.py @@ -230,6 +230,14 @@ async def get_components( components_with_related_type, ) + # Reconcile the download state of download required components against the + # filesystem before returning. Downloads happen in the worker process, so + # the in memory registry flag can be stale; a fresh check (a cheap folder + # stat per downloadable component) keeps the list truthful. + for comp_name, component_dict in selected_components.items(): + if getattr(component_dict.get("class"), "REQUIRES_DOWNLOAD", False): + component_registry.refresh_download_status(comp_name) + return [ _filter_by_language(_delete_class(component_dict), accept_language) for component_dict in selected_components.values() diff --git a/tests/back/api/test_components_api.py b/tests/back/api/test_components_api.py index 26cd14699..6f598d304 100644 --- a/tests/back/api/test_components_api.py +++ b/tests/back/api/test_components_api.py @@ -164,6 +164,7 @@ def test_get_component_by_id(client: TestClient): "description": "Task 1.", "display_name": "Test Task 1", "color": "#795548", + "downloaded": True, } response = client.get("/api/v1/component/TestTask2/") @@ -182,6 +183,7 @@ def test_get_component_by_id(client: TestClient): "description": "Task 2.", "display_name": None, "color": None, + "downloaded": True, } response = client.get("/api/v1/component/TestDataloader1/") @@ -199,6 +201,7 @@ def test_get_component_by_id(client: TestClient): "description": None, "display_name": None, "color": None, + "downloaded": True, } @@ -292,6 +295,7 @@ def test_get_components_select_only_tasks(client: TestClient): "description": "Task 1.", "display_name": "Test Task 1", "color": "#795548", + "downloaded": True, }, { "name": "TestTask2", @@ -307,6 +311,7 @@ def test_get_components_select_only_tasks(client: TestClient): "description": "Task 2.", "display_name": None, "color": None, + "downloaded": True, }, ] @@ -330,6 +335,7 @@ def test_get_components_select_only_dataloaders(client: TestClient): "description": None, "display_name": None, "color": None, + "downloaded": True, }, { "name": "TestDataloader2", @@ -344,6 +350,7 @@ def test_get_components_select_only_dataloaders(client: TestClient): "description": None, "display_name": None, "color": None, + "downloaded": True, }, { "name": "TestDataloader3", @@ -358,6 +365,7 @@ def test_get_components_select_only_dataloaders(client: TestClient): "description": None, "display_name": None, "color": None, + "downloaded": True, }, ] @@ -438,6 +446,7 @@ def test_get_components_ignore_models(client: TestClient): "description": "Task 1.", "display_name": "Test Task 1", "color": "#795548", + "downloaded": True, }, { "name": "TestTask2", @@ -453,6 +462,7 @@ def test_get_components_ignore_models(client: TestClient): "description": "Task 2.", "display_name": None, "color": None, + "downloaded": True, }, { "name": "TestDataloader1", @@ -467,6 +477,7 @@ def test_get_components_ignore_models(client: TestClient): "description": None, "display_name": None, "color": None, + "downloaded": True, }, { "name": "TestDataloader2", @@ -481,6 +492,7 @@ def test_get_components_ignore_models(client: TestClient): "description": None, "display_name": None, "color": None, + "downloaded": True, }, { "name": "TestDataloader3", @@ -495,6 +507,7 @@ def test_get_components_ignore_models(client: TestClient): "description": None, "display_name": None, "color": None, + "downloaded": True, }, ] @@ -517,6 +530,7 @@ def test_get_components_ignore_tasks_and_models(client: TestClient): "description": None, "display_name": None, "color": None, + "downloaded": True, }, { "name": "TestDataloader2", @@ -531,6 +545,7 @@ def test_get_components_ignore_tasks_and_models(client: TestClient): "description": None, "display_name": None, "color": None, + "downloaded": True, }, { "name": "TestDataloader3", @@ -545,6 +560,7 @@ def test_get_components_ignore_tasks_and_models(client: TestClient): "description": None, "display_name": None, "color": None, + "downloaded": True, }, ] @@ -608,6 +624,7 @@ def test_get_components_related_inverse_relation(client: TestClient): "description": "Task 1.", "display_name": "Test Task 1", "color": "#795548", + "downloaded": True, } ] @@ -653,6 +670,7 @@ def test_get_components_dataloader_component_parent(client: TestClient): "description": None, "display_name": None, "color": None, + "downloaded": True, }, { "name": "TestDataloader2", @@ -667,6 +685,7 @@ def test_get_components_dataloader_component_parent(client: TestClient): "description": None, "display_name": None, "color": None, + "downloaded": True, }, ] @@ -706,6 +725,7 @@ def test_get_components_by_type_and_task(client: TestClient): "description": None, "display_name": None, "color": None, + "downloaded": True, }, { "name": "TestDataloader2", @@ -720,6 +740,7 @@ def test_get_components_by_type_and_task(client: TestClient): "description": None, "display_name": None, "color": None, + "downloaded": True, }, ] @@ -758,6 +779,7 @@ def test_get_components_select_and_ignore_by_type(client: TestClient): "description": None, "display_name": None, "color": None, + "downloaded": True, }, { "name": "TestDataloader2", @@ -772,6 +794,7 @@ def test_get_components_select_and_ignore_by_type(client: TestClient): "description": None, "display_name": None, "color": None, + "downloaded": True, }, { "name": "TestDataloader3", @@ -786,6 +809,7 @@ def test_get_components_select_and_ignore_by_type(client: TestClient): "description": None, "display_name": None, "color": None, + "downloaded": True, }, ] @@ -811,6 +835,7 @@ def test_get_components_select_type_and_parent(client: TestClient): "description": None, "display_name": None, "color": None, + "downloaded": True, }, { "name": "TestDataloader2", @@ -825,5 +850,6 @@ def test_get_components_select_type_and_parent(client: TestClient): "description": None, "display_name": None, "color": None, + "downloaded": True, }, ] From 70b0e618eeddcab1b01fb0e9ae691e496baa1d26 Mon Sep 17 00:00:00 2001 From: Irozuku Date: Thu, 2 Jul 2026 09:44:27 -0400 Subject: [PATCH 35/89] fix: keep download controls mounted and add inline delete for models --- .../components/custom/ComponentSelector.jsx | 10 ++- .../src/components/models/ModelsRightBar.jsx | 70 +++++++++---------- .../models/model/ComponentDownloadControl.jsx | 46 +++++++++++- .../model/ComponentDownloadControl.test.jsx | 37 +++++++++- .../components/models/model/ModelListItem.jsx | 14 ++++ 5 files changed, 130 insertions(+), 47 deletions(-) diff --git a/DashAI/front/src/components/custom/ComponentSelector.jsx b/DashAI/front/src/components/custom/ComponentSelector.jsx index 97be244b8..997b9790d 100644 --- a/DashAI/front/src/components/custom/ComponentSelector.jsx +++ b/DashAI/front/src/components/custom/ComponentSelector.jsx @@ -109,8 +109,8 @@ function ComponentSelector({ const renderCard = (component) => { const isSelected = selected?.name === component.name; const icon = getIcon?.(component); - const needsDownload = - Boolean(component.metadata?.requires_download) && !component.downloaded; + const requiresDownload = Boolean(component.metadata?.requires_download); + const needsDownload = requiresDownload && !component.downloaded; const isCsvComponent = tourDataFor && (tourDataMatchFn @@ -181,13 +181,11 @@ function ComponentSelector({ /> )} - {needsDownload && ( + {requiresDownload && ( e.stopPropagation()}> { - if (isDownloaded) onDownloadChange?.(component); - }} + onStatusChange={() => onDownloadChange?.(component)} /> )} diff --git a/DashAI/front/src/components/models/ModelsRightBar.jsx b/DashAI/front/src/components/models/ModelsRightBar.jsx index b4864cd07..85be2c054 100644 --- a/DashAI/front/src/components/models/ModelsRightBar.jsx +++ b/DashAI/front/src/components/models/ModelsRightBar.jsx @@ -243,46 +243,40 @@ export default function ModelsRightBar({ onToggle }) { ) : ( {filteredModels.map((model, index) => { - const needsDownload = - Boolean(model.metadata?.requires_download) && - !model.downloaded; + const requiresDownload = Boolean( + model.metadata?.requires_download, + ); + const needsDownload = requiresDownload && !model.downloaded; return ( - - handleModelClick(model) - } - data-tour={index === 0 ? "first-model" : undefined} - /> - {needsDownload && ( - { - if (isDownloaded) fetchModels(); - }} - /> - )} - + model={ + needsDownload + ? { + ...model, + tooltip: t( + "common:componentDownload.mustDownload", + ), + } + : model + } + disabled={needsDownload} + onClick={ + needsDownload + ? undefined + : () => handleModelClick(model) + } + data-tour={index === 0 ? "first-model" : undefined} + action={ + requiresDownload ? ( + fetchModels()} + /> + ) : null + } + /> ); })} diff --git a/DashAI/front/src/components/models/model/ComponentDownloadControl.jsx b/DashAI/front/src/components/models/model/ComponentDownloadControl.jsx index bcc6e207c..d751d43dc 100644 --- a/DashAI/front/src/components/models/model/ComponentDownloadControl.jsx +++ b/DashAI/front/src/components/models/model/ComponentDownloadControl.jsx @@ -1,5 +1,13 @@ import React, { useState, useEffect, useRef } from "react"; -import { Box, Button, LinearProgress, Typography } from "@mui/material"; +import { + Box, + Button, + CircularProgress, + IconButton, + LinearProgress, + Tooltip, + Typography, +} from "@mui/material"; import DownloadIcon from "@mui/icons-material/Download"; import DeleteIcon from "@mui/icons-material/Delete"; import { useTranslation } from "react-i18next"; @@ -18,7 +26,11 @@ const formatSize = (bytes) => { return `${Math.round(mb)} MB`; }; -const ComponentDownloadControl = ({ component, onStatusChange }) => { +const ComponentDownloadControl = ({ + component, + onStatusChange, + compact = false, +}) => { const { t } = useTranslation(["common"]); const { enqueueSnackbar } = useSnackbar(); const meta = component.metadata || {}; @@ -86,6 +98,36 @@ const ComponentDownloadControl = ({ component, onStatusChange }) => { } }; + if (compact) { + if (downloading) { + return ( + + + + ); + } + if (downloaded) { + return ( + + + + + + ); + } + return ( + + + + + + ); + } + if (downloading) { return ( diff --git a/DashAI/front/src/components/models/model/ComponentDownloadControl.test.jsx b/DashAI/front/src/components/models/model/ComponentDownloadControl.test.jsx index 6c90798e3..9bd72583a 100644 --- a/DashAI/front/src/components/models/model/ComponentDownloadControl.test.jsx +++ b/DashAI/front/src/components/models/model/ComponentDownloadControl.test.jsx @@ -16,7 +16,10 @@ jest.mock("../../../utils/jobPoller", () => ({ })); import ComponentDownloadControl from "./ComponentDownloadControl"; -import { downloadComponent } from "../../../api/component"; +import { + downloadComponent, + deleteComponentDownload, +} from "../../../api/component"; const component = { name: "OpusMtEnRoaTransformer", @@ -38,4 +41,36 @@ describe("ComponentDownloadControl", () => { expect(downloadComponent).toHaveBeenCalledWith("OpusMtEnRoaTransformer"), ); }); + + it("compact mode triggers download from an icon button", async () => { + renderWithProviders( + {}} + />, + ); + const button = await screen.findByRole("button", { name: /download/i }); + fireEvent.click(button); + await waitFor(() => + expect(downloadComponent).toHaveBeenCalledWith("OpusMtEnRoaTransformer"), + ); + }); + + it("shows a delete control for a downloaded component and deletes it", async () => { + renderWithProviders( + {}} + />, + ); + const button = await screen.findByRole("button", { name: /delete/i }); + fireEvent.click(button); + await waitFor(() => + expect(deleteComponentDownload).toHaveBeenCalledWith( + "OpusMtEnRoaTransformer", + ), + ); + }); }); diff --git a/DashAI/front/src/components/models/model/ModelListItem.jsx b/DashAI/front/src/components/models/model/ModelListItem.jsx index c846e9baa..dd7a07d29 100644 --- a/DashAI/front/src/components/models/model/ModelListItem.jsx +++ b/DashAI/front/src/components/models/model/ModelListItem.jsx @@ -9,6 +9,7 @@ export default function ModelListItem({ model, disabled = false, onClick, + action = null, ...props }) { const theme = useTheme(); @@ -152,6 +153,19 @@ export default function ModelListItem({ {model.display_name || model.name} + + {/* Trailing action (e.g. download/delete control) */} + {action && ( + e.stopPropagation()} + onMouseDown={(e) => e.stopPropagation()} + onDragStart={(e) => e.stopPropagation()} + sx={{ flexShrink: 0, display: "flex", alignItems: "center" }} + > + {action} + + )} {!disabled && ( From 823bf0b44051838201fca6adf70b073168f5286a Mon Sep 17 00:00:00 2001 From: Irozuku Date: Thu, 2 Jul 2026 09:49:47 -0400 Subject: [PATCH 36/89] feat: notify on successful component download deletion --- .../src/components/models/model/ComponentDownloadControl.jsx | 3 +++ DashAI/front/src/utils/i18n/locales/de/common.json | 1 + DashAI/front/src/utils/i18n/locales/en/common.json | 1 + DashAI/front/src/utils/i18n/locales/es/common.json | 1 + DashAI/front/src/utils/i18n/locales/pt/common.json | 1 + DashAI/front/src/utils/i18n/locales/zh/common.json | 1 + 6 files changed, 8 insertions(+) diff --git a/DashAI/front/src/components/models/model/ComponentDownloadControl.jsx b/DashAI/front/src/components/models/model/ComponentDownloadControl.jsx index d751d43dc..4ef233238 100644 --- a/DashAI/front/src/components/models/model/ComponentDownloadControl.jsx +++ b/DashAI/front/src/components/models/model/ComponentDownloadControl.jsx @@ -91,6 +91,9 @@ const ComponentDownloadControl = ({ try { await deleteComponentDownload(component.name); finish(false); + enqueueSnackbar(t("common:componentDownload.deleted"), { + variant: "success", + }); } catch { enqueueSnackbar(t("common:componentDownload.failed"), { variant: "error", diff --git a/DashAI/front/src/utils/i18n/locales/de/common.json b/DashAI/front/src/utils/i18n/locales/de/common.json index 2f777cdbb..07e9df0e7 100644 --- a/DashAI/front/src/utils/i18n/locales/de/common.json +++ b/DashAI/front/src/utils/i18n/locales/de/common.json @@ -165,6 +165,7 @@ "delete": "Download loeschen", "downloading": "Wird heruntergeladen...", "done": "Komponente heruntergeladen", + "deleted": "Download geloescht", "failed": "Download der Komponente fehlgeschlagen", "mustDownload": "Dieses Modell muss vor der Nutzung heruntergeladen werden" }, diff --git a/DashAI/front/src/utils/i18n/locales/en/common.json b/DashAI/front/src/utils/i18n/locales/en/common.json index 0220e806f..478d80af9 100644 --- a/DashAI/front/src/utils/i18n/locales/en/common.json +++ b/DashAI/front/src/utils/i18n/locales/en/common.json @@ -165,6 +165,7 @@ "delete": "Delete download", "downloading": "Downloading...", "done": "Component downloaded", + "deleted": "Download deleted", "failed": "Component download failed", "mustDownload": "This model must be downloaded before use" }, diff --git a/DashAI/front/src/utils/i18n/locales/es/common.json b/DashAI/front/src/utils/i18n/locales/es/common.json index dbe4bbf16..f22b31b60 100644 --- a/DashAI/front/src/utils/i18n/locales/es/common.json +++ b/DashAI/front/src/utils/i18n/locales/es/common.json @@ -165,6 +165,7 @@ "delete": "Eliminar descarga", "downloading": "Descargando...", "done": "Componente descargado", + "deleted": "Descarga eliminada", "failed": "La descarga del componente ha fallado", "mustDownload": "Este modelo debe descargarse antes de usarlo" }, diff --git a/DashAI/front/src/utils/i18n/locales/pt/common.json b/DashAI/front/src/utils/i18n/locales/pt/common.json index e52110881..583d45dec 100644 --- a/DashAI/front/src/utils/i18n/locales/pt/common.json +++ b/DashAI/front/src/utils/i18n/locales/pt/common.json @@ -165,6 +165,7 @@ "delete": "Remover download", "downloading": "Baixando...", "done": "Componente baixado", + "deleted": "Download removido", "failed": "Falha ao baixar o componente", "mustDownload": "Este modelo precisa ser baixado antes de usar" }, diff --git a/DashAI/front/src/utils/i18n/locales/zh/common.json b/DashAI/front/src/utils/i18n/locales/zh/common.json index d0623fb42..02985d5fb 100644 --- a/DashAI/front/src/utils/i18n/locales/zh/common.json +++ b/DashAI/front/src/utils/i18n/locales/zh/common.json @@ -165,6 +165,7 @@ "delete": "删除下载", "downloading": "下载中...", "done": "组件已下载", + "deleted": "下载已删除", "failed": "组件下载失败", "mustDownload": "使用前必须先下载此模型" }, From e89a65bfc326b1327ffb641a188694acf1d32450 Mon Sep 17 00:00:00 2001 From: Irozuku Date: Thu, 2 Jul 2026 09:55:28 -0400 Subject: [PATCH 37/89] refactor: remove session info button from GenerativeChat component --- .../front/src/components/generative/GenerativeChat.jsx | 10 ---------- 1 file changed, 10 deletions(-) diff --git a/DashAI/front/src/components/generative/GenerativeChat.jsx b/DashAI/front/src/components/generative/GenerativeChat.jsx index 6609cb543..798c0bdb2 100644 --- a/DashAI/front/src/components/generative/GenerativeChat.jsx +++ b/DashAI/front/src/components/generative/GenerativeChat.jsx @@ -268,16 +268,6 @@ export default function GenerativeChat() { setParamsVersion((v) => v + 1); }} /> - setSessionInfoVisible(true)}> - - From b0388c9a3d3504bb9aad131a0d33c7ee1afb5112 Mon Sep 17 00:00:00 2001 From: Irozuku Date: Thu, 2 Jul 2026 10:30:53 -0400 Subject: [PATCH 38/89] feat: gate generative sessions on model download and log model switches --- ...1b0_add_model_name_to_parameter_history.py | 28 ++++ .../api_v1/endpoints/generative_session.py | 86 +++++++--- DashAI/back/dependencies/database/models.py | 4 + .../generative/CreateSessionContext.jsx | 23 ++- .../components/generative/GenerativeChat.jsx | 147 +++++++++++++++--- .../components/generative/ModelSwitcher.jsx | 12 +- .../src/utils/i18n/locales/de/generative.json | 2 + .../src/utils/i18n/locales/en/generative.json | 2 + .../src/utils/i18n/locales/es/generative.json | 2 + .../src/utils/i18n/locales/pt/generative.json | 2 + .../src/utils/i18n/locales/zh/generative.json | 2 + .../test_generative_session_download_gate.py | 49 +++++- 12 files changed, 301 insertions(+), 58 deletions(-) create mode 100644 DashAI/alembic/versions/a7d2c9e4f1b0_add_model_name_to_parameter_history.py diff --git a/DashAI/alembic/versions/a7d2c9e4f1b0_add_model_name_to_parameter_history.py b/DashAI/alembic/versions/a7d2c9e4f1b0_add_model_name_to_parameter_history.py new file mode 100644 index 000000000..ca9604eae --- /dev/null +++ b/DashAI/alembic/versions/a7d2c9e4f1b0_add_model_name_to_parameter_history.py @@ -0,0 +1,28 @@ +"""Add model_name to parameter_history + +Revision ID: a7d2c9e4f1b0 +Revises: f1a2b3c4d5e6 +Create Date: 2026-07-02 00:00:00.000000 + +""" + +from typing import Sequence, Union + +import sqlalchemy as sa + +from alembic import op + +revision: str = "a7d2c9e4f1b0" +down_revision: Union[str, None] = "f1a2b3c4d5e6" +branch_labels: Union[str, Sequence[str], None] = None +depends_on: Union[str, Sequence[str], None] = None + + +def upgrade() -> None: + with op.batch_alter_table("parameter_history", schema=None) as batch_op: + batch_op.add_column(sa.Column("model_name", sa.String(), nullable=True)) + + +def downgrade() -> None: + with op.batch_alter_table("parameter_history", schema=None) as batch_op: + batch_op.drop_column("model_name") diff --git a/DashAI/back/api/api_v1/endpoints/generative_session.py b/DashAI/back/api/api_v1/endpoints/generative_session.py index 6c59fb3d6..d4ec8d305 100644 --- a/DashAI/back/api/api_v1/endpoints/generative_session.py +++ b/DashAI/back/api/api_v1/endpoints/generative_session.py @@ -115,6 +115,7 @@ async def upload_generative_session( session_params_entry = GenerativeSessionParameterHistory( session_id=session.id, parameters=session.parameters, + model_name=session.model_name, modified_at=datetime.now(), ) db.add(session_params_entry) @@ -357,6 +358,7 @@ async def update_generative_session( model is unknown, not a generative model, or not yet downloaded. """ from DashAI.back.models.base_generative_model import BaseGenerativeModel + with session_factory() as db: try: session = db.get(GenerativeSession, session_id) @@ -395,8 +397,10 @@ async def update_generative_session( if description is not None: setattr(session, "description", description) - # Validate and apply a model change if provided - if model_name is not None: + # Validate and apply a model change if provided. A model may be + # selected even when it is not downloaded yet; the chat blocks input + # and offers a download until the weights become available. + if model_name is not None and model_name != session.model_name: try: model_class = component_registry[model_name]["class"] except KeyError as e: @@ -409,15 +413,37 @@ async def update_generative_session( status_code=status.HTTP_400_BAD_REQUEST, detail=f"Model {model_name} is not a valid generative model.", ) - entry = component_registry[model_name] - if getattr( - entry["class"], "REQUIRES_DOWNLOAD", False - ) and not entry.get("downloaded", False): - raise HTTPException( - status_code=status.HTTP_409_CONFLICT, - detail=f"Model {model_name} must be downloaded before use.", + + # Resolve the parameters for the new model: reuse the most + # recent parameters used for it in this session, else fall back + # to the model's schema defaults (its field placeholders). + last_used = ( + db.query(GenerativeSessionParameterHistory) + .filter( + GenerativeSessionParameterHistory.session_id == session_id, + GenerativeSessionParameterHistory.model_name == model_name, + ) + .order_by(GenerativeSessionParameterHistory.modified_at.desc()) + .first() + ) + if last_used is not None: + new_parameters = last_used.parameters + else: + properties = model_class.get_schema().get("properties", {}) + new_parameters = { + key: prop.get("placeholder") for key, prop in properties.items() + } + + session.model_name = model_name + session.parameters = new_parameters + db.add( + GenerativeSessionParameterHistory( + session_id=session.id, + parameters=new_parameters, + model_name=model_name, + modified_at=datetime.now(), ) - setattr(session, "model_name", model_name) + ) if name is not None or description is not None or model_name is not None: session.last_modified = datetime.now() @@ -489,6 +515,7 @@ async def update_generative_session_params( session_params_entry = GenerativeSessionParameterHistory( session_id=session.id, parameters=updated_parameters, + model_name=session.model_name, modified_at=datetime.now(), ) db.add(session_params_entry) @@ -617,26 +644,42 @@ async def get_parameter_history_entry( ) parameters_history = [p.__dict__ for p in parameters_history] + if not parameters_history: + return [] events = [] prev_params = parameters_history[0]["parameters"] + prev_model = parameters_history[0].get("model_name") for i in range(1, len(parameters_history)): curr = parameters_history[i] curr_params = curr["parameters"] + curr_model = curr.get("model_name") changes = [] - for key in curr_params: - old_val = prev_params.get(key) - new_val = curr_params[key] - if old_val != new_val: - changes.append( - { - "parameter": key, - "oldValue": old_val, - "newValue": new_val, - } - ) + # A model switch resets parameters to the new model's own + # values, so the raw parameter diff would be noise; report only + # the model change for that entry. + if curr_model and prev_model and curr_model != prev_model: + changes.append( + { + "parameter": "model", + "oldValue": prev_model, + "newValue": curr_model, + } + ) + else: + for key in curr_params: + old_val = prev_params.get(key) + new_val = curr_params[key] + if old_val != new_val: + changes.append( + { + "parameter": key, + "oldValue": old_val, + "newValue": new_val, + } + ) events.append( { @@ -646,6 +689,7 @@ async def get_parameter_history_entry( } ) prev_params = curr_params + prev_model = curr_model return events diff --git a/DashAI/back/dependencies/database/models.py b/DashAI/back/dependencies/database/models.py index cdf57d968..20b547eb6 100644 --- a/DashAI/back/dependencies/database/models.py +++ b/DashAI/back/dependencies/database/models.py @@ -658,6 +658,10 @@ class GenerativeSessionParameterHistory(Base): nullable=False, ) parameters: Mapped[JSON] = mapped_column(JSON, nullable=False) + # Model active when this snapshot was taken. Nullable so pre-migration rows + # remain valid; the parameters-history derivation only emits a model-change + # event when two consecutive snapshots both carry a model name that differs. + model_name: Mapped[str] = mapped_column(String, nullable=True) modified_at: Mapped[DateTime] = mapped_column( DateTime, default=datetime.now, diff --git a/DashAI/front/src/components/generative/CreateSessionContext.jsx b/DashAI/front/src/components/generative/CreateSessionContext.jsx index 34bdd62ce..5aa188a24 100644 --- a/DashAI/front/src/components/generative/CreateSessionContext.jsx +++ b/DashAI/front/src/components/generative/CreateSessionContext.jsx @@ -192,14 +192,35 @@ export function CreateSessionProvider({ children }) { [existingSessions], ); + // A model whose download was removed can no longer be used to create a + // session, so it must not stay selected. + const isUnavailable = (model) => + Boolean(model?.metadata?.requires_download) && !model?.downloaded; + // Sync selectedModel from URL param on load and after language-triggered // model refetch so display_name / description reflect the active language. + // If the URL points at a model that is no longer downloaded, drop back to + // the model selection step. useEffect(() => { if (!modelName || models.length === 0) return; const match = models.find((m) => m.name === modelName); - if (match) handleSelectModel(match); + if (!match) return; + if (isUnavailable(match)) { + setSelectedModel(null); + navigate("/app/generative/sessions/new"); + } else { + handleSelectModel(match); + } }, [modelName, models]); + // On the selection step, deselect a model if its download disappears (e.g. + // deleted from the inline control) after a models refetch. + useEffect(() => { + if (!selectedModel) return; + const match = models.find((m) => m.name === selectedModel.name); + if (match && isUnavailable(match)) setSelectedModel(null); + }, [models]); + const handleNext = () => { if (step === 0 && selectedModel) navigate(`/app/generative/sessions/new/${selectedModel.name}`); diff --git a/DashAI/front/src/components/generative/GenerativeChat.jsx b/DashAI/front/src/components/generative/GenerativeChat.jsx index 798c0bdb2..016a9a631 100644 --- a/DashAI/front/src/components/generative/GenerativeChat.jsx +++ b/DashAI/front/src/components/generative/GenerativeChat.jsx @@ -14,8 +14,14 @@ import { postProcess } from "../../api/process"; import { enqueueGenerativeProcessJob } from "../../api/job"; import { startJobQueue } from "../../api/job"; import { getHistoryBySessionId, getSessionById } from "../../api/session"; +import { + getComponentById, + getComponentDownloadStatus, +} from "../../api/component"; +import { getRelatedComponents } from "../../api/generativeTask"; import InfoSessionModal from "./InfoSessionModal"; import ModelSwitcher from "./ModelSwitcher"; +import ComponentDownloadControl from "../models/model/ComponentDownloadControl"; import { useSnackbar } from "notistack"; import { MediaInput } from "./MediaInput"; import { Trans, useTranslation } from "react-i18next"; @@ -48,6 +54,8 @@ export default function GenerativeChat() { const [showScrollButton, setShowScrollButton] = useState(false); const [sessionInfo, setSessionInfo] = useState(null); const [sessionInfoVisible, setSessionInfoVisible] = useState(false); + const [modelComponent, setModelComponent] = useState(null); + const [modelsByName, setModelsByName] = useState({}); const { enqueueSnackbar } = useSnackbar(); const { t } = useTranslation(["generative"]); const tourContext = useTourContext(); @@ -80,6 +88,49 @@ export default function GenerativeChat() { }); }; + // Resolve the session model's metadata plus a reconciled download status, so + // the chat can block input and offer a download when the weights are missing + // (e.g. after switching to a not downloaded model or deleting its download). + const modelName = sessionInfo?.model_name; + const refreshModelStatus = () => { + if (!modelName) { + setModelComponent(null); + return; + } + Promise.all([ + getComponentById(modelName), + getComponentDownloadStatus(modelName), + ]) + .then(([component, status]) => { + setModelComponent({ ...component, downloaded: status.downloaded }); + }) + .catch(() => setModelComponent(null)); + }; + + useEffect(() => { + refreshModelStatus(); + }, [modelName, paramsVersion]); + + // Map component name -> display name for the task's models, used to render + // model change history events with friendly names instead of class names. + useEffect(() => { + const currentTaskName = sessionInfo?.task_name; + if (!currentTaskName) return; + getRelatedComponents(currentTaskName) + .then((components) => { + const map = {}; + (components || []).forEach((c) => { + map[c.name] = c.display_name || c.name; + }); + setModelsByName(map); + }) + .catch(() => setModelsByName({})); + }, [sessionInfo?.task_name]); + + const modelBlocked = + Boolean(modelComponent?.metadata?.requires_download) && + !modelComponent?.downloaded; + const getMessages = () => { getProcessesBySessionId(sessionId).then((response) => { setIsLoadingMessage(false); @@ -202,23 +253,37 @@ export default function GenerativeChat() { }); let historyObject = history.map((entry) => { + const isModelChange = entry.changes.some((c) => c.parameter === "model"); return { type: "history", timestamp: entry.timestamp, id: entry.id, - changedMessage: entry.changes.map((change) => ( - - {change.parameter}: {change.oldValue}{" "} - {change.newValue}{" "} - - )), + isModelChange, + changedMessage: entry.changes.map((change) => { + const isModel = change.parameter === "model"; + const label = isModel + ? t("generative:label.sessionModel") + : change.parameter; + const oldValue = isModel + ? modelsByName[change.oldValue] || change.oldValue + : change.oldValue; + const newValue = isModel + ? modelsByName[change.newValue] || change.newValue + : change.newValue; + return ( + + {label}: {oldValue} {" "} + {newValue}{" "} + + ); + }), }; }); @@ -313,8 +378,17 @@ export default function GenerativeChat() { > {message.type === "history" ? ( - - Parameters updated: {message.changedMessage} + + {message.isModelChange + ? "Model changed: " + : "Parameters updated: "} + {message.changedMessage} ) : ( @@ -364,15 +438,40 @@ export default function GenerativeChat() { )} - {/* Chat input */} - { - handleSendMessage(input); - }} - isLoading={isLoadingMessage} - inputsCardinality={inputsCardinality} - /> + {/* Chat input, or a download prompt when the model is not available */} + {modelBlocked ? ( + + + {t("generative:label.modelNotDownloaded")} + + refreshModelStatus()} + /> + + ) : ( + { + handleSendMessage(input); + }} + isLoading={isLoadingMessage} + inputsCardinality={inputsCardinality} + /> + )} {/* Session Info Modal */} {sessionInfo && ( diff --git a/DashAI/front/src/components/generative/ModelSwitcher.jsx b/DashAI/front/src/components/generative/ModelSwitcher.jsx index 0a25dbc9c..54308112c 100644 --- a/DashAI/front/src/components/generative/ModelSwitcher.jsx +++ b/DashAI/front/src/components/generative/ModelSwitcher.jsx @@ -74,18 +74,16 @@ export default function ModelSwitcher({ sx={{ minWidth: 200 }} > {options.map((model) => { - const needsDownload = + // A not-downloaded model is still selectable; the chat blocks input + // and offers the download once it becomes the session's model. + const notDownloaded = Boolean(model.metadata?.requires_download) && !model.downloaded && model.name !== currentModelName; return ( - + {model.display_name || model.name} - {needsDownload + {notDownloaded ? ` (${t("generative:label.downloadRequired")})` : ""} diff --git a/DashAI/front/src/utils/i18n/locales/de/generative.json b/DashAI/front/src/utils/i18n/locales/de/generative.json index 701767f7c..bf88e144f 100644 --- a/DashAI/front/src/utils/i18n/locales/de/generative.json +++ b/DashAI/front/src/utils/i18n/locales/de/generative.json @@ -39,6 +39,8 @@ "startBySelectingATask": "Beginnen Sie mit der Auswahl einer generativen Aufgabe.", "noSessionsFound": "Keine Sitzungen gefunden", "parameterChangeEvent": "Parameter aktualisiert: <1>", + "modelChangeEvent": "Modell geändert: <1>", + "modelNotDownloaded": "Das Modell dieser Sitzung ist noch nicht heruntergeladen. Laden Sie es herunter, um fortzufahren.", "parameterChangeHistory": "Parameteränderungsverlauf für die aktuelle Sitzung", "searchSessions": "Sitzungen suchen", "selectGenerativeTask": "Generative Aufgabe für neue Sitzung auswählen", diff --git a/DashAI/front/src/utils/i18n/locales/en/generative.json b/DashAI/front/src/utils/i18n/locales/en/generative.json index 9bd5ddb0c..3c2ade281 100644 --- a/DashAI/front/src/utils/i18n/locales/en/generative.json +++ b/DashAI/front/src/utils/i18n/locales/en/generative.json @@ -39,6 +39,8 @@ "startBySelectingATask": "Start by selecting a generative task.", "noSessionsFound": "No sessions found", "parameterChangeEvent": "Parameters updated: <1>", + "modelChangeEvent": "Model changed: <1>", + "modelNotDownloaded": "This session's model is not downloaded yet. Download it to continue.", "parameterChangeHistory": "Parameter change history for the current session", "searchSessions": "Search Sessions", "selectGenerativeTask": "Select a generative task to start a new session", diff --git a/DashAI/front/src/utils/i18n/locales/es/generative.json b/DashAI/front/src/utils/i18n/locales/es/generative.json index 08f4bdefe..fcf27f01a 100644 --- a/DashAI/front/src/utils/i18n/locales/es/generative.json +++ b/DashAI/front/src/utils/i18n/locales/es/generative.json @@ -39,6 +39,8 @@ "startBySelectingATask": "Comienza seleccionando una tarea generativa.", "noSessionsFound": "No se encontraron sesiones", "parameterChangeEvent": "Parámetros actualizados: <1>", + "modelChangeEvent": "Modelo cambiado: <1>", + "modelNotDownloaded": "El modelo de esta sesión aún no está descargado. Descárgalo para continuar.", "parameterChangeHistory": "Historial de cambios de parámetros para la sesión actual", "searchSessions": "Buscar Sesiones", "selectGenerativeTask": "Seleccione una tarea generativa para comenzar una nueva sesión", diff --git a/DashAI/front/src/utils/i18n/locales/pt/generative.json b/DashAI/front/src/utils/i18n/locales/pt/generative.json index 69f867d4b..06643287f 100644 --- a/DashAI/front/src/utils/i18n/locales/pt/generative.json +++ b/DashAI/front/src/utils/i18n/locales/pt/generative.json @@ -39,6 +39,8 @@ "startBySelectingATask": "Comece selecionando uma tarefa generativa.", "noSessionsFound": "Nenhuma sessão encontrada", "parameterChangeEvent": "Parâmetros atualizados: <1>", + "modelChangeEvent": "Modelo alterado: <1>", + "modelNotDownloaded": "O modelo desta sessão ainda não foi baixado. Baixe-o para continuar.", "parameterChangeHistory": "Histórico de alterações de parâmetros para a sessão atual", "searchSessions": "Buscar Sessões", "selectGenerativeTask": "Selecione uma tarefa generativa para começar uma nova sessão", diff --git a/DashAI/front/src/utils/i18n/locales/zh/generative.json b/DashAI/front/src/utils/i18n/locales/zh/generative.json index 6f3b9870c..093dfda13 100644 --- a/DashAI/front/src/utils/i18n/locales/zh/generative.json +++ b/DashAI/front/src/utils/i18n/locales/zh/generative.json @@ -39,6 +39,8 @@ "startBySelectingATask": "从选择生成式任务开始。", "noSessionsFound": "未找到会话", "parameterChangeEvent": "参数已更新：<1>", + "modelChangeEvent": "模型已更改：<1>", + "modelNotDownloaded": "该会话的模型尚未下载。请先下载后再继续。", "parameterChangeHistory": "当前会话的参数变更历史", "searchSessions": "搜索会话", "selectGenerativeTask": "选择生成式任务以开始新会话", diff --git a/tests/back/api/test_generative_session_download_gate.py b/tests/back/api/test_generative_session_download_gate.py index 57de0d995..6881df992 100644 --- a/tests/back/api/test_generative_session_download_gate.py +++ b/tests/back/api/test_generative_session_download_gate.py @@ -100,9 +100,14 @@ def _create_sd_session(client, name): ) -def test_change_session_model_to_undownloaded_returns_409(client): - """Switching a session to a not-downloaded model must return HTTP 409.""" - created = _create_sd_session(client, "gen-switch-409") +def test_change_session_model_to_undownloaded_succeeds(client): + """Switching to a not-downloaded model is allowed; the chat gates its use. + + A user may point a session at any registered generative model even if its + weights are not present yet; the download is offered from the chat instead + of being blocked at switch time. + """ + created = _create_sd_session(client, "gen-switch-undownloaded") assert created.status_code == 201 session_id = created.json()["id"] @@ -110,8 +115,8 @@ def test_change_session_model_to_undownloaded_returns_409(client): f"/api/v1/generative-session/{session_id}", params={"model_name": "Qwen25_15BInstruct"}, ) - assert resp.status_code == 409 - assert "download" in resp.json()["detail"].lower() + assert resp.status_code == 200 + assert resp.json()["model_name"] == "Qwen25_15BInstruct" client.delete(f"/api/v1/generative-session/{session_id}") @@ -146,3 +151,37 @@ def test_change_session_model_valid_returns_200(client): assert resp.json()["model_name"] == "StableDiffusionV2Model" client.delete(f"/api/v1/generative-session/{session_id}") + + +def test_switch_model_resets_params_and_records_history(client): + """A switch resets params to the new model's defaults and logs the change.""" + created = _create_sd_session(client, "gen-switch-history") + assert created.status_code == 201 + session_id = created.json()["id"] + + resp = client.patch( + f"/api/v1/generative-session/{session_id}", + params={"model_name": "Qwen25_15BInstruct"}, + ) + assert resp.status_code == 200 + params = resp.json()["parameters"] + # Parameters were reset to the target model's own fields, not carried over + # from the Stable Diffusion session. + assert "num_inference_steps" not in params + assert "max_tokens" in params + + history = client.get(f"/api/v1/generative-session/parameters-history/{session_id}") + assert history.status_code == 200 + model_changes = [ + change + for event in history.json() + for change in event["changes"] + if change["parameter"] == "model" + ] + assert { + "parameter": "model", + "oldValue": "StableDiffusionV2Model", + "newValue": "Qwen25_15BInstruct", + } in model_changes + + client.delete(f"/api/v1/generative-session/{session_id}") From a54c780e1c1db1f5d2c78127c907cceb90f32178 Mon Sep 17 00:00:00 2001 From: Irozuku Date: Thu, 2 Jul 2026 14:07:57 -0400 Subject: [PATCH 39/89] feat: make text classification transformers downloadable --- .../models/hugging_face/albert_transformer.py | 1 + .../base_text_classification_transformer.py | 58 +++++++++++++++++-- .../models/hugging_face/bert_transformer.py | 1 + .../models/hugging_face/bertin_transformer.py | 1 + .../models/hugging_face/beto_transformer.py | 1 + .../hugging_face/deberta_v3_transformer.py | 1 + .../hugging_face/distilbert_transformer.py | 1 + .../hugging_face/electra_transformer.py | 1 + .../models/hugging_face/minilm_transformer.py | 1 + .../hugging_face/modernbert_transformer.py | 1 + .../multilingual_bert_transformer.py | 1 + .../hugging_face/roberta_transformer.py | 1 + .../hugging_face/xlm_roberta_transformer.py | 1 + .../models/hugging_face/xlnet_transformer.py | 1 + ...st_base_text_classification_transformer.py | 5 +- .../test_text_classification_downloadable.py | 45 ++++++++++++++ 16 files changed, 116 insertions(+), 5 deletions(-) create mode 100644 tests/back/models/test_text_classification_downloadable.py diff --git a/DashAI/back/models/hugging_face/albert_transformer.py b/DashAI/back/models/hugging_face/albert_transformer.py index a01622ab0..104deb58a 100644 --- a/DashAI/back/models/hugging_face/albert_transformer.py +++ b/DashAI/back/models/hugging_face/albert_transformer.py @@ -59,4 +59,5 @@ class AlbertTransformer(HuggingFaceTextClassificationTransformer): ICON: str = "Speed" SCHEMA = DistilBertTransformerSchema MODEL_NAME: str = "albert-base-v2" + DOWNLOAD_SIZE_BYTES: int = 47_000_000 TEMP_CHECKPOINT_DIR: str = "DashAI/back/user_models/temp_checkpoints_albert" diff --git a/DashAI/back/models/hugging_face/base_text_classification_transformer.py b/DashAI/back/models/hugging_face/base_text_classification_transformer.py index 741207eae..dc839c699 100644 --- a/DashAI/back/models/hugging_face/base_text_classification_transformer.py +++ b/DashAI/back/models/hugging_face/base_text_classification_transformer.py @@ -7,10 +7,11 @@ """ from pathlib import Path -from typing import TYPE_CHECKING, Union +from typing import TYPE_CHECKING, Optional, Union from sklearn.exceptions import NotFittedError +from DashAI.back.dependencies.downloads.downloadable import HFDownloadableMixin from DashAI.back.models.text_classification_model import TextClassificationModel from DashAI.back.models.utils import ( GPU_OR_CPU_PLACEHOLDER, @@ -22,7 +23,9 @@ from DashAI.back.dataloaders.classes.dashai_dataset import DashAIDataset -class HuggingFaceTextClassificationTransformer(TextClassificationModel): +class HuggingFaceTextClassificationTransformer( + HFDownloadableMixin, TextClassificationModel +): """Base implementation for Hugging Face text classification wrappers. Subclasses are expected to define at least ``MODEL_NAME`` and optionally @@ -44,8 +47,51 @@ class HuggingFaceTextClassificationTransformer(TextClassificationModel): "DashAI/back/user_models/temp_checkpoints_hf_text_classification" ) MAX_TOKEN_LENGTH: int = 512 + # Approximate on-disk size of the pretrained checkpoint; subclasses override + # with a value closer to their specific model for the download UI. + DOWNLOAD_SIZE_BYTES: int = 450_000_000 - def __init__(self, model=None, **kwargs): + @classmethod + def hf_repos(cls): + """Derive the single HuggingFace repo from the subclass ``MODEL_NAME``. + + Returns + ------- + list of tuple of (str, str) + A single ``(repo_id, repo_type)`` pair derived from ``MODEL_NAME``, + or an empty list when ``MODEL_NAME`` is not set. + """ + return [(cls.MODEL_NAME, "model")] if cls.MODEL_NAME else [] + + def _pretrained_source(self, pretrained_dir: Optional[str]) -> str: + """Resolve where to load the tokenizer and weights from. + + Prefers an explicit ``pretrained_dir`` (a saved run), then the local + component download folder when the weights are present, and finally + falls back to the Hugging Face Hub repo id. Downloading is enforced by + the run/session gates before real use; the Hub fallback keeps direct + instantiation working when nothing has been downloaded. + + Parameters + ---------- + pretrained_dir : str or None + Directory of a previously saved run, if any. + + Returns + ------- + str + A path or repo id accepted by ``from_pretrained``. + """ + if pretrained_dir: + return pretrained_dir + try: + if self.is_downloaded(): + return str(self._repo_dir(self.MODEL_NAME)) + except Exception: + pass + return self.MODEL_NAME + + def __init__(self, model=None, pretrained_dir: Optional[str] = None, **kwargs): """Initialize the transformer model. The process includes the instantiation of the pretrained model and the @@ -77,7 +123,7 @@ def __init__(self, model=None, **kwargs): f"{self.__class__.__name__} must define a non-empty MODEL_NAME." ) - self.model_name = self.MODEL_NAME + self.model_name = self._pretrained_source(pretrained_dir) self.tokenizer = AutoTokenizer.from_pretrained(self.model_name) self.log_train_every_n_epochs = kwargs.get("log_train_every_n_epochs", 1) @@ -375,6 +421,9 @@ def save(self, filename: Union[str, "Path"]) -> None: save_dir.mkdir(parents=True, exist_ok=True) self.model.save_pretrained(save_dir) + # Persist the tokenizer alongside the weights so a saved run is + # self-contained and does not depend on the component download folder. + self.tokenizer.save_pretrained(save_dir) config = AutoConfig.from_pretrained(save_dir) config.custom_params = { "num_train_epochs": self.training_args_params.get("num_train_epochs"), @@ -419,6 +468,7 @@ def load( loaded_model = cls( model=model, + pretrained_dir=str(filename), num_labels=custom_params.get("num_labels"), num_train_epochs=custom_params.get("num_train_epochs", 2), batch_size=custom_params.get("batch_size", 16), diff --git a/DashAI/back/models/hugging_face/bert_transformer.py b/DashAI/back/models/hugging_face/bert_transformer.py index 7fddc261f..d22451107 100644 --- a/DashAI/back/models/hugging_face/bert_transformer.py +++ b/DashAI/back/models/hugging_face/bert_transformer.py @@ -58,4 +58,5 @@ class BertTransformer(HuggingFaceTextClassificationTransformer): ICON: str = "Psychology" SCHEMA = DistilBertTransformerSchema MODEL_NAME: str = "bert-base-uncased" + DOWNLOAD_SIZE_BYTES: int = 440_000_000 TEMP_CHECKPOINT_DIR: str = "DashAI/back/user_models/temp_checkpoints_bert" diff --git a/DashAI/back/models/hugging_face/bertin_transformer.py b/DashAI/back/models/hugging_face/bertin_transformer.py index 693004535..109a7b649 100644 --- a/DashAI/back/models/hugging_face/bertin_transformer.py +++ b/DashAI/back/models/hugging_face/bertin_transformer.py @@ -58,4 +58,5 @@ class BertinTransformer(HuggingFaceTextClassificationTransformer): ICON: str = "RecordVoiceOver" SCHEMA = DistilBertTransformerSchema MODEL_NAME: str = "bertin-project/bertin-roberta-base-spanish" + DOWNLOAD_SIZE_BYTES: int = 500_000_000 TEMP_CHECKPOINT_DIR: str = "DashAI/back/user_models/temp_checkpoints_bertin" diff --git a/DashAI/back/models/hugging_face/beto_transformer.py b/DashAI/back/models/hugging_face/beto_transformer.py index cddfb45a8..db8440718 100644 --- a/DashAI/back/models/hugging_face/beto_transformer.py +++ b/DashAI/back/models/hugging_face/beto_transformer.py @@ -58,4 +58,5 @@ class BetoTransformer(HuggingFaceTextClassificationTransformer): ICON: str = "RecordVoiceOver" SCHEMA = DistilBertTransformerSchema MODEL_NAME: str = "dccuchile/bert-base-spanish-wwm-cased" + DOWNLOAD_SIZE_BYTES: int = 440_000_000 TEMP_CHECKPOINT_DIR: str = "DashAI/back/user_models/temp_checkpoints_beto" diff --git a/DashAI/back/models/hugging_face/deberta_v3_transformer.py b/DashAI/back/models/hugging_face/deberta_v3_transformer.py index da258deb9..271fa0f76 100644 --- a/DashAI/back/models/hugging_face/deberta_v3_transformer.py +++ b/DashAI/back/models/hugging_face/deberta_v3_transformer.py @@ -61,4 +61,5 @@ class DebertaV3Transformer(HuggingFaceTextClassificationTransformer): ICON: str = "Psychology" SCHEMA = DebertaV3TransformerSchema MODEL_NAME: str = "microsoft/deberta-v3-base" + DOWNLOAD_SIZE_BYTES: int = 440_000_000 TEMP_CHECKPOINT_DIR: str = "DashAI/back/user_models/temp_checkpoints_deberta_v3" diff --git a/DashAI/back/models/hugging_face/distilbert_transformer.py b/DashAI/back/models/hugging_face/distilbert_transformer.py index f43df4752..334ce67ec 100644 --- a/DashAI/back/models/hugging_face/distilbert_transformer.py +++ b/DashAI/back/models/hugging_face/distilbert_transformer.py @@ -305,4 +305,5 @@ class DistilBertTransformer(HuggingFaceTextClassificationTransformer): ICON: str = "Psychology" SCHEMA = DistilBertTransformerSchema MODEL_NAME: str = "distilbert-base-uncased" + DOWNLOAD_SIZE_BYTES: int = 270_000_000 TEMP_CHECKPOINT_DIR: str = "DashAI/back/user_models/temp_checkpoints_distilbert" diff --git a/DashAI/back/models/hugging_face/electra_transformer.py b/DashAI/back/models/hugging_face/electra_transformer.py index 7371a6f2d..00e26919e 100644 --- a/DashAI/back/models/hugging_face/electra_transformer.py +++ b/DashAI/back/models/hugging_face/electra_transformer.py @@ -58,4 +58,5 @@ class ElectraTransformer(HuggingFaceTextClassificationTransformer): ICON: str = "ElectricBolt" SCHEMA = DistilBertTransformerSchema MODEL_NAME: str = "google/electra-small-discriminator" + DOWNLOAD_SIZE_BYTES: int = 54_000_000 TEMP_CHECKPOINT_DIR: str = "DashAI/back/user_models/temp_checkpoints_electra" diff --git a/DashAI/back/models/hugging_face/minilm_transformer.py b/DashAI/back/models/hugging_face/minilm_transformer.py index 6d8e75a87..23a323b4d 100644 --- a/DashAI/back/models/hugging_face/minilm_transformer.py +++ b/DashAI/back/models/hugging_face/minilm_transformer.py @@ -58,4 +58,5 @@ class MiniLMTransformer(HuggingFaceTextClassificationTransformer): ICON: str = "Speed" SCHEMA = DistilBertTransformerSchema MODEL_NAME: str = "microsoft/MiniLM-L12-H384-uncased" + DOWNLOAD_SIZE_BYTES: int = 130_000_000 TEMP_CHECKPOINT_DIR: str = "DashAI/back/user_models/temp_checkpoints_minilm" diff --git a/DashAI/back/models/hugging_face/modernbert_transformer.py b/DashAI/back/models/hugging_face/modernbert_transformer.py index ae2be1fff..ad7e26033 100644 --- a/DashAI/back/models/hugging_face/modernbert_transformer.py +++ b/DashAI/back/models/hugging_face/modernbert_transformer.py @@ -55,5 +55,6 @@ class ModernBertTransformer(HuggingFaceTextClassificationTransformer): ICON: str = "Psychology" SCHEMA = ModernBertTransformerSchema MODEL_NAME: str = "answerdotai/ModernBERT-base" + DOWNLOAD_SIZE_BYTES: int = 600_000_000 MAX_TOKEN_LENGTH: int = 8192 TEMP_CHECKPOINT_DIR: str = "DashAI/back/user_models/temp_checkpoints_modernbert" diff --git a/DashAI/back/models/hugging_face/multilingual_bert_transformer.py b/DashAI/back/models/hugging_face/multilingual_bert_transformer.py index bd1e44bfe..d012a2e8b 100644 --- a/DashAI/back/models/hugging_face/multilingual_bert_transformer.py +++ b/DashAI/back/models/hugging_face/multilingual_bert_transformer.py @@ -59,6 +59,7 @@ class MultilingualBertTransformer(HuggingFaceTextClassificationTransformer): ICON: str = "Translate" SCHEMA = DistilBertTransformerSchema MODEL_NAME: str = "bert-base-multilingual-cased" + DOWNLOAD_SIZE_BYTES: int = 680_000_000 TEMP_CHECKPOINT_DIR: str = ( "DashAI/back/user_models/temp_checkpoints_multilingual_bert" ) diff --git a/DashAI/back/models/hugging_face/roberta_transformer.py b/DashAI/back/models/hugging_face/roberta_transformer.py index 5f65ea004..8e8fd40f1 100644 --- a/DashAI/back/models/hugging_face/roberta_transformer.py +++ b/DashAI/back/models/hugging_face/roberta_transformer.py @@ -58,4 +58,5 @@ class RobertaTransformer(HuggingFaceTextClassificationTransformer): ICON: str = "SmartToy" SCHEMA = DistilBertTransformerSchema MODEL_NAME: str = "roberta-base" + DOWNLOAD_SIZE_BYTES: int = 500_000_000 TEMP_CHECKPOINT_DIR: str = "DashAI/back/user_models/temp_checkpoints_roberta" diff --git a/DashAI/back/models/hugging_face/xlm_roberta_transformer.py b/DashAI/back/models/hugging_face/xlm_roberta_transformer.py index 6e04652be..51722452e 100644 --- a/DashAI/back/models/hugging_face/xlm_roberta_transformer.py +++ b/DashAI/back/models/hugging_face/xlm_roberta_transformer.py @@ -62,4 +62,5 @@ class XlmRobertaTransformer(HuggingFaceTextClassificationTransformer): ICON: str = "Language" SCHEMA = DistilBertTransformerSchema MODEL_NAME: str = "xlm-roberta-base" + DOWNLOAD_SIZE_BYTES: int = 1_100_000_000 TEMP_CHECKPOINT_DIR: str = "DashAI/back/user_models/temp_checkpoints_xlm_roberta" diff --git a/DashAI/back/models/hugging_face/xlnet_transformer.py b/DashAI/back/models/hugging_face/xlnet_transformer.py index 892157c7e..0bc110548 100644 --- a/DashAI/back/models/hugging_face/xlnet_transformer.py +++ b/DashAI/back/models/hugging_face/xlnet_transformer.py @@ -58,4 +58,5 @@ class XlnetTransformer(HuggingFaceTextClassificationTransformer): ICON: str = "AutoAwesome" SCHEMA = DistilBertTransformerSchema MODEL_NAME: str = "xlnet-base-cased" + DOWNLOAD_SIZE_BYTES: int = 470_000_000 TEMP_CHECKPOINT_DIR: str = "DashAI/back/user_models/temp_checkpoints_xlnet" diff --git a/tests/back/models/test_base_text_classification_transformer.py b/tests/back/models/test_base_text_classification_transformer.py index 223fdee92..b3d879ced 100644 --- a/tests/back/models/test_base_text_classification_transformer.py +++ b/tests/back/models/test_base_text_classification_transformer.py @@ -4,7 +4,10 @@ class DummyTokenizer: - pass + def save_pretrained(self, save_directory): + save_path = Path(save_directory) + save_path.mkdir(parents=True, exist_ok=True) + (save_path / "tokenizer.json").write_text("{}", encoding="utf-8") class DummyConfig: diff --git a/tests/back/models/test_text_classification_downloadable.py b/tests/back/models/test_text_classification_downloadable.py new file mode 100644 index 000000000..243a05762 --- /dev/null +++ b/tests/back/models/test_text_classification_downloadable.py @@ -0,0 +1,45 @@ +"""Tests that text classification transformers expose download metadata.""" + +import pytest +from kink import di + +from DashAI.back.dependencies.downloads.downloadable import HFDownloadableMixin +from DashAI.back.models.hugging_face.distilbert_transformer import DistilBertTransformer + + +@pytest.fixture +def component_root(tmp_path): + """Inject a temporary COMPONENT_PATH into the kink DI container.""" + sentinel = object() + old = di["config"] if "config" in di else sentinel # noqa: SIM401 + di["config"] = {"COMPONENT_PATH": str(tmp_path)} + yield tmp_path + if old is sentinel: + del di["config"] + else: + di["config"] = old + + +def test_text_classification_is_downloadable(): + assert issubclass(DistilBertTransformer, HFDownloadableMixin) + assert DistilBertTransformer.REQUIRES_DOWNLOAD is True + assert DistilBertTransformer.DOWNLOAD_SIZE_BYTES is not None + + +def test_hf_repos_derived_from_model_name(): + assert DistilBertTransformer.hf_repos() == [("distilbert-base-uncased", "model")] + + +def test_metadata_flags_download(): + meta = DistilBertTransformer.get_metadata() + assert meta["requires_download"] is True + assert meta["download_size_bytes"] == DistilBertTransformer.DOWNLOAD_SIZE_BYTES + + +def test_is_downloaded_uses_component_dir(component_root): + assert DistilBertTransformer.is_downloaded() is False + + repo_dir = component_root / "DistilBertTransformer" / "distilbert-base-uncased" + repo_dir.mkdir(parents=True) + (repo_dir / "config.json").write_text("{}") + assert DistilBertTransformer.is_downloaded() is True From 4f3a4a3cf719fa6d86268bc873b01a3ffd272220 Mon Sep 17 00:00:00 2001 From: Irozuku Date: Thu, 2 Jul 2026 14:13:16 -0400 Subject: [PATCH 40/89] feat: make translation transformers downloadable --- .../dependencies/downloads/downloadable.py | 50 ++++++++++++++++ .../models/hugging_face/m2m100_transformer.py | 13 ++++- .../models/hugging_face/nllb_transformer.py | 14 ++++- .../hugging_face/t5_small_transformer.py | 13 ++++- tests/back/models/test_nllb_transformer.py | 5 ++ .../models/test_translation_downloadable.py | 58 +++++++++++++++++++ 6 files changed, 144 insertions(+), 9 deletions(-) create mode 100644 tests/back/models/test_translation_downloadable.py diff --git a/DashAI/back/dependencies/downloads/downloadable.py b/DashAI/back/dependencies/downloads/downloadable.py index 33bef1a25..5aa4fbe0b 100644 --- a/DashAI/back/dependencies/downloads/downloadable.py +++ b/DashAI/back/dependencies/downloads/downloadable.py @@ -183,3 +183,53 @@ def download(cls, report: Optional[ProgressReporter] = None) -> None: snapshot_download( repo_id=rid, repo_type=rtype, local_dir=str(target), **kwargs ) + + +class HFPretrainedDownloadMixin(HFDownloadableMixin): + """HuggingFace mixin for models built around a single ``MODEL_NAME`` repo. + + Covers the common ``from_pretrained`` case: the repo is derived from the + subclass ``MODEL_NAME`` and ``_pretrained_source`` returns where to load + from (a saved run, the local download folder, or the Hub as a fallback). + """ + + MODEL_NAME: str = "" + + @classmethod + def hf_repos(cls): + """Derive the single repo entry from ``MODEL_NAME``. + + Returns + ------- + list of tuple of (str, str) + ``[(MODEL_NAME, "model")]`` or an empty list when unset. + """ + return [(cls.MODEL_NAME, "model")] if cls.MODEL_NAME else [] + + def _pretrained_source(self, pretrained_dir: Optional[str] = None) -> str: + """Resolve where ``from_pretrained`` should load from. + + Prefers an explicit ``pretrained_dir`` (a saved run), then the local + component download folder when present, and finally falls back to the + Hub repo id. Downloading is enforced by the run/session gates before + real use; the Hub fallback keeps direct instantiation working when + nothing has been downloaded. + + Parameters + ---------- + pretrained_dir : str or None + Directory of a previously saved run, if any. + + Returns + ------- + str + A path or repo id accepted by ``from_pretrained``. + """ + if pretrained_dir: + return pretrained_dir + try: + if self.is_downloaded(): + return str(self._repo_dir(self.MODEL_NAME)) + except Exception: + pass + return self.MODEL_NAME diff --git a/DashAI/back/models/hugging_face/m2m100_transformer.py b/DashAI/back/models/hugging_face/m2m100_transformer.py index 264e0a49f..99c023d00 100644 --- a/DashAI/back/models/hugging_face/m2m100_transformer.py +++ b/DashAI/back/models/hugging_face/m2m100_transformer.py @@ -9,6 +9,9 @@ from DashAI.back.core.schema_fields import schema_field, string_field from DashAI.back.core.utils import MultilingualString +from DashAI.back.dependencies.downloads.downloadable import ( + HFPretrainedDownloadMixin, +) from DashAI.back.models.hugging_face.opus_mt_en_es_transformer import ( OpusMtEnESTransformerSchema, ) @@ -97,7 +100,7 @@ class M2M100TransformerSchema(OpusMtEnESTransformerSchema): ) # type: ignore -class M2M100Transformer(TranslationModel): +class M2M100Transformer(HFPretrainedDownloadMixin, TranslationModel): """M2M100 multilingual seq2seq model for configurable language-pair translation. Fine-tunes the ``facebook/m2m100_418M`` checkpoint from Meta AI. The base @@ -154,13 +157,15 @@ class M2M100Transformer(TranslationModel): ) COLOR: str = "#6A1B9A" ICON: str = "Language" + MODEL_NAME: str = "facebook/m2m100_418M" + DOWNLOAD_SIZE_BYTES: int = 1_900_000_000 - def __init__(self, model=None, **kwargs): + def __init__(self, model=None, pretrained_dir=None, **kwargs): kwargs = self.validate_and_transform(kwargs) from transformers import AutoTokenizer - self.model_name = "facebook/m2m100_418M" + self.model_name = self._pretrained_source(pretrained_dir) self.tokenizer = AutoTokenizer.from_pretrained(self.model_name) self.source_language = kwargs.get("source_language", "en") @@ -352,6 +357,7 @@ def save(self, filename: Union[str, "Path"]) -> None: save_dir.mkdir(parents=True, exist_ok=True) self.model.save_pretrained(save_dir) + self.tokenizer.save_pretrained(save_dir) config = AutoConfig.from_pretrained(save_dir) config.custom_params = { "num_train_epochs": self.training_args.get("num_train_epochs"), @@ -376,6 +382,7 @@ def load(cls, filename: Union[str, "Path"]): loaded_model = cls( model=model, + pretrained_dir=str(filename), num_train_epochs=custom_params.get("num_train_epochs"), batch_size=custom_params.get("batch_size"), learning_rate=custom_params.get("learning_rate"), diff --git a/DashAI/back/models/hugging_face/nllb_transformer.py b/DashAI/back/models/hugging_face/nllb_transformer.py index ab7a5556c..6628ba9d6 100644 --- a/DashAI/back/models/hugging_face/nllb_transformer.py +++ b/DashAI/back/models/hugging_face/nllb_transformer.py @@ -9,6 +9,9 @@ from DashAI.back.core.schema_fields import schema_field, string_field from DashAI.back.core.utils import MultilingualString +from DashAI.back.dependencies.downloads.downloadable import ( + HFPretrainedDownloadMixin, +) from DashAI.back.models.hugging_face.opus_mt_en_es_transformer import ( OpusMtEnESTransformerSchema, ) @@ -121,7 +124,7 @@ class NllbTransformerSchema(OpusMtEnESTransformerSchema): ) # type: ignore -class NllbTransformer(TranslationModel): +class NllbTransformer(HFPretrainedDownloadMixin, TranslationModel): """Pretrained transformer for configurable multilingual translation. This model fine-tunes the ``facebook/nllb-200-distilled-600M`` checkpoint from @@ -201,7 +204,10 @@ def _resolve_language_token_id(self, language_code: str, field_name: str) -> int raise ValueError(f"Unsupported {field_name} '{language_code}'.") - def __init__(self, model=None, **kwargs): + MODEL_NAME: str = "facebook/nllb-200-distilled-600M" + DOWNLOAD_SIZE_BYTES: int = 2_400_000_000 + + def __init__(self, model=None, pretrained_dir=None, **kwargs): """Initialize the NLLB tokenizer and model. Downloads the ``facebook/nllb-200-distilled-600M`` tokenizer and, @@ -231,7 +237,7 @@ def __init__(self, model=None, **kwargs): from transformers import AutoTokenizer - self.model_name = "facebook/nllb-200-distilled-600M" + self.model_name = self._pretrained_source(pretrained_dir) self.tokenizer = AutoTokenizer.from_pretrained(self.model_name) self.source_language = kwargs.get("source_language", "spa_Latn") @@ -533,6 +539,7 @@ def save(self, filename: Union[str, "Path"]) -> None: save_dir.mkdir(parents=True, exist_ok=True) self.model.save_pretrained(save_dir) + self.tokenizer.save_pretrained(save_dir) config = AutoConfig.from_pretrained(save_dir) config.custom_params = { "num_train_epochs": self.training_args.get("num_train_epochs"), @@ -573,6 +580,7 @@ def load(cls, filename: Union[str, "Path"]): loaded_model = cls( model=model, + pretrained_dir=str(filename), num_train_epochs=custom_params.get("num_train_epochs"), batch_size=custom_params.get("batch_size"), learning_rate=custom_params.get("learning_rate"), diff --git a/DashAI/back/models/hugging_face/t5_small_transformer.py b/DashAI/back/models/hugging_face/t5_small_transformer.py index 91373a064..366de45b5 100644 --- a/DashAI/back/models/hugging_face/t5_small_transformer.py +++ b/DashAI/back/models/hugging_face/t5_small_transformer.py @@ -9,6 +9,9 @@ from DashAI.back.core.schema_fields import enum_field, schema_field from DashAI.back.core.utils import MultilingualString +from DashAI.back.dependencies.downloads.downloadable import ( + HFPretrainedDownloadMixin, +) from DashAI.back.models.hugging_face.opus_mt_en_es_transformer import ( OpusMtEnESTransformerSchema, ) @@ -73,7 +76,7 @@ class T5SmallTransformerSchema(OpusMtEnESTransformerSchema): ) # type: ignore -class T5SmallTransformer(TranslationModel): +class T5SmallTransformer(HFPretrainedDownloadMixin, TranslationModel): """T5-small seq2seq model for English-to-{German, French, Romanian} translation. Fine-tunes the ``t5-small`` checkpoint from Google. Translation direction is @@ -127,13 +130,15 @@ class T5SmallTransformer(TranslationModel): ) COLOR: str = "#00695C" ICON: str = "Language" + MODEL_NAME: str = "t5-small" + DOWNLOAD_SIZE_BYTES: int = 240_000_000 - def __init__(self, model=None, **kwargs): + def __init__(self, model=None, pretrained_dir=None, **kwargs): kwargs = self.validate_and_transform(kwargs) from transformers import AutoTokenizer - self.model_name = "t5-small" + self.model_name = self._pretrained_source(pretrained_dir) self.tokenizer = AutoTokenizer.from_pretrained(self.model_name) self.target_language = kwargs.get("target_language", "German") @@ -318,6 +323,7 @@ def save(self, filename: Union[str, "Path"]) -> None: save_dir.mkdir(parents=True, exist_ok=True) self.model.save_pretrained(save_dir) + self.tokenizer.save_pretrained(save_dir) config = AutoConfig.from_pretrained(save_dir) config.custom_params = { "num_train_epochs": self.training_args.get("num_train_epochs"), @@ -341,6 +347,7 @@ def load(cls, filename: Union[str, "Path"]): loaded_model = cls( model=model, + pretrained_dir=str(filename), num_train_epochs=custom_params.get("num_train_epochs"), batch_size=custom_params.get("batch_size"), learning_rate=custom_params.get("learning_rate"), diff --git a/tests/back/models/test_nllb_transformer.py b/tests/back/models/test_nllb_transformer.py index 772e0b8f3..70f79c32b 100644 --- a/tests/back/models/test_nllb_transformer.py +++ b/tests/back/models/test_nllb_transformer.py @@ -27,6 +27,11 @@ def decode(self, token_ids, skip_special_tokens=True): def convert_tokens_to_ids(self, token): return self.lang_code_to_id.get(token, 0) + def save_pretrained(self, save_directory): + save_path = Path(save_directory) + save_path.mkdir(parents=True, exist_ok=True) + (save_path / "tokenizer.json").write_text("{}", encoding="utf-8") + class DummyNllbTokenizerNoLangMap: def __init__(self): diff --git a/tests/back/models/test_translation_downloadable.py b/tests/back/models/test_translation_downloadable.py new file mode 100644 index 000000000..cba69d495 --- /dev/null +++ b/tests/back/models/test_translation_downloadable.py @@ -0,0 +1,58 @@ +"""Tests that translation transformers expose download metadata.""" + +import pytest +from kink import di + +from DashAI.back.dependencies.downloads.downloadable import ( + HFPretrainedDownloadMixin, +) +from DashAI.back.models.hugging_face.m2m100_transformer import M2M100Transformer +from DashAI.back.models.hugging_face.nllb_transformer import NllbTransformer +from DashAI.back.models.hugging_face.t5_small_transformer import T5SmallTransformer + +_CASES = [ + (M2M100Transformer, "facebook/m2m100_418M"), + (NllbTransformer, "facebook/nllb-200-distilled-600M"), + (T5SmallTransformer, "t5-small"), +] + + +@pytest.fixture +def component_root(tmp_path): + """Inject a temporary COMPONENT_PATH into the kink DI container.""" + sentinel = object() + old = di["config"] if "config" in di else sentinel # noqa: SIM401 + di["config"] = {"COMPONENT_PATH": str(tmp_path)} + yield tmp_path + if old is sentinel: + del di["config"] + else: + di["config"] = old + + +@pytest.mark.parametrize(("model_cls", "repo_id"), _CASES) +def test_translation_is_downloadable(model_cls, repo_id): + assert issubclass(model_cls, HFPretrainedDownloadMixin) + assert model_cls.REQUIRES_DOWNLOAD is True + assert model_cls.DOWNLOAD_SIZE_BYTES is not None + assert model_cls.hf_repos() == [(repo_id, "model")] + + +@pytest.mark.parametrize(("model_cls", "repo_id"), _CASES) +def test_translation_metadata_flags_download(model_cls, repo_id): + meta = model_cls.get_metadata() + assert meta["requires_download"] is True + assert meta["download_size_bytes"] == model_cls.DOWNLOAD_SIZE_BYTES + + +@pytest.mark.parametrize(("model_cls", "repo_id"), _CASES) +def test_translation_is_downloaded_uses_component_dir( + model_cls, repo_id, component_root +): + assert model_cls.is_downloaded() is False + + leaf = repo_id.split("/")[-1] + repo_dir = component_root / model_cls.__name__ / leaf + repo_dir.mkdir(parents=True) + (repo_dir / "config.json").write_text("{}") + assert model_cls.is_downloaded() is True From 9099cccdc26c117beff8c0c39901482e3b79ab24 Mon Sep 17 00:00:00 2001 From: Irozuku Date: Thu, 2 Jul 2026 14:22:28 -0400 Subject: [PATCH 41/89] feat: make torchvision image classifiers downloadable --- .../dependencies/downloads/downloadable.py | 86 +++++++++++++++++++ .../efficientnet_b0_image_classifier.py | 15 +++- .../back/models/resnet18_image_classifier.py | 13 ++- .../back/models/resnet50_image_classifier.py | 13 ++- .../test_image_classification_downloadable.py | 54 ++++++++++++ 5 files changed, 175 insertions(+), 6 deletions(-) create mode 100644 tests/back/models/test_image_classification_downloadable.py diff --git a/DashAI/back/dependencies/downloads/downloadable.py b/DashAI/back/dependencies/downloads/downloadable.py index 5aa4fbe0b..56112910f 100644 --- a/DashAI/back/dependencies/downloads/downloadable.py +++ b/DashAI/back/dependencies/downloads/downloadable.py @@ -233,3 +233,89 @@ def _pretrained_source(self, pretrained_dir: Optional[str] = None) -> str: except Exception: pass return self.MODEL_NAME + + +class TorchvisionDownloadMixin(DownloadableMixin): + """Downloadable mixin for torchvision models with ImageNet-pretrained weights. + + torchvision fetches pretrained weights into a global ``torch.hub`` cache. + This mixin redirects that cache to ``component_dir()`` so the weights are + stored and gated like any other downloadable component. Subclasses build + their backbone inside :meth:`local_hub` so the pretrained weights are read + from (and written to) the component's own folder. + + .. note:: + The download provides the ImageNet weights used when the model is built + with ``pretrained=True`` (the default). Training a fresh model with + ``pretrained=False`` needs no weights but is still gated as a + download-required component. + """ + + @classmethod + def _weights(cls): + """Return the torchvision weights enum member to download. + + Returns + ------- + torchvision.models.WeightsEnum + The pretrained weights descriptor whose file is fetched. + """ + raise NotImplementedError + + @classmethod + def _checkpoints_dir(cls): + """Return the directory where torch.hub stores downloaded checkpoints.""" + return cls.component_dir() / "checkpoints" + + @classmethod + def local_hub(cls): + """Context manager that points ``torch.hub`` at ``component_dir()``. + + Returns + ------- + contextlib.AbstractContextManager + A context that temporarily sets the torch hub directory to this + component's folder and restores the previous value on exit. + """ + import contextlib + + import torch + + @contextlib.contextmanager + def _ctx(): + old = torch.hub.get_dir() + cls.component_dir().mkdir(parents=True, exist_ok=True) + torch.hub.set_dir(str(cls.component_dir())) + try: + yield + finally: + torch.hub.set_dir(old) + + return _ctx() + + @classmethod + def is_downloaded(cls) -> bool: + """Return whether the pretrained weights file is present locally. + + Returns + ------- + bool + ``True`` when the component's ``checkpoints`` folder exists and is + non-empty. + """ + ckpt = cls._checkpoints_dir() + return ckpt.is_dir() and any(ckpt.iterdir()) + + @classmethod + def download(cls, report: Optional[ProgressReporter] = None) -> None: + """Fetch the pretrained weights into ``component_dir()``. + + Parameters + ---------- + report : ProgressReporter, optional + Callback invoked with an indeterminate progress message. + """ + if report is not None: + report(None, f"Downloading {cls.__name__} weights") + with cls.local_hub(): + cls._weights().get_state_dict(progress=False) diff --git a/DashAI/back/models/efficientnet_b0_image_classifier.py b/DashAI/back/models/efficientnet_b0_image_classifier.py index a6be1bcb2..56517f400 100644 --- a/DashAI/back/models/efficientnet_b0_image_classifier.py +++ b/DashAI/back/models/efficientnet_b0_image_classifier.py @@ -1,13 +1,16 @@ """EfficientNet-B0 image classifier for DashAI.""" from DashAI.back.core.utils import MultilingualString +from DashAI.back.dependencies.downloads.downloadable import TorchvisionDownloadMixin from DashAI.back.models.base_torchvision_image_classifier import ( TorchvisionImageClassifier, TorchvisionImageClassifierSchema, ) -class EfficientNetB0ImageClassifier(TorchvisionImageClassifier): +class EfficientNetB0ImageClassifier( + TorchvisionDownloadMixin, TorchvisionImageClassifier +): """EfficientNet-B0 image classifier (Tan & Le, 2019). Compact baseline of the EfficientNet family, which scales network width, @@ -53,13 +56,21 @@ class EfficientNetB0ImageClassifier(TorchvisionImageClassifier): ) COLOR: str = "#00838F" ICON: str = "Speed" + DOWNLOAD_SIZE_BYTES: int = 21_000_000 + + @classmethod + def _weights(cls): + from torchvision.models import EfficientNet_B0_Weights + + return EfficientNet_B0_Weights.DEFAULT def _build_backbone(self, num_classes: int, pretrained: bool): import torch.nn as nn from torchvision.models import EfficientNet_B0_Weights, efficientnet_b0 weights = EfficientNet_B0_Weights.DEFAULT if pretrained else None - model = efficientnet_b0(weights=weights) + with self.local_hub(): + model = efficientnet_b0(weights=weights) in_features = model.classifier[1].in_features model.classifier = nn.Sequential( nn.Dropout(self.dropout_rate), diff --git a/DashAI/back/models/resnet18_image_classifier.py b/DashAI/back/models/resnet18_image_classifier.py index aed89a626..bcc8fee43 100644 --- a/DashAI/back/models/resnet18_image_classifier.py +++ b/DashAI/back/models/resnet18_image_classifier.py @@ -1,13 +1,14 @@ """ResNet-18 image classifier for DashAI.""" from DashAI.back.core.utils import MultilingualString +from DashAI.back.dependencies.downloads.downloadable import TorchvisionDownloadMixin from DashAI.back.models.base_torchvision_image_classifier import ( TorchvisionImageClassifier, TorchvisionImageClassifierSchema, ) -class ResNet18ImageClassifier(TorchvisionImageClassifier): +class ResNet18ImageClassifier(TorchvisionDownloadMixin, TorchvisionImageClassifier): """ResNet-18 image classifier (He et al., 2015). 18-layer residual network with skip connections that solve the vanishing @@ -52,13 +53,21 @@ class ResNet18ImageClassifier(TorchvisionImageClassifier): ) COLOR: str = "#2E7D32" ICON: str = "AccountTree" + DOWNLOAD_SIZE_BYTES: int = 47_000_000 + + @classmethod + def _weights(cls): + from torchvision.models import ResNet18_Weights + + return ResNet18_Weights.DEFAULT def _build_backbone(self, num_classes: int, pretrained: bool): import torch.nn as nn from torchvision.models import ResNet18_Weights, resnet18 weights = ResNet18_Weights.DEFAULT if pretrained else None - model = resnet18(weights=weights) + with self.local_hub(): + model = resnet18(weights=weights) in_features = model.fc.in_features model.fc = nn.Sequential( nn.Dropout(self.dropout_rate), diff --git a/DashAI/back/models/resnet50_image_classifier.py b/DashAI/back/models/resnet50_image_classifier.py index 92f07a662..c556a3960 100644 --- a/DashAI/back/models/resnet50_image_classifier.py +++ b/DashAI/back/models/resnet50_image_classifier.py @@ -1,13 +1,14 @@ """ResNet-50 image classifier for DashAI.""" from DashAI.back.core.utils import MultilingualString +from DashAI.back.dependencies.downloads.downloadable import TorchvisionDownloadMixin from DashAI.back.models.base_torchvision_image_classifier import ( TorchvisionImageClassifier, TorchvisionImageClassifierSchema, ) -class ResNet50ImageClassifier(TorchvisionImageClassifier): +class ResNet50ImageClassifier(TorchvisionDownloadMixin, TorchvisionImageClassifier): """ResNet-50 image classifier (He et al., 2015). 50-layer residual network using bottleneck blocks. Deeper and more @@ -54,13 +55,21 @@ class ResNet50ImageClassifier(TorchvisionImageClassifier): ) COLOR: str = "#1B5E20" ICON: str = "AccountTree" + DOWNLOAD_SIZE_BYTES: int = 100_000_000 + + @classmethod + def _weights(cls): + from torchvision.models import ResNet50_Weights + + return ResNet50_Weights.DEFAULT def _build_backbone(self, num_classes: int, pretrained: bool): import torch.nn as nn from torchvision.models import ResNet50_Weights, resnet50 weights = ResNet50_Weights.DEFAULT if pretrained else None - model = resnet50(weights=weights) + with self.local_hub(): + model = resnet50(weights=weights) in_features = model.fc.in_features model.fc = nn.Sequential( nn.Dropout(self.dropout_rate), diff --git a/tests/back/models/test_image_classification_downloadable.py b/tests/back/models/test_image_classification_downloadable.py new file mode 100644 index 000000000..df315d2c1 --- /dev/null +++ b/tests/back/models/test_image_classification_downloadable.py @@ -0,0 +1,54 @@ +"""Tests that torchvision image classifiers expose download metadata.""" + +import pytest +from kink import di + +from DashAI.back.dependencies.downloads.downloadable import TorchvisionDownloadMixin +from DashAI.back.models.efficientnet_b0_image_classifier import ( + EfficientNetB0ImageClassifier, +) +from DashAI.back.models.resnet18_image_classifier import ResNet18ImageClassifier +from DashAI.back.models.resnet50_image_classifier import ResNet50ImageClassifier + +_CASES = [ + ResNet18ImageClassifier, + ResNet50ImageClassifier, + EfficientNetB0ImageClassifier, +] + + +@pytest.fixture +def component_root(tmp_path): + """Inject a temporary COMPONENT_PATH into the kink DI container.""" + sentinel = object() + old = di["config"] if "config" in di else sentinel # noqa: SIM401 + di["config"] = {"COMPONENT_PATH": str(tmp_path)} + yield tmp_path + if old is sentinel: + del di["config"] + else: + di["config"] = old + + +@pytest.mark.parametrize("model_cls", _CASES) +def test_image_classifier_is_downloadable(model_cls): + assert issubclass(model_cls, TorchvisionDownloadMixin) + assert model_cls.REQUIRES_DOWNLOAD is True + assert model_cls.DOWNLOAD_SIZE_BYTES is not None + + +@pytest.mark.parametrize("model_cls", _CASES) +def test_image_classifier_metadata_flags_download(model_cls): + meta = model_cls.get_metadata() + assert meta["requires_download"] is True + assert meta["download_size_bytes"] == model_cls.DOWNLOAD_SIZE_BYTES + + +@pytest.mark.parametrize("model_cls", _CASES) +def test_is_downloaded_uses_checkpoints_dir(model_cls, component_root): + assert model_cls.is_downloaded() is False + + ckpt = component_root / model_cls.__name__ / "checkpoints" + ckpt.mkdir(parents=True) + (ckpt / "weights.pth").write_bytes(b"w") + assert model_cls.is_downloaded() is True From 3886d58f2ae582894f38f206fe309bbb1dcbfac1 Mon Sep 17 00:00:00 2001 From: Irozuku Date: Thu, 2 Jul 2026 14:33:33 -0400 Subject: [PATCH 42/89] fix: reconcile model download gate against the filesystem --- .../back/api/api_v1/endpoints/generative_session.py | 11 ++++++----- DashAI/back/api/api_v1/endpoints/runs.py | 11 ++++++----- .../back/api/test_generative_session_download_gate.py | 5 +++++ tests/back/api/test_run_download_gate.py | 5 +++++ 4 files changed, 22 insertions(+), 10 deletions(-) diff --git a/DashAI/back/api/api_v1/endpoints/generative_session.py b/DashAI/back/api/api_v1/endpoints/generative_session.py index d4ec8d305..503b8c72d 100644 --- a/DashAI/back/api/api_v1/endpoints/generative_session.py +++ b/DashAI/back/api/api_v1/endpoints/generative_session.py @@ -47,11 +47,12 @@ async def upload_generative_session( detail=f"Model {params.model_name} is not registered.", ) from e - # Guard: model requires download but has not been downloaded -> 409 - entry = component_registry[params.model_name] - if getattr(entry["class"], "REQUIRES_DOWNLOAD", False) and not entry.get( - "downloaded", False - ): + # Guard: model requires download but has not been downloaded -> 409. + # Reconcile against the filesystem so a model downloaded after startup + # (in the worker process) is recognised without an API restart. + if getattr( + model_class, "REQUIRES_DOWNLOAD", False + ) and not component_registry.refresh_download_status(params.model_name): raise HTTPException( status_code=status.HTTP_409_CONFLICT, detail=( diff --git a/DashAI/back/api/api_v1/endpoints/runs.py b/DashAI/back/api/api_v1/endpoints/runs.py index 0c4dd2119..cdc61958e 100644 --- a/DashAI/back/api/api_v1/endpoints/runs.py +++ b/DashAI/back/api/api_v1/endpoints/runs.py @@ -325,17 +325,18 @@ async def upload_run( status_code=status.HTTP_404_NOT_FOUND, detail="Model session not found", ) - # REQUIRES_DOWNLOAD is read from the class (static contract); - # "downloaded" is read from the registry dict (runtime state, Task 4). + # REQUIRES_DOWNLOAD is the static contract; the download state is + # reconciled against the filesystem so a model downloaded after + # startup (in the worker process) is recognised without a restart. if params.model_name not in component_registry: raise HTTPException( status_code=status.HTTP_422_UNPROCESSABLE_ENTITY, detail=f"Unknown model '{params.model_name}'", ) entry = component_registry[params.model_name] - if getattr(entry["class"], "REQUIRES_DOWNLOAD", False) and not entry.get( - "downloaded", False - ): + if getattr( + entry["class"], "REQUIRES_DOWNLOAD", False + ) and not component_registry.refresh_download_status(params.model_name): raise HTTPException( status_code=status.HTTP_409_CONFLICT, detail=( diff --git a/tests/back/api/test_generative_session_download_gate.py b/tests/back/api/test_generative_session_download_gate.py index 6881df992..0ceaf1a03 100644 --- a/tests/back/api/test_generative_session_download_gate.py +++ b/tests/back/api/test_generative_session_download_gate.py @@ -36,6 +36,11 @@ def __getitem__(self, name): def get_components_by_types(self, select=None, ignore=None): return self._real.get_components_by_types(select=select, ignore=ignore) + def refresh_download_status(self, name): + if name == "FakeDownloadableGenerativeModel": + return _FakeDownloadableGenerativeModel.is_downloaded() + return self._real.refresh_download_status(name) + def __contains__(self, name): return name == "FakeDownloadableGenerativeModel" or name in self._real diff --git a/tests/back/api/test_run_download_gate.py b/tests/back/api/test_run_download_gate.py index e8e89ab9b..52af92061 100644 --- a/tests/back/api/test_run_download_gate.py +++ b/tests/back/api/test_run_download_gate.py @@ -40,6 +40,11 @@ def __getitem__(self, name): def get_components_by_types(self, select=None, ignore=None): return self._real.get_components_by_types(select=select, ignore=ignore) + def refresh_download_status(self, name): + if name == "FakeDownloadableModel": + return _FakeDownloadableModel.is_downloaded() + return self._real.refresh_download_status(name) + def __contains__(self, name): return name == "FakeDownloadableModel" or name in self._real From 592385e897099562df23d186d82acc82810515ab Mon Sep 17 00:00:00 2001 From: Irozuku Date: Thu, 2 Jul 2026 14:38:30 -0400 Subject: [PATCH 43/89] feat: split Stable Diffusion 2 into downloadable per-checkpoint models --- DashAI/back/initial_components.py | 10 +- .../hugging_face/stable_diffusion_v2_model.py | 169 +++++++++++------- tests/back/api/conftest.py | 17 ++ .../test_generative_session_download_gate.py | 8 +- tests/back/api/test_process_api.py | 2 +- tests/back/api/test_session_api.py | 10 +- .../test_stable_diffusion2_downloadable.py | 58 ++++++ 7 files changed, 201 insertions(+), 73 deletions(-) create mode 100644 tests/back/models/test_stable_diffusion2_downloadable.py diff --git a/DashAI/back/initial_components.py b/DashAI/back/initial_components.py index 056769619..689d007c9 100644 --- a/DashAI/back/initial_components.py +++ b/DashAI/back/initial_components.py @@ -210,7 +210,10 @@ StableDiffusionXLV1ControlNet, ) from DashAI.back.models.hugging_face.stable_diffusion_v2_model import ( - StableDiffusionV2Model, + StableDiffusion2, + StableDiffusion2_512, + StableDiffusion21, + StableDiffusion21_512, ) from DashAI.back.models.hugging_face.stable_diffusion_v3_model import ( StableDiffusionV3Model, @@ -398,7 +401,10 @@ def get_initial_components(): SGDClassifier, SmolLM2_360MInstruct, SmolLM2_17BInstruct, - StableDiffusionV2Model, + StableDiffusion2, + StableDiffusion2_512, + StableDiffusion21, + StableDiffusion21_512, StableDiffusionV3Model, StableDiffusionXLModel, StableDiffusionXLV1ControlNet, diff --git a/DashAI/back/models/hugging_face/stable_diffusion_v2_model.py b/DashAI/back/models/hugging_face/stable_diffusion_v2_model.py index 8354499f2..7d9230f46 100644 --- a/DashAI/back/models/hugging_face/stable_diffusion_v2_model.py +++ b/DashAI/back/models/hugging_face/stable_diffusion_v2_model.py @@ -9,6 +9,9 @@ ) from DashAI.back.core.schema_fields.base_schema import BaseSchema from DashAI.back.core.utils import MultilingualString +from DashAI.back.dependencies.downloads.downloadable import ( + HFPretrainedDownloadMixin, +) from DashAI.back.models.text_to_image_generation_model import ( TextToImageGenerationTaskModel, ) @@ -26,65 +29,6 @@ class StableDiffusionSchema(BaseSchema): ``StableDiffusionV2Model``. """ - model_name: schema_field( - enum_field( - enum=[ - "sd2-community/stable-diffusion-2", - "sd2-community/stable-diffusion-2-base", - "sd2-community/stable-diffusion-2-1", - "sd2-community/stable-diffusion-2-1-base", - ] - ), - placeholder="sd2-community/stable-diffusion-2", - description=MultilingualString( - en=( - "The specific Stable Diffusion 2.x checkpoint to load. " - "The '-base' variants are trained at 512x512 px and are faster; " - "the nonbase variants target 768x768 px and produce sharper detail. " - "The '2-1' variants are fine-tuned further " - "and generally outperform '2'." - ), - es=( - "El checkpoint específico de Stable Diffusion 2.x a cargar. " - "Las variantes '-base' se entrenan a 512x512 px y son más rápidas; " - "las variantes sin '-base' apuntan a 768x768 px " - "y producen mayor detalle. " - "Las variantes '2-1' están más ajustadas " - "y generalmente superan a '2'." - ), - pt=( - "O checkpoint específico do Stable Diffusion 2.x a carregar. " - "As variantes '-base' são treinadas a 512x512 px e são mais rápidas; " - "as variantes sem '-base' visam 768x768 px " - "e produzem maior detalhe. " - "As variantes '2-1' são mais ajustadas " - "e geralmente superam a '2'." - ), - de=( - "Der zu ladende spezifische Stable Diffusion 2.x-Checkpoint. " - "Die '-base'-Varianten werden bei 512x512 px trainiert und sind " - "schneller; " - "die Nicht-base-Varianten zielen auf 768x768 px ab " - "und liefern schärfere Details. " - "Die '2-1'-Varianten sind weiter feinabgestimmt " - "und übertreffen '2' in der Regel." - ), - zh=( - "要加载的 Stable Diffusion 2.x 检查点。" - "'-base' 变体在 512x512 像素下训练，速度更快；" - "非 base 变体目标分辨率为 768x768 像素，细节更清晰。" - "'2-1' 变体经过进一步微调，通常优于 '2'。" - ), - ), - alias=MultilingualString( - en="Model name", - es="Nombre del modelo", - pt="Nome do modelo", - de="Modellname", - zh="模型名称", - ), - ) # type: ignore - negative_prompt: Optional[ schema_field( string_field(), @@ -403,7 +347,9 @@ class StableDiffusionSchema(BaseSchema): ) # type: ignore -class StableDiffusionV2Model(TextToImageGenerationTaskModel): +class StableDiffusion2GenerationModel( + HFPretrainedDownloadMixin, TextToImageGenerationTaskModel +): """Latent diffusion model for high resolution text-to-image generation. Wraps the Stable Diffusion 2.x family of checkpoints released by @@ -431,6 +377,7 @@ class StableDiffusionV2Model(TextToImageGenerationTaskModel): """ SCHEMA = StableDiffusionSchema + MODEL_NAME: str = "" COLOR: str = "#1565c0" DISPLAY_NAME: str = MultilingualString( en="Stable Diffusion V2", @@ -538,7 +485,7 @@ def __init__(self, **kwargs): self.device = ( f"cuda:{DEVICE_TO_IDX.get(kwargs.get('device'))}" if use_gpu else "cpu" ) - self.model_name = kwargs.get("model_name", "sd2-community/stable-diffusion-2") + self.model_name = self._pretrained_source(None) self.model = DiffusionPipeline.from_pretrained( self.model_name, @@ -589,3 +536,103 @@ def generate(self, input: str) -> List[Any]: output = self.model(**params) return output.images + + +class StableDiffusion2(StableDiffusion2GenerationModel): + """768px Stable Diffusion 2 checkpoint. + + Downloads its checkpoint into the component's own download folder. + """ + + MODEL_NAME: str = "sd2-community/stable-diffusion-2" + # Full fp32 diffusers pipeline (text encoder + U-Net + VAE) is ~5 GB. + DOWNLOAD_SIZE_BYTES: int = 5_200_000_000 + DISPLAY_NAME = MultilingualString( + en="Stable Diffusion 2", + es="Stable Diffusion 2", + pt="Stable Diffusion 2", + de="Stable Diffusion 2", + zh="Stable Diffusion 2", + ) + DESCRIPTION = MultilingualString( + en="768px Stable Diffusion 2 checkpoint.", + es="768px Stable Diffusion 2 checkpoint.", + pt="768px Stable Diffusion 2 checkpoint.", + de="768px Stable Diffusion 2 checkpoint.", + zh="768px Stable Diffusion 2 checkpoint.", + ) + + +class StableDiffusion2_512(StableDiffusion2GenerationModel): # noqa: N801 + """512px base Stable Diffusion 2 checkpoint (faster). + + Downloads its checkpoint into the component's own download folder. + """ + + MODEL_NAME: str = "sd2-community/stable-diffusion-2-base" + # Full fp32 diffusers pipeline (text encoder + U-Net + VAE) is ~5 GB. + DOWNLOAD_SIZE_BYTES: int = 5_200_000_000 + DISPLAY_NAME = MultilingualString( + en="Stable Diffusion 2 (512px)", + es="Stable Diffusion 2 (512px)", + pt="Stable Diffusion 2 (512px)", + de="Stable Diffusion 2 (512px)", + zh="Stable Diffusion 2 (512px)", + ) + DESCRIPTION = MultilingualString( + en="512px base Stable Diffusion 2 checkpoint (faster).", + es="512px base Stable Diffusion 2 checkpoint (faster).", + pt="512px base Stable Diffusion 2 checkpoint (faster).", + de="512px base Stable Diffusion 2 checkpoint (faster).", + zh="512px base Stable Diffusion 2 checkpoint (faster).", + ) + + +class StableDiffusion21(StableDiffusion2GenerationModel): + """768px Stable Diffusion 2.1 checkpoint (further fine-tuned). + + Downloads its checkpoint into the component's own download folder. + """ + + MODEL_NAME: str = "sd2-community/stable-diffusion-2-1" + # Full fp32 diffusers pipeline (text encoder + U-Net + VAE) is ~5 GB. + DOWNLOAD_SIZE_BYTES: int = 5_200_000_000 + DISPLAY_NAME = MultilingualString( + en="Stable Diffusion 2.1", + es="Stable Diffusion 2.1", + pt="Stable Diffusion 2.1", + de="Stable Diffusion 2.1", + zh="Stable Diffusion 2.1", + ) + DESCRIPTION = MultilingualString( + en="768px Stable Diffusion 2.1 checkpoint (further fine-tuned).", + es="768px Stable Diffusion 2.1 checkpoint (further fine-tuned).", + pt="768px Stable Diffusion 2.1 checkpoint (further fine-tuned).", + de="768px Stable Diffusion 2.1 checkpoint (further fine-tuned).", + zh="768px Stable Diffusion 2.1 checkpoint (further fine-tuned).", + ) + + +class StableDiffusion21_512(StableDiffusion2GenerationModel): # noqa: N801 + """512px base Stable Diffusion 2.1 checkpoint. + + Downloads its checkpoint into the component's own download folder. + """ + + MODEL_NAME: str = "sd2-community/stable-diffusion-2-1-base" + # Full fp32 diffusers pipeline (text encoder + U-Net + VAE) is ~5 GB. + DOWNLOAD_SIZE_BYTES: int = 5_200_000_000 + DISPLAY_NAME = MultilingualString( + en="Stable Diffusion 2.1 (512px)", + es="Stable Diffusion 2.1 (512px)", + pt="Stable Diffusion 2.1 (512px)", + de="Stable Diffusion 2.1 (512px)", + zh="Stable Diffusion 2.1 (512px)", + ) + DESCRIPTION = MultilingualString( + en="512px base Stable Diffusion 2.1 checkpoint.", + es="512px base Stable Diffusion 2.1 checkpoint.", + pt="512px base Stable Diffusion 2.1 checkpoint.", + de="512px base Stable Diffusion 2.1 checkpoint.", + zh="512px base Stable Diffusion 2.1 checkpoint.", + ) diff --git a/tests/back/api/conftest.py b/tests/back/api/conftest.py index f9dc1ed9b..074e299f1 100644 --- a/tests/back/api/conftest.py +++ b/tests/back/api/conftest.py @@ -46,6 +46,23 @@ def client(test_path: Path): remove_dir_with_retry(app.container._services["config"]["LOCAL_PATH"]) +@pytest.fixture(scope="module", autouse=True) +def _mark_stable_diffusion2_downloaded(client): + """Make ``StableDiffusion2`` appear downloaded for session/run tests. + + All generative models now require a download, so the download gate would + reject session/run creation. Creating the component's repo folder lets the + filesystem-reconciled gate treat it as available without fetching weights. + """ + config = client.app.container._services["config"] + repo_dir = ( + Path(config["COMPONENT_PATH"]) / "StableDiffusion2" / "stable-diffusion-2" + ) + repo_dir.mkdir(parents=True, exist_ok=True) + (repo_dir / "config.json").write_text("{}", encoding="utf-8") + return + + @pytest.fixture(name="dataset_1", scope="module") def create_dataset_1(client) -> Dataset: """Create testing dataset 1 using job system.""" diff --git a/tests/back/api/test_generative_session_download_gate.py b/tests/back/api/test_generative_session_download_gate.py index 0ceaf1a03..c454ca785 100644 --- a/tests/back/api/test_generative_session_download_gate.py +++ b/tests/back/api/test_generative_session_download_gate.py @@ -96,7 +96,7 @@ def _create_sd_session(client, name): return client.post( "/api/v1/generative-session/", json={ - "model_name": "StableDiffusionV2Model", + "model_name": "StableDiffusion2", "task_name": "TextToImageGenerationTask", "parameters": _SD_PARAMS, "name": name, @@ -150,10 +150,10 @@ def test_change_session_model_valid_returns_200(client): resp = client.patch( f"/api/v1/generative-session/{session_id}", - params={"model_name": "StableDiffusionV2Model"}, + params={"model_name": "StableDiffusion2"}, ) assert resp.status_code == 200 - assert resp.json()["model_name"] == "StableDiffusionV2Model" + assert resp.json()["model_name"] == "StableDiffusion2" client.delete(f"/api/v1/generative-session/{session_id}") @@ -185,7 +185,7 @@ def test_switch_model_resets_params_and_records_history(client): ] assert { "parameter": "model", - "oldValue": "StableDiffusionV2Model", + "oldValue": "StableDiffusion2", "newValue": "Qwen25_15BInstruct", } in model_changes diff --git a/tests/back/api/test_process_api.py b/tests/back/api/test_process_api.py index 7a65cbea5..35b1fcfd7 100644 --- a/tests/back/api/test_process_api.py +++ b/tests/back/api/test_process_api.py @@ -6,7 +6,7 @@ def session(client: TestClient): """Create a valid session for process tests.""" params = { - "model_name": "StableDiffusionV2Model", + "model_name": "StableDiffusion2", "task_name": "TextToImageGenerationTask", "parameters": { "num_inference_steps": 1, diff --git a/tests/back/api/test_session_api.py b/tests/back/api/test_session_api.py index ece8d95cc..bb8e2b411 100644 --- a/tests/back/api/test_session_api.py +++ b/tests/back/api/test_session_api.py @@ -6,7 +6,7 @@ def create_session_1(client: TestClient): """Create testing session 1 using job system.""" params = { - "model_name": "StableDiffusionV2Model", + "model_name": "StableDiffusion2", "task_name": "TextToImageGenerationTask", "parameters": { "num_inference_steps": 1, @@ -56,7 +56,7 @@ def create_session_2(client: TestClient): def create_session_3(client: TestClient): """Create testing session 3 using a non-download-required model.""" params = { - "model_name": "StableDiffusionV2Model", + "model_name": "StableDiffusion2", "task_name": "TextToImageGenerationTask", "parameters": { "num_inference_steps": 1, @@ -85,7 +85,7 @@ def create_session_3(client: TestClient): def create_session_4(client: TestClient): """Create testing session 4 with an invalid task (valid model).""" params = { - "model_name": "StableDiffusionV2Model", + "model_name": "StableDiffusion2", "task_name": "SomeTask", "parameters": { "num_inference_steps": 1, @@ -116,7 +116,7 @@ def test_create_session(response_1): data = response_1.json() assert data["id"] is not None, "Session ID is missing" assert data["name"] == "session_1", "Session name does not match" - assert data["model_name"] == "StableDiffusionV2Model", "Model name does not match" + assert data["model_name"] == "StableDiffusion2", "Model name does not match" assert data["task_name"] == "TextToImageGenerationTask", "Task name does not match" @@ -135,7 +135,7 @@ def test_get_session_by_id(client: TestClient, response_1): data = response.json() assert data["id"] == session_id, "Retrieved session ID does not match" assert data["name"] == "session_1", "Session name does not match" - assert data["model_name"] == "StableDiffusionV2Model", "Model name does not match" + assert data["model_name"] == "StableDiffusion2", "Model name does not match" assert data["task_name"] == "TextToImageGenerationTask", "Task name does not match" diff --git a/tests/back/models/test_stable_diffusion2_downloadable.py b/tests/back/models/test_stable_diffusion2_downloadable.py new file mode 100644 index 000000000..b86b179c6 --- /dev/null +++ b/tests/back/models/test_stable_diffusion2_downloadable.py @@ -0,0 +1,58 @@ +"""Tests that Stable Diffusion 2 per-checkpoint models expose download metadata.""" + +import pytest +from kink import di + +from DashAI.back.dependencies.downloads.downloadable import HFPretrainedDownloadMixin +from DashAI.back.models.hugging_face.stable_diffusion_v2_model import ( + StableDiffusion2, + StableDiffusion2_512, + StableDiffusion21, + StableDiffusion21_512, +) + +_CASES = [ + (StableDiffusion2, "sd2-community/stable-diffusion-2"), + (StableDiffusion2_512, "sd2-community/stable-diffusion-2-base"), + (StableDiffusion21, "sd2-community/stable-diffusion-2-1"), + (StableDiffusion21_512, "sd2-community/stable-diffusion-2-1-base"), +] + + +@pytest.fixture +def component_root(tmp_path): + """Inject a temporary COMPONENT_PATH into the kink DI container.""" + sentinel = object() + old = di["config"] if "config" in di else sentinel # noqa: SIM401 + di["config"] = {"COMPONENT_PATH": str(tmp_path)} + yield tmp_path + if old is sentinel: + del di["config"] + else: + di["config"] = old + + +@pytest.mark.parametrize(("model_cls", "repo_id"), _CASES) +def test_sd2_is_downloadable(model_cls, repo_id): + assert issubclass(model_cls, HFPretrainedDownloadMixin) + assert model_cls.REQUIRES_DOWNLOAD is True + assert model_cls.DOWNLOAD_SIZE_BYTES is not None + assert model_cls.hf_repos() == [(repo_id, "model")] + + +@pytest.mark.parametrize(("model_cls", "repo_id"), _CASES) +def test_sd2_metadata_flags_download(model_cls, repo_id): + meta = model_cls.get_metadata() + assert meta["requires_download"] is True + assert meta["download_size_bytes"] == model_cls.DOWNLOAD_SIZE_BYTES + + +@pytest.mark.parametrize(("model_cls", "repo_id"), _CASES) +def test_sd2_is_downloaded_uses_component_dir(model_cls, repo_id, component_root): + assert model_cls.is_downloaded() is False + + leaf = repo_id.split("/")[-1] + repo_dir = component_root / model_cls.__name__ / leaf + repo_dir.mkdir(parents=True) + (repo_dir / "config.json").write_text("{}") + assert model_cls.is_downloaded() is True From 1ac62679b465e5cd7e370dcf189b1c48bf85730f Mon Sep 17 00:00:00 2001 From: Irozuku Date: Thu, 2 Jul 2026 14:41:40 -0400 Subject: [PATCH 44/89] feat: split SDXL and make SDXL-Turbo downloadable per checkpoint --- DashAI/back/initial_components.py | 6 +- .../models/hugging_face/sdxl_turbo_model.py | 10 +- .../hugging_face/stable_diffusion_xl_model.py | 116 +++++++++--------- tests/back/models/test_sdxl_downloadable.py | 31 +++++ 4 files changed, 101 insertions(+), 62 deletions(-) create mode 100644 tests/back/models/test_sdxl_downloadable.py diff --git a/DashAI/back/initial_components.py b/DashAI/back/initial_components.py index 689d007c9..ff31ea330 100644 --- a/DashAI/back/initial_components.py +++ b/DashAI/back/initial_components.py @@ -219,7 +219,8 @@ StableDiffusionV3Model, ) from DashAI.back.models.hugging_face.stable_diffusion_xl_model import ( - StableDiffusionXLModel, + RealVisXLV4, + StableDiffusionXL, ) from DashAI.back.models.hugging_face.t5_small_transformer import T5SmallTransformer from DashAI.back.models.hugging_face.tongyi_z_image_model import TongyiZImageModel @@ -406,7 +407,8 @@ def get_initial_components(): StableDiffusion21, StableDiffusion21_512, StableDiffusionV3Model, - StableDiffusionXLModel, + StableDiffusionXL, + RealVisXLV4, StableDiffusionXLV1ControlNet, SVC, SVR, diff --git a/DashAI/back/models/hugging_face/sdxl_turbo_model.py b/DashAI/back/models/hugging_face/sdxl_turbo_model.py index 918d90aed..6d7a16f52 100644 --- a/DashAI/back/models/hugging_face/sdxl_turbo_model.py +++ b/DashAI/back/models/hugging_face/sdxl_turbo_model.py @@ -8,6 +8,9 @@ ) from DashAI.back.core.schema_fields.base_schema import BaseSchema from DashAI.back.core.utils import MultilingualString +from DashAI.back.dependencies.downloads.downloadable import ( + HFPretrainedDownloadMixin, +) from DashAI.back.models.text_to_image_generation_model import ( TextToImageGenerationTaskModel, ) @@ -299,7 +302,7 @@ class SDXLTurboSchema(BaseSchema): ) # type: ignore -class SDXLTurboModel(TextToImageGenerationTaskModel): +class SDXLTurboModel(HFPretrainedDownloadMixin, TextToImageGenerationTaskModel): """Distilled SDXL model for near real time text-to-image generation. Wraps ``stabilityai/sdxl-turbo``, a version of Stable Diffusion XL @@ -323,6 +326,9 @@ class SDXLTurboModel(TextToImageGenerationTaskModel): """ SCHEMA = SDXLTurboSchema + MODEL_NAME: str = "stabilityai/sdxl-turbo" + # SDXL-Turbo diffusers pipeline is ~7 GB. + DOWNLOAD_SIZE_BYTES: int = 7_000_000_000 COLOR: str = "#b71c1c" DISPLAY_NAME: str = MultilingualString( en="SDXL Turbo", @@ -423,7 +429,7 @@ def __init__(self, **kwargs): ) self.model = StableDiffusionXLPipeline.from_pretrained( - "stabilityai/sdxl-turbo", + self._pretrained_source(None), torch_dtype=torch.float16 if use_gpu else torch.float32, variant="fp16" if use_gpu else None, ).to(self.device) diff --git a/DashAI/back/models/hugging_face/stable_diffusion_xl_model.py b/DashAI/back/models/hugging_face/stable_diffusion_xl_model.py index 40aec1b3b..6a7944911 100644 --- a/DashAI/back/models/hugging_face/stable_diffusion_xl_model.py +++ b/DashAI/back/models/hugging_face/stable_diffusion_xl_model.py @@ -9,6 +9,9 @@ ) from DashAI.back.core.schema_fields.base_schema import BaseSchema from DashAI.back.core.utils import MultilingualString +from DashAI.back.dependencies.downloads.downloadable import ( + HFPretrainedDownloadMixin, +) from DashAI.back.models.text_to_image_generation_model import ( TextToImageGenerationTaskModel, ) @@ -26,60 +29,6 @@ class StableDiffusionXLSchema(BaseSchema): ``StableDiffusionXLModel``. """ - model_name: schema_field( - enum_field( - enum=[ - "stabilityai/stable-diffusion-xl-base-1.0", - "SG161222/RealVisXL_V4.0", - ] - ), - placeholder="stabilityai/stable-diffusion-xl-base-1.0", - description=MultilingualString( - en=( - "The Stable Diffusion XL checkpoint to load. " - "'stable-diffusion-xl-base-1.0' is the official base model trained " - "at 1024x1024 px for high-quality photorealistic generation. " - "'RealVisXL_V4.0' is a popular community fine-tune of SDXL " - "optimized for realistic portraits and photography." - ), - es=( - "El checkpoint Stable Diffusion XL a cargar. " - "'stable-diffusion-xl-base-1.0' es el modelo base oficial entrenado " - "a 1024x1024 px para generación fotorrealista de alta calidad. " - "'RealVisXL_V4.0' es un popular fine-tune comunitario de SDXL " - "optimizado para retratos realistas y fotografía." - ), - pt=( - "O checkpoint Stable Diffusion XL a carregar. " - "'stable-diffusion-xl-base-1.0' é o modelo base oficial treinado " - "a 1024x1024 px para geração fotorrealista de alta qualidade. " - "'RealVisXL_V4.0' é um popular fine-tune comunitário do SDXL " - "otimizado para retratos realistas e fotografia." - ), - de=( - "Der zu ladende Stable Diffusion XL-Checkpoint. " - "'stable-diffusion-xl-base-1.0' ist das offizielle Basismodell, " - "bei 1024x1024 px für hochwertige fotorealistische Generierung " - "trainiert. " - "'RealVisXL_V4.0' ist ein beliebter Community-Fine-Tune von SDXL, " - "optimiert für realistische Porträts und Fotografie." - ), - zh=( - "要加载的 Stable Diffusion XL 检查点。" - "'stable-diffusion-xl-base-1.0' 是官方基础模型，" - "在 1024x1024px 下训练，用于高质量写实图像生成。" - "'RealVisXL_V4.0' 是针对写实人像和摄影优化的热门社区微调版本。" - ), - ), - alias=MultilingualString( - en="Model name", - es="Nombre del modelo", - pt="Nome do modelo", - de="Modellname", - zh="模型名称", - ), - ) # type: ignore - negative_prompt: Optional[ schema_field( string_field(), @@ -393,7 +342,9 @@ class StableDiffusionXLSchema(BaseSchema): ) # type: ignore -class StableDiffusionXLModel(TextToImageGenerationTaskModel): +class StableDiffusionXLGenerationModel( + HFPretrainedDownloadMixin, TextToImageGenerationTaskModel +): """Latent diffusion model for high-resolution 1024 px text-to-image generation. Wraps Stable Diffusion XL (SDXL) checkpoints. SDXL scales the standard @@ -415,6 +366,7 @@ class StableDiffusionXLModel(TextToImageGenerationTaskModel): """ SCHEMA = StableDiffusionXLSchema + MODEL_NAME: str = "" COLOR: str = "#0d47a1" DISPLAY_NAME: str = MultilingualString( en="Stable Diffusion XL", @@ -483,9 +435,7 @@ def __init__(self, **kwargs): self.device = ( f"cuda:{DEVICE_TO_IDX.get(kwargs.get('device'))}" if use_gpu else "cpu" ) - self.model_name = kwargs.get( - "model_name", "stabilityai/stable-diffusion-xl-base-1.0" - ) + self.model_name = self._pretrained_source(None) self.model = StableDiffusionXLPipeline.from_pretrained( self.model_name, @@ -534,3 +484,53 @@ def generate(self, input: str) -> List[Any]: output = self.model(**params) return output.images + + +class StableDiffusionXL(StableDiffusionXLGenerationModel): + """Stable Diffusion XL base 1.0 checkpoint. + + Downloads its checkpoint into the component's own download folder. + """ + + MODEL_NAME: str = "stabilityai/stable-diffusion-xl-base-1.0" + # SDXL diffusers pipeline (base + refiner-less) is ~7 GB. + DOWNLOAD_SIZE_BYTES: int = 7_000_000_000 + DISPLAY_NAME = MultilingualString( + en="Stable Diffusion XL", + es="Stable Diffusion XL", + pt="Stable Diffusion XL", + de="Stable Diffusion XL", + zh="Stable Diffusion XL", + ) + DESCRIPTION = MultilingualString( + en="Stable Diffusion XL base 1.0 checkpoint.", + es="Stable Diffusion XL base 1.0 checkpoint.", + pt="Stable Diffusion XL base 1.0 checkpoint.", + de="Stable Diffusion XL base 1.0 checkpoint.", + zh="Stable Diffusion XL base 1.0 checkpoint.", + ) + + +class RealVisXLV4(StableDiffusionXLGenerationModel): + """RealVisXL V4.0 photorealistic SDXL checkpoint. + + Downloads its checkpoint into the component's own download folder. + """ + + MODEL_NAME: str = "SG161222/RealVisXL_V4.0" + # SDXL diffusers pipeline (base + refiner-less) is ~7 GB. + DOWNLOAD_SIZE_BYTES: int = 7_000_000_000 + DISPLAY_NAME = MultilingualString( + en="RealVisXL V4.0", + es="RealVisXL V4.0", + pt="RealVisXL V4.0", + de="RealVisXL V4.0", + zh="RealVisXL V4.0", + ) + DESCRIPTION = MultilingualString( + en="RealVisXL V4.0 photorealistic SDXL checkpoint.", + es="RealVisXL V4.0 photorealistic SDXL checkpoint.", + pt="RealVisXL V4.0 photorealistic SDXL checkpoint.", + de="RealVisXL V4.0 photorealistic SDXL checkpoint.", + zh="RealVisXL V4.0 photorealistic SDXL checkpoint.", + ) diff --git a/tests/back/models/test_sdxl_downloadable.py b/tests/back/models/test_sdxl_downloadable.py new file mode 100644 index 000000000..3239fd924 --- /dev/null +++ b/tests/back/models/test_sdxl_downloadable.py @@ -0,0 +1,31 @@ +"""Tests that SDXL / SDXL-Turbo models expose download metadata.""" + +import pytest + +from DashAI.back.dependencies.downloads.downloadable import HFPretrainedDownloadMixin +from DashAI.back.models.hugging_face.sdxl_turbo_model import SDXLTurboModel +from DashAI.back.models.hugging_face.stable_diffusion_xl_model import ( + RealVisXLV4, + StableDiffusionXL, +) + +_CASES = [ + (StableDiffusionXL, "stabilityai/stable-diffusion-xl-base-1.0"), + (RealVisXLV4, "SG161222/RealVisXL_V4.0"), + (SDXLTurboModel, "stabilityai/sdxl-turbo"), +] + + +@pytest.mark.parametrize(("model_cls", "repo_id"), _CASES) +def test_sdxl_is_downloadable(model_cls, repo_id): + assert issubclass(model_cls, HFPretrainedDownloadMixin) + assert model_cls.REQUIRES_DOWNLOAD is True + assert model_cls.DOWNLOAD_SIZE_BYTES is not None + assert model_cls.hf_repos() == [(repo_id, "model")] + + +@pytest.mark.parametrize(("model_cls", "repo_id"), _CASES) +def test_sdxl_metadata_flags_download(model_cls, repo_id): + meta = model_cls.get_metadata() + assert meta["requires_download"] is True + assert meta["download_size_bytes"] == model_cls.DOWNLOAD_SIZE_BYTES From 4f64bda10b7f7a6aaac55d6c33cc6a4732c45e06 Mon Sep 17 00:00:00 2001 From: Irozuku Date: Thu, 2 Jul 2026 14:43:41 -0400 Subject: [PATCH 45/89] feat: split PixArt-Sigma and Tongyi Z-Image into downloadable checkpoints --- DashAI/back/initial_components.py | 16 ++- .../models/hugging_face/pixart_sigma_model.py | 118 +++++++++--------- .../hugging_face/tongyi_z_image_model.py | 113 +++++++++-------- .../models/test_pixart_tongyi_downloadable.py | 35 ++++++ 4 files changed, 159 insertions(+), 123 deletions(-) create mode 100644 tests/back/models/test_pixart_tongyi_downloadable.py diff --git a/DashAI/back/initial_components.py b/DashAI/back/initial_components.py index ff31ea330..19a319fef 100644 --- a/DashAI/back/initial_components.py +++ b/DashAI/back/initial_components.py @@ -183,7 +183,10 @@ from DashAI.back.models.hugging_face.opus_mt_fr_en_transformer import ( OpusMtFrEnTransformer, ) -from DashAI.back.models.hugging_face.pixart_sigma_model import PixArtSigmaModel +from DashAI.back.models.hugging_face.pixart_sigma_model import ( + PixArtSigma512, + PixArtSigma1024, +) from DashAI.back.models.hugging_face.qwen_model import ( Qwen25_05BInstruct, Qwen25_15BInstruct, @@ -223,7 +226,10 @@ StableDiffusionXL, ) from DashAI.back.models.hugging_face.t5_small_transformer import T5SmallTransformer -from DashAI.back.models.hugging_face.tongyi_z_image_model import TongyiZImageModel +from DashAI.back.models.hugging_face.tongyi_z_image_model import ( + TongyiZImage, + TongyiZImageTurbo, +) from DashAI.back.models.hugging_face.xlm_roberta_transformer import ( XlmRobertaTransformer, ) @@ -387,7 +393,8 @@ def get_initial_components(): OpusMtEnPtTransformer, OpusMtEsENTransformer, OpusMtFrEnTransformer, - PixArtSigmaModel, + PixArtSigma1024, + PixArtSigma512, Qwen25_05BInstruct, Qwen25_15BInstruct, RandomForestClassifier, @@ -414,7 +421,8 @@ def get_initial_components(): SVR, T5SmallTransformer, TfIdfLogRegTextClassificationModel, - TongyiZImageModel, + TongyiZImage, + TongyiZImageTurbo, XlmRobertaTransformer, XlnetTransformer, MLPImageClassifier, diff --git a/DashAI/back/models/hugging_face/pixart_sigma_model.py b/DashAI/back/models/hugging_face/pixart_sigma_model.py index e4ca5b6dd..b03ecbb96 100644 --- a/DashAI/back/models/hugging_face/pixart_sigma_model.py +++ b/DashAI/back/models/hugging_face/pixart_sigma_model.py @@ -9,6 +9,9 @@ ) from DashAI.back.core.schema_fields.base_schema import BaseSchema from DashAI.back.core.utils import MultilingualString +from DashAI.back.dependencies.downloads.downloadable import ( + HFPretrainedDownloadMixin, +) from DashAI.back.models.text_to_image_generation_model import ( TextToImageGenerationTaskModel, ) @@ -26,64 +29,6 @@ class PixArtSigmaSchema(BaseSchema): ``PixArtSigmaModel``. """ - model_name: schema_field( - enum_field( - enum=[ - "PixArt-alpha/PixArt-Sigma-XL-2-1024-MS", - "PixArt-alpha/PixArt-Sigma-XL-2-512-MS", - ] - ), - placeholder="PixArt-alpha/PixArt-Sigma-XL-2-1024-MS", - description=MultilingualString( - en=( - "The PixArt-Sigma checkpoint to load. " - "'PixArt-Sigma-XL-2-1024-MS' is the high resolution variant " - "trained at 1024px with multiscale support, delivering the best " - "image quality. " - "'PixArt-Sigma-XL-2-512-MS' is the 512px variant, faster and lighter " - "while still producing sharp results." - ), - es=( - "El checkpoint PixArt-Sigma a cargar. " - "'PixArt-Sigma-XL-2-1024-MS' es la variante de alta resolución " - "entrenada a 1024px con soporte multiescala, entregando la mejor " - "calidad de imagen. " - "'PixArt-Sigma-XL-2-512-MS' es la variante de 512px, más rápida y " - "ligera manteniendo resultados nítidos." - ), - pt=( - "O checkpoint PixArt-Sigma a carregar. " - "'PixArt-Sigma-XL-2-1024-MS' é a variante de alta resolução " - "treinada a 1024px com suporte multiescala, entregando a melhor " - "qualidade de imagem. " - "'PixArt-Sigma-XL-2-512-MS' é a variante de 512px, mais rápida e " - "leve, mantendo resultados nítidos." - ), - de=( - "Der zu ladende PixArt-Sigma-Checkpoint. " - "'PixArt-Sigma-XL-2-1024-MS' ist die hochauflösende Variante, " - "bei 1024px mit Multi-Skalen-Unterstützung trainiert und liefert " - "die beste Bildqualität. " - "'PixArt-Sigma-XL-2-512-MS' ist die 512px-Variante, schneller und " - "leichter bei dennoch scharfen Ergebnissen." - ), - zh=( - "要加载的 PixArt-Sigma 检查点。" - "'PixArt-Sigma-XL-2-1024-MS' 是以 1024px 训练的高分辨率变体，" - "支持多尺度，图像质量最佳。" - "'PixArt-Sigma-XL-2-512-MS' 是 512px 变体，速度更快、更轻量，" - "同样能产生清晰效果。" - ), - ), - alias=MultilingualString( - en="Model name", - es="Nombre del modelo", - pt="Nome do modelo", - de="Modellname", - zh="模型名称", - ), - ) # type: ignore - negative_prompt: Optional[ schema_field( string_field(), @@ -376,7 +321,9 @@ class PixArtSigmaSchema(BaseSchema): ) # type: ignore -class PixArtSigmaModel(TextToImageGenerationTaskModel): +class PixArtSigmaGenerationModel( + HFPretrainedDownloadMixin, TextToImageGenerationTaskModel +): """Diffusion Transformer model for high efficiency text-to-image generation. Wraps the PixArt-Sigma pipeline, which replaces the U-Net backbone used @@ -398,6 +345,7 @@ class PixArtSigmaModel(TextToImageGenerationTaskModel): """ SCHEMA = PixArtSigmaSchema + MODEL_NAME: str = "" COLOR: str = "#6a1b9a" DISPLAY_NAME: str = MultilingualString( en="PixArt-Sigma", @@ -510,9 +458,7 @@ def __init__(self, **kwargs): self.device = ( f"cuda:{DEVICE_TO_IDX.get(kwargs.get('device'))}" if use_gpu else "cpu" ) - self.model_name = kwargs.get( - "model_name", "PixArt-alpha/PixArt-Sigma-XL-2-1024-MS" - ) + self._pretrained_source(None) self.model = PixArtSigmaPipeline.from_pretrained( self.model_name, @@ -557,3 +503,51 @@ def generate(self, input: str) -> List[Any]: num_images_per_prompt=self.num_images_per_prompt, ) return output.images + + +class PixArtSigma1024(PixArtSigmaGenerationModel): + """PixArt-Sigma XL 1024px checkpoint. + + Downloads its checkpoint into the component's own download folder. + """ + + MODEL_NAME: str = "PixArt-alpha/PixArt-Sigma-XL-2-1024-MS" + DOWNLOAD_SIZE_BYTES: int = 2500000000 + DISPLAY_NAME = MultilingualString( + en="PixArt-Sigma 1024", + es="PixArt-Sigma 1024", + pt="PixArt-Sigma 1024", + de="PixArt-Sigma 1024", + zh="PixArt-Sigma 1024", + ) + DESCRIPTION = MultilingualString( + en="PixArt-Sigma XL 1024px checkpoint.", + es="PixArt-Sigma XL 1024px checkpoint.", + pt="PixArt-Sigma XL 1024px checkpoint.", + de="PixArt-Sigma XL 1024px checkpoint.", + zh="PixArt-Sigma XL 1024px checkpoint.", + ) + + +class PixArtSigma512(PixArtSigmaGenerationModel): + """PixArt-Sigma XL 512px checkpoint (faster). + + Downloads its checkpoint into the component's own download folder. + """ + + MODEL_NAME: str = "PixArt-alpha/PixArt-Sigma-XL-2-512-MS" + DOWNLOAD_SIZE_BYTES: int = 2500000000 + DISPLAY_NAME = MultilingualString( + en="PixArt-Sigma 512", + es="PixArt-Sigma 512", + pt="PixArt-Sigma 512", + de="PixArt-Sigma 512", + zh="PixArt-Sigma 512", + ) + DESCRIPTION = MultilingualString( + en="PixArt-Sigma XL 512px checkpoint (faster).", + es="PixArt-Sigma XL 512px checkpoint (faster).", + pt="PixArt-Sigma XL 512px checkpoint (faster).", + de="PixArt-Sigma XL 512px checkpoint (faster).", + zh="PixArt-Sigma XL 512px checkpoint (faster).", + ) diff --git a/DashAI/back/models/hugging_face/tongyi_z_image_model.py b/DashAI/back/models/hugging_face/tongyi_z_image_model.py index e1af35586..94bb0bebc 100644 --- a/DashAI/back/models/hugging_face/tongyi_z_image_model.py +++ b/DashAI/back/models/hugging_face/tongyi_z_image_model.py @@ -9,6 +9,9 @@ ) from DashAI.back.core.schema_fields.base_schema import BaseSchema from DashAI.back.core.utils import MultilingualString +from DashAI.back.dependencies.downloads.downloadable import ( + HFPretrainedDownloadMixin, +) from DashAI.back.models.text_to_image_generation_model import ( TextToImageGenerationTaskModel, ) @@ -26,61 +29,6 @@ class TongyiZImageSchema(BaseSchema): ``TongyiZImageModel``. """ - model_name: schema_field( - enum_field(enum=["Tongyi-MAI/Z-Image", "Tongyi-MAI/Z-Image-Turbo"]), - placeholder="Tongyi-MAI/Z-Image", - description=MultilingualString( - en=( - "The Tongyi Z-Image checkpoint to load. " - "'Tongyi-Z-Image' is Alibaba's 6B parameter text-to-image model " - "using a unique S3-DiT (Sparse Spatial-Spectral Diffusion Transformer) " - "architecture, one of the most downloaded models on " - "Hugging Face. It outperforms previous open source state of the art " - "models at a fraction of their parameter count." - ), - es=( - "El checkpoint Tongyi Z-Image a cargar. " - "'Tongyi-Z-Image' es el modelo de texto a imagen de 6B parámetros de " - "Alibaba que usa una arquitectura S3-DiT única (Sparse " - "Spatial-Spectral Diffusion Transformer), uno de los " - "más descargados en " - "Hugging Face. Supera a modelos de última generación anteriores con " - "una fracción de su cantidad de parámetros." - ), - pt=( - "O checkpoint Tongyi Z-Image a carregar. " - "'Tongyi-Z-Image' é o modelo de texto para imagem de 6B parâmetros " - "da Alibaba que usa uma arquitetura S3-DiT única (Sparse " - "Spatial-Spectral Diffusion Transformer), um dos " - "mais baixados no " - "Hugging Face. Supera modelos anteriores de última geração com " - "uma fração de sua quantidade de parâmetros." - ), - de=( - "Der zu ladende Tongyi Z-Image-Checkpoint. " - "'Tongyi-Z-Image' ist Alibabas 6B-Parameter-Text-zu-Bild-Modell " - "mit einer einzigartigen S3-DiT-Architektur (Sparse Spatial-Spectral " - "Diffusion Transformer), eines der am häufigsten heruntergeladenen " - "Modelle auf Hugging Face. Es übertrifft frühere Open-Source-Modelle " - "auf dem neuesten Stand bei einem Bruchteil deren Parameteranzahl." - ), - zh=( - "要加载的 Tongyi Z-Image 检查点。" - "'Tongyi-Z-Image' 是阿里巴巴的 60 亿参数文本到图像模型，" - "采用独特的 S3-DiT 架构（稀疏空间-频谱扩散变换器），" - "是 Hugging Face 上下载量最高的模型之一。" - "以更少的参数量超越了此前的开源最先进模型。" - ), - ), - alias=MultilingualString( - en="Model name", - es="Nombre del modelo", - pt="Nome do modelo", - de="Modellname", - zh="模型名称", - ), - ) # type: ignore - negative_prompt: Optional[ schema_field( string_field(), @@ -360,7 +308,9 @@ class TongyiZImageSchema(BaseSchema): ) # type: ignore -class TongyiZImageModel(TextToImageGenerationTaskModel): +class TongyiZImageGenerationModel( + HFPretrainedDownloadMixin, TextToImageGenerationTaskModel +): """Tongyi Z-Image S3-DiT model for high quality text-to-image generation. Wraps Alibaba's 6B parameter Tongyi Z-Image pipeline. The model uses a @@ -376,6 +326,7 @@ class TongyiZImageModel(TextToImageGenerationTaskModel): """ SCHEMA = TongyiZImageSchema + MODEL_NAME: str = "" COLOR: str = "#e65100" DISPLAY_NAME: str = MultilingualString( en="Tongyi Z-Image", @@ -443,7 +394,7 @@ def __init__(self, **kwargs): ) self.model = DiffusionPipeline.from_pretrained( - kwargs.get("model_name"), + self._pretrained_source(None), torch_dtype=torch.float16 if use_gpu else torch.float32, ).to(self.device) @@ -485,3 +436,51 @@ def generate(self, input: str) -> List[Any]: num_images_per_prompt=self.num_images_per_prompt, ) return output.images + + +class TongyiZImage(TongyiZImageGenerationModel): + """Tongyi Z-Image text-to-image checkpoint. + + Downloads its checkpoint into the component's own download folder. + """ + + MODEL_NAME: str = "Tongyi-MAI/Z-Image" + DOWNLOAD_SIZE_BYTES: int = 8000000000 + DISPLAY_NAME = MultilingualString( + en="Tongyi Z-Image", + es="Tongyi Z-Image", + pt="Tongyi Z-Image", + de="Tongyi Z-Image", + zh="Tongyi Z-Image", + ) + DESCRIPTION = MultilingualString( + en="Tongyi Z-Image text-to-image checkpoint.", + es="Tongyi Z-Image text-to-image checkpoint.", + pt="Tongyi Z-Image text-to-image checkpoint.", + de="Tongyi Z-Image text-to-image checkpoint.", + zh="Tongyi Z-Image text-to-image checkpoint.", + ) + + +class TongyiZImageTurbo(TongyiZImageGenerationModel): + """Tongyi Z-Image Turbo fast checkpoint. + + Downloads its checkpoint into the component's own download folder. + """ + + MODEL_NAME: str = "Tongyi-MAI/Z-Image-Turbo" + DOWNLOAD_SIZE_BYTES: int = 8000000000 + DISPLAY_NAME = MultilingualString( + en="Tongyi Z-Image Turbo", + es="Tongyi Z-Image Turbo", + pt="Tongyi Z-Image Turbo", + de="Tongyi Z-Image Turbo", + zh="Tongyi Z-Image Turbo", + ) + DESCRIPTION = MultilingualString( + en="Tongyi Z-Image Turbo fast checkpoint.", + es="Tongyi Z-Image Turbo fast checkpoint.", + pt="Tongyi Z-Image Turbo fast checkpoint.", + de="Tongyi Z-Image Turbo fast checkpoint.", + zh="Tongyi Z-Image Turbo fast checkpoint.", + ) diff --git a/tests/back/models/test_pixart_tongyi_downloadable.py b/tests/back/models/test_pixart_tongyi_downloadable.py new file mode 100644 index 000000000..1a646268a --- /dev/null +++ b/tests/back/models/test_pixart_tongyi_downloadable.py @@ -0,0 +1,35 @@ +"""Tests that PixArt-Sigma / Tongyi Z-Image models expose download metadata.""" + +import pytest + +from DashAI.back.dependencies.downloads.downloadable import HFPretrainedDownloadMixin +from DashAI.back.models.hugging_face.pixart_sigma_model import ( + PixArtSigma512, + PixArtSigma1024, +) +from DashAI.back.models.hugging_face.tongyi_z_image_model import ( + TongyiZImage, + TongyiZImageTurbo, +) + +_CASES = [ + (PixArtSigma1024, "PixArt-alpha/PixArt-Sigma-XL-2-1024-MS"), + (PixArtSigma512, "PixArt-alpha/PixArt-Sigma-XL-2-512-MS"), + (TongyiZImage, "Tongyi-MAI/Z-Image"), + (TongyiZImageTurbo, "Tongyi-MAI/Z-Image-Turbo"), +] + + +@pytest.mark.parametrize(("model_cls", "repo_id"), _CASES) +def test_is_downloadable(model_cls, repo_id): + assert issubclass(model_cls, HFPretrainedDownloadMixin) + assert model_cls.REQUIRES_DOWNLOAD is True + assert model_cls.DOWNLOAD_SIZE_BYTES is not None + assert model_cls.hf_repos() == [(repo_id, "model")] + + +@pytest.mark.parametrize(("model_cls", "repo_id"), _CASES) +def test_metadata_flags_download(model_cls, repo_id): + meta = model_cls.get_metadata() + assert meta["requires_download"] is True + assert meta["download_size_bytes"] == model_cls.DOWNLOAD_SIZE_BYTES From 0f72825391c21e0cfff7c963c4e61d473e680294 Mon Sep 17 00:00:00 2001 From: Irozuku Date: Thu, 2 Jul 2026 14:45:53 -0400 Subject: [PATCH 46/89] feat: split Stable Diffusion 3 into downloadable per-checkpoint models --- DashAI/back/initial_components.py | 10 +- .../hugging_face/stable_diffusion_v3_model.py | 177 +++++++++++------- tests/back/models/test_sd3_downloadable.py | 33 ++++ 3 files changed, 153 insertions(+), 67 deletions(-) create mode 100644 tests/back/models/test_sd3_downloadable.py diff --git a/DashAI/back/initial_components.py b/DashAI/back/initial_components.py index 19a319fef..55be38ecb 100644 --- a/DashAI/back/initial_components.py +++ b/DashAI/back/initial_components.py @@ -219,7 +219,10 @@ StableDiffusion21_512, ) from DashAI.back.models.hugging_face.stable_diffusion_v3_model import ( - StableDiffusionV3Model, + StableDiffusion3Medium, + StableDiffusion35Large, + StableDiffusion35LargeTurbo, + StableDiffusion35Medium, ) from DashAI.back.models.hugging_face.stable_diffusion_xl_model import ( RealVisXLV4, @@ -413,7 +416,10 @@ def get_initial_components(): StableDiffusion2_512, StableDiffusion21, StableDiffusion21_512, - StableDiffusionV3Model, + StableDiffusion3Medium, + StableDiffusion35Medium, + StableDiffusion35Large, + StableDiffusion35LargeTurbo, StableDiffusionXL, RealVisXLV4, StableDiffusionXLV1ControlNet, diff --git a/DashAI/back/models/hugging_face/stable_diffusion_v3_model.py b/DashAI/back/models/hugging_face/stable_diffusion_v3_model.py index 34685fd51..979cc79b9 100644 --- a/DashAI/back/models/hugging_face/stable_diffusion_v3_model.py +++ b/DashAI/back/models/hugging_face/stable_diffusion_v3_model.py @@ -9,6 +9,9 @@ ) from DashAI.back.core.schema_fields.base_schema import BaseSchema from DashAI.back.core.utils import MultilingualString +from DashAI.back.dependencies.downloads.downloadable import ( + HFPretrainedDownloadMixin, +) from DashAI.back.models.text_to_image_generation_model import ( TextToImageGenerationTaskModel, ) @@ -26,67 +29,6 @@ class StableDiffusionSchema(BaseSchema): (``num_images_per_prompt``) for ``StableDiffusionV3Model``. """ - model_name: schema_field( - enum_field( - enum=[ - "stabilityai/stable-diffusion-3-medium-diffusers", - "stabilityai/stable-diffusion-3.5-medium", - "stabilityai/stable-diffusion-3.5-large", - "stabilityai/stable-diffusion-3.5-large-turbo", - ] - ), - placeholder="stabilityai/stable-diffusion-3-medium-diffusers", - description=MultilingualString( - en=( - "The SD3/SD3.5 checkpoint to load. 'sd-3-medium' is the baseline " - "2B-parameter model. 'sd-3.5-medium' improves quality at similar " - "speed. 'sd-3.5-large' (8B) delivers the highest quality but needs " - "more VRAM. 'sd-3.5-large-turbo' is a distilled large model that " - "requires far fewer steps (4-8) for fast high-quality generation. " - "All variants target 1024x1024 px natively." - ), - es=( - "El checkpoint SD3/SD3.5 a cargar. 'sd-3-medium' es el modelo base " - "de 2B parámetros. 'sd-3.5-medium' mejora la calidad a velocidad " - "similar. 'sd-3.5-large' (8B) ofrece la mayor calidad pero necesita " - "más VRAM. 'sd-3.5-large-turbo' es un modelo large destilado que " - "requiere muchos menos pasos (4-8) para generación rápida de alta " - "calidad. Todas las variantes apuntan a 1024x1024 px de forma nativa." - ), - pt=( - "O checkpoint SD3/SD3.5 a carregar. 'sd-3-medium' é o modelo base " - "de 2B parâmetros. 'sd-3.5-medium' melhora a qualidade a velocidade " - "similar. 'sd-3.5-large' (8B) oferece a maior qualidade mas precisa " - "de mais VRAM. 'sd-3.5-large-turbo' é um modelo large destilado que " - "requer muito menos passos (4-8) para geração rápida de alta " - "qualidade. " - "Todas as variantes visam 1024x1024 px nativamente." - ), - de=( - "Der zu ladende SD3/SD3.5-Checkpoint. 'sd-3-medium' ist das " - "2B-Parameter-Basismodell. 'sd-3.5-medium' verbessert die Qualität " - "bei ähnlicher Geschwindigkeit. 'sd-3.5-large' (8B) liefert die höchste" - "Qualität, benötigt aber mehr VRAM. 'sd-3.5-large-turbo' ist ein " - "destilliertes Large-Modell, das deutlich weniger Schritte (4-8) für " - "schnelle hochwertige Generierung benötigt. " - "Alle Varianten zielen nativ auf 1024x1024 px ab." - ), - zh=( - "要加载的 SD3/SD3.5 检查点。'sd-3-medium' 是 2B 参数基准模型。" - "'sd-3.5-medium' 以相近速度提升质量。'sd-3.5-large'（8B）质量最高但" - "需要更多显存。'sd-3.5-large-turbo' 是蒸馏版大模型，仅需 4-8 步即可" - "快速生成高质量图像。所有变体原生目标分辨率为 1024x1024 像素。" - ), - ), - alias=MultilingualString( - en="Model name", - es="Nombre del modelo", - pt="Nome do modelo", - de="Modellname", - zh="模型名称", - ), - ) # type: ignore - huggingface_key: schema_field( string_field(), placeholder="", @@ -454,7 +396,9 @@ class StableDiffusionSchema(BaseSchema): ) # type: ignore -class StableDiffusionV3Model(TextToImageGenerationTaskModel): +class StableDiffusion3GenerationModel( + HFPretrainedDownloadMixin, TextToImageGenerationTaskModel +): """Multimodal Diffusion Transformer model for high-quality text-to-image generation. Wraps the Stable Diffusion 3 and 3.5 family of checkpoints from @@ -477,6 +421,7 @@ class StableDiffusionV3Model(TextToImageGenerationTaskModel): """ SCHEMA = StableDiffusionSchema + MODEL_NAME: str = "" COLOR: str = "#6a1b9a" DISPLAY_NAME: str = MultilingualString( en="Stable Diffusion V3", @@ -544,9 +489,7 @@ def __init__(self, **kwargs): self.device = ( f"cuda:{DEVICE_TO_IDX.get(kwargs.get('device'))}" if use_gpu else "cpu" ) - self.model_name = kwargs.get( - "model_name", "stabilityai/stable-diffusion-3-medium-diffusers" - ) + self.model_name = self._pretrained_source(None) self.huggingface_key = kwargs.get("huggingface_key") if self.huggingface_key: @@ -613,3 +556,107 @@ def generate(self, input: str) -> List[Any]: output = self.model(**params) return output.images + + +class StableDiffusion3Medium(StableDiffusion3GenerationModel): + """Stable Diffusion 3 Medium checkpoint (gated). + + Downloads its checkpoint into the component's own download folder. This is + a gated Hugging Face repo; downloading requires prior authentication + (an HF token in the environment). + """ + + MODEL_NAME: str = "stabilityai/stable-diffusion-3-medium-diffusers" + DOWNLOAD_SIZE_BYTES: int = 5500000000 + DISPLAY_NAME = MultilingualString( + en="Stable Diffusion 3 Medium", + es="Stable Diffusion 3 Medium", + pt="Stable Diffusion 3 Medium", + de="Stable Diffusion 3 Medium", + zh="Stable Diffusion 3 Medium", + ) + DESCRIPTION = MultilingualString( + en="Stable Diffusion 3 Medium checkpoint (gated).", + es="Stable Diffusion 3 Medium checkpoint (gated).", + pt="Stable Diffusion 3 Medium checkpoint (gated).", + de="Stable Diffusion 3 Medium checkpoint (gated).", + zh="Stable Diffusion 3 Medium checkpoint (gated).", + ) + + +class StableDiffusion35Medium(StableDiffusion3GenerationModel): + """Stable Diffusion 3.5 Medium checkpoint (gated). + + Downloads its checkpoint into the component's own download folder. This is + a gated Hugging Face repo; downloading requires prior authentication + (an HF token in the environment). + """ + + MODEL_NAME: str = "stabilityai/stable-diffusion-3.5-medium" + DOWNLOAD_SIZE_BYTES: int = 10000000000 + DISPLAY_NAME = MultilingualString( + en="Stable Diffusion 3.5 Medium", + es="Stable Diffusion 3.5 Medium", + pt="Stable Diffusion 3.5 Medium", + de="Stable Diffusion 3.5 Medium", + zh="Stable Diffusion 3.5 Medium", + ) + DESCRIPTION = MultilingualString( + en="Stable Diffusion 3.5 Medium checkpoint (gated).", + es="Stable Diffusion 3.5 Medium checkpoint (gated).", + pt="Stable Diffusion 3.5 Medium checkpoint (gated).", + de="Stable Diffusion 3.5 Medium checkpoint (gated).", + zh="Stable Diffusion 3.5 Medium checkpoint (gated).", + ) + + +class StableDiffusion35Large(StableDiffusion3GenerationModel): + """Stable Diffusion 3.5 Large checkpoint (gated). + + Downloads its checkpoint into the component's own download folder. This is + a gated Hugging Face repo; downloading requires prior authentication + (an HF token in the environment). + """ + + MODEL_NAME: str = "stabilityai/stable-diffusion-3.5-large" + DOWNLOAD_SIZE_BYTES: int = 16000000000 + DISPLAY_NAME = MultilingualString( + en="Stable Diffusion 3.5 Large", + es="Stable Diffusion 3.5 Large", + pt="Stable Diffusion 3.5 Large", + de="Stable Diffusion 3.5 Large", + zh="Stable Diffusion 3.5 Large", + ) + DESCRIPTION = MultilingualString( + en="Stable Diffusion 3.5 Large checkpoint (gated).", + es="Stable Diffusion 3.5 Large checkpoint (gated).", + pt="Stable Diffusion 3.5 Large checkpoint (gated).", + de="Stable Diffusion 3.5 Large checkpoint (gated).", + zh="Stable Diffusion 3.5 Large checkpoint (gated).", + ) + + +class StableDiffusion35LargeTurbo(StableDiffusion3GenerationModel): + """Stable Diffusion 3.5 Large Turbo checkpoint (gated). + + Downloads its checkpoint into the component's own download folder. This is + a gated Hugging Face repo; downloading requires prior authentication + (an HF token in the environment). + """ + + MODEL_NAME: str = "stabilityai/stable-diffusion-3.5-large-turbo" + DOWNLOAD_SIZE_BYTES: int = 16000000000 + DISPLAY_NAME = MultilingualString( + en="Stable Diffusion 3.5 Large Turbo", + es="Stable Diffusion 3.5 Large Turbo", + pt="Stable Diffusion 3.5 Large Turbo", + de="Stable Diffusion 3.5 Large Turbo", + zh="Stable Diffusion 3.5 Large Turbo", + ) + DESCRIPTION = MultilingualString( + en="Stable Diffusion 3.5 Large Turbo checkpoint (gated).", + es="Stable Diffusion 3.5 Large Turbo checkpoint (gated).", + pt="Stable Diffusion 3.5 Large Turbo checkpoint (gated).", + de="Stable Diffusion 3.5 Large Turbo checkpoint (gated).", + zh="Stable Diffusion 3.5 Large Turbo checkpoint (gated).", + ) diff --git a/tests/back/models/test_sd3_downloadable.py b/tests/back/models/test_sd3_downloadable.py new file mode 100644 index 000000000..4e3a3bbd3 --- /dev/null +++ b/tests/back/models/test_sd3_downloadable.py @@ -0,0 +1,33 @@ +"""Tests that Stable Diffusion 3 per-checkpoint models expose download metadata.""" + +import pytest + +from DashAI.back.dependencies.downloads.downloadable import HFPretrainedDownloadMixin +from DashAI.back.models.hugging_face.stable_diffusion_v3_model import ( + StableDiffusion3Medium, + StableDiffusion35Large, + StableDiffusion35LargeTurbo, + StableDiffusion35Medium, +) + +_CASES = [ + (StableDiffusion3Medium, "stabilityai/stable-diffusion-3-medium-diffusers"), + (StableDiffusion35Medium, "stabilityai/stable-diffusion-3.5-medium"), + (StableDiffusion35Large, "stabilityai/stable-diffusion-3.5-large"), + (StableDiffusion35LargeTurbo, "stabilityai/stable-diffusion-3.5-large-turbo"), +] + + +@pytest.mark.parametrize(("model_cls", "repo_id"), _CASES) +def test_sd3_is_downloadable(model_cls, repo_id): + assert issubclass(model_cls, HFPretrainedDownloadMixin) + assert model_cls.REQUIRES_DOWNLOAD is True + assert model_cls.DOWNLOAD_SIZE_BYTES is not None + assert model_cls.hf_repos() == [(repo_id, "model")] + + +@pytest.mark.parametrize(("model_cls", "repo_id"), _CASES) +def test_sd3_metadata_flags_download(model_cls, repo_id): + meta = model_cls.get_metadata() + assert meta["requires_download"] is True + assert meta["download_size_bytes"] == model_cls.DOWNLOAD_SIZE_BYTES From beef665862abef30a4357aeffd0db98df1f37a02 Mon Sep 17 00:00:00 2001 From: Irozuku Date: Thu, 2 Jul 2026 14:49:48 -0400 Subject: [PATCH 47/89] feat: make ControlNet models downloadable with per-repo local loading --- .../dependencies/downloads/downloadable.py | 26 ++++++ .../sd15_depth_controlnet_model.py | 14 +++- .../hugging_face/sd15_hed_controlnet_model.py | 14 +++- .../sd15_openpose_controlnet_model.py | 14 +++- .../sdxl_canny_controlnet_model.py | 17 +++- .../stable_diffusion_v1_depth_controlnet.py | 17 +++- .../models/test_controlnet_downloadable.py | 79 +++++++++++++++++++ 7 files changed, 164 insertions(+), 17 deletions(-) create mode 100644 tests/back/models/test_controlnet_downloadable.py diff --git a/DashAI/back/dependencies/downloads/downloadable.py b/DashAI/back/dependencies/downloads/downloadable.py index 56112910f..2334b25a0 100644 --- a/DashAI/back/dependencies/downloads/downloadable.py +++ b/DashAI/back/dependencies/downloads/downloadable.py @@ -158,6 +158,32 @@ def is_downloaded(cls) -> bool: for rid, *_ in repos ) + @classmethod + def _local_or_repo(cls, repo_id: str) -> str: + """Return the local dir for a repo if downloaded, else the repo id. + + Lets multi-repo components (e.g. ControlNet pipelines) load each repo + from the component's own download folder when present, falling back to + the Hub otherwise. Downloading is enforced by the run/session gates. + + Parameters + ---------- + repo_id : str + HuggingFace repo identifier. + + Returns + ------- + str + A local path (when the repo is present) or ``repo_id``. + """ + try: + target = cls._repo_dir(repo_id) + if target.is_dir() and any(target.iterdir()): + return str(target) + except Exception: + pass + return repo_id + @classmethod def download(cls, report: Optional[ProgressReporter] = None) -> None: """Download all repos listed in ``hf_repos()`` into ``component_dir()``. diff --git a/DashAI/back/models/hugging_face/sd15_depth_controlnet_model.py b/DashAI/back/models/hugging_face/sd15_depth_controlnet_model.py index 900bfb44f..a81a9459a 100644 --- a/DashAI/back/models/hugging_face/sd15_depth_controlnet_model.py +++ b/DashAI/back/models/hugging_face/sd15_depth_controlnet_model.py @@ -8,6 +8,9 @@ ) from DashAI.back.core.schema_fields.base_schema import BaseSchema from DashAI.back.core.utils import MultilingualString +from DashAI.back.dependencies.downloads.downloadable import ( + HFDownloadableMixin, +) from DashAI.back.models.controlnet_model import ControlNetModel as BaseControlNetModel from DashAI.back.models.utils import DEVICE_ENUM, DEVICE_PLACEHOLDER, DEVICE_TO_IDX @@ -239,7 +242,7 @@ def get_depth_map_sd15(image, device): return image -class SD15DepthControlNetModel(BaseControlNetModel): +class SD15DepthControlNetModel(HFDownloadableMixin, BaseControlNetModel): """Depth-conditioned ControlNet pipeline built on Stable Diffusion 1.5. Takes an input image and a text prompt. A depth map is estimated from the @@ -257,6 +260,11 @@ class SD15DepthControlNetModel(BaseControlNetModel): """ SCHEMA = SD15DepthControlNetSchema + HF_REPOS = [ + ("runwayml/stable-diffusion-v1-5", "model"), + ("lllyasviel/sd-controlnet-depth", "model"), + ] + DOWNLOAD_SIZE_BYTES = 5400000000 COLOR: str = "#4e342e" DISPLAY_NAME: str = MultilingualString( en="SD 1.5 Depth ControlNet", @@ -349,12 +357,12 @@ def __init__(self, **kwargs: Any): ) controlnet = ControlNetModel.from_pretrained( - "lllyasviel/sd-controlnet-depth", + self._local_or_repo("lllyasviel/sd-controlnet-depth"), torch_dtype=torch.float32 if self.device == "cpu" else torch.float16, ).to(self.device) self.pipe = StableDiffusionControlNetPipeline.from_pretrained( - "runwayml/stable-diffusion-v1-5", + self._local_or_repo("runwayml/stable-diffusion-v1-5"), controlnet=controlnet, torch_dtype=torch.float32 if self.device == "cpu" else torch.float16, ).to(self.device) diff --git a/DashAI/back/models/hugging_face/sd15_hed_controlnet_model.py b/DashAI/back/models/hugging_face/sd15_hed_controlnet_model.py index c04273c9b..54bfc6516 100644 --- a/DashAI/back/models/hugging_face/sd15_hed_controlnet_model.py +++ b/DashAI/back/models/hugging_face/sd15_hed_controlnet_model.py @@ -8,6 +8,9 @@ ) from DashAI.back.core.schema_fields.base_schema import BaseSchema from DashAI.back.core.utils import MultilingualString +from DashAI.back.dependencies.downloads.downloadable import ( + HFDownloadableMixin, +) from DashAI.back.models.controlnet_model import ControlNetModel as BaseControlNetModel from DashAI.back.models.utils import DEVICE_ENUM, DEVICE_PLACEHOLDER, DEVICE_TO_IDX @@ -166,7 +169,7 @@ class SD15HEDControlNetSchema(BaseSchema): ) # type: ignore -class SD15HEDControlNetModel(BaseControlNetModel): +class SD15HEDControlNetModel(HFDownloadableMixin, BaseControlNetModel): """HED soft-edge-conditioned ControlNet pipeline built on Stable Diffusion 1.5. Takes an input image and a text prompt. Soft edge maps are extracted from @@ -188,6 +191,11 @@ class SD15HEDControlNetModel(BaseControlNetModel): """ SCHEMA = SD15HEDControlNetSchema + HF_REPOS = [ + ("runwayml/stable-diffusion-v1-5", "model"), + ("lllyasviel/sd-controlnet-hed", "model"), + ] + DOWNLOAD_SIZE_BYTES = 5400000000 COLOR: str = "#006064" DISPLAY_NAME: str = MultilingualString( en="SD 1.5 HED ControlNet", @@ -305,12 +313,12 @@ def __init__(self, **kwargs: Any): self.hed_detector = HEDdetector.from_pretrained("lllyasviel/Annotators") controlnet = ControlNetModel.from_pretrained( - "lllyasviel/sd-controlnet-hed", + self._local_or_repo("lllyasviel/sd-controlnet-hed"), torch_dtype=torch.float32 if self.device == "cpu" else torch.float16, ).to(self.device) self.pipe = StableDiffusionControlNetPipeline.from_pretrained( - "runwayml/stable-diffusion-v1-5", + self._local_or_repo("runwayml/stable-diffusion-v1-5"), controlnet=controlnet, torch_dtype=torch.float32 if self.device == "cpu" else torch.float16, ).to(self.device) diff --git a/DashAI/back/models/hugging_face/sd15_openpose_controlnet_model.py b/DashAI/back/models/hugging_face/sd15_openpose_controlnet_model.py index 4b575a2a1..fcee46933 100644 --- a/DashAI/back/models/hugging_face/sd15_openpose_controlnet_model.py +++ b/DashAI/back/models/hugging_face/sd15_openpose_controlnet_model.py @@ -8,6 +8,9 @@ ) from DashAI.back.core.schema_fields.base_schema import BaseSchema from DashAI.back.core.utils import MultilingualString +from DashAI.back.dependencies.downloads.downloadable import ( + HFDownloadableMixin, +) from DashAI.back.models.controlnet_model import ControlNetModel as BaseControlNetModel from DashAI.back.models.utils import DEVICE_ENUM, DEVICE_PLACEHOLDER, DEVICE_TO_IDX @@ -162,7 +165,7 @@ class SD15OpenPoseControlNetSchema(BaseSchema): ) # type: ignore -class SD15OpenPoseControlNetModel(BaseControlNetModel): +class SD15OpenPoseControlNetModel(HFDownloadableMixin, BaseControlNetModel): """OpenPose-conditioned ControlNet pipeline built on Stable Diffusion 1.5. Takes an input image and a text prompt. Human body keypoints and skeleton @@ -183,6 +186,11 @@ class SD15OpenPoseControlNetModel(BaseControlNetModel): """ SCHEMA = SD15OpenPoseControlNetSchema + HF_REPOS = [ + ("runwayml/stable-diffusion-v1-5", "model"), + ("lllyasviel/sd-controlnet-openpose", "model"), + ] + DOWNLOAD_SIZE_BYTES = 5400000000 COLOR: str = "#880e4f" DISPLAY_NAME: str = MultilingualString( en="SD 1.5 OpenPose ControlNet", @@ -298,12 +306,12 @@ def __init__(self, **kwargs: Any): self.pose_detector = OpenposeDetector.from_pretrained("lllyasviel/Annotators") controlnet = ControlNetModel.from_pretrained( - "lllyasviel/sd-controlnet-openpose", + self._local_or_repo("lllyasviel/sd-controlnet-openpose"), torch_dtype=torch.float32 if self.device == "cpu" else torch.float16, ).to(self.device) self.pipe = StableDiffusionControlNetPipeline.from_pretrained( - "runwayml/stable-diffusion-v1-5", + self._local_or_repo("runwayml/stable-diffusion-v1-5"), controlnet=controlnet, torch_dtype=torch.float32 if self.device == "cpu" else torch.float16, ).to(self.device) diff --git a/DashAI/back/models/hugging_face/sdxl_canny_controlnet_model.py b/DashAI/back/models/hugging_face/sdxl_canny_controlnet_model.py index ba0cb3594..c7b603a84 100644 --- a/DashAI/back/models/hugging_face/sdxl_canny_controlnet_model.py +++ b/DashAI/back/models/hugging_face/sdxl_canny_controlnet_model.py @@ -8,6 +8,9 @@ ) from DashAI.back.core.schema_fields.base_schema import BaseSchema from DashAI.back.core.utils import MultilingualString +from DashAI.back.dependencies.downloads.downloadable import ( + HFDownloadableMixin, +) from DashAI.back.models.controlnet_model import ControlNetModel as BaseControlNetModel from DashAI.back.models.utils import DEVICE_ENUM, DEVICE_PLACEHOLDER, DEVICE_TO_IDX @@ -267,7 +270,7 @@ def get_canny_image( return Image.fromarray(edges_rgb) -class SDXLCannyControlNetModel(BaseControlNetModel): +class SDXLCannyControlNetModel(HFDownloadableMixin, BaseControlNetModel): """Canny-edge-conditioned ControlNet pipeline built on Stable Diffusion XL 1.0. Takes an input image and a text prompt. Canny edge maps are extracted using @@ -288,6 +291,12 @@ class SDXLCannyControlNetModel(BaseControlNetModel): """ SCHEMA = SDXLCannyControlNetSchema + HF_REPOS = [ + ("stabilityai/stable-diffusion-xl-base-1.0", "model"), + ("diffusers/controlnet-canny-sdxl-1.0", "model"), + ("madebyollin/sdxl-vae-fp16-fix", "model"), + ] + DOWNLOAD_SIZE_BYTES = 10000000000 COLOR: str = "#1a237e" DISPLAY_NAME: str = MultilingualString( en="SDXL Canny ControlNet", @@ -402,17 +411,17 @@ def __init__(self, **kwargs: Any): ) controlnet = ControlNetModel.from_pretrained( - "diffusers/controlnet-canny-sdxl-1.0", + self._local_or_repo("diffusers/controlnet-canny-sdxl-1.0"), torch_dtype=torch.float32 if self.device == "cpu" else torch.float16, ).to(self.device) vae = AutoencoderKL.from_pretrained( - "madebyollin/sdxl-vae-fp16-fix", + self._local_or_repo("madebyollin/sdxl-vae-fp16-fix"), torch_dtype=torch.float32 if self.device == "cpu" else torch.float16, ).to(self.device) self.pipe = StableDiffusionXLControlNetPipeline.from_pretrained( - "stabilityai/stable-diffusion-xl-base-1.0", + self._local_or_repo("stabilityai/stable-diffusion-xl-base-1.0"), controlnet=controlnet, vae=vae, torch_dtype=torch.float32 if self.device == "cpu" else torch.float16, diff --git a/DashAI/back/models/hugging_face/stable_diffusion_v1_depth_controlnet.py b/DashAI/back/models/hugging_face/stable_diffusion_v1_depth_controlnet.py index 1482b1da5..589d8aa43 100644 --- a/DashAI/back/models/hugging_face/stable_diffusion_v1_depth_controlnet.py +++ b/DashAI/back/models/hugging_face/stable_diffusion_v1_depth_controlnet.py @@ -8,6 +8,9 @@ ) from DashAI.back.core.schema_fields.base_schema import BaseSchema from DashAI.back.core.utils import MultilingualString +from DashAI.back.dependencies.downloads.downloadable import ( + HFDownloadableMixin, +) from DashAI.back.models.controlnet_model import ControlNetModel as BaseControlNetModel from DashAI.back.models.utils import DEVICE_ENUM, DEVICE_PLACEHOLDER, DEVICE_TO_IDX @@ -205,11 +208,17 @@ def get_depth_map(image, device): return image -class StableDiffusionXLV1ControlNet(BaseControlNetModel): +class StableDiffusionXLV1ControlNet(HFDownloadableMixin, BaseControlNetModel): """A wrapper implementation of ControlNet with depth preprocessing and stable diffusion xl 1.0 as pipeline.""" SCHEMA = StableDiffusionXLV1ControlNetSchema + HF_REPOS = [ + ("stabilityai/stable-diffusion-xl-base-1.0", "model"), + ("diffusers/controlnet-depth-sdxl-1.0-small", "model"), + ("madebyollin/sdxl-vae-fp16-fix", "model"), + ] + DOWNLOAD_SIZE_BYTES = 10000000000 COLOR: str = "#e65100" DISPLAY_NAME: str = MultilingualString( en="Stable Diffusion XL V1 ControlNet", @@ -313,19 +322,19 @@ def __init__(self, **kwargs: Any): ) self.controlnet = ControlNetModel.from_pretrained( - "diffusers/controlnet-depth-sdxl-1.0-small", + self._local_or_repo("diffusers/controlnet-depth-sdxl-1.0-small"), variant="fp16", use_safetensors=True, torch_dtype=torch.float32 if self.device == "cpu" else torch.float16, ).to(self.device) self.vae = AutoencoderKL.from_pretrained( - "madebyollin/sdxl-vae-fp16-fix", + self._local_or_repo("madebyollin/sdxl-vae-fp16-fix"), torch_dtype=torch.float32 if self.device == "cpu" else torch.float16, ).to(self.device) self.pipe = StableDiffusionXLControlNetPipeline.from_pretrained( - "stabilityai/stable-diffusion-xl-base-1.0", + self._local_or_repo("stabilityai/stable-diffusion-xl-base-1.0"), controlnet=self.controlnet, vae=self.vae, variant="fp16", diff --git a/tests/back/models/test_controlnet_downloadable.py b/tests/back/models/test_controlnet_downloadable.py new file mode 100644 index 000000000..cb32f6764 --- /dev/null +++ b/tests/back/models/test_controlnet_downloadable.py @@ -0,0 +1,79 @@ +"""Tests that ControlNet models expose multi-repo download metadata.""" + +import pytest +from kink import di + +from DashAI.back.dependencies.downloads.downloadable import HFDownloadableMixin +from DashAI.back.models.hugging_face.sd15_depth_controlnet_model import ( + SD15DepthControlNetModel, +) +from DashAI.back.models.hugging_face.sd15_hed_controlnet_model import ( + SD15HEDControlNetModel, +) +from DashAI.back.models.hugging_face.sd15_openpose_controlnet_model import ( + SD15OpenPoseControlNetModel, +) +from DashAI.back.models.hugging_face.sdxl_canny_controlnet_model import ( + SDXLCannyControlNetModel, +) +from DashAI.back.models.hugging_face.stable_diffusion_v1_depth_controlnet import ( + StableDiffusionXLV1ControlNet, +) + +_CLASSES = [ + SD15DepthControlNetModel, + SD15HEDControlNetModel, + SD15OpenPoseControlNetModel, + SDXLCannyControlNetModel, + StableDiffusionXLV1ControlNet, +] + + +@pytest.fixture +def component_root(tmp_path): + """Inject a temporary COMPONENT_PATH into the kink DI container.""" + sentinel = object() + old = di["config"] if "config" in di else sentinel # noqa: SIM401 + di["config"] = {"COMPONENT_PATH": str(tmp_path)} + yield tmp_path + if old is sentinel: + del di["config"] + else: + di["config"] = old + + +@pytest.mark.parametrize("model_cls", _CLASSES) +def test_controlnet_is_downloadable(model_cls): + assert issubclass(model_cls, HFDownloadableMixin) + assert model_cls.REQUIRES_DOWNLOAD is True + assert model_cls.DOWNLOAD_SIZE_BYTES is not None + # Each ControlNet pulls several repos (base + controlnet [+ vae]). + assert len(model_cls.hf_repos()) >= 2 + + +@pytest.mark.parametrize("model_cls", _CLASSES) +def test_controlnet_metadata_flags_download(model_cls): + meta = model_cls.get_metadata() + assert meta["requires_download"] is True + assert meta["download_size_bytes"] == model_cls.DOWNLOAD_SIZE_BYTES + + +def test_controlnet_is_downloaded_requires_all_repos(component_root): + """is_downloaded must be True only when every repo dir is present.""" + model_cls = SD15DepthControlNetModel + repos = [rid for rid, *_ in model_cls.hf_repos()] + + assert model_cls.is_downloaded() is False + + # Only the first repo present -> still not downloaded. + first = component_root / model_cls.__name__ / repos[0].split("/")[-1] + first.mkdir(parents=True) + (first / "config.json").write_text("{}") + assert model_cls.is_downloaded() is False + + # All repos present -> downloaded. + for rid in repos[1:]: + d = component_root / model_cls.__name__ / rid.split("/")[-1] + d.mkdir(parents=True) + (d / "config.json").write_text("{}") + assert model_cls.is_downloaded() is True From ec3203a42add36e096b41f51bec8dc60dd214027 Mon Sep 17 00:00:00 2001 From: Irozuku Date: Thu, 2 Jul 2026 15:08:52 -0400 Subject: [PATCH 48/89] feat: download ControlNet preprocessors into the component folder --- .../hugging_face/sd15_depth_controlnet_model.py | 15 ++++++++------- .../hugging_face/sd15_hed_controlnet_model.py | 7 +++++-- .../sd15_openpose_controlnet_model.py | 7 +++++-- .../stable_diffusion_v1_depth_controlnet.py | 15 ++++++++------- tests/back/models/test_controlnet_downloadable.py | 15 +++++++++++++++ 5 files changed, 41 insertions(+), 18 deletions(-) diff --git a/DashAI/back/models/hugging_face/sd15_depth_controlnet_model.py b/DashAI/back/models/hugging_face/sd15_depth_controlnet_model.py index a81a9459a..cead31cfb 100644 --- a/DashAI/back/models/hugging_face/sd15_depth_controlnet_model.py +++ b/DashAI/back/models/hugging_face/sd15_depth_controlnet_model.py @@ -180,7 +180,7 @@ class SD15DepthControlNetSchema(BaseSchema): ) # type: ignore -def get_depth_map_sd15(image, device): +def get_depth_map_sd15(image, device, model_source="Intel/dpt-hybrid-midas"): """Convert an input image to a normalised depth map for SD 1.5 ControlNet. Uses Intel's DPT-Hybrid-MiDaS model to estimate per-pixel depth, then @@ -208,10 +208,8 @@ def get_depth_map_sd15(image, device): from PIL import Image from transformers import DPTForDepthEstimation, DPTImageProcessor - depth_estimator = DPTForDepthEstimation.from_pretrained( - "Intel/dpt-hybrid-midas" - ).to(device) - feature_extractor = DPTImageProcessor.from_pretrained("Intel/dpt-hybrid-midas") + depth_estimator = DPTForDepthEstimation.from_pretrained(model_source).to(device) + feature_extractor = DPTImageProcessor.from_pretrained(model_source) pixel_values = feature_extractor(images=image, return_tensors="pt").pixel_values.to( device @@ -263,8 +261,9 @@ class SD15DepthControlNetModel(HFDownloadableMixin, BaseControlNetModel): HF_REPOS = [ ("runwayml/stable-diffusion-v1-5", "model"), ("lllyasviel/sd-controlnet-depth", "model"), + ("Intel/dpt-hybrid-midas", "model"), ] - DOWNLOAD_SIZE_BYTES = 5400000000 + DOWNLOAD_SIZE_BYTES = 5900000000 COLOR: str = "#4e342e" DISPLAY_NAME: str = MultilingualString( en="SD 1.5 Depth ControlNet", @@ -390,7 +389,9 @@ def generate(self, input: Tuple["Image.Image", str]) -> List[Any]: image = input[0] prompt = input[1] - depth_map = get_depth_map_sd15(image, self.device) + depth_map = get_depth_map_sd15( + image, self.device, self._local_or_repo("Intel/dpt-hybrid-midas") + ) output = self.pipe( prompt=prompt, image=depth_map, diff --git a/DashAI/back/models/hugging_face/sd15_hed_controlnet_model.py b/DashAI/back/models/hugging_face/sd15_hed_controlnet_model.py index 54bfc6516..461a25a4d 100644 --- a/DashAI/back/models/hugging_face/sd15_hed_controlnet_model.py +++ b/DashAI/back/models/hugging_face/sd15_hed_controlnet_model.py @@ -194,8 +194,9 @@ class SD15HEDControlNetModel(HFDownloadableMixin, BaseControlNetModel): HF_REPOS = [ ("runwayml/stable-diffusion-v1-5", "model"), ("lllyasviel/sd-controlnet-hed", "model"), + ("lllyasviel/Annotators", "model"), ] - DOWNLOAD_SIZE_BYTES = 5400000000 + DOWNLOAD_SIZE_BYTES = 7400000000 COLOR: str = "#006064" DISPLAY_NAME: str = MultilingualString( en="SD 1.5 HED ControlNet", @@ -310,7 +311,9 @@ def __init__(self, **kwargs: Any): f"cuda:{DEVICE_TO_IDX.get(kwargs.get('device'))}" if use_gpu else "cpu" ) - self.hed_detector = HEDdetector.from_pretrained("lllyasviel/Annotators") + self.hed_detector = HEDdetector.from_pretrained( + self._local_or_repo("lllyasviel/Annotators") + ) controlnet = ControlNetModel.from_pretrained( self._local_or_repo("lllyasviel/sd-controlnet-hed"), diff --git a/DashAI/back/models/hugging_face/sd15_openpose_controlnet_model.py b/DashAI/back/models/hugging_face/sd15_openpose_controlnet_model.py index fcee46933..1036adf31 100644 --- a/DashAI/back/models/hugging_face/sd15_openpose_controlnet_model.py +++ b/DashAI/back/models/hugging_face/sd15_openpose_controlnet_model.py @@ -189,8 +189,9 @@ class SD15OpenPoseControlNetModel(HFDownloadableMixin, BaseControlNetModel): HF_REPOS = [ ("runwayml/stable-diffusion-v1-5", "model"), ("lllyasviel/sd-controlnet-openpose", "model"), + ("lllyasviel/Annotators", "model"), ] - DOWNLOAD_SIZE_BYTES = 5400000000 + DOWNLOAD_SIZE_BYTES = 7400000000 COLOR: str = "#880e4f" DISPLAY_NAME: str = MultilingualString( en="SD 1.5 OpenPose ControlNet", @@ -303,7 +304,9 @@ def __init__(self, **kwargs: Any): f"cuda:{DEVICE_TO_IDX.get(kwargs.get('device'))}" if use_gpu else "cpu" ) - self.pose_detector = OpenposeDetector.from_pretrained("lllyasviel/Annotators") + self.pose_detector = OpenposeDetector.from_pretrained( + self._local_or_repo("lllyasviel/Annotators") + ) controlnet = ControlNetModel.from_pretrained( self._local_or_repo("lllyasviel/sd-controlnet-openpose"), diff --git a/DashAI/back/models/hugging_face/stable_diffusion_v1_depth_controlnet.py b/DashAI/back/models/hugging_face/stable_diffusion_v1_depth_controlnet.py index 589d8aa43..5965f1c43 100644 --- a/DashAI/back/models/hugging_face/stable_diffusion_v1_depth_controlnet.py +++ b/DashAI/back/models/hugging_face/stable_diffusion_v1_depth_controlnet.py @@ -156,7 +156,7 @@ class StableDiffusionXLV1ControlNetSchema(BaseSchema): ) # type: ignore -def get_depth_map(image, device): +def get_depth_map(image, device, model_source="Intel/dpt-hybrid-midas"): """Convert an input image to a normalised depth map for SDXL ControlNet. Uses Intel's DPT-Hybrid-MiDaS model to estimate per-pixel depth, then @@ -182,10 +182,8 @@ def get_depth_map(image, device): from PIL import Image from transformers import DPTForDepthEstimation, DPTImageProcessor - depth_estimator = DPTForDepthEstimation.from_pretrained( - "Intel/dpt-hybrid-midas" - ).to(device) - feature_extractor = DPTImageProcessor.from_pretrained("Intel/dpt-hybrid-midas") + depth_estimator = DPTForDepthEstimation.from_pretrained(model_source).to(device) + feature_extractor = DPTImageProcessor.from_pretrained(model_source) image = feature_extractor(images=image, return_tensors="pt").pixel_values.to(device) @@ -217,8 +215,9 @@ class StableDiffusionXLV1ControlNet(HFDownloadableMixin, BaseControlNetModel): ("stabilityai/stable-diffusion-xl-base-1.0", "model"), ("diffusers/controlnet-depth-sdxl-1.0-small", "model"), ("madebyollin/sdxl-vae-fp16-fix", "model"), + ("Intel/dpt-hybrid-midas", "model"), ] - DOWNLOAD_SIZE_BYTES = 10000000000 + DOWNLOAD_SIZE_BYTES = 10500000000 COLOR: str = "#e65100" DISPLAY_NAME: str = MultilingualString( en="Stable Diffusion XL V1 ControlNet", @@ -364,7 +363,9 @@ def generate(self, input: Tuple[Any, str]) -> List[Any]: image = input[0] prompt = input[1] - depth_map = get_depth_map(image, self.device) + depth_map = get_depth_map( + image, self.device, self._local_or_repo("Intel/dpt-hybrid-midas") + ) image = self.pipe( prompt=prompt, image=depth_map, diff --git a/tests/back/models/test_controlnet_downloadable.py b/tests/back/models/test_controlnet_downloadable.py index cb32f6764..461884509 100644 --- a/tests/back/models/test_controlnet_downloadable.py +++ b/tests/back/models/test_controlnet_downloadable.py @@ -51,6 +51,21 @@ def test_controlnet_is_downloadable(model_cls): assert len(model_cls.hf_repos()) >= 2 +# Preprocessor repos are fetched by download() (not from the Hub at run time), +# so they must be part of hf_repos() for the models that use one. +_PREPROCESSOR_REPOS = { + SD15DepthControlNetModel: "Intel/dpt-hybrid-midas", + SD15HEDControlNetModel: "lllyasviel/Annotators", + SD15OpenPoseControlNetModel: "lllyasviel/Annotators", + StableDiffusionXLV1ControlNet: "Intel/dpt-hybrid-midas", +} + + +@pytest.mark.parametrize(("model_cls", "repo"), _PREPROCESSOR_REPOS.items()) +def test_controlnet_preprocessor_included(model_cls, repo): + assert repo in [rid for rid, *_ in model_cls.hf_repos()] + + @pytest.mark.parametrize("model_cls", _CLASSES) def test_controlnet_metadata_flags_download(model_cls): meta = model_cls.get_metadata() From 2ea09b12f7217a1b065d5be2a105297656b2eed1 Mon Sep 17 00:00:00 2001 From: Irozuku Date: Fri, 3 Jul 2026 08:58:29 -0400 Subject: [PATCH 49/89] feat: add detailed per-checkpoint descriptions with repo links --- .../models/hugging_face/pixart_sigma_model.py | 84 +++++++- .../hugging_face/stable_diffusion_v2_model.py | 170 +++++++++++++-- .../hugging_face/stable_diffusion_v3_model.py | 200 ++++++++++++++++-- .../hugging_face/stable_diffusion_xl_model.py | 85 +++++++- .../hugging_face/tongyi_z_image_model.py | 78 ++++++- 5 files changed, 547 insertions(+), 70 deletions(-) diff --git a/DashAI/back/models/hugging_face/pixart_sigma_model.py b/DashAI/back/models/hugging_face/pixart_sigma_model.py index b03ecbb96..484228ea8 100644 --- a/DashAI/back/models/hugging_face/pixart_sigma_model.py +++ b/DashAI/back/models/hugging_face/pixart_sigma_model.py @@ -521,11 +521,45 @@ class PixArtSigma1024(PixArtSigmaGenerationModel): zh="PixArt-Sigma 1024", ) DESCRIPTION = MultilingualString( - en="PixArt-Sigma XL 1024px checkpoint.", - es="PixArt-Sigma XL 1024px checkpoint.", - pt="PixArt-Sigma XL 1024px checkpoint.", - de="PixArt-Sigma XL 1024px checkpoint.", - zh="PixArt-Sigma XL 1024px checkpoint.", + en=( + "PixArt-Sigma XL by PixArt-alpha, a diffusion transformer (DiT) " + "text-to-image model that reaches quality comparable to larger diffusion " + "models with far fewer parameters. This checkpoint generates at " + "1024x1024 px. Weights are downloaded into the component's own folder. " + "Model page: https://huggingface.co/PixArt-alpha/PixArt-Sigma-XL-2-1024-M" + "S" + ), + es=( + "PixArt-Sigma XL de PixArt-alpha, un modelo de texto a imagen basado en " + "transformer de difusión (DiT) que alcanza una calidad comparable a " + "modelos de difusión más grandes con muchos menos parámetros. Este " + "checkpoint genera a 1024x1024 px. Los pesos se descargan en la carpeta " + "propia del componente. Página del modelo: " + "https://huggingface.co/PixArt-alpha/PixArt-Sigma-XL-2-1024-MS" + ), + pt=( + "PixArt-Sigma XL de PixArt-alpha, um modelo de texto para imagem baseado " + "em transformer de difusão (DiT) que atinge qualidade comparável a " + "modelos de difusão maiores com muito menos parâmetros. Este " + "checkpoint gera a 1024x1024 px. Os pesos são baixados na pasta " + "própria do componente. Página do modelo: " + "https://huggingface.co/PixArt-alpha/PixArt-Sigma-XL-2-1024-MS" + ), + de=( + "PixArt-Sigma XL von PixArt-alpha, ein Text-zu-Bild-Modell auf Basis " + "eines Diffusion-Transformers (DiT), das mit weit weniger Parametern " + "eine Qualität vergleichbar mit größeren Diffusionsmodellen erreicht. " + "Dieser Checkpoint erzeugt Bilder mit 1024x1024 px. Die Gewichte werden " + "in den eigenen Ordner der Komponente heruntergeladen. Modellseite: " + "https://huggingface.co/PixArt-alpha/PixArt-Sigma-XL-2-1024-MS" + ), + zh=( + "PixArt-alpha 推出的 PixArt-Sigma XL，是一种基于扩散 " + "Transformer（DiT）的文本到图像模型，以远更少的参数量达到可媲美更大扩散模" + "型的质量。该检查点以 1024x1024 " + "像素生成。权重会下载到该组件自己的文件夹中。 模型页面： " + "https://huggingface.co/PixArt-alpha/PixArt-Sigma-XL-2-1024-MS" + ), ) @@ -545,9 +579,39 @@ class PixArtSigma512(PixArtSigmaGenerationModel): zh="PixArt-Sigma 512", ) DESCRIPTION = MultilingualString( - en="PixArt-Sigma XL 512px checkpoint (faster).", - es="PixArt-Sigma XL 512px checkpoint (faster).", - pt="PixArt-Sigma XL 512px checkpoint (faster).", - de="PixArt-Sigma XL 512px checkpoint (faster).", - zh="PixArt-Sigma XL 512px checkpoint (faster).", + en=( + "PixArt-Sigma XL by PixArt-alpha at 512x512 px, a diffusion transformer " + "(DiT) text-to-image model. The lower resolution makes it faster and " + "lighter than the 1024 px variant. Weights are downloaded into the " + "component's own folder. Model page: https://huggingface.co/PixArt-alpha/" + "PixArt-Sigma-XL-2-512-MS" + ), + es=( + "PixArt-Sigma XL de PixArt-alpha a 512x512 px, un modelo de texto a " + "imagen basado en transformer de difusión (DiT). La menor resolución " + "lo hace más rápido y ligero que la variante de 1024 px. Los pesos se " + "descargan en la carpeta propia del componente. Página del modelo: " + "https://huggingface.co/PixArt-alpha/PixArt-Sigma-XL-2-512-MS" + ), + pt=( + "PixArt-Sigma XL de PixArt-alpha a 512x512 px, um modelo de texto para " + "imagem baseado em transformer de difusão (DiT). A menor resolução o " + "torna mais rápido e leve que a variante de 1024 px. Os pesos são " + "baixados na pasta própria do componente. Página do modelo: " + "https://huggingface.co/PixArt-alpha/PixArt-Sigma-XL-2-512-MS" + ), + de=( + "PixArt-Sigma XL von PixArt-alpha bei 512x512 px, ein " + "Text-zu-Bild-Modell auf Basis eines Diffusion-Transformers (DiT). Die " + "geringere Auflösung macht es schneller und leichter als die " + "1024-px-Variante. Die Gewichte werden in den eigenen Ordner der " + "Komponente heruntergeladen. Modellseite: " + "https://huggingface.co/PixArt-alpha/PixArt-Sigma-XL-2-512-MS" + ), + zh=( + "PixArt-alpha 推出的 PixArt-Sigma XL，分辨率为 512x512 " + "像素，是一种基于扩散 Transformer（DiT）的文本到图像模型。较低的分辨率使" + "其比 1024 像素变体更快、更轻量。权重会下载到该组件自己的文件夹中。 " + "模型页面： https://huggingface.co/PixArt-alpha/PixArt-Sigma-XL-2-512-MS" + ), ) diff --git a/DashAI/back/models/hugging_face/stable_diffusion_v2_model.py b/DashAI/back/models/hugging_face/stable_diffusion_v2_model.py index 7d9230f46..55ff085d9 100644 --- a/DashAI/back/models/hugging_face/stable_diffusion_v2_model.py +++ b/DashAI/back/models/hugging_face/stable_diffusion_v2_model.py @@ -555,11 +555,46 @@ class StableDiffusion2(StableDiffusion2GenerationModel): zh="Stable Diffusion 2", ) DESCRIPTION = MultilingualString( - en="768px Stable Diffusion 2 checkpoint.", - es="768px Stable Diffusion 2 checkpoint.", - pt="768px Stable Diffusion 2 checkpoint.", - de="768px Stable Diffusion 2 checkpoint.", - zh="768px Stable Diffusion 2 checkpoint.", + en=( + "Stable Diffusion 2 by Stability AI, a latent text-to-image diffusion " + "model conditioned on OpenCLIP text embeddings. This checkpoint is " + "trained at 768x768 px and produces sharp, high-detail images. Weights " + "are downloaded into the component's own folder from the sd2-community " + "mirror. Model page: https://huggingface.co/sd2-community/stable-diffusio" + "n-2" + ), + es=( + "Stable Diffusion 2 de Stability AI, un modelo de difusión latente de " + "texto a imagen condicionado en embeddings de texto OpenCLIP. Este " + "checkpoint se entrena a 768x768 px y produce imágenes nítidas y muy " + "detalladas. Los pesos se descargan en la carpeta propia del componente " + "desde el espejo sd2-community. Página del modelo: " + "https://huggingface.co/sd2-community/stable-diffusion-2" + ), + pt=( + "Stable Diffusion 2 da Stability AI, um modelo de difusão latente de " + "texto para imagem condicionado em embeddings de texto OpenCLIP. Este " + "checkpoint é treinado a 768x768 px e produz imagens nítidas e com " + "muitos detalhes. Os pesos são baixados na pasta própria do componente " + "a partir do espelho sd2-community. Página do modelo: " + "https://huggingface.co/sd2-community/stable-diffusion-2" + ), + de=( + "Stable Diffusion 2 von Stability AI, ein latentes " + "Text-zu-Bild-Diffusionsmodell, das auf OpenCLIP-Texteinbettungen " + "konditioniert ist. Dieser Checkpoint wird bei 768x768 px trainiert und " + "erzeugt scharfe, detailreiche Bilder. Die Gewichte werden aus dem " + "sd2-community-Spiegel in den eigenen Ordner der Komponente " + "heruntergeladen. Modellseite: https://huggingface.co/sd2-community/stabl" + "e-diffusion-2" + ), + zh=( + "Stability AI 推出的 Stable Diffusion 2，是一种以 OpenCLIP " + "文本嵌入为条件的潜在文本到图像扩散模型。该检查点在 768x768 " + "像素下训练，可生成清晰且细节丰富的图像。权重会从 sd2-community " + "镜像下载到该组件自己的文件夹中。 模型页面： https://huggingface.co/sd2-c" + "ommunity/stable-diffusion-2" + ), ) @@ -580,11 +615,41 @@ class StableDiffusion2_512(StableDiffusion2GenerationModel): # noqa: N801 zh="Stable Diffusion 2 (512px)", ) DESCRIPTION = MultilingualString( - en="512px base Stable Diffusion 2 checkpoint (faster).", - es="512px base Stable Diffusion 2 checkpoint (faster).", - pt="512px base Stable Diffusion 2 checkpoint (faster).", - de="512px base Stable Diffusion 2 checkpoint (faster).", - zh="512px base Stable Diffusion 2 checkpoint (faster).", + en=( + "Stable Diffusion 2 base checkpoint by Stability AI, trained at 512x512 " + "px. It is faster and uses less memory than the 768 px variant, making " + "it a good choice for rapid prototyping. Weights are downloaded into the " + "component's own folder from the sd2-community mirror. Model page: " + "https://huggingface.co/sd2-community/stable-diffusion-2-base" + ), + es=( + "Checkpoint base de Stable Diffusion 2 de Stability AI, entrenado a " + "512x512 px. Es más rápido y usa menos memoria que la variante de 768 " + "px, ideal para prototipado rápido. Los pesos se descargan en la " + "carpeta propia del componente desde el espejo sd2-community. Página " + "del modelo: https://huggingface.co/sd2-community/stable-diffusion-2-base" + ), + pt=( + "Checkpoint base do Stable Diffusion 2 da Stability AI, treinado a " + "512x512 px. É mais rápido e usa menos memória que a variante de 768 " + "px, ideal para prototipagem rápida. Os pesos são baixados na pasta " + "própria do componente a partir do espelho sd2-community. Página do " + "modelo: https://huggingface.co/sd2-community/stable-diffusion-2-base" + ), + de=( + "Stable Diffusion 2 Basis-Checkpoint von Stability AI, trainiert bei " + "512x512 px. Er ist schneller und benötigt weniger Speicher als die " + "768-px-Variante und eignet sich gut für schnelles Prototyping. Die " + "Gewichte werden aus dem sd2-community-Spiegel in den eigenen Ordner der " + "Komponente heruntergeladen. Modellseite: " + "https://huggingface.co/sd2-community/stable-diffusion-2-base" + ), + zh=( + "Stability AI 推出的 Stable Diffusion 2 基础检查点，在 512x512 " + "像素下训练。相比 768 像素变体速度更快、显存占用更低，非常适合快速原型设" + "计。权重会从 sd2-community 镜像下载到该组件自己的文件夹中。 模型页面： " + "https://huggingface.co/sd2-community/stable-diffusion-2-base" + ), ) @@ -605,11 +670,42 @@ class StableDiffusion21(StableDiffusion2GenerationModel): zh="Stable Diffusion 2.1", ) DESCRIPTION = MultilingualString( - en="768px Stable Diffusion 2.1 checkpoint (further fine-tuned).", - es="768px Stable Diffusion 2.1 checkpoint (further fine-tuned).", - pt="768px Stable Diffusion 2.1 checkpoint (further fine-tuned).", - de="768px Stable Diffusion 2.1 checkpoint (further fine-tuned).", - zh="768px Stable Diffusion 2.1 checkpoint (further fine-tuned).", + en=( + "Stable Diffusion 2.1 by Stability AI, a further fine-tuned revision of " + "the 2.x family trained at 768x768 px. It generally produces cleaner, " + "more coherent images than the original 2.0. Weights are downloaded into " + "the component's own folder from the sd2-community mirror. Model page: " + "https://huggingface.co/sd2-community/stable-diffusion-2-1" + ), + es=( + "Stable Diffusion 2.1 de Stability AI, una revisión más ajustada de la " + "familia 2.x entrenada a 768x768 px. Suele producir imágenes más " + "limpias y coherentes que la 2.0 original. Los pesos se descargan en la " + "carpeta propia del componente desde el espejo sd2-community. Página " + "del modelo: https://huggingface.co/sd2-community/stable-diffusion-2-1" + ), + pt=( + "Stable Diffusion 2.1 da Stability AI, uma revisão mais ajustada da " + "família 2.x treinada a 768x768 px. Costuma produzir imagens mais " + "limpas e coerentes que a 2.0 original. Os pesos são baixados na pasta " + "própria do componente a partir do espelho sd2-community. Página do " + "modelo: https://huggingface.co/sd2-community/stable-diffusion-2-1" + ), + de=( + "Stable Diffusion 2.1 von Stability AI, eine weiter feinabgestimmte " + "Überarbeitung der 2.x-Familie, trainiert bei 768x768 px. Sie erzeugt " + "in der Regel sauberere, kohärentere Bilder als die ursprüngliche 2.0. " + "Die Gewichte werden aus dem sd2-community-Spiegel in den eigenen Ordner " + "der Komponente heruntergeladen. Modellseite: " + "https://huggingface.co/sd2-community/stable-diffusion-2-1" + ), + zh=( + "Stability AI 推出的 Stable Diffusion 2.1，是 2.x " + "系列的进一步微调版本，在 768x768 像素下训练。通常比原始的 2.0 " + "生成更干净、更连贯的图像。权重会从 sd2-community " + "镜像下载到该组件自己的文件夹中。 模型页面： https://huggingface.co/sd2-c" + "ommunity/stable-diffusion-2-1" + ), ) @@ -630,9 +726,43 @@ class StableDiffusion21_512(StableDiffusion2GenerationModel): # noqa: N801 zh="Stable Diffusion 2.1 (512px)", ) DESCRIPTION = MultilingualString( - en="512px base Stable Diffusion 2.1 checkpoint.", - es="512px base Stable Diffusion 2.1 checkpoint.", - pt="512px base Stable Diffusion 2.1 checkpoint.", - de="512px base Stable Diffusion 2.1 checkpoint.", - zh="512px base Stable Diffusion 2.1 checkpoint.", + en=( + "Stable Diffusion 2.1 base checkpoint by Stability AI, trained at " + "512x512 px. It combines the 2.1 fine-tuning improvements with the lower " + "memory footprint and faster generation of the 512 px base models. " + "Weights are downloaded into the component's own folder from the " + "sd2-community mirror. Model page: https://huggingface.co/sd2-community/s" + "table-diffusion-2-1-base" + ), + es=( + "Checkpoint base de Stable Diffusion 2.1 de Stability AI, entrenado a " + "512x512 px. Combina las mejoras de ajuste de la 2.1 con el menor " + "consumo de memoria y la generación más rápida de los modelos base de " + "512 px. Los pesos se descargan en la carpeta propia del componente " + "desde el espejo sd2-community. Página del modelo: " + "https://huggingface.co/sd2-community/stable-diffusion-2-1-base" + ), + pt=( + "Checkpoint base do Stable Diffusion 2.1 da Stability AI, treinado a " + "512x512 px. Combina as melhorias de ajuste da 2.1 com o menor uso de " + "memória e a geração mais rápida dos modelos base de 512 px. Os " + "pesos são baixados na pasta própria do componente a partir do espelho " + "sd2-community. Página do modelo: https://huggingface.co/sd2-community/s" + "table-diffusion-2-1-base" + ), + de=( + "Stable Diffusion 2.1 Basis-Checkpoint von Stability AI, trainiert bei " + "512x512 px. Er verbindet die Feinabstimmungs-Verbesserungen von 2.1 mit " + "dem geringeren Speicherbedarf und der schnelleren Generierung der " + "512-px-Basismodelle. Die Gewichte werden aus dem sd2-community-Spiegel " + "in den eigenen Ordner der Komponente heruntergeladen. Modellseite: " + "https://huggingface.co/sd2-community/stable-diffusion-2-1-base" + ), + zh=( + "Stability AI 推出的 Stable Diffusion 2.1 基础检查点，在 512x512 " + "像素下训练。它将 2.1 的微调改进与 512 " + "像素基础模型更低的显存占用和更快的生成速度相结合。权重会从 " + "sd2-community 镜像下载到该组件自己的文件夹中。 模型页面： " + "https://huggingface.co/sd2-community/stable-diffusion-2-1-base" + ), ) diff --git a/DashAI/back/models/hugging_face/stable_diffusion_v3_model.py b/DashAI/back/models/hugging_face/stable_diffusion_v3_model.py index 979cc79b9..3913d6134 100644 --- a/DashAI/back/models/hugging_face/stable_diffusion_v3_model.py +++ b/DashAI/back/models/hugging_face/stable_diffusion_v3_model.py @@ -576,11 +576,52 @@ class StableDiffusion3Medium(StableDiffusion3GenerationModel): zh="Stable Diffusion 3 Medium", ) DESCRIPTION = MultilingualString( - en="Stable Diffusion 3 Medium checkpoint (gated).", - es="Stable Diffusion 3 Medium checkpoint (gated).", - pt="Stable Diffusion 3 Medium checkpoint (gated).", - de="Stable Diffusion 3 Medium checkpoint (gated).", - zh="Stable Diffusion 3 Medium checkpoint (gated).", + en=( + "Stable Diffusion 3 Medium by Stability AI, built on the Multimodal " + "Diffusion Transformer (MMDiT) architecture with markedly improved text " + "rendering and prompt adherence over SD2. This is a gated Hugging Face " + "repo, so downloading requires prior authentication with an access " + "token. Weights are downloaded into the component's own folder. Model " + "page: https://huggingface.co/stabilityai/stable-diffusion-3-medium-diffu" + "sers" + ), + es=( + "Stable Diffusion 3 Medium de Stability AI, construido sobre la " + "arquitectura Multimodal Diffusion Transformer (MMDiT) con una " + "representación de texto y adherencia al prompt notablemente mejores " + "que SD2. Es un repositorio restringido de Hugging Face, por lo que la " + "descarga requiere autenticación previa con un token de acceso. Los " + "pesos se descargan en la carpeta propia del componente. Página del " + "modelo: https://huggingface.co/stabilityai/stable-diffusion-3-medium-dif" + "fusers" + ), + pt=( + "Stable Diffusion 3 Medium da Stability AI, construído sobre a " + "arquitetura Multimodal Diffusion Transformer (MMDiT) com renderização " + "de texto e aderência ao prompt bem melhores que o SD2. É um " + "repositório restrito do Hugging Face, portanto o download requer " + "autenticação prévia com um token de acesso. Os pesos são baixados " + "na pasta própria do componente. Página do modelo: " + "https://huggingface.co/stabilityai/stable-diffusion-3-medium-diffusers" + ), + de=( + "Stable Diffusion 3 Medium von Stability AI, basierend auf der " + "Multimodal-Diffusion-Transformer-Architektur (MMDiT) mit deutlich " + "verbesserter Textwiedergabe und Prompt-Treue gegenüber SD2. Dies ist " + "ein zugangsbeschränktes Hugging-Face-Repository, daher erfordert der " + "Download eine vorherige Authentifizierung mit einem Zugriffstoken. Die " + "Gewichte werden in den eigenen Ordner der Komponente heruntergeladen. " + "Modellseite: https://huggingface.co/stabilityai/stable-diffusion-3-mediu" + "m-diffusers" + ), + zh=( + "Stability AI 推出的 Stable Diffusion 3 Medium，基于多模态扩散 " + "Transformer（MMDiT）架构，相比 SD2 " + "在文本渲染和提示词遵循方面有显著提升。这是一个受限的 Hugging Face " + "仓库，因此下载前需要使用访问令牌进行身份验证。权重会下载到该组件自己的文" + "件夹中。 模型页面： https://huggingface.co/stabilityai/stable-diffusion-" + "3-medium-diffusers" + ), ) @@ -602,11 +643,50 @@ class StableDiffusion35Medium(StableDiffusion3GenerationModel): zh="Stable Diffusion 3.5 Medium", ) DESCRIPTION = MultilingualString( - en="Stable Diffusion 3.5 Medium checkpoint (gated).", - es="Stable Diffusion 3.5 Medium checkpoint (gated).", - pt="Stable Diffusion 3.5 Medium checkpoint (gated).", - de="Stable Diffusion 3.5 Medium checkpoint (gated).", - zh="Stable Diffusion 3.5 Medium checkpoint (gated).", + en=( + "Stable Diffusion 3.5 Medium by Stability AI, an updated MMDiT model " + "that balances image quality against hardware requirements, running " + "comfortably on consumer GPUs. This is a gated Hugging Face repo, so " + "downloading requires prior authentication with an access token. Weights " + "are downloaded into the component's own folder. Model page: " + "https://huggingface.co/stabilityai/stable-diffusion-3.5-medium" + ), + es=( + "Stable Diffusion 3.5 Medium de Stability AI, un modelo MMDiT " + "actualizado que equilibra la calidad de imagen con los requisitos de " + "hardware y funciona bien en GPUs de consumo. Es un repositorio " + "restringido de Hugging Face, por lo que la descarga requiere " + "autenticación previa con un token de acceso. Los pesos se descargan en " + "la carpeta propia del componente. Página del modelo: " + "https://huggingface.co/stabilityai/stable-diffusion-3.5-medium" + ), + pt=( + "Stable Diffusion 3.5 Medium da Stability AI, um modelo MMDiT atualizado " + "que equilibra a qualidade da imagem com os requisitos de hardware e " + "roda bem em GPUs de consumo. É um repositório restrito do Hugging " + "Face, portanto o download requer autenticação prévia com um token de " + "acesso. Os pesos são baixados na pasta própria do componente. Página " + "do modelo: https://huggingface.co/stabilityai/stable-diffusion-3.5-mediu" + "m" + ), + de=( + "Stable Diffusion 3.5 Medium von Stability AI, ein aktualisiertes " + "MMDiT-Modell, das Bildqualität und Hardwareanforderungen ausbalanciert " + "und komfortabel auf Consumer-GPUs läuft. Dies ist ein " + "zugangsbeschränktes Hugging-Face-Repository, daher erfordert der " + "Download eine vorherige Authentifizierung mit einem Zugriffstoken. Die " + "Gewichte werden in den eigenen Ordner der Komponente heruntergeladen. " + "Modellseite: https://huggingface.co/stabilityai/stable-diffusion-3.5-med" + "ium" + ), + zh=( + "Stability AI 推出的 Stable Diffusion 3.5 Medium，是更新的 MMDiT " + "模型，在图像质量与硬件需求之间取得平衡，可在消费级 GPU " + "上流畅运行。这是一个受限的 Hugging Face " + "仓库，因此下载前需要使用访问令牌进行身份验证。权重会下载到该组件自己的文" + "件夹中。 模型页面： https://huggingface.co/stabilityai/stable-diffusion-" + "3.5-medium" + ), ) @@ -628,11 +708,51 @@ class StableDiffusion35Large(StableDiffusion3GenerationModel): zh="Stable Diffusion 3.5 Large", ) DESCRIPTION = MultilingualString( - en="Stable Diffusion 3.5 Large checkpoint (gated).", - es="Stable Diffusion 3.5 Large checkpoint (gated).", - pt="Stable Diffusion 3.5 Large checkpoint (gated).", - de="Stable Diffusion 3.5 Large checkpoint (gated).", - zh="Stable Diffusion 3.5 Large checkpoint (gated).", + en=( + "Stable Diffusion 3.5 Large by Stability AI, the highest quality MMDiT " + "model in the 3.5 family, offering the strongest detail and prompt " + "adherence at the cost of more memory and slower generation. This is a " + "gated Hugging Face repo, so downloading requires prior authentication " + "with an access token. Weights are downloaded into the component's own " + "folder. Model page: https://huggingface.co/stabilityai/stable-diffusion-" + "3.5-large" + ), + es=( + "Stable Diffusion 3.5 Large de Stability AI, el modelo MMDiT de mayor " + "calidad de la familia 3.5, con el mejor detalle y adherencia al prompt " + "a costa de más memoria y una generación más lenta. Es un repositorio " + "restringido de Hugging Face, por lo que la descarga requiere " + "autenticación previa con un token de acceso. Los pesos se descargan en " + "la carpeta propia del componente. Página del modelo: " + "https://huggingface.co/stabilityai/stable-diffusion-3.5-large" + ), + pt=( + "Stable Diffusion 3.5 Large da Stability AI, o modelo MMDiT de maior " + "qualidade da família 3.5, oferecendo o melhor detalhe e aderência ao " + "prompt ao custo de mais memória e geração mais lenta. É um " + "repositório restrito do Hugging Face, portanto o download requer " + "autenticação prévia com um token de acesso. Os pesos são baixados " + "na pasta própria do componente. Página do modelo: " + "https://huggingface.co/stabilityai/stable-diffusion-3.5-large" + ), + de=( + "Stable Diffusion 3.5 Large von Stability AI, das qualitativ " + "hochwertigste MMDiT-Modell der 3.5-Familie, das beste Detailtreue und " + "Prompt-Treue bietet, allerdings auf Kosten von mehr Speicher und " + "langsamerer Generierung. Dies ist ein zugangsbeschränktes " + "Hugging-Face-Repository, daher erfordert der Download eine vorherige " + "Authentifizierung mit einem Zugriffstoken. Die Gewichte werden in den " + "eigenen Ordner der Komponente heruntergeladen. Modellseite: " + "https://huggingface.co/stabilityai/stable-diffusion-3.5-large" + ), + zh=( + "Stability AI 推出的 Stable Diffusion 3.5 Large，是 3.5 系列中质量最高的 " + "MMDiT 模型，提供最强的细节和提示词遵循能力，代价是更高的显存占用和更慢的" + "生成速度。这是一个受限的 Hugging Face " + "仓库，因此下载前需要使用访问令牌进行身份验证。权重会下载到该组件自己的文" + "件夹中。 模型页面： https://huggingface.co/stabilityai/stable-diffusion-" + "3.5-large" + ), ) @@ -654,9 +774,49 @@ class StableDiffusion35LargeTurbo(StableDiffusion3GenerationModel): zh="Stable Diffusion 3.5 Large Turbo", ) DESCRIPTION = MultilingualString( - en="Stable Diffusion 3.5 Large Turbo checkpoint (gated).", - es="Stable Diffusion 3.5 Large Turbo checkpoint (gated).", - pt="Stable Diffusion 3.5 Large Turbo checkpoint (gated).", - de="Stable Diffusion 3.5 Large Turbo checkpoint (gated).", - zh="Stable Diffusion 3.5 Large Turbo checkpoint (gated).", + en=( + "Stable Diffusion 3.5 Large Turbo by Stability AI, a distilled version " + "of 3.5 Large that produces high quality images in only a handful of " + "denoising steps for much faster generation. This is a gated Hugging " + "Face repo, so downloading requires prior authentication with an access " + "token. Weights are downloaded into the component's own folder. Model " + "page: https://huggingface.co/stabilityai/stable-diffusion-3.5-large-turb" + "o" + ), + es=( + "Stable Diffusion 3.5 Large Turbo de Stability AI, una versión " + "destilada de 3.5 Large que produce imágenes de alta calidad en apenas " + "unos pocos pasos de denoising para una generación mucho más rápida. " + "Es un repositorio restringido de Hugging Face, por lo que la descarga " + "requiere autenticación previa con un token de acceso. Los pesos se " + "descargan en la carpeta propia del componente. Página del modelo: " + "https://huggingface.co/stabilityai/stable-diffusion-3.5-large-turbo" + ), + pt=( + "Stable Diffusion 3.5 Large Turbo da Stability AI, uma versão destilada " + "do 3.5 Large que produz imagens de alta qualidade em apenas alguns " + "passos de denoising para uma geração muito mais rápida. É um " + "repositório restrito do Hugging Face, portanto o download requer " + "autenticação prévia com um token de acesso. Os pesos são baixados " + "na pasta própria do componente. Página do modelo: " + "https://huggingface.co/stabilityai/stable-diffusion-3.5-large-turbo" + ), + de=( + "Stable Diffusion 3.5 Large Turbo von Stability AI, eine destillierte " + "Version von 3.5 Large, die hochwertige Bilder in nur wenigen " + "Entrauschungsschritten für eine deutlich schnellere Generierung " + "erzeugt. Dies ist ein zugangsbeschränktes Hugging-Face-Repository, " + "daher erfordert der Download eine vorherige Authentifizierung mit einem " + "Zugriffstoken. Die Gewichte werden in den eigenen Ordner der Komponente " + "heruntergeladen. Modellseite: https://huggingface.co/stabilityai/stable-" + "diffusion-3.5-large-turbo" + ), + zh=( + "Stability AI 推出的 Stable Diffusion 3.5 Large Turbo，是 3.5 Large " + "的蒸馏版本，仅需少数几个去噪步骤即可生成高质量图像，从而大幅加快生成速度" + "。这是一个受限的 Hugging Face " + "仓库，因此下载前需要使用访问令牌进行身份验证。权重会下载到该组件自己的文" + "件夹中。 模型页面： https://huggingface.co/stabilityai/stable-diffusion-" + "3.5-large-turbo" + ), ) diff --git a/DashAI/back/models/hugging_face/stable_diffusion_xl_model.py b/DashAI/back/models/hugging_face/stable_diffusion_xl_model.py index 6a7944911..da0ec4463 100644 --- a/DashAI/back/models/hugging_face/stable_diffusion_xl_model.py +++ b/DashAI/back/models/hugging_face/stable_diffusion_xl_model.py @@ -503,11 +503,46 @@ class StableDiffusionXL(StableDiffusionXLGenerationModel): zh="Stable Diffusion XL", ) DESCRIPTION = MultilingualString( - en="Stable Diffusion XL base 1.0 checkpoint.", - es="Stable Diffusion XL base 1.0 checkpoint.", - pt="Stable Diffusion XL base 1.0 checkpoint.", - de="Stable Diffusion XL base 1.0 checkpoint.", - zh="Stable Diffusion XL base 1.0 checkpoint.", + en=( + "Stable Diffusion XL base 1.0 by Stability AI, a large latent diffusion " + "model that uses two text encoders and a 1024x1024 px native resolution " + "for high fidelity results. It is well suited to detailed, " + "photorealistic and artistic prompts. Weights are downloaded into the " + "component's own folder. Model page: https://huggingface.co/stabilityai/s" + "table-diffusion-xl-base-1.0" + ), + es=( + "Stable Diffusion XL base 1.0 de Stability AI, un modelo de difusión " + "latente grande que usa dos codificadores de texto y una resolución " + "nativa de 1024x1024 px para resultados de alta fidelidad. Es adecuado " + "para prompts detallados, fotorrealistas y artísticos. Los pesos se " + "descargan en la carpeta propia del componente. Página del modelo: " + "https://huggingface.co/stabilityai/stable-diffusion-xl-base-1.0" + ), + pt=( + "Stable Diffusion XL base 1.0 da Stability AI, um grande modelo de " + "difusão latente que usa dois codificadores de texto e resolução " + "nativa de 1024x1024 px para resultados de alta fidelidade. É adequado " + "para prompts detalhados, fotorrealistas e artísticos. Os pesos são " + "baixados na pasta própria do componente. Página do modelo: " + "https://huggingface.co/stabilityai/stable-diffusion-xl-base-1.0" + ), + de=( + "Stable Diffusion XL Basis 1.0 von Stability AI, ein großes latentes " + "Diffusionsmodell mit zwei Textencodern und einer nativen Auflösung von " + "1024x1024 px für Ergebnisse mit hoher Detailtreue. Es eignet sich gut " + "für detaillierte, fotorealistische und künstlerische Prompts. Die " + "Gewichte werden in den eigenen Ordner der Komponente heruntergeladen. " + "Modellseite: https://huggingface.co/stabilityai/stable-diffusion-xl-base" + "-1.0" + ), + zh=( + "Stability AI 推出的 Stable Diffusion XL base " + "1.0，是一种大型潜在扩散模型，使用两个文本编码器和 1024x1024 " + "像素的原生分辨率以获得高保真结果。非常适合细致、写实和艺术性的提示词。权" + "重会下载到该组件自己的文件夹中。 模型页面： https://huggingface.co/stabi" + "lityai/stable-diffusion-xl-base-1.0" + ), ) @@ -528,9 +563,39 @@ class RealVisXLV4(StableDiffusionXLGenerationModel): zh="RealVisXL V4.0", ) DESCRIPTION = MultilingualString( - en="RealVisXL V4.0 photorealistic SDXL checkpoint.", - es="RealVisXL V4.0 photorealistic SDXL checkpoint.", - pt="RealVisXL V4.0 photorealistic SDXL checkpoint.", - de="RealVisXL V4.0 photorealistic SDXL checkpoint.", - zh="RealVisXL V4.0 photorealistic SDXL checkpoint.", + en=( + "RealVisXL V4.0 by SG161222, a community fine-tune of Stable Diffusion " + "XL focused on photorealism. It excels at lifelike portraits, lighting " + "and textures while remaining compatible with the SDXL pipeline. Weights " + "are downloaded into the component's own folder. Model page: " + "https://huggingface.co/SG161222/RealVisXL_V4.0" + ), + es=( + "RealVisXL V4.0 de SG161222, un ajuste comunitario de Stable Diffusion " + "XL enfocado en el fotorrealismo. Destaca en retratos, iluminación y " + "texturas realistas, manteniéndose compatible con el pipeline de SDXL. " + "Los pesos se descargan en la carpeta propia del componente. Página del " + "modelo: https://huggingface.co/SG161222/RealVisXL_V4.0" + ), + pt=( + "RealVisXL V4.0 de SG161222, um ajuste comunitário do Stable Diffusion " + "XL focado em fotorrealismo. Destaca-se em retratos, iluminação e " + "texturas realistas, mantendo compatibilidade com o pipeline do SDXL. Os " + "pesos são baixados na pasta própria do componente. Página do modelo: " + "https://huggingface.co/SG161222/RealVisXL_V4.0" + ), + de=( + "RealVisXL V4.0 von SG161222, eine Community-Feinabstimmung von Stable " + "Diffusion XL mit Fokus auf Fotorealismus. Es glänzt bei lebensechten " + "Porträts, Beleuchtung und Texturen und bleibt mit der SDXL-Pipeline " + "kompatibel. Die Gewichte werden in den eigenen Ordner der Komponente " + "heruntergeladen. Modellseite: https://huggingface.co/SG161222/RealVisXL_" + "V4.0" + ), + zh=( + "SG161222 推出的 RealVisXL V4.0，是 Stable Diffusion XL " + "的社区微调版本，专注于照片级真实感。擅长逼真的人像、光照和纹理，同时保持" + "与 SDXL 流水线的兼容。权重会下载到该组件自己的文件夹中。 模型页面： " + "https://huggingface.co/SG161222/RealVisXL_V4.0" + ), ) diff --git a/DashAI/back/models/hugging_face/tongyi_z_image_model.py b/DashAI/back/models/hugging_face/tongyi_z_image_model.py index 94bb0bebc..4f3cd9bfe 100644 --- a/DashAI/back/models/hugging_face/tongyi_z_image_model.py +++ b/DashAI/back/models/hugging_face/tongyi_z_image_model.py @@ -454,11 +454,40 @@ class TongyiZImage(TongyiZImageGenerationModel): zh="Tongyi Z-Image", ) DESCRIPTION = MultilingualString( - en="Tongyi Z-Image text-to-image checkpoint.", - es="Tongyi Z-Image text-to-image checkpoint.", - pt="Tongyi Z-Image text-to-image checkpoint.", - de="Tongyi Z-Image text-to-image checkpoint.", - zh="Tongyi Z-Image text-to-image checkpoint.", + en=( + "Z-Image by Alibaba's Tongyi lab, a modern text-to-image diffusion model " + "with strong prompt following and multilingual support. This is the " + "standard, full-quality checkpoint. Weights are downloaded into the " + "component's own folder. Model page: https://huggingface.co/Tongyi-MAI/Z-" + "Image" + ), + es=( + "Z-Image del laboratorio Tongyi de Alibaba, un modelo moderno de " + "difusión de texto a imagen con buen seguimiento de prompts y soporte " + "multilingüe. Este es el checkpoint estándar de máxima calidad. Los " + "pesos se descargan en la carpeta propia del componente. Página del " + "modelo: https://huggingface.co/Tongyi-MAI/Z-Image" + ), + pt=( + "Z-Image do laboratório Tongyi da Alibaba, um modelo moderno de " + "difusão de texto para imagem com bom seguimento de prompts e suporte " + "multilíngue. Este é o checkpoint padrão de qualidade máxima. Os " + "pesos são baixados na pasta própria do componente. Página do modelo: " + "https://huggingface.co/Tongyi-MAI/Z-Image" + ), + de=( + "Z-Image aus Alibabas Tongyi-Labor, ein modernes " + "Text-zu-Bild-Diffusionsmodell mit guter Prompt-Befolgung und " + "mehrsprachiger Unterstützung. Dies ist der standardmäßige Checkpoint " + "in voller Qualität. Die Gewichte werden in den eigenen Ordner der " + "Komponente heruntergeladen. Modellseite: " + "https://huggingface.co/Tongyi-MAI/Z-Image" + ), + zh=( + "阿里巴巴通义实验室推出的 Z-Image，是一种现代文本到图像扩散模型，具有出色" + "的提示词遵循能力和多语言支持。这是标准的全质量检查点。权重会下载到该组件" + "自己的文件夹中。 模型页面： https://huggingface.co/Tongyi-MAI/Z-Image" + ), ) @@ -478,9 +507,38 @@ class TongyiZImageTurbo(TongyiZImageGenerationModel): zh="Tongyi Z-Image Turbo", ) DESCRIPTION = MultilingualString( - en="Tongyi Z-Image Turbo fast checkpoint.", - es="Tongyi Z-Image Turbo fast checkpoint.", - pt="Tongyi Z-Image Turbo fast checkpoint.", - de="Tongyi Z-Image Turbo fast checkpoint.", - zh="Tongyi Z-Image Turbo fast checkpoint.", + en=( + "Z-Image Turbo by Alibaba's Tongyi lab, a distilled variant of Z-Image " + "that generates images in far fewer denoising steps. It trades a little " + "quality for much faster generation. Weights are downloaded into the " + "component's own folder. Model page: https://huggingface.co/Tongyi-MAI/Z-" + "Image-Turbo" + ), + es=( + "Z-Image Turbo del laboratorio Tongyi de Alibaba, una variante destilada " + "de Z-Image que genera imágenes en muchos menos pasos de denoising. " + "Sacrifica algo de calidad a cambio de una generación mucho más " + "rápida. Los pesos se descargan en la carpeta propia del componente. " + "Página del modelo: https://huggingface.co/Tongyi-MAI/Z-Image-Turbo" + ), + pt=( + "Z-Image Turbo do laboratório Tongyi da Alibaba, uma variante destilada " + "do Z-Image que gera imagens em muito menos passos de denoising. Troca " + "um pouco de qualidade por uma geração muito mais rápida. Os pesos " + "são baixados na pasta própria do componente. Página do modelo: " + "https://huggingface.co/Tongyi-MAI/Z-Image-Turbo" + ), + de=( + "Z-Image Turbo aus Alibabas Tongyi-Labor, eine destillierte Variante von " + "Z-Image, die Bilder in weit weniger Entrauschungsschritten erzeugt. Sie " + "opfert etwas Qualität für eine deutlich schnellere Generierung. Die " + "Gewichte werden in den eigenen Ordner der Komponente heruntergeladen. " + "Modellseite: https://huggingface.co/Tongyi-MAI/Z-Image-Turbo" + ), + zh=( + "阿里巴巴通义实验室推出的 Z-Image Turbo，是 Z-Image " + "的蒸馏变体，可用更少的去噪步骤生成图像。以少量质量换取快得多的生成速度。" + "权重会下载到该组件自己的文件夹中。 模型页面： https://huggingface.co/Ton" + "gyi-MAI/Z-Image-Turbo" + ), ) From c63b9c355f67904f88ad9d86e59692dd4b3b2ce9 Mon Sep 17 00:00:00 2001 From: Irozuku Date: Fri, 3 Jul 2026 09:07:20 -0400 Subject: [PATCH 50/89] fix: update model download state in place to avoid list scroll reset --- .../src/components/custom/ComponentSelector.jsx | 4 +++- .../components/generative/CreateSessionCenter.jsx | 6 ++++-- .../components/generative/CreateSessionContext.jsx | 12 ++++++++++++ 3 files changed, 19 insertions(+), 3 deletions(-) diff --git a/DashAI/front/src/components/custom/ComponentSelector.jsx b/DashAI/front/src/components/custom/ComponentSelector.jsx index 997b9790d..d3c00539d 100644 --- a/DashAI/front/src/components/custom/ComponentSelector.jsx +++ b/DashAI/front/src/components/custom/ComponentSelector.jsx @@ -185,7 +185,9 @@ function ComponentSelector({ e.stopPropagation()}> onDownloadChange?.(component)} + onStatusChange={(isDownloaded) => + onDownloadChange?.(component, isDownloaded) + } /> )} diff --git a/DashAI/front/src/components/generative/CreateSessionCenter.jsx b/DashAI/front/src/components/generative/CreateSessionCenter.jsx index f15ce47ea..6a455aa80 100644 --- a/DashAI/front/src/components/generative/CreateSessionCenter.jsx +++ b/DashAI/front/src/components/generative/CreateSessionCenter.jsx @@ -29,7 +29,7 @@ export default function CreateSessionCenter() { step, models, loadingModels, - refetchModels, + markModelDownloaded, selectedModel, handleSelectModel, formik, @@ -127,7 +127,9 @@ export default function CreateSessionCenter() { components={models} selected={selectedModel} onSelect={handleSelectModelWithTour} - onDownloadChange={() => refetchModels()} + onDownloadChange={(model, isDownloaded) => + markModelDownloaded(model.name, isDownloaded) + } categoryKey="task_display_name" searchPlaceholder={t("generative:label.searchModels")} tourDataFor={tourContext?.run ? "model-card-qwen" : null} diff --git a/DashAI/front/src/components/generative/CreateSessionContext.jsx b/DashAI/front/src/components/generative/CreateSessionContext.jsx index 5aa188a24..3706e054c 100644 --- a/DashAI/front/src/components/generative/CreateSessionContext.jsx +++ b/DashAI/front/src/components/generative/CreateSessionContext.jsx @@ -235,11 +235,23 @@ export function CreateSessionProvider({ children }) { formik.submitForm(); }; + // Flip a single model's downloaded flag in place. Used when an inline + // download/delete finishes so the list updates without a full refetch + // (which would swap in the loading spinner and reset the scroll position). + const markModelDownloaded = useCallback((name, isDownloaded) => { + setModels((prev) => + prev.map((m) => + m.name === name ? { ...m, downloaded: isDownloaded } : m, + ), + ); + }, []); + const value = { step, models, loadingModels, refetchModels: loadModels, + markModelDownloaded, selectedModel, handleSelectModel, formik, From eb36dfbe4d7de5c9bb504e51ac93f496bf3411ab Mon Sep 17 00:00:00 2001 From: Irozuku Date: Fri, 3 Jul 2026 11:18:13 -0400 Subject: [PATCH 51/89] feat: add collector for nested downloadable components Walk a parameters dict for components selected as another component's parameter (at any depth) and report which still need downloading, reusing the same unwrap logic ModelFactory uses to build the model graph. --- DashAI/back/dependencies/downloads/nested.py | 143 ++++++++++++++++++ .../dependencies/test_nested_downloads.py | 103 +++++++++++++ 2 files changed, 246 insertions(+) create mode 100644 DashAI/back/dependencies/downloads/nested.py create mode 100644 tests/back/dependencies/test_nested_downloads.py diff --git a/DashAI/back/dependencies/downloads/nested.py b/DashAI/back/dependencies/downloads/nested.py new file mode 100644 index 000000000..45b65c973 --- /dev/null +++ b/DashAI/back/dependencies/downloads/nested.py @@ -0,0 +1,143 @@ +"""Discover nested downloadable components inside a parameters dict. + +A DashAI component may take another component as a parameter (a +``component_field``). That child may itself require a download, and the child +may in turn nest further components. This module walks a parameters dict the +same way :class:`~DashAI.back.models.model_factory.ModelFactory` does when it +instantiates the model graph, so the set of components it reports matches the +set that would actually be built. + +Two value shapes are handled, mirroring ``ModelFactory._process_param``: + +* the canonical ``{"component": , "params": {...}}`` descriptor, and +* the frontend-wrapped ``{"properties": {"component": ..., "params": ...}}``. +""" + +from typing import Any, Dict, Iterator, List, Optional, Tuple + + +def _unwrap(value: Any) -> Any: + """Strip the single-key ``properties`` wrapper the frontend adds. + + Parameters + ---------- + value : Any + A parameter value as stored in a parameters dict. + + Returns + ------- + Any + ``value["properties"]`` when ``value`` is a ``{"properties": ...}`` + wrapper, otherwise ``value`` unchanged. + """ + if isinstance(value, dict) and "properties" in value and len(value) == 1: + return value["properties"] + return value + + +def _iter_value( + value: Any, parent: Optional[str] +) -> Iterator[Tuple[str, Optional[str]]]: + """Yield ``(component_name, parent_name)`` for a single parameter value. + + Recurses into the selected component's own parameters so components nested + at any depth are reported. + + Parameters + ---------- + value : Any + A parameter value, possibly a nested component descriptor. + parent : str or None + Name of the component that owns this value, used as the ``parent`` of + any component found directly inside it. + + Yields + ------ + tuple of (str, str or None) + The nested component name and the name of its enclosing component. + """ + value = _unwrap(value) + if not (isinstance(value, dict) and "component" in value): + return + + parent_component_name = value["component"] + inner = value.get("params", {}).get("comp", {}) + if inner == {}: + name = parent_component_name + params = value.get("params", {}) + else: + name = inner.get("component") + params = inner.get("params", {}) + + if name: + yield name, parent + if isinstance(params, dict): + for sub_value in params.values(): + yield from _iter_value(sub_value, name) + + +def iter_config_components( + parameters: Dict[str, Any], +) -> Iterator[Tuple[str, Optional[str]]]: + """Yield ``(component_name, parent_name)`` for every nested component. + + Parameters + ---------- + parameters : dict + A parameters dict as produced by the DashAI configuration UI. + + Yields + ------ + tuple of (str, str or None) + Each nested component name paired with its enclosing component name + (``None`` at the top level). + """ + for value in parameters.values(): + yield from _iter_value(value, None) + + +def missing_downloads( + parameters: Dict[str, Any], + component_registry, +) -> List[Dict[str, Any]]: + """Return metadata for nested components that still need downloading. + + Each candidate is reconciled against the filesystem via + ``refresh_download_status`` so a component downloaded after startup (in the + worker process) is recognised without an API restart. + + Parameters + ---------- + parameters : dict + A parameters dict as produced by the DashAI configuration UI. + component_registry : ComponentRegistry + The registry used to resolve component classes and download state. + + Returns + ------- + list of dict + One entry per not-yet-downloaded nested component, each with + ``name``, ``parent``, and ``download_size_bytes`` keys. Empty when + every nested download-required component is already present. + """ + missing: List[Dict[str, Any]] = [] + seen = set() + for name, parent in iter_config_components(parameters): + if name in seen or name not in component_registry: + continue + seen.add(name) + component_class = component_registry[name]["class"] + if not getattr(component_class, "REQUIRES_DOWNLOAD", False): + continue + if component_registry.refresh_download_status(name): + continue + missing.append( + { + "name": name, + "parent": parent, + "download_size_bytes": getattr( + component_class, "DOWNLOAD_SIZE_BYTES", None + ), + } + ) + return missing diff --git a/tests/back/dependencies/test_nested_downloads.py b/tests/back/dependencies/test_nested_downloads.py new file mode 100644 index 000000000..9b196566f --- /dev/null +++ b/tests/back/dependencies/test_nested_downloads.py @@ -0,0 +1,103 @@ +"""Tests for nested downloadable-component discovery.""" + +from DashAI.back.dependencies.downloads.nested import ( + iter_config_components, + missing_downloads, +) + + +class _Comp: + """Stand-in component class carrying only download-related attributes.""" + + def __init__(self, requires, size=None): + self.REQUIRES_DOWNLOAD = requires + self.DOWNLOAD_SIZE_BYTES = size + + +class _FakeRegistry: + """Minimal registry: maps names to component classes and download state.""" + + def __init__(self, classes, downloaded): + self._classes = classes + self._downloaded = downloaded + + def __contains__(self, name): + return name in self._classes + + def __getitem__(self, name): + return {"class": self._classes[name]} + + def refresh_download_status(self, name): + return self._downloaded.get(name, True) + + +def test_iter_flat_component(): + params = {"tabular_classifier": {"component": "SVC", "params": {}}} + assert list(iter_config_components(params)) == [("SVC", None)] + + +def test_iter_unwraps_properties(): + params = { + "tabular_classifier": { + "properties": {"component": "SVC", "params": {}}, + } + } + assert list(iter_config_components(params)) == [("SVC", None)] + + +def test_iter_comp_wrapper(): + params = { + "tabular_classifier": { + "component": "BagOfWords", + "params": {"comp": {"component": "SVC", "params": {}}}, + } + } + assert list(iter_config_components(params)) == [("SVC", None)] + + +def test_iter_nested_depth(): + params = { + "outer": { + "component": "Wrapper", + "params": {"inner": {"component": "SVC", "params": {}}}, + } + } + names = list(iter_config_components(params)) + assert names == [("Wrapper", None), ("SVC", "Wrapper")] + + +def test_iter_ignores_primitives_and_fixed_values(): + params = {"n": 5, "alpha": {"fixed_value": 0.1}} + assert list(iter_config_components(params)) == [] + + +def test_missing_downloads_reports_undownloaded(): + classes = { + "SVC": _Comp(requires=False), + "BigNet": _Comp(requires=True, size=42), + } + downloaded = {"BigNet": False} + reg = _FakeRegistry(classes, downloaded) + params = { + "a": {"component": "SVC", "params": {}}, + "b": {"component": "BigNet", "params": {}}, + } + missing = missing_downloads(params, reg) + assert missing == [{"name": "BigNet", "parent": None, "download_size_bytes": 42}] + + +def test_missing_downloads_empty_when_all_present(): + classes = {"BigNet": _Comp(requires=True, size=1)} + reg = _FakeRegistry(classes, {"BigNet": True}) + params = {"b": {"component": "BigNet", "params": {}}} + assert missing_downloads(params, reg) == [] + + +def test_missing_downloads_dedupes(): + classes = {"BigNet": _Comp(requires=True, size=1)} + reg = _FakeRegistry(classes, {"BigNet": False}) + params = { + "a": {"component": "BigNet", "params": {}}, + "b": {"component": "BigNet", "params": {}}, + } + assert len(missing_downloads(params, reg)) == 1 From 25fa0bca4834316f799187d2ea84f52f8e8515d8 Mon Sep 17 00:00:00 2001 From: Irozuku Date: Fri, 3 Jul 2026 11:18:27 -0400 Subject: [PATCH 52/89] feat: block runs and sessions on undownloaded nested components The run creation, generative session upload, and model training gates now reject with the full list of nested components that still need downloading, not just the top level model. --- .../back/api/api_v1/endpoints/generative_session.py | 13 +++++++++++++ DashAI/back/api/api_v1/endpoints/runs.py | 12 ++++++++++++ DashAI/back/job/model_job.py | 8 ++++++++ 3 files changed, 33 insertions(+) diff --git a/DashAI/back/api/api_v1/endpoints/generative_session.py b/DashAI/back/api/api_v1/endpoints/generative_session.py index 503b8c72d..b67c5359c 100644 --- a/DashAI/back/api/api_v1/endpoints/generative_session.py +++ b/DashAI/back/api/api_v1/endpoints/generative_session.py @@ -15,6 +15,7 @@ GenerativeSessionParameterHistory, ProcessData, ) +from DashAI.back.dependencies.downloads.nested import missing_downloads if TYPE_CHECKING: from sqlalchemy.orm import sessionmaker @@ -60,6 +61,18 @@ async def upload_generative_session( ), ) + # A parameter may select another component that itself needs + # downloading; block until every nested one is present. + nested_missing = missing_downloads(params.parameters, component_registry) + if nested_missing: + names = ", ".join(m["name"] for m in nested_missing) + raise HTTPException( + status_code=status.HTTP_409_CONFLICT, + detail=( + f"These components must be downloaded before use: {names}." + ), + ) + # Check if the model is a subclass of GenerativeModel if not issubclass(model_class, BaseGenerativeModel): raise HTTPException( diff --git a/DashAI/back/api/api_v1/endpoints/runs.py b/DashAI/back/api/api_v1/endpoints/runs.py index cdc61958e..c9ea385bb 100644 --- a/DashAI/back/api/api_v1/endpoints/runs.py +++ b/DashAI/back/api/api_v1/endpoints/runs.py @@ -17,6 +17,7 @@ Run, RunStatus, ) +from DashAI.back.dependencies.downloads.nested import missing_downloads from DashAI.back.services.scoring_service import ScoringService if TYPE_CHECKING: @@ -343,6 +344,17 @@ async def upload_run( f"Model {params.model_name} must be downloaded before use." ), ) + # A parameter may select another component (e.g. a classifier) that + # itself needs downloading; block until every nested one is present. + nested_missing = missing_downloads(params.parameters, component_registry) + if nested_missing: + names = ", ".join(m["name"] for m in nested_missing) + raise HTTPException( + status_code=status.HTTP_409_CONFLICT, + detail=( + f"These components must be downloaded before use: {names}." + ), + ) run = Run( model_session_id=params.model_session_id, model_name=params.model_name, diff --git a/DashAI/back/job/model_job.py b/DashAI/back/job/model_job.py index b3829f54f..3819f9f34 100644 --- a/DashAI/back/job/model_job.py +++ b/DashAI/back/job/model_job.py @@ -7,6 +7,7 @@ from DashAI.back.core.enums.metrics import LevelEnum, SplitEnum from DashAI.back.dependencies.database.models import Dataset, Metric, ModelSession, Run +from DashAI.back.dependencies.downloads.nested import missing_downloads from DashAI.back.job.base_job import BaseJob, JobError from DashAI.back.metrics.base_metric import BaseMetric from DashAI.back.models.base_model import BaseModel @@ -229,6 +230,13 @@ def run( f"Model {run.model_name} is not downloaded. " "Download it before training." ) + nested_missing = missing_downloads(run.parameters, component_registry) + if nested_missing: + names = ", ".join(m["name"] for m in nested_missing) + raise JobError( + "These components are not downloaded. " + f"Download them before training: {names}." + ) try: factory = ModelFactory( run_model_class, From 990697bea6730e9143bc00432a0c84619695954e Mon Sep 17 00:00:00 2001 From: Irozuku Date: Fri, 3 Jul 2026 11:18:46 -0400 Subject: [PATCH 53/89] feat: add endpoint to resolve required nested downloads POST /components/downloads/required returns the nested components a configuration still needs downloaded, so the frontend does not re-walk the parameter tree in JS. --- .../back/api/api_v1/endpoints/components.py | 80 +++++++++++++++++++ .../api/test_components_download_required.py | 79 ++++++++++++++++++ 2 files changed, 159 insertions(+) create mode 100644 tests/back/api/test_components_download_required.py diff --git a/DashAI/back/api/api_v1/endpoints/components.py b/DashAI/back/api/api_v1/endpoints/components.py index 64b59ec84..5d7a4a534 100644 --- a/DashAI/back/api/api_v1/endpoints/components.py +++ b/DashAI/back/api/api_v1/endpoints/components.py @@ -7,6 +7,7 @@ from fastapi.exceptions import HTTPException from fastapi.responses import StreamingResponse from kink import di, inject +from pydantic import BaseModel from typing_extensions import Annotated from DashAI.back.core.utils import MultilingualString @@ -344,6 +345,85 @@ async def delete_component_download( return Response(status_code=status.HTTP_204_NO_CONTENT) +class RequiredDownloadsParams(BaseModel): + """Request body for resolving nested download-required components. + + Attributes + ---------- + model_name : str or None + The parent component being configured. When set and it still needs a + download, it is included in the result so the caller can gate on a + single list. + parameters : dict + The parameters dict as produced by the configuration UI. + """ + + model_name: Union[str, None] = None + parameters: Dict[str, Any] = {} + + +@router.post("/downloads/required") +@inject +async def get_required_downloads( + params: RequiredDownloadsParams, + accept_language: str | None = Header(default=None), + component_registry: "ComponentRegistry" = Depends(lambda: di["component_registry"]), +) -> List[Dict[str, Any]]: + """Return the components a configuration still needs downloaded. + + Walks the ``parameters`` dict for nested components (a component selected as + another component's parameter) and, optionally, checks the parent + ``model_name`` itself. Each component is reconciled against the filesystem so + the answer reflects downloads finished after startup. + + Parameters + ---------- + params : RequiredDownloadsParams + The parent ``model_name`` (optional) and its ``parameters`` dict. + accept_language : str | None + The 'Accept-Language' header used to localize display names. + component_registry : ComponentRegistry + Registry that resolves component classes and download state. + + Returns + ------- + list[dict] + One entry per not-yet-downloaded component, each with ``name``, + ``display_name``, ``parent``, and ``download_size_bytes``. + """ + from DashAI.back.dependencies.downloads.nested import missing_downloads + + missing = missing_downloads(params.parameters, component_registry) + + # Optionally fold in the parent model so callers can gate on one list. + if params.model_name and params.model_name in component_registry: + parent_class = component_registry[params.model_name]["class"] + if getattr( + parent_class, "REQUIRES_DOWNLOAD", False + ) and not component_registry.refresh_download_status(params.model_name): + missing.insert( + 0, + { + "name": params.model_name, + "parent": None, + "download_size_bytes": getattr( + parent_class, "DOWNLOAD_SIZE_BYTES", None + ), + }, + ) + + def _localized_name(name: str) -> str: + display = component_registry[name].get("display_name") + if isinstance(display, MultilingualString): + lang = (accept_language or "en").split("-")[0].lower() + return display.get(lang) + return display or name + + return [ + {**entry, "display_name": _localized_name(entry["name"])} for entry in missing + ] + + @router.get("/{id}/") @inject def get_component_by_id( diff --git a/tests/back/api/test_components_download_required.py b/tests/back/api/test_components_download_required.py new file mode 100644 index 000000000..a7b63544b --- /dev/null +++ b/tests/back/api/test_components_download_required.py @@ -0,0 +1,79 @@ +"""Tests for the nested download-resolution endpoint.""" + +import pytest +from fastapi.testclient import TestClient + +from DashAI.back.dependencies.downloads.downloadable import DownloadableMixin +from DashAI.back.dependencies.registry import ComponentRegistry +from DashAI.back.models.base_model import BaseModel + +URL = "/api/v1/component/downloads/required" + + +class PlainModel(BaseModel): + """A model with no download requirement.""" + + @classmethod + def get_schema(cls) -> dict: + return {} + + def save(self, filename=None): ... + + def load(self, filename): ... + + +class DownloadableModel(DownloadableMixin, BaseModel): + """A nested-selectable model that reports as not downloaded.""" + + DOWNLOAD_SIZE_BYTES = 123 + DESCRIPTION = "Downloadable" + DISPLAY_NAME = "Downloadable Model" + + @classmethod + def is_downloaded(cls) -> bool: + return False + + @classmethod + def get_schema(cls) -> dict: + return {} + + def save(self, filename=None): ... + + def load(self, filename): ... + + +@pytest.fixture(autouse=True) +def _registry(client, monkeypatch): + registry = ComponentRegistry(initial_components=[PlainModel, DownloadableModel]) + monkeypatch.setitem(client.app.container._services, "component_registry", registry) + return registry + + +def test_required_downloads_reports_nested(client: TestClient): + body = { + "parameters": { + "clf": {"component": "DownloadableModel", "params": {}}, + } + } + response = client.post(URL, json=body) + assert response.status_code == 200 + data = response.json() + assert len(data) == 1 + assert data[0]["name"] == "DownloadableModel" + assert data[0]["download_size_bytes"] == 123 + assert data[0]["display_name"] == "Downloadable Model" + + +def test_required_downloads_empty_for_plain(client: TestClient): + body = {"parameters": {"clf": {"component": "PlainModel", "params": {}}}} + response = client.post(URL, json=body) + assert response.status_code == 200 + assert response.json() == [] + + +def test_required_downloads_includes_parent_model(client: TestClient): + body = {"model_name": "DownloadableModel", "parameters": {}} + response = client.post(URL, json=body) + assert response.status_code == 200 + data = response.json() + assert [d["name"] for d in data] == ["DownloadableModel"] From 70ea125fc4a04417d16726ecc8a491e5590c3d97 Mon Sep 17 00:00:00 2001 From: Irozuku Date: Fri, 3 Jul 2026 11:18:59 -0400 Subject: [PATCH 54/89] feat: add dummy downloadable classifier for nested download testing A test-only tabular classifier that requires a download and itself exposes a nested classifier field, used to exercise the nested download flow at depth in the UI. --- DashAI/back/initial_components.py | 4 + .../dummy_downloadable_classifier.py | 144 ++++++++++++++++++ 2 files changed, 148 insertions(+) create mode 100644 DashAI/back/models/scikit_learn/dummy_downloadable_classifier.py diff --git a/DashAI/back/initial_components.py b/DashAI/back/initial_components.py index 55be38ecb..eaa6d6ee4 100644 --- a/DashAI/back/initial_components.py +++ b/DashAI/back/initial_components.py @@ -257,6 +257,9 @@ DecisionTreeRegression, ) from DashAI.back.models.scikit_learn.dummy_classifier import DummyClassifier +from DashAI.back.models.scikit_learn.dummy_downloadable_classifier import ( + DummyDownloadableClassifier, +) from DashAI.back.models.scikit_learn.elastic_net_regression import ElasticNetRegression from DashAI.back.models.scikit_learn.extra_trees_classifier import ExtraTreesClassifier from DashAI.back.models.scikit_learn.extra_trees_regression import ExtraTreesRegression @@ -424,6 +427,7 @@ def get_initial_components(): RealVisXLV4, StableDiffusionXLV1ControlNet, SVC, + DummyDownloadableClassifier, SVR, T5SmallTransformer, TfIdfLogRegTextClassificationModel, diff --git a/DashAI/back/models/scikit_learn/dummy_downloadable_classifier.py b/DashAI/back/models/scikit_learn/dummy_downloadable_classifier.py new file mode 100644 index 000000000..c69909ec3 --- /dev/null +++ b/DashAI/back/models/scikit_learn/dummy_downloadable_classifier.py @@ -0,0 +1,144 @@ +import time +from typing import Optional + +from DashAI.back.core.schema_fields import ( + BaseSchema, + component_field, + schema_field, +) +from DashAI.back.core.utils import MultilingualString +from DashAI.back.dependencies.downloads.downloadable import ( + DownloadableMixin, + ProgressReporter, +) +from DashAI.back.models.scikit_learn.svc import SVC + + +class DummyDownloadableClassifierSchema(BaseSchema): + """Schema for the dummy classifier, exposing one nested classifier field. + + The ``nested_classifier`` parameter is a component field so a second + (possibly download-required) tabular classifier can be selected inside + this one, letting the nested-download flow be exercised at depth. + """ + + nested_classifier: schema_field( + component_field(parent="TabularClassificationModel"), + placeholder={"component": "SVC", "params": {}}, + description=MultilingualString( + en="A nested tabular classifier, used to test nested downloads.", + es="Un clasificador tabular anidado, para probar descargas anidadas.", + pt="Um classificador tabular aninhado, para testar downloads aninhados.", + de=( + "Ein verschachtelter tabellarischer Klassifikator zum Testen " + "verschachtelter Downloads." + ), + zh="嵌套的表格分类器，用于测试嵌套下载。", + ), + alias=MultilingualString( + en="Nested classifier", + es="Clasificador anidado", + pt="Classificador aninhado", + de="Verschachtelter Klassifikator", + zh="嵌套分类器", + ), + ) # type: ignore + + +class DummyDownloadableClassifier(DownloadableMixin, SVC): + """A fake download-required tabular classifier for UI testing. + + Behaves exactly like :class:`SVC` at train time but is flagged as + requiring a download so the inline download control appears when it is + selected as another component's parameter (e.g. the Bag-of-Words tabular + classifier). Its ``download`` writes a marker file instead of fetching any + real artifact, so the download/delete flow can be exercised end to end + without network access. + """ + + SCHEMA = DummyDownloadableClassifierSchema + DOWNLOAD_SIZE_BYTES = 256 * 1024 * 1024 + COLOR = "#B39DDB" + ICON = "Timeline" + DISPLAY_NAME = MultilingualString( + en="Dummy Downloadable Classifier", + es="Clasificador Descargable de Prueba", + pt="Classificador Baixavel de Teste", + de="Dummy Herunterladbarer Klassifikator", + zh="虚拟可下载分类器", + ) + DESCRIPTION = MultilingualString( + en=( + "A test-only classifier that requires a download. It trains like an " + "SVM but is used to preview the inline download control when picking " + "a nested component." + ), + es=( + "Un clasificador solo de prueba que requiere descarga. Entrena como " + "una SVM, pero sirve para previsualizar el control de descarga en " + "linea al elegir un componente anidado." + ), + pt=( + "Um classificador apenas de teste que requer download. Treina como " + "uma SVM, mas serve para pre-visualizar o controle de download em " + "linha ao escolher um componente aninhado." + ), + de=( + "Ein reiner Testklassifikator, der einen Download erfordert. Er " + "trainiert wie eine SVM, dient aber zur Vorschau des Inline-" + "Download-Steuerelements bei der Auswahl einer verschachtelten " + "Komponente." + ), + zh=( + "仅用于测试的分类器，需要下载。它像 SVM 一样训练，" + "用于在选择嵌套组件时预览内联下载控件。" + ), + ) + + def __init__(self, **kwargs): + """Store the nested classifier and forward the rest to ``SVC``. + + Parameters + ---------- + **kwargs : dict + May include ``nested_classifier`` (an instantiated tabular + classifier), which is kept as an attribute and not passed to the + underlying sklearn estimator. + """ + self.nested_classifier = kwargs.pop("nested_classifier", None) + super().__init__(**kwargs) + + @classmethod + def is_downloaded(cls) -> bool: + """Return whether the marker file is present. + + Returns + ------- + bool + ``True`` when ``component_dir()`` exists and is non-empty. + """ + directory = cls.component_dir() + return directory.is_dir() and any(directory.iterdir()) + + @classmethod + def download(cls, report: Optional[ProgressReporter] = None) -> None: + """Write a marker file to simulate a download. + + A short delay is inserted so the downloading state is visible in the + UI. No real artifact is fetched. + + Parameters + ---------- + report : ProgressReporter, optional + Callback invoked with progress fractions and phase messages. + """ + directory = cls.component_dir() + directory.mkdir(parents=True, exist_ok=True) + steps = 4 + for step in range(steps): + if report is not None: + report(step / steps, "Downloading dummy weights") + time.sleep(1) + (directory / "weights.marker").write_text("dummy", encoding="utf-8") + if report is not None: + report(1.0, "Done") From 9c00e4326a74651e810d5e572d61e8327dc8f719 Mon Sep 17 00:00:00 2001 From: Irozuku Date: Fri, 3 Jul 2026 11:19:24 -0400 Subject: [PATCH 55/89] feat: show inline download control in nested component selectors When a component selected as another component's parameter requires a download, render the download control inline in the selector. Cache the downloaded flag in place so it persists across model switches. --- .../configurableObject/Inputs/ClassInput.jsx | 24 +++++++++++++++++++ .../shared/FormSchemaFieldWithParent.jsx | 19 +++++++++++++-- .../shared/FormSchemaModelSelect.jsx | 19 +++++++++++++-- DashAI/front/src/hooks/useModelParents.js | 20 ++++++++++++++-- 4 files changed, 76 insertions(+), 6 deletions(-) diff --git a/DashAI/front/src/components/configurableObject/Inputs/ClassInput.jsx b/DashAI/front/src/components/configurableObject/Inputs/ClassInput.jsx index 5750881f8..c46b51fb3 100644 --- a/DashAI/front/src/components/configurableObject/Inputs/ClassInput.jsx +++ b/DashAI/front/src/components/configurableObject/Inputs/ClassInput.jsx @@ -3,6 +3,7 @@ import PropTypes from "prop-types"; import FormTooltip from "../FormTooltip"; import { Input } from "./InputStyles"; import { + Box, IconButton, MenuItem, Dialog, @@ -15,6 +16,7 @@ import { } from "@mui/material"; import SettingsIcon from "@mui/icons-material/Settings"; import Subform from "../Subform"; +import ComponentDownloadControl from "../../models/model/ComponentDownloadControl"; import { getDefaultValues } from "../../../utils/values"; import { getModelSchema as getModelSchemaRequest, @@ -140,6 +142,28 @@ function ClassInput({ + {(() => { + const selectedComponent = options.find( + (option) => option.name === selectedOption, + ); + return selectedComponent?.metadata?.requires_download ? ( + + + setOptions((prev) => + prev.map((option) => + option.name === selectedComponent.name + ? { ...option, downloaded: isDownloaded } + : option, + ), + ) + } + /> + + ) : null; + })()} + {/* Button to show the modal that contains the subform */} diff --git a/DashAI/front/src/components/shared/FormSchemaFieldWithParent.jsx b/DashAI/front/src/components/shared/FormSchemaFieldWithParent.jsx index 5c35bf03c..96f23a9e4 100644 --- a/DashAI/front/src/components/shared/FormSchemaFieldWithParent.jsx +++ b/DashAI/front/src/components/shared/FormSchemaFieldWithParent.jsx @@ -1,4 +1,4 @@ -import { MenuItem, Tooltip, IconButton } from "@mui/material"; +import { Box, MenuItem, Tooltip, IconButton } from "@mui/material"; import PropTypes from "prop-types"; import { Input } from "../configurableObject/Inputs/InputStyles"; import React from "react"; @@ -12,6 +12,7 @@ import useModelParents from "../../hooks/useModelParents"; import { Settings } from "@mui/icons-material"; import { useTranslation } from "react-i18next"; import FormSchemaFieldCard from "./FormSchemaFieldCard"; +import ComponentDownloadControl from "../models/model/ComponentDownloadControl"; /** * Renders a parent-model selector field as a card. @@ -26,11 +27,15 @@ function FormSchemaFieldWithParent({ errorMessage, }) { const { addProperty, getModelFromCurrentProperty } = useFormSchemaStore(); - const { models } = useModelParents({ + const { models, markDownloaded } = useModelParents({ parent: field.value?.properties.component, }); const { t } = useTranslation(["common"]); + const selectedComponent = models?.find( + (model) => model.name === getModelFromCurrentProperty(name), + ); + const handleOnChange = async (event) => { const model = models?.find((model) => model.name === event.target.value); const { initialValues } = generateYupSchema( @@ -79,6 +84,16 @@ function FormSchemaFieldWithParent({ ))} + {selectedComponent?.metadata?.requires_download && ( + + + markDownloaded(selectedComponent.name, isDownloaded) + } + /> + + )} ); } diff --git a/DashAI/front/src/components/shared/FormSchemaModelSelect.jsx b/DashAI/front/src/components/shared/FormSchemaModelSelect.jsx index cb8c3fd89..ad9ab3967 100644 --- a/DashAI/front/src/components/shared/FormSchemaModelSelect.jsx +++ b/DashAI/front/src/components/shared/FormSchemaModelSelect.jsx @@ -1,4 +1,4 @@ -import { FormControl, MenuItem } from "@mui/material"; +import { Box, FormControl, MenuItem } from "@mui/material"; import React from "react"; import useModelParents from "../../hooks/useModelParents"; import { Input } from "../configurableObject/Inputs/InputStyles"; @@ -10,6 +10,7 @@ import { } from "../../utils/schema"; import PropTypes from "prop-types"; import { useTranslation } from "react-i18next"; +import ComponentDownloadControl from "../models/model/ComponentDownloadControl"; /** * This component is a select input for the models of a parent model @@ -19,7 +20,7 @@ import { useTranslation } from "react-i18next"; */ function FormSchemaModelSelect({ parent, selectedModel, onChange }) { - const { models } = useModelParents({ parent }); + const { models, markDownloaded } = useModelParents({ parent }); const { handleUpdateSchema } = useFormSchemaStore(); const { t } = useTranslation(["common"]); @@ -27,6 +28,10 @@ function FormSchemaModelSelect({ parent, selectedModel, onChange }) { return null; } + const selectedComponent = models.find( + (model) => model.name === selectedModel, + ); + const handleOnChange = async (event) => { const model = models.find((model) => model.name === event.target.value); const { initialValues } = generateYupSchema( @@ -52,6 +57,16 @@ function FormSchemaModelSelect({ parent, selectedModel, onChange }) { ))} + {selectedComponent?.metadata?.requires_download && ( + + + markDownloaded(selectedComponent.name, isDownloaded) + } + /> + + )} ); } diff --git a/DashAI/front/src/hooks/useModelParents.js b/DashAI/front/src/hooks/useModelParents.js index 2a6ab93b1..d187552e4 100644 --- a/DashAI/front/src/hooks/useModelParents.js +++ b/DashAI/front/src/hooks/useModelParents.js @@ -1,4 +1,4 @@ -import { useEffect, useState } from "react"; +import { useCallback, useEffect, useState } from "react"; import { getComponents } from "../api/component"; /* @@ -30,5 +30,21 @@ export default function useModelParents({ parent }) { } }, [parent]); - return { models, loading }; + // Flip a single model's downloaded flag in place so an inline download/delete + // is reflected in the cached list. Without this, switching models and coming + // back would re-mount the download control from the stale (not-downloaded) + // flag and show the Download button again. + const markDownloaded = useCallback((name, isDownloaded) => { + setModels((prev) => + prev + ? prev.map((model) => + model.name === name + ? { ...model, downloaded: isDownloaded } + : model, + ) + : prev, + ); + }, []); + + return { models, loading, markDownloaded }; } From 0a679d7e8bf9811a681b5d66ca6cc649c1a49a44 Mon Sep 17 00:00:00 2001 From: Irozuku Date: Fri, 3 Jul 2026 11:19:38 -0400 Subject: [PATCH 56/89] fix: sync download state across all controls for a component Download state is a global per-component fact, but the same component can render at several nesting levels. A module level pub/sub plus a status cache keeps every mounted control in sync and lets remounts pick up the latest state. --- .../models/model/ComponentDownloadControl.jsx | 66 +++++++++++++++++-- 1 file changed, 60 insertions(+), 6 deletions(-) diff --git a/DashAI/front/src/components/models/model/ComponentDownloadControl.jsx b/DashAI/front/src/components/models/model/ComponentDownloadControl.jsx index 4ef233238..8ad18a8b1 100644 --- a/DashAI/front/src/components/models/model/ComponentDownloadControl.jsx +++ b/DashAI/front/src/components/models/model/ComponentDownloadControl.jsx @@ -26,6 +26,45 @@ const formatSize = (bytes) => { return `${Math.round(mb)} MB`; }; +// Download state is a global, per-component fact (a component's artifacts are +// either on disk or not). The same component can be rendered by several +// controls at once (e.g. the same model selected at multiple nesting levels). +// This module-level pub/sub keeps every mounted control for a given component +// name in sync, and the cache lets a freshly mounted control pick up the +// latest known state instead of the (possibly stale) prop. +const downloadListeners = new Map(); // name -> Set<(state) => void> +const downloadStateCache = new Map(); // name -> { downloading, downloaded } +const anyChangeListeners = new Set(); // (name, state) => void + +const subscribeDownloadState = (name, listener) => { + let listeners = downloadListeners.get(name); + if (!listeners) { + listeners = new Set(); + downloadListeners.set(name, listeners); + } + listeners.add(listener); + return () => { + listeners.delete(listener); + }; +}; + +// Subscribe to every download/delete regardless of component name. Lets a +// container (e.g. a config dialog) re-check which nested components still need +// downloading after an inline control finishes. +export const subscribeAnyDownloadState = (listener) => { + anyChangeListeners.add(listener); + return () => { + anyChangeListeners.delete(listener); + }; +}; + +const broadcastDownloadState = (name, state) => { + downloadStateCache.set(name, { ...downloadStateCache.get(name), ...state }); + const listeners = downloadListeners.get(name); + if (listeners) listeners.forEach((listener) => listener(state)); + anyChangeListeners.forEach((listener) => listener(name, state)); +}; + const ComponentDownloadControl = ({ component, onStatusChange, @@ -34,14 +73,27 @@ const ComponentDownloadControl = ({ const { t } = useTranslation(["common"]); const { enqueueSnackbar } = useSnackbar(); const meta = component.metadata || {}; - const [downloaded, setDownloaded] = useState(Boolean(component.downloaded)); - const [downloading, setDownloading] = useState(false); + const cached = downloadStateCache.get(component.name); + const [downloaded, setDownloaded] = useState( + cached?.downloaded ?? Boolean(component.downloaded), + ); + const [downloading, setDownloading] = useState(cached?.downloading ?? false); const pollerIdRef = useRef(null); useEffect(() => { - setDownloaded(Boolean(component.downloaded)); + const known = downloadStateCache.get(component.name); + setDownloaded(known?.downloaded ?? Boolean(component.downloaded)); + setDownloading(known?.downloading ?? false); }, [component.name, component.downloaded]); + // Mirror download/delete triggered by any other control for this component. + useEffect(() => { + return subscribeDownloadState(component.name, (state) => { + if (state.downloading !== undefined) setDownloading(state.downloading); + if (state.downloaded !== undefined) setDownloaded(state.downloaded); + }); + }, [component.name]); + useEffect(() => { return () => { if (pollerIdRef.current != null) stopJobPolling(pollerIdRef.current); @@ -51,13 +103,15 @@ const ComponentDownloadControl = ({ if (!meta.requires_download) return null; const finish = (isDownloaded) => { - setDownloading(false); - setDownloaded(isDownloaded); + broadcastDownloadState(component.name, { + downloading: false, + downloaded: isDownloaded, + }); if (onStatusChange) onStatusChange(isDownloaded); }; const handleDownload = async () => { - setDownloading(true); + broadcastDownloadState(component.name, { downloading: true }); try { const { id } = await downloadComponent(component.name); pollerIdRef.current = id; From 9b51b6570277410bac875f1c7dcd24a12574f189 Mon Sep 17 00:00:00 2001 From: Irozuku Date: Fri, 3 Jul 2026 11:19:48 -0400 Subject: [PATCH 57/89] fix: allow selecting a component nested two levels deep getModelFromCurrentProperty did not unwrap the last hop of the property chain, so a component field nested inside another component read undefined and its select never reflected a selection. --- DashAI/front/src/contexts/schema.js | 15 +++++++++------ 1 file changed, 9 insertions(+), 6 deletions(-) diff --git a/DashAI/front/src/contexts/schema.js b/DashAI/front/src/contexts/schema.js index ab9e4da5f..a42b22c77 100644 --- a/DashAI/front/src/contexts/schema.js +++ b/DashAI/front/src/contexts/schema.js @@ -154,14 +154,17 @@ export const useFormSchemaStore = () => { if (properties.length === 0) return getModelFromSubform(formValues[property]); + // Walk down the property chain, unwrapping each subform into its params + // map. This must unwrap the last property too: the current view's params + // map is what holds the field being rendered. Without unwrapping the last + // hop, a component field nested inside another component (depth >= 2) reads + // undefined and its select never reflects a selection. let params = null; for (const prop of properties) { - if (params === null) { - params = formValues[prop.key]; - continue; - } - - params = getParamsFromSubform(params)[prop.key]; + params = + params === null + ? getParamsFromSubform(formValues[prop.key]) + : getParamsFromSubform(params[prop.key]); } return getModelFromSubform(params[property]); }; From a1ef348a941f8ad7948a71e797491096fb58afd6 Mon Sep 17 00:00:00 2001 From: Irozuku Date: Fri, 3 Jul 2026 11:20:03 -0400 Subject: [PATCH 58/89] fix: add root crumb so nested config can return to the top model The breadcrumbs never rendered a clickable root, so a nested submodel had no direct way back to the top level model. Add a root crumb that pops every nested property. --- .../shared/FormSchemaBreadScrumbs.jsx | 27 +++++++++++++++++-- .../shared/FormSchemaWithSelectedModel.jsx | 2 +- 2 files changed, 26 insertions(+), 3 deletions(-) diff --git a/DashAI/front/src/components/shared/FormSchemaBreadScrumbs.jsx b/DashAI/front/src/components/shared/FormSchemaBreadScrumbs.jsx index c6b608302..f0249a24b 100644 --- a/DashAI/front/src/components/shared/FormSchemaBreadScrumbs.jsx +++ b/DashAI/front/src/components/shared/FormSchemaBreadScrumbs.jsx @@ -1,22 +1,40 @@ import React from "react"; +import PropTypes from "prop-types"; import Breadcrumbs from "@mui/material/Breadcrumbs"; import Typography from "@mui/material/Typography"; import Link from "@mui/material/Link"; import { useTheme } from "@mui/material/styles"; +import { useTranslation } from "react-i18next"; import { useFormSchemaStore } from "../../contexts/schema"; /** * This component is the breadcrumbs for the form schema + * @param {string} rootLabel - Label for the root (top level) model crumb */ -function FormSchemaBreadScrumbs() { +function FormSchemaBreadScrumbs({ rootLabel }) { const theme = useTheme(); + const { t } = useTranslation(["common"]); const { properties, removeLastProperty } = useFormSchemaStore(); const handleRemoveLastProperty = (index) => { removeLastProperty(properties.length - 1 - index); }; + // Root crumb: pops every nested property to return to the top level model. + const rootCrumb = ( + removeLastProperty(properties.length)} + sx={{ background: "none", border: "none", cursor: "pointer" }} + > + {rootLabel || t("common:model")} + + ); + const linkedProperties = properties .slice(0, properties.length - 1) .map((property, index) => ( @@ -33,7 +51,8 @@ function FormSchemaBreadScrumbs() { )); return ( - + + {rootCrumb} {linkedProperties} {properties[properties.length - 1]?.label} @@ -42,4 +61,8 @@ function FormSchemaBreadScrumbs() { ); } +FormSchemaBreadScrumbs.propTypes = { + rootLabel: PropTypes.string, +}; + export default FormSchemaBreadScrumbs; diff --git a/DashAI/front/src/components/shared/FormSchemaWithSelectedModel.jsx b/DashAI/front/src/components/shared/FormSchemaWithSelectedModel.jsx index 67357e678..41721ec87 100644 --- a/DashAI/front/src/components/shared/FormSchemaWithSelectedModel.jsx +++ b/DashAI/front/src/components/shared/FormSchemaWithSelectedModel.jsx @@ -74,7 +74,7 @@ function FormSchemaWithSelectedModel({ > {Boolean(propertyData?.parent) && ( <> - + Date: Fri, 3 Jul 2026 11:20:28 -0400 Subject: [PATCH 59/89] feat: disable Next when a nested component is not downloaded The add model dialog resolves required nested downloads and disables the Next button until they are present, with a tooltip listing them. It re-checks after an inline download finishes anywhere. --- DashAI/front/src/api/component.ts | 21 +++++++++ .../src/components/models/AddModelDialog.jsx | 45 ++++++++++++++++++- .../src/utils/i18n/locales/de/common.json | 3 +- .../src/utils/i18n/locales/en/common.json | 3 +- .../src/utils/i18n/locales/es/common.json | 3 +- .../src/utils/i18n/locales/pt/common.json | 3 +- .../src/utils/i18n/locales/zh/common.json | 3 +- 7 files changed, 74 insertions(+), 7 deletions(-) diff --git a/DashAI/front/src/api/component.ts b/DashAI/front/src/api/component.ts index a4008831c..0e1a42fd7 100644 --- a/DashAI/front/src/api/component.ts +++ b/DashAI/front/src/api/component.ts @@ -77,3 +77,24 @@ export const getComponentDownloadStatus = async ( }>(`/v1/component/${name}/download`); return response.data; }; + +export interface RequiredDownload { + name: string; + display_name: string; + parent: string | null; + download_size_bytes: number | null; +} + +// Resolve which components a configuration still needs downloaded. Walks +// nested component parameters server-side (a component selected as another +// component's parameter) and optionally checks the parent model itself. +export const getRequiredDownloads = async ( + parameters: Record, + modelName?: string, +): Promise => { + const response = await api.post( + `/v1/component/downloads/required`, + { model_name: modelName ?? null, parameters }, + ); + return response.data; +}; diff --git a/DashAI/front/src/components/models/AddModelDialog.jsx b/DashAI/front/src/components/models/AddModelDialog.jsx index 2bd46d5c1..144a57e91 100644 --- a/DashAI/front/src/components/models/AddModelDialog.jsx +++ b/DashAI/front/src/components/models/AddModelDialog.jsx @@ -24,6 +24,8 @@ import ModelsTableSelectMetric from "./modelSession/ModelsTableSelectMetric"; import useSchema from "../../hooks/useSchema"; import { generateSequentialName } from "../../utils/nameGenerator"; import { createRun } from "../../api/run"; +import { getRequiredDownloads } from "../../api/component"; +import { subscribeAnyDownloadState } from "./model/ComponentDownloadControl"; import { useTranslation } from "react-i18next"; import { useTourContext } from "../tour/TourProvider"; import { checkIfHaveOptimazers } from "../../utils/schema"; @@ -54,6 +56,7 @@ function AddModelDialog({ const [goalMetric, setGoalMetric] = useState(""); const [hasLoadedInitialParams, setHasLoadedInitialParams] = useState(false); const [modelDownloaded, setModelDownloaded] = useState(true); + const [missingNested, setMissingNested] = useState([]); const { t } = useTranslation(["models", "common"]); const { defaultValues: defaultModelParams } = useSchema({ @@ -106,6 +109,37 @@ function AddModelDialog({ setModelDownloaded(!requiresDownload || isDownloaded); }, [preselectedModelObject]); + // Block advancing while any component selected inside the model parameters + // still needs downloading. The check walks the nested parameters server-side + // and re-runs after an inline download/delete finishes anywhere. + useEffect(() => { + if ( + !open || + activeStep !== 0 || + !modelParameters || + Object.keys(modelParameters).length === 0 + ) { + setMissingNested([]); + return; + } + let cancelled = false; + const check = async () => { + try { + const missing = await getRequiredDownloads(modelParameters); + if (!cancelled) setMissingNested(missing); + } catch { + if (!cancelled) setMissingNested([]); + } + }; + const timer = setTimeout(check, 300); + const unsubscribe = subscribeAnyDownloadState(() => check()); + return () => { + cancelled = true; + clearTimeout(timer); + unsubscribe(); + }; + }, [open, activeStep, JSON.stringify(modelParameters)]); + useEffect(() => { if ( selectedModel && @@ -387,7 +421,13 @@ function AddModelDialog({ preselectedModelObject?.metadata?.requires_download && !modelDownloaded ? t("common:componentDownload.mustDownload") - : "" + : activeStep === 0 && missingNested.length > 0 + ? t("common:componentDownload.mustDownloadNested", { + names: missingNested + .map((m) => m.display_name || m.name) + .join(", "), + }) + : "" } > @@ -403,7 +443,8 @@ function AddModelDialog({ Boolean( preselectedModelObject?.metadata?.requires_download, ) && - !modelDownloaded) + !modelDownloaded) || + (activeStep === 0 && missingNested.length > 0) } > {activeStep === steps.length - 1 diff --git a/DashAI/front/src/utils/i18n/locales/de/common.json b/DashAI/front/src/utils/i18n/locales/de/common.json index 07e9df0e7..06964acb6 100644 --- a/DashAI/front/src/utils/i18n/locales/de/common.json +++ b/DashAI/front/src/utils/i18n/locales/de/common.json @@ -167,7 +167,8 @@ "done": "Komponente heruntergeladen", "deleted": "Download geloescht", "failed": "Download der Komponente fehlgeschlagen", - "mustDownload": "Dieses Modell muss vor der Nutzung heruntergeladen werden" + "mustDownload": "Dieses Modell muss vor der Nutzung heruntergeladen werden", + "mustDownloadNested": "Diese Komponenten müssen zuerst heruntergeladen werden: {{names}}" }, "jobQueue": { "title": "Aufgabenwarteschlange", diff --git a/DashAI/front/src/utils/i18n/locales/en/common.json b/DashAI/front/src/utils/i18n/locales/en/common.json index 478d80af9..2b37fd4b4 100644 --- a/DashAI/front/src/utils/i18n/locales/en/common.json +++ b/DashAI/front/src/utils/i18n/locales/en/common.json @@ -167,7 +167,8 @@ "done": "Component downloaded", "deleted": "Download deleted", "failed": "Component download failed", - "mustDownload": "This model must be downloaded before use" + "mustDownload": "This model must be downloaded before use", + "mustDownloadNested": "These components must be downloaded first: {{names}}" }, "jobQueue": { "title": "Job Queue", diff --git a/DashAI/front/src/utils/i18n/locales/es/common.json b/DashAI/front/src/utils/i18n/locales/es/common.json index f22b31b60..33ea478e0 100644 --- a/DashAI/front/src/utils/i18n/locales/es/common.json +++ b/DashAI/front/src/utils/i18n/locales/es/common.json @@ -167,7 +167,8 @@ "done": "Componente descargado", "deleted": "Descarga eliminada", "failed": "La descarga del componente ha fallado", - "mustDownload": "Este modelo debe descargarse antes de usarlo" + "mustDownload": "Este modelo debe descargarse antes de usarlo", + "mustDownloadNested": "Estos componentes deben descargarse primero: {{names}}" }, "jobQueue": { "title": "Cola de trabajos", diff --git a/DashAI/front/src/utils/i18n/locales/pt/common.json b/DashAI/front/src/utils/i18n/locales/pt/common.json index 583d45dec..827390169 100644 --- a/DashAI/front/src/utils/i18n/locales/pt/common.json +++ b/DashAI/front/src/utils/i18n/locales/pt/common.json @@ -167,7 +167,8 @@ "done": "Componente baixado", "deleted": "Download removido", "failed": "Falha ao baixar o componente", - "mustDownload": "Este modelo precisa ser baixado antes de usar" + "mustDownload": "Este modelo precisa ser baixado antes de usar", + "mustDownloadNested": "Estes componentes precisam ser baixados primeiro: {{names}}" }, "jobQueue": { "title": "Fila de trabalhos", diff --git a/DashAI/front/src/utils/i18n/locales/zh/common.json b/DashAI/front/src/utils/i18n/locales/zh/common.json index 02985d5fb..14b8fcba3 100644 --- a/DashAI/front/src/utils/i18n/locales/zh/common.json +++ b/DashAI/front/src/utils/i18n/locales/zh/common.json @@ -167,7 +167,8 @@ "done": "组件已下载", "deleted": "下载已删除", "failed": "组件下载失败", - "mustDownload": "使用前必须先下载此模型" + "mustDownload": "使用前必须先下载此模型", + "mustDownloadNested": "必须先下载这些组件：{{names}}" }, "jobQueue": { "title": "任务队列", From 3ad94d9904b8fd107983e6b31ed5e14ce1aed936 Mon Sep 17 00:00:00 2001 From: Irozuku Date: Fri, 3 Jul 2026 11:23:47 -0400 Subject: [PATCH 60/89] fix: keep download button at normal color on undownloaded cards Dim only the card content, not the whole card, so the inline download control does not inherit the card's dimmed opacity. --- .../src/components/custom/ComponentSelector.jsx | 13 +++++++++++-- 1 file changed, 11 insertions(+), 2 deletions(-) diff --git a/DashAI/front/src/components/custom/ComponentSelector.jsx b/DashAI/front/src/components/custom/ComponentSelector.jsx index d3c00539d..3a1843e90 100644 --- a/DashAI/front/src/components/custom/ComponentSelector.jsx +++ b/DashAI/front/src/components/custom/ComponentSelector.jsx @@ -131,14 +131,23 @@ function ComponentSelector({ border: 1, borderColor: isSelected ? "primary.main" : "divider", bgcolor: isSelected ? "action.selected" : "background.paper", - opacity: needsDownload ? 0.6 : 1, transition: "border-color 0.15s, background 0.15s", "&:hover": { borderColor: needsDownload ? "divider" : "secondary.main", }, }} > - + {/* Dim only the card content while a download is required, so the + download control below keeps its normal color (CSS opacity on the + card would otherwise cap the button's opacity too). */} + {icon && ( Date: Fri, 3 Jul 2026 11:31:46 -0400 Subject: [PATCH 61/89] feat: allow selecting undownloaded components to preview their description Component cards stay clickable when a download is required (keeping the dimmed content and inline download button), so their description can be viewed. The generative Next and Create buttons stay disabled until the selected model is downloaded, and an undownloaded model is no longer auto-deselected. --- .../src/components/custom/ComponentSelector.jsx | 6 +++--- .../components/generative/CreateSessionCenter.jsx | 15 +++++++++++++-- .../generative/CreateSessionContext.jsx | 7 ++++--- 3 files changed, 20 insertions(+), 8 deletions(-) diff --git a/DashAI/front/src/components/custom/ComponentSelector.jsx b/DashAI/front/src/components/custom/ComponentSelector.jsx index 3a1843e90..c2a399dec 100644 --- a/DashAI/front/src/components/custom/ComponentSelector.jsx +++ b/DashAI/front/src/components/custom/ComponentSelector.jsx @@ -120,20 +120,20 @@ function ComponentSelector({ handleSelect(component)} + onClick={() => handleSelect(component)} data-tour={isCsvComponent ? tourDataFor : undefined} sx={{ p: 3, display: "flex", flexDirection: "column", gap: 3, - cursor: needsDownload ? "not-allowed" : "pointer", + cursor: "pointer", border: 1, borderColor: isSelected ? "primary.main" : "divider", bgcolor: isSelected ? "action.selected" : "background.paper", transition: "border-color 0.15s, background 0.15s", "&:hover": { - borderColor: needsDownload ? "divider" : "secondary.main", + borderColor: "secondary.main", }, }} > diff --git a/DashAI/front/src/components/generative/CreateSessionCenter.jsx b/DashAI/front/src/components/generative/CreateSessionCenter.jsx index 6a455aa80..f43d35f2c 100644 --- a/DashAI/front/src/components/generative/CreateSessionCenter.jsx +++ b/DashAI/front/src/components/generative/CreateSessionCenter.jsx @@ -68,9 +68,20 @@ export default function CreateSessionCenter() { } }, [step]); - const canGoNext = !!selectedModel; + // Read the download status from the (in place updated) models list so the + // gate reacts to an inline download without needing selectedModel to change. + const selectedModelState = + models.find((m) => m.name === selectedModel?.name) || selectedModel; + const selectedNeedsDownload = + Boolean(selectedModelState?.metadata?.requires_download) && + !selectedModelState?.downloaded; + + const canGoNext = !!selectedModel && !selectedNeedsDownload; const canCreate = - !!selectedModel && !!formik.values.name?.trim() && !submitting; + !!selectedModel && + !selectedNeedsDownload && + !!formik.values.name?.trim() && + !submitting; return ( { if (!selectedModel) return; const match = models.find((m) => m.name === selectedModel.name); - if (match && isUnavailable(match)) setSelectedModel(null); + if (!match) setSelectedModel(null); }, [models]); const handleNext = () => { From 67c2899b3032b5baa180eb864e84081775c069ec Mon Sep 17 00:00:00 2001 From: Irozuku Date: Fri, 3 Jul 2026 14:33:46 -0400 Subject: [PATCH 62/89] refactor: simplify download control and use a plain status icon in the models side bar Extract useComponentDownloadState and deleteComponent from ComponentDownloadControl and drop its compact variant, leaving a single button used by the selectors. The models side bar now downloads on row click and shows a tooltip-free status icon (download, spinner, or delete). --- .../src/components/models/ModelsRightBar.jsx | 34 ++- .../models/model/ComponentDownloadControl.jsx | 193 ++++++++------- .../model/ComponentDownloadControl.test.jsx | 27 +-- .../models/model/ModelDownloadStatusIcon.jsx | 78 +++++++ .../model/ModelDownloadStatusIcon.test.jsx | 50 ++++ .../components/models/model/ModelListItem.jsx | 220 ++++++++---------- 6 files changed, 343 insertions(+), 259 deletions(-) create mode 100644 DashAI/front/src/components/models/model/ModelDownloadStatusIcon.jsx create mode 100644 DashAI/front/src/components/models/model/ModelDownloadStatusIcon.test.jsx diff --git a/DashAI/front/src/components/models/ModelsRightBar.jsx b/DashAI/front/src/components/models/ModelsRightBar.jsx index 85be2c054..bb00620c7 100644 --- a/DashAI/front/src/components/models/ModelsRightBar.jsx +++ b/DashAI/front/src/components/models/ModelsRightBar.jsx @@ -7,7 +7,8 @@ import { useSnackbar } from "notistack"; import SideBar from "../threeSectionLayout/panelContainers/SideBar"; import { getComponents } from "../../api/component"; import ModelListItem from "./model/ModelListItem"; -import ComponentDownloadControl from "./model/ComponentDownloadControl"; +import { startComponentDownload } from "./model/ComponentDownloadControl"; +import ModelDownloadStatusIcon from "./model/ModelDownloadStatusIcon"; import { useTranslation } from "react-i18next"; import { useTourContext } from "../tour/TourProvider"; import { useModels } from "./ModelsContext"; @@ -93,6 +94,12 @@ export default function ModelsRightBar({ onToggle }) { // A download-required model that has not been downloaded cannot be // configured; it is blocked in the list with an inline download control. if (model.metadata?.requires_download && !model.downloaded) { + startComponentDownload({ + component: model, + enqueueSnackbar, + t, + onStatusChange: () => fetchModels(), + }); return; } selectModel(model); @@ -250,29 +257,16 @@ export default function ModelsRightBar({ onToggle }) { return ( handleModelClick(model) - } + onClick={() => handleModelClick(model)} + onDisabledClick={() => handleModelClick(model)} data-tour={index === 0 ? "first-model" : undefined} action={ requiresDownload ? ( - fetchModels()} + fetchModels()} /> ) : null } diff --git a/DashAI/front/src/components/models/model/ComponentDownloadControl.jsx b/DashAI/front/src/components/models/model/ComponentDownloadControl.jsx index 8ad18a8b1..7cf74afde 100644 --- a/DashAI/front/src/components/models/model/ComponentDownloadControl.jsx +++ b/DashAI/front/src/components/models/model/ComponentDownloadControl.jsx @@ -1,13 +1,5 @@ -import React, { useState, useEffect, useRef } from "react"; -import { - Box, - Button, - CircularProgress, - IconButton, - LinearProgress, - Tooltip, - Typography, -} from "@mui/material"; +import React, { useState, useEffect } from "react"; +import { Box, Button, LinearProgress, Typography } from "@mui/material"; import DownloadIcon from "@mui/icons-material/Download"; import DeleteIcon from "@mui/icons-material/Delete"; import { useTranslation } from "react-i18next"; @@ -35,6 +27,7 @@ const formatSize = (bytes) => { const downloadListeners = new Map(); // name -> Set<(state) => void> const downloadStateCache = new Map(); // name -> { downloading, downloaded } const anyChangeListeners = new Set(); // (name, state) => void +const activePollers = new Map(); // name -> poller id const subscribeDownloadState = (name, listener) => { let listeners = downloadListeners.get(name); @@ -65,20 +58,71 @@ const broadcastDownloadState = (name, state) => { anyChangeListeners.forEach((listener) => listener(name, state)); }; -const ComponentDownloadControl = ({ +export const stopComponentDownloadPolling = (componentName) => { + const pollerId = activePollers.get(componentName); + if (pollerId != null) { + stopJobPolling(pollerId); + activePollers.delete(componentName); + } +}; + +export const startComponentDownload = async ({ component, + enqueueSnackbar, + t, onStatusChange, - compact = false, }) => { - const { t } = useTranslation(["common"]); - const { enqueueSnackbar } = useSnackbar(); - const meta = component.metadata || {}; + broadcastDownloadState(component.name, { downloading: true }); + try { + const { id } = await downloadComponent(component.name); + activePollers.set(component.name, id); + startJobPolling( + id, + async () => { + activePollers.delete(component.name); + const status = await getComponentDownloadStatus(component.name); + broadcastDownloadState(component.name, { + downloading: false, + downloaded: status.downloaded, + }); + if (onStatusChange) onStatusChange(status.downloaded); + enqueueSnackbar(t("common:componentDownload.done"), { + variant: "success", + }); + }, + () => { + activePollers.delete(component.name); + broadcastDownloadState(component.name, { + downloading: false, + downloaded: false, + }); + if (onStatusChange) onStatusChange(false); + enqueueSnackbar(t("common:componentDownload.failed"), { + variant: "error", + }); + }, + ); + } catch (e) { + broadcastDownloadState(component.name, { + downloading: false, + downloaded: false, + }); + if (onStatusChange) onStatusChange(false); + enqueueSnackbar(t("common:componentDownload.failed"), { + variant: "error", + }); + } +}; + +// Subscribe a component to the shared download state. Returns the live +// { downloaded, downloading } flags, kept in sync across every mounted control +// for the same component name. +export const useComponentDownloadState = (component) => { const cached = downloadStateCache.get(component.name); const [downloaded, setDownloaded] = useState( cached?.downloaded ?? Boolean(component.downloaded), ); const [downloading, setDownloading] = useState(cached?.downloading ?? false); - const pollerIdRef = useRef(null); useEffect(() => { const known = downloadStateCache.get(component.name); @@ -86,7 +130,6 @@ const ComponentDownloadControl = ({ setDownloading(known?.downloading ?? false); }, [component.name, component.downloaded]); - // Mirror download/delete triggered by any other control for this component. useEffect(() => { return subscribeDownloadState(component.name, (state) => { if (state.downloading !== undefined) setDownloading(state.downloading); @@ -94,96 +137,46 @@ const ComponentDownloadControl = ({ }); }, [component.name]); - useEffect(() => { - return () => { - if (pollerIdRef.current != null) stopJobPolling(pollerIdRef.current); - }; - }, []); - - if (!meta.requires_download) return null; + return { downloaded, downloading }; +}; - const finish = (isDownloaded) => { +// Delete a component's download and broadcast the new state to every control. +export const deleteComponent = async ({ + component, + enqueueSnackbar, + t, + onStatusChange, +}) => { + try { + await deleteComponentDownload(component.name); broadcastDownloadState(component.name, { downloading: false, - downloaded: isDownloaded, + downloaded: false, }); - if (onStatusChange) onStatusChange(isDownloaded); - }; + if (onStatusChange) onStatusChange(false); + enqueueSnackbar(t("common:componentDownload.deleted"), { + variant: "success", + }); + } catch { + enqueueSnackbar(t("common:componentDownload.failed"), { + variant: "error", + }); + } +}; - const handleDownload = async () => { - broadcastDownloadState(component.name, { downloading: true }); - try { - const { id } = await downloadComponent(component.name); - pollerIdRef.current = id; - startJobPolling( - id, - async () => { - pollerIdRef.current = null; - const status = await getComponentDownloadStatus(component.name); - finish(status.downloaded); - enqueueSnackbar(t("common:componentDownload.done"), { - variant: "success", - }); - }, - () => { - pollerIdRef.current = null; - finish(false); - enqueueSnackbar(t("common:componentDownload.failed"), { - variant: "error", - }); - }, - ); - } catch (e) { - finish(false); - enqueueSnackbar(t("common:componentDownload.failed"), { - variant: "error", - }); - } - }; +const ComponentDownloadControl = ({ component, onStatusChange }) => { + const { t } = useTranslation(["common"]); + const { enqueueSnackbar } = useSnackbar(); + const meta = component.metadata || {}; + const { downloaded, downloading } = useComponentDownloadState(component); - const handleDelete = async () => { - try { - await deleteComponentDownload(component.name); - finish(false); - enqueueSnackbar(t("common:componentDownload.deleted"), { - variant: "success", - }); - } catch { - enqueueSnackbar(t("common:componentDownload.failed"), { - variant: "error", - }); - } - }; + if (!meta.requires_download) return null; - if (compact) { - if (downloading) { - return ( - - - - ); - } - if (downloaded) { - return ( - - - - - - ); - } - return ( - - - - - - ); - } + const handleDownload = () => + startComponentDownload({ component, enqueueSnackbar, t, onStatusChange }); + + const handleDelete = () => + deleteComponent({ component, enqueueSnackbar, t, onStatusChange }); if (downloading) { return ( diff --git a/DashAI/front/src/components/models/model/ComponentDownloadControl.test.jsx b/DashAI/front/src/components/models/model/ComponentDownloadControl.test.jsx index 9bd72583a..e8388119f 100644 --- a/DashAI/front/src/components/models/model/ComponentDownloadControl.test.jsx +++ b/DashAI/front/src/components/models/model/ComponentDownloadControl.test.jsx @@ -42,26 +42,17 @@ describe("ComponentDownloadControl", () => { ); }); - it("compact mode triggers download from an icon button", async () => { - renderWithProviders( - {}} - />, - ); - const button = await screen.findByRole("button", { name: /download/i }); - fireEvent.click(button); - await waitFor(() => - expect(downloadComponent).toHaveBeenCalledWith("OpusMtEnRoaTransformer"), - ); - }); - it("shows a delete control for a downloaded component and deletes it", async () => { + // A distinct name avoids the module-level download-state cache carrying + // over from the download test above. + const downloadedComponent = { + ...component, + name: "OpusMtEnRoaTransformerDownloaded", + downloaded: true, + }; renderWithProviders( {}} />, ); @@ -69,7 +60,7 @@ describe("ComponentDownloadControl", () => { fireEvent.click(button); await waitFor(() => expect(deleteComponentDownload).toHaveBeenCalledWith( - "OpusMtEnRoaTransformer", + "OpusMtEnRoaTransformerDownloaded", ), ); }); diff --git a/DashAI/front/src/components/models/model/ModelDownloadStatusIcon.jsx b/DashAI/front/src/components/models/model/ModelDownloadStatusIcon.jsx new file mode 100644 index 000000000..9fd77a39b --- /dev/null +++ b/DashAI/front/src/components/models/model/ModelDownloadStatusIcon.jsx @@ -0,0 +1,78 @@ +import React from "react"; +import PropTypes from "prop-types"; +import { Box, CircularProgress } from "@mui/material"; +import DownloadIcon from "@mui/icons-material/Download"; +import DeleteIcon from "@mui/icons-material/Delete"; +import { useSnackbar } from "notistack"; +import { useTranslation } from "react-i18next"; +import { + useComponentDownloadState, + deleteComponent, +} from "./ComponentDownloadControl"; + +/** + * Compact, tooltip-free download status shown at the end of a model row in the + * models side bar. The row click starts the download, so this is only an + * indicator: a spinner while downloading, a delete icon for a downloaded model, + * and a plain download icon otherwise. + * @param {object} model - The model component dict. + * @param {function} onChanged - Called after a delete so the list can refresh. + */ +export default function ModelDownloadStatusIcon({ model, onChanged }) { + const { t } = useTranslation(["common"]); + const { enqueueSnackbar } = useSnackbar(); + const { downloaded, downloading } = useComponentDownloadState(model); + + if (!model.metadata?.requires_download) return null; + + if (downloading) { + return ; + } + + if (downloaded) { + return ( + { + e.stopPropagation(); + deleteComponent({ + component: model, + enqueueSnackbar, + t, + onStatusChange: onChanged, + }); + }} + sx={{ + display: "flex", + alignItems: "center", + color: "error.main", + cursor: "pointer", + }} + > + + + ); + } + + // The row click handles the download; the icon is a non-interactive hint. + return ( + + + + ); +} + +ModelDownloadStatusIcon.propTypes = { + model: PropTypes.object.isRequired, + onChanged: PropTypes.func, +}; diff --git a/DashAI/front/src/components/models/model/ModelDownloadStatusIcon.test.jsx b/DashAI/front/src/components/models/model/ModelDownloadStatusIcon.test.jsx new file mode 100644 index 000000000..05e7827e7 --- /dev/null +++ b/DashAI/front/src/components/models/model/ModelDownloadStatusIcon.test.jsx @@ -0,0 +1,50 @@ +import React from "react"; +import { screen, fireEvent, waitFor } from "@testing-library/react"; +import { renderWithProviders } from "../../../test-utils/renderWithProviders"; + +jest.mock("../../../api/component", () => ({ + downloadComponent: jest.fn(() => Promise.resolve({ id: "job-1" })), + deleteComponentDownload: jest.fn(() => Promise.resolve()), + getComponentDownloadStatus: jest.fn(() => + Promise.resolve({ downloaded: false, requires_download: true }), + ), +})); +jest.mock("../../../utils/jobPoller", () => ({ + startJobPolling: jest.fn(), + stopJobPolling: jest.fn(), + subscribeJobs: jest.fn(() => () => {}), +})); + +import ModelDownloadStatusIcon from "./ModelDownloadStatusIcon"; +import { deleteComponentDownload } from "../../../api/component"; + +const model = { + name: "DummyDownloadableClassifier", + downloaded: false, + metadata: { requires_download: true, download_size_bytes: 268435456 }, +}; + +describe("ModelDownloadStatusIcon", () => { + it("renders no interactive control for an undownloaded model", () => { + renderWithProviders( + {}} />, + ); + expect(screen.queryByRole("button")).toBeNull(); + }); + + it("deletes a downloaded model from a plain delete icon", async () => { + renderWithProviders( + {}} + />, + ); + const del = await screen.findByRole("button", { name: /delete/i }); + fireEvent.click(del); + await waitFor(() => + expect(deleteComponentDownload).toHaveBeenCalledWith( + "DummyDownloadableClassifier", + ), + ); + }); +}); diff --git a/DashAI/front/src/components/models/model/ModelListItem.jsx b/DashAI/front/src/components/models/model/ModelListItem.jsx index dd7a07d29..d04e47212 100644 --- a/DashAI/front/src/components/models/model/ModelListItem.jsx +++ b/DashAI/front/src/components/models/model/ModelListItem.jsx @@ -1,5 +1,5 @@ import React, { useState } from "react"; -import { Box, Typography, Tooltip } from "@mui/material"; +import { Box, Typography } from "@mui/material"; import { useTheme } from "@mui/material/styles"; import HoverModelInfo from "./HoverModelInfo"; import { ModelIcon } from "./ModelIcon"; @@ -9,6 +9,7 @@ export default function ModelListItem({ model, disabled = false, onClick, + onDisabledClick, action = null, ...props }) { @@ -17,10 +18,8 @@ export default function ModelListItem({ const [hoveredModel, setHoveredModel] = useState(null); const handleMouseEnter = (event, model) => { - if (!disabled) { - setAnchorEl(event.currentTarget); - setHoveredModel(model); - } + setAnchorEl(event.currentTarget); + setHoveredModel(model); }; const handleMouseLeave = () => { @@ -28,6 +27,16 @@ export default function ModelListItem({ setHoveredModel(null); }; + const handleCardClick = (event) => { + if (disabled) { + if (onDisabledClick) onDisabledClick(event); + return; + } + if (onClick) onClick(event); + }; + + const isClickable = Boolean(onClick || onDisabledClick); + // Get color and icon from metadata or use defaults const color = model.color || model.metadata?.color || theme.palette.text.secondary; @@ -35,146 +44,115 @@ export default function ModelListItem({ return ( <> - { + e.dataTransfer.setData( + "application/x-dashai-model", + JSON.stringify(model), + ); + e.dataTransfer.effectAllowed = "copy"; + setCustomDragImage(e); + } + : undefined + } + onMouseEnter={(e) => handleMouseEnter(e, model)} + onMouseLeave={handleMouseLeave} + onClick={handleCardClick} + {...props} + sx={{ + display: "flex", + alignItems: "center", + gap: 3, + p: 3, + bgcolor: disabled ? theme.palette.ui.disabled : theme.palette.ui.box, + border: `1px solid ${theme.palette.ui.border}`, + borderRadius: 1, + cursor: isClickable ? "pointer" : "default", + transition: "all 0.2s", + position: "relative", + "&:hover": { + bgcolor: disabled + ? theme.palette.ui.disabled + : theme.palette.action.hover, + borderColor: disabled ? theme.palette.ui.border : color, + transform: disabled || !isClickable ? "none" : "translateX(4px)", }, + "&::after": disabled + ? { + content: '""', + position: "absolute", + inset: 0, + borderRadius: 1, + pointerEvents: "none", + background: + "repeating-linear-gradient(45deg, transparent, transparent 10px, rgba(0, 0, 0, 0.1) 10px, rgba(0, 0, 0, 0.1) 20px)", + } + : {}, }} > + {/* Icon */} { - e.dataTransfer.setData( - "application/x-dashai-model", - JSON.stringify(model), - ); - e.dataTransfer.effectAllowed = "copy"; - setCustomDragImage(e); - } - : undefined - } - onMouseEnter={(e) => handleMouseEnter(e, model)} - onMouseLeave={handleMouseLeave} - onClick={disabled ? null : onClick} - {...props} sx={{ display: "flex", alignItems: "center", - gap: 3, - p: 3, + justifyContent: "center", + width: 36, + height: 36, + borderRadius: 1, bgcolor: disabled ? theme.palette.ui.disabled - : theme.palette.ui.box, - border: `1px solid ${theme.palette.ui.border}`, - borderRadius: 1, - cursor: disabled ? "not-allowed" : "grab", - transition: "all 0.2s", - opacity: disabled ? 0.5 : 1, - filter: disabled ? "grayscale(0.6)" : "none", - position: "relative", - "&:hover": { - bgcolor: disabled - ? theme.palette.ui.disabled - : theme.palette.action.hover, - borderColor: disabled ? theme.palette.ui.border : color, - transform: disabled ? "none" : "translateX(4px)", - }, - "&::after": disabled - ? { - content: '""', - position: "absolute", - inset: 0, - borderRadius: 1, - pointerEvents: "none", - background: - "repeating-linear-gradient(45deg, transparent, transparent 10px, rgba(0, 0, 0, 0.1) 10px, rgba(0, 0, 0, 0.1) 20px)", - } - : {}, + : theme.palette.ui.border, + color: disabled + ? theme.palette.text.disabled + : theme.palette.text.primary, + flexShrink: 0, }} > - {/* Icon */} - + + + {/* Content */} + + - - + {model.display_name || model.name} + + - {/* Content */} - - - {model.display_name || model.name} - + {/* Trailing action (e.g. download/delete control) */} + {action && ( + + {action} - - {/* Trailing action (e.g. download/delete control) */} - {action && ( - e.stopPropagation()} - onMouseDown={(e) => e.stopPropagation()} - onDragStart={(e) => e.stopPropagation()} - sx={{ flexShrink: 0, display: "flex", alignItems: "center" }} - > - {action} - - )} - - - {!disabled && ( + )} + + { - )} + } ); } From 999fda185f9d75c0f599c3c839ce29588474adad Mon Sep 17 00:00:00 2001 From: Irozuku Date: Fri, 3 Jul 2026 14:34:04 -0400 Subject: [PATCH 63/89] feat: show model description and download size on hover in the side bar The hover popover now shows the description for every model, including download-required ones, and appends the download size when present. --- .../models/model/HoverModelInfo.jsx | 26 ++++++++++++++++++- .../src/utils/i18n/locales/de/custom.json | 1 + .../src/utils/i18n/locales/en/custom.json | 1 + .../src/utils/i18n/locales/es/custom.json | 1 + .../src/utils/i18n/locales/pt/custom.json | 1 + .../src/utils/i18n/locales/zh/custom.json | 1 + 6 files changed, 30 insertions(+), 1 deletion(-) diff --git a/DashAI/front/src/components/models/model/HoverModelInfo.jsx b/DashAI/front/src/components/models/model/HoverModelInfo.jsx index 9173ab61f..60a5ba0fb 100644 --- a/DashAI/front/src/components/models/model/HoverModelInfo.jsx +++ b/DashAI/front/src/components/models/model/HoverModelInfo.jsx @@ -3,13 +3,24 @@ import { Box, Typography, Popover } from "@mui/material"; import { useTranslation } from "react-i18next"; import { useTheme } from "@mui/material/styles"; +const formatSize = (bytes) => { + if (bytes == null) return null; + const mb = bytes / 1024 / 1024; + if (mb >= 1024) return `${(mb / 1024).toFixed(1)} GB`; + return `${Math.round(mb)} MB`; +}; + export default function HoverModelInfo({ anchorEl, hoveredModel, handleMouseLeave, }) { - const { t } = useTranslation(["common"]); + const { t } = useTranslation(["common", "custom"]); const theme = useTheme(); + const size = formatSize( + hoveredModel?.metadata?.download_size_bytes || + hoveredModel?.download_size_bytes, + ); return ( + + {size && ( + + {t("custom:modelSize", { size })} + + )} )} diff --git a/DashAI/front/src/utils/i18n/locales/de/custom.json b/DashAI/front/src/utils/i18n/locales/de/custom.json index cda354d45..3e848d208 100644 --- a/DashAI/front/src/utils/i18n/locales/de/custom.json +++ b/DashAI/front/src/utils/i18n/locales/de/custom.json @@ -3,6 +3,7 @@ "selectAnItemToShowInfo": "Element auswählen, um die Beschreibung zu sehen.", "selectInferenceMethods": "Wählen Sie die anzuwendenden Inferenzmethoden", "search": "Suchen", + "modelSize": "Größe: {{size}}", "noItemsFound": "Keine Komponenten gefunden", "tryAdjustingSearch": "Versuchen Sie, Ihre Suche oder Filter anzupassen", "componentsAvailable_one": "{{count}} Komponente verfügbar", diff --git a/DashAI/front/src/utils/i18n/locales/en/custom.json b/DashAI/front/src/utils/i18n/locales/en/custom.json index 7ca90eba3..3e8644117 100644 --- a/DashAI/front/src/utils/i18n/locales/en/custom.json +++ b/DashAI/front/src/utils/i18n/locales/en/custom.json @@ -3,6 +3,7 @@ "selectAnItemToShowInfo": "Select an item to see the description.", "selectInferenceMethods": "Select the inference methods you want to apply", "search": "Search", + "modelSize": "Size: {{size}}", "noItemsFound": "No components found", "tryAdjustingSearch": "Try adjusting your search or filters", "componentsAvailable_one": "{{count}} component available", diff --git a/DashAI/front/src/utils/i18n/locales/es/custom.json b/DashAI/front/src/utils/i18n/locales/es/custom.json index c2e1699b5..ad33b4e0b 100644 --- a/DashAI/front/src/utils/i18n/locales/es/custom.json +++ b/DashAI/front/src/utils/i18n/locales/es/custom.json @@ -3,6 +3,7 @@ "selectAnItemToShowInfo": "Seleccione un elemento para ver la descripción.", "selectInferenceMethods": "Seleccione los métodos de inferencia que desea aplicar", "search": "Buscar", + "modelSize": "Tamaño: {{size}}", "noItemsFound": "No se encontraron componentes", "tryAdjustingSearch": "Intenta ajustar tu búsqueda o filtros", "componentsAvailable_one": "{{count}} componente disponible", diff --git a/DashAI/front/src/utils/i18n/locales/pt/custom.json b/DashAI/front/src/utils/i18n/locales/pt/custom.json index 38a1242a2..054d5b567 100644 --- a/DashAI/front/src/utils/i18n/locales/pt/custom.json +++ b/DashAI/front/src/utils/i18n/locales/pt/custom.json @@ -3,6 +3,7 @@ "selectAnItemToShowInfo": "Selecione um elemento para ver a descrição.", "selectInferenceMethods": "Selecione os métodos de inferência que deseja aplicar", "search": "Buscar", + "modelSize": "Tamanho: {{size}}", "noItemsFound": "Nenhum componente encontrado", "tryAdjustingSearch": "Tente ajustar sua busca ou filtros", "componentsAvailable_one": "{{count}} componente disponível", diff --git a/DashAI/front/src/utils/i18n/locales/zh/custom.json b/DashAI/front/src/utils/i18n/locales/zh/custom.json index 3f033a19d..2b09eb0ea 100644 --- a/DashAI/front/src/utils/i18n/locales/zh/custom.json +++ b/DashAI/front/src/utils/i18n/locales/zh/custom.json @@ -3,6 +3,7 @@ "selectAnItemToShowInfo": "选择一个项目以查看描述。", "selectInferenceMethods": "选择您想要应用的推理方法", "search": "搜索", + "modelSize": "大小：{{size}}", "noItemsFound": "未找到组件", "tryAdjustingSearch": "尝试调整搜索条件或筛选器", "componentsAvailable_one": "{{count}} 个可用组件", From ed784a3967659bdd691a09f617c6ade08729de43 Mon Sep 17 00:00:00 2001 From: Irozuku Date: Fri, 3 Jul 2026 14:39:47 -0400 Subject: [PATCH 64/89] fix: correct model descriptions to reflect explicit downloads Transformer and translation model descriptions said weights download on first use, which no longer holds now that downloads are explicit. Update all five languages to state the weights must be downloaded before use. --- DashAI/back/models/hugging_face/albert_transformer.py | 10 +++++----- .../base_text_classification_transformer.py | 4 ++-- DashAI/back/models/hugging_face/bert_transformer.py | 10 +++++----- DashAI/back/models/hugging_face/bertin_transformer.py | 10 +++++----- DashAI/back/models/hugging_face/beto_transformer.py | 10 +++++----- DashAI/back/models/hugging_face/electra_transformer.py | 10 +++++----- DashAI/back/models/hugging_face/m2m100_transformer.py | 10 +++++----- DashAI/back/models/hugging_face/minilm_transformer.py | 10 +++++----- .../hugging_face/multilingual_bert_transformer.py | 10 +++++----- .../models/hugging_face/opus_mt_en_de_transformer.py | 10 +++++----- .../models/hugging_face/opus_mt_en_es_transformer.py | 10 +++++----- .../models/hugging_face/opus_mt_en_fr_transformer.py | 10 +++++----- .../models/hugging_face/opus_mt_en_pt_transformer.py | 10 +++++----- .../models/hugging_face/opus_mt_es_en_transformer.py | 10 +++++----- .../models/hugging_face/opus_mt_fr_en_transformer.py | 10 +++++----- DashAI/back/models/hugging_face/roberta_transformer.py | 10 +++++----- .../back/models/hugging_face/t5_small_transformer.py | 10 +++++----- .../models/hugging_face/xlm_roberta_transformer.py | 10 +++++----- DashAI/back/models/hugging_face/xlnet_transformer.py | 10 +++++----- 19 files changed, 92 insertions(+), 92 deletions(-) diff --git a/DashAI/back/models/hugging_face/albert_transformer.py b/DashAI/back/models/hugging_face/albert_transformer.py index 104deb58a..8fd152db0 100644 --- a/DashAI/back/models/hugging_face/albert_transformer.py +++ b/DashAI/back/models/hugging_face/albert_transformer.py @@ -34,25 +34,25 @@ class AlbertTransformer(HuggingFaceTextClassificationTransformer): DESCRIPTION: str = MultilingualString( en=( "Parameter efficient BERT variant for English text classification. " - "Downloads weights from Hugging Face on first use (internet required)." + "Download its weights from Hugging Face before use (internet required)." ), es=( "Variante de BERT eficiente en parámetros para clasificación en inglés. " - "Descarga pesos de Hugging Face en el primer uso (requiere internet)." + "Descarga sus pesos de Hugging Face antes de usarlo (requiere internet)." ), pt=( "Variante do BERT eficiente em parâmetros para classificação de " "texto em inglês. " - "Baixa os pesos do Hugging Face no primeiro uso (requer internet)." + "Baixe seus pesos do Hugging Face antes de usar (requer internet)." ), de=( "Parametereffiziente BERT-Variante für englische Textklassifikation. " - "Lädt Gewichte von Hugging Face bei der ersten Verwendung herunter " + "Lädt die Gewichte vor der Nutzung von Hugging Face herunter " "(Internet erforderlich)." ), zh=( "参数高效的 BERT 变体，用于英文文本分类。" - "首次使用时从 Hugging Face 下载权重（需要网络）。" + "使用前需从 Hugging Face 下载权重（需要网络）。" ), ) COLOR: str = "#00838F" diff --git a/DashAI/back/models/hugging_face/base_text_classification_transformer.py b/DashAI/back/models/hugging_face/base_text_classification_transformer.py index dc839c699..a2f9ffbf5 100644 --- a/DashAI/back/models/hugging_face/base_text_classification_transformer.py +++ b/DashAI/back/models/hugging_face/base_text_classification_transformer.py @@ -38,8 +38,8 @@ class HuggingFaceTextClassificationTransformer( - Save/load utilities that preserve custom training parameters. .. note:: - Requires internet access on first use to download pretrained weights - from the Hugging Face Hub. + The pretrained weights must be downloaded from the Hugging Face Hub + (internet access required) before the model can be used. """ MODEL_NAME: str = "" diff --git a/DashAI/back/models/hugging_face/bert_transformer.py b/DashAI/back/models/hugging_face/bert_transformer.py index d22451107..8b2121fea 100644 --- a/DashAI/back/models/hugging_face/bert_transformer.py +++ b/DashAI/back/models/hugging_face/bert_transformer.py @@ -34,24 +34,24 @@ class BertTransformer(HuggingFaceTextClassificationTransformer): DESCRIPTION: str = MultilingualString( en=( "Bidirectional BERT model for English text classification. " - "Downloads weights from Hugging Face on first use (internet required)." + "Download its weights from Hugging Face before use (internet required)." ), es=( "Modelo BERT bidireccional para clasificación de texto en inglés. " - "Descarga pesos de Hugging Face en el primer uso (requiere internet)." + "Descarga sus pesos de Hugging Face antes de usarlo (requiere internet)." ), pt=( "Modelo BERT bidirecional para classificação de texto em inglês. " - "Baixa os pesos do Hugging Face no primeiro uso (requer internet)." + "Baixe seus pesos do Hugging Face antes de usar (requer internet)." ), de=( "Bidirektionales BERT-Modell für englische Textklassifikation. " - "Lädt Gewichte von Hugging Face bei der ersten Verwendung herunter " + "Lädt die Gewichte vor der Nutzung von Hugging Face herunter " "(Internet erforderlich)." ), zh=( "双向 BERT 模型，用于英文文本分类。" - "首次使用时从 Hugging Face 下载权重（需要网络）。" + "使用前需从 Hugging Face 下载权重（需要网络）。" ), ) COLOR: str = "#1565C0" diff --git a/DashAI/back/models/hugging_face/bertin_transformer.py b/DashAI/back/models/hugging_face/bertin_transformer.py index 109a7b649..2a4cfb608 100644 --- a/DashAI/back/models/hugging_face/bertin_transformer.py +++ b/DashAI/back/models/hugging_face/bertin_transformer.py @@ -34,24 +34,24 @@ class BertinTransformer(HuggingFaceTextClassificationTransformer): DESCRIPTION: str = MultilingualString( en=( "Spanish RoBERTa (BERTIN) pretrained on large Spanish corpora. " - "Downloads weights from Hugging Face on first use (internet required)." + "Download its weights from Hugging Face before use (internet required)." ), es=( "RoBERTa en español (BERTIN) preentrenada en grandes corpus en español. " - "Descarga pesos de Hugging Face en el primer uso (requiere internet)." + "Descarga sus pesos de Hugging Face antes de usarlo (requiere internet)." ), pt=( "RoBERTa em espanhol (BERTIN) pré-treinada em grandes corpus em espanhol. " - "Baixa os pesos do Hugging Face no primeiro uso (requer internet)." + "Baixe seus pesos do Hugging Face antes de usar (requer internet)." ), de=( "Spanisches RoBERTa (BERTIN) vortrainiert auf großen spanischen Korpora. " - "Lädt Gewichte von Hugging Face bei der ersten Verwendung herunter " + "Lädt die Gewichte vor der Nutzung von Hugging Face herunter " "(Internet erforderlich)." ), zh=( "在大型西班牙语语料库上预训练的 RoBERTa（BERTIN）模型。" - "首次使用时从 Hugging Face 下载权重（需要网络）。" + "使用前需从 Hugging Face 下载权重（需要网络）。" ), ) COLOR: str = "#AD1457" diff --git a/DashAI/back/models/hugging_face/beto_transformer.py b/DashAI/back/models/hugging_face/beto_transformer.py index db8440718..4c6212137 100644 --- a/DashAI/back/models/hugging_face/beto_transformer.py +++ b/DashAI/back/models/hugging_face/beto_transformer.py @@ -34,24 +34,24 @@ class BetoTransformer(HuggingFaceTextClassificationTransformer): DESCRIPTION: str = MultilingualString( en=( "Spanish BERT (BETO) pretrained on Spanish corpora. " - "Downloads weights from Hugging Face on first use (internet required)." + "Download its weights from Hugging Face before use (internet required)." ), es=( "BERT en español (BETO) preentrenado en corpus en español. " - "Descarga pesos de Hugging Face en el primer uso (requiere internet)." + "Descarga sus pesos de Hugging Face antes de usarlo (requiere internet)." ), pt=( "BERT em espanhol (BETO) pré-treinado em corpus em espanhol. " - "Baixa os pesos do Hugging Face no primeiro uso (requer internet)." + "Baixe seus pesos do Hugging Face antes de usar (requer internet)." ), de=( "Spanisches BERT (BETO) vortrainiert auf spanischen Korpora. " - "Lädt Gewichte von Hugging Face bei der ersten Verwendung herunter " + "Lädt die Gewichte vor der Nutzung von Hugging Face herunter " "(Internet erforderlich)." ), zh=( "在西班牙语语料库上预训练的 BERT（BETO）模型。" - "首次使用时从 Hugging Face 下载权重（需要网络）。" + "使用前需从 Hugging Face 下载权重（需要网络）。" ), ) COLOR: str = "#C62828" diff --git a/DashAI/back/models/hugging_face/electra_transformer.py b/DashAI/back/models/hugging_face/electra_transformer.py index 00e26919e..83867a4d7 100644 --- a/DashAI/back/models/hugging_face/electra_transformer.py +++ b/DashAI/back/models/hugging_face/electra_transformer.py @@ -34,24 +34,24 @@ class ElectraTransformer(HuggingFaceTextClassificationTransformer): DESCRIPTION: str = MultilingualString( en=( "Sample efficient ELECTRA discriminator for text classification. " - "Downloads weights from Hugging Face on first use (internet required)." + "Download its weights from Hugging Face before use (internet required)." ), es=( "Discriminador ELECTRA eficiente en muestras para clasificación de texto. " - "Descarga pesos de Hugging Face en el primer uso (requiere internet)." + "Descarga sus pesos de Hugging Face antes de usarlo (requiere internet)." ), pt=( "Discriminador ELECTRA eficiente em amostras para classificação de texto. " - "Baixa os pesos do Hugging Face no primeiro uso (requer internet)." + "Baixe seus pesos do Hugging Face antes de usar (requer internet)." ), de=( "Stichprobeneffizienter ELECTRA-Diskriminator für Textklassifikation. " - "Lädt Gewichte von Hugging Face bei der ersten Verwendung herunter " + "Lädt die Gewichte vor der Nutzung von Hugging Face herunter " "(Internet erforderlich)." ), zh=( "样本高效的 ELECTRA 判别器，用于文本分类。" - "首次使用时从 Hugging Face 下载权重（需要网络）。" + "使用前需从 Hugging Face 下载权重（需要网络）。" ), ) COLOR: str = "#558B2F" diff --git a/DashAI/back/models/hugging_face/m2m100_transformer.py b/DashAI/back/models/hugging_face/m2m100_transformer.py index 99c023d00..d5eedddd5 100644 --- a/DashAI/back/models/hugging_face/m2m100_transformer.py +++ b/DashAI/back/models/hugging_face/m2m100_transformer.py @@ -131,28 +131,28 @@ class M2M100Transformer(HFPretrainedDownloadMixin, TranslationModel): en=( "Facebook M2M-100 model for direct translation across 100 languages " "using ISO 639-1 codes. " - "Downloads weights from Hugging Face on first use (internet required)." + "Download its weights from Hugging Face before use (internet required)." ), es=( "Modelo M2M-100 de Facebook para traducción directa entre 100 idiomas " "usando códigos ISO 639-1. " - "Descarga pesos de Hugging Face en el primer uso (requiere internet)." + "Descarga sus pesos de Hugging Face antes de usarlo (requiere internet)." ), pt=( "Modelo M2M-100 do Facebook para tradução direta entre 100 idiomas " "usando códigos ISO 639-1. " - "Baixa os pesos do Hugging Face no primeiro uso (requer internet)." + "Baixe seus pesos do Hugging Face antes de usar (requer internet)." ), de=( "Facebook M2M-100-Modell für direkte Übersetzung zwischen 100 Sprachen " "mit ISO 639-1-Codes. " - "Lädt Gewichte von Hugging Face bei der ersten Verwendung herunter " + "Lädt die Gewichte vor der Nutzung von Hugging Face herunter " "(Internet erforderlich)." ), zh=( "Facebook M2M-100 模型，使用 ISO 639-1 代码支持" " 100 种语言之间的直接翻译。" - "首次使用时从 Hugging Face 下载权重（需要网络）。" + "使用前需从 Hugging Face 下载权重（需要网络）。" ), ) COLOR: str = "#6A1B9A" diff --git a/DashAI/back/models/hugging_face/minilm_transformer.py b/DashAI/back/models/hugging_face/minilm_transformer.py index 23a323b4d..818e1a6a0 100644 --- a/DashAI/back/models/hugging_face/minilm_transformer.py +++ b/DashAI/back/models/hugging_face/minilm_transformer.py @@ -34,24 +34,24 @@ class MiniLMTransformer(HuggingFaceTextClassificationTransformer): DESCRIPTION: str = MultilingualString( en=( "Compact, fast MiniLM model for efficient text classification. " - "Downloads weights from Hugging Face on first use (internet required)." + "Download its weights from Hugging Face before use (internet required)." ), es=( "Modelo MiniLM compacto y rápido para clasificación de texto eficiente. " - "Descarga pesos de Hugging Face en el primer uso (requiere internet)." + "Descarga sus pesos de Hugging Face antes de usarlo (requiere internet)." ), pt=( "Modelo MiniLM compacto e rápido para classificação de texto eficiente. " - "Baixa os pesos do Hugging Face no primeiro uso (requer internet)." + "Baixe seus pesos do Hugging Face antes de usar (requer internet)." ), de=( "Kompaktes, schnelles MiniLM-Modell für effiziente Textklassifikation. " - "Lädt Gewichte von Hugging Face bei der ersten Verwendung herunter " + "Lädt die Gewichte vor der Nutzung von Hugging Face herunter " "(Internet erforderlich)." ), zh=( "紧凑快速的 MiniLM 模型，用于高效文本分类。" - "首次使用时从 Hugging Face 下载权重（需要网络）。" + "使用前需从 Hugging Face 下载权重（需要网络）。" ), ) COLOR: str = "#0277BD" diff --git a/DashAI/back/models/hugging_face/multilingual_bert_transformer.py b/DashAI/back/models/hugging_face/multilingual_bert_transformer.py index d012a2e8b..916f6a08f 100644 --- a/DashAI/back/models/hugging_face/multilingual_bert_transformer.py +++ b/DashAI/back/models/hugging_face/multilingual_bert_transformer.py @@ -33,26 +33,26 @@ class MultilingualBertTransformer(HuggingFaceTextClassificationTransformer): DESCRIPTION: str = MultilingualString( en=( "BERT pretrained on 104 languages for multilingual text classification. " - "Downloads weights from Hugging Face on first use (internet required)." + "Download its weights from Hugging Face before use (internet required)." ), es=( "BERT preentrenado en 104 idiomas para clasificación de texto " "multilingüe. " - "Descarga pesos de Hugging Face en el primer uso (requiere internet)." + "Descarga sus pesos de Hugging Face antes de usarlo (requiere internet)." ), pt=( "BERT pré-treinado em 104 idiomas para classificação de texto " "multilingual. " - "Baixa os pesos do Hugging Face no primeiro uso (requer internet)." + "Baixe seus pesos do Hugging Face antes de usar (requer internet)." ), de=( "BERT vortrainiert auf 104 Sprachen für mehrsprachige Textklassifikation. " - "Lädt Gewichte von Hugging Face bei der ersten Verwendung herunter " + "Lädt die Gewichte vor der Nutzung von Hugging Face herunter " "(Internet erforderlich)." ), zh=( "在 104 种语言上预训练的 BERT，用于多语言文本分类。" - "首次使用时从 Hugging Face 下载权重（需要网络）。" + "使用前需从 Hugging Face 下载权重（需要网络）。" ), ) COLOR: str = "#283593" diff --git a/DashAI/back/models/hugging_face/opus_mt_en_de_transformer.py b/DashAI/back/models/hugging_face/opus_mt_en_de_transformer.py index acc031fa9..5950dc5ed 100644 --- a/DashAI/back/models/hugging_face/opus_mt_en_de_transformer.py +++ b/DashAI/back/models/hugging_face/opus_mt_en_de_transformer.py @@ -39,24 +39,24 @@ class OpusMtEnDeTransformer(OpusMtTransformerMixin): DESCRIPTION: str = MultilingualString( en=( "Pretrained transformer for English to German translation. " - "Downloads weights from Hugging Face on first use (internet required)." + "Download its weights from Hugging Face before use (internet required)." ), es=( "Transformer preentrenado para traducción inglés-alemán. " - "Descarga pesos de Hugging Face en el primer uso (requiere internet)." + "Descarga sus pesos de Hugging Face antes de usarlo (requiere internet)." ), pt=( "Transformer pré-treinado para tradução inglês-alemão. " - "Baixa os pesos do Hugging Face no primeiro uso (requer internet)." + "Baixe seus pesos do Hugging Face antes de usar (requer internet)." ), de=( "Vortrainierter Transformer für Englisch-Deutsch-Übersetzung. " - "Lädt Gewichte von Hugging Face bei der ersten Verwendung herunter " + "Lädt die Gewichte vor der Nutzung von Hugging Face herunter " "(Internet erforderlich)." ), zh=( "用于英语到德语翻译的预训练 Transformer。" - "首次使用时从 Hugging Face 下载权重（需要网络）。" + "使用前需从 Hugging Face 下载权重（需要网络）。" ), ) COLOR: str = "#455A64" diff --git a/DashAI/back/models/hugging_face/opus_mt_en_es_transformer.py b/DashAI/back/models/hugging_face/opus_mt_en_es_transformer.py index 82f827345..203c11fbb 100644 --- a/DashAI/back/models/hugging_face/opus_mt_en_es_transformer.py +++ b/DashAI/back/models/hugging_face/opus_mt_en_es_transformer.py @@ -280,24 +280,24 @@ class OpusMtEnESTransformer(OpusMtTransformerMixin): DESCRIPTION: str = MultilingualString( en=( "Pretrained transformer for English to Spanish translation. " - "Downloads weights from Hugging Face on first use (internet required)." + "Download its weights from Hugging Face before use (internet required)." ), es=( "Transformer preentrenado para traducción inglés-español. " - "Descarga pesos de Hugging Face en el primer uso (requiere internet)." + "Descarga sus pesos de Hugging Face antes de usarlo (requiere internet)." ), pt=( "Transformer pré-treinado para tradução inglês-espanhol. " - "Baixa os pesos do Hugging Face no primeiro uso (requer internet)." + "Baixe seus pesos do Hugging Face antes de usar (requer internet)." ), de=( "Vortrainierter Transformer für Englisch-Spanisch-Übersetzung. " - "Lädt Gewichte von Hugging Face bei der ersten Verwendung herunter " + "Lädt die Gewichte vor der Nutzung von Hugging Face herunter " "(Internet erforderlich)." ), zh=( "用于英语到西班牙语翻译的预训练 Transformer。" - "首次使用时从 Hugging Face 下载权重（需要网络）。" + "使用前需从 Hugging Face 下载权重（需要网络）。" ), ) COLOR: str = "#FFA500" diff --git a/DashAI/back/models/hugging_face/opus_mt_en_fr_transformer.py b/DashAI/back/models/hugging_face/opus_mt_en_fr_transformer.py index ba45b6b3d..ad026f97f 100644 --- a/DashAI/back/models/hugging_face/opus_mt_en_fr_transformer.py +++ b/DashAI/back/models/hugging_face/opus_mt_en_fr_transformer.py @@ -39,24 +39,24 @@ class OpusMtEnFrTransformer(OpusMtTransformerMixin): DESCRIPTION: str = MultilingualString( en=( "Pretrained transformer for English to French translation. " - "Downloads weights from Hugging Face on first use (internet required)." + "Download its weights from Hugging Face before use (internet required)." ), es=( "Transformer preentrenado para traducción inglés-francés. " - "Descarga pesos de Hugging Face en el primer uso (requiere internet)." + "Descarga sus pesos de Hugging Face antes de usarlo (requiere internet)." ), pt=( "Transformer pré-treinado para tradução inglês-francês. " - "Baixa os pesos do Hugging Face no primeiro uso (requer internet)." + "Baixe seus pesos do Hugging Face antes de usar (requer internet)." ), de=( "Vortrainierter Transformer für Englisch-Französisch-Übersetzung. " - "Lädt Gewichte von Hugging Face bei der ersten Verwendung herunter " + "Lädt die Gewichte vor der Nutzung von Hugging Face herunter " "(Internet erforderlich)." ), zh=( "用于英语到法语翻译的预训练 Transformer。" - "首次使用时从 Hugging Face 下载权重（需要网络）。" + "使用前需从 Hugging Face 下载权重（需要网络）。" ), ) COLOR: str = "#1976D2" diff --git a/DashAI/back/models/hugging_face/opus_mt_en_pt_transformer.py b/DashAI/back/models/hugging_face/opus_mt_en_pt_transformer.py index 17a1f2076..1837c7ec2 100644 --- a/DashAI/back/models/hugging_face/opus_mt_en_pt_transformer.py +++ b/DashAI/back/models/hugging_face/opus_mt_en_pt_transformer.py @@ -39,24 +39,24 @@ class OpusMtEnPtTransformer(OpusMtTransformerMixin): DESCRIPTION: str = MultilingualString( en=( "Pretrained transformer for English to Portuguese translation. " - "Downloads weights from Hugging Face on first use (internet required)." + "Download its weights from Hugging Face before use (internet required)." ), es=( "Transformer preentrenado para traducción inglés-portugués. " - "Descarga pesos de Hugging Face en el primer uso (requiere internet)." + "Descarga sus pesos de Hugging Face antes de usarlo (requiere internet)." ), pt=( "Transformer pré-treinado para tradução inglês-português. " - "Baixa os pesos do Hugging Face no primeiro uso (requer internet)." + "Baixe seus pesos do Hugging Face antes de usar (requer internet)." ), de=( "Vortrainierter Transformer für Englisch-Portugiesisch-Übersetzung. " - "Lädt Gewichte von Hugging Face bei der ersten Verwendung herunter " + "Lädt die Gewichte vor der Nutzung von Hugging Face herunter " "(Internet erforderlich)." ), zh=( "用于英语到葡萄牙语翻译的预训练 Transformer。" - "首次使用时从 Hugging Face 下载权重（需要网络）。" + "使用前需从 Hugging Face 下载权重（需要网络）。" ), ) COLOR: str = "#2E7D32" diff --git a/DashAI/back/models/hugging_face/opus_mt_es_en_transformer.py b/DashAI/back/models/hugging_face/opus_mt_es_en_transformer.py index c7bcc2e81..4fadc8e1e 100644 --- a/DashAI/back/models/hugging_face/opus_mt_es_en_transformer.py +++ b/DashAI/back/models/hugging_face/opus_mt_es_en_transformer.py @@ -42,24 +42,24 @@ class OpusMtEsENTransformer(OpusMtTransformerMixin): DESCRIPTION: str = MultilingualString( en=( "Pretrained transformer for Spanish to English translation. " - "Downloads weights from Hugging Face on first use (internet required)." + "Download its weights from Hugging Face before use (internet required)." ), es=( "Transformer preentrenado para traducción español-inglés. " - "Descarga pesos de Hugging Face en el primer uso (requiere internet)." + "Descarga sus pesos de Hugging Face antes de usarlo (requiere internet)." ), pt=( "Transformer pré-treinado para tradução espanhol-inglês. " - "Baixa os pesos do Hugging Face no primeiro uso (requer internet)." + "Baixe seus pesos do Hugging Face antes de usar (requer internet)." ), de=( "Vortrainierter Transformer für Spanisch-Englisch-Übersetzung. " - "Lädt Gewichte von Hugging Face bei der ersten Verwendung herunter " + "Lädt die Gewichte vor der Nutzung von Hugging Face herunter " "(Internet erforderlich)." ), zh=( "用于西班牙语到英语翻译的预训练 Transformer。" - "首次使用时从 Hugging Face 下载权重（需要网络）。" + "使用前需从 Hugging Face 下载权重（需要网络）。" ), ) COLOR: str = "#FF8A65" diff --git a/DashAI/back/models/hugging_face/opus_mt_fr_en_transformer.py b/DashAI/back/models/hugging_face/opus_mt_fr_en_transformer.py index 903a3896c..ec793edbb 100644 --- a/DashAI/back/models/hugging_face/opus_mt_fr_en_transformer.py +++ b/DashAI/back/models/hugging_face/opus_mt_fr_en_transformer.py @@ -39,24 +39,24 @@ class OpusMtFrEnTransformer(OpusMtTransformerMixin): DESCRIPTION: str = MultilingualString( en=( "Pretrained transformer for French to English translation. " - "Downloads weights from Hugging Face on first use (internet required)." + "Download its weights from Hugging Face before use (internet required)." ), es=( "Transformer preentrenado para traducción francés-inglés. " - "Descarga pesos de Hugging Face en el primer uso (requiere internet)." + "Descarga sus pesos de Hugging Face antes de usarlo (requiere internet)." ), pt=( "Transformer pré-treinado para tradução francês-inglês. " - "Baixa os pesos do Hugging Face no primeiro uso (requer internet)." + "Baixe seus pesos do Hugging Face antes de usar (requer internet)." ), de=( "Vortrainierter Transformer für Französisch-Englisch-Übersetzung. " - "Lädt Gewichte von Hugging Face bei der ersten Verwendung herunter " + "Lädt die Gewichte vor der Nutzung von Hugging Face herunter " "(Internet erforderlich)." ), zh=( "用于法语到英语翻译的预训练 Transformer。" - "首次使用时从 Hugging Face 下载权重（需要网络）。" + "使用前需从 Hugging Face 下载权重（需要网络）。" ), ) COLOR: str = "#0097A7" diff --git a/DashAI/back/models/hugging_face/roberta_transformer.py b/DashAI/back/models/hugging_face/roberta_transformer.py index 8e8fd40f1..e3383a926 100644 --- a/DashAI/back/models/hugging_face/roberta_transformer.py +++ b/DashAI/back/models/hugging_face/roberta_transformer.py @@ -34,24 +34,24 @@ class RobertaTransformer(HuggingFaceTextClassificationTransformer): DESCRIPTION: str = MultilingualString( en=( "Robustly optimised BERT for English text classification. " - "Downloads weights from Hugging Face on first use (internet required)." + "Download its weights from Hugging Face before use (internet required)." ), es=( "BERT optimizado robustamente para clasificación de texto en inglés. " - "Descarga pesos de Hugging Face en el primer uso (requiere internet)." + "Descarga sus pesos de Hugging Face antes de usarlo (requiere internet)." ), pt=( "BERT otimizado robustamente para classificação de texto em inglês. " - "Baixa os pesos do Hugging Face no primeiro uso (requer internet)." + "Baixe seus pesos do Hugging Face antes de usar (requer internet)." ), de=( "Robust optimiertes BERT für englische Textklassifikation. " - "Lädt Gewichte von Hugging Face bei der ersten Verwendung herunter " + "Lädt die Gewichte vor der Nutzung von Hugging Face herunter " "(Internet erforderlich)." ), zh=( "经过鲁棒优化的 BERT，用于英文文本分类。" - "首次使用时从 Hugging Face 下载权重（需要网络）。" + "使用前需从 Hugging Face 下载权重（需要网络）。" ), ) COLOR: str = "#E65100" diff --git a/DashAI/back/models/hugging_face/t5_small_transformer.py b/DashAI/back/models/hugging_face/t5_small_transformer.py index 366de45b5..997199351 100644 --- a/DashAI/back/models/hugging_face/t5_small_transformer.py +++ b/DashAI/back/models/hugging_face/t5_small_transformer.py @@ -105,27 +105,27 @@ class T5SmallTransformer(HFPretrainedDownloadMixin, TranslationModel): en=( "Google T5-small model for English-to-{German, French, Romanian} " "translation using task prefixes. " - "Downloads weights from Hugging Face on first use (internet required)." + "Download its weights from Hugging Face before use (internet required)." ), es=( "Modelo T5-small de Google para traducción inglés-{alemán, francés, " "rumano} usando prefijos de tarea. " - "Descarga pesos de Hugging Face en el primer uso (requiere internet)." + "Descarga sus pesos de Hugging Face antes de usarlo (requiere internet)." ), pt=( "Modelo T5-small do Google para tradução inglês-{alemão, francês, " "romeno} usando prefixos de tarefa. " - "Baixa os pesos do Hugging Face no primeiro uso (requer internet)." + "Baixe seus pesos do Hugging Face antes de usar (requer internet)." ), de=( "Google T5-small-Modell für Englisch-zu-{Deutsch, Französisch, Rumänisch}-" "Übersetzung mit Aufgabenpräfixen. " - "Lädt Gewichte von Hugging Face bei der ersten Verwendung herunter " + "Lädt die Gewichte vor der Nutzung von Hugging Face herunter " "(Internet erforderlich)." ), zh=( "谷歌 T5-small 模型，通过任务前缀实现英语到德语/法语/罗马尼亚语翻译。" - "首次使用时从 Hugging Face 下载权重（需要网络）。" + "使用前需从 Hugging Face 下载权重（需要网络）。" ), ) COLOR: str = "#00695C" diff --git a/DashAI/back/models/hugging_face/xlm_roberta_transformer.py b/DashAI/back/models/hugging_face/xlm_roberta_transformer.py index 51722452e..caf06a6a4 100644 --- a/DashAI/back/models/hugging_face/xlm_roberta_transformer.py +++ b/DashAI/back/models/hugging_face/xlm_roberta_transformer.py @@ -35,27 +35,27 @@ class XlmRobertaTransformer(HuggingFaceTextClassificationTransformer): en=( "Multilingual RoBERTa for crosslingual text classification " "(100 languages). " - "Downloads weights from Hugging Face on first use (internet required)." + "Download its weights from Hugging Face before use (internet required)." ), es=( "RoBERTa multilingüe para clasificación de texto entre idiomas " "(100 idiomas). " - "Descarga pesos de Hugging Face en el primer uso (requiere internet)." + "Descarga sus pesos de Hugging Face antes de usarlo (requiere internet)." ), pt=( "RoBERTa multilingual para classificação de texto entre idiomas " "(100 idiomas). " - "Baixa os pesos do Hugging Face no primeiro uso (requer internet)." + "Baixe seus pesos do Hugging Face antes de usar (requer internet)." ), de=( "Mehrsprachiges RoBERTa für sprachübergreifende Textklassifikation " "(100 Sprachen). " - "Lädt Gewichte von Hugging Face bei der ersten Verwendung herunter " + "Lädt die Gewichte vor der Nutzung von Hugging Face herunter " "(Internet erforderlich)." ), zh=( "支持跨语言文本分类的多语言 RoBERTa（100 种语言）。" - "首次使用时从 Hugging Face 下载权重（需要网络）。" + "使用前需从 Hugging Face 下载权重（需要网络）。" ), ) COLOR: str = "#6A1B9A" diff --git a/DashAI/back/models/hugging_face/xlnet_transformer.py b/DashAI/back/models/hugging_face/xlnet_transformer.py index 0bc110548..9a7f59050 100644 --- a/DashAI/back/models/hugging_face/xlnet_transformer.py +++ b/DashAI/back/models/hugging_face/xlnet_transformer.py @@ -34,24 +34,24 @@ class XlnetTransformer(HuggingFaceTextClassificationTransformer): DESCRIPTION: str = MultilingualString( en=( "Autoregressive XLNet model for English text classification. " - "Downloads weights from Hugging Face on first use (internet required)." + "Download its weights from Hugging Face before use (internet required)." ), es=( "Modelo XLNet autorregresivo para clasificación de texto en inglés. " - "Descarga pesos de Hugging Face en el primer uso (requiere internet)." + "Descarga sus pesos de Hugging Face antes de usarlo (requiere internet)." ), pt=( "Modelo XLNet autorregressivo para classificação de texto em inglês. " - "Baixa os pesos do Hugging Face no primeiro uso (requer internet)." + "Baixe seus pesos do Hugging Face antes de usar (requer internet)." ), de=( "Autoregressives XLNet-Modell für englische Textklassifikation. " - "Lädt Gewichte von Hugging Face bei der ersten Verwendung herunter " + "Lädt die Gewichte vor der Nutzung von Hugging Face herunter " "(Internet erforderlich)." ), zh=( "自回归 XLNet 模型，用于英文文本分类。" - "首次使用时从 Hugging Face 下载权重（需要网络）。" + "使用前需从 Hugging Face 下载权重（需要网络）。" ), ) COLOR: str = "#37474F" From d6afec323f3da0e77095a168e670968c8ea5e7b6 Mon Sep 17 00:00:00 2001 From: Irozuku Date: Fri, 3 Jul 2026 14:47:04 -0400 Subject: [PATCH 65/89] fix: refresh models side bar when any download finishes The row disabled state derives from the fetched model list, so a download that completes after navigating back (started by a previous mount) left the row disabled. Subscribe to the shared download state and refetch on completion or delete. --- .../src/components/models/ModelsRightBar.jsx | 16 +++++++++++++++- 1 file changed, 15 insertions(+), 1 deletion(-) diff --git a/DashAI/front/src/components/models/ModelsRightBar.jsx b/DashAI/front/src/components/models/ModelsRightBar.jsx index bb00620c7..513c6daf2 100644 --- a/DashAI/front/src/components/models/ModelsRightBar.jsx +++ b/DashAI/front/src/components/models/ModelsRightBar.jsx @@ -7,7 +7,10 @@ import { useSnackbar } from "notistack"; import SideBar from "../threeSectionLayout/panelContainers/SideBar"; import { getComponents } from "../../api/component"; import ModelListItem from "./model/ModelListItem"; -import { startComponentDownload } from "./model/ComponentDownloadControl"; +import { + startComponentDownload, + subscribeAnyDownloadState, +} from "./model/ComponentDownloadControl"; import ModelDownloadStatusIcon from "./model/ModelDownloadStatusIcon"; import { useTranslation } from "react-i18next"; import { useTourContext } from "../tour/TourProvider"; @@ -66,6 +69,17 @@ export default function ModelsRightBar({ onToggle }) { } }, [session, fetchModels]); + // Refetch when any download finishes (or is deleted) so the row's disabled + // state, which is derived from the fetched model list, stays in sync. The + // download may have been started by a previous mount of this panel, so we + // cannot rely on that download's onStatusChange callback firing here. + useEffect(() => { + if (!session) return undefined; + return subscribeAnyDownloadState((_name, state) => { + if (state.downloaded !== undefined) fetchModels(); + }); + }, [session, fetchModels]); + // Filter models based on search useEffect(() => { if (searchQuery.trim() === "") { From 8813305cce6dcbb299dbaeb8aa214c7968e8d29c Mon Sep 17 00:00:00 2001 From: Irozuku Date: Fri, 3 Jul 2026 14:53:06 -0400 Subject: [PATCH 66/89] fix: keep model row disabled while its download is in progress The backend can report a download as present mid-download (its folder is already non-empty), which un-disabled the row while the icon still showed a spinner. Drive the row's disabled state from the shared live download state, treating an in-progress download as not ready so the row and icon always agree. --- .../src/components/models/ModelsRightBar.jsx | 104 ++++++++++++------ 1 file changed, 68 insertions(+), 36 deletions(-) diff --git a/DashAI/front/src/components/models/ModelsRightBar.jsx b/DashAI/front/src/components/models/ModelsRightBar.jsx index 513c6daf2..797da998b 100644 --- a/DashAI/front/src/components/models/ModelsRightBar.jsx +++ b/DashAI/front/src/components/models/ModelsRightBar.jsx @@ -10,8 +10,50 @@ import ModelListItem from "./model/ModelListItem"; import { startComponentDownload, subscribeAnyDownloadState, + useComponentDownloadState, } from "./model/ComponentDownloadControl"; import ModelDownloadStatusIcon from "./model/ModelDownloadStatusIcon"; + +/** + * A single model row whose disabled state and download icon both derive from + * the shared live download state, so they never disagree. While a download is + * in progress the row stays disabled even if the backend already reports the + * (partially written) files as present. + */ +function ModelRow({ model, onUse, onDownload, onChanged, dataTour }) { + const requiresDownload = Boolean(model.metadata?.requires_download); + const { downloaded, downloading } = useComponentDownloadState(model); + const ready = !requiresDownload || (downloaded && !downloading); + + const handleClick = () => { + if (downloading) return; + if (ready) onUse(model); + else onDownload(model); + }; + + return ( + + ) : null + } + /> + ); +} + +ModelRow.propTypes = { + model: PropTypes.object.isRequired, + onUse: PropTypes.func.isRequired, + onDownload: PropTypes.func.isRequired, + onChanged: PropTypes.func, + dataTour: PropTypes.string, +}; import { useTranslation } from "react-i18next"; import { useTourContext } from "../tour/TourProvider"; import { useModels } from "./ModelsContext"; @@ -98,24 +140,13 @@ export default function ModelsRightBar({ onToggle }) { const tourContext = useTourContext(); - const handleModelClick = (model) => { + const handleUseModel = (model) => { if (!session) { enqueueSnackbar(t("models:error.selectSessionFirst"), { variant: "warning", }); return; } - // A download-required model that has not been downloaded cannot be - // configured; it is blocked in the list with an inline download control. - if (model.metadata?.requires_download && !model.downloaded) { - startComponentDownload({ - component: model, - enqueueSnackbar, - t, - onStatusChange: () => fetchModels(), - }); - return; - } selectModel(model); if (tourContext?.run && tourContext?.stepIndex === 2) { const waitForElement = () => { @@ -130,6 +161,21 @@ export default function ModelsRightBar({ onToggle }) { } }; + const handleDownloadModel = (model) => { + if (!session) { + enqueueSnackbar(t("models:error.selectSessionFirst"), { + variant: "warning", + }); + return; + } + startComponentDownload({ + component: model, + enqueueSnackbar, + t, + onStatusChange: () => fetchModels(), + }); + }; + if (sessionRightContent) { return ( @@ -263,30 +309,16 @@ export default function ModelsRightBar({ onToggle }) { ) : ( - {filteredModels.map((model, index) => { - const requiresDownload = Boolean( - model.metadata?.requires_download, - ); - const needsDownload = requiresDownload && !model.downloaded; - return ( - handleModelClick(model)} - onDisabledClick={() => handleModelClick(model)} - data-tour={index === 0 ? "first-model" : undefined} - action={ - requiresDownload ? ( - fetchModels()} - /> - ) : null - } - /> - ); - })} + {filteredModels.map((model, index) => ( + + ))} )} From 640212b1878b0037cf95e7359225997774e80134 Mon Sep 17 00:00:00 2001 From: Irozuku Date: Fri, 3 Jul 2026 14:56:24 -0400 Subject: [PATCH 67/89] fix: avoid scroll reset when a model download finishes or is deleted Refetching the model list flipped loading, swapping the list for a spinner and resetting scroll. Update the affected model's downloaded flag in place from the shared download-state subscription instead, keeping the list mounted and the scroll position intact. --- .../src/components/models/ModelsRightBar.jsx | 39 ++++++++++--------- 1 file changed, 20 insertions(+), 19 deletions(-) diff --git a/DashAI/front/src/components/models/ModelsRightBar.jsx b/DashAI/front/src/components/models/ModelsRightBar.jsx index 797da998b..98a680fca 100644 --- a/DashAI/front/src/components/models/ModelsRightBar.jsx +++ b/DashAI/front/src/components/models/ModelsRightBar.jsx @@ -20,7 +20,7 @@ import ModelDownloadStatusIcon from "./model/ModelDownloadStatusIcon"; * in progress the row stays disabled even if the backend already reports the * (partially written) files as present. */ -function ModelRow({ model, onUse, onDownload, onChanged, dataTour }) { +function ModelRow({ model, onUse, onDownload, dataTour }) { const requiresDownload = Boolean(model.metadata?.requires_download); const { downloaded, downloading } = useComponentDownloadState(model); const ready = !requiresDownload || (downloaded && !downloading); @@ -39,9 +39,7 @@ function ModelRow({ model, onUse, onDownload, onChanged, dataTour }) { onDisabledClick={handleClick} data-tour={dataTour} action={ - requiresDownload ? ( - - ) : null + requiresDownload ? : null } /> ); @@ -51,7 +49,6 @@ ModelRow.propTypes = { model: PropTypes.object.isRequired, onUse: PropTypes.func.isRequired, onDownload: PropTypes.func.isRequired, - onChanged: PropTypes.func, dataTour: PropTypes.string, }; import { useTranslation } from "react-i18next"; @@ -111,16 +108,23 @@ export default function ModelsRightBar({ onToggle }) { } }, [session, fetchModels]); - // Refetch when any download finishes (or is deleted) so the row's disabled - // state, which is derived from the fetched model list, stays in sync. The - // download may have been started by a previous mount of this panel, so we - // cannot rely on that download's onStatusChange callback firing here. + // When any download finishes (or is deleted) update just that model's flag in + // place. A full refetch would flip `loading`, swap the list for a spinner and + // reset the scroll position; an in-place update keeps the list mounted and + // keeps `downloaded` accurate for the model passed on to the config dialog. useEffect(() => { if (!session) return undefined; - return subscribeAnyDownloadState((_name, state) => { - if (state.downloaded !== undefined) fetchModels(); + return subscribeAnyDownloadState((name, state) => { + if (state.downloaded === undefined) return; + setModels((prev) => + prev.map((model) => + model.name === name + ? { ...model, downloaded: state.downloaded } + : model, + ), + ); }); - }, [session, fetchModels]); + }, [session]); // Filter models based on search useEffect(() => { @@ -168,12 +172,10 @@ export default function ModelsRightBar({ onToggle }) { }); return; } - startComponentDownload({ - component: model, - enqueueSnackbar, - t, - onStatusChange: () => fetchModels(), - }); + // Completion is reflected by the shared download-state subscription above, + // which updates the model's flag in place without a scroll-resetting + // refetch. + startComponentDownload({ component: model, enqueueSnackbar, t }); }; if (sessionRightContent) { @@ -315,7 +317,6 @@ export default function ModelsRightBar({ onToggle }) { model={model} onUse={handleUseModel} onDownload={handleDownloadModel} - onChanged={fetchModels} dataTour={index === 0 ? "first-model" : undefined} /> ))} From 0ba2b4b80e4c624087d9034ce36537dd74c4617a Mon Sep 17 00:00:00 2001 From: Irozuku Date: Fri, 3 Jul 2026 15:02:34 -0400 Subject: [PATCH 68/89] fix: render the model delete control as an icon button The delete affordance in the models side bar is now a Material IconButton (hover and ripple feedback) instead of a bare clickable icon; the download hint stays a plain icon. --- .../models/model/ModelDownloadStatusIcon.jsx | 16 +++++----------- 1 file changed, 5 insertions(+), 11 deletions(-) diff --git a/DashAI/front/src/components/models/model/ModelDownloadStatusIcon.jsx b/DashAI/front/src/components/models/model/ModelDownloadStatusIcon.jsx index 9fd77a39b..3c0920385 100644 --- a/DashAI/front/src/components/models/model/ModelDownloadStatusIcon.jsx +++ b/DashAI/front/src/components/models/model/ModelDownloadStatusIcon.jsx @@ -1,6 +1,6 @@ import React from "react"; import PropTypes from "prop-types"; -import { Box, CircularProgress } from "@mui/material"; +import { Box, CircularProgress, IconButton } from "@mui/material"; import DownloadIcon from "@mui/icons-material/Download"; import DeleteIcon from "@mui/icons-material/Delete"; import { useSnackbar } from "notistack"; @@ -31,9 +31,9 @@ export default function ModelDownloadStatusIcon({ model, onChanged }) { if (downloaded) { return ( - { e.stopPropagation(); @@ -44,15 +44,9 @@ export default function ModelDownloadStatusIcon({ model, onChanged }) { onStatusChange: onChanged, }); }} - sx={{ - display: "flex", - alignItems: "center", - color: "error.main", - cursor: "pointer", - }} > - + ); } From 38552a58d4533ee873c16d214ba746c908206168 Mon Sep 17 00:00:00 2001 From: Irozuku Date: Fri, 3 Jul 2026 15:04:07 -0400 Subject: [PATCH 69/89] feat: add a tooltip to the model delete icon button --- .../models/model/ModelDownloadStatusIcon.jsx | 36 ++++++++++--------- 1 file changed, 19 insertions(+), 17 deletions(-) diff --git a/DashAI/front/src/components/models/model/ModelDownloadStatusIcon.jsx b/DashAI/front/src/components/models/model/ModelDownloadStatusIcon.jsx index 3c0920385..07eea8edb 100644 --- a/DashAI/front/src/components/models/model/ModelDownloadStatusIcon.jsx +++ b/DashAI/front/src/components/models/model/ModelDownloadStatusIcon.jsx @@ -1,6 +1,6 @@ import React from "react"; import PropTypes from "prop-types"; -import { Box, CircularProgress, IconButton } from "@mui/material"; +import { Box, CircularProgress, IconButton, Tooltip } from "@mui/material"; import DownloadIcon from "@mui/icons-material/Download"; import DeleteIcon from "@mui/icons-material/Delete"; import { useSnackbar } from "notistack"; @@ -31,22 +31,24 @@ export default function ModelDownloadStatusIcon({ model, onChanged }) { if (downloaded) { return ( - { - e.stopPropagation(); - deleteComponent({ - component: model, - enqueueSnackbar, - t, - onStatusChange: onChanged, - }); - }} - > - - + + { + e.stopPropagation(); + deleteComponent({ + component: model, + enqueueSnackbar, + t, + onStatusChange: onChanged, + }); + }} + > + + + ); } From 5932951f8707f010670d8cc10e7060c2052bf05b Mon Sep 17 00:00:00 2001 From: Irozuku Date: Fri, 3 Jul 2026 15:08:51 -0400 Subject: [PATCH 70/89] fix: disable train and retrain when the run model is not downloaded A run whose model still needs downloading cannot be trained, so the Train/Retrain button is disabled with a tooltip until the model is downloaded. The state comes from the shared live download state so an inline download re-enables it. --- .../front/src/components/models/RunCard.jsx | 63 ++++++++++++------- 1 file changed, 39 insertions(+), 24 deletions(-) diff --git a/DashAI/front/src/components/models/RunCard.jsx b/DashAI/front/src/components/models/RunCard.jsx index f4f2b18f3..22c7a1ea9 100644 --- a/DashAI/front/src/components/models/RunCard.jsx +++ b/DashAI/front/src/components/models/RunCard.jsx @@ -45,6 +45,7 @@ import FormSchemaContainer from "../shared/FormSchemaContainer"; import OptimizationTableSelectOptimizer from "./modelSession/OptimizationTableSelectOptimizer"; import ModelsTableSelectMetric from "./modelSession/ModelsTableSelectMetric"; import useSchema from "../../hooks/useSchema"; +import { useComponentDownloadState } from "./model/ComponentDownloadControl"; import { updateRunParameters, getRunOperationsCount } from "../../api/run"; import RetrainConfirmDialog from "./RetrainConfirmDialog"; import { renderParamValue } from "./ModelParamBlock"; @@ -252,6 +253,15 @@ function RunCard({ const model = models.find((m) => m.name === run.model_name); const modelDisplayName = model?.display_name || run.model_name; + // A download-required model must be downloaded before it can be trained. + // Track the live download state so the button reflects an inline download. + const { downloaded, downloading } = useComponentDownloadState( + model || { name: run.model_name }, + ); + const modelNotDownloaded = + Boolean(model?.metadata?.requires_download) && + !(downloaded && !downloading); + const getStatusColor = (status) => { switch (status) { case 0: @@ -373,32 +383,37 @@ function RunCard({ {canTrain && ( 0 || - operationsCount.predictions > 0) - ? t("models:message.retrainWillResetOperations", { - explainersCount: operationsCount.explainers, - predictionsCount: operationsCount.predictions, - }) - : "" + modelNotDownloaded + ? t("common:componentDownload.mustDownload") + : run.status === 3 && + operationsCount && + (operationsCount.explainers > 0 || + operationsCount.predictions > 0) + ? t("models:message.retrainWillResetOperations", { + explainersCount: operationsCount.explainers, + predictionsCount: operationsCount.predictions, + }) + : "" } > - + + + )} {isRunning && ( From b2e868667c1e409c613dc4be2c76c7153b906a77 Mon Sep 17 00:00:00 2001 From: Irozuku Date: Fri, 3 Jul 2026 15:12:52 -0400 Subject: [PATCH 71/89] fix: skip undownloaded models in Run All Run All trained every not-started run, including ones whose model still needs downloading. It now trains only runs whose model is ready (downloaded or download-free), reading the live download state, and warns when any were skipped. --- .../models/SessionVisualization.jsx | 36 ++++++++++++++++--- .../models/model/ComponentDownloadControl.jsx | 5 +++ .../src/utils/i18n/locales/de/models.json | 1 + .../src/utils/i18n/locales/en/models.json | 1 + .../src/utils/i18n/locales/es/models.json | 1 + .../src/utils/i18n/locales/pt/models.json | 1 + .../src/utils/i18n/locales/zh/models.json | 1 + 7 files changed, 42 insertions(+), 4 deletions(-) diff --git a/DashAI/front/src/components/models/SessionVisualization.jsx b/DashAI/front/src/components/models/SessionVisualization.jsx index 2fd1e0da3..6fbd987c3 100644 --- a/DashAI/front/src/components/models/SessionVisualization.jsx +++ b/DashAI/front/src/components/models/SessionVisualization.jsx @@ -23,9 +23,11 @@ import { import ModelComparisonTable from "./ModelComparisonTable"; import RunCard from "./RunCard"; import { getComponents } from "../../api/component"; +import { getComponentDownloadState } from "./model/ComponentDownloadControl"; import ResultsGraphs from "../../pages/results/components/ResultsGraphs"; import RetrainConfirmDialog from "./RetrainConfirmDialog"; import { useTranslation } from "react-i18next"; +import { useSnackbar } from "notistack"; import { useModels } from "./ModelsContext"; import { useTourContext } from "../tour/TourProvider"; @@ -42,6 +44,7 @@ export default function SessionVisualization() { const [explainerRefreshTrigger, setExplainerRefreshTrigger] = useState(0); const isResizing = React.useRef(false); const { t } = useTranslation(["models", "common"]); + const { enqueueSnackbar } = useSnackbar(); const sessionTourContext = useTourContext(); const { @@ -199,6 +202,34 @@ export default function SessionVisualization() { } }; + // True when a run's model is ready to train: it either needs no download or + // its download is present and not in progress (live state overrides the + // possibly stale fetched flag). + const isRunModelReady = React.useCallback( + (run) => { + const model = models.find((m) => m.name === run.model_name); + if (!model?.metadata?.requires_download) return true; + const cached = getComponentDownloadState(run.model_name); + const downloaded = cached?.downloaded ?? Boolean(model.downloaded); + const downloading = Boolean(cached?.downloading); + return downloaded && !downloading; + }, + [models], + ); + + // Train every not-started run whose model is downloaded, skipping (and warning + // about) any whose model still needs downloading. + const handleRunAll = () => { + const notStarted = runs.filter((r) => r.status === 0); + const ready = notStarted.filter(isRunModelReady); + ready.forEach((run) => onTrain(run)); + if (ready.length < notStarted.length) { + enqueueSnackbar(t("models:message.skippedUndownloadedRuns"), { + variant: "warning", + }); + } + }; + const handleMouseMove = React.useCallback((e) => { if (isResizing.current) { const details = document.querySelector("[data-accordion-details]"); @@ -444,10 +475,7 @@ export default function SessionVisualization() { variant="contained" size="small" startIcon={} - onClick={() => { - const notStartedRuns = runs.filter((r) => r.status === 0); - notStartedRuns.forEach((run) => onTrain(run)); - }} + onClick={handleRunAll} > {t("models:button.runAll")} diff --git a/DashAI/front/src/components/models/model/ComponentDownloadControl.jsx b/DashAI/front/src/components/models/model/ComponentDownloadControl.jsx index 7cf74afde..414e1422a 100644 --- a/DashAI/front/src/components/models/model/ComponentDownloadControl.jsx +++ b/DashAI/front/src/components/models/model/ComponentDownloadControl.jsx @@ -41,6 +41,11 @@ const subscribeDownloadState = (name, listener) => { }; }; +// Read the last known download state for a component name, or undefined if it +// has not been tracked this session. Lets non-hook call sites (e.g. a bulk +// "run all" handler) consult the live state without subscribing. +export const getComponentDownloadState = (name) => downloadStateCache.get(name); + // Subscribe to every download/delete regardless of component name. Lets a // container (e.g. a config dialog) re-check which nested components still need // downloading after an inline control finishes. diff --git a/DashAI/front/src/utils/i18n/locales/de/models.json b/DashAI/front/src/utils/i18n/locales/de/models.json index c79b662c4..0ad723411 100644 --- a/DashAI/front/src/utils/i18n/locales/de/models.json +++ b/DashAI/front/src/utils/i18n/locales/de/models.json @@ -177,6 +177,7 @@ "chartType": "Diagrammtyp" }, "message": { + "skippedUndownloadedRuns": "Durchläufe mit nicht heruntergeladenem Modell wurden übersprungen. Lade das Modell zuerst herunter.", "allRunsCompleted": "{{experiment}} hat alle Durchläufe abgeschlossen.", "confirmDeleteRun": "Sind Sie sicher, dass Sie diesen Durchlauf löschen möchten? Diese Aktion kann nicht rückgängig gemacht werden.", "editingParametersWarning": "Das Bearbeiten von Parametern setzt den Status dieses Durchlaufs auf 'Nicht gestartet' zurück. Alle vorhandenen Metriken und Ergebnisse gehen verloren.", diff --git a/DashAI/front/src/utils/i18n/locales/en/models.json b/DashAI/front/src/utils/i18n/locales/en/models.json index 3770033e7..9a9c1bfc6 100644 --- a/DashAI/front/src/utils/i18n/locales/en/models.json +++ b/DashAI/front/src/utils/i18n/locales/en/models.json @@ -177,6 +177,7 @@ "chartType": "Chart type" }, "message": { + "skippedUndownloadedRuns": "Skipped runs whose model is not downloaded. Download the model first.", "allRunsCompleted": "{{experiment}} has completed all its runs.", "confirmDeleteRun": "Are you sure you want to delete this run? This action cannot be undone.", "editingParametersWarning": "Editing parameters will reset this run's status to 'Not Started'. All existing metrics and results will be lost.", diff --git a/DashAI/front/src/utils/i18n/locales/es/models.json b/DashAI/front/src/utils/i18n/locales/es/models.json index 7ede341d8..03ae785df 100644 --- a/DashAI/front/src/utils/i18n/locales/es/models.json +++ b/DashAI/front/src/utils/i18n/locales/es/models.json @@ -179,6 +179,7 @@ "chartType": "Tipo de gráfico" }, "message": { + "skippedUndownloadedRuns": "Se omitieron las ejecuciones cuyo modelo no está descargado. Descarga el modelo primero.", "allRunsCompleted": "{{experiment}} ha completado todas sus ejecuciones.", "confirmDeleteRun": "¿Está seguro de que desea eliminar esta ejecución? Esta acción no se puede deshacer.", "editingParametersWarning": "Al editar los parámetros se restablecerá el estado de esta ejecución a 'No Iniciado'. Todas las métricas y resultados existentes se perderán.", diff --git a/DashAI/front/src/utils/i18n/locales/pt/models.json b/DashAI/front/src/utils/i18n/locales/pt/models.json index 50f551a17..0de61f6fa 100644 --- a/DashAI/front/src/utils/i18n/locales/pt/models.json +++ b/DashAI/front/src/utils/i18n/locales/pt/models.json @@ -179,6 +179,7 @@ "chartType": "Tipo de gráfico" }, "message": { + "skippedUndownloadedRuns": "Execuções cujo modelo não está baixado foram ignoradas. Baixe o modelo primeiro.", "allRunsCompleted": "{{experiment}} concluiu todas as suas execuções.", "confirmDeleteRun": "Tem certeza de que deseja excluir esta execução? Esta ação não pode ser desfeita.", "editingParametersWarning": "Ao editar os parâmetros, o estado desta execução será redefinido para 'Não Iniciado'. Todas as métricas e resultados existentes serão perdidos.", diff --git a/DashAI/front/src/utils/i18n/locales/zh/models.json b/DashAI/front/src/utils/i18n/locales/zh/models.json index 666fdfe2f..23fc67646 100644 --- a/DashAI/front/src/utils/i18n/locales/zh/models.json +++ b/DashAI/front/src/utils/i18n/locales/zh/models.json @@ -177,6 +177,7 @@ "chartType": "图表类型" }, "message": { + "skippedUndownloadedRuns": "已跳过模型未下载的运行。请先下载模型。", "allRunsCompleted": "{{experiment}} 已完成所有运行。", "confirmDeleteRun": "确定要删除此运行吗？此操作无法撤销。", "editingParametersWarning": "编辑参数将把此运行的状态重置为「未开始」。所有现有指标和结果将丢失。", From 89c16cb3ce3212efbe4f31392cc08e3c4335c4b1 Mon Sep 17 00:00:00 2001 From: Irozuku Date: Fri, 3 Jul 2026 15:24:54 -0400 Subject: [PATCH 72/89] fix: block generative session while its model is still downloading The chat input used a one-time reconciled download status that reports true mid-download (partial files), so a still-downloading model was usable. Both the input gate and the model switcher labels now read the shared live download state, blocking while downloading and updating the moment a download finishes. --- .../components/generative/GenerativeChat.jsx | 11 ++++++++-- .../components/generative/ModelSwitcher.jsx | 21 +++++++++++++++++-- 2 files changed, 28 insertions(+), 4 deletions(-) diff --git a/DashAI/front/src/components/generative/GenerativeChat.jsx b/DashAI/front/src/components/generative/GenerativeChat.jsx index 016a9a631..6fa754c9d 100644 --- a/DashAI/front/src/components/generative/GenerativeChat.jsx +++ b/DashAI/front/src/components/generative/GenerativeChat.jsx @@ -21,7 +21,9 @@ import { import { getRelatedComponents } from "../../api/generativeTask"; import InfoSessionModal from "./InfoSessionModal"; import ModelSwitcher from "./ModelSwitcher"; -import ComponentDownloadControl from "../models/model/ComponentDownloadControl"; +import ComponentDownloadControl, { + useComponentDownloadState, +} from "../models/model/ComponentDownloadControl"; import { useSnackbar } from "notistack"; import { MediaInput } from "./MediaInput"; import { Trans, useTranslation } from "react-i18next"; @@ -127,9 +129,14 @@ export default function GenerativeChat() { .catch(() => setModelsByName({})); }, [sessionInfo?.task_name]); + // Use the live download state so an in-progress download keeps the input + // blocked even when the backend already reports the (partial) files as + // present, and unblocks the moment the download actually finishes. + const { downloaded: liveDownloaded, downloading: liveDownloading } = + useComponentDownloadState(modelComponent || { name: modelName || "" }); const modelBlocked = Boolean(modelComponent?.metadata?.requires_download) && - !modelComponent?.downloaded; + !(liveDownloaded && !liveDownloading); const getMessages = () => { getProcessesBySessionId(sessionId).then((response) => { diff --git a/DashAI/front/src/components/generative/ModelSwitcher.jsx b/DashAI/front/src/components/generative/ModelSwitcher.jsx index 54308112c..7a517b2d0 100644 --- a/DashAI/front/src/components/generative/ModelSwitcher.jsx +++ b/DashAI/front/src/components/generative/ModelSwitcher.jsx @@ -5,6 +5,10 @@ import { useSnackbar } from "notistack"; import { useTranslation } from "react-i18next"; import { getRelatedComponents } from "../../api/generativeTask"; import { updateGenerativeSession } from "../../api/session"; +import { + getComponentDownloadState, + subscribeAnyDownloadState, +} from "../models/model/ComponentDownloadControl"; /** * Session-level model switcher: lets the user change the model used by a @@ -21,6 +25,9 @@ export default function ModelSwitcher({ const { enqueueSnackbar } = useSnackbar(); const [models, setModels] = useState([]); const [saving, setSaving] = useState(false); + // Bump to re-render when any download state changes so the labels reflect + // downloads that finished after the model list was fetched. + const [, setDownloadVersion] = useState(0); useEffect(() => { if (!taskName) return; @@ -29,6 +36,11 @@ export default function ModelSwitcher({ .catch(() => setModels([])); }, [taskName]); + useEffect( + () => subscribeAnyDownloadState(() => setDownloadVersion((v) => v + 1)), + [], + ); + const handleChange = async (event) => { const newModel = event.target.value; if (!newModel || newModel === currentModelName) return; @@ -75,10 +87,15 @@ export default function ModelSwitcher({ > {options.map((model) => { // A not-downloaded model is still selectable; the chat blocks input - // and offers the download once it becomes the session's model. + // and offers the download once it becomes the session's model. Read + // the live download state so a finished download drops the label and + // an in-progress one keeps it despite a premature backend flag. + const cached = getComponentDownloadState(model.name); + const downloaded = cached?.downloaded ?? model.downloaded; + const downloading = Boolean(cached?.downloading); const notDownloaded = Boolean(model.metadata?.requires_download) && - !model.downloaded && + !(downloaded && !downloading) && model.name !== currentModelName; return ( From 383ea0a2b426ab4a73486c674034dd939583caad Mon Sep 17 00:00:00 2001 From: Irozuku Date: Fri, 3 Jul 2026 15:39:57 -0400 Subject: [PATCH 73/89] fix: stop reporting freshly queued jobs as failed The job poller's final flush reported every non-finished watched job as an error when the queue briefly looked empty, so a download queued right after a delete showed a false failure. It now only errors a job that vanished or is actually in an error state, and keeps watching pending or running jobs. --- DashAI/front/src/utils/jobPoller.js | 18 ++++++++++++++---- 1 file changed, 14 insertions(+), 4 deletions(-) diff --git a/DashAI/front/src/utils/jobPoller.js b/DashAI/front/src/utils/jobPoller.js index 2a981e372..5f0fc7dba 100644 --- a/DashAI/front/src/utils/jobPoller.js +++ b/DashAI/front/src/utils/jobPoller.js @@ -125,13 +125,23 @@ async function pollJobs() { const job = allJobs.find((j) => j.id === jobId); if (job?.status === "finished") { if (watcher.onSuccess) watcher.onSuccess(job); - } else { + stopJobPolling(jobId); + } else if (job?.status === "error") { + if (watcher.onError) watcher.onError(job); + stopJobPolling(jobId); + } else if (!job) { + // The job vanished from the queue entirely; treat it as failed. if (watcher.onError) - watcher.onError(job || { id: jobId, status: "deleted" }); + watcher.onError({ id: jobId, status: "deleted" }); + stopJobPolling(jobId); } + // A job still pending or running is left watched and rechecked on the + // next poll, so a freshly queued job is never reported as an error + // just because the queue briefly looked empty. + } + if (state.jobWatchers.size === 0) { + stopJobPoller(); } - state.jobWatchers.clear(); - stopJobPoller(); return; } catch (e) { console.error("[JobPoller] Final flush failed, will retry:", e); From 90ffb3fd2ec63bfd94d7e56d220c8274428d10c0 Mon Sep 17 00:00:00 2001 From: Irozuku Date: Fri, 3 Jul 2026 15:50:39 -0400 Subject: [PATCH 74/89] chore: remove dummy downloadable classifier used for nested download testing --- DashAI/back/initial_components.py | 4 - .../dummy_downloadable_classifier.py | 144 ------------------ 2 files changed, 148 deletions(-) delete mode 100644 DashAI/back/models/scikit_learn/dummy_downloadable_classifier.py diff --git a/DashAI/back/initial_components.py b/DashAI/back/initial_components.py index eaa6d6ee4..55be38ecb 100644 --- a/DashAI/back/initial_components.py +++ b/DashAI/back/initial_components.py @@ -257,9 +257,6 @@ DecisionTreeRegression, ) from DashAI.back.models.scikit_learn.dummy_classifier import DummyClassifier -from DashAI.back.models.scikit_learn.dummy_downloadable_classifier import ( - DummyDownloadableClassifier, -) from DashAI.back.models.scikit_learn.elastic_net_regression import ElasticNetRegression from DashAI.back.models.scikit_learn.extra_trees_classifier import ExtraTreesClassifier from DashAI.back.models.scikit_learn.extra_trees_regression import ExtraTreesRegression @@ -427,7 +424,6 @@ def get_initial_components(): RealVisXLV4, StableDiffusionXLV1ControlNet, SVC, - DummyDownloadableClassifier, SVR, T5SmallTransformer, TfIdfLogRegTextClassificationModel, diff --git a/DashAI/back/models/scikit_learn/dummy_downloadable_classifier.py b/DashAI/back/models/scikit_learn/dummy_downloadable_classifier.py deleted file mode 100644 index c69909ec3..000000000 --- a/DashAI/back/models/scikit_learn/dummy_downloadable_classifier.py +++ /dev/null @@ -1,144 +0,0 @@ -import time -from typing import Optional - -from DashAI.back.core.schema_fields import ( - BaseSchema, - component_field, - schema_field, -) -from DashAI.back.core.utils import MultilingualString -from DashAI.back.dependencies.downloads.downloadable import ( - DownloadableMixin, - ProgressReporter, -) -from DashAI.back.models.scikit_learn.svc import SVC - - -class DummyDownloadableClassifierSchema(BaseSchema): - """Schema for the dummy classifier, exposing one nested classifier field. - - The ``nested_classifier`` parameter is a component field so a second - (possibly download-required) tabular classifier can be selected inside - this one, letting the nested-download flow be exercised at depth. - """ - - nested_classifier: schema_field( - component_field(parent="TabularClassificationModel"), - placeholder={"component": "SVC", "params": {}}, - description=MultilingualString( - en="A nested tabular classifier, used to test nested downloads.", - es="Un clasificador tabular anidado, para probar descargas anidadas.", - pt="Um classificador tabular aninhado, para testar downloads aninhados.", - de=( - "Ein verschachtelter tabellarischer Klassifikator zum Testen " - "verschachtelter Downloads." - ), - zh="嵌套的表格分类器，用于测试嵌套下载。", - ), - alias=MultilingualString( - en="Nested classifier", - es="Clasificador anidado", - pt="Classificador aninhado", - de="Verschachtelter Klassifikator", - zh="嵌套分类器", - ), - ) # type: ignore - - -class DummyDownloadableClassifier(DownloadableMixin, SVC): - """A fake download-required tabular classifier for UI testing. - - Behaves exactly like :class:`SVC` at train time but is flagged as - requiring a download so the inline download control appears when it is - selected as another component's parameter (e.g. the Bag-of-Words tabular - classifier). Its ``download`` writes a marker file instead of fetching any - real artifact, so the download/delete flow can be exercised end to end - without network access. - """ - - SCHEMA = DummyDownloadableClassifierSchema - DOWNLOAD_SIZE_BYTES = 256 * 1024 * 1024 - COLOR = "#B39DDB" - ICON = "Timeline" - DISPLAY_NAME = MultilingualString( - en="Dummy Downloadable Classifier", - es="Clasificador Descargable de Prueba", - pt="Classificador Baixavel de Teste", - de="Dummy Herunterladbarer Klassifikator", - zh="虚拟可下载分类器", - ) - DESCRIPTION = MultilingualString( - en=( - "A test-only classifier that requires a download. It trains like an " - "SVM but is used to preview the inline download control when picking " - "a nested component." - ), - es=( - "Un clasificador solo de prueba que requiere descarga. Entrena como " - "una SVM, pero sirve para previsualizar el control de descarga en " - "linea al elegir un componente anidado." - ), - pt=( - "Um classificador apenas de teste que requer download. Treina como " - "uma SVM, mas serve para pre-visualizar o controle de download em " - "linha ao escolher um componente aninhado." - ), - de=( - "Ein reiner Testklassifikator, der einen Download erfordert. Er " - "trainiert wie eine SVM, dient aber zur Vorschau des Inline-" - "Download-Steuerelements bei der Auswahl einer verschachtelten " - "Komponente." - ), - zh=( - "仅用于测试的分类器，需要下载。它像 SVM 一样训练，" - "用于在选择嵌套组件时预览内联下载控件。" - ), - ) - - def __init__(self, **kwargs): - """Store the nested classifier and forward the rest to ``SVC``. - - Parameters - ---------- - **kwargs : dict - May include ``nested_classifier`` (an instantiated tabular - classifier), which is kept as an attribute and not passed to the - underlying sklearn estimator. - """ - self.nested_classifier = kwargs.pop("nested_classifier", None) - super().__init__(**kwargs) - - @classmethod - def is_downloaded(cls) -> bool: - """Return whether the marker file is present. - - Returns - ------- - bool - ``True`` when ``component_dir()`` exists and is non-empty. - """ - directory = cls.component_dir() - return directory.is_dir() and any(directory.iterdir()) - - @classmethod - def download(cls, report: Optional[ProgressReporter] = None) -> None: - """Write a marker file to simulate a download. - - A short delay is inserted so the downloading state is visible in the - UI. No real artifact is fetched. - - Parameters - ---------- - report : ProgressReporter, optional - Callback invoked with progress fractions and phase messages. - """ - directory = cls.component_dir() - directory.mkdir(parents=True, exist_ok=True) - steps = 4 - for step in range(steps): - if report is not None: - report(step / steps, "Downloading dummy weights") - time.sleep(1) - (directory / "weights.marker").write_text("dummy", encoding="utf-8") - if report is not None: - report(1.0, "Done") From e71d100295c5d65013279d4a0d25d5ae22ed78d0 Mon Sep 17 00:00:00 2001 From: Irozuku Date: Fri, 3 Jul 2026 15:56:40 -0400 Subject: [PATCH 75/89] test: drop dummy model name from download icon test fixture --- .../components/models/model/ModelDownloadStatusIcon.test.jsx | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/DashAI/front/src/components/models/model/ModelDownloadStatusIcon.test.jsx b/DashAI/front/src/components/models/model/ModelDownloadStatusIcon.test.jsx index 05e7827e7..c80b99a26 100644 --- a/DashAI/front/src/components/models/model/ModelDownloadStatusIcon.test.jsx +++ b/DashAI/front/src/components/models/model/ModelDownloadStatusIcon.test.jsx @@ -19,7 +19,7 @@ import ModelDownloadStatusIcon from "./ModelDownloadStatusIcon"; import { deleteComponentDownload } from "../../../api/component"; const model = { - name: "DummyDownloadableClassifier", + name: "DownloadableTestModel", downloaded: false, metadata: { requires_download: true, download_size_bytes: 268435456 }, }; @@ -43,7 +43,7 @@ describe("ModelDownloadStatusIcon", () => { fireEvent.click(del); await waitFor(() => expect(deleteComponentDownload).toHaveBeenCalledWith( - "DummyDownloadableClassifier", + "DownloadableTestModel", ), ); }); From 962ef3324fe1bdb874cebda4307d7f4c3db46ca6 Mon Sep 17 00:00:00 2001 From: Irozuku Date: Thu, 9 Jul 2026 09:38:07 -0400 Subject: [PATCH 76/89] fix: set accurate model download sizes to match the actual download DOWNLOAD_SIZE_BYTES were rough estimates far below what is actually fetched: snapshot_download pulls the whole HF repo (every weight format and, for diffusion, all fp16/fp32 variants). Update each model's declared size to the exact repo byte total, verified against an on-disk download. --- DashAI/back/models/hugging_face/albert_transformer.py | 2 +- DashAI/back/models/hugging_face/bert_transformer.py | 2 +- DashAI/back/models/hugging_face/bertin_transformer.py | 2 +- DashAI/back/models/hugging_face/beto_transformer.py | 2 +- DashAI/back/models/hugging_face/deberta_v3_transformer.py | 2 +- DashAI/back/models/hugging_face/distilbert_transformer.py | 2 +- DashAI/back/models/hugging_face/electra_transformer.py | 2 +- DashAI/back/models/hugging_face/llama_model.py | 6 +++--- DashAI/back/models/hugging_face/m2m100_transformer.py | 2 +- DashAI/back/models/hugging_face/minilm_transformer.py | 2 +- DashAI/back/models/hugging_face/mistral_model.py | 4 ++-- DashAI/back/models/hugging_face/mixtral_model.py | 4 ++-- DashAI/back/models/hugging_face/modernbert_transformer.py | 2 +- .../models/hugging_face/multilingual_bert_transformer.py | 2 +- DashAI/back/models/hugging_face/nllb_transformer.py | 2 +- .../back/models/hugging_face/opus_mt_en_de_transformer.py | 1 + .../back/models/hugging_face/opus_mt_en_es_transformer.py | 1 + .../back/models/hugging_face/opus_mt_en_fr_transformer.py | 1 + .../back/models/hugging_face/opus_mt_es_en_transformer.py | 1 + .../back/models/hugging_face/opus_mt_fr_en_transformer.py | 1 + DashAI/back/models/hugging_face/pixart_sigma_model.py | 4 ++-- DashAI/back/models/hugging_face/qwen_model.py | 4 ++-- DashAI/back/models/hugging_face/roberta_transformer.py | 2 +- .../models/hugging_face/sd15_depth_controlnet_model.py | 2 +- .../back/models/hugging_face/sd15_hed_controlnet_model.py | 2 +- .../models/hugging_face/sd15_openpose_controlnet_model.py | 2 +- .../models/hugging_face/sdxl_canny_controlnet_model.py | 2 +- DashAI/back/models/hugging_face/sdxl_turbo_model.py | 2 +- DashAI/back/models/hugging_face/smol_lm_model.py | 4 ++-- .../hugging_face/stable_diffusion_v1_depth_controlnet.py | 2 +- .../back/models/hugging_face/stable_diffusion_v2_model.py | 8 ++++---- .../back/models/hugging_face/stable_diffusion_v3_model.py | 8 ++++---- .../back/models/hugging_face/stable_diffusion_xl_model.py | 4 ++-- DashAI/back/models/hugging_face/t5_small_transformer.py | 2 +- DashAI/back/models/hugging_face/tongyi_z_image_model.py | 4 ++-- .../back/models/hugging_face/xlm_roberta_transformer.py | 2 +- DashAI/back/models/hugging_face/xlnet_transformer.py | 2 +- 37 files changed, 52 insertions(+), 47 deletions(-) diff --git a/DashAI/back/models/hugging_face/albert_transformer.py b/DashAI/back/models/hugging_face/albert_transformer.py index 8fd152db0..1ace38ad1 100644 --- a/DashAI/back/models/hugging_face/albert_transformer.py +++ b/DashAI/back/models/hugging_face/albert_transformer.py @@ -59,5 +59,5 @@ class AlbertTransformer(HuggingFaceTextClassificationTransformer): ICON: str = "Speed" SCHEMA = DistilBertTransformerSchema MODEL_NAME: str = "albert-base-v2" - DOWNLOAD_SIZE_BYTES: int = 47_000_000 + DOWNLOAD_SIZE_BYTES: int = 330556087 TEMP_CHECKPOINT_DIR: str = "DashAI/back/user_models/temp_checkpoints_albert" diff --git a/DashAI/back/models/hugging_face/bert_transformer.py b/DashAI/back/models/hugging_face/bert_transformer.py index 8b2121fea..594fbc0e3 100644 --- a/DashAI/back/models/hugging_face/bert_transformer.py +++ b/DashAI/back/models/hugging_face/bert_transformer.py @@ -58,5 +58,5 @@ class BertTransformer(HuggingFaceTextClassificationTransformer): ICON: str = "Psychology" SCHEMA = DistilBertTransformerSchema MODEL_NAME: str = "bert-base-uncased" - DOWNLOAD_SIZE_BYTES: int = 440_000_000 + DOWNLOAD_SIZE_BYTES: int = 3454102158 TEMP_CHECKPOINT_DIR: str = "DashAI/back/user_models/temp_checkpoints_bert" diff --git a/DashAI/back/models/hugging_face/bertin_transformer.py b/DashAI/back/models/hugging_face/bertin_transformer.py index 2a4cfb608..984b1e5e2 100644 --- a/DashAI/back/models/hugging_face/bertin_transformer.py +++ b/DashAI/back/models/hugging_face/bertin_transformer.py @@ -58,5 +58,5 @@ class BertinTransformer(HuggingFaceTextClassificationTransformer): ICON: str = "RecordVoiceOver" SCHEMA = DistilBertTransformerSchema MODEL_NAME: str = "bertin-project/bertin-roberta-base-spanish" - DOWNLOAD_SIZE_BYTES: int = 500_000_000 + DOWNLOAD_SIZE_BYTES: int = 1510386247 TEMP_CHECKPOINT_DIR: str = "DashAI/back/user_models/temp_checkpoints_bertin" diff --git a/DashAI/back/models/hugging_face/beto_transformer.py b/DashAI/back/models/hugging_face/beto_transformer.py index 4c6212137..78a549c05 100644 --- a/DashAI/back/models/hugging_face/beto_transformer.py +++ b/DashAI/back/models/hugging_face/beto_transformer.py @@ -58,5 +58,5 @@ class BetoTransformer(HuggingFaceTextClassificationTransformer): ICON: str = "RecordVoiceOver" SCHEMA = DistilBertTransformerSchema MODEL_NAME: str = "dccuchile/bert-base-spanish-wwm-cased" - DOWNLOAD_SIZE_BYTES: int = 440_000_000 + DOWNLOAD_SIZE_BYTES: int = 1416527075 TEMP_CHECKPOINT_DIR: str = "DashAI/back/user_models/temp_checkpoints_beto" diff --git a/DashAI/back/models/hugging_face/deberta_v3_transformer.py b/DashAI/back/models/hugging_face/deberta_v3_transformer.py index 271fa0f76..71684667d 100644 --- a/DashAI/back/models/hugging_face/deberta_v3_transformer.py +++ b/DashAI/back/models/hugging_face/deberta_v3_transformer.py @@ -61,5 +61,5 @@ class DebertaV3Transformer(HuggingFaceTextClassificationTransformer): ICON: str = "Psychology" SCHEMA = DebertaV3TransformerSchema MODEL_NAME: str = "microsoft/deberta-v3-base" - DOWNLOAD_SIZE_BYTES: int = 440_000_000 + DOWNLOAD_SIZE_BYTES: int = 1851424112 TEMP_CHECKPOINT_DIR: str = "DashAI/back/user_models/temp_checkpoints_deberta_v3" diff --git a/DashAI/back/models/hugging_face/distilbert_transformer.py b/DashAI/back/models/hugging_face/distilbert_transformer.py index 334ce67ec..84695f532 100644 --- a/DashAI/back/models/hugging_face/distilbert_transformer.py +++ b/DashAI/back/models/hugging_face/distilbert_transformer.py @@ -305,5 +305,5 @@ class DistilBertTransformer(HuggingFaceTextClassificationTransformer): ICON: str = "Psychology" SCHEMA = DistilBertTransformerSchema MODEL_NAME: str = "distilbert-base-uncased" - DOWNLOAD_SIZE_BYTES: int = 270_000_000 + DOWNLOAD_SIZE_BYTES: int = 1529742866 TEMP_CHECKPOINT_DIR: str = "DashAI/back/user_models/temp_checkpoints_distilbert" diff --git a/DashAI/back/models/hugging_face/electra_transformer.py b/DashAI/back/models/hugging_face/electra_transformer.py index 83867a4d7..088bd2aa7 100644 --- a/DashAI/back/models/hugging_face/electra_transformer.py +++ b/DashAI/back/models/hugging_face/electra_transformer.py @@ -58,5 +58,5 @@ class ElectraTransformer(HuggingFaceTextClassificationTransformer): ICON: str = "ElectricBolt" SCHEMA = DistilBertTransformerSchema MODEL_NAME: str = "google/electra-small-discriminator" - DOWNLOAD_SIZE_BYTES: int = 54_000_000 + DOWNLOAD_SIZE_BYTES: int = 163615837 TEMP_CHECKPOINT_DIR: str = "DashAI/back/user_models/temp_checkpoints_electra" diff --git a/DashAI/back/models/hugging_face/llama_model.py b/DashAI/back/models/hugging_face/llama_model.py index 170817f39..fd59f4c7c 100644 --- a/DashAI/back/models/hugging_face/llama_model.py +++ b/DashAI/back/models/hugging_face/llama_model.py @@ -21,7 +21,7 @@ class Llama31_8BInstruct(GGUFTextGenerationModel): # noqa: N801 REPO_ID = "bartowski/Meta-Llama-3.1-8B-Instruct-GGUF" GGUF_PATTERN = "*Q4_K_M.gguf" - DOWNLOAD_SIZE_BYTES = 4_900_000_000 + DOWNLOAD_SIZE_BYTES = 4920739232 SCHEMA = GGUFTextGenerationSchema COLOR: str = "#1a237e" DISPLAY_NAME = MultilingualString( @@ -84,7 +84,7 @@ class Llama32_1BInstruct(GGUFTextGenerationModel): # noqa: N801 REPO_ID = "bartowski/Llama-3.2-1B-Instruct-GGUF" GGUF_PATTERN = "*Q4_K_M.gguf" - DOWNLOAD_SIZE_BYTES = 800_000_000 + DOWNLOAD_SIZE_BYTES = 807694464 SCHEMA = GGUFTextGenerationSchema COLOR: str = "#1a237e" DISPLAY_NAME = MultilingualString( @@ -142,7 +142,7 @@ class Llama32_3BInstruct(GGUFTextGenerationModel): # noqa: N801 REPO_ID = "bartowski/Llama-3.2-3B-Instruct-GGUF" GGUF_PATTERN = "*Q4_K_M.gguf" - DOWNLOAD_SIZE_BYTES = 2_000_000_000 + DOWNLOAD_SIZE_BYTES = 2019377696 SCHEMA = GGUFTextGenerationSchema COLOR: str = "#1a237e" DISPLAY_NAME = MultilingualString( diff --git a/DashAI/back/models/hugging_face/m2m100_transformer.py b/DashAI/back/models/hugging_face/m2m100_transformer.py index d5eedddd5..601e5d93b 100644 --- a/DashAI/back/models/hugging_face/m2m100_transformer.py +++ b/DashAI/back/models/hugging_face/m2m100_transformer.py @@ -158,7 +158,7 @@ class M2M100Transformer(HFPretrainedDownloadMixin, TranslationModel): COLOR: str = "#6A1B9A" ICON: str = "Language" MODEL_NAME: str = "facebook/m2m100_418M" - DOWNLOAD_SIZE_BYTES: int = 1_900_000_000 + DOWNLOAD_SIZE_BYTES: int = 3877717593 def __init__(self, model=None, pretrained_dir=None, **kwargs): kwargs = self.validate_and_transform(kwargs) diff --git a/DashAI/back/models/hugging_face/minilm_transformer.py b/DashAI/back/models/hugging_face/minilm_transformer.py index 818e1a6a0..8888c8bc9 100644 --- a/DashAI/back/models/hugging_face/minilm_transformer.py +++ b/DashAI/back/models/hugging_face/minilm_transformer.py @@ -58,5 +58,5 @@ class MiniLMTransformer(HuggingFaceTextClassificationTransformer): ICON: str = "Speed" SCHEMA = DistilBertTransformerSchema MODEL_NAME: str = "microsoft/MiniLM-L12-H384-uncased" - DOWNLOAD_SIZE_BYTES: int = 130_000_000 + DOWNLOAD_SIZE_BYTES: int = 400889386 TEMP_CHECKPOINT_DIR: str = "DashAI/back/user_models/temp_checkpoints_minilm" diff --git a/DashAI/back/models/hugging_face/mistral_model.py b/DashAI/back/models/hugging_face/mistral_model.py index 5b07d1996..52c38adea 100644 --- a/DashAI/back/models/hugging_face/mistral_model.py +++ b/DashAI/back/models/hugging_face/mistral_model.py @@ -21,7 +21,7 @@ class Mistral7BInstructV03(GGUFTextGenerationModel): REPO_ID = "bartowski/Mistral-7B-Instruct-v0.3-GGUF" GGUF_PATTERN = "*Q4_K_M.gguf" - DOWNLOAD_SIZE_BYTES = 4_400_000_000 + DOWNLOAD_SIZE_BYTES = 4372812000 SCHEMA = GGUFTextGenerationSchema COLOR: str = "#ff6f00" DISPLAY_NAME = MultilingualString( @@ -79,7 +79,7 @@ class MistralNemoInstruct2407(GGUFTextGenerationModel): REPO_ID = "bartowski/Mistral-Nemo-Instruct-2407-GGUF" GGUF_PATTERN = "*Q4_K_M.gguf" - DOWNLOAD_SIZE_BYTES = 7_100_000_000 + DOWNLOAD_SIZE_BYTES = 7477208192 SCHEMA = GGUFTextGenerationSchema COLOR: str = "#ff6f00" DISPLAY_NAME = MultilingualString( diff --git a/DashAI/back/models/hugging_face/mixtral_model.py b/DashAI/back/models/hugging_face/mixtral_model.py index 978818d87..78f70b4bc 100644 --- a/DashAI/back/models/hugging_face/mixtral_model.py +++ b/DashAI/back/models/hugging_face/mixtral_model.py @@ -30,7 +30,7 @@ class Mixtral8x7BInstructQ4KM(GGUFTextGenerationModel): REPO_ID = _MIXTRAL_REPO GGUF_PATTERN = "*Q4_K_M.gguf" - DOWNLOAD_SIZE_BYTES = 26_000_000_000 + DOWNLOAD_SIZE_BYTES = 28448468384 SCHEMA = GGUFTextGenerationSchema COLOR: str = "#4a148c" DISPLAY_NAME = MultilingualString( @@ -91,7 +91,7 @@ class Mixtral8x7BInstructQ2K(GGUFTextGenerationModel): REPO_ID = _MIXTRAL_REPO GGUF_PATTERN = "*Q2_K.gguf" - DOWNLOAD_SIZE_BYTES = 16_000_000_000 + DOWNLOAD_SIZE_BYTES = 17311231392 SCHEMA = GGUFTextGenerationSchema COLOR: str = "#4a148c" DISPLAY_NAME = MultilingualString( diff --git a/DashAI/back/models/hugging_face/modernbert_transformer.py b/DashAI/back/models/hugging_face/modernbert_transformer.py index ad7e26033..a087fc310 100644 --- a/DashAI/back/models/hugging_face/modernbert_transformer.py +++ b/DashAI/back/models/hugging_face/modernbert_transformer.py @@ -55,6 +55,6 @@ class ModernBertTransformer(HuggingFaceTextClassificationTransformer): ICON: str = "Psychology" SCHEMA = ModernBertTransformerSchema MODEL_NAME: str = "answerdotai/ModernBERT-base" - DOWNLOAD_SIZE_BYTES: int = 600_000_000 + DOWNLOAD_SIZE_BYTES: int = 3134313772 MAX_TOKEN_LENGTH: int = 8192 TEMP_CHECKPOINT_DIR: str = "DashAI/back/user_models/temp_checkpoints_modernbert" diff --git a/DashAI/back/models/hugging_face/multilingual_bert_transformer.py b/DashAI/back/models/hugging_face/multilingual_bert_transformer.py index 916f6a08f..5e3574f8f 100644 --- a/DashAI/back/models/hugging_face/multilingual_bert_transformer.py +++ b/DashAI/back/models/hugging_face/multilingual_bert_transformer.py @@ -59,7 +59,7 @@ class MultilingualBertTransformer(HuggingFaceTextClassificationTransformer): ICON: str = "Translate" SCHEMA = DistilBertTransformerSchema MODEL_NAME: str = "bert-base-multilingual-cased" - DOWNLOAD_SIZE_BYTES: int = 680_000_000 + DOWNLOAD_SIZE_BYTES: int = 3226865011 TEMP_CHECKPOINT_DIR: str = ( "DashAI/back/user_models/temp_checkpoints_multilingual_bert" ) diff --git a/DashAI/back/models/hugging_face/nllb_transformer.py b/DashAI/back/models/hugging_face/nllb_transformer.py index 6628ba9d6..9fbe23e81 100644 --- a/DashAI/back/models/hugging_face/nllb_transformer.py +++ b/DashAI/back/models/hugging_face/nllb_transformer.py @@ -205,7 +205,7 @@ def _resolve_language_token_id(self, language_code: str, field_name: str) -> int raise ValueError(f"Unsupported {field_name} '{language_code}'.") MODEL_NAME: str = "facebook/nllb-200-distilled-600M" - DOWNLOAD_SIZE_BYTES: int = 2_400_000_000 + DOWNLOAD_SIZE_BYTES: int = 2482655255 def __init__(self, model=None, pretrained_dir=None, **kwargs): """Initialize the NLLB tokenizer and model. diff --git a/DashAI/back/models/hugging_face/opus_mt_en_de_transformer.py b/DashAI/back/models/hugging_face/opus_mt_en_de_transformer.py index 5950dc5ed..0c5854f6b 100644 --- a/DashAI/back/models/hugging_face/opus_mt_en_de_transformer.py +++ b/DashAI/back/models/hugging_face/opus_mt_en_de_transformer.py @@ -29,6 +29,7 @@ class OpusMtEnDeTransformer(OpusMtTransformerMixin): MODEL_NAME: str = "Helsinki-NLP/opus-mt-en-de" TEMP_CHECKPOINT_DIR: str = "DashAI/back/user_models/temp_checkpoints_opus-mt-en-de" SCHEMA = OpusMtEnDeTransformerSchema + DOWNLOAD_SIZE_BYTES = 1430871510 DISPLAY_NAME: str = MultilingualString( en="Opus MT En-De Transformer", es="Transformer Opus MT En-De", diff --git a/DashAI/back/models/hugging_face/opus_mt_en_es_transformer.py b/DashAI/back/models/hugging_face/opus_mt_en_es_transformer.py index 203c11fbb..2c9cf281a 100644 --- a/DashAI/back/models/hugging_face/opus_mt_en_es_transformer.py +++ b/DashAI/back/models/hugging_face/opus_mt_en_es_transformer.py @@ -270,6 +270,7 @@ class OpusMtEnESTransformer(OpusMtTransformerMixin): MODEL_NAME: str = "Helsinki-NLP/opus-mt-en-es" TEMP_CHECKPOINT_DIR: str = "DashAI/back/user_models/temp_checkpoints_opus-mt-en-es" SCHEMA = OpusMtEnESTransformerSchema + DOWNLOAD_SIZE_BYTES = 937836389 DISPLAY_NAME: str = MultilingualString( en="Opus MT En-Es Transformer", es="Transformer Opus MT En-Es", diff --git a/DashAI/back/models/hugging_face/opus_mt_en_fr_transformer.py b/DashAI/back/models/hugging_face/opus_mt_en_fr_transformer.py index ad026f97f..9c409d3ad 100644 --- a/DashAI/back/models/hugging_face/opus_mt_en_fr_transformer.py +++ b/DashAI/back/models/hugging_face/opus_mt_en_fr_transformer.py @@ -29,6 +29,7 @@ class OpusMtEnFrTransformer(OpusMtTransformerMixin): MODEL_NAME: str = "Helsinki-NLP/opus-mt-en-fr" TEMP_CHECKPOINT_DIR: str = "DashAI/back/user_models/temp_checkpoints_opus-mt-en-fr" SCHEMA = OpusMtEnFrTransformerSchema + DOWNLOAD_SIZE_BYTES = 903735972 DISPLAY_NAME: str = MultilingualString( en="Opus MT En-Fr Transformer", es="Transformer Opus MT En-Fr", diff --git a/DashAI/back/models/hugging_face/opus_mt_es_en_transformer.py b/DashAI/back/models/hugging_face/opus_mt_es_en_transformer.py index 4fadc8e1e..96b71168c 100644 --- a/DashAI/back/models/hugging_face/opus_mt_es_en_transformer.py +++ b/DashAI/back/models/hugging_face/opus_mt_es_en_transformer.py @@ -32,6 +32,7 @@ class OpusMtEsENTransformer(OpusMtTransformerMixin): MODEL_NAME: str = "Helsinki-NLP/opus-mt-es-en" TEMP_CHECKPOINT_DIR: str = "DashAI/back/user_models/temp_checkpoints_opus-mt-es-en" SCHEMA = OpusMtEsENTransformerSchema + DOWNLOAD_SIZE_BYTES = 627891360 DISPLAY_NAME: str = MultilingualString( en="Opus MT Es-En Transformer", es="Transformer Opus MT Es-En", diff --git a/DashAI/back/models/hugging_face/opus_mt_fr_en_transformer.py b/DashAI/back/models/hugging_face/opus_mt_fr_en_transformer.py index ec793edbb..fcb7e397c 100644 --- a/DashAI/back/models/hugging_face/opus_mt_fr_en_transformer.py +++ b/DashAI/back/models/hugging_face/opus_mt_fr_en_transformer.py @@ -29,6 +29,7 @@ class OpusMtFrEnTransformer(OpusMtTransformerMixin): MODEL_NAME: str = "Helsinki-NLP/opus-mt-fr-en" TEMP_CHECKPOINT_DIR: str = "DashAI/back/user_models/temp_checkpoints_opus-mt-fr-en" SCHEMA = OpusMtFrEnTransformerSchema + DOWNLOAD_SIZE_BYTES = 1204539675 DISPLAY_NAME: str = MultilingualString( en="Opus MT Fr-En Transformer", es="Transformer Opus MT Fr-En", diff --git a/DashAI/back/models/hugging_face/pixart_sigma_model.py b/DashAI/back/models/hugging_face/pixart_sigma_model.py index 484228ea8..7d4596380 100644 --- a/DashAI/back/models/hugging_face/pixart_sigma_model.py +++ b/DashAI/back/models/hugging_face/pixart_sigma_model.py @@ -512,7 +512,7 @@ class PixArtSigma1024(PixArtSigmaGenerationModel): """ MODEL_NAME: str = "PixArt-alpha/PixArt-Sigma-XL-2-1024-MS" - DOWNLOAD_SIZE_BYTES: int = 2500000000 + DOWNLOAD_SIZE_BYTES: int = 21832490389 DISPLAY_NAME = MultilingualString( en="PixArt-Sigma 1024", es="PixArt-Sigma 1024", @@ -570,7 +570,7 @@ class PixArtSigma512(PixArtSigmaGenerationModel): """ MODEL_NAME: str = "PixArt-alpha/PixArt-Sigma-XL-2-512-MS" - DOWNLOAD_SIZE_BYTES: int = 2500000000 + DOWNLOAD_SIZE_BYTES: int = 2447758901 DISPLAY_NAME = MultilingualString( en="PixArt-Sigma 512", es="PixArt-Sigma 512", diff --git a/DashAI/back/models/hugging_face/qwen_model.py b/DashAI/back/models/hugging_face/qwen_model.py index 56afdcb6c..e37706a18 100644 --- a/DashAI/back/models/hugging_face/qwen_model.py +++ b/DashAI/back/models/hugging_face/qwen_model.py @@ -23,7 +23,7 @@ class Qwen25_05BInstruct(GGUFTextGenerationModel): # noqa: N801 REPO_ID = "Qwen/Qwen2.5-0.5B-Instruct-GGUF" GGUF_PATTERN = "*8_0.gguf" - DOWNLOAD_SIZE_BYTES = 700_000_000 + DOWNLOAD_SIZE_BYTES = 675710816 SCHEMA = GGUFTextGenerationSchema COLOR: str = "#2e7d32" DISPLAY_NAME = MultilingualString( @@ -88,7 +88,7 @@ class Qwen25_15BInstruct(GGUFTextGenerationModel): # noqa: N801 REPO_ID = "Qwen/Qwen2.5-1.5B-Instruct-GGUF" GGUF_PATTERN = "*8_0.gguf" - DOWNLOAD_SIZE_BYTES = 1_900_000_000 + DOWNLOAD_SIZE_BYTES = 1894532128 SCHEMA = GGUFTextGenerationSchema COLOR: str = "#2e7d32" DISPLAY_NAME = MultilingualString( diff --git a/DashAI/back/models/hugging_face/roberta_transformer.py b/DashAI/back/models/hugging_face/roberta_transformer.py index e3383a926..2a9414640 100644 --- a/DashAI/back/models/hugging_face/roberta_transformer.py +++ b/DashAI/back/models/hugging_face/roberta_transformer.py @@ -58,5 +58,5 @@ class RobertaTransformer(HuggingFaceTextClassificationTransformer): ICON: str = "SmartToy" SCHEMA = DistilBertTransformerSchema MODEL_NAME: str = "roberta-base" - DOWNLOAD_SIZE_BYTES: int = 500_000_000 + DOWNLOAD_SIZE_BYTES: int = 2815192007 TEMP_CHECKPOINT_DIR: str = "DashAI/back/user_models/temp_checkpoints_roberta" diff --git a/DashAI/back/models/hugging_face/sd15_depth_controlnet_model.py b/DashAI/back/models/hugging_face/sd15_depth_controlnet_model.py index cead31cfb..206aa3d76 100644 --- a/DashAI/back/models/hugging_face/sd15_depth_controlnet_model.py +++ b/DashAI/back/models/hugging_face/sd15_depth_controlnet_model.py @@ -263,7 +263,7 @@ class SD15DepthControlNetModel(HFDownloadableMixin, BaseControlNetModel): ("lllyasviel/sd-controlnet-depth", "model"), ("Intel/dpt-hybrid-midas", "model"), ] - DOWNLOAD_SIZE_BYTES = 5900000000 + DOWNLOAD_SIZE_BYTES = 50640633273 COLOR: str = "#4e342e" DISPLAY_NAME: str = MultilingualString( en="SD 1.5 Depth ControlNet", diff --git a/DashAI/back/models/hugging_face/sd15_hed_controlnet_model.py b/DashAI/back/models/hugging_face/sd15_hed_controlnet_model.py index 461a25a4d..c1647539f 100644 --- a/DashAI/back/models/hugging_face/sd15_hed_controlnet_model.py +++ b/DashAI/back/models/hugging_face/sd15_hed_controlnet_model.py @@ -196,7 +196,7 @@ class SD15HEDControlNetModel(HFDownloadableMixin, BaseControlNetModel): ("lllyasviel/sd-controlnet-hed", "model"), ("lllyasviel/Annotators", "model"), ] - DOWNLOAD_SIZE_BYTES = 7400000000 + DOWNLOAD_SIZE_BYTES = 60738022767 COLOR: str = "#006064" DISPLAY_NAME: str = MultilingualString( en="SD 1.5 HED ControlNet", diff --git a/DashAI/back/models/hugging_face/sd15_openpose_controlnet_model.py b/DashAI/back/models/hugging_face/sd15_openpose_controlnet_model.py index 1036adf31..6a37aa83e 100644 --- a/DashAI/back/models/hugging_face/sd15_openpose_controlnet_model.py +++ b/DashAI/back/models/hugging_face/sd15_openpose_controlnet_model.py @@ -191,7 +191,7 @@ class SD15OpenPoseControlNetModel(HFDownloadableMixin, BaseControlNetModel): ("lllyasviel/sd-controlnet-openpose", "model"), ("lllyasviel/Annotators", "model"), ] - DOWNLOAD_SIZE_BYTES = 7400000000 + DOWNLOAD_SIZE_BYTES = 60737694276 COLOR: str = "#880e4f" DISPLAY_NAME: str = MultilingualString( en="SD 1.5 OpenPose ControlNet", diff --git a/DashAI/back/models/hugging_face/sdxl_canny_controlnet_model.py b/DashAI/back/models/hugging_face/sdxl_canny_controlnet_model.py index c7b603a84..8f5e7e0a1 100644 --- a/DashAI/back/models/hugging_face/sdxl_canny_controlnet_model.py +++ b/DashAI/back/models/hugging_face/sdxl_canny_controlnet_model.py @@ -296,7 +296,7 @@ class SDXLCannyControlNetModel(HFDownloadableMixin, BaseControlNetModel): ("diffusers/controlnet-canny-sdxl-1.0", "model"), ("madebyollin/sdxl-vae-fp16-fix", "model"), ] - DOWNLOAD_SIZE_BYTES = 10000000000 + DOWNLOAD_SIZE_BYTES = 93308273070 COLOR: str = "#1a237e" DISPLAY_NAME: str = MultilingualString( en="SDXL Canny ControlNet", diff --git a/DashAI/back/models/hugging_face/sdxl_turbo_model.py b/DashAI/back/models/hugging_face/sdxl_turbo_model.py index 6d7a16f52..0e65a33ee 100644 --- a/DashAI/back/models/hugging_face/sdxl_turbo_model.py +++ b/DashAI/back/models/hugging_face/sdxl_turbo_model.py @@ -328,7 +328,7 @@ class SDXLTurboModel(HFPretrainedDownloadMixin, TextToImageGenerationTaskModel): SCHEMA = SDXLTurboSchema MODEL_NAME: str = "stabilityai/sdxl-turbo" # SDXL-Turbo diffusers pipeline is ~7 GB. - DOWNLOAD_SIZE_BYTES: int = 7_000_000_000 + DOWNLOAD_SIZE_BYTES: int = 55516176914 COLOR: str = "#b71c1c" DISPLAY_NAME: str = MultilingualString( en="SDXL Turbo", diff --git a/DashAI/back/models/hugging_face/smol_lm_model.py b/DashAI/back/models/hugging_face/smol_lm_model.py index b2999369d..e550a19b9 100644 --- a/DashAI/back/models/hugging_face/smol_lm_model.py +++ b/DashAI/back/models/hugging_face/smol_lm_model.py @@ -21,7 +21,7 @@ class SmolLM2_360MInstruct(GGUFTextGenerationModel): # noqa: N801 REPO_ID = "HuggingFaceTB/SmolLM2-360M-Instruct-GGUF" GGUF_PATTERN = "*q8_0.gguf" - DOWNLOAD_SIZE_BYTES = 400_000_000 + DOWNLOAD_SIZE_BYTES = 386404992 SCHEMA = GGUFTextGenerationSchema COLOR: str = "#00695c" DISPLAY_NAME = MultilingualString( @@ -79,7 +79,7 @@ class SmolLM2_17BInstruct(GGUFTextGenerationModel): # noqa: N801 REPO_ID = "HuggingFaceTB/SmolLM2-1.7B-Instruct-GGUF" GGUF_PATTERN = "*q4_k_m.gguf" - DOWNLOAD_SIZE_BYTES = 1_100_000_000 + DOWNLOAD_SIZE_BYTES = 1055609536 SCHEMA = GGUFTextGenerationSchema COLOR: str = "#00695c" DISPLAY_NAME = MultilingualString( diff --git a/DashAI/back/models/hugging_face/stable_diffusion_v1_depth_controlnet.py b/DashAI/back/models/hugging_face/stable_diffusion_v1_depth_controlnet.py index 5965f1c43..b21cfdfbd 100644 --- a/DashAI/back/models/hugging_face/stable_diffusion_v1_depth_controlnet.py +++ b/DashAI/back/models/hugging_face/stable_diffusion_v1_depth_controlnet.py @@ -217,7 +217,7 @@ class StableDiffusionXLV1ControlNet(HFDownloadableMixin, BaseControlNetModel): ("madebyollin/sdxl-vae-fp16-fix", "model"), ("Intel/dpt-hybrid-midas", "model"), ] - DOWNLOAD_SIZE_BYTES = 10500000000 + DOWNLOAD_SIZE_BYTES = 80696728644 COLOR: str = "#e65100" DISPLAY_NAME: str = MultilingualString( en="Stable Diffusion XL V1 ControlNet", diff --git a/DashAI/back/models/hugging_face/stable_diffusion_v2_model.py b/DashAI/back/models/hugging_face/stable_diffusion_v2_model.py index 55ff085d9..3ed411604 100644 --- a/DashAI/back/models/hugging_face/stable_diffusion_v2_model.py +++ b/DashAI/back/models/hugging_face/stable_diffusion_v2_model.py @@ -546,7 +546,7 @@ class StableDiffusion2(StableDiffusion2GenerationModel): MODEL_NAME: str = "sd2-community/stable-diffusion-2" # Full fp32 diffusers pipeline (text encoder + U-Net + VAE) is ~5 GB. - DOWNLOAD_SIZE_BYTES: int = 5_200_000_000 + DOWNLOAD_SIZE_BYTES: int = 25911933905 DISPLAY_NAME = MultilingualString( en="Stable Diffusion 2", es="Stable Diffusion 2", @@ -606,7 +606,7 @@ class StableDiffusion2_512(StableDiffusion2GenerationModel): # noqa: N801 MODEL_NAME: str = "sd2-community/stable-diffusion-2-base" # Full fp32 diffusers pipeline (text encoder + U-Net + VAE) is ~5 GB. - DOWNLOAD_SIZE_BYTES: int = 5_200_000_000 + DOWNLOAD_SIZE_BYTES: int = 25911843836 DISPLAY_NAME = MultilingualString( en="Stable Diffusion 2 (512px)", es="Stable Diffusion 2 (512px)", @@ -661,7 +661,7 @@ class StableDiffusion21(StableDiffusion2GenerationModel): MODEL_NAME: str = "sd2-community/stable-diffusion-2-1" # Full fp32 diffusers pipeline (text encoder + U-Net + VAE) is ~5 GB. - DOWNLOAD_SIZE_BYTES: int = 5_200_000_000 + DOWNLOAD_SIZE_BYTES: int = 36341303572 DISPLAY_NAME = MultilingualString( en="Stable Diffusion 2.1", es="Stable Diffusion 2.1", @@ -717,7 +717,7 @@ class StableDiffusion21_512(StableDiffusion2GenerationModel): # noqa: N801 MODEL_NAME: str = "sd2-community/stable-diffusion-2-1-base" # Full fp32 diffusers pipeline (text encoder + U-Net + VAE) is ~5 GB. - DOWNLOAD_SIZE_BYTES: int = 5_200_000_000 + DOWNLOAD_SIZE_BYTES: int = 36341275775 DISPLAY_NAME = MultilingualString( en="Stable Diffusion 2.1 (512px)", es="Stable Diffusion 2.1 (512px)", diff --git a/DashAI/back/models/hugging_face/stable_diffusion_v3_model.py b/DashAI/back/models/hugging_face/stable_diffusion_v3_model.py index 3913d6134..ca8daef20 100644 --- a/DashAI/back/models/hugging_face/stable_diffusion_v3_model.py +++ b/DashAI/back/models/hugging_face/stable_diffusion_v3_model.py @@ -567,7 +567,7 @@ class StableDiffusion3Medium(StableDiffusion3GenerationModel): """ MODEL_NAME: str = "stabilityai/stable-diffusion-3-medium-diffusers" - DOWNLOAD_SIZE_BYTES: int = 5500000000 + DOWNLOAD_SIZE_BYTES: int = 31012147557 DISPLAY_NAME = MultilingualString( en="Stable Diffusion 3 Medium", es="Stable Diffusion 3 Medium", @@ -634,7 +634,7 @@ class StableDiffusion35Medium(StableDiffusion3GenerationModel): """ MODEL_NAME: str = "stabilityai/stable-diffusion-3.5-medium" - DOWNLOAD_SIZE_BYTES: int = 10000000000 + DOWNLOAD_SIZE_BYTES: int = 48861581303 DISPLAY_NAME = MultilingualString( en="Stable Diffusion 3.5 Medium", es="Stable Diffusion 3.5 Medium", @@ -699,7 +699,7 @@ class StableDiffusion35Large(StableDiffusion3GenerationModel): """ MODEL_NAME: str = "stabilityai/stable-diffusion-3.5-large" - DOWNLOAD_SIZE_BYTES: int = 16000000000 + DOWNLOAD_SIZE_BYTES: int = 71585723216 DISPLAY_NAME = MultilingualString( en="Stable Diffusion 3.5 Large", es="Stable Diffusion 3.5 Large", @@ -765,7 +765,7 @@ class StableDiffusion35LargeTurbo(StableDiffusion3GenerationModel): """ MODEL_NAME: str = "stabilityai/stable-diffusion-3.5-large-turbo" - DOWNLOAD_SIZE_BYTES: int = 16000000000 + DOWNLOAD_SIZE_BYTES: int = 71582971259 DISPLAY_NAME = MultilingualString( en="Stable Diffusion 3.5 Large Turbo", es="Stable Diffusion 3.5 Large Turbo", diff --git a/DashAI/back/models/hugging_face/stable_diffusion_xl_model.py b/DashAI/back/models/hugging_face/stable_diffusion_xl_model.py index da0ec4463..5c690e363 100644 --- a/DashAI/back/models/hugging_face/stable_diffusion_xl_model.py +++ b/DashAI/back/models/hugging_face/stable_diffusion_xl_model.py @@ -494,7 +494,7 @@ class StableDiffusionXL(StableDiffusionXLGenerationModel): MODEL_NAME: str = "stabilityai/stable-diffusion-xl-base-1.0" # SDXL diffusers pipeline (base + refiner-less) is ~7 GB. - DOWNLOAD_SIZE_BYTES: int = 7_000_000_000 + DOWNLOAD_SIZE_BYTES: int = 76912765291 DISPLAY_NAME = MultilingualString( en="Stable Diffusion XL", es="Stable Diffusion XL", @@ -554,7 +554,7 @@ class RealVisXLV4(StableDiffusionXLGenerationModel): MODEL_NAME: str = "SG161222/RealVisXL_V4.0" # SDXL diffusers pipeline (base + refiner-less) is ~7 GB. - DOWNLOAD_SIZE_BYTES: int = 7_000_000_000 + DOWNLOAD_SIZE_BYTES: int = 27754923233 DISPLAY_NAME = MultilingualString( en="RealVisXL V4.0", es="RealVisXL V4.0", diff --git a/DashAI/back/models/hugging_face/t5_small_transformer.py b/DashAI/back/models/hugging_face/t5_small_transformer.py index 997199351..5a061c63b 100644 --- a/DashAI/back/models/hugging_face/t5_small_transformer.py +++ b/DashAI/back/models/hugging_face/t5_small_transformer.py @@ -131,7 +131,7 @@ class T5SmallTransformer(HFPretrainedDownloadMixin, TranslationModel): COLOR: str = "#00695C" ICON: str = "Language" MODEL_NAME: str = "t5-small" - DOWNLOAD_SIZE_BYTES: int = 240_000_000 + DOWNLOAD_SIZE_BYTES: int = 2246973515 def __init__(self, model=None, pretrained_dir=None, **kwargs): kwargs = self.validate_and_transform(kwargs) diff --git a/DashAI/back/models/hugging_face/tongyi_z_image_model.py b/DashAI/back/models/hugging_face/tongyi_z_image_model.py index 4f3cd9bfe..30ac8f3c8 100644 --- a/DashAI/back/models/hugging_face/tongyi_z_image_model.py +++ b/DashAI/back/models/hugging_face/tongyi_z_image_model.py @@ -445,7 +445,7 @@ class TongyiZImage(TongyiZImageGenerationModel): """ MODEL_NAME: str = "Tongyi-MAI/Z-Image" - DOWNLOAD_SIZE_BYTES: int = 8000000000 + DOWNLOAD_SIZE_BYTES: int = 20547479575 DISPLAY_NAME = MultilingualString( en="Tongyi Z-Image", es="Tongyi Z-Image", @@ -498,7 +498,7 @@ class TongyiZImageTurbo(TongyiZImageGenerationModel): """ MODEL_NAME: str = "Tongyi-MAI/Z-Image-Turbo" - DOWNLOAD_SIZE_BYTES: int = 8000000000 + DOWNLOAD_SIZE_BYTES: int = 32899667397 DISPLAY_NAME = MultilingualString( en="Tongyi Z-Image Turbo", es="Tongyi Z-Image Turbo", diff --git a/DashAI/back/models/hugging_face/xlm_roberta_transformer.py b/DashAI/back/models/hugging_face/xlm_roberta_transformer.py index caf06a6a4..12d1941ab 100644 --- a/DashAI/back/models/hugging_face/xlm_roberta_transformer.py +++ b/DashAI/back/models/hugging_face/xlm_roberta_transformer.py @@ -62,5 +62,5 @@ class XlmRobertaTransformer(HuggingFaceTextClassificationTransformer): ICON: str = "Language" SCHEMA = DistilBertTransformerSchema MODEL_NAME: str = "xlm-roberta-base" - DOWNLOAD_SIZE_BYTES: int = 1_100_000_000 + DOWNLOAD_SIZE_BYTES: int = 6352430498 TEMP_CHECKPOINT_DIR: str = "DashAI/back/user_models/temp_checkpoints_xlm_roberta" diff --git a/DashAI/back/models/hugging_face/xlnet_transformer.py b/DashAI/back/models/hugging_face/xlnet_transformer.py index 9a7f59050..63a914b2b 100644 --- a/DashAI/back/models/hugging_face/xlnet_transformer.py +++ b/DashAI/back/models/hugging_face/xlnet_transformer.py @@ -58,5 +58,5 @@ class XlnetTransformer(HuggingFaceTextClassificationTransformer): ICON: str = "AutoAwesome" SCHEMA = DistilBertTransformerSchema MODEL_NAME: str = "xlnet-base-cased" - DOWNLOAD_SIZE_BYTES: int = 470_000_000 + DOWNLOAD_SIZE_BYTES: int = 1600067104 TEMP_CHECKPOINT_DIR: str = "DashAI/back/user_models/temp_checkpoints_xlnet" From b50f17fba71534ed562e6cfb109075e5f1e8e560 Mon Sep 17 00:00:00 2001 From: Irozuku Date: Thu, 9 Jul 2026 09:38:24 -0400 Subject: [PATCH 77/89] fix: adapt Opus En-Roa and Roa-En models to the download system These two Opus-MT models were merged in still describing weights as downloading on first use and inheriting the generic 300 MB size. Correct their descriptions to state weights must be downloaded before use, and set their real download sizes (1.17 GB / 1.21 GB). --- .../hugging_face/opus_mt_en_roa_transformer.py | 17 +++++++++-------- .../hugging_face/opus_mt_roa_en_transformer.py | 17 +++++++++-------- 2 files changed, 18 insertions(+), 16 deletions(-) diff --git a/DashAI/back/models/hugging_face/opus_mt_en_roa_transformer.py b/DashAI/back/models/hugging_face/opus_mt_en_roa_transformer.py index 3c68ff48b..9845bf231 100644 --- a/DashAI/back/models/hugging_face/opus_mt_en_roa_transformer.py +++ b/DashAI/back/models/hugging_face/opus_mt_en_roa_transformer.py @@ -81,6 +81,7 @@ class OpusMtEnRoaTransformer(OpusMtTransformerMixin): MODEL_NAME: str = "Helsinki-NLP/opus-mt-en-roa" TEMP_CHECKPOINT_DIR: str = "DashAI/back/user_models/temp_checkpoints_opus-mt-en-roa" SCHEMA = OpusMtEnRoaTransformerSchema + DOWNLOAD_SIZE_BYTES = 1171638932 DISPLAY_NAME: str = MultilingualString( en="Opus MT En-Roa Transformer", es="Transformer Opus MT En-Roa", @@ -92,32 +93,32 @@ class OpusMtEnRoaTransformer(OpusMtTransformerMixin): en=( "Pretrained transformer for English to Romance translation " "(Portuguese, Spanish, French, Italian, Romanian, Catalan, " - "Galician), selected via the target language parameter. Downloads " - "weights from Hugging Face on first use (internet required)." + "Galician), selected via the target language parameter. Download " + "its weights from Hugging Face before use (internet required)." ), es=( "Transformer preentrenado para traducción del inglés a lenguas " "romances (portugués, español, francés, italiano, rumano, catalán, " "gallego), seleccionadas con el parámetro de idioma de destino. " - "Descarga pesos de Hugging Face en el primer uso (requiere internet)." + "Descarga sus pesos de Hugging Face antes de usarlo (requiere internet)." ), pt=( "Transformer pré-treinado para tradução do inglês para línguas " "românicas (português, espanhol, francês, italiano, romeno, catalão, " - "galego), selecionadas pelo parâmetro de idioma de destino. Baixa os " - "pesos do Hugging Face no primeiro uso (requer internet)." + "galego), selecionadas pelo parâmetro de idioma de destino. Baixe " + "seus pesos do Hugging Face antes de usar (requer internet)." ), de=( "Vortrainierter Transformer für die Übersetzung von Englisch in " "romanische Sprachen (Portugiesisch, Spanisch, Französisch, " "Italienisch, Rumänisch, Katalanisch, Galicisch), ausgewählt über " - "den Zielsprachenparameter. Lädt Gewichte von Hugging Face bei der " - "ersten Verwendung herunter (Internet erforderlich)." + "den Zielsprachenparameter. Lädt die Gewichte vor der Nutzung von " + "Hugging Face herunter (Internet erforderlich)." ), zh=( "用于英语到罗曼语翻译的预训练 Transformer（葡萄牙语、西班牙语、法语、" "意大利语、罗马尼亚语、加泰罗尼亚语、加利西亚语），通过目标语言参数选择。" - "首次使用时从 Hugging Face 下载权重（需要网络）。" + "使用前需从 Hugging Face 下载权重（需要网络）。" ), ) COLOR: str = "#5E35B1" diff --git a/DashAI/back/models/hugging_face/opus_mt_roa_en_transformer.py b/DashAI/back/models/hugging_face/opus_mt_roa_en_transformer.py index 34b9e818c..bfb2f43b0 100644 --- a/DashAI/back/models/hugging_face/opus_mt_roa_en_transformer.py +++ b/DashAI/back/models/hugging_face/opus_mt_roa_en_transformer.py @@ -31,6 +31,7 @@ class OpusMtRoaEnTransformer(OpusMtTransformerMixin): MODEL_NAME: str = "Helsinki-NLP/opus-mt-roa-en" TEMP_CHECKPOINT_DIR: str = "DashAI/back/user_models/temp_checkpoints_opus-mt-roa-en" SCHEMA = OpusMtRoaEnTransformerSchema + DOWNLOAD_SIZE_BYTES = 1206056595 DISPLAY_NAME: str = MultilingualString( en="Opus MT Roa-En Transformer", es="Transformer Opus MT Roa-En", @@ -41,28 +42,28 @@ class OpusMtRoaEnTransformer(OpusMtTransformerMixin): DESCRIPTION: str = MultilingualString( en=( "Pretrained transformer for Romance to English translation " - "(includes Portuguese to English). Downloads weights from Hugging " - "Face on first use (internet required)." + "(includes Portuguese to English). Download its weights from " + "Hugging Face before use (internet required)." ), es=( "Transformer preentrenado para traducción de lenguas romances al " - "inglés (incluye portugués a inglés). Descarga pesos de Hugging " - "Face en el primer uso (requiere internet)." + "inglés (incluye portugués a inglés). Descarga sus pesos de " + "Hugging Face antes de usarlo (requiere internet)." ), pt=( "Transformer pré-treinado para tradução de línguas românicas para o " - "inglês (inclui português para inglês). Baixa os pesos do Hugging " - "Face no primeiro uso (requer internet)." + "inglês (inclui português para inglês). Baixe seus pesos do " + "Hugging Face antes de usar (requer internet)." ), de=( "Vortrainierter Transformer für die Übersetzung romanischer Sprachen " "ins Englische (einschließlich Portugiesisch nach Englisch). Lädt " - "Gewichte von Hugging Face bei der ersten Verwendung herunter " + "die Gewichte vor der Nutzung von Hugging Face herunter " "(Internet erforderlich)." ), zh=( "用于罗曼语到英语翻译的预训练 Transformer（包括葡萄牙语到英语）。" - "首次使用时从 Hugging Face 下载权重（需要网络）。" + "使用前需从 Hugging Face 下载权重（需要网络）。" ), ) COLOR: str = "#00796B" From cdc8e0df13263228a4ff83a4cf0004fa8e04c155 Mon Sep 17 00:00:00 2001 From: Irozuku Date: Thu, 9 Jul 2026 09:49:28 -0400 Subject: [PATCH 78/89] feat: skip alternate framework weights when downloading HF models HF repos ship TensorFlow, Flax, Rust, ONNX, OpenVINO and CoreML copies of the weights that from_pretrained never uses here, so downloads were several times larger than needed. Add HF_IGNORE_PATTERNS (applied to every snapshot_download) to skip them, keeping both .bin and .safetensors so fine-tuning and inference still load, and update DOWNLOAD_SIZE_BYTES to the slimmer real sizes (e.g. Bert 3.4GB->882MB, T5 2.2GB->486MB, opus ~1.2GB->~300MB, SDXL 77GB->35GB). --- .../dependencies/downloads/downloadable.py | 23 +++++++++++++++++++ .../models/hugging_face/albert_transformer.py | 2 +- .../models/hugging_face/bert_transformer.py | 2 +- .../models/hugging_face/bertin_transformer.py | 2 +- .../models/hugging_face/beto_transformer.py | 2 +- .../hugging_face/deberta_v3_transformer.py | 2 +- .../hugging_face/distilbert_transformer.py | 2 +- .../hugging_face/electra_transformer.py | 2 +- .../models/hugging_face/m2m100_transformer.py | 2 +- .../models/hugging_face/minilm_transformer.py | 2 +- .../hugging_face/modernbert_transformer.py | 2 +- .../multilingual_bert_transformer.py | 2 +- .../hugging_face/opus_mt_en_de_transformer.py | 2 +- .../hugging_face/opus_mt_en_es_transformer.py | 2 +- .../hugging_face/opus_mt_en_fr_transformer.py | 2 +- .../opus_mt_en_roa_transformer.py | 2 +- .../hugging_face/opus_mt_es_en_transformer.py | 2 +- .../hugging_face/opus_mt_fr_en_transformer.py | 2 +- .../opus_mt_roa_en_transformer.py | 2 +- .../hugging_face/roberta_transformer.py | 2 +- .../sdxl_canny_controlnet_model.py | 2 +- .../models/hugging_face/sdxl_turbo_model.py | 2 +- .../stable_diffusion_v1_depth_controlnet.py | 2 +- .../hugging_face/stable_diffusion_xl_model.py | 2 +- .../hugging_face/t5_small_transformer.py | 2 +- .../hugging_face/xlm_roberta_transformer.py | 2 +- .../models/hugging_face/xlnet_transformer.py | 2 +- 27 files changed, 49 insertions(+), 26 deletions(-) diff --git a/DashAI/back/dependencies/downloads/downloadable.py b/DashAI/back/dependencies/downloads/downloadable.py index 2334b25a0..5b94f43a1 100644 --- a/DashAI/back/dependencies/downloads/downloadable.py +++ b/DashAI/back/dependencies/downloads/downloadable.py @@ -75,9 +75,30 @@ class HFDownloadableMixin(DownloadableMixin): * ``(repo_id, repo_type)`` -- full snapshot download (original behavior). * ``(repo_id, repo_type, allow_patterns)`` -- partial download; only files matching the glob patterns in ``allow_patterns`` are fetched. + + ``HF_IGNORE_PATTERNS`` is applied to every repo download to skip the + non-PyTorch weight formats HuggingFace repos ship alongside the PyTorch / + safetensors weights (TensorFlow, Flax, Rust, ONNX, OpenVINO, CoreML). These + are never used by ``from_pretrained`` here, so dropping them shrinks the + download without affecting fine-tuning or inference. It keeps both ``.bin`` + and ``.safetensors`` so any model still has a loadable weight. """ HF_REPOS: List[Union[Tuple[str, str], Tuple[str, str, List[str]]]] = [] + #: Alternate-framework artifacts to skip on every download (``*`` matches + #: path separators in ``huggingface_hub`` glob semantics, so these match at + #: any depth, e.g. ``unet/diffusion_flax_model.msgpack``). + HF_IGNORE_PATTERNS: Optional[List[str]] = [ + "*.h5", + "*.msgpack", + "*.ot", + "*.onnx", + "*.onnx_data", + "*.tflite", + "*.mlmodel", + "*openvino*", + "*coreml*", + ] @classmethod def hf_repos(cls) -> List[Union[Tuple[str, str], Tuple[str, str, List[str]]]]: @@ -206,6 +227,8 @@ def download(cls, report: Optional[ProgressReporter] = None) -> None: kwargs = {} if allow_patterns is not None: kwargs["allow_patterns"] = allow_patterns + if cls.HF_IGNORE_PATTERNS: + kwargs["ignore_patterns"] = list(cls.HF_IGNORE_PATTERNS) snapshot_download( repo_id=rid, repo_type=rtype, local_dir=str(target), **kwargs ) diff --git a/DashAI/back/models/hugging_face/albert_transformer.py b/DashAI/back/models/hugging_face/albert_transformer.py index 1ace38ad1..25b686d0a 100644 --- a/DashAI/back/models/hugging_face/albert_transformer.py +++ b/DashAI/back/models/hugging_face/albert_transformer.py @@ -59,5 +59,5 @@ class AlbertTransformer(HuggingFaceTextClassificationTransformer): ICON: str = "Speed" SCHEMA = DistilBertTransformerSchema MODEL_NAME: str = "albert-base-v2" - DOWNLOAD_SIZE_BYTES: int = 330556087 + DOWNLOAD_SIZE_BYTES: int = 96833451 TEMP_CHECKPOINT_DIR: str = "DashAI/back/user_models/temp_checkpoints_albert" diff --git a/DashAI/back/models/hugging_face/bert_transformer.py b/DashAI/back/models/hugging_face/bert_transformer.py index 594fbc0e3..508448e3d 100644 --- a/DashAI/back/models/hugging_face/bert_transformer.py +++ b/DashAI/back/models/hugging_face/bert_transformer.py @@ -58,5 +58,5 @@ class BertTransformer(HuggingFaceTextClassificationTransformer): ICON: str = "Psychology" SCHEMA = DistilBertTransformerSchema MODEL_NAME: str = "bert-base-uncased" - DOWNLOAD_SIZE_BYTES: int = 3454102158 + DOWNLOAD_SIZE_BYTES: int = 881643453 TEMP_CHECKPOINT_DIR: str = "DashAI/back/user_models/temp_checkpoints_bert" diff --git a/DashAI/back/models/hugging_face/bertin_transformer.py b/DashAI/back/models/hugging_face/bertin_transformer.py index 984b1e5e2..f3b24ced4 100644 --- a/DashAI/back/models/hugging_face/bertin_transformer.py +++ b/DashAI/back/models/hugging_face/bertin_transformer.py @@ -58,5 +58,5 @@ class BertinTransformer(HuggingFaceTextClassificationTransformer): ICON: str = "RecordVoiceOver" SCHEMA = DistilBertTransformerSchema MODEL_NAME: str = "bertin-project/bertin-roberta-base-spanish" - DOWNLOAD_SIZE_BYTES: int = 1510386247 + DOWNLOAD_SIZE_BYTES: int = 1011598492 TEMP_CHECKPOINT_DIR: str = "DashAI/back/user_models/temp_checkpoints_bertin" diff --git a/DashAI/back/models/hugging_face/beto_transformer.py b/DashAI/back/models/hugging_face/beto_transformer.py index 78a549c05..3e4ecb1bb 100644 --- a/DashAI/back/models/hugging_face/beto_transformer.py +++ b/DashAI/back/models/hugging_face/beto_transformer.py @@ -58,5 +58,5 @@ class BetoTransformer(HuggingFaceTextClassificationTransformer): ICON: str = "RecordVoiceOver" SCHEMA = DistilBertTransformerSchema MODEL_NAME: str = "dccuchile/bert-base-spanish-wwm-cased" - DOWNLOAD_SIZE_BYTES: int = 1416527075 + DOWNLOAD_SIZE_BYTES: int = 440350800 TEMP_CHECKPOINT_DIR: str = "DashAI/back/user_models/temp_checkpoints_beto" diff --git a/DashAI/back/models/hugging_face/deberta_v3_transformer.py b/DashAI/back/models/hugging_face/deberta_v3_transformer.py index 71684667d..03bd82c5a 100644 --- a/DashAI/back/models/hugging_face/deberta_v3_transformer.py +++ b/DashAI/back/models/hugging_face/deberta_v3_transformer.py @@ -61,5 +61,5 @@ class DebertaV3Transformer(HuggingFaceTextClassificationTransformer): ICON: str = "Psychology" SCHEMA = DebertaV3TransformerSchema MODEL_NAME: str = "microsoft/deberta-v3-base" - DOWNLOAD_SIZE_BYTES: int = 1851424112 + DOWNLOAD_SIZE_BYTES: int = 373616107 TEMP_CHECKPOINT_DIR: str = "DashAI/back/user_models/temp_checkpoints_deberta_v3" diff --git a/DashAI/back/models/hugging_face/distilbert_transformer.py b/DashAI/back/models/hugging_face/distilbert_transformer.py index 84695f532..ebdff496d 100644 --- a/DashAI/back/models/hugging_face/distilbert_transformer.py +++ b/DashAI/back/models/hugging_face/distilbert_transformer.py @@ -305,5 +305,5 @@ class DistilBertTransformer(HuggingFaceTextClassificationTransformer): ICON: str = "Psychology" SCHEMA = DistilBertTransformerSchema MODEL_NAME: str = "distilbert-base-uncased" - DOWNLOAD_SIZE_BYTES: int = 1529742866 + DOWNLOAD_SIZE_BYTES: int = 536641210 TEMP_CHECKPOINT_DIR: str = "DashAI/back/user_models/temp_checkpoints_distilbert" diff --git a/DashAI/back/models/hugging_face/electra_transformer.py b/DashAI/back/models/hugging_face/electra_transformer.py index 088bd2aa7..43e8e65dc 100644 --- a/DashAI/back/models/hugging_face/electra_transformer.py +++ b/DashAI/back/models/hugging_face/electra_transformer.py @@ -58,5 +58,5 @@ class ElectraTransformer(HuggingFaceTextClassificationTransformer): ICON: str = "ElectricBolt" SCHEMA = DistilBertTransformerSchema MODEL_NAME: str = "google/electra-small-discriminator" - DOWNLOAD_SIZE_BYTES: int = 163615837 + DOWNLOAD_SIZE_BYTES: int = 54946248 TEMP_CHECKPOINT_DIR: str = "DashAI/back/user_models/temp_checkpoints_electra" diff --git a/DashAI/back/models/hugging_face/m2m100_transformer.py b/DashAI/back/models/hugging_face/m2m100_transformer.py index 601e5d93b..de5c618cb 100644 --- a/DashAI/back/models/hugging_face/m2m100_transformer.py +++ b/DashAI/back/models/hugging_face/m2m100_transformer.py @@ -158,7 +158,7 @@ class M2M100Transformer(HFPretrainedDownloadMixin, TranslationModel): COLOR: str = "#6A1B9A" ICON: str = "Language" MODEL_NAME: str = "facebook/m2m100_418M" - DOWNLOAD_SIZE_BYTES: int = 3877717593 + DOWNLOAD_SIZE_BYTES: int = 1941936305 def __init__(self, model=None, pretrained_dir=None, **kwargs): kwargs = self.validate_and_transform(kwargs) diff --git a/DashAI/back/models/hugging_face/minilm_transformer.py b/DashAI/back/models/hugging_face/minilm_transformer.py index 8888c8bc9..c9990e4e5 100644 --- a/DashAI/back/models/hugging_face/minilm_transformer.py +++ b/DashAI/back/models/hugging_face/minilm_transformer.py @@ -58,5 +58,5 @@ class MiniLMTransformer(HuggingFaceTextClassificationTransformer): ICON: str = "Speed" SCHEMA = DistilBertTransformerSchema MODEL_NAME: str = "microsoft/MiniLM-L12-H384-uncased" - DOWNLOAD_SIZE_BYTES: int = 400889386 + DOWNLOAD_SIZE_BYTES: int = 133721893 TEMP_CHECKPOINT_DIR: str = "DashAI/back/user_models/temp_checkpoints_minilm" diff --git a/DashAI/back/models/hugging_face/modernbert_transformer.py b/DashAI/back/models/hugging_face/modernbert_transformer.py index a087fc310..af58e5b2a 100644 --- a/DashAI/back/models/hugging_face/modernbert_transformer.py +++ b/DashAI/back/models/hugging_face/modernbert_transformer.py @@ -55,6 +55,6 @@ class ModernBertTransformer(HuggingFaceTextClassificationTransformer): ICON: str = "Psychology" SCHEMA = ModernBertTransformerSchema MODEL_NAME: str = "answerdotai/ModernBERT-base" - DOWNLOAD_SIZE_BYTES: int = 3134313772 + DOWNLOAD_SIZE_BYTES: int = 1199464688 MAX_TOKEN_LENGTH: int = 8192 TEMP_CHECKPOINT_DIR: str = "DashAI/back/user_models/temp_checkpoints_modernbert" diff --git a/DashAI/back/models/hugging_face/multilingual_bert_transformer.py b/DashAI/back/models/hugging_face/multilingual_bert_transformer.py index 5e3574f8f..a3359d142 100644 --- a/DashAI/back/models/hugging_face/multilingual_bert_transformer.py +++ b/DashAI/back/models/hugging_face/multilingual_bert_transformer.py @@ -59,7 +59,7 @@ class MultilingualBertTransformer(HuggingFaceTextClassificationTransformer): ICON: str = "Translate" SCHEMA = DistilBertTransformerSchema MODEL_NAME: str = "bert-base-multilingual-cased" - DOWNLOAD_SIZE_BYTES: int = 3226865011 + DOWNLOAD_SIZE_BYTES: int = 1431570300 TEMP_CHECKPOINT_DIR: str = ( "DashAI/back/user_models/temp_checkpoints_multilingual_bert" ) diff --git a/DashAI/back/models/hugging_face/opus_mt_en_de_transformer.py b/DashAI/back/models/hugging_face/opus_mt_en_de_transformer.py index 0c5854f6b..8007c5e97 100644 --- a/DashAI/back/models/hugging_face/opus_mt_en_de_transformer.py +++ b/DashAI/back/models/hugging_face/opus_mt_en_de_transformer.py @@ -29,7 +29,7 @@ class OpusMtEnDeTransformer(OpusMtTransformerMixin): MODEL_NAME: str = "Helsinki-NLP/opus-mt-en-de" TEMP_CHECKPOINT_DIR: str = "DashAI/back/user_models/temp_checkpoints_opus-mt-en-de" SCHEMA = OpusMtEnDeTransformerSchema - DOWNLOAD_SIZE_BYTES = 1430871510 + DOWNLOAD_SIZE_BYTES = 300772148 DISPLAY_NAME: str = MultilingualString( en="Opus MT En-De Transformer", es="Transformer Opus MT En-De", diff --git a/DashAI/back/models/hugging_face/opus_mt_en_es_transformer.py b/DashAI/back/models/hugging_face/opus_mt_en_es_transformer.py index 2c9cf281a..3ed97502d 100644 --- a/DashAI/back/models/hugging_face/opus_mt_en_es_transformer.py +++ b/DashAI/back/models/hugging_face/opus_mt_en_es_transformer.py @@ -270,7 +270,7 @@ class OpusMtEnESTransformer(OpusMtTransformerMixin): MODEL_NAME: str = "Helsinki-NLP/opus-mt-en-es" TEMP_CHECKPOINT_DIR: str = "DashAI/back/user_models/temp_checkpoints_opus-mt-en-es" SCHEMA = OpusMtEnESTransformerSchema - DOWNLOAD_SIZE_BYTES = 937836389 + DOWNLOAD_SIZE_BYTES = 315310815 DISPLAY_NAME: str = MultilingualString( en="Opus MT En-Es Transformer", es="Transformer Opus MT En-Es", diff --git a/DashAI/back/models/hugging_face/opus_mt_en_fr_transformer.py b/DashAI/back/models/hugging_face/opus_mt_en_fr_transformer.py index 9c409d3ad..18c791f16 100644 --- a/DashAI/back/models/hugging_face/opus_mt_en_fr_transformer.py +++ b/DashAI/back/models/hugging_face/opus_mt_en_fr_transformer.py @@ -29,7 +29,7 @@ class OpusMtEnFrTransformer(OpusMtTransformerMixin): MODEL_NAME: str = "Helsinki-NLP/opus-mt-en-fr" TEMP_CHECKPOINT_DIR: str = "DashAI/back/user_models/temp_checkpoints_opus-mt-en-fr" SCHEMA = OpusMtEnFrTransformerSchema - DOWNLOAD_SIZE_BYTES = 903735972 + DOWNLOAD_SIZE_BYTES = 303750994 DISPLAY_NAME: str = MultilingualString( en="Opus MT En-Fr Transformer", es="Transformer Opus MT En-Fr", diff --git a/DashAI/back/models/hugging_face/opus_mt_en_roa_transformer.py b/DashAI/back/models/hugging_face/opus_mt_en_roa_transformer.py index 9845bf231..c6a090c5d 100644 --- a/DashAI/back/models/hugging_face/opus_mt_en_roa_transformer.py +++ b/DashAI/back/models/hugging_face/opus_mt_en_roa_transformer.py @@ -81,7 +81,7 @@ class OpusMtEnRoaTransformer(OpusMtTransformerMixin): MODEL_NAME: str = "Helsinki-NLP/opus-mt-en-roa" TEMP_CHECKPOINT_DIR: str = "DashAI/back/user_models/temp_checkpoints_opus-mt-en-roa" SCHEMA = OpusMtEnRoaTransformerSchema - DOWNLOAD_SIZE_BYTES = 1171638932 + DOWNLOAD_SIZE_BYTES = 297844640 DISPLAY_NAME: str = MultilingualString( en="Opus MT En-Roa Transformer", es="Transformer Opus MT En-Roa", diff --git a/DashAI/back/models/hugging_face/opus_mt_es_en_transformer.py b/DashAI/back/models/hugging_face/opus_mt_es_en_transformer.py index 96b71168c..43d077460 100644 --- a/DashAI/back/models/hugging_face/opus_mt_es_en_transformer.py +++ b/DashAI/back/models/hugging_face/opus_mt_es_en_transformer.py @@ -32,7 +32,7 @@ class OpusMtEsENTransformer(OpusMtTransformerMixin): MODEL_NAME: str = "Helsinki-NLP/opus-mt-es-en" TEMP_CHECKPOINT_DIR: str = "DashAI/back/user_models/temp_checkpoints_opus-mt-es-en" SCHEMA = OpusMtEsENTransformerSchema - DOWNLOAD_SIZE_BYTES = 627891360 + DOWNLOAD_SIZE_BYTES = 315310760 DISPLAY_NAME: str = MultilingualString( en="Opus MT Es-En Transformer", es="Transformer Opus MT Es-En", diff --git a/DashAI/back/models/hugging_face/opus_mt_fr_en_transformer.py b/DashAI/back/models/hugging_face/opus_mt_fr_en_transformer.py index fcb7e397c..94f77975b 100644 --- a/DashAI/back/models/hugging_face/opus_mt_fr_en_transformer.py +++ b/DashAI/back/models/hugging_face/opus_mt_fr_en_transformer.py @@ -29,7 +29,7 @@ class OpusMtFrEnTransformer(OpusMtTransformerMixin): MODEL_NAME: str = "Helsinki-NLP/opus-mt-fr-en" TEMP_CHECKPOINT_DIR: str = "DashAI/back/user_models/temp_checkpoints_opus-mt-fr-en" SCHEMA = OpusMtFrEnTransformerSchema - DOWNLOAD_SIZE_BYTES = 1204539675 + DOWNLOAD_SIZE_BYTES = 604554697 DISPLAY_NAME: str = MultilingualString( en="Opus MT Fr-En Transformer", es="Transformer Opus MT Fr-En", diff --git a/DashAI/back/models/hugging_face/opus_mt_roa_en_transformer.py b/DashAI/back/models/hugging_face/opus_mt_roa_en_transformer.py index bfb2f43b0..5d1f189f0 100644 --- a/DashAI/back/models/hugging_face/opus_mt_roa_en_transformer.py +++ b/DashAI/back/models/hugging_face/opus_mt_roa_en_transformer.py @@ -31,7 +31,7 @@ class OpusMtRoaEnTransformer(OpusMtTransformerMixin): MODEL_NAME: str = "Helsinki-NLP/opus-mt-roa-en" TEMP_CHECKPOINT_DIR: str = "DashAI/back/user_models/temp_checkpoints_opus-mt-roa-en" SCHEMA = OpusMtRoaEnTransformerSchema - DOWNLOAD_SIZE_BYTES = 1206056595 + DOWNLOAD_SIZE_BYTES = 315135823 DISPLAY_NAME: str = MultilingualString( en="Opus MT Roa-En Transformer", es="Transformer Opus MT Roa-En", diff --git a/DashAI/back/models/hugging_face/roberta_transformer.py b/DashAI/back/models/hugging_face/roberta_transformer.py index 2a9414640..434a3bfd4 100644 --- a/DashAI/back/models/hugging_face/roberta_transformer.py +++ b/DashAI/back/models/hugging_face/roberta_transformer.py @@ -58,5 +58,5 @@ class RobertaTransformer(HuggingFaceTextClassificationTransformer): ICON: str = "SmartToy" SCHEMA = DistilBertTransformerSchema MODEL_NAME: str = "roberta-base" - DOWNLOAD_SIZE_BYTES: int = 2815192007 + DOWNLOAD_SIZE_BYTES: int = 1003342916 TEMP_CHECKPOINT_DIR: str = "DashAI/back/user_models/temp_checkpoints_roberta" diff --git a/DashAI/back/models/hugging_face/sdxl_canny_controlnet_model.py b/DashAI/back/models/hugging_face/sdxl_canny_controlnet_model.py index 8f5e7e0a1..83a5bd1d3 100644 --- a/DashAI/back/models/hugging_face/sdxl_canny_controlnet_model.py +++ b/DashAI/back/models/hugging_face/sdxl_canny_controlnet_model.py @@ -296,7 +296,7 @@ class SDXLCannyControlNetModel(HFDownloadableMixin, BaseControlNetModel): ("diffusers/controlnet-canny-sdxl-1.0", "model"), ("madebyollin/sdxl-vae-fp16-fix", "model"), ] - DOWNLOAD_SIZE_BYTES = 93308273070 + DOWNLOAD_SIZE_BYTES = 51644917846 COLOR: str = "#1a237e" DISPLAY_NAME: str = MultilingualString( en="SDXL Canny ControlNet", diff --git a/DashAI/back/models/hugging_face/sdxl_turbo_model.py b/DashAI/back/models/hugging_face/sdxl_turbo_model.py index 0e65a33ee..0b7c8cca7 100644 --- a/DashAI/back/models/hugging_face/sdxl_turbo_model.py +++ b/DashAI/back/models/hugging_face/sdxl_turbo_model.py @@ -328,7 +328,7 @@ class SDXLTurboModel(HFPretrainedDownloadMixin, TextToImageGenerationTaskModel): SCHEMA = SDXLTurboSchema MODEL_NAME: str = "stabilityai/sdxl-turbo" # SDXL-Turbo diffusers pipeline is ~7 GB. - DOWNLOAD_SIZE_BYTES: int = 55516176914 + DOWNLOAD_SIZE_BYTES: int = 41631892171 COLOR: str = "#b71c1c" DISPLAY_NAME: str = MultilingualString( en="SDXL Turbo", diff --git a/DashAI/back/models/hugging_face/stable_diffusion_v1_depth_controlnet.py b/DashAI/back/models/hugging_face/stable_diffusion_v1_depth_controlnet.py index b21cfdfbd..03dec56b8 100644 --- a/DashAI/back/models/hugging_face/stable_diffusion_v1_depth_controlnet.py +++ b/DashAI/back/models/hugging_face/stable_diffusion_v1_depth_controlnet.py @@ -217,7 +217,7 @@ class StableDiffusionXLV1ControlNet(HFDownloadableMixin, BaseControlNetModel): ("madebyollin/sdxl-vae-fp16-fix", "model"), ("Intel/dpt-hybrid-midas", "model"), ] - DOWNLOAD_SIZE_BYTES = 80696728644 + DOWNLOAD_SIZE_BYTES = 39033373420 COLOR: str = "#e65100" DISPLAY_NAME: str = MultilingualString( en="Stable Diffusion XL V1 ControlNet", diff --git a/DashAI/back/models/hugging_face/stable_diffusion_xl_model.py b/DashAI/back/models/hugging_face/stable_diffusion_xl_model.py index 5c690e363..a5fe8e81c 100644 --- a/DashAI/back/models/hugging_face/stable_diffusion_xl_model.py +++ b/DashAI/back/models/hugging_face/stable_diffusion_xl_model.py @@ -494,7 +494,7 @@ class StableDiffusionXL(StableDiffusionXLGenerationModel): MODEL_NAME: str = "stabilityai/stable-diffusion-xl-base-1.0" # SDXL diffusers pipeline (base + refiner-less) is ~7 GB. - DOWNLOAD_SIZE_BYTES: int = 76912765291 + DOWNLOAD_SIZE_BYTES: int = 35249410067 DISPLAY_NAME = MultilingualString( en="Stable Diffusion XL", es="Stable Diffusion XL", diff --git a/DashAI/back/models/hugging_face/t5_small_transformer.py b/DashAI/back/models/hugging_face/t5_small_transformer.py index 5a061c63b..37a1ba8f6 100644 --- a/DashAI/back/models/hugging_face/t5_small_transformer.py +++ b/DashAI/back/models/hugging_face/t5_small_transformer.py @@ -131,7 +131,7 @@ class T5SmallTransformer(HFPretrainedDownloadMixin, TranslationModel): COLOR: str = "#00695C" ICON: str = "Language" MODEL_NAME: str = "t5-small" - DOWNLOAD_SIZE_BYTES: int = 2246973515 + DOWNLOAD_SIZE_BYTES: int = 486302401 def __init__(self, model=None, pretrained_dir=None, **kwargs): kwargs = self.validate_and_transform(kwargs) diff --git a/DashAI/back/models/hugging_face/xlm_roberta_transformer.py b/DashAI/back/models/hugging_face/xlm_roberta_transformer.py index 12d1941ab..f6cb74142 100644 --- a/DashAI/back/models/hugging_face/xlm_roberta_transformer.py +++ b/DashAI/back/models/hugging_face/xlm_roberta_transformer.py @@ -62,5 +62,5 @@ class XlmRobertaTransformer(HuggingFaceTextClassificationTransformer): ICON: str = "Language" SCHEMA = DistilBertTransformerSchema MODEL_NAME: str = "xlm-roberta-base" - DOWNLOAD_SIZE_BYTES: int = 6352430498 + DOWNLOAD_SIZE_BYTES: int = 2245330190 TEMP_CHECKPOINT_DIR: str = "DashAI/back/user_models/temp_checkpoints_xlm_roberta" diff --git a/DashAI/back/models/hugging_face/xlnet_transformer.py b/DashAI/back/models/hugging_face/xlnet_transformer.py index 63a914b2b..56b05428b 100644 --- a/DashAI/back/models/hugging_face/xlnet_transformer.py +++ b/DashAI/back/models/hugging_face/xlnet_transformer.py @@ -58,5 +58,5 @@ class XlnetTransformer(HuggingFaceTextClassificationTransformer): ICON: str = "AutoAwesome" SCHEMA = DistilBertTransformerSchema MODEL_NAME: str = "xlnet-base-cased" - DOWNLOAD_SIZE_BYTES: int = 1600067104 + DOWNLOAD_SIZE_BYTES: int = 469226606 TEMP_CHECKPOINT_DIR: str = "DashAI/back/user_models/temp_checkpoints_xlnet" From f51d7556a43a013ea597880606f7016887afc052 Mon Sep 17 00:00:00 2001 From: Irozuku Date: Thu, 9 Jul 2026 09:52:49 -0400 Subject: [PATCH 79/89] fix: force classic transfer path for model downloads The Xet transfer backend returns a 404 on its read-token endpoint for some repos (e.g. bert-base-uncased), aborting the whole snapshot download. Disable Xet before downloading so it uses the reliable HTTP/LFS path. --- DashAI/back/dependencies/downloads/downloadable.py | 13 +++++++++++++ 1 file changed, 13 insertions(+) diff --git a/DashAI/back/dependencies/downloads/downloadable.py b/DashAI/back/dependencies/downloads/downloadable.py index 5b94f43a1..555ede709 100644 --- a/DashAI/back/dependencies/downloads/downloadable.py +++ b/DashAI/back/dependencies/downloads/downloadable.py @@ -7,6 +7,7 @@ """ import logging +import os import pathlib import shutil from typing import Callable, List, Optional, Tuple, Union @@ -216,6 +217,18 @@ def download(cls, report: Optional[ProgressReporter] = None) -> None: ``report(None, "Downloading ")``. ``None`` means no progress reporting. """ + # Force the classic HTTP/LFS transfer path. The Xet backend can return + # a 404 on its read-token endpoint for some repos (e.g. bert-base- + # uncased), which aborts the whole download; the classic path is + # slower but reliable. + os.environ["HF_HUB_DISABLE_XET"] = "1" + try: + from huggingface_hub import constants as hf_constants + + hf_constants.HF_HUB_DISABLE_XET = True + except Exception: + pass + for entry in cls.hf_repos(): rid, rtype, allow_patterns = cls._unpack_entry(entry) target = cls._repo_dir(rid) From 12ea6734244509997000ee9322e78177b758817a Mon Sep 17 00:00:00 2001 From: Irozuku Date: Thu, 9 Jul 2026 10:39:42 -0400 Subject: [PATCH 80/89] fix: resolve PixArt-Sigma model source before loading PixArtSigma __init__ discarded the _pretrained_source(None) result and then read self.model_name, which was never set, raising AttributeError on generation. Assign the resolved source to self.model_name like the other diffusion models. --- DashAI/back/models/hugging_face/pixart_sigma_model.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/DashAI/back/models/hugging_face/pixart_sigma_model.py b/DashAI/back/models/hugging_face/pixart_sigma_model.py index 7d4596380..2fd17b261 100644 --- a/DashAI/back/models/hugging_face/pixart_sigma_model.py +++ b/DashAI/back/models/hugging_face/pixart_sigma_model.py @@ -458,7 +458,7 @@ def __init__(self, **kwargs): self.device = ( f"cuda:{DEVICE_TO_IDX.get(kwargs.get('device'))}" if use_gpu else "cpu" ) - self._pretrained_source(None) + self.model_name = self._pretrained_source(None) self.model = PixArtSigmaPipeline.from_pretrained( self.model_name, From 9263176de94a3a6ee14660fa67272f018623e3da Mon Sep 17 00:00:00 2001 From: Irozuku Date: Thu, 9 Jul 2026 10:50:43 -0400 Subject: [PATCH 81/89] fix: load PixArt-Sigma 512 from its transformer plus the 1024 pipeline The 512 checkpoint has no model_index.json (it ships only the transformer), so PixArtSigmaPipeline.from_pretrained on it failed with 'no file named model_index.json'. Download both the 512 (transformer) and 1024 (full pipeline) repos and build the pipeline from the 1024 scaffold with the 512 transformer injected, the standard PixArt-Sigma 512 recipe. Size updated to the combined download. --- .../models/hugging_face/pixart_sigma_model.py | 84 +++++++++++++++++-- 1 file changed, 75 insertions(+), 9 deletions(-) diff --git a/DashAI/back/models/hugging_face/pixart_sigma_model.py b/DashAI/back/models/hugging_face/pixart_sigma_model.py index 2fd17b261..dc855743f 100644 --- a/DashAI/back/models/hugging_face/pixart_sigma_model.py +++ b/DashAI/back/models/hugging_face/pixart_sigma_model.py @@ -450,20 +450,13 @@ def __init__(self, **kwargs): num_images_per_prompt : int Number of images to generate per prompt call. """ - import torch - from diffusers import PixArtSigmaPipeline - kwargs = self.validate_and_transform(kwargs) use_gpu = DEVICE_TO_IDX.get(kwargs.get("device")) >= 0 self.device = ( f"cuda:{DEVICE_TO_IDX.get(kwargs.get('device'))}" if use_gpu else "cpu" ) - self.model_name = self._pretrained_source(None) - self.model = PixArtSigmaPipeline.from_pretrained( - self.model_name, - torch_dtype=torch.float16 if use_gpu else torch.float32, - ).to(self.device) + self.model = self._build_pipeline(use_gpu).to(self.device) self.negative_prompt = kwargs.get("negative_prompt") self.num_inference_steps = kwargs.get("num_inference_steps") @@ -473,6 +466,32 @@ def __init__(self, **kwargs): self.height = kwargs.get("height") self.num_images_per_prompt = kwargs.get("num_images_per_prompt") + def _build_pipeline(self, use_gpu: bool): + """Load the PixArt-Sigma pipeline from the downloaded checkpoint. + + The default loads a self-contained checkpoint (one that ships a full + ``model_index.json`` pipeline, e.g. the 1024px variant). Subclasses + whose checkpoint contains only the transformer override this. + + Parameters + ---------- + use_gpu : bool + Whether a GPU is available (selects float16 vs float32). + + Returns + ------- + diffusers.PixArtSigmaPipeline + The loaded pipeline (not yet moved to a device). + """ + import torch + from diffusers import PixArtSigmaPipeline + + self.model_name = self._pretrained_source(None) + return PixArtSigmaPipeline.from_pretrained( + self.model_name, + torch_dtype=torch.float16 if use_gpu else torch.float32, + ) + def generate(self, input: str) -> List[Any]: """Generate images from a text prompt. @@ -570,7 +589,11 @@ class PixArtSigma512(PixArtSigmaGenerationModel): """ MODEL_NAME: str = "PixArt-alpha/PixArt-Sigma-XL-2-512-MS" - DOWNLOAD_SIZE_BYTES: int = 2447758901 + # The 512 checkpoint ships only the transformer; the T5 text encoder, VAE, + # scheduler and tokenizer are loaded from the 1024 checkpoint, so both repos + # are downloaded. Size is the sum of the two. + PIPELINE_REPO: str = "PixArt-alpha/PixArt-Sigma-XL-2-1024-MS" + DOWNLOAD_SIZE_BYTES: int = 24280249290 DISPLAY_NAME = MultilingualString( en="PixArt-Sigma 512", es="PixArt-Sigma 512", @@ -615,3 +638,46 @@ class PixArtSigma512(PixArtSigmaGenerationModel): "模型页面： https://huggingface.co/PixArt-alpha/PixArt-Sigma-XL-2-512-MS" ), ) + + @classmethod + def hf_repos(cls): + """Download the 512 transformer repo and the 1024 pipeline repo. + + Returns + ------- + list of tuple of (str, str) + The 512 checkpoint (transformer only) and the 1024 checkpoint + (full pipeline: T5, VAE, scheduler, tokenizer). + """ + return [(cls.MODEL_NAME, "model"), (cls.PIPELINE_REPO, "model")] + + def _build_pipeline(self, use_gpu: bool): + """Load the 512 transformer into the 1024 pipeline scaffold. + + The 512 repo has no ``model_index.json`` (it ships only the + transformer), so the pipeline is loaded from the 1024 repo with the + 512 transformer injected. + + Parameters + ---------- + use_gpu : bool + Whether a GPU is available (selects float16 vs float32). + + Returns + ------- + diffusers.PixArtSigmaPipeline + The loaded pipeline (not yet moved to a device). + """ + import torch + from diffusers import PixArtSigmaPipeline, Transformer2DModel + + dtype = torch.float16 if use_gpu else torch.float32 + transformer_dir = str(self._repo_dir(self.MODEL_NAME)) + pipeline_dir = str(self._repo_dir(self.PIPELINE_REPO)) + self.model_name = pipeline_dir + transformer = Transformer2DModel.from_pretrained( + transformer_dir, subfolder="transformer", torch_dtype=dtype + ) + return PixArtSigmaPipeline.from_pretrained( + pipeline_dir, transformer=transformer, torch_dtype=dtype + ) From 82d00698b269d1723ef35c66fba5c26b78945121 Mon Sep 17 00:00:00 2001 From: Irozuku Date: Thu, 9 Jul 2026 11:00:14 -0400 Subject: [PATCH 82/89] refactor: merge PixArt-Sigma 512 and 1024 into one component The 512 checkpoint has no standalone pipeline and needs the 1024 scaffold, so keeping them as separate components meant downloading the shared T5/VAE twice. Replace PixArtSigma512 and PixArtSigma1024 with a single PixArtSigma that downloads both repos once and picks the checkpoint (1024 or 512) via a schema parameter, injecting the 512 transformer into the 1024 pipeline when selected. --- DashAI/back/initial_components.py | 8 +- .../models/hugging_face/pixart_sigma_model.py | 256 ++++++------------ .../models/test_pixart_tongyi_downloadable.py | 30 +- 3 files changed, 104 insertions(+), 190 deletions(-) diff --git a/DashAI/back/initial_components.py b/DashAI/back/initial_components.py index c02f29c22..019ecca95 100644 --- a/DashAI/back/initial_components.py +++ b/DashAI/back/initial_components.py @@ -187,10 +187,7 @@ from DashAI.back.models.hugging_face.opus_mt_roa_en_transformer import ( OpusMtRoaEnTransformer, ) -from DashAI.back.models.hugging_face.pixart_sigma_model import ( - PixArtSigma512, - PixArtSigma1024, -) +from DashAI.back.models.hugging_face.pixart_sigma_model import PixArtSigma from DashAI.back.models.hugging_face.qwen_model import ( Qwen25_05BInstruct, Qwen25_15BInstruct, @@ -402,8 +399,7 @@ def get_initial_components(): OpusMtEnRoaTransformer, OpusMtEsENTransformer, OpusMtFrEnTransformer, - PixArtSigma1024, - PixArtSigma512, + PixArtSigma, Qwen25_05BInstruct, Qwen25_15BInstruct, OpusMtRoaEnTransformer, diff --git a/DashAI/back/models/hugging_face/pixart_sigma_model.py b/DashAI/back/models/hugging_face/pixart_sigma_model.py index dc855743f..20a5f02fb 100644 --- a/DashAI/back/models/hugging_face/pixart_sigma_model.py +++ b/DashAI/back/models/hugging_face/pixart_sigma_model.py @@ -21,7 +21,7 @@ class PixArtSigmaSchema(BaseSchema): """Configuration schema for PixArt-Sigma text-to-image generation. - Configures the checkpoint variant (``model_name``), prompt conditioning + Configures the checkpoint (``checkpoint``), prompt conditioning (``negative_prompt``), denoising schedule (``num_inference_steps``), classifier free guidance strength (``guidance_scale``), output dimensions (``width``, ``height``), reproducibility (``seed``), hardware target @@ -29,6 +29,46 @@ class PixArtSigmaSchema(BaseSchema): ``PixArtSigmaModel``. """ + checkpoint: schema_field( + enum_field(enum=["1024", "512"]), + placeholder="1024", + description=MultilingualString( + en=( + "Which PixArt-Sigma checkpoint to use: '1024' for best quality " + "at 1024x1024 px, or '512' for a faster, lighter model at " + "512x512 px. Both checkpoints are downloaded together." + ), + es=( + "Qué checkpoint de PixArt-Sigma usar: '1024' para mejor calidad " + "a 1024x1024 px, o '512' para un modelo más rápido y ligero a " + "512x512 px. Ambos checkpoints se descargan juntos." + ), + pt=( + "Qual checkpoint do PixArt-Sigma usar: '1024' para melhor " + "qualidade a 1024x1024 px, ou '512' para um modelo mais rápido " + "e leve a 512x512 px. Ambos os checkpoints são baixados juntos." + ), + de=( + "Welcher PixArt-Sigma-Checkpoint verwendet wird: '1024' für " + "beste Qualität bei 1024x1024 px oder '512' für ein schnelleres, " + "leichteres Modell bei 512x512 px. Beide Checkpoints werden " + "zusammen heruntergeladen." + ), + zh=( + "使用哪个 PixArt-Sigma 检查点：'1024' 表示 1024x1024 " + "像素的最佳质量，'512' 表示 512x512 像素更快更轻量的模型。" + "两个检查点会一起下载。" + ), + ), + alias=MultilingualString( + en="Checkpoint", + es="Checkpoint", + pt="Checkpoint", + de="Checkpoint", + zh="检查点", + ), + ) # type: ignore + negative_prompt: Optional[ schema_field( string_field(), @@ -321,9 +361,7 @@ class PixArtSigmaSchema(BaseSchema): ) # type: ignore -class PixArtSigmaGenerationModel( - HFPretrainedDownloadMixin, TextToImageGenerationTaskModel -): +class PixArtSigma(HFPretrainedDownloadMixin, TextToImageGenerationTaskModel): """Diffusion Transformer model for high efficiency text-to-image generation. Wraps the PixArt-Sigma pipeline, which replaces the U-Net backbone used @@ -345,7 +383,12 @@ class PixArtSigmaGenerationModel( """ SCHEMA = PixArtSigmaSchema - MODEL_NAME: str = "" + # The 1024 checkpoint is a full pipeline (T5, VAE, scheduler, tokenizer). + # The 512 checkpoint ships only a transformer, injected into this pipeline + # when the 512 variant is selected, so both repos are downloaded together. + MODEL_NAME: str = "PixArt-alpha/PixArt-Sigma-XL-2-1024-MS" + TRANSFORMER_512_REPO: str = "PixArt-alpha/PixArt-Sigma-XL-2-512-MS" + DOWNLOAD_SIZE_BYTES: int = 24280249290 COLOR: str = "#6a1b9a" DISPLAY_NAME: str = MultilingualString( en="PixArt-Sigma", @@ -456,6 +499,7 @@ def __init__(self, **kwargs): f"cuda:{DEVICE_TO_IDX.get(kwargs.get('device'))}" if use_gpu else "cpu" ) + self.checkpoint = kwargs.get("checkpoint") self.model = self._build_pipeline(use_gpu).to(self.device) self.negative_prompt = kwargs.get("negative_prompt") @@ -466,12 +510,26 @@ def __init__(self, **kwargs): self.height = kwargs.get("height") self.num_images_per_prompt = kwargs.get("num_images_per_prompt") + @classmethod + def hf_repos(cls): + """Download both the 1024 (full pipeline) and 512 (transformer) repos. + + Returns + ------- + list of tuple of (str, str) + The 1024 checkpoint (T5, VAE, scheduler, tokenizer, transformer) + and the 512 checkpoint (transformer only), so either variant can be + used after a single download. + """ + return [(cls.MODEL_NAME, "model"), (cls.TRANSFORMER_512_REPO, "model")] + def _build_pipeline(self, use_gpu: bool): - """Load the PixArt-Sigma pipeline from the downloaded checkpoint. + """Load the PixArt-Sigma pipeline for the selected checkpoint. - The default loads a self-contained checkpoint (one that ships a full - ``model_index.json`` pipeline, e.g. the 1024px variant). Subclasses - whose checkpoint contains only the transformer override this. + The pipeline scaffold (T5, VAE, scheduler, tokenizer) always comes from + the 1024 repo. For the ``"1024"`` checkpoint its own transformer is + used; for ``"512"`` the transformer from the 512 repo is injected (the + 512 repo has no ``model_index.json`` and cannot be loaded on its own). Parameters ---------- @@ -486,11 +544,22 @@ def _build_pipeline(self, use_gpu: bool): import torch from diffusers import PixArtSigmaPipeline - self.model_name = self._pretrained_source(None) - return PixArtSigmaPipeline.from_pretrained( - self.model_name, - torch_dtype=torch.float16 if use_gpu else torch.float32, - ) + dtype = torch.float16 if use_gpu else torch.float32 + pipeline_dir = str(self._repo_dir(self.MODEL_NAME)) + self.model_name = pipeline_dir + + if self.checkpoint == "512": + from diffusers import Transformer2DModel + + transformer_dir = str(self._repo_dir(self.TRANSFORMER_512_REPO)) + transformer = Transformer2DModel.from_pretrained( + transformer_dir, subfolder="transformer", torch_dtype=dtype + ) + return PixArtSigmaPipeline.from_pretrained( + pipeline_dir, transformer=transformer, torch_dtype=dtype + ) + + return PixArtSigmaPipeline.from_pretrained(pipeline_dir, torch_dtype=dtype) def generate(self, input: str) -> List[Any]: """Generate images from a text prompt. @@ -522,162 +591,3 @@ def generate(self, input: str) -> List[Any]: num_images_per_prompt=self.num_images_per_prompt, ) return output.images - - -class PixArtSigma1024(PixArtSigmaGenerationModel): - """PixArt-Sigma XL 1024px checkpoint. - - Downloads its checkpoint into the component's own download folder. - """ - - MODEL_NAME: str = "PixArt-alpha/PixArt-Sigma-XL-2-1024-MS" - DOWNLOAD_SIZE_BYTES: int = 21832490389 - DISPLAY_NAME = MultilingualString( - en="PixArt-Sigma 1024", - es="PixArt-Sigma 1024", - pt="PixArt-Sigma 1024", - de="PixArt-Sigma 1024", - zh="PixArt-Sigma 1024", - ) - DESCRIPTION = MultilingualString( - en=( - "PixArt-Sigma XL by PixArt-alpha, a diffusion transformer (DiT) " - "text-to-image model that reaches quality comparable to larger diffusion " - "models with far fewer parameters. This checkpoint generates at " - "1024x1024 px. Weights are downloaded into the component's own folder. " - "Model page: https://huggingface.co/PixArt-alpha/PixArt-Sigma-XL-2-1024-M" - "S" - ), - es=( - "PixArt-Sigma XL de PixArt-alpha, un modelo de texto a imagen basado en " - "transformer de difusión (DiT) que alcanza una calidad comparable a " - "modelos de difusión más grandes con muchos menos parámetros. Este " - "checkpoint genera a 1024x1024 px. Los pesos se descargan en la carpeta " - "propia del componente. Página del modelo: " - "https://huggingface.co/PixArt-alpha/PixArt-Sigma-XL-2-1024-MS" - ), - pt=( - "PixArt-Sigma XL de PixArt-alpha, um modelo de texto para imagem baseado " - "em transformer de difusão (DiT) que atinge qualidade comparável a " - "modelos de difusão maiores com muito menos parâmetros. Este " - "checkpoint gera a 1024x1024 px. Os pesos são baixados na pasta " - "própria do componente. Página do modelo: " - "https://huggingface.co/PixArt-alpha/PixArt-Sigma-XL-2-1024-MS" - ), - de=( - "PixArt-Sigma XL von PixArt-alpha, ein Text-zu-Bild-Modell auf Basis " - "eines Diffusion-Transformers (DiT), das mit weit weniger Parametern " - "eine Qualität vergleichbar mit größeren Diffusionsmodellen erreicht. " - "Dieser Checkpoint erzeugt Bilder mit 1024x1024 px. Die Gewichte werden " - "in den eigenen Ordner der Komponente heruntergeladen. Modellseite: " - "https://huggingface.co/PixArt-alpha/PixArt-Sigma-XL-2-1024-MS" - ), - zh=( - "PixArt-alpha 推出的 PixArt-Sigma XL，是一种基于扩散 " - "Transformer（DiT）的文本到图像模型，以远更少的参数量达到可媲美更大扩散模" - "型的质量。该检查点以 1024x1024 " - "像素生成。权重会下载到该组件自己的文件夹中。 模型页面： " - "https://huggingface.co/PixArt-alpha/PixArt-Sigma-XL-2-1024-MS" - ), - ) - - -class PixArtSigma512(PixArtSigmaGenerationModel): - """PixArt-Sigma XL 512px checkpoint (faster). - - Downloads its checkpoint into the component's own download folder. - """ - - MODEL_NAME: str = "PixArt-alpha/PixArt-Sigma-XL-2-512-MS" - # The 512 checkpoint ships only the transformer; the T5 text encoder, VAE, - # scheduler and tokenizer are loaded from the 1024 checkpoint, so both repos - # are downloaded. Size is the sum of the two. - PIPELINE_REPO: str = "PixArt-alpha/PixArt-Sigma-XL-2-1024-MS" - DOWNLOAD_SIZE_BYTES: int = 24280249290 - DISPLAY_NAME = MultilingualString( - en="PixArt-Sigma 512", - es="PixArt-Sigma 512", - pt="PixArt-Sigma 512", - de="PixArt-Sigma 512", - zh="PixArt-Sigma 512", - ) - DESCRIPTION = MultilingualString( - en=( - "PixArt-Sigma XL by PixArt-alpha at 512x512 px, a diffusion transformer " - "(DiT) text-to-image model. The lower resolution makes it faster and " - "lighter than the 1024 px variant. Weights are downloaded into the " - "component's own folder. Model page: https://huggingface.co/PixArt-alpha/" - "PixArt-Sigma-XL-2-512-MS" - ), - es=( - "PixArt-Sigma XL de PixArt-alpha a 512x512 px, un modelo de texto a " - "imagen basado en transformer de difusión (DiT). La menor resolución " - "lo hace más rápido y ligero que la variante de 1024 px. Los pesos se " - "descargan en la carpeta propia del componente. Página del modelo: " - "https://huggingface.co/PixArt-alpha/PixArt-Sigma-XL-2-512-MS" - ), - pt=( - "PixArt-Sigma XL de PixArt-alpha a 512x512 px, um modelo de texto para " - "imagem baseado em transformer de difusão (DiT). A menor resolução o " - "torna mais rápido e leve que a variante de 1024 px. Os pesos são " - "baixados na pasta própria do componente. Página do modelo: " - "https://huggingface.co/PixArt-alpha/PixArt-Sigma-XL-2-512-MS" - ), - de=( - "PixArt-Sigma XL von PixArt-alpha bei 512x512 px, ein " - "Text-zu-Bild-Modell auf Basis eines Diffusion-Transformers (DiT). Die " - "geringere Auflösung macht es schneller und leichter als die " - "1024-px-Variante. Die Gewichte werden in den eigenen Ordner der " - "Komponente heruntergeladen. Modellseite: " - "https://huggingface.co/PixArt-alpha/PixArt-Sigma-XL-2-512-MS" - ), - zh=( - "PixArt-alpha 推出的 PixArt-Sigma XL，分辨率为 512x512 " - "像素，是一种基于扩散 Transformer（DiT）的文本到图像模型。较低的分辨率使" - "其比 1024 像素变体更快、更轻量。权重会下载到该组件自己的文件夹中。 " - "模型页面： https://huggingface.co/PixArt-alpha/PixArt-Sigma-XL-2-512-MS" - ), - ) - - @classmethod - def hf_repos(cls): - """Download the 512 transformer repo and the 1024 pipeline repo. - - Returns - ------- - list of tuple of (str, str) - The 512 checkpoint (transformer only) and the 1024 checkpoint - (full pipeline: T5, VAE, scheduler, tokenizer). - """ - return [(cls.MODEL_NAME, "model"), (cls.PIPELINE_REPO, "model")] - - def _build_pipeline(self, use_gpu: bool): - """Load the 512 transformer into the 1024 pipeline scaffold. - - The 512 repo has no ``model_index.json`` (it ships only the - transformer), so the pipeline is loaded from the 1024 repo with the - 512 transformer injected. - - Parameters - ---------- - use_gpu : bool - Whether a GPU is available (selects float16 vs float32). - - Returns - ------- - diffusers.PixArtSigmaPipeline - The loaded pipeline (not yet moved to a device). - """ - import torch - from diffusers import PixArtSigmaPipeline, Transformer2DModel - - dtype = torch.float16 if use_gpu else torch.float32 - transformer_dir = str(self._repo_dir(self.MODEL_NAME)) - pipeline_dir = str(self._repo_dir(self.PIPELINE_REPO)) - self.model_name = pipeline_dir - transformer = Transformer2DModel.from_pretrained( - transformer_dir, subfolder="transformer", torch_dtype=dtype - ) - return PixArtSigmaPipeline.from_pretrained( - pipeline_dir, transformer=transformer, torch_dtype=dtype - ) diff --git a/tests/back/models/test_pixart_tongyi_downloadable.py b/tests/back/models/test_pixart_tongyi_downloadable.py index 1a646268a..982b9d7a7 100644 --- a/tests/back/models/test_pixart_tongyi_downloadable.py +++ b/tests/back/models/test_pixart_tongyi_downloadable.py @@ -3,33 +3,41 @@ import pytest from DashAI.back.dependencies.downloads.downloadable import HFPretrainedDownloadMixin -from DashAI.back.models.hugging_face.pixart_sigma_model import ( - PixArtSigma512, - PixArtSigma1024, -) +from DashAI.back.models.hugging_face.pixart_sigma_model import PixArtSigma from DashAI.back.models.hugging_face.tongyi_z_image_model import ( TongyiZImage, TongyiZImageTurbo, ) -_CASES = [ - (PixArtSigma1024, "PixArt-alpha/PixArt-Sigma-XL-2-1024-MS"), - (PixArtSigma512, "PixArt-alpha/PixArt-Sigma-XL-2-512-MS"), +# Single-repo checkpoints: hf_repos() is exactly one (repo_id, "model"). +_SINGLE_REPO_CASES = [ (TongyiZImage, "Tongyi-MAI/Z-Image"), (TongyiZImageTurbo, "Tongyi-MAI/Z-Image-Turbo"), ] +_ALL = [PixArtSigma, TongyiZImage, TongyiZImageTurbo] + -@pytest.mark.parametrize(("model_cls", "repo_id"), _CASES) -def test_is_downloadable(model_cls, repo_id): +@pytest.mark.parametrize(("model_cls", "repo_id"), _SINGLE_REPO_CASES) +def test_single_repo_is_downloadable(model_cls, repo_id): assert issubclass(model_cls, HFPretrainedDownloadMixin) assert model_cls.REQUIRES_DOWNLOAD is True assert model_cls.DOWNLOAD_SIZE_BYTES is not None assert model_cls.hf_repos() == [(repo_id, "model")] -@pytest.mark.parametrize(("model_cls", "repo_id"), _CASES) -def test_metadata_flags_download(model_cls, repo_id): +def test_pixart_sigma_downloads_both_checkpoints(): + """PixArt-Sigma downloads the 1024 pipeline and the 512 transformer.""" + assert issubclass(PixArtSigma, HFPretrainedDownloadMixin) + assert PixArtSigma.REQUIRES_DOWNLOAD is True + assert PixArtSigma.DOWNLOAD_SIZE_BYTES is not None + repos = [repo_id for repo_id, *_ in PixArtSigma.hf_repos()] + assert "PixArt-alpha/PixArt-Sigma-XL-2-1024-MS" in repos + assert "PixArt-alpha/PixArt-Sigma-XL-2-512-MS" in repos + + +@pytest.mark.parametrize("model_cls", _ALL) +def test_metadata_flags_download(model_cls): meta = model_cls.get_metadata() assert meta["requires_download"] is True assert meta["download_size_bytes"] == model_cls.DOWNLOAD_SIZE_BYTES From 03275860e71e443cda69f53c94c8e2a3b1f91fed Mon Sep 17 00:00:00 2001 From: Irozuku Date: Thu, 9 Jul 2026 14:57:16 -0400 Subject: [PATCH 83/89] feat: confirm before deleting a component download Deleting a downloaded component from the selector button or the models side bar icon now opens a confirmation modal instead of deleting immediately, reusing DeleteConfirmationModal with a per-component message (all five locales). --- .../models/model/ComponentDownloadControl.jsx | 31 +++++++++++---- .../model/ComponentDownloadControl.test.jsx | 8 +++- .../models/model/ModelDownloadStatusIcon.jsx | 39 +++++++++++++------ .../model/ModelDownloadStatusIcon.test.jsx | 8 +++- .../src/utils/i18n/locales/de/common.json | 3 +- .../src/utils/i18n/locales/en/common.json | 3 +- .../src/utils/i18n/locales/es/common.json | 3 +- .../src/utils/i18n/locales/pt/common.json | 3 +- .../src/utils/i18n/locales/zh/common.json | 3 +- 9 files changed, 73 insertions(+), 28 deletions(-) diff --git a/DashAI/front/src/components/models/model/ComponentDownloadControl.jsx b/DashAI/front/src/components/models/model/ComponentDownloadControl.jsx index 414e1422a..4b7fa91c6 100644 --- a/DashAI/front/src/components/models/model/ComponentDownloadControl.jsx +++ b/DashAI/front/src/components/models/model/ComponentDownloadControl.jsx @@ -10,6 +10,7 @@ import { getComponentDownloadStatus, } from "../../../api/component"; import { startJobPolling, stopJobPolling } from "../../../utils/jobPoller"; +import DeleteConfirmationModal from "../../threeSectionLayout/DeleteConfirmationModal"; const formatSize = (bytes) => { if (bytes == null) return ""; @@ -174,6 +175,7 @@ const ComponentDownloadControl = ({ component, onStatusChange }) => { const { enqueueSnackbar } = useSnackbar(); const meta = component.metadata || {}; const { downloaded, downloading } = useComponentDownloadState(component); + const [confirmOpen, setConfirmOpen] = useState(false); if (!meta.requires_download) return null; @@ -196,14 +198,27 @@ const ComponentDownloadControl = ({ component, onStatusChange }) => { if (downloaded) { return ( - + <> + + setConfirmOpen(false)} + onConfirm={() => { + setConfirmOpen(false); + handleDelete(); + }} + content={t("common:componentDownload.confirmDelete", { + name: component.display_name || component.name, + })} + /> + ); } diff --git a/DashAI/front/src/components/models/model/ComponentDownloadControl.test.jsx b/DashAI/front/src/components/models/model/ComponentDownloadControl.test.jsx index e8388119f..645c7f585 100644 --- a/DashAI/front/src/components/models/model/ComponentDownloadControl.test.jsx +++ b/DashAI/front/src/components/models/model/ComponentDownloadControl.test.jsx @@ -56,8 +56,14 @@ describe("ComponentDownloadControl", () => { onStatusChange={() => {}} />, ); - const button = await screen.findByRole("button", { name: /delete/i }); + const button = await screen.findByRole("button", { + name: "Delete download", + }); fireEvent.click(button); + // Deletion only happens after confirming in the modal. + expect(deleteComponentDownload).not.toHaveBeenCalled(); + const confirm = await screen.findByRole("button", { name: "Delete" }); + fireEvent.click(confirm); await waitFor(() => expect(deleteComponentDownload).toHaveBeenCalledWith( "OpusMtEnRoaTransformerDownloaded", diff --git a/DashAI/front/src/components/models/model/ModelDownloadStatusIcon.jsx b/DashAI/front/src/components/models/model/ModelDownloadStatusIcon.jsx index 07eea8edb..975c57360 100644 --- a/DashAI/front/src/components/models/model/ModelDownloadStatusIcon.jsx +++ b/DashAI/front/src/components/models/model/ModelDownloadStatusIcon.jsx @@ -1,4 +1,4 @@ -import React from "react"; +import React, { useState } from "react"; import PropTypes from "prop-types"; import { Box, CircularProgress, IconButton, Tooltip } from "@mui/material"; import DownloadIcon from "@mui/icons-material/Download"; @@ -9,6 +9,7 @@ import { useComponentDownloadState, deleteComponent, } from "./ComponentDownloadControl"; +import DeleteConfirmationModal from "../../threeSectionLayout/DeleteConfirmationModal"; /** * Compact, tooltip-free download status shown at the end of a model row in the @@ -22,6 +23,7 @@ export default function ModelDownloadStatusIcon({ model, onChanged }) { const { t } = useTranslation(["common"]); const { enqueueSnackbar } = useSnackbar(); const { downloaded, downloading } = useComponentDownloadState(model); + const [confirmOpen, setConfirmOpen] = useState(false); if (!model.metadata?.requires_download) return null; @@ -31,13 +33,25 @@ export default function ModelDownloadStatusIcon({ model, onChanged }) { if (downloaded) { return ( - - { - e.stopPropagation(); + <> + + { + e.stopPropagation(); + setConfirmOpen(true); + }} + > + + + + setConfirmOpen(false)} + onConfirm={() => { + setConfirmOpen(false); deleteComponent({ component: model, enqueueSnackbar, @@ -45,10 +59,11 @@ export default function ModelDownloadStatusIcon({ model, onChanged }) { onStatusChange: onChanged, }); }} - > - - - + content={t("common:componentDownload.confirmDelete", { + name: model.display_name || model.name, + })} + /> + ); } diff --git a/DashAI/front/src/components/models/model/ModelDownloadStatusIcon.test.jsx b/DashAI/front/src/components/models/model/ModelDownloadStatusIcon.test.jsx index c80b99a26..6c17a3b16 100644 --- a/DashAI/front/src/components/models/model/ModelDownloadStatusIcon.test.jsx +++ b/DashAI/front/src/components/models/model/ModelDownloadStatusIcon.test.jsx @@ -32,15 +32,19 @@ describe("ModelDownloadStatusIcon", () => { expect(screen.queryByRole("button")).toBeNull(); }); - it("deletes a downloaded model from a plain delete icon", async () => { + it("deletes a downloaded model after confirming in the modal", async () => { renderWithProviders( {}} />, ); - const del = await screen.findByRole("button", { name: /delete/i }); + const del = await screen.findByRole("button", { name: "Delete download" }); fireEvent.click(del); + // Deletion only happens after confirming in the modal. + expect(deleteComponentDownload).not.toHaveBeenCalled(); + const confirm = await screen.findByRole("button", { name: "Delete" }); + fireEvent.click(confirm); await waitFor(() => expect(deleteComponentDownload).toHaveBeenCalledWith( "DownloadableTestModel", diff --git a/DashAI/front/src/utils/i18n/locales/de/common.json b/DashAI/front/src/utils/i18n/locales/de/common.json index 06964acb6..343e7fbda 100644 --- a/DashAI/front/src/utils/i18n/locales/de/common.json +++ b/DashAI/front/src/utils/i18n/locales/de/common.json @@ -168,7 +168,8 @@ "deleted": "Download geloescht", "failed": "Download der Komponente fehlgeschlagen", "mustDownload": "Dieses Modell muss vor der Nutzung heruntergeladen werden", - "mustDownloadNested": "Diese Komponenten müssen zuerst heruntergeladen werden: {{names}}" + "mustDownloadNested": "Diese Komponenten müssen zuerst heruntergeladen werden: {{names}}", + "confirmDelete": "Die heruntergeladenen Dateien für {{name}} löschen? Du kannst es später erneut herunterladen." }, "jobQueue": { "title": "Aufgabenwarteschlange", diff --git a/DashAI/front/src/utils/i18n/locales/en/common.json b/DashAI/front/src/utils/i18n/locales/en/common.json index 2b37fd4b4..a636c54fa 100644 --- a/DashAI/front/src/utils/i18n/locales/en/common.json +++ b/DashAI/front/src/utils/i18n/locales/en/common.json @@ -168,7 +168,8 @@ "deleted": "Download deleted", "failed": "Component download failed", "mustDownload": "This model must be downloaded before use", - "mustDownloadNested": "These components must be downloaded first: {{names}}" + "mustDownloadNested": "These components must be downloaded first: {{names}}", + "confirmDelete": "Delete the downloaded files for {{name}}? You can download it again later." }, "jobQueue": { "title": "Job Queue", diff --git a/DashAI/front/src/utils/i18n/locales/es/common.json b/DashAI/front/src/utils/i18n/locales/es/common.json index 33ea478e0..5fc1775e8 100644 --- a/DashAI/front/src/utils/i18n/locales/es/common.json +++ b/DashAI/front/src/utils/i18n/locales/es/common.json @@ -168,7 +168,8 @@ "deleted": "Descarga eliminada", "failed": "La descarga del componente ha fallado", "mustDownload": "Este modelo debe descargarse antes de usarlo", - "mustDownloadNested": "Estos componentes deben descargarse primero: {{names}}" + "mustDownloadNested": "Estos componentes deben descargarse primero: {{names}}", + "confirmDelete": "¿Eliminar los archivos descargados de {{name}}? Podrás descargarlo de nuevo más tarde." }, "jobQueue": { "title": "Cola de trabajos", diff --git a/DashAI/front/src/utils/i18n/locales/pt/common.json b/DashAI/front/src/utils/i18n/locales/pt/common.json index 827390169..bde408bde 100644 --- a/DashAI/front/src/utils/i18n/locales/pt/common.json +++ b/DashAI/front/src/utils/i18n/locales/pt/common.json @@ -168,7 +168,8 @@ "deleted": "Download removido", "failed": "Falha ao baixar o componente", "mustDownload": "Este modelo precisa ser baixado antes de usar", - "mustDownloadNested": "Estes componentes precisam ser baixados primeiro: {{names}}" + "mustDownloadNested": "Estes componentes precisam ser baixados primeiro: {{names}}", + "confirmDelete": "Excluir os arquivos baixados de {{name}}? Você poderá baixá-lo novamente mais tarde." }, "jobQueue": { "title": "Fila de trabalhos", diff --git a/DashAI/front/src/utils/i18n/locales/zh/common.json b/DashAI/front/src/utils/i18n/locales/zh/common.json index 14b8fcba3..6f42fda97 100644 --- a/DashAI/front/src/utils/i18n/locales/zh/common.json +++ b/DashAI/front/src/utils/i18n/locales/zh/common.json @@ -168,7 +168,8 @@ "deleted": "下载已删除", "failed": "组件下载失败", "mustDownload": "使用前必须先下载此模型", - "mustDownloadNested": "必须先下载这些组件：{{names}}" + "mustDownloadNested": "必须先下载这些组件：{{names}}", + "confirmDelete": "删除 {{name}} 的已下载文件？之后可以重新下载。" }, "jobQueue": { "title": "任务队列", From 6199d2bb48972c628c110ad5561f5a7c31d8b0fa Mon Sep 17 00:00:00 2001 From: Irozuku Date: Thu, 9 Jul 2026 16:21:48 -0400 Subject: [PATCH 84/89] fix: stop delete confirmation clicks from triggering the row underneath React bubbles events through the component tree even for portaled modals, so confirming a component deletion from within a clickable model row also fired the row's onClick and opened the configure dialog. Stop propagation on the modal content. --- .../components/threeSectionLayout/DeleteConfirmationModal.jsx | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/DashAI/front/src/components/threeSectionLayout/DeleteConfirmationModal.jsx b/DashAI/front/src/components/threeSectionLayout/DeleteConfirmationModal.jsx index 688f73345..93210664b 100644 --- a/DashAI/front/src/components/threeSectionLayout/DeleteConfirmationModal.jsx +++ b/DashAI/front/src/components/threeSectionLayout/DeleteConfirmationModal.jsx @@ -26,6 +26,10 @@ export default function DeleteConfirmationModal({ slotProps={{ transition: { onExited } }} > e.stopPropagation()} sx={{ position: "absolute", top: "50%", From 47da57fd3961302fcc300ce787e24ec343854dfa Mon Sep 17 00:00:00 2001 From: Irozuku Date: Thu, 9 Jul 2026 17:06:52 -0400 Subject: [PATCH 85/89] feat: show download size on the delete button The delete button in the component download control now shows the download size (Delete download ({{size}})), keeping the size visible after a component is downloaded. Added deleteWithSize to all five locales. --- .../components/models/model/ComponentDownloadControl.jsx | 6 +++++- .../models/model/ComponentDownloadControl.test.jsx | 2 +- DashAI/front/src/utils/i18n/locales/de/common.json | 1 + DashAI/front/src/utils/i18n/locales/en/common.json | 1 + DashAI/front/src/utils/i18n/locales/es/common.json | 1 + DashAI/front/src/utils/i18n/locales/pt/common.json | 1 + DashAI/front/src/utils/i18n/locales/zh/common.json | 1 + 7 files changed, 11 insertions(+), 2 deletions(-) diff --git a/DashAI/front/src/components/models/model/ComponentDownloadControl.jsx b/DashAI/front/src/components/models/model/ComponentDownloadControl.jsx index 4b7fa91c6..6d9c44cff 100644 --- a/DashAI/front/src/components/models/model/ComponentDownloadControl.jsx +++ b/DashAI/front/src/components/models/model/ComponentDownloadControl.jsx @@ -205,7 +205,11 @@ const ComponentDownloadControl = ({ component, onStatusChange }) => { startIcon={} onClick={() => setConfirmOpen(true)} > - {t("common:componentDownload.delete")} + {meta.download_size_bytes != null + ? t("common:componentDownload.deleteWithSize", { + size: formatSize(meta.download_size_bytes), + }) + : t("common:componentDownload.delete")} { />, ); const button = await screen.findByRole("button", { - name: "Delete download", + name: /delete download/i, }); fireEvent.click(button); // Deletion only happens after confirming in the modal. diff --git a/DashAI/front/src/utils/i18n/locales/de/common.json b/DashAI/front/src/utils/i18n/locales/de/common.json index 343e7fbda..190b20453 100644 --- a/DashAI/front/src/utils/i18n/locales/de/common.json +++ b/DashAI/front/src/utils/i18n/locales/de/common.json @@ -163,6 +163,7 @@ "componentDownload": { "download": "Herunterladen ({{size}})", "delete": "Download loeschen", + "deleteWithSize": "Download loeschen ({{size}})", "downloading": "Wird heruntergeladen...", "done": "Komponente heruntergeladen", "deleted": "Download geloescht", diff --git a/DashAI/front/src/utils/i18n/locales/en/common.json b/DashAI/front/src/utils/i18n/locales/en/common.json index a636c54fa..333032630 100644 --- a/DashAI/front/src/utils/i18n/locales/en/common.json +++ b/DashAI/front/src/utils/i18n/locales/en/common.json @@ -163,6 +163,7 @@ "componentDownload": { "download": "Download ({{size}})", "delete": "Delete download", + "deleteWithSize": "Delete download ({{size}})", "downloading": "Downloading...", "done": "Component downloaded", "deleted": "Download deleted", diff --git a/DashAI/front/src/utils/i18n/locales/es/common.json b/DashAI/front/src/utils/i18n/locales/es/common.json index 5fc1775e8..e4d787aac 100644 --- a/DashAI/front/src/utils/i18n/locales/es/common.json +++ b/DashAI/front/src/utils/i18n/locales/es/common.json @@ -163,6 +163,7 @@ "componentDownload": { "download": "Descargar ({{size}})", "delete": "Eliminar descarga", + "deleteWithSize": "Eliminar descarga ({{size}})", "downloading": "Descargando...", "done": "Componente descargado", "deleted": "Descarga eliminada", diff --git a/DashAI/front/src/utils/i18n/locales/pt/common.json b/DashAI/front/src/utils/i18n/locales/pt/common.json index bde408bde..6f9eabf1c 100644 --- a/DashAI/front/src/utils/i18n/locales/pt/common.json +++ b/DashAI/front/src/utils/i18n/locales/pt/common.json @@ -163,6 +163,7 @@ "componentDownload": { "download": "Baixar ({{size}})", "delete": "Remover download", + "deleteWithSize": "Remover download ({{size}})", "downloading": "Baixando...", "done": "Componente baixado", "deleted": "Download removido", diff --git a/DashAI/front/src/utils/i18n/locales/zh/common.json b/DashAI/front/src/utils/i18n/locales/zh/common.json index 6f42fda97..cdc49b6c0 100644 --- a/DashAI/front/src/utils/i18n/locales/zh/common.json +++ b/DashAI/front/src/utils/i18n/locales/zh/common.json @@ -163,6 +163,7 @@ "componentDownload": { "download": "下载 ({{size}})", "delete": "删除下载", + "deleteWithSize": "删除下载（{{size}}）", "downloading": "下载中...", "done": "组件已下载", "deleted": "下载已删除", From 1918286d309577a16288e40caee6da90d28fb41d Mon Sep 17 00:00:00 2001 From: Irozuku Date: Fri, 10 Jul 2026 09:38:24 -0400 Subject: [PATCH 86/89] fix: close model hover popover when clicking a row action Clicking the delete icon opened the confirmation modal without the mouse leaving the row, so the hover popover stayed anchored and visible after deletion. Clear the hover on any row click via a capture phase handler, which fires even for the action icon that stops propagation. --- DashAI/front/src/components/models/model/ModelListItem.jsx | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/DashAI/front/src/components/models/model/ModelListItem.jsx b/DashAI/front/src/components/models/model/ModelListItem.jsx index d04e47212..252e84650 100644 --- a/DashAI/front/src/components/models/model/ModelListItem.jsx +++ b/DashAI/front/src/components/models/model/ModelListItem.jsx @@ -61,6 +61,11 @@ export default function ModelListItem({ } onMouseEnter={(e) => handleMouseEnter(e, model)} onMouseLeave={handleMouseLeave} + // Close the hover popover on any click in the row (capture phase, so it + // runs even for the action's icon which stops propagation). Otherwise a + // click that opens a modal (e.g. delete confirmation) leaves the popover + // stuck open since no mouseleave fires. + onClickCapture={handleMouseLeave} onClick={handleCardClick} {...props} sx={{ From d5bd97ab335072c85a4d6e5899fc12e22e2f4af9 Mon Sep 17 00:00:00 2001 From: Irozuku Date: Fri, 10 Jul 2026 09:38:30 -0400 Subject: [PATCH 87/89] fix: disable train in the comparison table when the model is not downloaded The per row train button in the model comparison table ignored download status. Disable it (with a tooltip) when the run's model still needs downloading, using the shared live download state so it enables once the download finishes. --- .../models/ModelComparisonTable.jsx | 33 +++++++++++++++++-- 1 file changed, 31 insertions(+), 2 deletions(-) diff --git a/DashAI/front/src/components/models/ModelComparisonTable.jsx b/DashAI/front/src/components/models/ModelComparisonTable.jsx index 7d8eed469..20233f935 100644 --- a/DashAI/front/src/components/models/ModelComparisonTable.jsx +++ b/DashAI/front/src/components/models/ModelComparisonTable.jsx @@ -19,6 +19,10 @@ import { useTranslation } from "react-i18next"; import { useTableLocalization } from "../../utils/useTableLocalization"; import api from "../../api/api"; import DeleteConfirmationModal from "../threeSectionLayout/DeleteConfirmationModal"; +import { + getComponentDownloadState, + subscribeAnyDownloadState, +} from "./model/ComponentDownloadControl"; /** * Compact comparison table showing all runs in a session. @@ -43,6 +47,24 @@ function ModelComparisonTable({ const [loadingScores, setLoadingScores] = useState(false); const [runs, setRuns] = useState(initialRuns); const [runToDelete, setRunToDelete] = useState(null); + // Bump to re-render when a download finishes so the train button enables. + const [, setDownloadVersion] = useState(0); + + useEffect( + () => subscribeAnyDownloadState(() => setDownloadVersion((v) => v + 1)), + [], + ); + + // A run is trainable only if its model needs no download or the download is + // present and not in progress (live state overrides a stale fetched flag). + const isModelReady = (modelName) => { + const model = models.find((m) => m.name === modelName); + if (!model?.metadata?.requires_download) return true; + const cached = getComponentDownloadState(modelName); + const downloaded = cached?.downloaded ?? Boolean(model.downloaded); + const downloading = Boolean(cached?.downloading); + return downloaded && !downloading; + }; const { t, i18n } = useTranslation(["models", "common"]); const theme = useTheme(); @@ -420,10 +442,17 @@ function ModelComparisonTable({ row.original.status === 3; const isRunning = row.original.status === 1 || row.original.status === 2; + const modelReady = isModelReady(row.original.model_name); return ( - + r.id === row.original.id)); }} - disabled={!canTrain} + disabled={!canTrain || !modelReady} color="primary" > From 1d1db7a60636814416ce89f8eb295d8df5079431 Mon Sep 17 00:00:00 2001 From: Irozuku Date: Fri, 10 Jul 2026 10:44:45 -0400 Subject: [PATCH 88/89] fix: improve hover feedback for not downloaded models in ModelListItem --- DashAI/front/src/components/models/model/ModelListItem.jsx | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/DashAI/front/src/components/models/model/ModelListItem.jsx b/DashAI/front/src/components/models/model/ModelListItem.jsx index 252e84650..ba774af61 100644 --- a/DashAI/front/src/components/models/model/ModelListItem.jsx +++ b/DashAI/front/src/components/models/model/ModelListItem.jsx @@ -83,8 +83,11 @@ export default function ModelListItem({ bgcolor: disabled ? theme.palette.ui.disabled : theme.palette.action.hover, - borderColor: disabled ? theme.palette.ui.border : color, - transform: disabled || !isClickable ? "none" : "translateX(4px)", + // A not downloaded row is disabled but still clickable to start the + // download, so give it the same hover feedback (border highlight + + // slide) instead of feeling stiff. + borderColor: color, + transform: !isClickable ? "none" : "translateX(4px)", }, "&::after": disabled ? { From 0f50088b8b5fdb7f92bc8bc8b441086eed8a2d08 Mon Sep 17 00:00:00 2001 From: Irozuku Date: Fri, 10 Jul 2026 12:45:54 -0400 Subject: [PATCH 89/89] test: fix download, opus, and job fixtures broken by download changes Update the snapshot_download assertions for the new ignore_patterns, add save_pretrained to the opus DummyTokenizer (the base save persists the tokenizer now), and make the jobs test registry module scoped so DummyModel is registered before the module scoped run fixtures create runs. --- tests/back/api/test_jobs.py | 33 +++++++++++-------- tests/back/downloads/test_downloadable.py | 2 ++ .../models/test_opus_mt_en_es_transformer.py | 3 ++ .../models/test_opus_mt_es_en_transformer.py | 3 ++ 4 files changed, 28 insertions(+), 13 deletions(-) diff --git a/tests/back/api/test_jobs.py b/tests/back/api/test_jobs.py index 9ac229fc2..31ef9f53c 100644 --- a/tests/back/api/test_jobs.py +++ b/tests/back/api/test_jobs.py @@ -78,10 +78,18 @@ def score(true_labels: list, probs_pred_labels: list): return 1 -@pytest.fixture(autouse=True, name="test_registry") -def setup_test_registry(client, monkeypatch: pytest.MonkeyPatch): - """Setup a test registry with test task, dataloader and model components.""" +@pytest.fixture(scope="module", name="test_registry", autouse=True) +def setup_test_registry(client): + """Setup a test registry with test task, dataloader and model components. + + Module-scoped (with a manual swap, since ``monkeypatch`` is function + scoped) so the test components are registered before the module-scoped + ``create_run`` / ``create_model_session`` fixtures run. + """ container = client.app.container + sentinel = object() + services = container._services + old = services.get("component_registry", sentinel) test_registry = ComponentRegistry( initial_components=[ @@ -94,13 +102,12 @@ def setup_test_registry(client, monkeypatch: pytest.MonkeyPatch): OptunaOptimizer, ] ) - - monkeypatch.setitem( - container._services, - "component_registry", - test_registry, - ) - return test_registry + services["component_registry"] = test_registry + yield test_registry + if old is sentinel: + del services["component_registry"] + else: + services["component_registry"] = old @pytest.fixture(scope="module", name="dataset_id") @@ -110,7 +117,7 @@ def dataset_id(dataset_1: Dataset) -> int: @pytest.fixture(scope="module", name="model_session_id", autouse=True) -def create_model_session(client: TestClient, dataset_id: int): +def create_model_session(client: TestClient, dataset_id: int, test_registry): container = client.app.container session = container["session_factory"] @@ -150,7 +157,7 @@ def create_model_session(client: TestClient, dataset_id: int): @pytest.fixture(scope="module", name="run_id", autouse=True) -def create_run(client: TestClient, model_session_id: int): +def create_run(client: TestClient, model_session_id: int, test_registry): response = client.post( "/api/v1/run/", json={ @@ -182,7 +189,7 @@ def create_run(client: TestClient, model_session_id: int): @pytest.fixture(scope="module", name="failed_run_id", autouse=True) -def create_failed_run(client: TestClient, model_session_id: int): +def create_failed_run(client: TestClient, model_session_id: int, test_registry): container = client.app.container session_factory = container["session_factory"] diff --git a/tests/back/downloads/test_downloadable.py b/tests/back/downloads/test_downloadable.py index 6013a276d..4597dc2a3 100644 --- a/tests/back/downloads/test_downloadable.py +++ b/tests/back/downloads/test_downloadable.py @@ -63,6 +63,7 @@ def test_download_fetches_into_component_dir(components_root): repo_id="owner/model-a", repo_type="model", local_dir=str(components_root / "_Dummy" / "model-a"), + ignore_patterns=list(_Dummy.HF_IGNORE_PATTERNS), ) assert calls[0] == (None, "Downloading owner/model-a") @@ -90,6 +91,7 @@ def test_download_3tuple_passes_allow_patterns(components_root): repo_type="model", local_dir=str(components_root / "_DummyPartial" / "model-a"), allow_patterns=["*8_0.gguf"], + ignore_patterns=list(_DummyPartial.HF_IGNORE_PATTERNS), ) diff --git a/tests/back/models/test_opus_mt_en_es_transformer.py b/tests/back/models/test_opus_mt_en_es_transformer.py index 9d3108c6b..a37ec6bba 100644 --- a/tests/back/models/test_opus_mt_en_es_transformer.py +++ b/tests/back/models/test_opus_mt_en_es_transformer.py @@ -12,6 +12,9 @@ def __call__(self, text, truncation=True, padding="max_length", max_length=512): del text, truncation, padding, max_length return {"input_ids": [1, 2, 3], "attention_mask": [1, 1, 1]} + def save_pretrained(self, save_directory): + Path(save_directory).mkdir(parents=True, exist_ok=True) + class DummySeq2SeqModel: def __init__(self): diff --git a/tests/back/models/test_opus_mt_es_en_transformer.py b/tests/back/models/test_opus_mt_es_en_transformer.py index 24f6a2de3..8850ebcd9 100644 --- a/tests/back/models/test_opus_mt_es_en_transformer.py +++ b/tests/back/models/test_opus_mt_es_en_transformer.py @@ -22,6 +22,9 @@ def decode(self, token_ids, skip_special_tokens=True): del token_ids, skip_special_tokens return "translated text" + def save_pretrained(self, save_directory): + Path(save_directory).mkdir(parents=True, exist_ok=True) + class DummySeq2SeqModel: def __init__(self):