Update test

test.py

@@ -1,12 +1,15 @@
 from app import generate_content
+from visual_anagrams.views import VIEW_MAP_NAMES
 
 if __name__ == "__main__":
-    print(generate_content(
-        "a painting of",
-        "vases",
-        "a sloth",
-        "Flip",
-        1,
-        0
-    ))
+    for view in VIEW_MAP_NAMES:
+        print(view)
+        print(generate_content(
+            "a painting of",
+            "vases",
+            "a sloth",
+            view,
+            1,
+            0
+        ))
     

visual_anagrams/samplers.py

@@ -42,7 +42,7 @@ def sample_stage_1(model,
         prompt_embeds.dtype,
         device,
         generator,
-    )
+    ).to('cuda')
     # ic(noisy_images.shape)
 
     for i, t in enumerate(tqdm(timesteps)):