RohitGandikota commited on
Commit
bc70d86
·
1 Parent(s): 16e9d06

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +3 -0
app.py CHANGED
@@ -5,6 +5,8 @@ from StableDiffuser import StableDiffuser
5
  from tqdm import tqdm
6
  from train import train
7
 
 
 
8
  model_map = {'Van Gogh' : 'models/vangogh.pt',
9
  'Pablo Picasso': 'models/pablopicasso.pt',
10
  'Car' : 'models/car.pt',
@@ -22,6 +24,7 @@ class Demo:
22
  def __init__(self) -> None:
23
 
24
  self.training = False
 
25
  print(f'Available GPUs: {torch.cuda.device_count()}')
26
  self.diffuser = StableDiffuser(scheduler='DDIM', seed=42).to('cuda:1').eval().half()
27
 
 
5
  from tqdm import tqdm
6
  from train import train
7
 
8
+ !nvidia-smi
9
+
10
  model_map = {'Van Gogh' : 'models/vangogh.pt',
11
  'Pablo Picasso': 'models/pablopicasso.pt',
12
  'Car' : 'models/car.pt',
 
24
  def __init__(self) -> None:
25
 
26
  self.training = False
27
+
28
  print(f'Available GPUs: {torch.cuda.device_count()}')
29
  self.diffuser = StableDiffuser(scheduler='DDIM', seed=42).to('cuda:1').eval().half()
30