main.py

import logging
import math
import os
from typing import Any, Dict, List, Optional, Tuple, Union
#from diffusers.models.controlnet import ControlNetConditioningEmbedding
from diffusers.models.controlnets.controlnet import ControlNetConditioningEmbedding
import torch
from torch import nn
import torch.nn.functional as F
import torch.utils.checkpoint
import transformers
from accelerate import Accelerator
from accelerate.logging import get_logger
from accelerate.utils import ProjectConfiguration, set_seed

from tqdm.auto import tqdm
from src.configs.stage2_config import args

import diffusers
from diffusers import (
    AutoencoderKL,
    DDPMScheduler,
)
from diffusers.optimization import get_scheduler
from diffusers.utils import check_min_version, is_wandb_available
from src.dataset.stage2_dataset import InpaintDataset, InpaintCollate_fn
from transformers import CLIPVisionModelWithProjection
from transformers import Dinov2Model
from src.models.stage2_inpaint_unet_2d_condition import Stage2_InapintUNet2DConditionModel



import glob
import os
import torch
from torch import nn
from PIL import Image, ImageOps
import numpy as np
from diffusers import UniPCMultistepScheduler
from src.models.stage2_inpaint_unet_2d_condition import Stage2_InapintUNet2DConditionModel

from torchvision import transforms
#from diffusers.models.controlnet import ControlNetConditioningEmbedding
from transformers import CLIPImageProcessor
from transformers import Dinov2Model
from diffusers import AutoencoderKL, DDPMScheduler, UNet2DConditionModel,ControlNetModel,DDIMScheduler
from src.pipelines.PCDMs_pipeline import PCDMsPipeline
#from single_extract_pose import inference_pose


import spaces
from libs.easy_dwpose import DWposeDetector
from libs.easy_dwpose.draw import draw_openpose
from PIL import Image
import cv2
import os
import gradio as gr
import rembg
import uuid
import gc
from numba import cuda
import requests
import json

from huggingface_hub import hf_hub_download, HfApi


# Inputs ===================================================================================================

input_img = "sm.png"
train_imgs = ["target.png"]
in_vid = "walk.mp4"
out_vid = 'out.mp4'

"""
train_steps = 100
inference_steps = 10
fps = 12
"""

debug = False
save_model = True
should_gen_vid = False
max_batch_size = 8


def save_temp_imgs(imgs):
    os.makedirs('temp', exist_ok=True)
    results = []

    api = HfApi()
    

    for i, img in enumerate(imgs):

        #img_name = 'temp/'+str(uuid.uuid4())+'.png'
        img_name = 'temp/'+str(i)+'.png'
        img.save(img_name)

        """
        url = 'https://tmpfiles.org/api/v1/upload'

        try:
            response = requests.post(url, files={'file': open(img_name, 'rb')})

            # Check for successful response (status code 200)
            response.raise_for_status()

            # Print the server's response
            print("Status Code:", response.status_code)

            data = response.json()
            print("Response JSON:", data)
            results.append(data['data']['url'])

        except requests.exceptions.RequestException as e:
            print(f"An error occurred: {e}")
        """

        results.append('https://huggingface.co/datasets/acmyu/KeyframesAIFiles/resolve/main/'+img_name)

    api.upload_file(
        path_or_fileobj='temp',
        path_in_repo='temp',
        repo_id="acmyu/KeyframesAIFiles",
        repo_type="dataset",
    )

    return results


def getThumbnails(imgs):
    thumbs = []
    thumb_size = (512, 512)
    for img in imgs:
        th = img.copy()
        th.thumbnail(thumb_size)
        thumbs.append(th)
    return thumbs


# Pose detection ==============================================================================================

def load_models():
    dwpose = DWposeDetector(device="cuda")
    rembg_session = rembg.new_session("u2netp")

    pcdms_model = hf_hub_download(repo_id="acmyu/PCDMs", filename="pcdms_ckpt.pt")
    
    # Load scheduler
    noise_scheduler = DDPMScheduler.from_pretrained("stabilityai/stable-diffusion-2-1-base", subfolder="scheduler")

    # Load model
    image_encoder_p = Dinov2Model.from_pretrained('facebook/dinov2-giant')
    image_encoder_g = CLIPVisionModelWithProjection.from_pretrained('laion/CLIP-ViT-H-14-laion2B-s32B-b79K')#("openai/clip-vit-base-patch32")

    vae = AutoencoderKL.from_pretrained("stabilityai/stable-diffusion-2-1-base", subfolder="vae")
    unet = Stage2_InapintUNet2DConditionModel.from_pretrained(
                "stabilityai/stable-diffusion-2-1-base", 
                torch_dtype=torch.float16,
                subfolder="unet",
                in_channels=9, 
                low_cpu_mem_usage=False, 
                ignore_mismatched_sizes=True)

    
    return dwpose, rembg_session, pcdms_model, noise_scheduler, image_encoder_p, image_encoder_g, vae, unet


#load_models()

def img_pad(img, tw, th, transparent=False):
    img.thumbnail((tw, th))
    if transparent:
        new_img = Image.new('RGBA', (tw, th), (0, 0, 0, 0))
    else:
        new_img = Image.new("RGB", (tw, th), (0, 0, 0))
    left = (tw - img.width) // 2
    top = (th - img.height) // 2
    new_img.paste(img, (left, top))
    return new_img


def resize_and_pad(img, target_img):
    tw, th = target_img.size
    w, h = img.size
    
    if tw/th > w/h:
        tw = int(th * w/h)
    elif tw/th < w/h:
        th = int(tw * h/w)
    
    img = img.resize((tw, th), Image.BICUBIC)

    tw, th = target_img.size

    return img_pad(img, tw, th)


