From b971b659e3ca4b186134bd1a56f21bf3db9b79ac Mon Sep 17 00:00:00 2001
From: Sean Sube <seansube@gmail.com>
Date: Thu, 26 Jan 2023 23:15:24 -0600
Subject: [PATCH] use a torch rng

---
 api/onnx_web/image.py   |  2 --
 api/onnx_web/upscale.py | 10 +++++++++-
 2 files changed, 9 insertions(+), 3 deletions(-)

diff --git a/api/onnx_web/image.py b/api/onnx_web/image.py
index 9a74ccf6..0bb53b7c 100644
--- a/api/onnx_web/image.py
+++ b/api/onnx_web/image.py
@@ -200,8 +200,6 @@ def process_tiles(
     tiles_y = height // tile
     total = tiles_x * tiles_y
 
-    rng = random.RandomState(0)
-
     for y in range(tiles_y):
         for x in range(tiles_x):
             idx = (y * tiles_x) + x
diff --git a/api/onnx_web/upscale.py b/api/onnx_web/upscale.py
index 2e463e92..a1df5c7f 100644
--- a/api/onnx_web/upscale.py
+++ b/api/onnx_web/upscale.py
@@ -11,6 +11,7 @@ from realesrgan import RealESRGANer
 from typing import Literal, Union
 
 import numpy as np
+import torch
 
 from .image import (
     process_tiles
@@ -146,8 +147,15 @@ def upscale_stable_diffusion(ctx: ServerContext, params: UpscaleParams, image: I
     # )
     # result = pipeline('', image=image)
 
+    generator = torch.manual_seed(0)
+    seed = generator.initial_seed()
+
     pipeline = StableDiffusionUpscalePipeline.from_pretrained('stabilityai/stable-diffusion-x4-upscaling')
-    upscale = lambda i: pipeline('an astronaut eating a hamburger', image=i, rng=np.random.RandomState(0)).images[0]
+    upscale = lambda i: pipeline(
+        'an astronaut eating a hamburger',
+        image=i,
+        generator=torch.manual_seed(initial_seed),
+    ).images[0]
     result = process_tiles(image, 128, 4, [upscale])
     return result