eliphatfs commited on
Commit ·
1caa49c
1
Parent(s): 71bd9a3
Use FP16 for SD.
Browse files
openshape/demo/sd_pc2img.py
CHANGED
|
@@ -13,8 +13,10 @@ class Wrapper(transformers.modeling_utils.PreTrainedModel):
|
|
| 13 |
return rst.ObjectProxy(image_embeds=x)
|
| 14 |
|
| 15 |
|
|
|
|
| 16 |
pipe = StableUnCLIPImg2ImgPipeline.from_pretrained(
|
| 17 |
"diffusers/stable-diffusion-2-1-unclip-i2i-l",
|
|
|
|
| 18 |
image_encoder = Wrapper()
|
| 19 |
)
|
| 20 |
if torch.cuda.is_available():
|
|
|
|
| 13 |
return rst.ObjectProxy(image_embeds=x)
|
| 14 |
|
| 15 |
|
| 16 |
+
half = torch.float16 if torch.cuda.is_available() else torch.bfloat16
|
| 17 |
pipe = StableUnCLIPImg2ImgPipeline.from_pretrained(
|
| 18 |
"diffusers/stable-diffusion-2-1-unclip-i2i-l",
|
| 19 |
+
torch_dtype=half, variation="fp16",
|
| 20 |
image_encoder = Wrapper()
|
| 21 |
)
|
| 22 |
if torch.cuda.is_available():
|