def remove_zero_pad(image):
    image = np.array(image)
    dummy = np.argwhere(image != 0) # assume blackground is zero
    max_y = dummy[:, 0].max()
    min_y = dummy[:, 0].min()
    min_x = dummy[:, 1].min()
    max_x = dummy[:, 1].max()
    crop_image = image[min_y:max_y, min_x:max_x]

    return Image.fromarray(crop_image)
    

def get_pose(img, dwpose, outfile, crop=False):
    #pil_image = Image.open("imgs/"+img).convert("RGB")
    #skeleton = dwpose(pil_image, output_type="np", include_hands=True, include_face=False)

    #img.thumbnail((512,512))
    out_img, pose = dwpose(img, include_hands=True, include_face=False)

    #print(pose['bodies'])
    
    if crop:
        bbox = out_img.getbbox()
        out_img = out_img.crop(bbox)
        out_img = ImageOps.expand(out_img, border=int(out_img.width*0.2), fill=(0,0,0))
    
    return out_img, pose


def extract_frames(video_path, fps):
    video_capture = cv2.VideoCapture(video_path)
    frame_count = 0
    frames = []
    
    fps_in = video_capture.get(cv2.CAP_PROP_FPS)
    fps_out = fps

    index_in = -1
    index_out = -1

    while True:
        success = video_capture.grab()
        if not success: break
        index_in += 1

        out_due = int(index_in / fps_in * fps_out)
        if out_due > index_out:
            success, frame = video_capture.retrieve()
            if not success: 
                break
            index_out += 1
            
            frame_count += 1
            frames.append(Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)))

    video_capture.release()
    print(f"Extracted {frame_count} frames")
    return frames


def removebg(img, rembg_session, transparent=False):
    
    if transparent:
        result = Image.new('RGBA', img.size, (0, 0, 0, 0))
    else:
        result = Image.new("RGB", img.size, "#ffffff")
    out = rembg.remove(img, session=rembg_session)
    result.paste(out, mask=out)
    return result


def prepare_inputs_train(images, bg_remove, dwpose, rembg_session):
    print("remove background", bg_remove)
    if bg_remove:
        images = [removebg(img, rembg_session) for img in images]

    in_img = images[0]
    in_pose, _ = get_pose(in_img, dwpose, "in_pose.png")
    train_poses = []
    train_imgs = [resize_and_pad(img, in_img) for img in images[1:]]
    
    for i, img in enumerate(train_imgs):
        train_pose, _ = get_pose(img, dwpose, "tr_pose"+str(i)+".png")
        train_poses.append(train_pose)
        
    return in_img, in_pose, train_imgs, train_poses
        

def prepare_inputs_inference(in_img, in_vid, fps, dwpose, rembg_session, bg_remove, resize_inputs, is_app=False):
    progress=gr.Progress(track_tqdm=True)

    print("prepare_inputs_inference")
    
    in_pose, _ = get_pose(in_img, dwpose, "in_pose.png")
    
    frames = extract_frames(in_vid, fps)
    print("remove background", bg_remove)
    if bg_remove:
        in_img = removebg(in_img, rembg_session)
        #frames = [removebg(img, rembg_session) for img in frames]
    if debug:
        for i, frame in enumerate(frames):
            frame.save("out/frame_"+str(i)+".png")
    
    print("vid: ", in_vid, fps)

    progress_bar = tqdm(range(len(frames)), initial=0, desc="Frames")
    target_poses = []
    target_poses_coords = []
    max_left = max_top = 999999
    max_right = max_bottom = 0
    it = frames
    if is_app:
        it = progress.tqdm(frames, desc="Pose Detection")
    for f in it:
        tpose, tpose_coords = get_pose(f, dwpose, "tar_pose"+str(len(target_poses))+".png")
        #print(tpose_coords)
        coords = {}
        for k in tpose_coords:
            if k == 'bodies_multi':
                coords['bodies'] = tpose_coords[k].tolist()
            elif k in ['hands']:
                coords[k] = tpose_coords[k].tolist()
            elif k in ['num_candidates']:
                coords[k] = tpose_coords[k]
        #print(coords)
        target_poses.append(tpose)
        target_poses_coords.append(json.dumps(coords))
        progress_bar.update(1)
        
        bbox = tpose.getbbox()
        left, top, right, bottom = bbox
        max_left = min(max_left, left)
        max_top = min(max_top, top)
        max_right = max(max_right, right)
        max_bottom = max(max_bottom, bottom)
        
    target_poses_cropped = []
    for tpose in target_poses:
        if resize_inputs:
            tpose = tpose.crop((max_left, max_top, max_right, max_bottom))
            tpose = ImageOps.expand(tpose, border=int(tpose.width*0.2), fill=(0,0,0))
            
            tpose = resize_and_pad(tpose, in_img)
        
        
        if debug:
            tpose.save("out/"+"tar_pose"+str(len(target_poses_cropped))+".png")
        target_poses_cropped.append(tpose)
        
    return in_img, target_poses_cropped, in_pose, target_poses_coords, frames


def prepare_inputs(images, in_vid, fps, bg_remove, dwpose, rembg_session, resize_inputs, is_app=False):
    
    in_img, in_pose, train_imgs, train_poses = prepare_inputs_train(images, bg_remove, dwpose, rembg_session)
    
    in_img, target_poses_cropped, _, _, _ = prepare_inputs_inference(in_img, in_vid, fps, dwpose, rembg_session, bg_remove, resize_inputs, is_app)
    
    
    return in_img, in_pose, train_imgs, train_poses, target_poses_cropped


# Training ===================================================================================================

# Will error if the minimal version of diffusers is not installed. Remove at your own risks.
check_min_version("0.18.0.dev0")

