Update glide_text2im/glide_util.py

glide_text2im/glide_util.py

@@ -44,7 +44,7 @@ def sample(
     uncond_ref = th.ones_like(cond_ref) 
     
     model_kwargs = {}
-    model_kwargs['ref'] =  th.cat([cond_ref, uncond_ref], 0).cuda()
+    model_kwargs['ref'] =  th.cat([cond_ref, uncond_ref], 0).to(device)
 
     def cfg_model_fn(x_t, ts, **kwargs):
         half = x_t[: len(x_t) // 2]
@@ -60,7 +60,7 @@ def sample(
 
 
     if upsample_enabled:
-        model_kwargs['low_res'] = prompt['low_res'].cuda()
+        model_kwargs['low_res'] = prompt['low_res'].to(device)
         noise = th.randn((batch_size, 3, side_y, side_x), device=device) * upsample_temp
         model_fn = glide_model # just use the base model, no need for CFG.
         model_kwargs['ref'] =  model_kwargs['ref'][:batch_size]