Update app.py

app.py

@@ -1,6 +1,6 @@
 import gradio as gr
 import torch
-from diffusers import StableDiffusionPipeline
+from diffusers import StableDiffusionPipeline, StableDiffusionInpaintPipeline
 from PIL import Image
 import numpy as np
 import cv2
@@ -8,38 +8,40 @@ import cv2
 # Device setup
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
-# Load Stable Diffusion model
-model_name = "CompVis/stable-diffusion-v1-4"
-def load_image_model():
+# Load model pipelines for different styles
+models = {
+    "fooocusv2": "CompVis/stable-diffusion-v1-4",
+    "SAI Anime": "CompVis/stable-diffusion-v1-4-anime"
+    # Add more models as needed
+}
+
+def load_image_model(style):
+    model_name = models.get(style, "CompVis/stable-diffusion-v1-4")
     return StableDiffusionPipeline.from_pretrained(model_name).to(device)
 
-image_model = load_image_model()
+image_model = load_image_model("fooocusv2")
+inpaint_model = StableDiffusionInpaintPipeline.from_pretrained("CompVis/stable-diffusion-v1-4-inpainting").to(device)
 
-def generate_image(prompt):
+def generate_image(prompt, style):
     try:
+        image_model = load_image_model(style)
         with torch.no_grad():
             image = image_model(prompt).images[0]
         return image, None
     except Exception as e:
         return None, f"Error generating image: {str(e)}"
 
-# Face swap function placeholder (replace with actual implementation)
-def swap_faces(image1, image2):
-    # Placeholder implementation, replace this with your actual face swapping code
-    return image1
-
 def face_swap(image1, image2):
     try:
         if image1 is None or image2 is None:
             return None, "Images for face swap are required"
         image1 = cv2.cvtColor(np.array(image1), cv2.COLOR_RGB2BGR)
         image2 = cv2.cvtColor(np.array(image2), cv2.COLOR_RGB2BGR)
-        swapped_image = swap_faces(image1, image2)
+        swapped_image = image1  # Placeholder implementation
         return cv2.cvtColor(swapped_image, cv2.COLOR_BGR2RGB), None
     except Exception as e:
         return None, f"Error during face swap: {str(e)}"
 
-# Upscaling function
 def upscale_image(image, scale_factor=2):
     try:
         if image is None:
@@ -52,17 +54,25 @@ def upscale_image(image, scale_factor=2):
     except Exception as e:
         return None, f"Error during upscaling: {str(e)}"
 
-# Gradio interface function
-def process_image(prompt, image1=None, image2=None, scale_factor=2):
+def inpaint_image(image, mask):
+    try:
+        if image is None or mask is None:
+            return None, "Image and mask are required for inpainting"
+        image = Image.fromarray(np.array(image))
+        mask = Image.fromarray(np.array(mask))
+        inpainted_image = inpaint_model(prompt="inpainting", image=image, mask_image=mask).images[0]
+        return inpainted_image, None
+    except Exception as e:
+        return None, f"Error during inpainting: {str(e)}"
+
+def process_image(prompt, style, image1=None, image2=None, mask=None, scale_factor=2):
     try:
         if prompt:
-            # Generate image from prompt
-            generated_image, error = generate_image(prompt)
+            generated_image, error = generate_image(prompt, style)
             if error:
                 return None, error
             return generated_image, None
         elif image1 and image2:
-            # Perform faceswap if two images are provided
             swapped_image, error = face_swap(image1, image2)
             if error:
                 return None, error
@@ -70,8 +80,13 @@ def process_image(prompt, image1=None, image2=None, scale_factor=2):
             if error:
                 return None, error
             return upscaled_image, None
+        elif image1 and mask:
+            inpainted_image, error = inpaint_image(image1, mask)
+            if error:
+                return None, error
+            return inpainted_image, None
         else:
-            return None, "Either a prompt or two images must be provided"
+            return None, "Either a prompt or two images or an image with a mask must be provided"
     except Exception as e:
         return None, f"Error in process_image function: {str(e)}"
 
@@ -80,16 +95,18 @@ iface = gr.Interface(
     fn=process_image,
     inputs=[
         gr.Textbox(label="Enter your prompt", placeholder="Type your prompt here..."),
-        gr.Image(label="Image 1 (for faceswap)", type="pil"),
-        gr.Image(label="Image 2 (for faceswap)", type="pil"),
-        gr.Slider(label="Upscale Factor", minimum=1, maximum=4, step=1, value=2)
+        gr.Dropdown(label="Select Style", choices=list(models.keys()), default="fooocusv2"),
+        gr.Image(label="Image 1 (for faceswap or inpainting)", type="pil", optional=True),
+        gr.Image(label="Image 2 (for faceswap)", type="pil", optional=True),
+        gr.Image(label="Mask (for inpainting)", type="pil", optional=True),
+        gr.Slider(label="Upscale Factor", minimum=1, maximum=4, step=1, value=2, optional=True)
     ],
     outputs=[
         gr.Image(label="Output Image"),
         gr.Textbox(label="Error Message", placeholder="Error details will appear here...")
     ],
     title="Fooocus Image Processing",
-    description="Generate images from prompts, swap faces, and upscale images."
+    description="Generate images from prompts with style selection, swap faces, upscale images, and perform inpainting."
 )
 
 if __name__ == "__main__":