logger = get_logger(__name__)
    

class ImageProjModel_p(torch.nn.Module):
    """SD model with image prompt"""

    def __init__(self, in_dim, hidden_dim, out_dim, dropout = 0.):
        super().__init__()

        self.net = nn.Sequential(
            nn.Linear(in_dim, hidden_dim),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.LayerNorm(hidden_dim),
            nn.Linear(hidden_dim, out_dim),
            nn.Dropout(dropout)
        )

    def forward(self, x): 
        return self.net(x)

class ImageProjModel_g(torch.nn.Module):
    """SD model with image prompt"""

    def __init__(self, in_dim, hidden_dim, out_dim, dropout = 0.):
        super().__init__()

        self.net = nn.Sequential(
            nn.Linear(in_dim, hidden_dim),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.LayerNorm(hidden_dim),
            nn.Linear(hidden_dim, out_dim),
            nn.Dropout(dropout)
        )

    def forward(self, x):  # b, 257,1280
        return self.net(x)


class SDModel(torch.nn.Module):
    """SD model with image prompt"""
    def __init__(self, unet) -> None:
        super().__init__()
        self.image_proj_model_p = ImageProjModel_p(in_dim=1536, hidden_dim=768, out_dim=1024)

        self.unet = unet
        self.pose_proj = ControlNetConditioningEmbedding(
            conditioning_embedding_channels=320,
            block_out_channels=(16, 32, 96, 256),
            conditioning_channels=3)


    def forward(self, noisy_latents, timesteps, simg_f_p, timg_f_g, pose_f):

        extra_image_embeddings_p = self.image_proj_model_p(simg_f_p)
        extra_image_embeddings_g = timg_f_g
        
        print(extra_image_embeddings_p.size())
        print(extra_image_embeddings_g.size())

        encoder_image_hidden_states = torch.cat([extra_image_embeddings_p ,extra_image_embeddings_g], dim=1)
        pose_cond = self.pose_proj(pose_f)

        pred_noise = self.unet(noisy_latents, timesteps, class_labels=timg_f_g, encoder_hidden_states=encoder_image_hidden_states,my_pose_cond=pose_cond).sample
        return pred_noise
    
def load_training_checkpoint(model, pcdms_model, tag=None, **kwargs):
    #model_sd = torch.load(load_dir, map_location="cpu")["module"]
    model_sd = torch.load(
        pcdms_model, 
        map_location="cpu"
    )["module"]
    

    image_proj_model_dict = {}
    pose_proj_dict = {}
    unet_dict = {}
    for k in model_sd.keys():
        if k.startswith("pose_proj"):
            pose_proj_dict[k.replace("pose_proj.", "")] = model_sd[k]

        elif k.startswith("image_proj_model_p"):
            image_proj_model_dict[k.replace("image_proj_model_p.", "")] = model_sd[k]
            
        elif k.startswith("image_proj_model."):
            image_proj_model_dict[k.replace("image_proj_model.", "")] = model_sd[k]


        elif k.startswith("unet"):
            unet_dict[k.replace("unet.", "")] = model_sd[k]
        else:
            print(k)
    
    model.pose_proj.load_state_dict(pose_proj_dict)
    model.image_proj_model_p.load_state_dict(image_proj_model_dict)
    model.unet.load_state_dict(unet_dict)
    
    return model, 0, 0


def checkpoint_model(checkpoint_folder, ckpt_id, model, epoch, last_global_step, **kwargs):
    """Utility function for checkpointing model + optimizer dictionaries
    The main purpose for this is to be able to resume training from that instant again
    """
    checkpoint_state_dict = {
        "epoch": epoch,
        "last_global_step": last_global_step,
    }
    # Add extra kwargs too
    checkpoint_state_dict.update(kwargs)

    success = model.save_checkpoint(checkpoint_folder, ckpt_id, checkpoint_state_dict)
    status_msg = f"checkpointing: checkpoint_folder={checkpoint_folder}, ckpt_id={ckpt_id}"
    if success:
        logging.info(f"Success {status_msg}")
    else:
        logging.warning(f"Failure {status_msg}")
    return


