change stable2 model and how stablexl loads, add reqs

pull/11/head
Konstantine Tsafatinos 2023-07-19 17:25:28 -04:00
parent 08da0681cd
commit ff0114d341
3 changed files with 5 additions and 6 deletions

View File

@ -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

2
skynet/constants.py 100644 → 100755
View File

@ -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'},

6
skynet/utils.py 100644 → 100755
View File

@ -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