robin-courant commited on
Commit
c395a63
1 Parent(s): c2d401a

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +2 -6
app.py CHANGED
@@ -109,10 +109,6 @@ def generate(
109
  sample_id = SAMPLE_IDS[LABEL_TO_IDS[sample_label]]
110
  seq_feat = diffuser.net.model.clip_sequential
111
 
112
- print(device)
113
- print(clip_model.device)
114
- print(diffuser.device)
115
-
116
  batch = get_batch(prompt, sample_id, clip_model, dataset, seq_feat, device)
117
  with torch.no_grad():
118
  out = diffuser.predict_step(batch, 0)
@@ -210,8 +206,8 @@ def launch_app(gen_fn: Callable):
210
  # ------------------------------------------------------------------------------------- #
211
 
212
  diffuser, clip_model, dataset, device = init("config")
213
- diffuser.to(device)
214
- clip_model.to(device)
215
 
216
  generate_sample = partial(
217
  generate,
 
109
  sample_id = SAMPLE_IDS[LABEL_TO_IDS[sample_label]]
110
  seq_feat = diffuser.net.model.clip_sequential
111
 
 
 
 
 
112
  batch = get_batch(prompt, sample_id, clip_model, dataset, seq_feat, device)
113
  with torch.no_grad():
114
  out = diffuser.predict_step(batch, 0)
 
206
  # ------------------------------------------------------------------------------------- #
207
 
208
  diffuser, clip_model, dataset, device = init("config")
209
+ diffuser.to("cuda")
210
+ clip_model.to("cuda")
211
 
212
  generate_sample = partial(
213
  generate,