Spaces:
Paused
Paused
import os | |
import gradio as gr | |
import spaces | |
import torch | |
import torch.nn as nn | |
from diffusers import EulerDiscreteScheduler, AutoencoderKL, UNet2DConditionModel | |
from huggingface_hub import hf_hub_download | |
from transformers import SiglipImageProcessor, SiglipVisionModel | |
from torchvision.io import read_image | |
import torchvision.transforms.v2 as transforms | |
from torchvision.utils import make_grid | |
class TryOffDiff(nn.Module): | |
def __init__(self): | |
super().__init__() | |
self.unet = UNet2DConditionModel.from_pretrained("CompVis/stable-diffusion-v1-4", subfolder="unet") | |
self.transformer = torch.nn.TransformerEncoderLayer(d_model=768, nhead=8, batch_first=True) | |
self.proj = nn.Linear(1024, 77) | |
self.norm = nn.LayerNorm(768) | |
def adapt_embeddings(self, x): | |
x = self.transformer(x) | |
x = self.proj(x.permute(0, 2, 1)).permute(0, 2, 1) | |
return self.norm(x) | |
def forward(self, noisy_latents, t, cond_emb): | |
cond_emb = self.adapt_embeddings(cond_emb) | |
return self.unet(noisy_latents, t, encoder_hidden_states=cond_emb).sample | |
class PadToSquare: | |
def __call__(self, img): | |
_, h, w = img.shape # Get the original dimensions | |
max_side = max(h, w) | |
pad_h = (max_side - h) // 2 | |
pad_w = (max_side - w) // 2 | |
padding = (pad_w, pad_h, max_side - w - pad_w, max_side - h - pad_h) | |
return transforms.functional.pad(img, padding, padding_mode="edge") | |
# Set device | |
device = "cuda" if torch.cuda.is_available() else "cpu" | |
# Initialize Image Encoder | |
img_processor = SiglipImageProcessor.from_pretrained( | |
"google/siglip-base-patch16-512", | |
do_resize=False, | |
do_rescale=False, | |
do_normalize=False | |
) | |
img_enc = SiglipVisionModel.from_pretrained("google/siglip-base-patch16-512").eval().to(device) | |
img_enc_transform = transforms.Compose([ | |
PadToSquare(), # Custom transform to pad the image to a square | |
transforms.Resize((512, 512)), | |
transforms.ToDtype(torch.float32, scale=True), | |
transforms.Normalize(mean=[0.5], std=[0.5]), | |
]) | |
# Load TryOffDiff Model | |
path_model = hf_hub_download( | |
repo_id="rizavelioglu/tryoffdiff", | |
filename="tryoffdiff.pth", # or one of ["ldm-1", "ldm-2", "ldm-3", ...], | |
force_download=False | |
) | |
path_scheduler = hf_hub_download( | |
repo_id="rizavelioglu/tryoffdiff", | |
filename="scheduler/scheduler_config.json", | |
force_download=False | |
) | |
net = TryOffDiff() | |
net.load_state_dict(torch.load(path_model, weights_only=False)) | |
net.eval().to(device) | |
# Initialize VAE (only Decoder will be used) | |
vae = AutoencoderKL.from_pretrained("CompVis/stable-diffusion-v1-4", subfolder="vae").eval().to(device) | |
torch.cuda.empty_cache() | |
# Define image generation function | |
def generate_image(input_image, seed=42, guidance_scale=2.0, num_inference_steps=50, is_upscale=False): | |
# Configure scheduler | |
scheduler = EulerDiscreteScheduler.from_pretrained(path_scheduler) | |
scheduler.is_scale_input_called = True # suppress warning | |
scheduler.set_timesteps(num_inference_steps) | |
# Set random seed | |
generator = torch.Generator(device=device).manual_seed(seed) | |
x = torch.randn(1, 4, 64, 64, generator=generator, device=device) | |
# Process input image | |
cond_image = img_enc_transform(read_image(input_image)) | |
inputs = {k: v.to(img_enc.device) for k, v in img_processor(images=cond_image, return_tensors="pt").items()} | |
cond_emb = img_enc(**inputs).last_hidden_state.to(device) | |
# Prepare unconditioned embeddings | |
uncond_emb = torch.zeros_like(cond_emb) if guidance_scale > 1 else None | |
# Denoising loop with mixed precision | |
with torch.autocast(device): | |
for t in scheduler.timesteps: | |
if guidance_scale > 1: | |
noise_pred = net( | |
torch.cat([x] * 2), t, torch.cat([uncond_emb, cond_emb]) | |
).chunk(2) | |
noise_pred = noise_pred[0] + guidance_scale * (noise_pred[1] - noise_pred[0]) | |
else: | |
noise_pred = net(x, t, cond_emb) | |
scheduler_output = scheduler.step(noise_pred, t, x) | |
x = scheduler_output.prev_sample | |
# Decode preds | |
decoded = vae.decode(1 / 0.18215 * scheduler_output.pred_original_sample).sample | |
images = (decoded / 2 + 0.5).cpu() | |
# Create grid | |
grid = make_grid(images, nrow=len(images), normalize=True, scale_each=True) | |
if is_upscale: | |
pass | |
else: | |
return transforms.ToPILImage()(grid) | |
title = "Virtual Try-Off Generator" | |
description = r""" | |
This is the demo of the paper <a href="https://arxiv.org/abs/2411.18350">TryOffDiff: Virtual-Try-Off via High-Fidelity Garment Reconstruction using Diffusion Models</a>. | |
<br>Upload an image of a clothed individual to generate a standardized garment image using TryOffDiff. | |
<br> Check out the <a href="https://rizavelioglu.github.io/tryoffdiff/">project page</a> for more information. | |
""" | |
article = r""" | |
Example images are sampled from the `VITON-HD-test` set, which the models did not see during training. | |
<br>**Citation** <br>If you find our work useful in your research, please consider giving a star ⭐ and | |
a citation: | |
``` | |
@article{velioglu2024tryoffdiff, | |
title = {TryOffDiff: Virtual-Try-Off via High-Fidelity Garment Reconstruction using Diffusion Models}, | |
author = {Velioglu, Riza and Bevandic, Petra and Chan, Robin and Hammer, Barbara}, | |
journal = {arXiv}, | |
year = {2024}, | |
note = {\url{https://doi.org/nt3n}} | |
} | |
``` | |
""" | |
examples = [[f"examples/{img_filename}", 42, 2.0, 20, False] for img_filename in sorted(os.listdir("examples/"))] | |
# Create Gradio App | |
demo = gr.Interface( | |
fn=generate_image, | |
inputs=[ | |
gr.Image(type="filepath", label="Reference Image"), | |
gr.Slider(value=42, minimum=0, maximum=1e6, step=1, label="Seed"), | |
gr.Slider(value=2.0, minimum=1, maximum=5, step=0.5, label="Guidance Scale(s)", info="No guidance applied at s=1, hence faster inference."), | |
gr.Slider(value=20, minimum=0, maximum=1000, step=10, label="# of Inference Steps"), | |
gr.Checkbox(value=False, label="Upscale Output") | |
], | |
outputs=gr.Image(type="pil", label="Generated Garment", height=512, width=512), | |
title=title, | |
description=description, | |
article=article, | |
examples=examples, | |
examples_per_page=4, | |
) | |
demo.launch() | |