LORA support

app.py

@@ -6,7 +6,11 @@ import gradio as gr
 import requests
 import spaces
 import torch
-from diffusers import AutoencoderKL, StableDiffusionImg2ImgPipeline
+from diffusers import (
+    AutoencoderKL,
+    AutoPipelineForImage2Image,
+    StableDiffusionImg2ImgPipeline,
+)
 from diffusers.pipelines.stable_diffusion.convert_from_ckpt import (
     download_from_original_stable_diffusion_ckpt,
 )
@@ -99,16 +103,38 @@ if os.path.exists(get_file_name("VAE")):
         use_safetensors=True,
     )
 
+model_type = model["type"]
 
-logger.debug(f"Loading pipeline")
+if model_type == "Checkpoint":
+    logger.debug(f"Loading pipeline for checkpoint")
 
-pipe = download_from_original_stable_diffusion_ckpt(
-    checkpoint_path_or_dict=get_file_name("Model"),
-    from_safetensors=True,
-    pipeline_class=StableDiffusionImg2ImgPipeline,
-    load_safety_checker=False,
-    **pipe_args,
-)
+    pipe = download_from_original_stable_diffusion_ckpt(
+        checkpoint_path_or_dict=get_file_name("Model"),
+        from_safetensors=True,
+        pipeline_class=StableDiffusionImg2ImgPipeline,
+        load_safety_checker=False,
+        **pipe_args,
+    )
+elif model_type == "LORA":
+    logger.debug(f"Loading pipeline for LORA")
+
+    base_model = model_version["baseModel"]
+
+    if base_model == "SD 1.5":
+        pipe = AutoPipelineForImage2Image.from_pretrained(
+            "stable-diffusion-v1-5/stable-diffusion-v1-5",
+            safety_checker=None,
+            requires_safety_checker=False,
+            torch_dtype=torch.float16,
+            use_safetensors=True,
+            variant="fp16",
+        )
+    else:
+        raise ValueError(f"Unsupported base model: {base_model}")
+
+    pipe.load_lora_weights(get_file_name("Model"), adapter_name=slugify(model_name))
+else:
+    raise ValueError(f"Unsupported model type: {model_type}")
 
 pipe = pipe.to("cuda")