app.py

#   vae:
#     class_path: src.models.vae.LatentVAE
#     init_args:
#       precompute: true
#       weight_path: /mnt/bn/wangshuai6/models/sd-vae-ft-ema/
#   denoiser:
#     class_path: src.models.denoiser.decoupled_improved_dit.DDT
#     init_args:
#       in_channels: 4
#       patch_size: 2
#       num_groups: 16
#       hidden_size: &hidden_dim 1152
#       num_blocks: 28
#       num_encoder_blocks: 22
#       num_classes: 1000
#   conditioner:
#     class_path: src.models.conditioner.LabelConditioner
#     init_args:
#       null_class: 1000
#   diffusion_sampler:
#     class_path: src.diffusion.stateful_flow_matching.sampling.EulerSampler
#     init_args:
#       num_steps: 250
#       guidance: 3.0
#       state_refresh_rate: 1
#       guidance_interval_min: 0.3
#       guidance_interval_max: 1.0
#       timeshift: 1.0
#       last_step: 0.04
#       scheduler: *scheduler
#       w_scheduler: src.diffusion.stateful_flow_matching.scheduling.LinearScheduler
#       guidance_fn: src.diffusion.base.guidance.simple_guidance_fn
#       step_fn: src.diffusion.stateful_flow_matching.sampling.ode_step_fn
import os
import torch
import spaces
import argparse
from omegaconf import OmegaConf
from src.models.vae import fp2uint8
from src.diffusion.base.guidance import simple_guidance_fn
from src.diffusion.stateful_flow_matching.sharing_sampling import EulerSampler
from src.diffusion.stateful_flow_matching.scheduling import LinearScheduler
from PIL import Image
import gradio as gr
from huggingface_hub import snapshot_download


def instantiate_class(config):
    kwargs = config.get("init_args", {})
    class_module, class_name = config["class_path"].rsplit(".", 1)
    module = __import__(class_module, fromlist=[class_name])
    args_class = getattr(module, class_name)
    return args_class(**kwargs)

def load_model(weight_dict, denosier):
    prefix = "ema_denoiser."
    for k, v in denoiser.state_dict().items():
        try:
            v.copy_(weight_dict["state_dict"][prefix + k])
        except:
            print(f"Failed to copy {prefix + k} to denoiser weight")
    return denoiser


class Pipeline:
    def __init__(self, vae, denoiser, conditioner, diffusion_sampler, resolution, classlabels2ids):
        self.vae = vae
        self.denoiser = denoiser
        self.conditioner = conditioner
        self.diffusion_sampler = diffusion_sampler
        self.resolution = resolution
        self.classlabels2ids = classlabels2ids
    @spaces.GPU
    @torch.no_grad()
    @torch.autocast(device_type="cuda", dtype=torch.bfloat16)
    def __call__(self, y, num_images, seed, num_steps, guidance, state_refresh_rate, guidance_interval_min, guidance_interval_max, timeshift):
        self.diffusion_sampler.num_steps = num_steps
        self.diffusion_sampler.guidance = guidance
        self.diffusion_sampler.state_refresh_rate = state_refresh_rate
        self.diffusion_sampler.guidance_interval_min = guidance_interval_min
        self.diffusion_sampler.guidance_interval_max = guidance_interval_max
        self.diffusion_sampler.timeshift = timeshift
        generator = torch.Generator(device="cpu").manual_seed(seed)
        xT = torch.randn((num_images, 4, self.resolution//8, self.resolution//8), device="cpu", dtype=torch.float32, generator=generator)
        xT = xT.to("cuda")
        with torch.no_grad():
            condition, uncondition = conditioner([self.classlabels2ids[y],]*num_images)
        # Sample images:
        samples = diffusion_sampler(denoiser, xT, condition, uncondition)
        samples = vae.decode(samples)
        # fp32 -1,1 -> uint8 0,255
        samples = fp2uint8(samples)
        samples = samples.permute(0, 2, 3, 1).cpu().numpy()
        images = []
        for i in range(num_images):
            image = Image.fromarray(samples[i])
            images.append(image)
        return images

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--config", type=str, default="configs/repa_improved_ddt_xlen22de6_512.yaml")
    parser.add_argument("--resolution", type=int, default=512)
    parser.add_argument("--model_id", type=str, default="MCG-NJU/DDT-XL-22en6de-R512")
    parser.add_argument("--ckpt_path", type=str, default="models")
    args = parser.parse_args()

    if not os.path.exists(args.ckpt_path):
        snapshot_download(repo_id=args.model_id, local_dir=args.ckpt_path)

    config = OmegaConf.load(args.config)
    vae_config = config.model.vae
    diffusion_sampler_config = config.model.diffusion_sampler
    denoiser_config = config.model.denoiser
    conditioner_config = config.model.conditioner

    vae = instantiate_class(vae_config)
    denoiser = instantiate_class(denoiser_config)
    conditioner = instantiate_class(conditioner_config)


    diffusion_sampler = EulerSampler(
       scheduler=LinearScheduler(),
       w_scheduler=LinearScheduler(),
       guidance_fn=simple_guidance_fn,
       num_steps=50,
       guidance=4.0,
       state_refresh_rate=1,
       guidance_interval_min=0.3,
       guidance_interval_max=1.0,
       timeshift=1.0
    )
    ckpt_path = os.path.join(args.ckpt_path, "model.ckpt")
    ckpt = torch.load(ckpt_path, map_location="cpu")
    denoiser = load_model(ckpt, denoiser)
    denoiser = denoiser.cuda()
    vae = vae.cuda()
    denoiser.eval()

    # read imagenet classlabels
    with open("imagenet_classlabels.txt", "r") as f:
        classlabels = f.readlines()
        classlabels = [label.strip() for label in classlabels]

    classlabels2id = {label: i for i, label in enumerate(classlabels)}
    id2classlabels = {i: label for i, label in enumerate(classlabels)}

    pipeline = Pipeline(vae, denoiser, conditioner, diffusion_sampler, args.resolution, classlabels2id)

    with gr.Blocks() as demo:
        gr.Markdown("DDT: Decoupled Diffusion Transformer on ImageNet 512x512")
        with gr.Row():
            with gr.Column(scale=1):
                num_steps = gr.Slider(minimum=1, maximum=100, step=1, label="num steps", value=50)
                guidance = gr.Slider(minimum=0.1, maximum=10.0, step=0.1, label="CFG", value=4.0)
                num_images = gr.Slider(minimum=1, maximum=10, step=1, label="num images", value=4)
                label = gr.Dropdown(choices=classlabels, value=id2classlabels[950], label="label")
                seed = gr.Slider(minimum=0, maximum=1000000, step=1, label="seed", value=0)
                state_refresh_rate = gr.Slider(minimum=1, maximum=10, step=1, label="encoder reuse", value=1)
                guidance_interval_min = gr.Slider(minimum=0.0, maximum=1.0, step=0.1, label="interval guidance min",
                                                  value=0.0)
                guidance_interval_max = gr.Slider(minimum=0.1, maximum=1.0, step=0.1, label="interval guidance max",
                                                  value=1.0)
                timeshift = gr.Slider(minimum=0.1, maximum=2.0, step=0.1, label="timeshift", value=1.0)
            with gr.Column(scale=2):
                btn = gr.Button("Generate")
                output = gr.Gallery(label="Images")

        btn.click(fn=pipeline,
                  inputs=[
                      label,
                      num_images,
                      seed,
                      num_steps,
                      guidance,
                      state_refresh_rate,
                      guidance_interval_min,
                      guidance_interval_max,
                      timeshift
                  ], outputs=[output])
    demo.launch()