From a51e3d8cc591f59b795783d68236e8cda10a900d Mon Sep 17 00:00:00 2001
From: Pedro Cuenca <pedro@huggingface.co>
Date: Sat, 14 Jan 2023 19:10:44 +0100
Subject: [PATCH] Populate model picker.

---
 Diffusion-macOS/PromptView.swift | 14 +++++++-------
 Diffusion/ModelInfo.swift        | 15 +++++++++++----
 2 files changed, 18 insertions(+), 11 deletions(-)

diff --git a/Diffusion-macOS/PromptView.swift b/Diffusion-macOS/PromptView.swift
index 55960d8..0de2107 100644
--- a/Diffusion-macOS/PromptView.swift
+++ b/Diffusion-macOS/PromptView.swift
@@ -7,14 +7,14 @@
 
 import SwiftUI
 import CompactSlider
-//import StableDiffusion
 
 
 struct PromptView: View {
     
-    var models = ["hf-default-model"]
+    static let models = ModelInfo.MODELS
+    static let modelNames = models.map { $0.modelVersion }
     
-    @State private var model = "hf-default-model"
+    @State private var model = ModelInfo.v2Base.modelVersion
     @State private var positivePrompt = ""
     @State private var negativePrompt = ""
     @State private var steps = 50.0
@@ -33,10 +33,10 @@ struct PromptView: View {
                 Group {
                     DisclosureGroup {
                         Picker("", selection: $model) {
-                                        ForEach(models, id: \.self) {
-                                            Text($0)
-                                        }
-                                    }
+                            ForEach(Self.modelNames, id: \.self) {
+                                Text($0)
+                            }
+                        }
                     } label: {
                         Label("Model", systemImage: "cpu").foregroundColor(.secondary)
                     }
diff --git a/Diffusion/ModelInfo.swift b/Diffusion/ModelInfo.swift
index 9365bb2..0171892 100644
--- a/Diffusion/ModelInfo.swift
+++ b/Diffusion/ModelInfo.swift
@@ -60,22 +60,29 @@ extension ModelInfo {
     // TODO: repo does not exist yet
     static let v14Base = ModelInfo(
         modelId: "pcuenq/coreml-stable-diffusion-v1-4",
-        modelVersion: "1.4"
+        modelVersion: "CompVis/stable-diffusion-v1-4"
     )
 
     static let v15Base = ModelInfo(
         modelId: "pcuenq/coreml-stable-diffusion-v1-5",
-        modelVersion: "1.5"
+        modelVersion: "runwayml/stable-diffusion-v1-5"
     )
     
     static let v2Base = ModelInfo(
         modelId: "pcuenq/coreml-stable-diffusion-2-base",
-        modelVersion: "2-base"
+        modelVersion: "stabilityai/stable-diffusion-2-base"
     )
 
     static let v21Base = ModelInfo(
         modelId: "pcuenq/coreml-stable-diffusion-2-1-base",
-        modelVersion: "2.1-base",
+        modelVersion: "stabilityai/stable-diffusion-2-1-base",
         supportsEncoder: true
     )
+    
+    static let MODELS = [
+        ModelInfo.v14Base,
+        ModelInfo.v15Base,
+        ModelInfo.v2Base,
+        ModelInfo.v21Base
+    ]
 }