Skip to content

Commit

Permalink
fix loader test
Browse files Browse the repository at this point in the history
  • Loading branch information
noskill committed Oct 1, 2024
1 parent 07f46f6 commit 225cdcd
Showing 1 changed file with 6 additions and 1 deletion.
7 changes: 6 additions & 1 deletion tests/pipe_test.py
Original file line number Diff line number Diff line change
Expand Up @@ -114,6 +114,11 @@ def test_loader(self):
loader = Loader()
model_id = self.get_model()
model_type = self.model_type()
device = torch.device('cpu')
if torch.cuda.is_available():
device = torch.device('cuda', 0)
if 'device' not in self.device_args:
self.device_args['device'] = device
classes = self.get_cls_by_type(MaskedIm2ImPipe)
# load inpainting pipe
cls = classes[model_type]
Expand Down Expand Up @@ -232,7 +237,7 @@ def test_lpw_turned_off(self):
seed = 49045438434843
params = dict(prompt=prompt,
negative_prompt="spherical",
generator=torch.Generator(pipe.pipe.device).manual_seed(seed))
generator=torch.Generator().manual_seed(seed))
image = pipe.gen(params)
image.save("cube_test_no_lpw.png")
params = dict(prompt=prompt + " , best quality, famous photo",
Expand Down

0 comments on commit 225cdcd

Please sign in to comment.