@spaces.GPU(duration=600)
def train(modelId, in_image, in_pose, train_images, train_poses, train_steps, pcdms_model, noise_scheduler, image_encoder_p, image_encoder_g, vae, unet, finetune=True, is_app=False):
    logging_dir = 'outputs/logging'
    print('start train')
    
    
    progress=gr.Progress(track_tqdm=True)

    accelerator = Accelerator(
        log_with=args.report_to,
        project_dir=logging_dir,
        mixed_precision=args.mixed_precision,
        gradient_accumulation_steps=args.gradient_accumulation_steps
    )

    # Make one log on every process with the configuration for debugging.
    #logging.basicConfig(
    #    format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
    #    datefmt="%m/%d/%Y %H:%M:%S",
    #    level=logging.INFO, )

    print(accelerator.state)
    if accelerator.is_local_main_process:
        transformers.utils.logging.set_verbosity_warning()
        diffusers.utils.logging.set_verbosity_info()
    else:
        transformers.utils.logging.set_verbosity_error()
        diffusers.utils.logging.set_verbosity_error()

    # If passed along, set the training seed now.
    set_seed(42)

    # Handle the repository creation
    if accelerator.is_main_process:
        os.makedirs('outputs', exist_ok=True)
    
    
    """
    unet = Stage2_InapintUNet2DConditionModel.from_pretrained("stabilityai/stable-diffusion-2-1-base", subfolder="unet",
                                                   in_channels=9, class_embed_type="projection" ,projection_class_embeddings_input_dim=1024,
                                                  low_cpu_mem_usage=False, ignore_mismatched_sizes=True)
    """
    image_encoder_p.requires_grad_(False)
    image_encoder_g.requires_grad_(False)
    vae.requires_grad_(False)

    sd_model = SDModel(unet=unet)
    sd_model.train()


    if args.gradient_checkpointing:
        sd_model.enable_gradient_checkpointing()


    # Enable TF32 for faster training on Ampere GPUs,
    # cf https://pytorch.org/docs/stable/notes/cuda.html#tensorfloat-32-tf32-on-ampere-devices
    if args.allow_tf32:
        torch.backends.cuda.matmul.allow_tf32 = True

    learning_rate = 1e-4
    train_batch_size = min(len(train_images), max_batch_size) #len(train_images) % 16
    

    # Optimizer creation
    params_to_optimize = sd_model.parameters()
    optimizer = torch.optim.AdamW(
        params_to_optimize,
        lr=learning_rate,
        betas=(args.adam_beta1, args.adam_beta2),
        weight_decay=args.adam_weight_decay,
        eps=args.adam_epsilon,
    )
    
    inputs = [{
        "source_image": in_image,
        "source_pose": in_pose,
        "target_image": timg,
        "target_pose": tpose,
    } for timg, tpose in zip(train_images, train_poses)]
    
    """
    inputs = {[
        "source_image": Image.open('imgs/sm.png'),
        "source_pose": Image.open('imgs/sm_pose.jpg'),
        "target_image": Image.open('imgs/target.png'),
        "target_pose": Image.open('imgs/target_pose.jpg'),
    ]}
    """
    
    #print(inputs)
    
    dataset = InpaintDataset(
        inputs, 
        'imgs/', 
        size=(args.img_width, args.img_height), # w h
        imgp_drop_rate=0.1,
        imgg_drop_rate=0.1,
    )

    """
    dataset = InpaintDataset(
        args.json_path,
        args.image_root_path,
        size=(args.img_width, args.img_height), # w h
        imgp_drop_rate=0.1,
        imgg_drop_rate=0.1,
    )
    """

    train_sampler = torch.utils.data.distributed.DistributedSampler(
        dataset, num_replicas=accelerator.num_processes, rank=accelerator.process_index, shuffle=True)

    train_dataloader = torch.utils.data.DataLoader(
        dataset,
        sampler=train_sampler,
        collate_fn=InpaintCollate_fn,
        batch_size=train_batch_size,
        num_workers=0,)
    

    # Scheduler and math around the number of training steps.
    overrode_max_train_steps = False
    num_update_steps_per_epoch = math.ceil(len(train_dataloader) / args.gradient_accumulation_steps)
    if args.max_train_steps is None:
        args.max_train_steps = args.num_train_epochs * num_update_steps_per_epoch
        overrode_max_train_steps = True
    args.max_train_steps = train_steps

    lr_scheduler = get_scheduler(
        args.lr_scheduler,
        optimizer=optimizer,
        num_warmup_steps=args.lr_warmup_steps * accelerator.num_processes,
        num_training_steps=args.max_train_steps * accelerator.num_processes,
        num_cycles=args.lr_num_cycles,
        power=args.lr_power,
    )

    # Prepare everything with our `accelerator`.
    sd_model, optimizer, train_dataloader, lr_scheduler = accelerator.prepare(sd_model, optimizer, train_dataloader, lr_scheduler)

    # For mixed precision training we cast the text_encoder and vae weights to half-precision
    # as these models are only used for inference, keeping weights in full precision is not required.
    weight_dtype = torch.float32
    """
    if accelerator.mixed_precision == "fp16":
        weight_dtype = torch.float16
    elif accelerator.mixed_precision == "bf16":
        weight_dtype = torch.bfloat16
    """

    # Move vae, unet and text_encoder to device and cast to weight_dtype
    vae.to(accelerator.device, dtype=weight_dtype)
    sd_model.unet.to(accelerator.device, dtype=weight_dtype)
    image_encoder_p.to(accelerator.device, dtype=weight_dtype)
    image_encoder_g.to(accelerator.device, dtype=weight_dtype)

    # We need to recalculate our total training steps as the size of the training dataloader may have changed.
    num_update_steps_per_epoch = math.ceil(len(train_dataloader) / args.gradient_accumulation_steps)
    if overrode_max_train_steps:
        args.max_train_steps = args.num_train_epochs * num_update_steps_per_epoch
    # Afterwards we recalculate our number of training epochs
    args.num_train_epochs = math.ceil(args.max_train_steps / num_update_steps_per_epoch)


    args.num_train_epochs = train_steps


    # Train!
    total_batch_size = (
            train_batch_size
            * accelerator.num_processes
            * args.gradient_accumulation_steps
    )

    print("***** Running training *****")
    print(f"  Num batches each epoch = {len(train_dataloader)}")
    print(f"  Num Epochs = {args.num_train_epochs}")
    print(f"  Instantaneous batch size per device = {train_batch_size}")
    print(
        f"  Total train batch size (w. parallel, distributed & accumulation) = {total_batch_size}"
    )
    print(f"  Gradient Accumulation steps = {args.gradient_accumulation_steps}")
    print(f"  Total optimization steps = {args.max_train_steps}")


    if args.resume_from_checkpoint:
        # New Code #
        # Loads the DeepSpeed checkpoint from the specified path
        prior_model, last_epoch, last_global_step = load_training_checkpoint(
            sd_model,
            pcdms_model,
            **{"load_optimizer_states": True, "load_lr_scheduler_states": True},
        )
        print(f"Resumed from checkpoint: {args.resume_from_checkpoint}, global step: {last_global_step}")
        starting_epoch = last_epoch
        global_steps = last_global_step
        sd_model = sd_model
    else:
        global_steps = 0
        starting_epoch = 0
        sd_model = sd_model

    progress_bar = tqdm(range(global_steps, args.max_train_steps), initial=global_steps, desc="Steps",
                        # Only show the progress bar once on each machine.
                        disable=not accelerator.is_local_main_process, )

    bsz = train_batch_size
    
    if not finetune or train_steps == 0:
        accelerator.wait_for_everyone()
        accelerator.end_training()
        return {k: v.cpu() for k, v in sd_model.state_dict().items()}


    it = range(starting_epoch, args.num_train_epochs)
    if is_app:
        it = progress.tqdm(it, desc="Fine-tuning")
    for epoch in it:
        for step, batch in enumerate(train_dataloader):
            with accelerator.accumulate(sd_model):
                with torch.no_grad():
                    # Convert images to latent space
                    latents = vae.encode(batch["source_target_image"].to(dtype=weight_dtype)).latent_dist.sample()
                    latents = latents * vae.config.scaling_factor

                    # Get the masked image latents
                    masked_latents = vae.encode(batch["vae_source_mask_image"].to(dtype=weight_dtype)).latent_dist.sample()
                    masked_latents = masked_latents * vae.config.scaling_factor

                    bsz = batch["target_image"].size(dim=0)

                    # mask
                    mask1 = torch.ones((bsz, 1, int(args.img_height / 8), int(args.img_width / 8))).to(accelerator.device, dtype=weight_dtype)
                    mask0 = torch.zeros((bsz, 1, int(args.img_height / 8), int(args.img_width / 8))).to(accelerator.device, dtype=weight_dtype)
                    mask = torch.cat([mask1, mask0], dim=3)
                    # Get the image embedding for conditioning
                    cond_image_feature_p = image_encoder_p(batch["source_image"].to(accelerator.device, dtype=weight_dtype))
                    cond_image_feature_p = (cond_image_feature_p.last_hidden_state)


                    cond_image_feature_g = image_encoder_g(batch["target_image"].to(accelerator.device, dtype=weight_dtype), ).image_embeds
                    cond_image_feature_g =cond_image_feature_g.unsqueeze(1)

                # Sample noise that we'll add to the latents
                noise = torch.randn_like(latents)
                if args.noise_offset:
                    # https://www.crosslabs.org//blog/diffusion-with-offset-noise
                    noise += args.noise_offset * torch.randn(
                        (latents.shape[0], latents.shape[1], 1, 1), device=latents.device
                    )

                # Sample a random timestep for each image
                #timesteps = torch.randint(0, noise_scheduler.config.num_train_timesteps, (train_batch_size,),device=latents.device, )
                timesteps = torch.randint(0, noise_scheduler.config.num_train_timesteps, (bsz,),device=latents.device, )
                timesteps = timesteps.long()

                

                # Add noise to the latents according to the noise magnitude at each timestep (this is the forward diffusion process)
                noisy_latents = noise_scheduler.add_noise(latents, noise, timesteps)

                #print(noisy_latents.size(), mask.size(), masked_latents.size())
                
                noisy_latents = torch.cat([noisy_latents, mask, masked_latents], dim=1)
                # Get the text embedding for conditioning
                

                cond_pose = batch["source_target_pose"].to(dtype=weight_dtype)
                
                #print(noisy_latents.size())
                #print(cond_image_feature_p.size())
                #print(cond_image_feature_g.size())
                #print(cond_pose.size())

                # Predict the noise residual
                model_pred = sd_model(noisy_latents, timesteps, cond_image_feature_p,cond_image_feature_g, cond_pose, )

                # Get the target for loss depending on the prediction type
                if noise_scheduler.config.prediction_type == "epsilon":
                    target = noise
                elif noise_scheduler.config.prediction_type == "v_prediction":
                    target = noise_scheduler.get_velocity(latents, noise, timesteps)
                else:
                    raise ValueError(
                        f"Unknown prediction type {noise_scheduler.config.prediction_type}"
                    )

                loss = F.mse_loss(model_pred.float(), target.float(), reduction="mean")

                accelerator.backward(loss)
                if accelerator.sync_gradients:
                    params_to_clip = sd_model.parameters()
                    accelerator.clip_grad_norm_(params_to_clip, args.max_grad_norm)
                optimizer.step()
                lr_scheduler.step()
                optimizer.zero_grad(set_to_none=args.set_grads_to_none)

            # Checks if the accelerator has performed an optimization step behind the scenes
            if accelerator.sync_gradients:
                global_steps += 1

            if global_steps >= args.max_train_steps:
                break
            

            logs = {"loss": loss.detach().item(), "lr": lr_scheduler.get_last_lr()[0]}
            print(logs)
            progress_bar.set_postfix(**logs)

        progress_bar.update(1)

    # Create the pipeline using  the trained modules and save it.
    accelerator.wait_for_everyone()
    accelerator.end_training()

    sd_model.unet.cpu()
    sd_model.cpu()
    del vae
    del image_encoder_p
    del image_encoder_g

    if save_model: #if global_steps % args.checkpointing_steps == 0 or global_steps == args.max_train_steps:
        print('saving', modelId)
        
        checkpoint_state_dict = {
            "epoch": 0,
            "module": {k: v.cpu() for k, v in sd_model.state_dict().items()}, #sd_model.state_dict(),
        }
        print(list(sd_model.state_dict().keys())[:20])
        torch.save(checkpoint_state_dict, modelId+".pt")
        
        del sd_model
        gc.collect()
        torch.cuda.empty_cache()
        print('done train')
        print(torch.cuda.memory_allocated()/1024**2)
        return
    
    del sd_model
    gc.collect()
    torch.cuda.empty_cache()    
    return {k: v.cpu() for k, v in sd_model.state_dict().items()}




