Fix warmup steps after JIT to actually work

app.py

@@ -15,10 +15,6 @@ pipe = DiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-xl-base-1
 pipe.scheduler = DDIMScheduler.from_config(pipe.scheduler.config)
 pipe = pipe.to("cuda")
 
-# optimize for speed
-pipe.unet = torch.compile(pipe.unet, mode="max-autotune", fullgraph=True) # hopefully this works on Ampere series GPU
-pipe(prompt="an astronaut riding a green horse", num_inference_steps=25) # force lengthy JIT compilation to happen ahead of time
-
 # watermarking helper functions. paraphrased from the reference impl of arXiv:2305.20030
 
 def circle_mask(size=128, r=16, x_offset=0, y_offset=0):
@@ -114,6 +110,10 @@ def detect(image):
 def generate(prompt):
     return pipe(prompt=prompt, num_inference_steps=25, latents=get_noise()).images[0]
 
+# optimize for speed
+pipe.unet = torch.compile(pipe.unet, mode="reduce-overhead", fullgraph=True)
+print(detect(generate("an astronaut riding a green horse"))) # warmup after jit
+
 # actual gradio demo
 
 def manager(input, progress=gr.Progress(track_tqdm=True)): # to prevent the queue from overloading