Update README.md
Browse files
README.md
CHANGED
|
@@ -55,7 +55,10 @@ Now that you have everything installed, here's how you can load the model and st
|
|
| 55 |
from diffusers import DiffusionPipeline
|
| 56 |
|
| 57 |
# Load the model from HuggingFace
|
| 58 |
-
|
|
|
|
|
|
|
|
|
|
| 59 |
|
| 60 |
# Provide your text prompt
|
| 61 |
prompt = "A sunset over a mountain range with orange and purple hues in the sky."
|
|
|
|
| 55 |
from diffusers import DiffusionPipeline
|
| 56 |
|
| 57 |
# Load the model from HuggingFace
|
| 58 |
+
|
| 59 |
+
device = "cuda" if torch.cuda.is_available() else "cpu"
|
| 60 |
+
pipe = DiffusionPipeline.from_pretrained("ImageInception/ArtifyAI-v1.0", torch_dtype=torch.float16 if device == "cuda" else torch.float32)
|
| 61 |
+
pipe = pipe.to(device)
|
| 62 |
|
| 63 |
# Provide your text prompt
|
| 64 |
prompt = "A sunset over a mountain range with orange and purple hues in the sky."
|