From 3a87e4c60142b56e8183d68f627abac95f2d928d Mon Sep 17 00:00:00 2001
From: Luke Mino-Altherr <luke@comfy.org>
Date: Thu, 13 Nov 2025 18:22:29 -0800
Subject: [PATCH] styling and fixes

---
 src/components/common/UrlInput.vue            |   3 +
 .../components/UploadModelConfirmation.vue    |  48 +++---
 .../assets/components/UploadModelDialog.vue   | 151 ++++++++----------
 .../assets/components/UploadModelProgress.vue |  55 +++----
 .../assets/components/UploadModelUrlInput.vue |  28 ++--
 .../assets/composables/useModelTypes.ts       |  94 +++++++++++
 src/platform/assets/services/assetService.ts  |  21 ++-
 7 files changed, 245 insertions(+), 155 deletions(-)
 create mode 100644 src/platform/assets/composables/useModelTypes.ts

diff --git a/src/components/common/UrlInput.vue b/src/components/common/UrlInput.vue
index df40ba0c4..b0dfbd60f 100644
--- a/src/components/common/UrlInput.vue
+++ b/src/components/common/UrlInput.vue
@@ -35,6 +35,7 @@ import { ValidationState } from '@/utils/validationUtil'
 const props = defineProps<{
   modelValue: string
   validateUrlFn?: (url: string) => Promise<boolean>
+  disableValidation?: boolean
 }>()
 
 const emit = defineEmits<{
@@ -101,6 +102,8 @@ const defaultValidateUrl = async (url: string): Promise<boolean> => {
 }
 
 const validateUrl = async (value: string) => {
+  if (props.disableValidation) return
+
   if (validationState.value === ValidationState.LOADING) return
 
   const url = cleanInput(value)
diff --git a/src/platform/assets/components/UploadModelConfirmation.vue b/src/platform/assets/components/UploadModelConfirmation.vue
index 23ae622b9..2471967f5 100644
--- a/src/platform/assets/components/UploadModelConfirmation.vue
+++ b/src/platform/assets/components/UploadModelConfirmation.vue
@@ -1,8 +1,8 @@
 <template>
-  <div class="flex flex-col gap-6">
+  <div class="flex flex-col gap-4">
     <!-- Model Info Section -->
     <div class="flex flex-col gap-2">
-      <p class="text-sm text-muted mb-0">
+      <p class="text-sm text-muted m-0">
         {{ $t('assetBrowser.modelAssociatedWithLink') }}
       </p>
       <p class="text-sm mt-0">
@@ -12,23 +12,14 @@
 
     <!-- Model Type Selection -->
     <div class="flex flex-col gap-2">
-      <label for="model-type" class="text-sm text-muted">
+      <label class="text-sm text-muted">
         {{ $t('assetBrowser.whatTypeOfModel') }}
       </label>
-      <select
-        id="model-type"
-        :value="modelValue"
-        class="w-full rounded-md border border-border bg-surface px-3 py-2 text-sm outline-none focus:border-primary focus:ring-1 focus:ring-primary"
-        @change="$emit('update:modelValue', ($event.target as HTMLSelectElement).value)"
-      >
-        <option
-          v-for="option in modelTypeOptions"
-          :key="option.value"
-          :value="option.value"
-        >
-          {{ option.label }}
-        </option>
-      </select>
+      <SingleSelect
+        v-model="selectedModelType"
+        :label="$t('assetBrowser.whatTypeOfModel')"
+        :options="modelTypes"
+      />
       <div class="flex items-center gap-2 text-sm text-muted">
         <i class="icon-[lucide--info]" />
         <span>{{ $t('assetBrowser.notSureLeaveAsIs') }}</span>
@@ -38,6 +29,11 @@
 </template>
 
 <script setup lang="ts">
+import { computed } from 'vue'
+
+import SingleSelect from '@/components/input/SingleSelect.vue'
+import { useModelTypes } from '@/platform/assets/composables/useModelTypes'
+
 interface ModelMetadata {
   content_length: number
   final_url: string
@@ -48,21 +44,19 @@ interface ModelMetadata {
   preview_url?: string
 }
 
-defineProps<{
+const props = defineProps<{
   modelValue: string
   metadata: ModelMetadata | null
 }>()
 
-defineEmits<{
+const emit = defineEmits<{
   'update:modelValue': [value: string]
 }>()
 
-const modelTypeOptions = [
-  { label: 'LoRA', value: 'lora' },
-  { label: 'Checkpoint', value: 'checkpoint' },
-  { label: 'Embedding', value: 'embedding' },
-  { label: 'VAE', value: 'vae' },
-  { label: 'Upscale Model', value: 'upscale_model' },
-  { label: 'ControlNet', value: 'controlnet' }
-]
+const { modelTypes } = useModelTypes()
+
+const selectedModelType = computed({
+  get: () => props.modelValue,
+  set: (value: string) => emit('update:modelValue', value)
+})
 </script>
diff --git a/src/platform/assets/components/UploadModelDialog.vue b/src/platform/assets/components/UploadModelDialog.vue
index 71791b434..0790bc972 100644
--- a/src/platform/assets/components/UploadModelDialog.vue
+++ b/src/platform/assets/components/UploadModelDialog.vue
@@ -1,10 +1,7 @@
 <template>
-  <div class="upload-model-dialog flex flex-col justify-between gap-6 p-4">
+  <div class="upload-model-dialog flex flex-col justify-between gap-6 p-4 pt-6">
     <!-- Step 1: Enter URL -->
-    <UploadModelUrlInput
-      v-if="currentStep === 1"
-      v-model="wizardData.url"
-    />
+    <UploadModelUrlInput v-if="currentStep === 1" v-model="wizardData.url" />
 
     <!-- Step 2: Confirm Metadata -->
     <UploadModelConfirmation
@@ -24,72 +21,69 @@
 
     <!-- Navigation Footer -->
     <div class="flex justify-end gap-2">
-      <button
+      <TextButton
         v-if="currentStep !== 1 && currentStep !== 3"
-        type="button"
-        :class="getButtonClasses('secondary')"
+        :label="$t('g.back')"
+        type="secondary"
+        size="md"
         :disabled="isFetchingMetadata || isUploading"
-        @click="goToPreviousStep"
-      >
-        {{ $t('g.back') }}
-      </button>
+        :on-click="goToPreviousStep"
+      />
       <span v-else />
 
-      <button
+      <IconTextButton
         v-if="currentStep === 1"
-        type="button"
-        :class="getButtonClasses('primary')"
-        :disabled="!canProceedStep1 || isFetchingMetadata"
-        @click="handleStep1Continue"
+        :label="$t('g.continue')"
+        type="primary"
+        size="md"
+        :disabled="!canFetchMetadata || isFetchingMetadata"
+        :on-click="handleFetchMetadata"
       >
-        <i
-          v-if="isFetchingMetadata"
-          class="icon-[lucide--loader-circle] mr-2 animate-spin"
-        />
-        {{ $t('g.continue') }}
-      </button>
-      <button
+        <template #icon>
+          <i
+            v-if="isFetchingMetadata"
+            class="icon-[lucide--loader-circle] animate-spin"
+          />
+        </template>
+      </IconTextButton>
+      <IconTextButton
         v-else-if="currentStep === 2"
-        type="button"
-        :class="getButtonClasses('primary')"
-        :disabled="!canProceedStep2 || isUploading"
-        @click="handleStep2Upload"
+        :label="$t('assetBrowser.upload')"
+        type="primary"
+        size="md"
+        :disabled="!canUploadModel || isUploading"
+        :on-click="handleUploadModel"
       >
-        <i
-          v-if="isUploading"
-          class="icon-[lucide--loader-circle] mr-2 animate-spin"
-        />
-        {{ $t('assetBrowser.upload') }}
-      </button>
-      <button
+        <template #icon>
+          <i
+            v-if="isUploading"
+            class="icon-[lucide--loader-circle] animate-spin"
+          />
+        </template>
+      </IconTextButton>
+      <TextButton
         v-else-if="currentStep === 3 && uploadStatus === 'success'"
-        type="button"
-        :class="getButtonClasses('primary')"
-        @click="handleClose"
-      >
-        {{ $t('assetBrowser.finish') }}
-      </button>
+        :label="$t('assetBrowser.finish')"
+        type="primary"
+        size="md"
+        :on-click="handleClose"
+      />
     </div>
   </div>
 </template>
 
 <script setup lang="ts">
-import { computed, ref } from 'vue'
-import { useI18n } from 'vue-i18n'
+import { computed, onMounted, ref } from 'vue'
 
-import { assetService } from '@/platform/assets/services/assetService'
+import IconTextButton from '@/components/button/IconTextButton.vue'
+import TextButton from '@/components/button/TextButton.vue'
 import UploadModelConfirmation from '@/platform/assets/components/UploadModelConfirmation.vue'
 import UploadModelProgress from '@/platform/assets/components/UploadModelProgress.vue'
 import UploadModelUrlInput from '@/platform/assets/components/UploadModelUrlInput.vue'
+import { useModelTypes } from '@/platform/assets/composables/useModelTypes'
+import { assetService } from '@/platform/assets/services/assetService'
 import { useDialogStore } from '@/stores/dialogStore'
-import {
-  getBaseButtonClasses,
-  getButtonSizeClasses,
-  getButtonTypeClasses
-} from '@/types/buttonTypes'
-import { cn } from '@/utils/tailwindUtil'
 
-const { t } = useI18n()
 const dialogStore = useDialogStore()
 
 const emit = defineEmits<{
@@ -122,38 +116,21 @@ const wizardData = ref<{
   tags: []
 })
 
-const selectedModelType = ref<string>('lora')
+const selectedModelType = ref<string>('loras')
 
-const modelTypeOptions = [
-  'lora',
-  'checkpoint',
-  'embedding',
-  'vae',
-  'upscale_model',
-  'controlnet'
-]
+const { modelTypes, fetchModelTypes } = useModelTypes()
 
-// Button styling helper
-function getButtonClasses(type: 'primary' | 'secondary') {
-  return cn(
-    getBaseButtonClasses(),
-    getButtonSizeClasses('md'),
-    getButtonTypeClasses(type)
-  )
-}
-
-// Step 1 validation
-const canProceedStep1 = computed(() => {
+// Validation
+const canFetchMetadata = computed(() => {
   return wizardData.value.url.trim().length > 0
 })
 
-// Step 2 validation
-const canProceedStep2 = computed(() => {
+const canUploadModel = computed(() => {
   return !!selectedModelType.value
 })
 
-async function handleStep1Continue() {
-  if (!canProceedStep1.value) return
+async function handleFetchMetadata() {
+  if (!canFetchMetadata.value) return
 
   isFetchingMetadata.value = true
   try {
@@ -167,8 +144,8 @@ async function handleStep1Continue() {
     if (metadata.tags && metadata.tags.length > 0) {
       wizardData.value.tags = metadata.tags
       // Try to detect model type from tags
-      const typeTag = metadata.tags.find(tag =>
-        modelTypeOptions.includes(tag)
+      const typeTag = metadata.tags.find((tag) =>
+        modelTypes.value.some((type) => type.value === tag)
       )
       if (typeTag) {
         selectedModelType.value = typeTag
@@ -178,22 +155,26 @@ async function handleStep1Continue() {
     currentStep.value = 2
   } catch (error) {
     console.error('Failed to retrieve metadata:', error)
-    uploadError.value = error instanceof Error ? error.message : 'Failed to retrieve metadata'
+    uploadError.value =
+      error instanceof Error ? error.message : 'Failed to retrieve metadata'
     // TODO: Show error toast to user
   } finally {
     isFetchingMetadata.value = false
   }
 }
 
-async function handleStep2Upload() {
-  if (!canProceedStep2.value) return
+async function handleUploadModel() {
+  if (!canUploadModel.value) return
 
   isUploading.value = true
   uploadStatus.value = 'uploading'
 
   try {
     const tags = ['models', selectedModelType.value]
-    const filename = wizardData.value.metadata?.filename || wizardData.value.metadata?.name || 'model'
+    const filename =
+      wizardData.value.metadata?.filename ||
+      wizardData.value.metadata?.name ||
+      'model'
 
     await assetService.uploadAssetFromUrl({
       url: wizardData.value.url,
@@ -212,12 +193,12 @@ async function handleStep2Upload() {
   } catch (error) {
     console.error('Failed to upload asset:', error)
     uploadStatus.value = 'error'
-    uploadError.value = error instanceof Error ? error.message : 'Failed to upload model'
+    uploadError.value =
+      error instanceof Error ? error.message : 'Failed to upload model'
     currentStep.value = 3
   } finally {
     isUploading.value = false
   }
-  
 }
 
 function goToPreviousStep() {
@@ -229,6 +210,10 @@ function goToPreviousStep() {
 function handleClose() {
   dialogStore.closeDialog({ key: 'upload-model' })
 }
+
+onMounted(() => {
+  fetchModelTypes()
+})
 </script>
 
 <style scoped>
diff --git a/src/platform/assets/components/UploadModelProgress.vue b/src/platform/assets/components/UploadModelProgress.vue
index e649c64fd..ba7e65b00 100644
--- a/src/platform/assets/components/UploadModelProgress.vue
+++ b/src/platform/assets/components/UploadModelProgress.vue
@@ -1,51 +1,40 @@
 <template>
-  <div class="flex flex-col gap-6">
+  <div class="flex flex-1 flex-col gap-6">
     <!-- Uploading State -->
     <div
       v-if="status === 'uploading'"
-      class="flex flex-col items-center justify-center gap-6 py-8"
+      class="flex flex-1 flex-col items-center justify-center gap-6"
     >
       <i
         class="icon-[lucide--loader-circle] animate-spin text-6xl text-primary"
       />
       <div class="text-center">
-        <h3 class="mb-2 text-lg font-semibold">
+        <p class="m-0 text-sm font-bold">
           {{ $t('assetBrowser.uploadingModel') }}
-        </h3>
+        </p>
       </div>
     </div>
 
     <!-- Success State -->
-    <div v-else-if="status === 'success'" class="flex flex-col gap-6">
-      <div class="flex items-start gap-2">
-        <h3 class="text-sm text-muted mb-0">
+    <div v-else-if="status === 'success'" class="flex flex-col gap-8">
+      <div class="flex flex-col gap-4">
+        <p class="text-sm text-muted m-0 font-bold">
           {{ $t('assetBrowser.modelUploaded') }} 🎉
-        </h3>
+        </p>
+        <p class="text-sm text-muted m-0">
+          {{ $t('assetBrowser.findInLibrary', { type: modelType }) }}
+        </p>
       </div>
 
-      <p class="text-sm text-muted">
-        {{ $t('assetBrowser.findInLibrary', { type: modelType.toUpperCase() }) }}
-      </p>
-
-      <div class="flex items-center gap-3 rounded-lg bg-surface-hover p-4">
-        <img
-          v-if="metadata?.preview_url"
-          :src="metadata.preview_url"
-          :alt="metadata?.name"
-          class="size-16 rounded object-cover"
-        />
-        <div
-          v-else
-          class="flex size-16 items-center justify-center rounded bg-surface"
-        >
-          <i class="icon-[lucide--image] text-2xl text-muted" />
-        </div>
-        <div class="flex flex-col gap-1">
-          <p class="text-base font-medium">
+      <div
+        class="flex flex-row items-start p-8 bg-node-component-surface rounded-lg"
+      >
+        <div class="flex flex-col justify-center items-start gap-1 flex-1">
+          <p class="text-sm m-0">
             {{ metadata?.name || metadata?.filename }}
           </p>
-          <p class="text-sm text-muted">
-            {{ modelType.toUpperCase() }}
+          <p class="text-sm text-muted m-0">
+            {{ modelType }}
           </p>
         </div>
       </div>
@@ -54,14 +43,14 @@
     <!-- Error State -->
     <div
       v-else-if="status === 'error'"
-      class="flex flex-col items-center justify-center gap-6 py-8"
+      class="flex flex-1 flex-col items-center justify-center gap-6"
     >
       <i class="icon-[lucide--x-circle] text-6xl text-red-500" />
       <div class="text-center">
-        <h3 class="mb-2 text-lg font-semibold">
+        <p class="m-0 text-sm font-bold">
           {{ $t('assetBrowser.uploadFailed') }}
-        </h3>
-        <p v-if="error" class="text-sm text-muted">
+        </p>
+        <p v-if="error" class="text-sm text-muted mb-0">
           {{ error }}
         </p>
       </div>
diff --git a/src/platform/assets/components/UploadModelUrlInput.vue b/src/platform/assets/components/UploadModelUrlInput.vue
index a0375411d..17d681b08 100644
--- a/src/platform/assets/components/UploadModelUrlInput.vue
+++ b/src/platform/assets/components/UploadModelUrlInput.vue
@@ -1,7 +1,7 @@
 <template>
-  <div class="flex flex-col gap-6">
+  <div class="flex flex-col gap-4">
     <div class="flex flex-col gap-2">
-      <p class="text-sm text-muted mb-0">
+      <p class="text-sm text-muted m-0">
         {{ $t('assetBrowser.uploadModelDescription1') }}
       </p>
       <ul class="list-disc space-y-1 pl-5 mt-0 text-sm text-muted">
@@ -11,16 +11,13 @@
     </div>
 
     <div class="flex flex-col gap-2">
-      <label for="civitai-link" class="text-sm text-muted mb-0">
+      <label class="text-sm text-muted mb-0">
         {{ $t('assetBrowser.civitaiLinkLabel') }}
       </label>
-      <input
-        id="civitai-link"
-        :value="modelValue"
-        type="text"
+      <UrlInput
+        v-model="url"
         :placeholder="$t('assetBrowser.civitaiLinkPlaceholder')"
-        class="w-full rounded-md border border-border bg-surface px-3 py-2 text-sm text-muted outline-none focus:border-primary focus:ring-1 focus:ring-primary"
-        @input="$emit('update:modelValue', ($event.target as HTMLInputElement).value)"
+        :disable-validation="true"
       />
       <p class="text-xs text-muted">
         {{ $t('assetBrowser.civitaiLinkExample') }}
@@ -30,11 +27,20 @@
 </template>
 
 <script setup lang="ts">
-defineProps<{
+import { computed } from 'vue'
+
+import UrlInput from '@/components/common/UrlInput.vue'
+
+const props = defineProps<{
   modelValue: string
 }>()
 
-defineEmits<{
+const emit = defineEmits<{
   'update:modelValue': [value: string]
 }>()
+
+const url = computed({
+  get: () => props.modelValue,
+  set: (value: string) => emit('update:modelValue', value)
+})
 </script>
diff --git a/src/platform/assets/composables/useModelTypes.ts b/src/platform/assets/composables/useModelTypes.ts
new file mode 100644
index 000000000..d180847f9
--- /dev/null
+++ b/src/platform/assets/composables/useModelTypes.ts
@@ -0,0 +1,94 @@
+import { ref } from 'vue'
+
+import { assetService } from '@/platform/assets/services/assetService'
+
+/**
+ * Format folder name to display name
+ * Converts "upscale_models" -> "Upscale Models"
+ * Converts "loras" -> "LoRAs"
+ */
+function formatDisplayName(folderName: string): string {
+  // Special cases for acronyms and proper nouns
+  const specialCases: Record<string, string> = {
+    loras: 'LoRAs',
+    ipadapter: 'IP-Adapter',
+    sams: 'SAMs',
+    clip_vision: 'CLIP Vision',
+    animatediff_motion_lora: 'AnimateDiff Motion LoRA',
+    animatediff_models: 'AnimateDiff Models',
+    vae: 'VAE',
+    sam2: 'SAM 2',
+    controlnet: 'ControlNet',
+    gligen: 'GLIGEN'
+  }
+
+  if (specialCases[folderName]) {
+    return specialCases[folderName]
+  }
+
+  return folderName
+    .split('_')
+    .map((word) => word.charAt(0).toUpperCase() + word.slice(1))
+    .join(' ')
+}
+
+interface ModelTypeOption {
+  name: string // Display name
+  value: string // Actual tag value
+}
+
+// Shared state across all instances
+const modelTypes = ref<ModelTypeOption[]>([])
+const isLoading = ref(false)
+const error = ref<string | null>(null)
+let fetchPromise: Promise<void> | null = null
+
+/**
+ * Composable for fetching and managing model types from the API
+ * Uses shared state to ensure data is only fetched once
+ */
+export function useModelTypes() {
+  /**
+   * Fetch model types from the API (only fetches once, subsequent calls reuse the same promise)
+   */
+  async function fetchModelTypes() {
+    // If already loaded, return immediately
+    if (modelTypes.value.length > 0) {
+      return
+    }
+
+    // If currently loading, return the existing promise
+    if (fetchPromise) {
+      return fetchPromise
+    }
+
+    isLoading.value = true
+    error.value = null
+
+    fetchPromise = (async () => {
+      try {
+        const response = await assetService.getModelTypes()
+        modelTypes.value = response.map((folder) => ({
+          name: formatDisplayName(folder.name),
+          value: folder.name
+        }))
+      } catch (err) {
+        error.value =
+          err instanceof Error ? err.message : 'Failed to fetch model types'
+        console.error('Failed to fetch model types:', err)
+      } finally {
+        isLoading.value = false
+        fetchPromise = null
+      }
+    })()
+
+    return fetchPromise
+  }
+
+  return {
+    modelTypes,
+    isLoading,
+    error,
+    fetchModelTypes
+  }
+}
diff --git a/src/platform/assets/services/assetService.ts b/src/platform/assets/services/assetService.ts
index 9f4613f42..621b41284 100644
--- a/src/platform/assets/services/assetService.ts
+++ b/src/platform/assets/services/assetService.ts
@@ -316,6 +316,24 @@ function createAssetService() {
     return await res.json()
   }
 
+  /**
+   * Gets available model types from the server
+   *
+   * @returns Promise<ModelFolder[]> - List of model types with their folder mappings
+   * @throws Error if request fails
+   */
+  async function getModelTypes(): Promise<ModelFolder[]> {
+    const res = await api.fetchApi('/experiment/models')
+
+    if (!res.ok) {
+      throw new Error(
+        `Failed to fetch model types: Server returned ${res.status}`
+      )
+    }
+
+    return await res.json()
+  }
+
   return {
     getAssetModelFolders,
     getAssetModels,
@@ -325,7 +343,8 @@ function createAssetService() {
     getAssetsByTag,
     deleteAsset,
     getAssetMetadata,
-    uploadAssetFromUrl
+    uploadAssetFromUrl,
+    getModelTypes
   }
 }