Update app.py

app.py

@@ -7,6 +7,7 @@ from diffusers import ControlNetModel
 import torch
 from PIL import Image
 from rembg import remove
+from diffusers import DiffusionPipeline, DDIMScheduler
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
 if torch.cuda.is_available():
@@ -54,7 +55,11 @@ def infer(
     ip_adapter_enabled=False,
     ip_adapter_scale=0.0,
     ip_adapter_image=None,
+    d_bckg=False,
+    ddim_use=False,
+    distill_vae=False,
     progress=gr.Progress(track_tqdm=True),
+    
 ):
     control_strength=float(control_strength)
     if randomize_seed:
@@ -69,7 +74,7 @@ def infer(
     if controlnet_enabled and control_image:
         controlnet_model = ControlNetModel.from_pretrained(CONTROLNET_MODES.get(control_mode))
         if model_id == "SD1.5 + lora Unet TextEncoder": 
-            pipe=StableDiffusionControlNetPipeline.from_pretrained("stable-deeffusion-v1-5/stable-diffusion-v1-5",controlnet=controlnet_model)
+            pipe=StableDiffusionControlNetPipeline.from_pretrained("stable-diffusion-v1-5/stable-diffusion-v1-5",controlnet=controlnet_model)
             pipe.unet = PeftModel.from_pretrained(pipe.unet, "um235/vCat_v2", subfolder="unet")
             pipe.text_encoder = PeftModel.from_pretrained(pipe.text_encoder, "um235/vCat_v2", subfolder="text_encoder")
         elif model_id == "SD1.5 + lora Unet TextEncoder" or model_id == "SD1.5 + lora Unet":
@@ -93,7 +98,7 @@ def infer(
             pipe.set_ip_adapter_scale(ip_adapter_scale)
     
     pipe.safety_checker = None
-           
+    if ddim_use: pipe.scheduler = DDIMScheduler.from_config(pipe.scheduler.config, rescale_betas_zero_snr=True)       
     pipe = pipe.to(device)
 
     image = pipe(
@@ -108,12 +113,13 @@ def infer(
         cross_attention_kwargs={"scale": lscale},
         controlnet_conditioning_scale=control_strength,
         ip_adapter_image=ip_adapter_image,
+        
     ).images[0]
     
-    #if d_bckg:
-        #image=remove(image)
+    if d_bckg:
+        image=remove(image)
         
-    #pipe.scheduler = DDIMScheduler.from_config(pipe.scheduler.config, rescale_betas_zero_snr=True)
+
     return image, seed
 
 
@@ -207,7 +213,7 @@ with gr.Blocks(css=css) as demo:
                     minimum=0.0,
                     maximum=2.0,
                     step=0.05,
-                    value=1.0,
+                    value=0.55,
                     visible=False,
                 )
 
@@ -295,6 +301,9 @@ with gr.Blocks(css=css) as demo:
             ip_adapter_enabled,
             ip_adapter_scale,
             ip_adapter_image,
+            d_bckg,
+            ddim_use,
+            distill_vae
         ],
         outputs=[result, seed],
     )
@@ -310,7 +319,5 @@ with gr.Blocks(css=css) as demo:
         inputs=[ip_adapter_enabled],
         outputs=[ip_adapter_scale, ip_adapter_image],
     )
-
-
 if __name__ == "__main__":
     demo.launch()