add seed

app_base.py

@@ -48,6 +48,7 @@ def create_demo() -> gr.Blocks:
                 upscale_prompt,
                 start_size=pre_upscale_start_size,
                 upscale_steps=pre_upscale_steps,
+                seed=seed,
             )
             input_image = input_image.resize((generate_size, generate_size))
         run_task_time, time_cost_str = get_time_cost(run_task_time, time_cost_str)
@@ -70,6 +71,7 @@ def create_demo() -> gr.Blocks:
             upscale_prompt,
             start_size=upscale_start_size,
             upscale_steps=upscale_steps,
+            seed=seed,
         )
         run_task_time, time_cost_str = get_time_cost(run_task_time, time_cost_str)
 

upscale.py

@@ -3,23 +3,30 @@ import torch
 from PIL import Image
 from diffusers import StableDiffusionUpscalePipeline
 
+from hidiffusion import apply_hidiffusion
+
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
 model_id = "stabilityai/stable-diffusion-x4-upscaler"
 upscale_pipe = StableDiffusionUpscalePipeline.from_pretrained(model_id, torch_dtype=torch.float16)
 upscale_pipe = upscale_pipe.to(device)
 
+apply_hidiffusion(upscale_pipe)
+
 def upscale_image(
     input_image: Image,
     prompt: str,
     start_size: int = 128,
     upscale_steps: int = 30,
+    seed: int = 42,
 ):
+    generator = torch.Generator().manual_seed(seed)
     input_image = input_image.resize((start_size, start_size))
     upscaled_image = upscale_pipe(
         prompt=prompt, 
         image=input_image,
         num_inference_steps=upscale_steps,
+        generator=generator,
     ).images[0]
     
     return upscaled_image