offload some work to CPU

modules/app.py

@@ -11,6 +11,13 @@ import time
 
 login(token=os.getenv("HF_TOKEN"))
 
+torch.set_float32_matmul_precision("high")
+
+torch._inductor.config.conv_1x1_as_mm = True
+torch._inductor.config.coordinate_descent_tuning = True
+torch._inductor.config.epilogue_fusion = False
+torch._inductor.config.coordinate_descent_check_all_directions = True
+
 app = FastAPI(debug=True)
 
 pipeline = StableDiffusion3Pipeline.from_pretrained(
@@ -42,6 +49,7 @@ def sd_on_gpu(prompt, negative_prompt, num_inference_steps, height, width, guida
         height=height,
         width=width,
         guidance_scale=guidance_scale,
+        generator=torch.manual_seed(1),
     )
     run_time = time.time() - start_time
     return response, run_time
@@ -59,7 +67,7 @@ play = gr.Interface(
         gr.Label(value="Image Generation Playground"),
         gr.Textbox(label="Prompt", lines=3),
         gr.Textbox(label="Negative Prompt", lines=2),
-        gr.Slider(label="Inference Steps", value=20, minimum=1, maximum=30, step=1),
+        gr.Slider(label="Inference Steps", value=13, minimum=1, maximum=30, step=1),
         gr.Number(label="Height", value=1024),
         gr.Number(label="Width", value=1024),
         gr.Slider(label="Guidance Scale", value=7, minimum=1, maximum=15, step=1)
@@ -73,7 +81,7 @@ stickers = gr.Interface(
         gr.Label(value="Sticker Optimization Console"),
         gr.Textbox(label="Prompt", lines=3),
         gr.Textbox(label="Negative Prompt", lines=2),
-        gr.Slider(label="Inference Steps", value=20, minimum=1, maximum=30, step=1),
+        gr.Slider(label="Inference Steps", value=13, minimum=1, maximum=30, step=1),
         gr.Number(label="Height", value=1024),
         gr.Number(label="Width", value=1024),
         gr.Slider(label="Guidance Scale", value=7, minimum=1, maximum=15, step=1)
@@ -81,5 +89,15 @@ stickers = gr.Interface(
     outputs=[gr.Gallery(), gr.Number(label="Total Generation Time"), gr.Number(label="GPU Time")],
 )
 
+@spaces.GPU(duration=10)
+def setup_pipe():
+    pipeline.transformer.to(memory_format=torch.channels_last)
+    pipeline.vae.to(memory_format=torch.channels_last)
+
+    pipeline.transformer = torch.compile(pipeline.transformer, mode="max-autotune", fullgraph=True)
+    pipeline.vae.decode = torch.compile(pipeline.vae.decode, mode="max-autotune", fullgraph=True)
+
+setup_pipe()
 app = gr.mount_gradio_app(app, play, path="/gradio")
-app = gr.mount_gradio_app(app, stickers, path="/stickers")
+app = gr.mount_gradio_app(app, stickers, path="/stickers")
+print('mounted')