File size: 3,033 Bytes
ca2d463 6057c81 ca2d463 8a51f48 6d80f3d 5690efa f881321 5690efa ca2d463 6057c81 8a51f48 6057c81 ca2d463 8a51f48 ca2d463 6057c81 ca2d463 8a51f48 6057c81 ca2d463 6057c81 ca2d463 6057c81 b7c5838 ca2d463 b7c5838 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 |
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 with the correct precision based on the device
if device == "cuda":
vae = AutoencoderKL.from_pretrained("madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float16) # Use fp16 for GPU
pipeline = AutoPipelineForInpainting.from_pretrained(
"diffusers/stable-diffusion-xl-1.0-inpainting-0.1",
vae=vae,
torch_dtype=torch.float16, # Use fp16 for GPU
variant="fp16", # Correct variant for GPU
use_safetensors=True
).to(device) # Ensure it uses the GPU
else:
vae = AutoencoderKL.from_pretrained("madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float32) # Use fp32 for CPU
pipeline = AutoPipelineForInpainting.from_pretrained(
"diffusers/stable-diffusion-xl-1.0-inpainting-0.1",
vae=vae,
torch_dtype=torch.float32, # Use fp32 for CPU
variant="fp16", # Use fp32 for CPU
use_safetensors=True
).to(device) # Ensure it uses the CPU if no GPU
# Define the inference function
def inpaint(person_image, garment_image, prompt):
# Preprocess the images by resizing them to 512x512
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=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 the generated image
# Set up the Gradio interface
demo = gr.Interface(
fn=inpaint,
inputs=[
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 with Segmentation",
description="Inpainting model for seamless garment transfer on segmented body image using Stable Diffusion XL.",
server_timeout=100, # Increase timeout duration to prevent session errors
)
demo.launch(share=True) # Enable share link for testing in a public domain
|