refactor: Enhance checkpoint download functionality with new modal and manager integration

static/js/checkpoints.js

@@ -3,6 +3,7 @@ import { initializeInfiniteScroll } from './utils/infiniteScroll.js';
 import { confirmDelete, closeDeleteModal } from './utils/modalUtils.js';
 import { createPageControls } from './components/controls/index.js';
 import { loadMoreCheckpoints } from './api/checkpointApi.js';
+import { CheckpointDownloadManager } from './managers/CheckpointDownloadManager.js';
 
 // Initialize the Checkpoints page
 class CheckpointsPageManager {
@@ -10,6 +11,9 @@ class CheckpointsPageManager {
         // Initialize page controls
         this.pageControls = createPageControls('checkpoints');
         
+        // Initialize checkpoint download manager
+        window.checkpointDownloadManager = new CheckpointDownloadManager();
+        
         // Expose only necessary functions to global scope
         this._exposeRequiredGlobalFunctions();
     }

static/js/components/controls/CheckpointsControls.js

@@ -2,6 +2,7 @@
 import { PageControls } from './PageControls.js';
 import { loadMoreCheckpoints, resetAndReload, refreshCheckpoints, fetchCivitai } from '../../api/checkpointApi.js';
 import { showToast } from '../../utils/uiHelpers.js';
+import { CheckpointDownloadManager } from '../../managers/CheckpointDownloadManager.js';
 
 /**
  * CheckpointsControls class - Extends PageControls for Checkpoint-specific functionality
@@ -11,6 +12,9 @@ export class CheckpointsControls extends PageControls {
         // Initialize with 'checkpoints' page type
         super('checkpoints');
         
+        // Initialize checkpoint download manager
+        this.downloadManager = new CheckpointDownloadManager();
+        
         // Register API methods specific to the Checkpoints page
         this.registerCheckpointsAPI();
     }
@@ -38,6 +42,11 @@ export class CheckpointsControls extends PageControls {
                 return await fetchCivitai();
             },
             
+            // Add show download modal functionality
+            showDownloadModal: () => {
+                this.downloadManager.showDownloadModal();
+            },
+            
             // No clearCustomFilter implementation is needed for checkpoints
             // as custom filters are currently only used for LoRAs
             clearCustomFilter: async () => {

static/js/components/controls/PageControls.js

@@ -103,13 +103,12 @@ export class PageControls {
             fetchButton.addEventListener('click', () => this.fetchFromCivitai());
         }
         
+        const downloadButton = document.querySelector('[data-action="download"]');
+        if (downloadButton) {
+            downloadButton.addEventListener('click', () => this.showDownloadModal());
+        }
+        
         if (this.pageType === 'loras') {
-            // Download button - LoRAs only
-            const downloadButton = document.querySelector('[data-action="download"]');
-            if (downloadButton) {
-                downloadButton.addEventListener('click', () => this.showDownloadModal());
-            }
-            
             // Bulk operations button - LoRAs only
             const bulkButton = document.querySelector('[data-action="bulk"]');
             if (bulkButton) {
@@ -349,14 +348,9 @@ export class PageControls {
     }
     
     /**
-     * Show download modal (LoRAs only)
+     * Show download modal
      */
     showDownloadModal() {
-        if (this.pageType !== 'loras' || !this.api) {
-            console.error('Download modal is only available for LoRAs');
-            return;
-        }
-        
         this.api.showDownloadModal();
     }
     

static/js/managers/CheckpointDownloadManager.js

@@ -0,0 +1,423 @@
+import { modalManager } from './ModalManager.js';
+import { showToast } from '../utils/uiHelpers.js';
+import { LoadingManager } from './LoadingManager.js';
+import { state } from '../state/index.js';
+import { resetAndReload } from '../api/checkpointApi.js';
+import { getStorageItem } from '../utils/storageHelpers.js';
+
+export class CheckpointDownloadManager {
+    constructor() {
+        this.currentVersion = null;
+        this.versions = [];
+        this.modelInfo = null;
+        this.modelVersionId = null;
+        
+        this.initialized = false;
+        this.selectedFolder = '';
+
+        this.loadingManager = new LoadingManager();
+        this.folderClickHandler = null;
+        this.updateTargetPath = this.updateTargetPath.bind(this);
+    }
+
+    showDownloadModal() {
+        console.log('Showing checkpoint download modal...');
+        if (!this.initialized) {
+            const modal = document.getElementById('checkpointDownloadModal');
+            if (!modal) {
+                console.error('Checkpoint download modal element not found');
+                return;
+            }
+            this.initialized = true;
+        }
+        
+        modalManager.showModal('checkpointDownloadModal', null, () => {
+            // Cleanup handler when modal closes
+            this.cleanupFolderBrowser();
+        });
+        this.resetSteps();
+    }
+
+    resetSteps() {
+        document.querySelectorAll('#checkpointDownloadModal .download-step').forEach(step => step.style.display = 'none');
+        document.getElementById('cpUrlStep').style.display = 'block';
+        document.getElementById('checkpointUrl').value = '';
+        document.getElementById('cpUrlError').textContent = '';
+        
+        // Clear new folder input
+        const newFolderInput = document.getElementById('cpNewFolder');
+        if (newFolderInput) {
+            newFolderInput.value = '';
+        }
+        
+        this.currentVersion = null;
+        this.versions = [];
+        this.modelInfo = null;
+        this.modelVersionId = null;
+        
+        // Clear selected folder and remove selection from UI
+        this.selectedFolder = '';
+        const folderBrowser = document.getElementById('cpFolderBrowser');
+        if (folderBrowser) {
+            folderBrowser.querySelectorAll('.folder-item').forEach(f => 
+                f.classList.remove('selected'));
+        }
+    }
+
+    async validateAndFetchVersions() {
+        const url = document.getElementById('checkpointUrl').value.trim();
+        const errorElement = document.getElementById('cpUrlError');
+        
+        try {
+            this.loadingManager.showSimpleLoading('Fetching model versions...');
+            
+            const modelId = this.extractModelId(url);
+            if (!modelId) {
+                throw new Error('Invalid Civitai URL format');
+            }
+
+            const response = await fetch(`/api/civitai/versions/${modelId}`);
+            if (!response.ok) {
+                throw new Error('Failed to fetch model versions');
+            }
+            
+            this.versions = await response.json();
+            if (!this.versions.length) {
+                throw new Error('No versions available for this model');
+            }
+            
+            // If we have a version ID from URL, pre-select it
+            if (this.modelVersionId) {
+                this.currentVersion = this.versions.find(v => v.id.toString() === this.modelVersionId);
+            }
+            
+            this.showVersionStep();
+        } catch (error) {
+            errorElement.textContent = error.message;
+        } finally {
+            this.loadingManager.hide();
+        }
+    }
+
+    extractModelId(url) {
+        const modelMatch = url.match(/civitai\.com\/models\/(\d+)/);
+        const versionMatch = url.match(/modelVersionId=(\d+)/);
+        
+        if (modelMatch) {
+            this.modelVersionId = versionMatch ? versionMatch[1] : null;
+            return modelMatch[1];
+        }
+        return null;
+    }
+
+    showVersionStep() {
+        document.getElementById('cpUrlStep').style.display = 'none';
+        document.getElementById('cpVersionStep').style.display = 'block';
+        
+        const versionList = document.getElementById('cpVersionList');
+        versionList.innerHTML = this.versions.map(version => {
+            const firstImage = version.images?.find(img => !img.url.endsWith('.mp4'));
+            const thumbnailUrl = firstImage ? firstImage.url : '/loras_static/images/no-preview.png';
+            
+            // Use version-level size or fallback to first file
+            const fileSize = version.modelSizeKB ? 
+                (version.modelSizeKB / 1024).toFixed(2) : 
+                (version.files[0]?.sizeKB / 1024).toFixed(2);
+            
+            // Use version-level existsLocally flag
+            const existsLocally = version.existsLocally;
+            const localPath = version.localPath;
+            
+            // Check if this is an early access version
+            const isEarlyAccess = version.availability === 'EarlyAccess';
+            
+            // Create early access badge if needed
+            let earlyAccessBadge = '';
+            if (isEarlyAccess) {
+                earlyAccessBadge = `
+                    <div class="early-access-badge" title="Early access required">
+                        <i class="fas fa-clock"></i> Early Access
+                    </div>
+                `;
+            }
+            
+            // Status badge for local models
+            const localStatus = existsLocally ? 
+                `<div class="local-badge">
+                    <i class="fas fa-check"></i> In Library
+                    <div class="local-path">${localPath || ''}</div>
+                 </div>` : '';
+
+            return `
+                <div class="version-item ${this.currentVersion?.id === version.id ? 'selected' : ''} 
+                     ${existsLocally ? 'exists-locally' : ''} 
+                     ${isEarlyAccess ? 'is-early-access' : ''}"
+                     onclick="checkpointDownloadManager.selectVersion('${version.id}')">
+                    <div class="version-thumbnail">
+                        <img src="${thumbnailUrl}" alt="Version preview">
+                    </div>
+                    <div class="version-content">
+                        <div class="version-header">
+                            <h3>${version.name}</h3>
+                            ${localStatus}
+                        </div>
+                        <div class="version-info">
+                            ${version.baseModel ? `<div class="base-model">${version.baseModel}</div>` : ''}
+                            ${earlyAccessBadge}
+                        </div>
+                        <div class="version-meta">
+                            <span><i class="fas fa-calendar"></i> ${new Date(version.createdAt).toLocaleDateString()}</span>
+                            <span><i class="fas fa-file-archive"></i> ${fileSize} MB</span>
+                        </div>
+                    </div>
+                </div>
+            `;
+        }).join('');
+        
+        // Update Next button state based on initial selection
+        this.updateNextButtonState();
+    }
+
+    selectVersion(versionId) {
+        this.currentVersion = this.versions.find(v => v.id.toString() === versionId.toString());
+        if (!this.currentVersion) return;
+
+        document.querySelectorAll('#cpVersionList .version-item').forEach(item => {
+            item.classList.toggle('selected', item.querySelector('h3').textContent === this.currentVersion.name);
+        });
+        
+        // Update Next button state after selection
+        this.updateNextButtonState();
+    }
+    
+    updateNextButtonState() {
+        const nextButton = document.querySelector('#cpVersionStep .primary-btn');
+        if (!nextButton) return;
+        
+        const existsLocally = this.currentVersion?.existsLocally;
+        
+        if (existsLocally) {
+            nextButton.disabled = true;
+            nextButton.classList.add('disabled');
+            nextButton.textContent = 'Already in Library';
+        } else {
+            nextButton.disabled = false;
+            nextButton.classList.remove('disabled');
+            nextButton.textContent = 'Next';
+        }
+    }
+
+    async proceedToLocation() {
+        if (!this.currentVersion) {
+            showToast('Please select a version', 'error');
+            return;
+        }
+        
+        // Double-check if the version exists locally
+        const existsLocally = this.currentVersion.existsLocally;
+        if (existsLocally) {
+            showToast('This version already exists in your library', 'info');
+            return;
+        }
+
+        document.getElementById('cpVersionStep').style.display = 'none';
+        document.getElementById('cpLocationStep').style.display = 'block';
+        
+        try {
+            // Use checkpoint roots endpoint instead of lora roots
+            const response = await fetch('/api/checkpoints/roots');
+            if (!response.ok) {
+                throw new Error('Failed to fetch checkpoint roots');
+            }
+            
+            const data = await response.json();
+            const checkpointRoot = document.getElementById('checkpointRoot');
+            checkpointRoot.innerHTML = data.roots.map(root => 
+                `<option value="${root}">${root}</option>`
+            ).join('');
+
+            // Set default checkpoint root if available
+            const defaultRoot = getStorageItem('settings', {}).default_checkpoints_root;
+            if (defaultRoot && data.roots.includes(defaultRoot)) {
+                checkpointRoot.value = defaultRoot;
+            }
+
+            // Initialize folder browser after loading roots
+            this.initializeFolderBrowser();
+        } catch (error) {
+            showToast(error.message, 'error');
+        }
+    }
+
+    backToUrl() {
+        document.getElementById('cpVersionStep').style.display = 'none';
+        document.getElementById('cpUrlStep').style.display = 'block';
+    }
+
+    backToVersions() {
+        document.getElementById('cpLocationStep').style.display = 'none';
+        document.getElementById('cpVersionStep').style.display = 'block';
+    }
+
+    async startDownload() {
+        const checkpointRoot = document.getElementById('checkpointRoot').value;
+        const newFolder = document.getElementById('cpNewFolder').value.trim();
+        
+        if (!checkpointRoot) {
+            showToast('Please select a checkpoint root directory', 'error');
+            return;
+        }
+
+        // Construct relative path
+        let targetFolder = '';
+        if (this.selectedFolder) {
+            targetFolder = this.selectedFolder;
+        }
+        if (newFolder) {
+            targetFolder = targetFolder ? 
+                `${targetFolder}/${newFolder}` : newFolder;
+        }
+
+        try {
+            const downloadUrl = this.currentVersion.downloadUrl;
+            if (!downloadUrl) {
+                throw new Error('No download URL available');
+            }
+
+            // Show enhanced loading with progress details
+            const updateProgress = this.loadingManager.showDownloadProgress(1);
+            updateProgress(0, 0, this.currentVersion.name);
+
+            // Setup WebSocket for progress updates
+            const wsProtocol = window.location.protocol === 'https:' ? 'wss://' : 'ws://';
+            const ws = new WebSocket(`${wsProtocol}${window.location.host}/ws/fetch-progress`);
+            
+            ws.onmessage = (event) => {
+                const data = JSON.parse(event.data);
+                if (data.status === 'progress') {
+                    // Update progress display with current progress
+                    updateProgress(data.progress, 0, this.currentVersion.name);
+                    
+                    // Add more detailed status messages based on progress
+                    if (data.progress < 3) {
+                        this.loadingManager.setStatus(`Preparing download...`);
+                    } else if (data.progress === 3) {
+                        this.loadingManager.setStatus(`Downloaded preview image`);
+                    } else if (data.progress > 3 && data.progress < 100) {
+                        this.loadingManager.setStatus(`Downloading checkpoint file`);
+                    } else {
+                        this.loadingManager.setStatus(`Finalizing download...`);
+                    }
+                }
+            };
+            
+            ws.onerror = (error) => {
+                console.error('WebSocket error:', error);
+                // Continue with download even if WebSocket fails
+            };
+
+            // Start download using checkpoint download endpoint
+            const response = await fetch('/api/checkpoints/download', {
+                method: 'POST',
+                headers: { 'Content-Type': 'application/json' },
+                body: JSON.stringify({
+                    download_url: downloadUrl,
+                    checkpoint_root: checkpointRoot,
+                    relative_path: targetFolder
+                })
+            });
+
+            if (!response.ok) {
+                throw new Error(await response.text());
+            }
+
+            showToast('Download completed successfully', 'success');
+            modalManager.closeModal('checkpointDownloadModal');
+            
+            // Update state and trigger reload with folder update
+            state.activeFolder = targetFolder;
+            await resetAndReload(true); // Pass true to update folders
+
+        } catch (error) {
+            showToast(error.message, 'error');
+        } finally {
+            this.loadingManager.hide();
+        }
+    }
+
+    initializeFolderBrowser() {
+        const folderBrowser = document.getElementById('cpFolderBrowser');
+        if (!folderBrowser) return;
+
+        // Cleanup existing handler if any
+        this.cleanupFolderBrowser();
+
+        // Create new handler
+        this.folderClickHandler = (event) => {
+            const folderItem = event.target.closest('.folder-item');
+            if (!folderItem) return;
+
+            if (folderItem.classList.contains('selected')) {
+                folderItem.classList.remove('selected');
+                this.selectedFolder = '';
+            } else {
+                folderBrowser.querySelectorAll('.folder-item').forEach(f => 
+                    f.classList.remove('selected'));
+                folderItem.classList.add('selected');
+                this.selectedFolder = folderItem.dataset.folder;
+            }
+            
+            // Update path display after folder selection
+            this.updateTargetPath();
+        };
+
+        // Add the new handler
+        folderBrowser.addEventListener('click', this.folderClickHandler);
+        
+        // Add event listeners for path updates
+        const checkpointRoot = document.getElementById('checkpointRoot');
+        const newFolder = document.getElementById('cpNewFolder');
+        
+        checkpointRoot.addEventListener('change', this.updateTargetPath);
+        newFolder.addEventListener('input', this.updateTargetPath);
+        
+        // Update initial path
+        this.updateTargetPath();
+    }
+
+    cleanupFolderBrowser() {
+        if (this.folderClickHandler) {
+            const folderBrowser = document.getElementById('cpFolderBrowser');
+            if (folderBrowser) {
+                folderBrowser.removeEventListener('click', this.folderClickHandler);
+                this.folderClickHandler = null;
+            }
+        }
+        
+        // Remove path update listeners
+        const checkpointRoot = document.getElementById('checkpointRoot');
+        const newFolder = document.getElementById('cpNewFolder');
+        
+        if (checkpointRoot) checkpointRoot.removeEventListener('change', this.updateTargetPath);
+        if (newFolder) newFolder.removeEventListener('input', this.updateTargetPath);
+    }
+    
+    updateTargetPath() {
+        const pathDisplay = document.getElementById('cpTargetPathDisplay');
+        const checkpointRoot = document.getElementById('checkpointRoot').value;
+        const newFolder = document.getElementById('cpNewFolder').value.trim();
+        
+        let fullPath = checkpointRoot || 'Select a checkpoint root directory';
+        
+        if (checkpointRoot) {
+            if (this.selectedFolder) {
+                fullPath += '/' + this.selectedFolder;
+            }
+            if (newFolder) {
+                fullPath += '/' + newFolder;
+            }
+        }
+
+        pathDisplay.innerHTML = `<span class="path-text">${fullPath}</span>`;
+    }
+}

static/js/managers/ModalManager.js

@@ -35,6 +35,19 @@ export class ModalManager {
                 closeOnOutsideClick: true
             });
         }
+
+        // Add checkpointDownloadModal registration
+        const checkpointDownloadModal = document.getElementById('checkpointDownloadModal');
+        if (checkpointDownloadModal) {
+            this.registerModal('checkpointDownloadModal', {
+                element: checkpointDownloadModal,
+                onClose: () => {
+                    this.getModal('checkpointDownloadModal').element.style.display = 'none';
+                    document.body.classList.remove('modal-open');
+                },
+                closeOnOutsideClick: true
+            });
+        }
         
         const deleteModal = document.getElementById('deleteModal');
         if (deleteModal) {