Update glide_text2im/glide_util.py

glide_text2im/glide_util.py

@@ -1,6 +1,6 @@
 import os
 from typing import Tuple
-from . import dist_util
+#from . import dist_util
 import PIL
 import numpy as np
 import torch as th
@@ -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).to(dist_util.dev())
+    model_kwargs['ref'] =  th.cat([cond_ref, uncond_ref], 0).cuda()
 
     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'].to(dist_util.dev())
+        model_kwargs['low_res'] = prompt['low_res'].cuda()
         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]