Cache pipeline before running on GPU

app.py

@@ -12,6 +12,18 @@ models = [
 ]
 DEFAULT_MODEL = "stabilityai/stable-diffusion-xl-refiner-1.0"
 
+loaded_models: set[str] = set()
+
+
+def load_pipeline(model):
+    pipeline_type = (
+        StableDiffusionInstructPix2PixPipeline
+        if model == "timbrooks/instruct-pix2pix"
+        else AutoPipelineForImage2Image
+    )
+
+    return pipeline_type.from_pretrained(model)
+
 
 def generate_image(
     model: str,
@@ -20,14 +32,8 @@ def generate_image(
     strength: float,
     progress,
 ):
-    pipeline_type = (
-        StableDiffusionInstructPix2PixPipeline
-        if model == "timbrooks/instruct-pix2pix"
-        else AutoPipelineForImage2Image
-    )
-
     logger.debug(f"Loading pipeline: {dict(model=model)}")
-    pipe = pipeline_type.from_pretrained(model).to("cuda")
+    pipe = load_pipeline(model).to("cuda")
 
     logger.debug(f"Generating image: {dict(prompt=prompt)}")
     additional_args = (
@@ -75,6 +81,13 @@ def generate(
     # Downscale the image
     init_image.thumbnail((1024, 1024))
 
+    # Cache the model files for the pipeline
+    if model not in loaded_models:
+        logger.debug(f"Caching pipeline: {dict(model=model)}")
+
+        load_pipeline(model)
+        loaded_models.add(model)
+
     gpu_runner = gpu_3min if model == "timbrooks/instruct-pix2pix" else gpu
 
     return gpu_runner(model, prompt, init_image, strength, progress)