# Pose-transfer ===================================================================================================


device = "cuda"

class ImageProjModel(torch.nn.Module):
    """SD model with image prompt"""
    def __init__(self, in_dim, hidden_dim, out_dim, dropout = 0.):
        super().__init__()

        self.net = nn.Sequential(
            nn.Linear(in_dim, hidden_dim),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.LayerNorm(hidden_dim),
            nn.Linear(hidden_dim, out_dim),
            nn.Dropout(dropout)
        )

    def forward(self, x):  
        return self.net(x)

def image_grid(imgs, rows, cols):
    assert len(imgs) == rows * cols
    w, h = imgs[0].size
    print(w, h)
    grid = Image.new("RGB", size=(cols * w, rows * h))
    grid_w, grid_h = grid.size

    for i, img in enumerate(imgs):
        grid.paste(img, box=(i % cols * w, i // cols * h))
    return grid

def load_mydict(modelId, finetuned_model):
    if save_model:
        model_ckpt_path = modelId+'.pt'
        model_sd = torch.load(model_ckpt_path, map_location="cpu")["module"]
    else:
        model_sd = finetuned_model #torch.load(model_ckpt_path, map_location="cpu")["module"]

    image_proj_model_dict = {}
    pose_proj_dict = {}
    unet_dict = {}
    for k in model_sd.keys():
        if k.startswith("pose_proj"):
            pose_proj_dict[k.replace("pose_proj.", "")] = model_sd[k]

        elif k.startswith("image_proj_model_p"):
            image_proj_model_dict[k.replace("image_proj_model_p.", "")] = model_sd[k]
        elif k.startswith("image_proj_model"):
            image_proj_model_dict[k.replace("image_proj_model.", "")] = model_sd[k]


        elif k.startswith("unet"):
            unet_dict[k.replace("unet.", "")] = model_sd[k]
        else:
            print(k)
    return image_proj_model_dict, pose_proj_dict, unet_dict



@spaces.GPU(duration=600)
def inference(modelId, in_image, in_pose, target_poses, inference_steps, finetuned_model, vae, unet, image_encoder, is_app=False):
    print('start inference')
    progress=gr.Progress(track_tqdm=True)
    
    if not save_model:
        finetuned_model = {k: v.cuda() for k, v in finetuned_model.items()}
    
    device = "cuda"
    pretrained_model_name_or_path ="stabilityai/stable-diffusion-2-1-base"
    image_encoder_path = "facebook/dinov2-giant"
    #model_ckpt_path = "./pcdms_ckpt.pt"   # ckpt path
    model_ckpt_path = modelId+'.pt'


    clip_image_processor = CLIPImageProcessor()
    img_transform = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize([0.5], [0.5]),
    ])

    generator = torch.Generator(device=device).manual_seed(42)
    
    """
    unet = Stage2_InapintUNet2DConditionModel.from_pretrained(pretrained_model_name_or_path, torch_dtype=torch.float16,subfolder="unet",in_channels=9, low_cpu_mem_usage=False, ignore_mismatched_sizes=True).to(device)
    vae = AutoencoderKL.from_pretrained(pretrained_model_name_or_path,subfolder="vae").to(device, dtype=torch.float16)
    image_encoder = Dinov2Model.from_pretrained(image_encoder_path).to(device, dtype=torch.float16)
    """
    noise_scheduler = DDIMScheduler(
        num_train_timesteps=1000,
        beta_start=0.00085,
        beta_end=0.012,
        beta_schedule="scaled_linear",
        clip_sample=False,
        set_alpha_to_one=False,
        steps_offset=1,
    )
    
    unet = unet.to(device, dtype=torch.float16)
    vae = vae.to(device, dtype=torch.float16)
    image_encoder = image_encoder.to(device, dtype=torch.float16)
    

    image_proj_model = ImageProjModel(in_dim=1536, hidden_dim=768, out_dim=1024).to(device).to(dtype=torch.float16)
    pose_proj_model = ControlNetConditioningEmbedding(
        conditioning_embedding_channels=320,
        block_out_channels=(16, 32, 96, 256),
        conditioning_channels=3).to(device).to(dtype=torch.float16)


    # load weight
    print('loading', modelId)
    image_proj_model_dict, pose_proj_dict, unet_dict = load_mydict(modelId, finetuned_model)
    print('loaded', modelId)
    image_proj_model.load_state_dict(image_proj_model_dict)
    pose_proj_model.load_state_dict(pose_proj_dict)
    unet.load_state_dict(unet_dict)
    
    
    pipe = PCDMsPipeline.from_pretrained("stabilityai/stable-diffusion-2-1-base", unet=unet,  torch_dtype=torch.float16, scheduler=noise_scheduler,feature_extractor=None,safety_checker=None).to(device)

    print('====================== model load finish ===================')

    results = []
    progress_bar = tqdm(range(len(target_poses)), initial=0, desc="Frames")
    
    
    it = target_poses
    if is_app:
        it = progress.tqdm(it, desc="Pose Transfer")
    for pose in it:

        num_samples = 1
        image_size = (512, 512)
        s_img_path = 'imgs/'+input_img # input image 1
        #target_pose_img = 'imgs/pose_'+str(n)+'.png' # input image 2
        
        #t_pose = inference_pose(target_pose_img, image_size=(image_size[1], image_size[0])).resize(image_size, Image.BICUBIC)
        #t_pose = Image.open(target_pose_img).convert("RGB").resize((image_size), Image.BICUBIC)
        t_pose = pose.convert("RGB").resize((image_size), Image.BICUBIC)
        #t_pose = resize_and_pad(pose.convert("RGB"))


        #s_img = Image.open(s_img_path)
        width_orig, height_orig = in_image.size
        s_img = in_image.convert("RGB").resize(image_size, Image.BICUBIC)
        #s_img = resize_and_pad(in_image.convert("RGB"))
        black_image = Image.new("RGB", s_img.size, (0, 0, 0)).resize(image_size, Image.BICUBIC)

        s_img_t_mask = Image.new("RGB", (s_img.width * 2, s_img.height))
        s_img_t_mask.paste(s_img, (0, 0))
        s_img_t_mask.paste(black_image, (s_img.width, 0))

        #s_pose = inference_pose(s_img_path, image_size=(image_size[1], image_size[0])).resize(image_size, Image.BICUBIC)
        #s_pose = Image.open('imgs/sm_pose.jpg').convert("RGB").resize(image_size, Image.BICUBIC)
        s_pose = in_pose.convert("RGB").resize(image_size, Image.BICUBIC)
        #s_pose = resize_and_pad(in_pose.convert("RGB"))
        print('source image width: {}, height: {}'.format(s_pose.width, s_pose.height))
        #t_pose = Image.open(target_pose_img).convert("RGB").resize((image_size), Image.BICUBIC)

        st_pose = Image.new("RGB", (s_pose.width * 2, s_pose.height))
        st_pose.paste(s_pose, (0, 0))
        st_pose.paste(t_pose, (s_pose.width, 0))


        clip_s_img = clip_image_processor(images=s_img, return_tensors="pt").pixel_values
        vae_image = torch.unsqueeze(img_transform(s_img_t_mask), 0)
        cond_st_pose = torch.unsqueeze(img_transform(st_pose), 0)

        mask1 = torch.ones((1, 1, int(image_size[0] / 8), int(image_size[1] / 8))).to(device, dtype=torch.float16)
        mask0 = torch.zeros((1, 1, int(image_size[0] / 8), int(image_size[1] / 8))).to(device, dtype=torch.float16)
        mask = torch.cat([mask1, mask0], dim=3)


        with torch.inference_mode():
            cond_pose = pose_proj_model(cond_st_pose.to(dtype=torch.float16, device=device))
            simg_mask_latents = pipe.vae.encode(vae_image.to(device, dtype=torch.float16)).latent_dist.sample()
            simg_mask_latents = simg_mask_latents * 0.18215

            images_embeds = image_encoder(clip_s_img.to(device, dtype=torch.float16)).last_hidden_state
            image_prompt_embeds = image_proj_model(images_embeds)
            uncond_image_prompt_embeds = image_proj_model(torch.zeros_like(images_embeds))

        bs_embed, seq_len, _ = image_prompt_embeds.shape
        image_prompt_embeds = image_prompt_embeds.repeat(1, num_samples, 1)
        image_prompt_embeds = image_prompt_embeds.view(bs_embed * num_samples, seq_len, -1)
        uncond_image_prompt_embeds = uncond_image_prompt_embeds.repeat(1, num_samples, 1)
        uncond_image_prompt_embeds = uncond_image_prompt_embeds.view(bs_embed * num_samples, seq_len, -1)

        output, _ = pipe(
            simg_mask_latents= simg_mask_latents,
            mask = mask,
            cond_pose = cond_pose,
            prompt_embeds=image_prompt_embeds,
            negative_prompt_embeds=uncond_image_prompt_embeds,
            height=image_size[1],
            width=image_size[0]*2,
            num_images_per_prompt=num_samples,
            guidance_scale=2.0,
            generator=generator,
            num_inference_steps=inference_steps,
        )

        output = output.images[-1]
        
        result = output.crop((image_size[0], 0, image_size[0] * 2, image_size[1]))
        result = result.resize((width_orig, height_orig), Image.BICUBIC)
        #result = remove_zero_pad(result)
        
        if debug:
            result.save('out/'+str(len(results))+'.png')
        results.append(result)
        progress_bar.update(1)

    del unet
    del vae
    del image_encoder
    del image_proj_model
    del pose_proj_model

    if not save_model:
        del finetuned_model

    gc.collect()
    torch.cuda.empty_cache()
    print(torch.cuda.memory_allocated()/1024**2)
        
    return results
        

