Create app.py

app.py

@@ -0,0 +1,229 @@
+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)