Updated App.py

app.py

@@ -2,53 +2,58 @@ import gradio as gr
 from PIL import Image
 from diffusers import AutoPipelineForInpainting, AutoencoderKL
 import torch
+from SegBody import segment_body  # Import the segmentation function
 
 # Check if CUDA is available and set the device accordingly
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
-# Load models
-vae = AutoencoderKL.from_pretrained("madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float32)  # Use float32 for CPU compatibility
+# Load models with fp16 variant for GPU compatibility
+vae = AutoencoderKL.from_pretrained("madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float16)
 pipeline = AutoPipelineForInpainting.from_pretrained(
     "diffusers/stable-diffusion-xl-1.0-inpainting-0.1", 
     vae=vae, 
-    torch_dtype=torch.float32,  # Use float32 for CPU compatibility
-    variant="fp32",  # Use fp32 for CPU
+    torch_dtype=torch.float16,  # Use fp16 for GPU
+    variant="fp16",  # Ensure you are using fp16 for GPU
     use_safetensors=True
 ).to(device)  # Ensure it uses the appropriate device (CPU or GPU)
 
 # Define the inference function
-def inpaint(prompt, image, mask_image, ip_image):
+def inpaint(person_image, garment_image, prompt):
     # Preprocess the images by resizing them to 512x512
-    image = image.convert("RGB").resize((512, 512))
-    mask_image = mask_image.convert("RGB").resize((512, 512))
-    ip_image = ip_image.convert("RGB").resize((512, 512))
+    person_image = person_image.convert("RGB").resize((512, 512))
+    garment_image = garment_image.convert("RGB").resize((512, 512))
+
+    # Use segment_body to generate the body mask for inpainting
+    seg_image, mask_image = segment_body(person_image, face=False)  # You can control face removal here (face=False)
+
+    # Resize mask to 512x512 to match the inpainting requirements
+    mask_image = mask_image.resize((512, 512))
 
     # Perform inpainting using the pipeline
     results = pipeline(
         prompt=prompt,
         negative_prompt="ugly, bad quality, bad anatomy",
-        image=image,
-        mask_image=mask_image,
-        ip_adapter_image=ip_image,
+        image=person_image,
+        mask_image=mask_image,  # Use the mask from segmentation
+        ip_adapter_image=garment_image,  # Garment image as the IP Adapter image
         strength=0.99,
         guidance_scale=8.0,
         num_inference_steps=100
     )
 
-    return results.images[0]
+    return results.images[0]  # Return the generated image
 
 # Set up the Gradio interface
 demo = gr.Interface(
     fn=inpaint,
     inputs=[
-        gr.Textbox(label="Prompt", placeholder="Enter the prompt for the model"),
-        gr.Image(type="pil", label="Input Image"),
-        gr.Image(type="pil", label="Mask Image"),
-        gr.Image(type="pil", label="IP Adapter Image")
+        gr.Image(type="pil", label="Person Image"),  # Input for person image
+        gr.Image(type="pil", label="Garment Image"),  # Input for garment image
+        gr.Textbox(label="Prompt", placeholder="Enter the prompt for the model")  # Text prompt for inpainting
     ],
     outputs=gr.Image(type="pil"),
-    title="Stable Diffusion Inpainting",
-    description="A model for inpainting and image editing using Stable Diffusion XL."
+    title="Stable Diffusion Inpainting with Segmentation",
+    description="Inpainting model for seamless garment transfer on segmented body image using Stable Diffusion XL."
 )
 
 demo.launch()