def gen_vid(frames, video_name, fps, codec):
    progress=gr.Progress(track_tqdm=True)
    
    frame = cv2.cvtColor(np.array(frames[0]), cv2.COLOR_RGB2BGR)
    height, width, layers = frame.shape

    #video = cv2.VideoWriter(video_name, 0, 1, (width,height))
    if codec == 'mp4':
        video = cv2.VideoWriter(video_name, cv2.VideoWriter_fourcc(*'mp4v'), fps, (width, height))
    else:
        video = cv2.VideoWriter(video_name, cv2.VideoWriter_fourcc(*'VP90'), fps, (width, height))

    for r in progress.tqdm(frames, desc="Creating video"):
        image = cv2.cvtColor(np.array(r), cv2.COLOR_RGB2BGR)
        video.write(image)

    #cv2.destroyAllWindows()
    #video.release()
    


def run(images, video_path, train_steps=100, inference_steps=10, fps=12, bg_remove=False, resize_inputs=True, finetune=True, is_app=False):
    print("==== Load Models ====")
    dwpose, rembg_session, pcdms_model, noise_scheduler, image_encoder_p, image_encoder_g, vae, unet = load_models()
    
    print("==== Pose Detection ====")
    in_img, in_pose, train_imgs, train_poses, target_poses = prepare_inputs(images, video_path, fps, bg_remove, dwpose, rembg_session, resize_inputs, is_app=is_app)
    
    if save_model:
        train("fine_tuned_pcdms", in_img, in_pose, train_imgs, train_poses, train_steps, pcdms_model, noise_scheduler, image_encoder_p, image_encoder_g, vae, unet, finetune, is_app)
        print('next')
        results = inference("fine_tuned_pcdms", in_img, in_pose, target_poses, inference_steps, None, vae, unet, image_encoder_p, is_app)
        
    else:    
        print("==== Finetuning ====")
        finetuned_model = train("fine_tuned_pcdms", in_img, in_pose, train_imgs, train_poses, train_steps, pcdms_model, noise_scheduler, image_encoder_p, image_encoder_g, vae, unet, finetune, is_app)
        
        print("==== Pose Transfer ====")
        results = inference("fine_tuned_pcdms", in_img, in_pose, target_poses, inference_steps, finetuned_model, vae, unet, image_encoder_p, is_app)

    return results


