Merge pull request #14 from jmkim0309/hotfix/ts_cpu

Hotfix/ts cpu

Author: saehoon sam kim

demo/product_demo.py

@@ -89,7 +89,7 @@ def __init__(
 def default_parser():
     parser = argparse.ArgumentParser()
     parser.add_argument("--root-dir", type=str, default=None)
-    parser.add_argument("--max_bsz", type=int, default=4)
+    parser.add_argument("--max_bsz", type=int, default=1)
     parser.add_argument(
         "--progressive", type=str, default="loop", choices=("loop", "stage", "final")
     )

karlo/modules/diffusion/respace.py

@@ -104,8 +104,9 @@ def condition_score(self, cond_fn, *args, **kwargs):
 
     def _wrap_model(self, model):
         def wrapped(x, ts, **kwargs):
+            ts_cpu = ts.detach().to("cpu")
             return model(
-                x, self.timestep_map[ts].to(device=ts.device, dtype=ts.dtype), **kwargs
+                x, self.timestep_map[ts_cpu].to(device=ts.device, dtype=ts.dtype), **kwargs
             )
 
         return wrapped