feat: add VAE download infrastructure with metadata registry #220
Conversation
There was a problem hiding this comment.
The latest main is using an updated process for downloading models where we keep track of the "artifacts" (currently only supporting an artifact of type HuggingfaceRepo) that are required for a pipeline. At the moment, the artifacts are listed here per pipeline, but we'll likely want to move towards each pipeline specifying its own artifacts instead of centralizing them in a single place.
I think we should align the way we download "extra" VAEs with this updated process because it gives us automatica downloads, download progress updates in the UI, support for uv run download_models --pipeline <PIPELINE_ID> simply by listing an artifact. And I'd prefer not to special case VAEs with their own API endpoints and/or code paths since they are really just another type of model.
This leaves the question of UX in the UI. I see in this PR that the UI checks if a specific VAE needs to be downloaded, but also updates the downloader code for each pipeline to download all the extra VAEs if the pipeline models are being downloaded. We could probably update the /api/v1/models/download and /api/v1/models/status endpoints to download specific models instead of only accepting a pipeline ID, but since the extra VAEs here are so small eg each of them <= 50 MB I'm in favor of just including all of them as artifacts for each pipeline in the hardcoded PIPELINE_ARTIFACTS constant in this PR and then properly design a way to download specific individual models with the API endpoints separately.
feat: add VAE download infrastructure with metadata registry Adds complete infrastructure for downloading and managing VAE models independently of pipeline downloads. This change has no effect on the current code. When the ability to select different VAE types is added, the system will check if the selected VAE is downloaded and prompt for download if missing. Backend: - Add /api/v1/vae/status endpoint to check VAE download status - Add /api/v1/vae/download endpoint to trigger VAE downloads - Add VAEMetadata dataclass and VAE_METADATA registry as single source of truth for VAE filenames and download sources - Add vae_file_exists() and get_vae_file_path() using metadata registry - Add download_vae() and download_downloadable_vaes() functions Frontend: - Add checkVaeStatus() and downloadVae() API functions - Extend DownloadDialog to show VAE-specific download prompts - Add VAE status checking before stream start in StreamPage - Add separate VAE download flow with polling for completion Prepares codebase for upcoming PRs adding additional VAE types. Signed-off-by: RyanOnTheInside <7623207+ryanontheinside@users.noreply.github.com>
Signed-off-by: RyanOnTheInside <7623207+ryanontheinside@users.noreply.github.com>
ryanontheinside
force-pushed
the
ryanontheinside/refactor/vae-download-infrastructure
branch
from
ba8e68f to
070292c
Compare
|
|
||
|
|
||
| # Single source of truth for VAE metadata | ||
| VAE_METADATA: dict[str, VAEMetadata] = { |
There was a problem hiding this comment.
It looks like we're missing lighttae here.
| """Metadata for a VAE type (filenames, download sources).""" | ||
|
|
||
| filename: str | ||
| download_repo: str | None = None # None = bundled with main model repo |
There was a problem hiding this comment.
Do we still use VAEMetadata in the later PRs?
| files=["lightvaew2_1.pth"], | ||
| ) | ||
|
|
||
| TAE_ARTIFACT = HuggingfaceRepoArtifact( |
There was a problem hiding this comment.
It looks like we're missing a lighttae artifact.
feat: add VAE download infrastructure with metadata registry
Adds complete infrastructure for downloading and managing VAE models independently of pipeline downloads. This change has no effect on the current code. When the ability to select different VAE types is added, the system will check if the selected VAE is downloaded and prompt for download if missing.
Backend:
Frontend:
Prepares codebase for upcoming PRs adding additional VAE types.