v1

app.py

@@ -0,0 +1,168 @@
+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
+@spaces.GPU(duration=75)
+@torch.no_grad()
+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()

requirements.txt

@@ -0,0 +1,10 @@
+torch>=2.4.0
+torchvision>=0.19.0
+diffusers>=0.29.2
+transformers>=4.37.2
+gradio>=5.7.0
+spaces>=0.30.4
+huggingface-hub>=0.26.2
+#torchvision>=0.20.1
+#diffusers>=0.31.0
+#transformers>=4.46.3