Update app.py

app.py

@@ -9,10 +9,11 @@ import numpy as np
 from PIL import Image
 import spaces
 import torch
-from diffusers import StableDiffusionPipeline, StableDiffusionXLPipeline, EulerAncestralDiscreteScheduler
+from diffusers import StableDiffusionPipeline, StableDiffusionInpaintPipeline, StableDiffusionXLPipeline, EulerAncestralDiscreteScheduler
+from diffusers.utils import load_image
 
 DESCRIPTION = """
-# [**Fluently** Playground](https://huggingface.co/fluently)
+# [Fluently Playground](https://huggingface.co/fluently)
 
 """
 if not torch.cuda.is_available():
@@ -51,6 +52,14 @@ if torch.cuda.is_available():
     pipe_epic.scheduler = EulerAncestralDiscreteScheduler.from_config(pipe_epic.scheduler.config)
     pipe_epic.to(device)
 
+    pipe_inpaint = StableDiffusionInpaintPipeline.from_pretrained(
+        "fluently/Fluently-v3-inpainting",
+        torch_dtype=torch.float16,
+        use_safetensors=True,
+    )
+    #pipe_inpaint.scheduler = EulerAncestralDiscreteScheduler.from_config(pipe_inpaint.scheduler.config)
+    pipe_inpaint.to(device)
+
     pipe_xl = StableDiffusionXLPipeline.from_pretrained(
         "fluently/Fluently-XL-v2",
         torch_dtype=torch.float16,
@@ -74,10 +83,20 @@ def randomize_seed_fn(seed: int, randomize_seed: bool) -> int:
         seed = random.randint(0, MAX_SEED)
     return seed
 
+def get_model(model):
+    if model == "Fluently v3 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)
+
 
 @spaces.GPU(enable_queue=True)
 def generate(
     model,
+    inpaint_image,
+    mask_image,
+    blur_factor,
+    strength,
     prompt: str,
     negative_prompt: str = "",
     use_negative_prompt: bool = False,
@@ -127,7 +146,7 @@ def generate(
             num_images_per_prompt=1,
             output_type="pil",
         ).images
-    else:
+    elif model == "Fluently XL v2":
         images = pipe_xl(
             prompt=prompt,
             negative_prompt=negative_prompt,
@@ -138,11 +157,28 @@ def generate(
             num_images_per_prompt=1,
             output_type="pil",
         ).images
+    else:
+        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=30,
+            strength=strength,
+            num_images_per_prompt=1,
+            output_type="pil",
+        ).images
 
     image_paths = [save_image(img) for img in images]
     print(image_paths)
     return image_paths, seed
 
+
+
         
 
 
@@ -173,10 +209,20 @@ with gr.Blocks(title="Fluently Playground", css=css) as demo:
     with gr.Row():
         model = gr.Radio(
             label="Model",
-            choices=["Fluently XL v2", "Fluently v3.5", "Fluently Anime", "Fluently Epic"],
+            choices=["Fluently XL v2", "Fluently v3.5", "Fluently Anime", "Fluently Epic", "Fluently v3 inpaint"],
             value="Fluently v3.5",
             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(
@@ -244,6 +290,13 @@ with gr.Blocks(title="Fluently 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(
@@ -255,6 +308,10 @@ with gr.Blocks(title="Fluently Playground", css=css) as demo:
         fn=generate,
         inputs=[
             model,
+            inpaint_image,
+            mask_image,
+            blur_factor,
+            strength,
             prompt,
             negative_prompt,
             use_negative_prompt,