Update README.md
Browse files
README.md
CHANGED
|
@@ -92,7 +92,7 @@ if torch.cuda.is_available():
|
|
| 92 |
elif torch.backends.mps.is_available():
|
| 93 |
# Apple Silicon
|
| 94 |
device = "mps"
|
| 95 |
-
dtype = torch.
|
| 96 |
else:
|
| 97 |
# CPU fallback (functional but very slow for this model)
|
| 98 |
device = "cpu"
|
|
@@ -104,6 +104,8 @@ pipe = ZImagePipeline.from_pretrained(
|
|
| 104 |
low_cpu_mem_usage=True,
|
| 105 |
)
|
| 106 |
|
|
|
|
|
|
|
| 107 |
if device == "cuda":
|
| 108 |
pipe.enable_model_cpu_offload()
|
| 109 |
else:
|
|
|
|
| 92 |
elif torch.backends.mps.is_available():
|
| 93 |
# Apple Silicon
|
| 94 |
device = "mps"
|
| 95 |
+
dtype = torch.bfloat16
|
| 96 |
else:
|
| 97 |
# CPU fallback (functional but very slow for this model)
|
| 98 |
device = "cpu"
|
|
|
|
| 104 |
low_cpu_mem_usage=True,
|
| 105 |
)
|
| 106 |
|
| 107 |
+
pipe.enable_attention_slicing()
|
| 108 |
+
|
| 109 |
if device == "cuda":
|
| 110 |
pipe.enable_model_cpu_offload()
|
| 111 |
else:
|