From 5f2d07917d1c0c53125bea973d59b2ac466bc8ad Mon Sep 17 00:00:00 2001
From: Sergey Borisov <stalkek7779@yandex.ru>
Date: Thu, 15 Jun 2023 21:30:15 +0300
Subject: [PATCH] Fix lora import, fix sd2 config, fix list models api

---
 invokeai/app/api/routers/models.py                           | 2 +-
 invokeai/backend/model_management/model_manager.py           | 1 -
 invokeai/backend/model_management/models/lora.py             | 1 +
 invokeai/backend/model_management/models/stable_diffusion.py | 2 ++
 4 files changed, 4 insertions(+), 2 deletions(-)

diff --git a/invokeai/app/api/routers/models.py b/invokeai/app/api/routers/models.py
index 4fc40a1eec7..f510279f183 100644
--- a/invokeai/app/api/routers/models.py
+++ b/invokeai/app/api/routers/models.py
@@ -80,7 +80,7 @@ async def list_models(
     ),
 ) -> ModelsList:
     """Gets a list of models"""
-    models_raw = ApiDependencies.invoker.services.model_manager.list_models(model_type)
+    models_raw = ApiDependencies.invoker.services.model_manager.list_models(base_model, model_type)
     models = parse_obj_as(ModelsList, { "models": models_raw })
     return models
 
diff --git a/invokeai/backend/model_management/model_manager.py b/invokeai/backend/model_management/model_manager.py
index fd0141c3f77..cbb319c6ea6 100644
--- a/invokeai/backend/model_management/model_manager.py
+++ b/invokeai/backend/model_management/model_manager.py
@@ -512,7 +512,6 @@ def list_models(
         named 'model-name', and model_manager.config to get the full OmegaConf
         object derived from models.yaml
         """
-        assert not(model_type is not None and base_model is None), "model_type must be provided with base_model"
 
         models = dict()
         for model_key in sorted(self.models, key=str.casefold):
diff --git a/invokeai/backend/model_management/models/lora.py b/invokeai/backend/model_management/models/lora.py
index c69677fd0c4..bcf3224ece1 100644
--- a/invokeai/backend/model_management/models/lora.py
+++ b/invokeai/backend/model_management/models/lora.py
@@ -1,3 +1,4 @@
+import os
 import torch
 from typing import Optional, Union, Literal
 from .base import (
diff --git a/invokeai/backend/model_management/models/stable_diffusion.py b/invokeai/backend/model_management/models/stable_diffusion.py
index 38890a703cf..bd519c88c80 100644
--- a/invokeai/backend/model_management/models/stable_diffusion.py
+++ b/invokeai/backend/model_management/models/stable_diffusion.py
@@ -123,6 +123,7 @@ class StableDiffusion2Model(DiffusersModel):
     class DiffusersConfig(ModelConfigBase):
         format: Literal["diffusers"]
         vae: Optional[str] = Field(None)
+        variant: ModelVariantType
         prediction_type: SchedulerPredictionType
         upcast_attention: bool
 
@@ -130,6 +131,7 @@ class CheckpointConfig(ModelConfigBase):
         format: Literal["checkpoint"]
         vae: Optional[str] = Field(None)
         config: Optional[str] = Field(None)
+        variant: ModelVariantType
         prediction_type: SchedulerPredictionType
         upcast_attention: bool