John Alexander
commited on
Commit
•
4de4026
1
Parent(s):
8d25670
Update app.py
Browse files
app.py
CHANGED
@@ -49,6 +49,8 @@ if SAFETY_CHECKER:
|
|
49 |
# Function
|
50 |
@spaces.GPU(enable_queue=True)
|
51 |
def generate_image(prompt, ckpt):
|
|
|
|
|
52 |
|
53 |
checkpoint = checkpoints[ckpt][0]
|
54 |
num_inference_steps = checkpoints[ckpt][1]
|
|
|
49 |
# Function
|
50 |
@spaces.GPU(enable_queue=True)
|
51 |
def generate_image(prompt, ckpt):
|
52 |
+
if torch.cuda.is_available():
|
53 |
+
pipe = StableDiffusionXLPipeline.from_pretrained(base, torch_dtype=torch.float16, variant="fp16").to("cuda") # Moved inside
|
54 |
|
55 |
checkpoint = checkpoints[ckpt][0]
|
56 |
num_inference_steps = checkpoints[ckpt][1]
|