pad results to specified resolution

app.py

@@ -1,44 +1,1196 @@
-from main import run_app, run_train, run_inference
+import logging
+import math
+import os
+from typing import Any, Dict, List, Optional, Tuple, Union
+from diffusers.models.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 easy_dwpose import DWposeDetector
 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 uuid
+
+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="cpu")
+    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 = 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
+
+
+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
+
 
-with gr.Blocks() as demo:
-    with gr.Row():
-        with gr.Column():
-            char_imgs = gr.Gallery(type="pil", label="Images of the Character")
-            mocap = gr.Video(label="Motion-Capture Video")
-            tr_steps = gr.Number(label="Training steps", value=10)
-            inf_steps = gr.Number(label="Inference steps", value=10)
-            fps = gr.Number(label="Output frame rate", value=12)
-            modelId = gr.Text(label="Model Id", value="fine_tuned_pcdms")
-            remove_bg = gr.Checkbox(label="Remove background", value=False)
-            resize_inputs = gr.Checkbox(label="Resize images to match video", value=True)
-            train_btn = gr.Button(value="Train")
-            inference_btn = gr.Button(value="Inference")
-            submit_btn = gr.Button(value="Generate")
-        with gr.Column():
-            animation = gr.Video(label="Result")
-            frames = gr.Gallery(type="pil", label="Frames", format="png")
-            frames_thumb = gr.Gallery(type="pil", label="Thumbnails", format="png")
+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]
 
-    submit_btn.click(
-        run_app, inputs=[char_imgs, mocap, tr_steps, inf_steps, fps, remove_bg, resize_inputs], outputs=[animation, frames]
+    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_poses.append(get_pose(img, dwpose, "tr_pose"+str(i)+".png"))
+        
+    return in_img, in_pose, train_imgs, train_poses
+        
+
+def prepare_inputs_inference(in_img, in_vid, fps, dwpose, rembg_session, bg_remove, resize='target', 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 = []
+    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 = get_pose(f, dwpose, "tar_pose"+str(len(target_poses))+".png")
+        target_poses.append(tpose)
+        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=='target':
+            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
+
+
+def prepare_inputs(images, in_vid, fps, bg_remove, dwpose, rembg_session, resize='target', 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, 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)
     
-    train_btn.click(
-        run_train, inputs=[char_imgs, tr_steps, modelId, remove_bg, resize_inputs], outputs=[]
+    
+    """
+    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,
     )
     
-    inference_btn.click(
-        run_inference, inputs=[char_imgs, mocap, tr_steps, inf_steps, fps, modelId, remove_bg, resize_inputs], outputs=[animation, frames, frames_thumb]
+    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 ====")
+    if resize_inputs:
+        resize = 'target'
+    else:
+        resize = 'none'
+    in_img, in_pose, train_imgs, train_poses, target_poses = prepare_inputs(images, video_path, fps, bg_remove, dwpose, rembg_session, resize=resize, 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 = prepare_inputs_inference(in_img, video_path, fps, dwpose, rembg_session, bg_remove, 'target', 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)
+    
+
+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
+"""
 
-demo.launch(share=True)
+"""
+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.",
+)
+"""
 
 
 

main.py

@@ -167,6 +167,17 @@ def load_models():
 
 #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
@@ -180,12 +191,8 @@ def resize_and_pad(img, target_img):
     img = img.resize((tw, th), Image.BICUBIC)
 
     tw, th = target_img.size
-    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
+    return img_pad(img, tw, th)
 
 
 def remove_zero_pad(image):
@@ -1105,7 +1112,7 @@ def run_train(images, train_steps=100, modelId="fine_tuned_pcdms", bg_remove=Tru
     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", bg_remove=True, resize_inputs=True):
+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
     
@@ -1130,7 +1137,9 @@ def run_inference(images, video_path, train_steps=100, inference_steps=10, fps=1
             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!")