VRAM: Generator to CPU

Signed-off-by: bghira <[email protected]>
bghira · Jul 8, 2023 · f115230 · f115230
1 parent c477f27
commit f115230
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/discord_tron_client/classes/image_manipulation/pipeline.py b/discord_tron_client/classes/image_manipulation/pipeline.py
@@ -396,7 +396,7 @@ def _get_generator(self, user_config: dict):
             self.seed = int(self.seed) + random.randint(-5, 5)
         elif int(self.seed) < 0:
             self.seed = random.randint(0, 2**32)
-        generator = torch.Generator(device=self.pipeline_manager.device)
+        generator = torch.Generator(device="cpu")
         generator.manual_seed(int(self.seed))
         logging.info(f"Seed: {self.seed}")
         return generator
@@ -478,7 +478,7 @@ def _refiner_pipeline(self, images: Image, user_config: dict, prompt: str = None
             # Get a random int:
             seed = random.randint(0, 2**32)
             new_images.append(pipe(
-                generator = torch.Generator().manual_seed(int(seed)),
+                generator = torch.Generator(device="cpu").manual_seed(int(seed)),
                 prompt=prompt,
                 negative_prompt=negative_prompt,
                 image=image,