Update app.py

app.py

@@ -15,6 +15,7 @@ import numpy as np
 MODELS = {
     "RealVisXL V5.0 Lightning": "SG161222/RealVisXL_V5.0_Lightning",
     "Lustify Lightning": "GraydientPlatformAPI/lustify-lightning",
+    "Lustify Inpaint": "andro-flock/LUSTIFY-SDXL-NSFW-checkpoint-v2-0-INPAINTING",
     "Juggernaut XL Lightning": "RunDiffusion/Juggernaut-XL-Lightning",
 }
 
@@ -54,9 +55,17 @@ pipe = StableDiffusionXLFillPipeline.from_pretrained(
     controlnet=model,
 )
 
+pipe.scheduler = TCDScheduler.from_config(pipe.scheduler.config)
+
 pipe.to("cuda")
 
-pipe.scheduler = TCDScheduler.from_config(pipe.scheduler.config)
+pipe_inpaint = StableDiffusionXLFillPipeline.from_pretrained(
+    "andro-flock/LUSTIFY-SDXL-NSFW-checkpoint-v2-0-INPAINTING",
+    torch_dtype=torch.float16,
+    vae=vae,
+    controlnet=model,
+)
+pipe_inpaint.to("cuda")
 
 @spaces.GPU(duration=16)
 def fill_image(prompt, image, model_selection, paste_back):
@@ -236,12 +245,15 @@ def inpaint(prompt, image, inpaint_model, paste_back):
             vae=vae,
             controlnet=model,
         ).to("cuda")
+    # if pipe.config.model_name == "Lustify Inpaint":
+        
     
     mask = Image.fromarray(image["mask"]).convert("L")
     image = Image.fromarray(image["image"])
     
     result = pipe(prompt=prompt, image=image, mask_image=mask).images[0]
-    
+    # result = pipe_inpaint(prompt=prompt, image=image, mask_image=mask).images[0]
+
     if paste_back:
         result.paste(image, (0, 0), Image.fromarray(255 - np.array(mask)))