💬 Commit message: Update 2026-02-14 07:13:36, 9 files, 312 lines

📁 Files changed: 9 📝 Lines changed: 312 • api.py • download_routes.py • index-BIJHSBBO.js • index-D2fa7i11.css • index-Dd0b3iAi.js • index-lqI0My76.css • index.html • client.ts • ModelCard.vue
2026-02-14 07:13:36 +01:00
parent c2f915c59c
commit 120c4c68b0
9 changed files with 277 additions and 31 deletions
@@ -3,6 +3,8 @@
 from __future__ import annotations

 import re
+import time
+from collections.abc import Callable
 from http import HTTPStatus
 from typing import TYPE_CHECKING, Any

@@ -26,6 +28,9 @@ from tensors.config import CIVITAI_API_BASE, CIVITAI_DOWNLOAD_BASE, BaseModel, M
 if TYPE_CHECKING:
    from rich.console import Console

+# Progress update throttle interval (seconds)
+_PROGRESS_UPDATE_INTERVAL = 0.25
+

 def _get_headers(api_key: str | None) -> dict[str, str]:
    """Get headers for CivitAI API requests."""
@@ -262,6 +267,98 @@ def _stream_download(
    return True


+# Type alias for progress callback: (downloaded_bytes, total_bytes, speed_bytes_per_sec)
+ProgressCallback = Callable[[int, int, float], None]
+
+
+def _stream_download_with_callback(
+    response: httpx.Response,
+    dest_path: Path,
+    mode: str,
+    initial_size: int,
+    on_progress: ProgressCallback | None = None,
+) -> bool:
+    """Stream download content to file with progress callback."""
+    content_length = response.headers.get("content-length")
+    total_size = int(content_length) + initial_size if content_length else 0
+    downloaded = initial_size
+    start_time = time.time()
+    last_time = start_time
+
+    with dest_path.open(mode) as f:
+        for chunk in response.iter_bytes(1024 * 1024):
+            f.write(chunk)
+            downloaded += len(chunk)
+
+            if on_progress:
+                now = time.time()
+                elapsed = now - start_time
+                speed = downloaded / elapsed if elapsed > 0 else 0
+                # Throttle updates
+                if now - last_time >= _PROGRESS_UPDATE_INTERVAL:
+                    on_progress(downloaded, total_size, speed)
+                    last_time = now
+
+    # Final progress update
+    if on_progress:
+        elapsed = time.time() - start_time
+        speed = downloaded / elapsed if elapsed > 0 else 0
+        on_progress(downloaded, total_size, speed)
+
+    return True
+
+
+def download_model_with_progress(
+    version_id: int,
+    dest_path: Path,
+    api_key: str | None,
+    on_progress: ProgressCallback | None = None,
+    resume: bool = True,
+) -> bool:
+    """Download a model from CivitAI with progress callback instead of console output."""
+    import logging  # noqa: PLC0415
+
+    logger = logging.getLogger(__name__)
+    url = f"{CIVITAI_DOWNLOAD_BASE}/{version_id}"
+    params: dict[str, str] = {}
+    if api_key:
+        params["token"] = api_key
+
+    # Set up resume
+    headers: dict[str, str] = {}
+    mode = "wb"
+    initial_size = 0
+
+    if resume and dest_path.exists():
+        initial_size = dest_path.stat().st_size
+        headers["Range"] = f"bytes={initial_size}-"
+        mode = "ab"
+        logger.info(f"Resuming download from {initial_size / (1024**2):.1f} MB")
+
+    try:
+        with httpx.stream(
+            "GET",
+            url,
+            params=params,
+            headers=headers,
+            follow_redirects=True,
+            timeout=httpx.Timeout(30.0, read=None),
+        ) as response:
+            if response.status_code == HTTPStatus.REQUESTED_RANGE_NOT_SATISFIABLE:
+                return True  # Already complete
+
+            response.raise_for_status()
+            dest_path = _get_dest_from_response(response, dest_path)
+            return _stream_download_with_callback(response, dest_path, mode, initial_size, on_progress)
+
+    except httpx.HTTPStatusError as e:
+        logger.error(f"Download error: HTTP {e.response.status_code}")
+        return False
+    except httpx.RequestError as e:
+        logger.error(f"Download error: {e}")
+        return False
+
+
 def download_model(
    version_id: int,
    dest_path: Path,
@@ -9,7 +9,7 @@ from typing import Any
 from fastapi import APIRouter, BackgroundTasks, HTTPException
 from pydantic import BaseModel as PydanticBaseModel

-from tensors.api import download_model, fetch_civitai_by_hash, fetch_civitai_model, fetch_civitai_model_version
+from tensors.api import download_model_with_progress, fetch_civitai_by_hash, fetch_civitai_model, fetch_civitai_model_version
 from tensors.config import MODELS_DIR, load_api_key
 from tensors.db import Database

@@ -89,6 +89,22 @@ def _get_output_dir(version_info: dict[str, Any], override: str | None) -> Path:
    return type_dirs.get(model_type, MODELS_DIR / "other")


+_KB = 1024
+_MB = _KB * 1024
+_GB = _MB * 1024
+
+
+def _format_size(size_bytes: int) -> str:
+    """Format bytes as human readable string."""
+    if size_bytes >= _GB:
+        return f"{size_bytes / _GB:.1f} GB"
+    if size_bytes >= _MB:
+        return f"{size_bytes / _MB:.1f} MB"
+    if size_bytes >= _KB:
+        return f"{size_bytes / _KB:.1f} KB"
+    return f"{size_bytes} B"
+
+
 def _do_download(
    version_id: int,
    dest_path: Path,
@@ -98,19 +114,27 @@ def _do_download(
    """Background task to perform the download."""
    try:
        _active_downloads[download_id]["status"] = "downloading"
+        _active_downloads[download_id]["downloaded"] = 0
+        _active_downloads[download_id]["total"] = 0
+        _active_downloads[download_id]["speed"] = 0
+        _active_downloads[download_id]["progress"] = 0

-        # Create a mock console for download progress
-        from io import StringIO  # noqa: PLC0415
+        def on_progress(downloaded: int, total: int, speed: float) -> None:
+            """Update progress in active downloads dict."""
+            _active_downloads[download_id]["downloaded"] = downloaded
+            _active_downloads[download_id]["total"] = total
+            _active_downloads[download_id]["speed"] = speed
+            _active_downloads[download_id]["downloaded_str"] = _format_size(downloaded)
+            _active_downloads[download_id]["total_str"] = _format_size(total) if total > 0 else "Unknown"
+            _active_downloads[download_id]["speed_str"] = f"{_format_size(int(speed))}/s"
+            if total > 0:
+                _active_downloads[download_id]["progress"] = round(100 * downloaded / total, 1)

-        from rich.console import Console  # noqa: PLC0415
-
-        output = StringIO()
-        console = Console(file=output, force_terminal=False)
-
-        success = download_model(version_id, dest_path, api_key, console, resume=True)
+        success = download_model_with_progress(version_id, dest_path, api_key, on_progress, resume=True)

        if success:
            _active_downloads[download_id]["status"] = "completed"
+            _active_downloads[download_id]["progress"] = 100
            _active_downloads[download_id]["path"] = str(dest_path)

            # Auto-scan and link the downloaded file
@@ -5,8 +5,8 @@
    <link rel="icon" type="image/svg+xml" href="/vite.svg" />
    <meta name="viewport" content="width=device-width, initial-scale=1.0" />
    <title>Tensors</title>
-    <script type="module" crossorigin src="/assets/index-Dd0b3iAi.js"></script>
-    <link rel="stylesheet" crossorigin href="/assets/index-D2fa7i11.css">
+    <script type="module" crossorigin src="/assets/index-BIJHSBBO.js"></script>
+    <link rel="stylesheet" crossorigin href="/assets/index-lqI0My76.css">
  </head>
  <body>
    <div id="app"></div>
@@ -94,9 +94,44 @@ export async function getCivitaiModel(id: number): Promise<CivitaiModel> {
 }

 // Download
-export async function downloadModel(modelId?: number, versionId?: number): Promise<{ ok: boolean }> {
+export interface DownloadResponse {
+  download_id: string
+  status: string
+  version_id: number
+  destination: string
+  model_name: string
+  version_name: string
+}
+
+export interface DownloadStatus {
+  id: string
+  version_id: number
+  status: 'queued' | 'downloading' | 'completed' | 'failed'
+  path: string
+  filename: string
+  model_name: string
+  version_name: string
+  downloaded?: number
+  total?: number
+  progress?: number
+  speed?: number
+  downloaded_str?: string
+  total_str?: string
+  speed_str?: string
+  error?: string
+}
+
+export async function downloadModel(modelId?: number, versionId?: number): Promise<DownloadResponse> {
  return fetchJson('/api/download', {
    method: 'POST',
    body: JSON.stringify({ model_id: modelId, version_id: versionId }),
  })
 }
+
+export async function getDownloadStatus(downloadId: string): Promise<DownloadStatus> {
+  return fetchJson(`/api/download/status/${downloadId}`)
+}
+
+export async function getActiveDownloads(): Promise<{ downloads: DownloadStatus[]; total: number }> {
+  return fetchJson('/api/download/active')
+}
@@ -1,6 +1,7 @@
 <script setup lang="ts">
-import { ref, computed } from 'vue'
+import { ref, computed, onUnmounted } from 'vue'
 import type { CivitaiModel } from '@/types'
+import type { DownloadStatus } from '@/api/client'
 import * as api from '@/api/client'

 const props = defineProps<{
@@ -11,6 +12,15 @@ const showDialog = ref(false)
 const loadingDetails = ref(false)
 const modelDetails = ref<CivitaiModel | null>(null)

+// Download tracking
+const activeDownload = ref<DownloadStatus | null>(null)
+const downloadingVersionId = ref<number | null>(null)
+let pollInterval: ReturnType<typeof setInterval> | null = null
+
+onUnmounted(() => {
+  if (pollInterval) clearInterval(pollInterval)
+})
+
 const previewImage = computed(() => {
  const version = props.model.modelVersions?.[0]
  return version?.images?.[0]?.url || ''
@@ -43,13 +53,41 @@ async function openDetails() {
 }

 async function downloadVersion(versionId: number) {
+  if (downloadingVersionId.value === versionId) return // Already downloading
  if (!confirm(`Download "${props.model.name}" to the server?`)) return

  try {
-    await api.downloadModel(undefined, versionId)
-    alert('Download started! Check server for progress.')
+    downloadingVersionId.value = versionId
+    const response = await api.downloadModel(undefined, versionId)
+
+    // Start polling for progress
+    const downloadId = response.download_id
+    pollInterval = setInterval(async () => {
+      try {
+        const status = await api.getDownloadStatus(downloadId)
+        activeDownload.value = status
+
+        if (status.status === 'completed' || status.status === 'failed') {
+          if (pollInterval) clearInterval(pollInterval)
+          pollInterval = null
+
+          if (status.status === 'failed') {
+            alert('Download failed: ' + (status.error || 'Unknown error'))
+          }
+
+          // Keep showing completed state briefly, then clear
+          setTimeout(() => {
+            activeDownload.value = null
+            downloadingVersionId.value = null
+          }, 3000)
+        }
+      } catch (e) {
+        console.error('Failed to get download status:', e)
+      }
+    }, 500)
  } catch (e: any) {
    alert('Download failed: ' + e.message)
+    downloadingVersionId.value = null
  }
 }
 </script>
@@ -105,7 +143,32 @@ async function downloadVersion(versionId: number) {
        by {{ model.creator?.username || 'Unknown' }}
      </span>
      <v-spacer />
+      <!-- Download progress or button -->
+      <template v-if="downloadingVersionId === model.modelVersions?.[0]?.id && activeDownload">
+        <div class="download-progress">
+          <v-progress-linear
+            :model-value="activeDownload.progress || 0"
+            :indeterminate="activeDownload.status === 'queued'"
+            :color="activeDownload.status === 'completed' ? 'success' : 'primary'"
+            height="6"
+            rounded
+          />
+          <span class="text-caption text-grey">
+            <template v-if="activeDownload.status === 'completed'">
+              Done!
+            </template>
+            <template v-else-if="activeDownload.status === 'queued'">
+              Queued...
+            </template>
+            <template v-else>
+              {{ activeDownload.downloaded_str }} / {{ activeDownload.total_str }}
+              <span class="text-primary ml-1">{{ activeDownload.speed_str }}</span>
+            </template>
+          </span>
+        </div>
+      </template>
      <v-btn
+        v-else
        size="small"
        color="primary"
        variant="flat"
@@ -191,7 +254,22 @@ async function downloadVersion(versionId: number) {
                <v-list-item-title>{{ version.name }}</v-list-item-title>
                <v-list-item-subtitle>{{ version.baseModel }}</v-list-item-subtitle>
                <template #append>
+                  <template v-if="downloadingVersionId === version.id && activeDownload">
+                    <div class="download-progress-dialog">
+                      <v-progress-linear
+                        :model-value="activeDownload.progress || 0"
+                        :indeterminate="activeDownload.status === 'queued'"
+                        :color="activeDownload.status === 'completed' ? 'success' : 'primary'"
+                        height="6"
+                        rounded
+                      />
+                      <span class="text-caption text-grey">
+                        {{ activeDownload.progress?.toFixed(0) || 0 }}%
+                      </span>
+                    </div>
+                  </template>
                  <v-btn
+                    v-else
                    size="small"
                    color="primary"
                    variant="flat"
@@ -224,4 +302,16 @@ async function downloadVersion(versionId: number) {
  -webkit-box-orient: vertical;
  overflow: hidden;
 }
+.download-progress {
+  display: flex;
+  flex-direction: column;
+  gap: 2px;
+  width: 120px;
+}
+.download-progress-dialog {
+  display: flex;
+  align-items: center;
+  gap: 8px;
+  width: 100px;
+}
 </style>