Spaces:

KheemDH
/

ViTON

Runtime error

ViTON / app.py

Change model precision to float32 for CPU compatibility

b7c5838 almost 2 years ago

3.03 kB

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

	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