tricktreat commited on
Commit
bbe4314
1 Parent(s): 7208b68

Update models_server.py

Browse files
Files changed (1) hide show
  1. models_server.py +4 -4
models_server.py CHANGED
@@ -170,8 +170,8 @@ def load_pipes(local_deployment):
170
  # "model": DiffusionPipeline.from_pretrained(f"{local_models}lambdalabs/sd-image-variations-diffusers"), #torch_dtype=torch.float16
171
  # "device": "cuda:0"
172
  # },
173
- "bdsqlsz/stable-diffusion-v1-5": {
174
- "model": DiffusionPipeline.from_pretrained(f"{local_models}bdsqlsz/stable-diffusion-v1-5"),
175
  "device": "cuda:0"
176
  },
177
  # "superb/wav2vec2-base-superb-ks": {
@@ -268,7 +268,7 @@ def load_pipes(local_deployment):
268
 
269
  controlnet = ControlNetModel.from_pretrained(f"{local_models}lllyasviel/sd-controlnet-canny", torch_dtype=torch.float16)
270
  controlnetpipe = StableDiffusionControlNetPipeline.from_pretrained(
271
- f"{local_models}bdsqlsz/stable-diffusion-v1-5", controlnet=controlnet, torch_dtype=torch.float16
272
  )
273
 
274
 
@@ -463,7 +463,7 @@ def models(model_id, data):
463
  result = {"generated text": generated_text}
464
 
465
  # text to image
466
- if model_id == "bdsqlsz/stable-diffusion-v1-5":
467
  file_name = str(uuid.uuid4())[:4]
468
  text = data["text"]
469
  out = pipe(prompt=text)
 
170
  # "model": DiffusionPipeline.from_pretrained(f"{local_models}lambdalabs/sd-image-variations-diffusers"), #torch_dtype=torch.float16
171
  # "device": "cuda:0"
172
  # },
173
+ "CompVis/stable-diffusion-v1-4": {
174
+ "model": DiffusionPipeline.from_pretrained(f"{local_models}CompVis/stable-diffusion-v1-4"),
175
  "device": "cuda:0"
176
  },
177
  # "superb/wav2vec2-base-superb-ks": {
 
268
 
269
  controlnet = ControlNetModel.from_pretrained(f"{local_models}lllyasviel/sd-controlnet-canny", torch_dtype=torch.float16)
270
  controlnetpipe = StableDiffusionControlNetPipeline.from_pretrained(
271
+ f"{local_models}CompVis/stable-diffusion-v1-4", controlnet=controlnet, torch_dtype=torch.float16
272
  )
273
 
274
 
 
463
  result = {"generated text": generated_text}
464
 
465
  # text to image
466
+ if model_id == "CompVis/stable-diffusion-v1-4":
467
  file_name = str(uuid.uuid4())[:4]
468
  text = data["text"]
469
  out = pipe(prompt=text)