def run_train(images, train_steps=100, modelId="fine_tuned_pcdms", bg_remove=True, resize_inputs=True):
    finetune=True
    is_app=True
    images = [img[0] for img in images]
    
    dwpose, rembg_session, pcdms_model, noise_scheduler, image_encoder_p, image_encoder_g, vae, unet = load_models()
    
    if resize_inputs:
        resize = 'target'
    else:
        resize = 'none'
        
    in_img, in_pose, train_imgs, train_poses = prepare_inputs_train(images, bg_remove, dwpose, rembg_session)
    
    train(modelId, in_img, in_pose, train_imgs, train_poses, train_steps, pcdms_model, noise_scheduler, image_encoder_p, image_encoder_g, vae, unet, finetune, is_app)


def run_inference(images, video_path, train_steps=100, inference_steps=10, fps=12, modelId="fine_tuned_pcdms", img_width=1920, img_height=1080, bg_remove=True, resize_inputs=True):
    finetune=True
    is_app=True
    
    
    dwpose, rembg_session, pcdms_model, noise_scheduler, image_encoder_p, image_encoder_g, vae, unet = load_models()
    
    if not os.path.exists(modelId+".pt"):
        run_train(images, train_steps, modelId, bg_remove, resize_inputs)
    
    images = [img[0] for img in images]
    in_img = images[0]

    in_img, target_poses, in_pose, target_poses_coords, orig_frames = prepare_inputs_inference(in_img, video_path, fps, dwpose, rembg_session, bg_remove, resize_inputs, is_app)
    
    results = inference(modelId, in_img, in_pose, target_poses, inference_steps, None, vae, unet, image_encoder_p, is_app)
    #urls = save_temp_imgs(results)

    if should_gen_vid:
        if debug:
            gen_vid(results, out_vid+'.mp4', fps, 'mp4')
        else:
            gen_vid(results, out_vid+'.webm', fps, 'webm')
    
    
    # postprocessing
    results = [removebg(img, rembg_session, True) for img in results]
    #results = [img_pad(img, img_width, img_height, True) for img in results]
    
    print("Done!")
    
    return out_vid+'.webm', results, getThumbnails(results), target_poses_coords, getThumbnails(orig_frames)
    

