sdart commited on
Commit
4b90f4d
1 Parent(s): 3206fc4

switch to gpu

Browse files
Files changed (2) hide show
  1. model.py +9 -0
  2. requirements.txt +2 -1
model.py CHANGED
@@ -7,8 +7,17 @@ import torch
7
  import PIL.Image
8
  import numpy as np
9
 
 
 
 
 
 
 
 
10
  device = "cuda"
11
 
 
 
12
  class Model:
13
  def __init__(self):
14
  modelID = "runwayml/stable-diffusion-v1-5"
 
7
  import PIL.Image
8
  import numpy as np
9
 
10
+ # Check environment
11
+ print(f"Is CUDA available: {torch.cuda.is_available()}")
12
+ # True
13
+ print(f"CUDA device: {torch.cuda.get_device_name(torch.cuda.current_device())}")
14
+ # Tesla T4
15
+
16
+
17
  device = "cuda"
18
 
19
+
20
+
21
  class Model:
22
  def __init__(self):
23
  modelID = "runwayml/stable-diffusion-v1-5"
requirements.txt CHANGED
@@ -1,8 +1,9 @@
 
 
1
  accelerate==0.18.0
2
  diffusers==0.16.0
3
  gradio==3.30.0
4
  safetensors==0.3.0
5
- torch==2.0.0
6
  torchvision==0.15.1
7
  transformers==4.28.1
8
  xformers==0.0.18
 
1
+ --extra-index-url https://download.pytorch.org/whl/cu113
2
+ torch
3
  accelerate==0.18.0
4
  diffusers==0.16.0
5
  gradio==3.30.0
6
  safetensors==0.3.0
 
7
  torchvision==0.15.1
8
  transformers==4.28.1
9
  xformers==0.0.18