From 29dde1643233b3c07f9384da6be27911e2bf9c58 Mon Sep 17 00:00:00 2001
From: wailovet <wailovet@163.com>
Date: Thu, 4 Jul 2024 22:00:23 +0800
Subject: [PATCH] 'update'

---
 hook_HYDiT_utils.py | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/hook_HYDiT_utils.py b/hook_HYDiT_utils.py
index 6b9d865..d5ac04b 100644
--- a/hook_HYDiT_utils.py
+++ b/hook_HYDiT_utils.py
@@ -77,6 +77,7 @@ def easy_sample_images(
         infer_steps=20,
         sampler='dpmpp_2m_karras',
         train_steps=0,
+        seed=0,
 ):
     from hydit.diffusion.pipeline import StableDiffusionPipeline
     from diffusers import schedulers
@@ -186,7 +187,8 @@ def easy_sample_images(
             freqs_cis_img = calc_rope(height, width)
 
             try:
-
+                generator = torch.Generator(device="cuda")
+                generator.manual_seed(seed)
                 samples = pipeline(
                     height=height,
                     width=width,
@@ -201,6 +203,7 @@ def easy_sample_images(
                     learn_sigma=args.learn_sigma,
                     freqs_cis_img=freqs_cis_img,
                     image_meta_size=image_meta_size,
+                    generator=generator,
                 )[0]
 
                 pass