From d09d92a717577e9f06af62e190765312981324b9 Mon Sep 17 00:00:00 2001
From: Luis Capelo <luis@modal.com>
Date: Tue, 18 Apr 2023 17:54:07 +0000
Subject: [PATCH 1/2] Pins torch and applies recommendation

Applies recommendation for loading models faster.
---
 .../stable_diffusion/stable_diffusion_cli.py  | 30 ++++++++++---------
 1 file changed, 16 insertions(+), 14 deletions(-)

diff --git a/06_gpu_and_ml/stable_diffusion/stable_diffusion_cli.py b/06_gpu_and_ml/stable_diffusion/stable_diffusion_cli.py
index 5a5ddac0e..e555a9b78 100644
--- a/06_gpu_and_ml/stable_diffusion/stable_diffusion_cli.py
+++ b/06_gpu_and_ml/stable_diffusion/stable_diffusion_cli.py
@@ -98,6 +98,7 @@ def download_models():
         "transformers",
         "triton",
         "safetensors",
+        "torch>=2.0"
     )
     .pip_install("xformers", pre=True)
     .run_function(
@@ -133,20 +134,21 @@ def __enter__(self):
 
         torch.backends.cuda.matmul.allow_tf32 = True
 
-        scheduler = diffusers.DPMSolverMultistepScheduler.from_pretrained(
-            cache_path,
-            subfolder="scheduler",
-            solver_order=2,
-            prediction_type="epsilon",
-            thresholding=False,
-            algorithm_type="dpmsolver++",
-            solver_type="midpoint",
-            denoise_final=True,  # important if steps are <= 10
-        )
-        self.pipe = diffusers.StableDiffusionPipeline.from_pretrained(
-            cache_path, scheduler=scheduler
-        ).to("cuda")
-        self.pipe.enable_xformers_memory_efficient_attention()
+        with torch.device("cuda"):
+            scheduler = diffusers.DPMSolverMultistepScheduler.from_pretrained(
+                cache_path,
+                subfolder="scheduler",
+                solver_order=2,
+                prediction_type="epsilon",
+                thresholding=False,
+                algorithm_type="dpmsolver++",
+                solver_type="midpoint",
+                denoise_final=True,  # important if steps are <= 10
+            )
+            self.pipe = diffusers.StableDiffusionPipeline.from_pretrained(
+                cache_path, scheduler=scheduler
+            ).to("cuda")
+            self.pipe.enable_xformers_memory_efficient_attention()
 
     @stub.function(gpu="A10G")
     def run_inference(

From 701a122f2770032d3dadcbb7f366a37a0a8f3949 Mon Sep 17 00:00:00 2001
From: Luis Capelo <luis@modal.com>
Date: Tue, 18 Apr 2023 17:58:41 +0000
Subject: [PATCH 2/2] Improves formatting

---
 06_gpu_and_ml/stable_diffusion/stable_diffusion_cli.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/06_gpu_and_ml/stable_diffusion/stable_diffusion_cli.py b/06_gpu_and_ml/stable_diffusion/stable_diffusion_cli.py
index e555a9b78..3e179f060 100644
--- a/06_gpu_and_ml/stable_diffusion/stable_diffusion_cli.py
+++ b/06_gpu_and_ml/stable_diffusion/stable_diffusion_cli.py
@@ -98,7 +98,7 @@ def download_models():
         "transformers",
         "triton",
         "safetensors",
-        "torch>=2.0"
+        "torch>=2.0",
     )
     .pip_install("xformers", pre=True)
     .run_function(