Update README.md
Browse files
README.md
CHANGED
|
@@ -56,7 +56,9 @@ Once you have your environment set up in Google Colab (or any Python environment
|
|
| 56 |
from diffusers import DiffusionPipeline
|
| 57 |
|
| 58 |
# Load the model from HuggingFace
|
| 59 |
-
|
|
|
|
|
|
|
| 60 |
|
| 61 |
# Provide your text prompt
|
| 62 |
prompt = "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k"
|
|
|
|
| 56 |
from diffusers import DiffusionPipeline
|
| 57 |
|
| 58 |
# Load the model from HuggingFace
|
| 59 |
+
device = "cuda" if torch.cuda.is_available() else "cpu"
|
| 60 |
+
pipe = DiffusionPipeline.from_pretrained("ImageInception/ArtifyAI-v1.1, torch_dtype=torch.float16 if device == "cuda" else torch.float32)
|
| 61 |
+
pipe = pipe.to(device)
|
| 62 |
|
| 63 |
# Provide your text prompt
|
| 64 |
prompt = "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k"
|