From ff0114d3414834ace6896c7117d7b88cc9f7cef3 Mon Sep 17 00:00:00 2001
From: Konstantine Tsafatinos <kosta.tsaf@gmail.com>
Date: Wed, 19 Jul 2023 17:25:28 -0400
Subject: [PATCH] change stable2 model and how stablexl loads, add reqs

---
 requirements.cuda.0.txt | 3 ++-
 skynet/constants.py     | 2 +-
 skynet/utils.py         | 6 ++----
 3 files changed, 5 insertions(+), 6 deletions(-)
 mode change 100644 => 100755 skynet/constants.py
 mode change 100644 => 100755 skynet/utils.py

diff --git a/requirements.cuda.0.txt b/requirements.cuda.0.txt
index e31de88..f796537 100644
--- a/requirements.cuda.0.txt
+++ b/requirements.cuda.0.txt
@@ -3,6 +3,7 @@ triton
 accelerate
 transformers
 huggingface_hub
-diffusers[torch]
+diffusers[torch]>=0.18.0
+invisible-watermark
 torch==1.13.0+cu117
 --extra-index-url https://download.pytorch.org/whl/cu117
diff --git a/skynet/constants.py b/skynet/constants.py
old mode 100644
new mode 100755
index 1bae4b1..2bb6d3b
--- a/skynet/constants.py
+++ b/skynet/constants.py
@@ -7,7 +7,7 @@ DOCKER_RUNTIME_CUDA = 'skynet:runtime-cuda'
 MODELS = {
     'prompthero/openjourney':                   { 'short': 'midj'},
     'runwayml/stable-diffusion-v1-5':           { 'short': 'stable'},
-    'stabilityai/stable-diffusion-2-1':         { 'short': 'stable2'},
+    'stabilityai/stable-diffusion-2-1-base':    { 'short': 'stable2'},
     'snowkidy/stable-diffusion-xl-base-0.9':    { 'short': 'stablexl'},
     'Linaqruf/anything-v3.0':                   { 'short': 'hdanime'},
     'hakurei/waifu-diffusion':                  { 'short': 'waifu'},
diff --git a/skynet/utils.py b/skynet/utils.py
old mode 100644
new mode 100755
index ad13a15..79932a5
--- a/skynet/utils.py
+++ b/skynet/utils.py
@@ -77,15 +77,13 @@ def pipeline_for(model: str, mem_fraction: float = 1.0, image=False) -> Diffusio
         'safety_checker': None
     }
 
-    if model == 'snowkidy/stable-diffusion-xl-base-0.9':
-        # TODO: figure out what this does
-        params['addition_embed_type'] = { 'text_time': None }
-
     if model == 'runwayml/stable-diffusion-v1-5':
         params['revision'] = 'fp16'
 
     if image:
         pipe_class = StableDiffusionImg2ImgPipeline
+    elif model == 'snowkidy/stable-diffusion-xl-base-0.9':
+        pipe_class = DiffusionPipeline
     else:
         pipe_class = StableDiffusionPipeline