SDXL base models and model versions

app.py

@@ -1,6 +1,6 @@
 import os
 import shutil
-from urllib.parse import urlparse
+from urllib.parse import parse_qs, urlparse
 
 import gradio as gr
 import requests
@@ -10,6 +10,7 @@ from diffusers import (
     AutoencoderKL,
     AutoPipelineForImage2Image,
     StableDiffusionImg2ImgPipeline,
+    StableDiffusionXLImg2ImgPipeline,
 )
 from diffusers.pipelines.stable_diffusion.convert_from_ckpt import (
     download_from_original_stable_diffusion_ckpt,
@@ -36,7 +37,16 @@ gpu_duration = int(os.environ.get("GPU_DURATION", 60))
 
 logger.debug(f"Loading model info for: {model_url}")
 
-model_id = int(urlparse(model_url).path.split("/")[2])
+model_url_parsed = urlparse(model_url)
+
+model_id = int(model_url_parsed.path.split("/")[2])
+
+model_version_id = parse_qs(model_url_parsed.query).get("modelVersionId")
+if model_version_id is not None:
+    model_version_id = int(model_version_id[0])
+
+    logger.debug(f"Model version id: {model_version_id}")
+
 r = requests.get(f"https://civitai.com/api/v1/models/{model_id}")
 try:
     r.raise_for_status()
@@ -49,7 +59,11 @@ model = r.json()
 
 logger.debug(f"Model info: {model}")
 
-model_version = model["modelVersions"][0]
+model_version = (
+    model["modelVersions"][0]
+    if model_version_id is None
+    else next(mv for mv in model["modelVersions"] if mv["id"] == model_version_id)
+)
 
 assert len(model_version["files"]) <= 2
 assert len({file["type"] for file in model_version["files"]}) == len(
@@ -92,26 +106,31 @@ for _ in thread_map(
 ):
     pass
 
-
-pipe_args = {}
-if os.path.exists(get_file_name("VAE")):
-    logger.debug(f"Loading VAE")
-
-    pipe_args["vae"] = AutoencoderKL.from_single_file(
-        get_file_name("VAE"),
-        torch_dtype=torch.float16,
-        use_safetensors=True,
-    )
-
 model_type = model["type"]
 
 if model_type == "Checkpoint":
     logger.debug(f"Loading pipeline for checkpoint")
 
+    pipe_args = {}
+    if os.path.exists(get_file_name("VAE")):
+        logger.debug(f"Loading VAE")
+
+        pipe_args["vae"] = AutoencoderKL.from_single_file(
+            get_file_name("VAE"),
+            torch_dtype=torch.float16,
+            use_safetensors=True,
+        )
+
+    base_model = model_version["baseModel"]
+    if base_model == "SD 1.5":
+        pipeline_class = StableDiffusionImg2ImgPipeline
+    elif base_model == "SDXL 1.0":
+        pipeline_class = StableDiffusionXLImg2ImgPipeline
+
     pipe = download_from_original_stable_diffusion_ckpt(
         checkpoint_path_or_dict=get_file_name("Model"),
         from_safetensors=True,
-        pipeline_class=StableDiffusionImg2ImgPipeline,
+        pipeline_class=pipeline_class,
         load_safety_checker=False,
         **pipe_args,
     )