Skip to content

Add model download management, job progress tracking, and generative model switching#764

Open
Irozuku wants to merge 90 commits into
developfrom
feat/component-download
Open

Add model download management, job progress tracking, and generative model switching#764
Irozuku wants to merge 90 commits into
developfrom
feat/component-download

Conversation

@Irozuku

@Irozuku Irozuku commented Jul 10, 2026

Copy link
Copy Markdown
Collaborator

Summary

Adds an explicit download system for every weight backed ML component (text classification, translation, image classification, GGUF LLMs, diffusion, ControlNet) and a per job progress bar in the job queue. Before this change these models silently pulled weights from HuggingFace on first use, with no visibility, no size information, and no way to manage them, and runs or sessions could start against weights that were not present. Now each component that requires a download is fetched on demand with progress, an accurate size, and a delete action, and runs, trainings, and generative sessions are blocked until the required weights are present, including components nested inside another component's parameters. The job queue also shows real progress (a determinate bar plus a phase message) for long running jobs instead of a static "running" label. Generative sessions can also change their model on the fly: the switcher lets you pick any compatible model (blocking input and offering a download when it is not present), records each model change and its parameters in the chat history, and restores the parameters last used for that model in the session or its defaults.

Component download in the model list and Component Selector (generative)

image image

Download gating on train and run

image

Per job progress bar in the job queue

image image image

Generative session with model selector

image image image

Type of Change

  • Backend change
  • Frontend change
  • CI / Workflow change
  • Build / Packaging change
  • Bug fix
  • Documentation

Changes (by file)

Download infrastructure (backend)

  • DashAI/back/dependencies/downloads/downloadable.py: DownloadableMixin plus HFDownloadableMixin, HFPretrainedDownloadMixin, and TorchvisionDownloadMixin. Downloads each repo into the component folder, uses allow_patterns and ignore_patterns to skip alternate framework weights (TensorFlow, Flax, Rust, ONNX, OpenVINO, CoreML), and forces the classic transfer path (Xet disabled).
  • DashAI/back/dependencies/downloads/nested.py: walks a parameters dict for components chosen as another component's parameter at any depth and reports which ones still need downloading.
  • DashAI/back/job/component_download_job.py, DashAI/back/job/base_job.py: the download job plus a report_progress helper.
  • DashAI/back/api/api_v1/endpoints/components.py: download, status, and delete endpoints, POST /components/downloads/required, and reconciliation of download status against the filesystem on read.
  • DashAI/back/api/api_v1/endpoints/runs.py, generative_session.py, job/model_job.py: gate run creation, session creation, and training on the model and its nested components being downloaded, returning 409 or JobError with the list of missing components.

Generative session model switching (backend)

  • DashAI/back/api/api_v1/endpoints/generative_session.py: allow changing a session's model, including to one that is not yet downloaded, record the change in the parameter history with the model name, emit a model change event, and restore the parameters last used for that model in the session or its defaults.
  • DashAI/back/dependencies/database/models.py and migration a7d2c9e4f1b0: add model_name to GenerativeSessionParameterHistory.

Models made downloadable (backend)

  • Text classification transformers, translation models (Opus MT including En Roa and Roa En, M2M100, NLLB, T5), torchvision image classifiers, GGUF LLMs (Llama, Mistral, Mixtral, Qwen, SmolLM, split per quantization or checkpoint), diffusion (SD2, SD3, SDXL, SDXL Turbo, PixArt Sigma, Tongyi), and ControlNet. Each declares REQUIRES_DOWNLOAD, DOWNLOAD_SIZE_BYTES, and its repos, and descriptions were corrected to state that weights must be downloaded before use.
  • DashAI/back/models/hugging_face/pixart_sigma_model.py: merged the 512 and 1024 variants into one PixArtSigma with a checkpoint selector, because the 512 repo ships only a transformer and reuses the 1024 pipeline.
  • DashAI/back/initial_components.py: registrations updated for the split and merged classes.

Per job progress bar

  • DashAI/back/dependencies/job_queues/huey_job_queue.py, base_job_queue.py, job/base_job.py: progress and progress_message columns on task_copy (with a migration for existing databases), a report_progress method, and pass through in /jobs/changes.
  • Instrumented ModelJob, DatasetJob, PredictJob, and ConverterJob with progress checkpoints.
  • DashAI/front/src/components/jobs/JobQueueWidget.jsx: a determinate or indeterminate bar plus a phase message for each job.

Download UI (frontend)

  • api/component.ts: download, delete, and status calls plus getRequiredDownloads.
  • components/models/model/ComponentDownloadControl.jsx: the download and delete control, with a module level pub/sub so every control for the same component stays in sync. ModelDownloadStatusIcon.jsx for the models side bar.
  • custom/ComponentSelector.jsx, ComponentDetailsPanel.jsx: inline download control and the size shown on the delete button.
  • shared/FormSchema* selectors and contexts/schema.js: inline download for nested components, a fix for selecting a component nested two levels deep, and a root breadcrumb.
  • models/ModelsRightBar.jsx, RunCard.jsx, ModelComparisonTable.jsx, SessionVisualization.jsx: disable Train, Retrain, and Run All until the model is downloaded (using live state), and update the list in place to avoid scroll resets.
  • generative/ModelSwitcher.jsx, GenerativeChat.jsx, CreateSession*: a session level switcher to change the model mid session, model change events shown in the chat history, and blocking of input and switching until the model is downloaded.
  • threeSectionLayout/DeleteConfirmationModal.jsx: delete confirmation that stops click propagation so confirming does not trigger the row underneath.
  • utils/jobPoller.js: stop reporting freshly queued jobs as failed.
  • i18n common.json and custom.json (en, es, pt, de, zh): keys for download, size, must download, and confirm delete.

Testing

  • Backend: uv run pytest tests/back/ covers the download mixin, the nested collector and /downloads/required, per model downloadable metadata, the run, session, and training gates, generative session model switching, and job progress reporting.
  • Frontend: yarn test covers the download control, the side bar icon, and the job queue progress bar.
  • Manual: download a model and watch the progress bar advance, try to train or use it before and after download, delete it through the confirmation,
  • Manual: switch a generative session's model and confirm the change appears in the chat history. Verifiy that the following models mantain the history of the conversation.

Notes

  • Download sizes reflect the full HuggingFace repo minus the alternate framework duplicates. Diffusion pipelines are still large because they keep both fp16 and fp32 safetensors, so an fp16 only pass is a possible follow up that would need per pipeline verification.
  • Registry names changed for split and merged models (for example PixArtSigma512 and PixArtSigma1024 became PixArtSigma, and the GGUF models split per quantization), so generative sessions that reference old names need to be recreated.
  • Downloads use the classic HTTP or LFS path with Xet disabled, after a 404 on the Xet read token endpoint for some repos.

Irozuku added 30 commits July 1, 2026 11:52
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.
…r 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.
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.
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.
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.
…nline 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.
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.
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.
… 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.
Irozuku added 23 commits July 3, 2026 15:02
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.
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.
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.
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.
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.
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.
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).
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).
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.
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.
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.
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.
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).
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.
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.
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.
…loaded

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.
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.
@Irozuku Irozuku added bug Something isn't working enhancement New feature or request front Frontend work back Backend work labels Jul 10, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

back Backend work bug Something isn't working enhancement New feature or request front Frontend work

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant