Add support for controlnet & sdxl checkpoint conversion

invokeai/app/api_app.py

@@ -203,7 +203,10 @@ def find_port(port: int):
                 return find_port(port=port + 1)
             else:
                 return port
-
+
+    from invokeai.backend.install.check_root import check_invokeai_root
+    check_invokeai_root(app_config)  # note, may exit with an exception if root not set up
+
     port = find_port(app_config.port)
     if port != app_config.port:
         logger.warn(f"Port {app_config.port} in use, using port {port}")

invokeai/backend/install/check_root.py

@@ -0,0 +1,31 @@
+"""
+Check that the invokeai_root is correctly configured and exit if not.
+"""
+import sys
+from invokeai.app.services.config import (
+    InvokeAIAppConfig,
+)
+
+def check_invokeai_root(config: InvokeAIAppConfig):
+    try:
+        assert config.model_conf_path.exists()
+        assert config.db_path.exists()
+        assert config.models_path.exists()
+        for model in [
+                'CLIP-ViT-bigG-14-laion2B-39B-b160k',
+                'bert-base-uncased',
+                'clip-vit-large-patch14',
+                'sd-vae-ft-mse',
+                'stable-diffusion-2-clip',
+                'stable-diffusion-safety-checker']:
+            assert (config.models_path / f'core/convert/{model}').exists()
+    except:
+        print()
+        print('== STARTUP ABORTED ==')
+        print('** One or more necessary files is missing from your InvokeAI root directory **')
+        print('** Please rerun the configuration script to fix this problem. **')
+        print('** From the launcher, selection option [7]. **')
+        print('** From the command line, activate the virtual environment and run "invokeai-configure --yes --skip-sd-weights" **')
+        input('Press any key to continue...')
+        sys.exit(0)
+

invokeai/backend/install/invokeai_configure.py

@@ -32,6 +32,7 @@
 from tqdm import tqdm
 from transformers import (
     CLIPTextModel,
+    CLIPTextConfig,
     CLIPTokenizer,
     AutoFeatureExtractor,
     BertTokenizerFast,
@@ -55,6 +56,7 @@
 from invokeai.backend.install.legacy_arg_parsing import legacy_parser
 from invokeai.backend.install.model_install_backend import (
     hf_download_from_pretrained,
+    hf_download_with_resume,
     InstallSelections,
     ModelInstall,
 )
@@ -204,6 +206,15 @@ def download_conversion_models():
         pipeline = CLIPTextModel.from_pretrained(repo_id, subfolder="text_encoder", **kwargs)
         pipeline.save_pretrained(target_dir / 'stable-diffusion-2-clip' / 'text_encoder', safe_serialization=True)
 
+        # sd-xl - tokenizer_2
+        repo_id = "laion/CLIP-ViT-bigG-14-laion2B-39B-b160k"
+        _, model_name = repo_id.split('/')
+        pipeline = CLIPTokenizer.from_pretrained(repo_id, **kwargs)
+        pipeline.save_pretrained(target_dir / model_name, safe_serialization=True)
+
+        pipeline = CLIPTextConfig.from_pretrained(repo_id, **kwargs)
+        pipeline.save_pretrained(target_dir / model_name, safe_serialization=True)
+
         # VAE
         logger.info('Downloading stable diffusion VAE')
         vae = AutoencoderKL.from_pretrained('stabilityai/sd-vae-ft-mse', **kwargs)

invokeai/backend/install/model_install_backend.py

@@ -58,7 +58,15 @@
             SchedulerPredictionType.Epsilon: 'v2-inpainting-inference.yaml',
             SchedulerPredictionType.VPrediction: 'v2-inpainting-inference-v.yaml',
         }
-    }
+    },
+
+    BaseModelType.StableDiffusionXL: {
+        ModelVariantType.Normal: 'sd_xl_base.yaml',
+    },
+
+    BaseModelType.StableDiffusionXLRefiner: {
+        ModelVariantType.Normal: 'sd_xl_refiner.yaml',
+    },
 }
 
 @dataclass
@@ -329,6 +337,7 @@ def _make_attributes(self, path: Path, info: ModelProbeInfo)->dict:
             description = str(description),
             model_format = info.format,
             )
+        legacy_conf = None
         if info.model_type == ModelType.Main:
             attributes.update(dict(variant = info.variant_type,))
             if info.format=="checkpoint":
@@ -343,11 +352,17 @@ def _make_attributes(self, path: Path, info: ModelProbeInfo)->dict:
                 except KeyError:
                     legacy_conf = Path(self.config.legacy_conf_dir, 'v1-inference.yaml')  # best guess
 
-                attributes.update(
-                    dict(
-                        config = str(legacy_conf)
-                    )
+        if info.model_type == ModelType.ControlNet and info.format=="checkpoint":
+            possible_conf = path.with_suffix('.yaml')
+            if possible_conf.exists():
+                legacy_conf = str(self.relative_to_root(possible_conf))
+
+        if legacy_conf:
+            attributes.update(
+                dict(
+                    config = str(legacy_conf)
                 )
+            )
         return attributes
 
     def relative_to_root(self, path: Path)->Path: