app.py

from diffusers import UNet2DConditionModel, AutoencoderKL, DDIMScheduler, AutoencoderTiny
from transformers import AutoTokenizer, CLIPTextModel, CLIPTextModelWithProjection
from accelerate import Accelerator
from huggingface_hub import hf_hub_download
import spaces
import gradio as gr    
import numpy as np
import torch
import time
import PIL

base = "stabilityai/stable-diffusion-xl-base-1.0"
repo_id = "tianweiy/DMD2"
subfolder = "model/sdxl/sdxl_cond999_8node_lr5e-7_denoising4step_diffusion1000_gan5e-3_guidance8_noinit_noode_backsim_scratch_checkpoint_model_019000"
filename = "pytorch_model.bin"


class ModelWrapper:
    def __init__(self, model_id, checkpoint_path, precision, image_resolution, latent_resolution, num_train_timesteps, conditioning_timestep, num_step, revision, accelerator):
        super().__init__()
        torch.set_grad_enabled(False)
        
        self.DTYPE = getattr(torch, precision)
        self.device = accelerator.device

        self.tokenizer_one = AutoTokenizer.from_pretrained(model_id, subfolder="tokenizer", revision=revision, use_fast=False)
        self.tokenizer_two = AutoTokenizer.from_pretrained(model_id, subfolder="tokenizer", revision=revision, use_fast=False)

        self.text_encoder = SDXLTextEncoder(model_id, revision, accelerator, dtype=self.DTYPE)

        self.vae = AutoencoderKL.from_pretrained(model_id, subfolder="vae").float().to(self.device)
        self.vae_dtype = torch.float32

        self.tiny_vae = AutoencoderTiny.from_pretrained("madebyollin/taesdxl", torch_dtype=self.DTYPE).to(self.device)
        self.tiny_vae_dtype = self.DTYPE

        self.model = self.create_generator(model_id, checkpoint_path).to(dtype=self.DTYPE).to(self.device)

        self.accelerator = accelerator
        self.image_resolution = image_resolution
        self.latent_resolution = latent_resolution
        self.num_train_timesteps = num_train_timesteps
        self.vae_downsample_ratio = image_resolution // latent_resolution
        self.conditioning_timestep = conditioning_timestep

        self.scheduler = DDIMScheduler.from_pretrained(model_id, subfolder="scheduler")
        self.alphas_cumprod = self.scheduler.alphas_cumprod.to(self.device)
        self.num_step = num_step

    def create_generator(self, model_id, checkpoint_path):
        generator = UNet2DConditionModel.from_pretrained(model_id, subfolder="unet").to(self.DTYPE)
        state_dict = torch.load(checkpoint_path, map_location="cpu")
        generator.load_state_dict(state_dict, strict=True)
        generator.requires_grad_(False)
        return generator

    def build_condition_input(self, height, width):
        original_size = (height, width)
        target_size = (height, width)
        crop_top_left = (0, 0)

        add_time_ids = list(original_size + crop_top_left + target_size)
        add_time_ids = torch.tensor([add_time_ids], device=self.device, dtype=self.DTYPE)
        return add_time_ids

    def _encode_prompt(self, prompt):
        text_input_ids_one = self.tokenizer_one([prompt], padding="max_length", max_length=self.tokenizer_one.model_max_length, truncation=True, return_tensors="pt").input_ids
        text_input_ids_two = self.tokenizer_two([prompt], padding="max_length", max_length=self.tokenizer_two.model_max_length, truncation=True, return_tensors="pt").input_ids

        prompt_dict = {
            'text_input_ids_one': text_input_ids_one.unsqueeze(0).to(self.device),
            'text_input_ids_two': text_input_ids_two.unsqueeze(0).to(self.device)
        }
        return prompt_dict

    @staticmethod
    def _get_time():
        torch.cuda.synchronize()
        return time.time()

    def sample(self, noise, unet_added_conditions, prompt_embed, fast_vae_decode):
        alphas_cumprod = self.scheduler.alphas_cumprod.to(self.device)

        if self.num_step == 1:
            all_timesteps = [self.conditioning_timestep]
            step_interval = 0 
        elif self.num_step == 4:
            all_timesteps = [999, 749, 499, 249]
            step_interval = 250 
        else:
            raise NotImplementedError()
        
        DTYPE = prompt_embed.dtype
        
        for constant in all_timesteps:
            current_timesteps = torch.ones(len(prompt_embed), device=self.device, dtype=torch.long) * constant
            eval_images = self.model(noise, current_timesteps, prompt_embed, added_cond_kwargs=unet_added_conditions).sample

            eval_images = get_x0_from_noise(noise, eval_images, alphas_cumprod, current_timesteps).to(self.DTYPE)

            next_timestep = current_timesteps - step_interval 
            noise = self.scheduler.add_noise(eval_images, torch.randn_like(eval_images), next_timestep).to(DTYPE)  

        if fast_vae_decode:
            eval_images = self.tiny_vae.decode(eval_images.to(self.tiny_vae_dtype) / self.tiny_vae.config.scaling_factor, return_dict=False)[0]
        else:
            eval_images = self.vae.decode(eval_images.to(self.vae_dtype) / self.vae.config.scaling_factor, return_dict=False)[0]
        eval_images = ((eval_images + 1.0) * 127.5).clamp(0, 255).to(torch.uint8).permute(0, 2, 3, 1)
        return eval_images 
        
    @spaces.GPU(enable_queue=True)
    @torch.no_grad()
    def inference(self, prompt, seed, height, width, num_images, fast_vae_decode):
        print("Running model inference...")

        if seed == -1:
            seed = np.random.randint(0, 1000000)

        generator = torch.manual_seed(seed)

        add_time_ids = self.build_condition_input(height, width).repeat(num_images, 1)

        noise = torch.randn(num_images, 4, height // self.vae_downsample_ratio, width // self.vae_downsample_ratio, generator=generator).to(device=self.device, dtype=self.DTYPE) 

        prompt_inputs = self._encode_prompt(prompt)
        
        start_time = self._get_time()

        prompt_embeds, pooled_prompt_embeds = self.text_encoder(prompt_inputs)

        batch_prompt_embeds, batch_pooled_prompt_embeds = (
            prompt_embeds.repeat(num_images, 1, 1),
            pooled_prompt_embeds.repeat(num_images, 1, 1)
        )

        unet_added_conditions = {
            "time_ids": add_time_ids,
            "text_embeds": batch_pooled_prompt_embeds.squeeze(1)
        }

        eval_images = self.sample(noise=noise, unet_added_conditions=unet_added_conditions, prompt_embed=batch_prompt_embeds, fast_vae_decode=fast_vae_decode)

        end_time = self._get_time()

        output_image_list = [] 
        for image in eval_images:
            output_image_list.append(PIL.Image.fromarray(image.cpu().numpy()))

        return output_image_list, f"Run successfully in {(end_time-start_time):.2f} seconds"

def get_x0_from_noise(sample, model_output, alphas_cumprod, timestep):
    alpha_prod_t = alphas_cumprod[timestep].reshape(-1, 1, 1, 1)
    beta_prod_t = 1 - alpha_prod_t

    pred_original_sample = (sample - beta_prod_t ** (0.5) * model_output) / alpha_prod_t ** (0.5)
    return pred_original_sample

class SDXLTextEncoder(torch.nn.Module):
    def __init__(self, model_id, revision, accelerator, dtype=torch.float32):
        super().__init__()

        self.text_encoder_one = CLIPTextModel.from_pretrained(model_id, subfolder="text_encoder", revision=revision).to(accelerator.device).to(dtype=dtype)
        self.text_encoder_two = CLIPTextModelWithProjection.from_pretrained(model_id, subfolder="text_encoder_2", revision=revision).to(accelerator.device).to(dtype=dtype)

        self.accelerator = accelerator

    def forward(self, batch):
        text_input_ids_one = batch['text_input_ids_one'].to(self.accelerator.device).squeeze(1)
        text_input_ids_two = batch['text_input_ids_two'].to(self.accelerator.device).squeeze(1)
        prompt_embeds_list = []

        for text_input_ids, text_encoder in zip([text_input_ids_one, text_input_ids_two], [self.text_encoder_one, self.text_encoder_two]):
            prompt_embeds = text_encoder(text_input_ids.to(text_encoder.device), output_hidden_states=True)

            pooled_prompt_embeds = prompt_embeds[0]

            prompt_embeds = prompt_embeds.hidden_states[-2]
            bs_embed, seq_len, _ = prompt_embeds.shape
            prompt_embeds = prompt_embeds.view(bs_embed, seq_len, -1)
            prompt_embeds_list.append(prompt_embeds)

        prompt_embeds = torch.cat(prompt_embeds_list, dim=-1)
        pooled_prompt_embeds = pooled_prompt_embeds.view(len(text_input_ids_one), -1)
        
        return prompt_embeds, pooled_prompt_embeds

def create_demo():
    TITLE = "# DMD2-SDXL Demo"
    model_id = "stabilityai/stable-diffusion-xl-base-1.0"
    checkpoint_path = hf_hub_download(repo_id=repo_id, subfolder=subfolder,filename=filename)
    precision = "float16"
    image_resolution = 1024
    latent_resolution = 128
    num_train_timesteps = 1000
    conditioning_timestep = 999
    num_step = 4
    revision = None

    torch.backends.cuda.matmul.allow_tf32 = True
    torch.backends.cudnn.allow_tf32 = True 

    accelerator = Accelerator()

    model = ModelWrapper(model_id, checkpoint_path, precision, image_resolution, latent_resolution, num_train_timesteps, conditioning_timestep, num_step, revision, accelerator)

    with gr.Blocks() as demo:
        gr.Markdown(TITLE)
        with gr.Row():
            with gr.Column():
                prompt = gr.Text(value="An oil painting of two rabbits in the style of American Gothic, wearing the same clothes as in the original.", label="Prompt")
                run_button = gr.Button("Run")
                with gr.Accordion(label="Advanced options", open=True):
                    seed = gr.Slider(label="Seed", minimum=-1, maximum=1000000, step=1, value=0)
                    num_images = gr.Slider(label="Number of generated images", minimum=1, maximum=16, step=1, value=16)
                    fast_vae_decode = gr.Checkbox(label="Use Tiny VAE for faster decoding", value=True)
                    height = gr.Slider(label="Image Height", minimum=512, maximum=1536, step=64, value=1024)
                    width = gr.Slider(label="Image Width", minimum=512, maximum=1536, step=64, value=1024)
            with gr.Column():
                result = gr.Gallery(label="Generated Images", show_label=False, elem_id="gallery", height=1024)
                error_message = gr.Text(label="Job Status")

        inputs = [prompt, seed, height, width, num_images, fast_vae_decode]
        run_button.click(fn=model.inference, inputs=inputs, outputs=[result, error_message], concurrency_limit=1)
    return demo

if __name__ == "__main__":
    demo = create_demo()
    demo.queue(api_open=False)
    demo.launch(show_error=True, share=True)