add pruna

server/pipelines/txt2img.py

@@ -12,6 +12,7 @@ from pydantic import BaseModel, Field
 from util import ParamsModel
 from PIL import Image
 from typing import List
+from pruna import SmashConfig, smash
 
 base_model = "SimianLuo/LCM_Dreamshaper_v7"
 taesd_model = "madebyollin/taesd"
@@ -86,6 +87,13 @@ class Pipeline:
                 taesd_model, torch_dtype=torch_dtype, use_safetensors=True
             ).to(device)
 
+        if args.pruna:
+            # Create and smash your model
+            smash_config = SmashConfig()
+            # smash_config["cacher"] = "deepcache"
+            smash_config["compiler"] = "stable_fast"
+            self.pipe = smash(model=self.pipe, smash_config=smash_config)
+
         if args.sfast:
             from sfast.compilers.stable_diffusion_pipeline_compiler import (
                 compile,

server/pipelines/txt2imgLora.py

@@ -12,6 +12,7 @@ from config import Args
 from pydantic import BaseModel, Field
 from util import ParamsModel
 from PIL import Image
+from pruna import SmashConfig, smash
 
 base_model = "wavymulder/Analog-Diffusion"
 lcm_lora_id = "latent-consistency/lcm-lora-sdv1-5"
@@ -93,6 +94,13 @@ class Pipeline:
                 taesd_model, torch_dtype=torch_dtype, use_safetensors=True
             ).to(device)
 
+        if args.pruna:
+            # Create and smash your model
+            smash_config = SmashConfig()
+            # smash_config["cacher"] = "deepcache"
+            smash_config["compiler"] = "stable_fast"
+            self.pipe = smash(model=self.pipe, smash_config=smash_config)
+
         self.pipe.scheduler = LCMScheduler.from_config(self.pipe.scheduler.config)
         self.pipe.set_progress_bar_config(disable=True)
         self.pipe.load_lora_weights(lcm_lora_id, adapter_name="lcm")