Update app.py

app.py

@@ -3,7 +3,7 @@ import numpy as np
 import random
 import os
 import torch
-from diffusers import StableDiffusionPipeline, ControlNetModel, StableDiffusionControlNetPipeline
+from diffusers import StableDiffusionPipeline, ControlNetModel, StableDiffusionControlNetPipeline, AutoencoderTiny, DDIMScheduler
 from diffusers.utils import load_image
 from peft import PeftModel, LoraConfig
 from rembg import remove
@@ -39,10 +39,26 @@ def infer(
     ip_adapter_checkbox=False,
     ip_adapter_scale=0.0,
     ip_adapter_image=None,
+
+    tiny_vae=False,
+    ddim=False,
+    
     del_background=False,
+    alpha_matting=False, 
+    alpha_matting_foreground_threshold=240, 
+    alpha_matting_background_threshold=10, 
+    alpha_matting_erode_size=10, 
+    post_process_mask=False,
+    
     progress=gr.Progress(track_tqdm=True),    
 ):  
-    ckpt_dir='./model_output'
+    if model_id == model_id_default:
+        ckpt_dir='./model_output'
+    elif 'base' in model_id:
+        ckpt_dir='./model_output_distilled_base'
+    else:
+        ckpt_dir='./model_output_distilled_small'
+        
     unet_sub_dir = os.path.join(ckpt_dir, "unet")
     text_encoder_sub_dir = os.path.join(ckpt_dir, "text_encoder")
 
@@ -106,6 +122,12 @@ def infer(
 
     pipe.unet.load_state_dict({k: lora_scale*v for k, v in pipe.unet.state_dict().items()})
     pipe.text_encoder.load_state_dict({k: lora_scale*v for k, v in pipe.text_encoder.state_dict().items()})
+
+    if tiny_vae:
+        pipe.vae = AutoencoderTiny.from_pretrained("madebyollin/taesd", torch_dtype=torch_dtype)
+
+    if ddim:
+        pipe.scheduler = DDIMScheduler.from_config(pipe.scheduler.config)
     
     if torch_dtype in (torch.float16, torch.bfloat16):
         pipe.unet.half()
@@ -119,7 +141,13 @@ def infer(
     pipe.to(device)
 
     if del_background:
-        return remove(pipe(**params).images[0])
+        return remove(pipe(**params).images[0],
+                      alpha_matting=alpha_matting, 
+                      alpha_matting_foreground_threshold=alpha_matting_foreground_threshold, 
+                      alpha_matting_background_threshold=alpha_matting_background_threshold, 
+                      alpha_matting_erode_size=alpha_matting_erode_size, 
+                      post_process_mask=post_process_mask
+                     )
     else:
         return pipe(**params).images[0]
 
@@ -139,12 +167,15 @@ with gr.Blocks(css=css, fill_height=True) as demo:
         gr.Markdown(" # Text-to-Image demo")
 
         with gr.Row():
-            model_id = gr.Textbox(
-                label="Model ID",
-                max_lines=1,
-                placeholder="Enter model id",
-                value=model_id_default,
-            )
+            model_id = gr.Dropdown(
+                    label="Model ID",
+                    choices=[model_id_default,
+                             "nota-ai/bk-sdm-v2-base",
+                             "nota-ai/bk-sdm-v2-small"],
+                    value=model_id_default,
+                    max_choices=1
+                )
+            
 
         prompt = gr.Textbox(
             label="Prompt",
@@ -190,11 +221,58 @@ with gr.Blocks(css=css, fill_height=True) as demo:
                 step=1,
                 value=20,  # Replace with defaults that work for your model
             )
+        with gr.Row():
+            tiny_vae = = gr.Checkbox(
+                label="Use AutoencoderTiny?",
+                value=False
+            )
+            ddim = = gr.Checkbox(
+                label="Use DDIMScheduler?",
+                value=False
+            )
+            
         with gr.Row():
             del_background = gr.Checkbox(
                 label="Delete background?",
                 value=False
             )
+            with gr.Column(visible=False) as rembg_params:
+                alpha_matting = gr.Checkbox(
+                    label="alpha_matting",
+                    value=False
+                )
+                with gr.Column(visible=False) as alpha_params:
+                    alpha_matting_foreground_threshold = gr.Slider(
+                        label="alpha_matting_foreground_threshold",
+                        minimum=0,
+                        maximum=255,
+                        step=1,
+                        value=240,  
+                    )
+                    alpha_matting_background_threshold = gr.Slider(
+                        label="alpha_matting_background_threshold",
+                        minimum=0,
+                        maximum=255,
+                        step=1,
+                        value=10,  
+                    )
+                    alpha_matting_erode_size = gr.Slider(
+                        label="alpha_matting_erode_size",
+                        minimum=0,
+                        maximum=100,
+                        step=1,
+                        value=10,  
+                    )
+                alpha_matting.change(
+                    fn=lambda x: gr.Row.update(visible=x),
+                    inputs=alpha_matting,
+                    outputs=alpha_params
+                )
+                post_process_mask = gr.Checkbox(
+                    label="post_process_mask",
+                    value=False
+                )
+            
         with gr.Row():
             controlnet_checkbox = gr.Checkbox(
                 label="ControlNet",
@@ -294,7 +372,14 @@ with gr.Blocks(css=css, fill_height=True) as demo:
             ip_adapter_checkbox,
             ip_adapter_scale,
             ip_adapter_image,  
+            tiny_vae,
+            ddim,
             del_background,
+            alpha_matting, 
+            alpha_matting_foreground_threshold, 
+            alpha_matting_background_threshold, 
+            alpha_matting_erode_size, 
+            post_process_mask,
         ],
         outputs=[result],
     )