Switch to SDXL and add PEFT

inference.py

@@ -43,8 +43,8 @@ class timer:
         print(f"{self.method} took {str(round(end - self.start, 2))}s")
 
 
-def load_models(model_id="Lykon/dreamshaper-7"):
-    from diffusers import AutoPipelineForImage2Image, LCMScheduler
+def load_models(model_id="stabilityai/stable-diffusion-xl-base-1.0"):
+    from diffusers import UNet2DConditionModel, AutoPipelineForImage2Image, LCMScheduler
     from diffusers.utils import load_image
 
     if not is_mac:
@@ -52,11 +52,14 @@ def load_models(model_id="Lykon/dreamshaper-7"):
 
     use_fp16 = should_use_fp16()
 
-    lcm_lora_id = "latent-consistency/lcm-lora-sdv1-5"
+    lora_id = "artificialguybr/LogoRedmond-LogoLoraForSDXL-V2"
+
+    unet = UNet2DConditionModel.from_pretrained("latent-consistency/lcm-sdxl", torch_dtype=torch.float16, variant="fp16")
 
     if use_fp16:
         pipe = AutoPipelineForImage2Image.from_pretrained(
             model_id,
+            unet=unet,
             cache_dir=cache_path,
             torch_dtype=torch.float16,
             variant="fp16",
@@ -65,13 +68,14 @@ def load_models(model_id="Lykon/dreamshaper-7"):
     else:
         pipe = AutoPipelineForImage2Image.from_pretrained(
             model_id,
+            unet=unet,
             cache_dir=cache_path,
             safety_checker=None
         )
 
     pipe.scheduler = LCMScheduler.from_config(pipe.scheduler.config)
 
-    pipe.load_lora_weights(lcm_lora_id)
+    pipe.load_lora_weights(lora_id)
     pipe.fuse_lora()
 
     device = "mps" if is_mac else "cuda"

requirements.txt

@@ -3,4 +3,5 @@ flask
 torch
 transformers
 accelerate
+peft
 pillow