rimjhimittal/final_final

1import gradio as gr2from diffusers import AutoPipelineForInpainting, AutoencoderKL3from diffusers.utils import load_image4import torch5from PIL import Image6import spaces7from SegBody import segment_body8 9 10# Load models11 12vae = AutoencoderKL.from_pretrained("madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float32)13pipeline = AutoPipelineForInpainting.from_pretrained(14    "diffusers/stable-diffusion-xl-1.0-inpainting-0.1", 15    vae=vae, 16    torch_dtype=torch.float32, 17    variant="fp16", 18    use_safetensors=True,19    device="cuda"20)21pipeline.load_ip_adapter("h94/IP-Adapter", subfolder="sdxl_models", weight_name="ip-adapter_sdxl.bin", low_cpu_mem_usage=True)22 23# Function to process images24def virtual_try_on(img, clothing, prompt, negative_prompt, ip_scale=1.0, strength=0.99, guidance_scale=7.5, steps=100):25    _, mask_img = segment_body(img, face=False)26    pipeline.set_ip_adapter_scale(ip_scale)27    images = pipeline(28        prompt=prompt,29        negative_prompt=negative_prompt,30        image=img,31        mask_image=mask_img,32        ip_adapter_image=clothing,33        strength=strength,34        guidance_scale=guidance_scale,35        num_inference_steps=steps,36    ).images37    return images[0]38 39@spaces.GPU(duration=120)40def process_images(image, ip_image):41    image = image.convert("RGB").resize((512, 512))42    ip_image = ip_image.convert("RGB").resize((512, 512))43    44    seg_image, mask_image = segment_body(image, face=False)45    mask_image.resize((512, 512)) 46    47 48    return virtual_try_on(img=image,49               clothing=ip_image,50               prompt="photorealistic, perfect body, beautiful skin, realistic skin, natural skin",51               negative_prompt="ugly, bad quality, bad anatomy, deformed body, deformed hands, deformed feet, deformed face, deformed clothing, deformed skin, bad skin, leggings, tights, stockings")52 53# Create the Gradio interface54interface = gr.Interface(55    fn=process_images, 56    inputs=[gr.Image(type="pil"), gr.Image(type="pil")], 57    outputs=gr.Image(type="pil"),58    title="Image Inpainting Demo",59    description="Upload two images for inpainting using Stable Diffusion XL."60)61 62# Launch the Gradio interface63interface.launch()