Update app.py

app.py

@@ -10,7 +10,8 @@ from PIL import Image
 import spaces
 import torch
 from diffusers import (
-    StableDiffusionXLPipeline, 
+    StableDiffusionXLPipeline,
+    StableDiffusionXLInpaintPipeline,
     DPMSolverMultistepScheduler
 )
 DESCRIPTION = """
@@ -34,6 +35,20 @@ pipe = StableDiffusionXLPipeline.from_pretrained(
 )
 pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)
 pipe.to('cuda')
+pipe_inpaint = StableDiffusionXLInpaintPipeline.from_pretrained(
+    "ehristoforu/Visionix-alpha-inpainting", 
+    torch_dtype=torch.float16,
+    use_safetensors=True,
+)
+pipe_inpaint.scheduler = DPMSolverMultistepScheduler.from_config(pipe_inpaint.scheduler.config)
+pipe_inpaint.to('cuda')
+
+def get_model(model):
+    if model == "Alpha inpaint":
+        return gr.update(visible=True), gr.update(visible=True), gr.update(visible=True), gr.update(visible=True), gr.update(visible=True)
+    else:
+        return gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), gr.update(visible=False)
+
 
 def save_image(img):
     unique_name = str(uuid.uuid4()) + ".png"
@@ -49,6 +64,11 @@ def randomize_seed_fn(seed: int, randomize_seed: bool) -> int:
 
 @spaces.GPU(durarion=50, enable_queue=True)
 def generate(
+    model,
+    inpaint_image,
+    mask_image,
+    blur_factor,
+    strength,
     prompt: str,
     negative_prompt: str = "",
     use_negative_prompt: bool = False,
@@ -65,16 +85,32 @@ def generate(
 
     if not use_negative_prompt:
         negative_prompt = ""  # type: ignore
-    images = pipe(
-        prompt=prompt,
-        negative_prompt=f"{negative_prompt}",     
-        width=width,
-        height=height,
-        guidance_scale=guidance_scale,
-        num_inference_steps=25,
-        num_images_per_prompt=1,
-        output_type="pil",
-    ).images
+    if model == "Alpha":
+        images = pipe(
+            prompt=prompt,
+            negative_prompt=f"{negative_prompt}",     
+            width=width,
+            height=height,
+            guidance_scale=guidance_scale,
+            num_inference_steps=25,
+            num_images_per_prompt=1,
+            output_type="pil",
+        ).images
+    elif model == "Alpha inpaint":
+        blurred_mask = pipe_inpaint.mask_processor.blur(mask_image, blur_factor=blur_factor)
+        images = pipe_inpaint(
+            prompt=prompt,
+            image=inpaint_image,
+            mask_image=blurred_mask,
+            negative_prompt=negative_prompt,
+            width=width,
+            height=height,
+            guidance_scale=guidance_scale,
+            num_inference_steps=25,
+            strength=strength,
+            num_images_per_prompt=1,
+            output_type="pil",
+        ).images
 
     image_paths = [save_image(img) for img in images]
     print(image_paths)
@@ -106,6 +142,25 @@ with gr.Blocks(title="Visionix Playground", css=css) as demo:
         elem_id="duplicate-button",
         visible=False,
     )
+
+    with gr.Row():
+        model = gr.Radio(
+            label="Model",
+            choices=["Alpha", "Alpha inpaint"],
+            value="Alpha",
+            interactive=True,
+        )
+
+    md_mask = gr.Markdown("""
+    ⚠️ To generate an inpaint mask, go [here](https://huggingface.co/spaces/stevhliu/inpaint-mask-maker).
+    """, visible=False)
+    inpaint_image = gr.Image(label="Inpaint Image", interactive=True, scale=5, visible=False, type="pil")
+    mask_image = gr.Image(label="Mask Image", interactive=True, scale=5, visible=False, type="pil")
+
+    blur_factor = gr.Slider(label="Mask Blur Factor", minimum=0, maximum=100, value=4, step=1, interactive=True, visible=False)
+    strength = gr.Slider(label="Denoising Strength", minimum=0.00, maximum=1.00, value=0.70, step=0.01, interactive=True, visible=False)
+    
+    
     with gr.Group():
         with gr.Row():
             prompt = gr.Text(
@@ -175,7 +230,13 @@ with gr.Blocks(title="Visionix Playground", css=css) as demo:
         outputs=negative_prompt,
         api_name=False,
     )
-    
+
+    model.change(
+        fn=get_model,
+        inputs=model,
+        outputs=[md_mask, inpaint_image, mask_image, blur_factor, strength],
+        api_name=False,
+    )
 
     gr.on(
         triggers=[