add distilled

app.py

@@ -6,7 +6,7 @@ from rembg import remove
 
 # import spaces #[uncomment to use ZeroGPU]
 from peft import PeftModel
-from diffusers import DiffusionPipeline, StableDiffusionPipeline, ControlNetModel, StableDiffusionControlNetPipeline
+from diffusers import DiffusionPipeline, StableDiffusionPipeline, ControlNetModel, StableDiffusionControlNetPipeline, AutoencoderTiny, DDIMScheduler
 from diffusers.utils import load_image
 import torch
 
@@ -26,7 +26,8 @@ MAX_IMAGE_SIZE = 512
 
 
 # @spaces.GPU #[uncomment to use ZeroGPU]
-def load_model(model_id, lora_strength, use_controlnet=False, control_mode="edge_detection", use_ip_adapter=False, control_strength_ip=0.0):
+def load_model(model_id, lora_strength, use_controlnet=False, control_mode="edge_detection", use_ip_adapter=False, control_strength_ip=0.0,
+               acceleration_mode=None):
     global pipe
     if pipe is not None:
         del pipe
@@ -75,8 +76,25 @@ def load_model(model_id, lora_strength, use_controlnet=False, control_mode="edge
                 )
                 pipe.unet = PeftModel.from_pretrained(pipe.unet, model_id, scaling=lora_strength, torch_dtype=torch_dtype)
             else:
-                pipe = StableDiffusionPipeline.from_pretrained("stable-diffusion-v1-5/stable-diffusion-v1-5", torch_dtype=torch_dtype)
-                pipe.unet = PeftModel.from_pretrained(pipe.unet, model_id, scaling=lora_strength)
+                if acceleration_mode is None:
+                    pipe = StableDiffusionPipeline.from_pretrained("stable-diffusion-v1-5/stable-diffusion-v1-5", torch_dtype=torch_dtype)
+                    pipe.unet = PeftModel.from_pretrained(pipe.unet, model_id, scaling=lora_strength)
+                elif acceleration_mode == "distilled":
+                    pipe = StableDiffusionPipeline.from_pretrained(
+                        "nota-ai/bk-sdm-small", torch_dtype=torch.float16, use_safetensors=True,
+                    )
+                elif acceleration_mode == "distilled + tiny":
+                    pipe = StableDiffusionPipeline.from_pretrained(
+                        "nota-ai/bk-sdm-small", torch_dtype=torch.float16, use_safetensors=True,
+                    )
+                    pipe.vae = AutoencoderTiny.from_pretrained(
+                        "sayakpaul/taesd-diffusers", torch_dtype=torch.float16, use_safetensors=True,
+                    )
+                elif acceleration_mode == "DDIM":
+                    pipe = StableDiffusionPipeline.from_pretrained("stable-diffusion-v1-5/stable-diffusion-v1-5", torch_dtype=torch_dtype)
+                    pipe.scheduler = DDIMScheduler.from_config(
+                        pipe.scheduler.config, rescale_betas_zero_snr=True, timestep_spacing="trailing"
+                    )
 
         if use_ip_adapter:
             pipe.load_ip_adapter("h94/IP-Adapter", subfolder="models", weight_name="ip-adapter_sd15.bin")
@@ -110,6 +128,7 @@ def infer(
         control_strength_ip,
         control_image_ip,
         use_rmbg,
+        acceleration_mode,
         progress=gr.Progress(track_tqdm=True),
 ):
     load_status = load_model(
@@ -118,7 +137,8 @@ def infer(
         use_controlnet,
         control_mode,
         use_ip_adapter,
-        control_strength_ip
+        control_strength_ip,
+        acceleration_mode
     )
     if randomize_seed:
         seed = random.randint(0, MAX_SEED)
@@ -278,6 +298,12 @@ with gr.Blocks(css=css) as demo:
 
         use_rmbg = gr.Checkbox(label="Delete background?", value=False)
 
+        use_acceleration = gr.Checkbox(label="Use accelerate model? (only for 1.5 SD!)", value=False)
+        with gr.Accordion("Acceleration Settings", open=True, visible=False) as acceleration_settings:
+            acceleration_mode = gr.Dropdown(label="Acceleration mode",
+                                            choices=["distilled", "distilled + tiny", "DDIM"],
+                                            value=None)
+
         gr.Examples(examples=examples, inputs=[prompt])
 
     gr.on(
@@ -301,7 +327,8 @@ with gr.Blocks(css=css) as demo:
             use_ip_adapter,
             control_strength_ip,
             control_image_ip,
-            use_rmbg
+            use_rmbg,
+            acceleration_mode
         ],
         outputs=[result, seed, model_status],
     )
@@ -323,5 +350,11 @@ with gr.Blocks(css=css) as demo:
         inputs=[use_rmbg]
     )
 
+    use_acceleration.change(
+        fn=lambda x: gr.update(visible=x, value=None),
+        inputs=[use_acceleration],
+        outputs=[acceleration_settings]
+    )
+
 if __name__ == "__main__":
     demo.launch()