SandaAbhishekSagar
commited on
Commit
·
ab3ad5b
1
Parent(s):
687bf24
revamped code
Browse files- image_generator.py +1 -4
image_generator.py
CHANGED
|
@@ -36,10 +36,7 @@ import torch
|
|
| 36 |
|
| 37 |
# Preload the model globally
|
| 38 |
device = "cuda" if torch.cuda.is_available() else "cpu"
|
| 39 |
-
|
| 40 |
-
"runwayml/stable-diffusion-v1-5"
|
| 41 |
-
)
|
| 42 |
-
pipeline = pipeline.to(device)
|
| 43 |
|
| 44 |
def generate_image(prompt):
|
| 45 |
"""Generate an image based on the input prompt."""
|
|
|
|
| 36 |
|
| 37 |
# Preload the model globally
|
| 38 |
device = "cuda" if torch.cuda.is_available() else "cpu"
|
| 39 |
+
pipeline = StableDiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5").to(device)
|
|
|
|
|
|
|
|
|
|
| 40 |
|
| 41 |
def generate_image(prompt):
|
| 42 |
"""Generate an image based on the input prompt."""
|