From d3c145f219a19e134f7858d0b1c0f813fae4d8f0 Mon Sep 17 00:00:00 2001
From: Pedro Cuenca <pedro@huggingface.co>
Date: Tue, 24 Jan 2023 14:58:23 +0100
Subject: [PATCH] Add v1-4 and ofa-sys v0 models.

---
 Diffusion/ModelInfo.swift | 12 +++++++++---
 1 file changed, 9 insertions(+), 3 deletions(-)

diff --git a/Diffusion/ModelInfo.swift b/Diffusion/ModelInfo.swift
index 88b68c1..3ffeb26 100644
--- a/Diffusion/ModelInfo.swift
+++ b/Diffusion/ModelInfo.swift
@@ -59,7 +59,7 @@ extension ModelInfo {
 extension ModelInfo {
     // TODO: repo does not exist yet
     static let v14Base = ModelInfo(
-        modelId: "pcuenq/coreml-stable-diffusion-v1-4",
+        modelId: "pcuenq/coreml-stable-diffusion-1-4",
         modelVersion: "CompVis/stable-diffusion-v1-4"
     )
 
@@ -79,11 +79,17 @@ extension ModelInfo {
         supportsEncoder: true
     )
     
+    static let ofaSmall = ModelInfo(
+        modelId: "pcuenq/coreml-small-stable-diffusion-v0",
+        modelVersion: "OFA-Sys/small-stable-diffusion-v0"
+    )
+
     static let MODELS = [
-//        ModelInfo.v14Base,
+        ModelInfo.v14Base,
         ModelInfo.v15Base,
         ModelInfo.v2Base,
-        ModelInfo.v21Base
+        ModelInfo.v21Base,
+        ModelInfo.ofaSmall
     ]
     
     static func from(modelVersion: String) -> ModelInfo? {