fix manual seed

singnet · Sep 4, 2024 · 88aa836 · 88aa836
1 parent cb369ab
commit 88aa836
Showing 1 changed file with 12 additions and 12 deletions.
diff --git a/tests/pipe_test.py b/tests/pipe_test.py
@@ -54,7 +54,7 @@ def test_basic_txt2im(self):
         seed = 49045438434843
         params = dict(prompt="a cube  planet, cube-shaped, space photo, masterpiece",
                       negative_prompt="spherical",
-                      generator=torch.cuda.manual_seed(seed))
+                      generator=torch.Generator(pipe.pipe.device).manual_seed(seed))
         image = pipe.gen(params)
         image.save("cube_test.png")
 
@@ -125,15 +125,15 @@ def test_img2img_basic(self):
         seed = 49045438434843
         pipe.setup(im, strength=0.7, steps=5, guidance_scale=3.3)
         self.assertEqual(3.3, pipe.pipe_params['guidance_scale'])
-        image = pipe.gen(dict(prompt="cube planet cartoon style", generator=torch.cuda.manual_seed(seed)))
+        image = pipe.gen(dict(prompt="cube planet cartoon style", generator=torch.Generator(pipe.pipe.device).manual_seed(seed)))
         image.save('test_img2img_basic.png')
         pipe.setup(im, strength=0.7, steps=5, guidance_scale=7.6)
-        image1 = pipe.gen(dict(prompt="cube planet cartoon style", generator=torch.cuda.manual_seed(seed)))
+        image1 = pipe.gen(dict(prompt="cube planet cartoon style", generator=torch.Generator(pipe.pipe.device).manual_seed(seed)))
         diff = self.compute_diff(image1, image)
         # check that difference is large
         self.assertGreater(diff, 1000)
         pipe.setup(im, strength=0.7, steps=5, guidance_scale=3.3)
-        image2 = pipe.gen(dict(prompt="cube planet cartoon style", generator=torch.cuda.manual_seed(seed)))
+        image2 = pipe.gen(dict(prompt="cube planet cartoon style", generator=torch.Generator(pipe.pipe.device).manual_seed(seed)))
         diff = self.compute_diff(image2, image)
         # check that difference is small
         self.assertLess(diff, 1)
@@ -153,15 +153,15 @@ def test_maskedimg2img_basic(self):
                scheduler=scheduler, clip_skip=0, blur=blur, blur_compose=3, steps=5, guidance_scale=7.6)
         pipe.setup(**param_3_3)
         self.assertEqual(3.3, pipe.pipe_params['guidance_scale'])
-        image = pipe.gen(dict(prompt="cube planet cartoon style", generator=torch.cuda.manual_seed(seed)))
+        image = pipe.gen(dict(prompt="cube planet cartoon style", generator=torch.Generator(pipe.pipe.device).manual_seed(seed)))
         image.save('test_img2img_basic.png')
         pipe.setup(**param_7_6)
-        image1 = pipe.gen(dict(prompt="cube planet cartoon style", generator=torch.cuda.manual_seed(seed)))
+        image1 = pipe.gen(dict(prompt="cube planet cartoon style", generator=torch.Generator(pipe.pipe.device).manual_seed(seed)))
         diff = self.compute_diff(image1, image)
         # check that difference is large
         self.assertGreater(diff, 1000)
         pipe.setup(**param_3_3)
-        image2 = pipe.gen(dict(prompt="cube planet cartoon style", generator=torch.cuda.manual_seed(seed)))
+        image2 = pipe.gen(dict(prompt="cube planet cartoon style", generator=torch.Generator(pipe.pipe.device).manual_seed(seed)))
         diff = self.compute_diff(image2, image)
         # check that difference is small
         self.assertLess(diff, 1)
@@ -177,12 +177,12 @@ def test_lpw(self):
         seed = 49045438434843
         params = dict(prompt=prompt,
                       negative_prompt="spherical",
-                      generator=torch.cuda.manual_seed(seed))
+                      generator=torch.Generator(pipe.pipe.device).manual_seed(seed))
         image = pipe.gen(params)
         image.save("cube_test_lpw.png")
         params = dict(prompt=prompt + " , best quality, famous photo",
                 negative_prompt="spherical",
-                generator=torch.cuda.manual_seed(seed))
+                generator=torch.Generator(pipe.pipe.device).manual_seed(seed))
         image1 = pipe.gen(params)
         image.save("cube_test_lpw1.png")
         diff = self.compute_diff(image1, image)
@@ -200,12 +200,12 @@ def test_lpw_turned_off(self):
         seed = 49045438434843
         params = dict(prompt=prompt,
                       negative_prompt="spherical",
-                      generator=torch.cuda.manual_seed(seed))
+                      generator=torch.Generator(pipe.pipe.device).manual_seed(seed))
         image = pipe.gen(params)
         image.save("cube_test_no_lpw.png")
         params = dict(prompt=prompt + " , best quality, famous photo",
                 negative_prompt="spherical",
-                generator=torch.cuda.manual_seed(seed))
+                generator=torch.Generator(pipe.pipe.device).manual_seed(seed))
         image1 = pipe.gen(params)
         image.save("cube_test_no_lpw1.png")
         diff = self.compute_diff(image1, image)
@@ -221,7 +221,7 @@ def test_controlnet(self):
         seed = 49045438434843
         params = dict(prompt="a cube  planet, cube-shaped, space photo, masterpiece",
                       negative_prompt="spherical",
-                      generator=torch.cuda.manual_seed(seed))
+                      generator=torch.Generator(pipe.pipe.device).manual_seed(seed))
         image = pipe.gen(params)
         image.save("cube_test.png")