def run_generate_frame(images, target_poses, train_steps=100, inference_steps=10, modelId="fine_tuned_pcdms", img_width=1920, img_height=1080, bg_remove=True, resize_inputs=True):
    finetune=True
    is_app=True

    print(target_poses)
    target_poses = json.loads(target_poses)
    target_poses = [draw_openpose(pose, height=img_height, width=img_width, include_hands=True, include_face=False) for pose in target_poses]
    target_poses[0].save('test.png')
    
    dwpose, rembg_session, pcdms_model, noise_scheduler, image_encoder_p, image_encoder_g, vae, unet = load_models()
    
    if not os.path.exists(modelId+".pt"):
        run_train(images, train_steps, modelId, bg_remove, resize_inputs)
    
    images = [img[0] for img in images]
    in_img = images[0]
    in_pose, _ = get_pose(in_img, dwpose, "in_pose.png")

    results = inference(modelId, in_img, in_pose, target_poses, inference_steps, None, vae, unet, image_encoder_p, is_app)
    #urls = save_temp_imgs(results)
    
    # postprocessing
    results = [removebg(img, rembg_session, True) for img in results]
    #results = [img_pad(img, img_width, img_height, True) for img in results]
    
    print("Done!")
    
    return results, getThumbnails(results)


def run_app(images, video_path, train_steps=100, inference_steps=10, fps=12, bg_remove=False, resize_inputs=True):
    
    images = [img[0] for img in images]
    
    results = run(images, video_path, train_steps, inference_steps, fps, bg_remove, resize_inputs, finetune=True, is_app=True)
    
    
    print("==== Video generation ====")
    out_vid = f"out_{uuid.uuid4()}"
    
    if debug:
        gen_vid(results, out_vid+'.mp4', fps, 'mp4')
    else:
        gen_vid(results, out_vid+'.webm', fps, 'webm')
    
    
    
    print("Done!")
    
    return out_vid+'.webm', results



"""
train_steps = 100
inference_steps = 10
fps = 12
"""

"""
iface = gr.Interface(
    fn=run,
    inputs=[
        gr.Gallery(type="pil", label="Images of the Character"),
        gr.Video(label="Motion-Capture Video"),
        gr.Number(label="Training steps", value=100),
        gr.Number(label="Inference steps", value=10),
        gr.Number(label="Output frame rate", value=12),
        gr.Checkbox(label="Remove background", value=False),
    ],
    outputs=[gr.Video(label="Result"), gr.Gallery(type="pil", label="Frames")],
    title="Keyframes AI",
    description="Upload images of your character and a motion-capture video to generate an animation of the character.",
)
"""