app.py

app.py

@@ -16,7 +16,7 @@ import os
 import time
 import numpy as np
 
-from utils.utils import load_cn_model, load_cn_config, load_tagger_model, resize_image_aspect_ratio, base_generation
+from utils.utils import load_cn_model, load_cn_config, load_tagger_model, load_lora_model, resize_image_aspect_ratio, base_generation
 from utils.prompt_analysis import PromptAnalysis
 
 path = os.getcwd()
@@ -24,10 +24,13 @@ cn_dir = f"{path}/controlnet"
 os.makedirs(cn_dir)
 tagger_dir = f"{path}/tagger"
 os.mkdir(tagger_dir)
+lora_dir = f"{path}/lora"
+os.mkdir(lora_dir)
 
 load_cn_model(cn_dir)
 load_cn_config(cn_dir)
 load_tagger_model(tagger_dir)
+load_lora_model(lora_dir)
 
 IS_SPACES_ZERO = os.environ.get("SPACES_ZERO_GPU", "0") == "1"
 IS_SPACE = os.environ.get("SPACE_ID", None) is not None
@@ -49,11 +52,16 @@ pipe = StableDiffusionXLControlNetImg2ImgPipeline.from_pretrained(
     model,
     controlnet=controlnet,
     torch_dtype=dtype,
-    variant="fp16",
     use_safetensors=True,
     scheduler=scheduler,
 )
 
+pipe.load_lora_weights(
+    lora_dir, 
+    weight_name="sdxl_BWLine.safetensors"
+)
+
+
 compel = Compel(
     tokenizer=[pipe.tokenizer, pipe.tokenizer_2],
     text_encoder=[pipe.text_encoder, pipe.text_encoder_2],

requirements.txt

@@ -19,6 +19,6 @@ hidiffusion==0.1.8
 spaces
 torch==2.2
 controlnet-aux==0.0.9
-onnx
-onnxruntime
-mediapipe
+onnx==1.16.1
+onnxruntime==1.18.0
+mediapipe==0.10.14

utils/utils.py

@@ -61,6 +61,28 @@ def load_tagger_model(model_dir):
             print(f'{file} already exists.')
 
 
+def load_lora_model(model_dir):
+  folder = model_dir
+  file_name = 'sdxl_BWLine.safetensors'
+  url = "https://huggingface.co/tori29umai/lineart/resolve/main/sdxl_BWLine.safetensors"
+
+  file_path = os.path.join(folder, file_name)
+  if not os.path.exists(file_path):
+    response = requests.get(url, stream=True)
+
+    total_size = int(response.headers.get('content-length', 0))
+    with open(file_path, 'wb') as f, tqdm(
+            desc=file_name,
+            total=total_size,
+            unit='iB',
+            unit_scale=True,
+            unit_divisor=1024,
+        ) as bar:
+        for data in response.iter_content(chunk_size=1024):
+            size = f.write(data)
+            bar.update(size)            
+
+
 def resize_image_aspect_ratio(image):
     # 元の画像サイズを取得
     original_width, original_height = image.size