mirror of https://github.com/skygpu/skynet.git
Fix stablexl pipeline
parent
4082adf184
commit
440bb015cd
|
@ -8,7 +8,7 @@ MODELS = {
|
||||||
'prompthero/openjourney': { 'short': 'midj'},
|
'prompthero/openjourney': { 'short': 'midj'},
|
||||||
'runwayml/stable-diffusion-v1-5': { 'short': 'stable'},
|
'runwayml/stable-diffusion-v1-5': { 'short': 'stable'},
|
||||||
'stabilityai/stable-diffusion-2-1-base': { 'short': 'stable2'},
|
'stabilityai/stable-diffusion-2-1-base': { 'short': 'stable2'},
|
||||||
'snowkidy/stable-diffusion-xl-base-1.0': { 'short': 'stablexl'},
|
'stabilityai/stable-diffusion-xl-base-1.0': { 'short': 'stablexl'},
|
||||||
'Linaqruf/anything-v3.0': { 'short': 'hdanime'},
|
'Linaqruf/anything-v3.0': { 'short': 'hdanime'},
|
||||||
'hakurei/waifu-diffusion': { 'short': 'waifu'},
|
'hakurei/waifu-diffusion': { 'short': 'waifu'},
|
||||||
'nitrosocke/Ghibli-Diffusion': { 'short': 'ghibli'},
|
'nitrosocke/Ghibli-Diffusion': { 'short': 'ghibli'},
|
||||||
|
|
|
@ -15,6 +15,8 @@ from PIL import Image
|
||||||
from basicsr.archs.rrdbnet_arch import RRDBNet
|
from basicsr.archs.rrdbnet_arch import RRDBNet
|
||||||
from diffusers import (
|
from diffusers import (
|
||||||
DiffusionPipeline,
|
DiffusionPipeline,
|
||||||
|
StableDiffusionXLPipeline,
|
||||||
|
StableDiffusionXLImg2ImgPipeline,
|
||||||
StableDiffusionPipeline,
|
StableDiffusionPipeline,
|
||||||
StableDiffusionImg2ImgPipeline,
|
StableDiffusionImg2ImgPipeline,
|
||||||
EulerAncestralDiscreteScheduler
|
EulerAncestralDiscreteScheduler
|
||||||
|
@ -80,10 +82,14 @@ def pipeline_for(model: str, mem_fraction: float = 1.0, image=False) -> Diffusio
|
||||||
if model == 'runwayml/stable-diffusion-v1-5':
|
if model == 'runwayml/stable-diffusion-v1-5':
|
||||||
params['revision'] = 'fp16'
|
params['revision'] = 'fp16'
|
||||||
|
|
||||||
|
if model == 'stabilityai/stable-diffusion-xl-base-1.0':
|
||||||
|
if image:
|
||||||
|
pipe_class = StableDiffusionXLImg2ImgPipeline
|
||||||
|
else:
|
||||||
|
pipe_class = StableDiffusionXLPipeline
|
||||||
|
else:
|
||||||
if image:
|
if image:
|
||||||
pipe_class = StableDiffusionImg2ImgPipeline
|
pipe_class = StableDiffusionImg2ImgPipeline
|
||||||
elif model == 'snowkidy/stable-diffusion-xl-base-0.9':
|
|
||||||
pipe_class = DiffusionPipeline
|
|
||||||
else:
|
else:
|
||||||
pipe_class = StableDiffusionPipeline
|
pipe_class = StableDiffusionPipeline
|
||||||
|
|
||||||
|
|
Loading…
Reference in New Issue