Update app.py

app.py

@@ -1,113 +1,51 @@
 import gradio as gr
 import torch
-from diffusers import StableDiffusionPipeline, StableDiffusionInpaintPipeline
-from PIL import Image
-import numpy as np
-import cv2
+from diffusers import DiffusionPipeline
 
-# Device setup
+# Set up device
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
-# 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
-}
+# Initialize the inpainting model
+try:
+    inpaint_model = DiffusionPipeline.from_pretrained("diffusers/stable-diffusion-xl-1.0-inpainting-0.1").to(device)
+except Exception as e:
+    print(f"Error initializing model: {e}")
 
-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("fooocusv2")
-inpaint_model = StableDiffusionInpaintPipeline.from_pretrained("CompVis/stable-diffusion-v1-4-inpainting").to(device)
-
-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)}"
-
-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 = 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)}"
-
-def upscale_image(image, scale_factor=2):
+def process_image(prompt, image, style, upscale_factor, inpaint):
     try:
-        if image is None:
-            return None, "Image for upscaling is required"
-        image = Image.fromarray(np.array(image))
-        width, height = image.size
-        new_size = (int(width * scale_factor), int(height * scale_factor))
-        upscaled_image = image.resize(new_size, Image.LANCZOS)
-        return upscaled_image, None
-    except Exception as e:
-        return None, f"Error during upscaling: {str(e)}"
-
-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:
-            generated_image, error = generate_image(prompt, style)
-            if error:
-                return None, error
-            return generated_image, None
-        elif image1 and image2:
-            swapped_image, error = face_swap(image1, image2)
-            if error:
-                return None, error
-            upscaled_image, error = upscale_image(swapped_image, scale_factor)
-            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
+        # Convert the image to the appropriate format if needed
+        if image is not None:
+            image = image.convert("RGB")
+
+        # Example placeholder logic for using the pipeline
+        # Here we assume the pipeline can handle both image and prompt; adjust as needed
+        if inpaint:
+            result = inpaint_model(prompt=prompt, image=image, guidance_scale=7.5)
         else:
-            return None, "Either a prompt or two images or an image with a mask must be provided"
+            result = inpaint_model(prompt=prompt, guidance_scale=7.5)
+        
+        return result.images[0]  # Return the first image from the result
     except Exception as e:
-        return None, f"Error in process_image function: {str(e)}"
-
-# Gradio interface setup
-iface = gr.Interface(
-    fn=process_image,
-    inputs=[
-        gr.Textbox(label="Enter your prompt", placeholder="Type your prompt here..."),
-        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 with style selection, swap faces, upscale images, and perform inpainting."
-)
-
-if __name__ == "__main__":
-    iface.launch(server_name="0.0.0.0", server_port=7860)
+        return f"Error in process_image function: {e}"
+
+# Define the Gradio interface
+with gr.Blocks() as demo:
+    with gr.Row():
+        with gr.Column():
+            prompt_input = gr.Textbox(label="Enter your prompt")
+            image_input = gr.Image(label="Image (for inpainting)", type="pil", optional=True)
+            style_input = gr.Dropdown(choices=["Fooocus Style", "SAI Anime"], label="Select Style")
+            upscale_input = gr.Slider(minimum=1, maximum=4, step=1, default=2, label="Upscale Factor")
+            inpaint_input = gr.Checkbox(label="Enable Inpainting")
+            
+            output = gr.Image(label="Generated Image")
+
+            generate_button = gr.Button("Generate Image")
+            generate_button.click(
+                process_image,
+                inputs=[prompt_input, image_input, style_input, upscale_input, inpaint_input],
+                outputs=output
+            )
+
+# Launch the interface
+demo.launch()