human_img_edit_gradio.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "473523df",
   "metadata": {},
   "source": [
    "### Example Jupyter Code for Generation Results with Gradio Demo\n",
    "\n",
    "```\n",
    "--------------------------------------------------------\n",
    "DisCo - Disentangled Control for Referring Human Dance Generation in Real World\n",
    "Licensed under The Apache-2.0 license License [see LICENSE for details]\n",
    "Tan Wang (TAN317@e.ntu.edu.sg)\n",
    "Work done during internship at Microsoft\n",
    "--------------------------------------------------------\n",
    "```\n",
    "\n",
    "Pls remember to change the path (model checkpoint, root dir, eval_save_filename, and so on) in the `manual_args`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "fdbf0c37",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/wangtan/.local/lib/python3.8/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
      "  from .autonotebook import tqdm as notebook_tqdm\n",
      "[2023-07-03 08:19:31 <wutils_ldm.py:150> <module>] <utils.py>: Deep Learning Utils @ Chenfei Wu\u001b[0m\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "WANDB_ENABLE: 0\n",
      "2023-07-03 08:19:31,989.989 262807:common.py:1785 setup_yaml(): python 3 env\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[2023-07-03 08:19:32 <wutils.py:132> <module>] <utils.py>: Deep Learning Utils @ Chenfei Wu\u001b[0m\n",
      "[2023-07-03 08:19:32 <wutils.py:132> <module>] <utils.py>: Deep Learning Utils @ Chenfei Wu\u001b[0m\n",
      "[2023-07-03 08:19:32 <__init__.py:71> BasicArgs] Detected unknown Node MICL-ZhangHWSvr1.\u001b[0m\n",
      "[2023-07-03 08:19:32 <__init__.py:71> BasicArgs] Detected unknown Node MICL-ZhangHWSvr1.\u001b[0m\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"0\"\n",
    "os.environ[\"WANDB_ENABLE\"] = \"0\"\n",
    "\n",
    "from utils.wutils_ldm import *\n",
    "from agent import Agent_LDM, WarmupLinearLR, WarmupLinearConstantLR\n",
    "import torch\n",
    "from config import BasicArgs\n",
    "from utils.lib import *\n",
    "# from utils.args import parse_with_cf\n",
    "from utils.dist import dist_init\n",
    "from dataset.tsv_dataset import make_data_sampler, make_batch_data_sampler\n",
    "from finetune_sdm_yaml import get_loader_info, make_data_loader\n",
    "torch.multiprocessing.set_sharing_strategy('file_system')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a8735d83",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "f0fcd501",
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2023-07-03 08:19:32,623.623 262807:__init__.py:51 _is_triton_available(): A matching Triton is not available, some optimizations will not be enabled.\n",
      "Error caught was: No module named 'triton'\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[2023-07-03 08:19:32 <507173574.py:71> <module>] Building models...\u001b[0m\n",
      "[2023-07-03 08:19:32 <507173574.py:71> <module>] Building models...\u001b[0m\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "distributed training ... presumbly debug with 1 GPU\n",
      "Using seed 42 for rank 0\n",
      "Using seed 42 for torch.cuda\n",
      "Loading pre-trained image_encoder from /home1/wangtan/code/ms_internship2/github_repo/run_test/diffusers/sd-image-variations-diffusers/image_encoder\n",
      "Loading pre-trained vae from /home1/wangtan/code/ms_internship2/github_repo/run_test/diffusers/sd-image-variations-diffusers/vae\n",
      "Loading pre-trained unet from /home1/wangtan/code/ms_internship2/github_repo/run_test/diffusers/sd-image-variations-diffusers/unet\n",
      "use sd vae to init the controlnet condition embedding\n",
      "Args: {'root_dir': '/home1/wangtan/code/ms_internship2/github_repo/run_test', 'cf': 'config/ref_attn_clip_combine_controlnet/app_demo_image_edit.py', 'pretrained_model': '/home/wangtan/data/disco/ft_checkpoint/moretiktok_nocfg/mp_rank_00_model_states.pt', 'pretrained_model_lora': None, 'pretrained_model_controlnet': None, 'debug': False, 'debug_seed': False, 'debug_dataloader': False, 'log_dir': '/home1/wangtan/code/ms_internship2/github_repo/run_test/exp/tiktok_ft', 'deepspeed': True, 'use_amp': False, 'seed': 42, 'fix_dist_seed': True, 'tiktok_data_root': 'keli/dataset/TikTok_dataset/', 'img_size': [256, 256], 'max_video_len': 1, 'debug_max_video_len': 1, 'conds': ['poses', 'masks'], 'gradient_checkpointing': True, 'find_unused_parameters': False, 'enable_xformers_memory_efficient_attention': True, 'trainable_modules': None, 'scale_factor': 0.18215, 'loss_target': 'noise', 'x0_steps': 200, 'pretrained_model_path': '/home1/wangtan/code/ms_internship2/github_repo/run_test/diffusers/sd-image-variations-diffusers', 'num_workers': 4, 'node_split_sampler': False, 'gradient_accumulate_steps': 1, 'max_grad_norm': -1, 'learning_rate': 0.0002, 'decay': 0.001, 'warmup_ratio': 0.1, 'max_train_samples': None, 'debug_max_train_samples': 100, 'drop_text': 1.0, 'local_train_batch_size': 32, 'epochs': 20, 'eval_step': 500.0, 'save_step': 500.0, 'do_train': False, 'train_yaml': './blob_dir/debug_output/video_sythesis/dataset/composite/train_webvid10m_a_54.yaml', 'resume': False, 'null_caption': False, 'refer_sdvae': True, 'controlnet_conditioning_scale_cond': 1.0, 'controlnet_conditioning_scale_ref': 1.0, 'unet_unfreeze_type': 'all', 'controlnet_attn': False, 'use_cfg': False, 'refer_clip_preprocess': False, 'refer_clip_proj': False, 'ref_null_caption': False, 'combine_clip_local': True, 'combine_use_mask': True, 'drop_ref': 0.0, 'my_adapter': False, 'pos_resize_img': False, 'fg_variation': 0.0, 'strong_aug_stage2': False, 'strong_rand_stage2': False, 'strong_aug_stage1': False, 'stage1_pretrain_path': None, 'stage2_only_pose': False, 'constant_lr': False, 'SD2_not_add_image_emb_noise': False, 'val_yaml': './blob_dir/debug_output/video_sythesis/dataset/composite/val_webvid10m_a.yaml', 'max_eval_samples': None, 'debug_max_eval_samples': 20, 'pose_normalize': False, 'normalize_by_1st_frm': False, 'local_eval_batch_size': 32, 'eval_visu': True, 'eval_visu_trainsample': False, 'eval_visu_imagefolder': False, 'eval_visu_changepose': False, 'eval_visu_changefore': False, 'eval_save_filename': 'try', 'eval_before_train': True, 'eval_scheduler': 'ddim', 'eval_enc_dec_only': False, 'num_inf_videos_per_prompt': 1, 'num_inference_steps': 50, 'guidance_scale': 3.0, 'stepwise_sample_depth': -1, 'interpolation': None, 'interpolate_mode': None, 'visu_save': False, 'freeze_pose': False, 'freeze_background': False, 'ft_img_num': 0, 'ft_one_ref_image': True, 'ft_iters': None, 's1': 1.0, 's2': 1.0, 'ft_idx': None, 'ref_mode': 'first', '__module__': 'mymodule', 'task_name': 'ref_attn_clip_combine_controlnet', 'method_name': 'app_demo_image_edit', 'dataset_cf': 'dataset/app_demo_human_image_edit_singleinput.py', 'img_full_size': [256, 256], 'fps': 5, 'data_dir': './blob_dir/data/mtp_vlp_ray/debug/debug_pretrain', 'debug_train_yaml': './blob_dir/data/mtp_vlp_ray/debug/debug_pretrain/composite/train_webvid2.5m_2.yaml', 'debug_val_yaml': './blob_dir/data/mtp_vlp_ray/debug/debug_pretrain/composite/val_webvid2.5m.yaml', 'web_data_root': '/datadrive_d/wangtan/azure_storage/vigstandard_data/linjli/debug_output/video_sythesis/dataset/Lindsey_0504_youtube/frames/single_person', 'sd15_path': '/home1/wangtan/code/ms_internship2/github_repo/run_test/diffusers/stable-diffusion-v1-5-2', 'freeze_unet': True, 'num_inf_images_per_prompt': 1, '__doc__': None, 'n_gpu': 1, 'local_size': 1, 'num_gpus': 1, 'distributed': True, 'num_nodes': 1, 'word_size': 1, 'local_rank': 0, 'rank': 0, 'node_id': 0, 'dist': True, 'nodes': 1, 'world_size': 1, 'train_batch_size': 32, 'eval_batch_size': 32}\n"
     ]
    }
   ],
   "source": [
    "from utils.args import sharedArgs\n",
    "manual_args = ['--cf', 'config/ref_attn_clip_combine_controlnet/app_demo_image_edit.py', '--eval_visu', 'True', '--root_dir', '/home1/wangtan/code/ms_internship2/github_repo/run_test', '--local_train_batch_size', '32', '--local_eval_batch_size', '32', '--log_dir', 'exp/tiktok_ft', '--epochs', '20', '--deepspeed', '--eval_step', '500',\n",
    "               '--save_step', '500', '--gradient_accumulate_steps', '1', '--learning_rate', '2e-4', '--fix_dist_seed', 'True', '--loss_target',\n",
    "               'noise', '--unet_unfreeze_type', 'all', '--guidance_scale', '3', '--refer_sdvae', 'True', '--ref_null_caption', 'False', '--combine_clip_local', 'True', '--combine_use_mask', 'True', '--conds', 'poses','masks', '--pretrained_model', '/home/wangtan/data/disco/ft_checkpoint/moretiktok_nocfg/mp_rank_00_model_states.pt', '--eval_save_filename', 'try']\n",
    "parsed_args = sharedArgs.parser.parse_args(args=manual_args)\n",
    "# args = sharedArgs.parser.parse_args(args=['--cf', 'config/ref_attn_clip_combine_controlnet/app_demo_image_edit.py', '--eval_visu', 'True'])\n",
    "# import pdb; pdb.set_trace()\n",
    "\n",
    "###### process the args #######\n",
    "if parsed_args.root_dir:\n",
    "    BasicArgs.root_dir = parsed_args.root_dir\n",
    "else:\n",
    "    parsed_args.root_dir = BasicArgs.root_dir\n",
    "parsed_args.pretrained_model_path = os.path.join(parsed_args.root_dir, parsed_args.pretrained_model_path)\n",
    "\n",
    "def parse_with_cf(parsed_args):\n",
    "    \"\"\"This function will set args based on the input config file.\n",
    "    (1) it only overwrites unset parameters,\n",
    "        i.e., these parameters not set from user command line input\n",
    "    (2) it also sets configs in the config file but declared in the parser\n",
    "    \"\"\"\n",
    "    # convert to EasyDict object,\n",
    "    # enabling access from attributes even for nested config\n",
    "    # e.g., args.train_datasets[0].name\n",
    "    args = edict(vars(parsed_args))\n",
    "    if os.path.exists(parsed_args.cf):\n",
    "        cf = import_filename(parsed_args.cf)\n",
    "        config_args = edict(vars(cf.Args))\n",
    "        override_keys = {arg[2:].split(\"=\")[0] for arg in manual_args\n",
    "                         if arg.startswith(\"--\")}\n",
    "        # import pdb;pdb.set_trace()\n",
    "        for k, v in config_args.items():\n",
    "            if k not in override_keys:\n",
    "                setattr(args, k, v)\n",
    "    else:\n",
    "        raise NotImplementedError('Config filename %s does not exist.' % args.cf)\n",
    "    return args\n",
    "\n",
    "args = parse_with_cf(parsed_args)\n",
    "        \n",
    "args.n_gpu = T.cuda.device_count() # local size\n",
    "args.local_size = args.n_gpu\n",
    "if args.root_dir not in args.log_dir:\n",
    "    args.log_dir = os.path.join(args.root_dir, args.log_dir)\n",
    "if args.stepwise_sample_depth == -1:\n",
    "    args.interpolation = None\n",
    "    args.interpolate_mode = None\n",
    "if args.interpolation != \"interpolate\":\n",
    "    args.interpolate_mode = None\n",
    "\n",
    "assert args.eval_step > 0, \"eval_step must be positive\"\n",
    "assert args.save_step > 0, \"save_step must be positive\"\n",
    "\n",
    "dist_init(args)\n",
    "args.dist = args.distributed\n",
    "args.nodes = args.num_nodes\n",
    "args.world_size = args.num_gpus\n",
    "args.train_batch_size = args.local_train_batch_size * args.world_size\n",
    "args.eval_batch_size = args.local_eval_batch_size * args.world_size\n",
    "#############################################\n",
    "\n",
    "cf = import_filename(args.cf)\n",
    "Net, inner_collect_fn = cf.Net, cf.inner_collect_fn\n",
    "\n",
    "dataset_cf = import_filename(args.dataset_cf)\n",
    "BaseDataset = dataset_cf.BaseDataset\n",
    "\n",
    "# args = update_args(parsed_args, args)\n",
    "\n",
    "# init models\n",
    "logger.info('Building models...')\n",
    "model = Net(args)\n",
    "print(f\"Args: {edict(vars(args))}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "019e44c6",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "21a550cd",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8f876d44",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "d1666198",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\u001b[33m[2023-07-03 08:19:47 <287074974.py:3> <module>] Do eval_visu...\u001b[0m\n",
      "\u001b[33m[2023-07-03 08:19:47 <287074974.py:3> <module>] Do eval_visu...\u001b[0m\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "number of samples: 25\n",
      "Specify the load model path, not use deepspeed but the pytorch original load func\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[2023-07-03 08:19:51 <wutils_ldm.py:456> file2data] Loaded data from /home/wangtan/data/disco/ft_checkpoint/moretiktok_nocfg/mp_rank_00_model_states.pt\u001b[0m\n",
      "[2023-07-03 08:19:51 <wutils_ldm.py:456> file2data] Loaded data from /home/wangtan/data/disco/ft_checkpoint/moretiktok_nocfg/mp_rank_00_model_states.pt\u001b[0m\n",
      "\u001b[33m[2023-07-03 08:19:54 <wutils_ldm.py:701> adaptively_load_state_dict] Strictly Loaded state_dict.\u001b[0m\n",
      "\u001b[33m[2023-07-03 08:19:54 <wutils_ldm.py:701> adaptively_load_state_dict] Strictly Loaded state_dict.\u001b[0m\n",
      "\u001b[33m[2023-07-03 08:19:54 <agent.py:808> load_checkpoint_for_deepspeed_diff_gpu] Loaded checkpoint /home/wangtan/data/disco/ft_checkpoint/moretiktok_nocfg/mp_rank_00_model_states.pt of global_step 7200\u001b[0m\n",
      "\u001b[33m[2023-07-03 08:19:54 <agent.py:808> load_checkpoint_for_deepspeed_diff_gpu] Loaded checkpoint /home/wangtan/data/disco/ft_checkpoint/moretiktok_nocfg/mp_rank_00_model_states.pt of global_step 7200\u001b[0m\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[2023-07-03 08:19:54,390] [INFO] [logging.py:69:log_dist] [Rank 0] DeepSpeed info: version=0.6.3, git-hash=unknown, git-branch=unknown\n",
      "[2023-07-03 08:19:54,729] [INFO] [engine.py:278:__init__] DeepSpeed Flops Profiler Enabled: False\n",
      "[2023-07-03 08:19:54,734] [INFO] [config.py:1059:print] DeepSpeedEngine configuration:\n",
      "[2023-07-03 08:19:54,736] [INFO] [config.py:1063:print]   activation_checkpointing_config  {\n",
      "    \"partition_activations\": false, \n",
      "    \"contiguous_memory_optimization\": false, \n",
      "    \"cpu_checkpointing\": false, \n",
      "    \"number_checkpoints\": null, \n",
      "    \"synchronize_checkpoint_boundary\": false, \n",
      "    \"profile\": false\n",
      "}\n",
      "[2023-07-03 08:19:54,737] [INFO] [config.py:1063:print]   aio_config ................... {'block_size': 1048576, 'queue_depth': 8, 'thread_count': 1, 'single_submit': False, 'overlap_events': True}\n",
      "[2023-07-03 08:19:54,738] [INFO] [config.py:1063:print]   amp_enabled .................. False\n",
      "[2023-07-03 08:19:54,739] [INFO] [config.py:1063:print]   amp_params ................... False\n",
      "[2023-07-03 08:19:54,740] [INFO] [config.py:1063:print]   autotuning_config ............ {\n",
      "    \"enabled\": false, \n",
      "    \"start_step\": null, \n",
      "    \"end_step\": null, \n",
      "    \"metric_path\": null, \n",
      "    \"arg_mappings\": null, \n",
      "    \"metric\": \"throughput\", \n",
      "    \"model_info\": null, \n",
      "    \"results_dir\": null, \n",
      "    \"exps_dir\": null, \n",
      "    \"overwrite\": true, \n",
      "    \"fast\": true, \n",
      "    \"start_profile_step\": 3, \n",
      "    \"end_profile_step\": 5, \n",
      "    \"tuner_type\": \"gridsearch\", \n",
      "    \"tuner_early_stopping\": 5, \n",
      "    \"tuner_num_trials\": 50, \n",
      "    \"model_info_path\": null, \n",
      "    \"mp_size\": 1, \n",
      "    \"max_train_batch_size\": null, \n",
      "    \"min_train_batch_size\": 1, \n",
      "    \"max_train_micro_batch_size_per_gpu\": 1.024000e+03, \n",
      "    \"min_train_micro_batch_size_per_gpu\": 1, \n",
      "    \"num_tuning_micro_batch_sizes\": 3\n",
      "}\n",
      "[2023-07-03 08:19:54,740] [INFO] [config.py:1063:print]   bfloat16_enabled ............. False\n",
      "[2023-07-03 08:19:54,741] [INFO] [config.py:1063:print]   checkpoint_tag_validation_enabled  True\n",
      "[2023-07-03 08:19:54,742] [INFO] [config.py:1063:print]   checkpoint_tag_validation_fail  False\n",
      "[2023-07-03 08:19:54,743] [INFO] [config.py:1063:print]   communication_data_type ...... None\n",
      "[2023-07-03 08:19:54,744] [INFO] [config.py:1063:print]   curriculum_enabled ........... False\n",
      "[2023-07-03 08:19:54,745] [INFO] [config.py:1063:print]   curriculum_params ............ False\n",
      "[2023-07-03 08:19:54,745] [INFO] [config.py:1063:print]   dataloader_drop_last ......... False\n",
      "[2023-07-03 08:19:54,746] [INFO] [config.py:1063:print]   disable_allgather ............ False\n",
      "[2023-07-03 08:19:54,747] [INFO] [config.py:1063:print]   dump_state ................... False\n",
      "[2023-07-03 08:19:54,748] [INFO] [config.py:1063:print]   dynamic_loss_scale_args ...... None\n",
      "[2023-07-03 08:19:54,748] [INFO] [config.py:1063:print]   eigenvalue_enabled ........... False\n",
      "[2023-07-03 08:19:54,749] [INFO] [config.py:1063:print]   eigenvalue_gas_boundary_resolution  1\n",
      "[2023-07-03 08:19:54,751] [INFO] [config.py:1063:print]   eigenvalue_layer_name ........ bert.encoder.layer\n",
      "[2023-07-03 08:19:54,752] [INFO] [config.py:1063:print]   eigenvalue_layer_num ......... 0\n",
      "[2023-07-03 08:19:54,753] [INFO] [config.py:1063:print]   eigenvalue_max_iter .......... 100\n",
      "[2023-07-03 08:19:54,753] [INFO] [config.py:1063:print]   eigenvalue_stability ......... 1e-06\n",
      "[2023-07-03 08:19:54,754] [INFO] [config.py:1063:print]   eigenvalue_tol ............... 0.01\n",
      "[2023-07-03 08:19:54,755] [INFO] [config.py:1063:print]   eigenvalue_verbose ........... False\n",
      "[2023-07-03 08:19:54,755] [INFO] [config.py:1063:print]   elasticity_enabled ........... False\n",
      "[2023-07-03 08:19:54,756] [INFO] [config.py:1063:print]   flops_profiler_config ........ {\n",
      "    \"enabled\": false, \n",
      "    \"profile_step\": 1, \n",
      "    \"module_depth\": -1, \n",
      "    \"top_modules\": 3, \n",
      "    \"detailed\": true, \n",
      "    \"output_file\": null\n",
      "}\n",
      "[2023-07-03 08:19:54,757] [INFO] [config.py:1063:print]   fp16_enabled ................. True\n",
      "[2023-07-03 08:19:54,757] [INFO] [config.py:1063:print]   fp16_master_weights_and_gradients  False\n",
      "[2023-07-03 08:19:54,758] [INFO] [config.py:1063:print]   fp16_mixed_quantize .......... False\n",
      "[2023-07-03 08:19:54,758] [INFO] [config.py:1063:print]   global_rank .................. 0\n",
      "[2023-07-03 08:19:54,759] [INFO] [config.py:1063:print]   gradient_accumulation_steps .. 1\n",
      "[2023-07-03 08:19:54,760] [INFO] [config.py:1063:print]   gradient_clipping ............ 0.0\n",
      "[2023-07-03 08:19:54,760] [INFO] [config.py:1063:print]   gradient_predivide_factor .... 1.0\n",
      "[2023-07-03 08:19:54,761] [INFO] [config.py:1063:print]   initial_dynamic_scale ........ 4294967296\n",
      "[2023-07-03 08:19:54,761] [INFO] [config.py:1063:print]   loss_scale ................... 0\n",
      "[2023-07-03 08:19:54,762] [INFO] [config.py:1063:print]   memory_breakdown ............. False\n",
      "[2023-07-03 08:19:54,763] [INFO] [config.py:1063:print]   optimizer_legacy_fusion ...... False\n",
      "[2023-07-03 08:19:54,763] [INFO] [config.py:1063:print]   optimizer_name ............... None\n",
      "[2023-07-03 08:19:54,764] [INFO] [config.py:1063:print]   optimizer_params ............. None\n",
      "[2023-07-03 08:19:54,764] [INFO] [config.py:1063:print]   pipeline ..................... {'stages': 'auto', 'partition': 'best', 'seed_layers': False, 'activation_checkpoint_interval': 0}\n",
      "[2023-07-03 08:19:54,765] [INFO] [config.py:1063:print]   pld_enabled .................. False\n",
      "[2023-07-03 08:19:54,765] [INFO] [config.py:1063:print]   pld_params ................... False\n",
      "[2023-07-03 08:19:54,766] [INFO] [config.py:1063:print]   prescale_gradients ........... False\n",
      "[2023-07-03 08:19:54,766] [INFO] [config.py:1063:print]   quantize_change_rate ......... 0.001\n",
      "[2023-07-03 08:19:54,767] [INFO] [config.py:1063:print]   quantize_groups .............. 1\n",
      "[2023-07-03 08:19:54,767] [INFO] [config.py:1063:print]   quantize_offset .............. 1000\n",
      "[2023-07-03 08:19:54,768] [INFO] [config.py:1063:print]   quantize_period .............. 1000\n",
      "[2023-07-03 08:19:54,768] [INFO] [config.py:1063:print]   quantize_rounding ............ 0\n",
      "[2023-07-03 08:19:54,769] [INFO] [config.py:1063:print]   quantize_start_bits .......... 16\n",
      "[2023-07-03 08:19:54,769] [INFO] [config.py:1063:print]   quantize_target_bits ......... 8\n",
      "[2023-07-03 08:19:54,770] [INFO] [config.py:1063:print]   quantize_training_enabled .... False\n",
      "[2023-07-03 08:19:54,770] [INFO] [config.py:1063:print]   quantize_type ................ 0\n",
      "[2023-07-03 08:19:54,771] [INFO] [config.py:1063:print]   quantize_verbose ............. False\n",
      "[2023-07-03 08:19:54,771] [INFO] [config.py:1063:print]   scheduler_name ............... None\n",
      "[2023-07-03 08:19:54,772] [INFO] [config.py:1063:print]   scheduler_params ............. None\n",
      "[2023-07-03 08:19:54,772] [INFO] [config.py:1063:print]   sparse_attention ............. None\n",
      "[2023-07-03 08:19:54,773] [INFO] [config.py:1063:print]   sparse_gradients_enabled ..... False\n",
      "[2023-07-03 08:19:54,773] [INFO] [config.py:1063:print]   steps_per_print .............. 10\n",
      "[2023-07-03 08:19:54,774] [INFO] [config.py:1063:print]   tensorboard_enabled .......... True\n",
      "[2023-07-03 08:19:54,774] [INFO] [config.py:1063:print]   tensorboard_job_name ......... tensorboard_log\n",
      "[2023-07-03 08:19:54,775] [INFO] [config.py:1063:print]   tensorboard_output_path ...... /home1/wangtan/code/ms_internship2/github_repo/run_test/exp/tiktok_ft\n",
      "[2023-07-03 08:19:54,775] [INFO] [config.py:1063:print]   train_batch_size ............. 32\n",
      "[2023-07-03 08:19:54,776] [INFO] [config.py:1063:print]   train_micro_batch_size_per_gpu  32\n",
      "[2023-07-03 08:19:54,776] [INFO] [config.py:1063:print]   use_quantizer_kernel ......... False\n",
      "[2023-07-03 08:19:54,776] [INFO] [config.py:1063:print]   wall_clock_breakdown ......... False\n",
      "[2023-07-03 08:19:54,777] [INFO] [config.py:1063:print]   world_size ................... 1\n",
      "[2023-07-03 08:19:54,777] [INFO] [config.py:1063:print]   zero_allow_untested_optimizer  False\n",
      "[2023-07-03 08:19:54,778] [INFO] [config.py:1063:print]   zero_config .................. {\n",
      "    \"stage\": 0, \n",
      "    \"contiguous_gradients\": true, \n",
      "    \"reduce_scatter\": true, \n",
      "    \"reduce_bucket_size\": 5.000000e+08, \n",
      "    \"allgather_partitions\": true, \n",
      "    \"allgather_bucket_size\": 5.000000e+08, \n",
      "    \"overlap_comm\": false, \n",
      "    \"load_from_fp32_weights\": true, \n",
      "    \"elastic_checkpoint\": false, \n",
      "    \"offload_param\": null, \n",
      "    \"offload_optimizer\": null, \n",
      "    \"sub_group_size\": 1.000000e+09, \n",
      "    \"prefetch_bucket_size\": 5.000000e+07, \n",
      "    \"param_persistence_threshold\": 1.000000e+05, \n",
      "    \"max_live_parameters\": 1.000000e+09, \n",
      "    \"max_reuse_distance\": 1.000000e+09, \n",
      "    \"gather_16bit_weights_on_model_save\": false, \n",
      "    \"ignore_unused_parameters\": true, \n",
      "    \"round_robin_gradients\": false, \n",
      "    \"legacy_stage1\": false\n",
      "}\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[2023-07-03 08:19:54,779] [INFO] [config.py:1063:print]   zero_enabled ................. False\n",
      "[2023-07-03 08:19:54,779] [INFO] [config.py:1063:print]   zero_optimization_stage ...... 0\n",
      "[2023-07-03 08:19:54,780] [INFO] [config.py:1065:print]   json = {\n",
      "    \"train_micro_batch_size_per_gpu\": 32, \n",
      "    \"gradient_accumulation_steps\": 1, \n",
      "    \"fp16\": {\n",
      "        \"enabled\": true\n",
      "    }, \n",
      "    \"flops_profiler\": {\n",
      "        \"enabled\": false, \n",
      "        \"profile_step\": 1, \n",
      "        \"module_depth\": -1, \n",
      "        \"top_modules\": 3, \n",
      "        \"detailed\": true\n",
      "    }, \n",
      "    \"tensorboard\": {\n",
      "        \"enabled\": true, \n",
      "        \"output_path\": \"/home1/wangtan/code/ms_internship2/github_repo/run_test/exp/tiktok_ft\", \n",
      "        \"job_name\": \"tensorboard_log\"\n",
      "    }\n",
      "}\n",
      "Using /home/wangtan/.cache/torch_extensions/py38_cu113 as PyTorch extensions root...\n",
      "Emitting ninja build file /home/wangtan/.cache/torch_extensions/py38_cu113/utils/build.ninja...\n",
      "Building extension module utils...\n",
      "Allowing ninja to set a default number of workers... (overridable by setting the environment variable MAX_JOBS=N)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[2023-07-03 08:19:56 <agent.py:237> prepare_dist_model] Successfully built models with {'trainable': 1581790996, 'frozen': 387620071, 'trainable_fp32': 0, 'trainalbe_fp16': 1581790996, 'frozen_fp32': 303966208, 'frozen_fp16': 83653863} parameters\u001b[0m\n",
      "[2023-07-03 08:19:56 <agent.py:237> prepare_dist_model] Successfully built models with {'trainable': 1581790996, 'frozen': 387620071, 'trainable_fp32': 0, 'trainalbe_fp16': 1581790996, 'frozen_fp32': 303966208, 'frozen_fp16': 83653863} parameters\u001b[0m\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "ninja: no work to do.\n",
      "Loading extension module utils...\n",
      "Time to load utils op: 0.8241767883300781 seconds\n",
      "Mode [all]: There are 686 modules in unet to be set as requires_grad=True.\n"
     ]
    }
   ],
   "source": [
    "### prepare the eval\n",
    "\n",
    "logger.warning(\"Do eval_visu...\")\n",
    "if getattr(args, 'refer_clip_preprocess', None):\n",
    "    eval_dataset = BaseDataset(args, args.val_yaml, split='val', preprocesser=model.feature_extractor)\n",
    "else:\n",
    "    eval_dataset = BaseDataset(args, args.val_yaml, split='val')\n",
    "eval_dataloader, eval_info = make_data_loader(\n",
    "    args, args.local_eval_batch_size, \n",
    "    eval_dataset)\n",
    "\n",
    "\n",
    "trainer = Agent_LDM(args=args, model=model)\n",
    "trainer.eval_demo_pre()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "4377766c",
   "metadata": {},
   "outputs": [],
   "source": [
    "def load_image(image):\n",
    "    if not image.mode == \"RGB\":\n",
    "        image = image.convert(\"RGB\")\n",
    "\n",
    "    return image\n",
    "\n",
    "\n",
    "\n",
    "@torch.no_grad()\n",
    "def inference(reference_fg, fg_mask, ref_bg_image, bg_mask, skeleton_img, *args, **kwargs):\n",
    "    reference_fg = load_image(reference_fg)\n",
    "    fg_mask = load_image(fg_mask)\n",
    "    ref_bg_image = load_image(ref_bg_image)\n",
    "    bg_mask = load_image(bg_mask)\n",
    "    skeleton_img = load_image(skeleton_img)\n",
    "    \n",
    "    input_data = [reference_fg, fg_mask, ref_bg_image, bg_mask, skeleton_img]\n",
    "    output_image = trainer.eval_demo_run(input_data, eval_dataset=eval_dataset)\n",
    "    return output_image\n",
    "\n",
    "@torch.no_grad()\n",
    "def inference_masked(reference_fg, ref_bg_image, skeleton_img, *args, **kwargs):\n",
    "    reference_fg = load_image(reference_fg)\n",
    "    ref_bg_image = load_image(ref_bg_image)\n",
    "    skeleton_img = load_image(skeleton_img)\n",
    "    \n",
    "    input_data = [reference_fg, ref_bg_image, skeleton_img]\n",
    "    output_image = trainer.eval_demo_run_masked(input_data, eval_dataset=eval_dataset)\n",
    "    return output_image"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0580155a",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "035a1454",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "d9c3b8ff",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/wangtan/anaconda3/envs/disco/lib/python3.8/site-packages/gradio/layouts.py:80: UserWarning: The `style` method is deprecated. Please set these arguments in the constructor instead.\n",
      "  warnings.warn(\n",
      "/home/wangtan/anaconda3/envs/disco/lib/python3.8/site-packages/gradio/blocks.py:261: UserWarning: api_name load_example already exists, using load_example_1\n",
      "  warnings.warn(f\"api_name {api_name} already exists, using {api_name_}\")\n",
      "/home/wangtan/anaconda3/envs/disco/lib/python3.8/site-packages/gradio/blocks.py:261: UserWarning: api_name load_example already exists, using load_example_2\n",
      "  warnings.warn(f\"api_name {api_name} already exists, using {api_name_}\")\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Running on local URL:  http://127.0.0.1:7861\n",
      "Running on public URL: https://e4ec2a4f33f5ee858d.gradio.live\n",
      "\n",
      "This share link expires in 72 hours. For free permanent hosting and GPU upgrades, run `gradio deploy` from Terminal to deploy to Spaces (https://huggingface.co/spaces)\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<div><iframe src=\"https://e4ec2a4f33f5ee858d.gradio.live\" width=\"100%\" height=\"500\" allow=\"autoplay; camera; microphone; clipboard-read; clipboard-write;\" frameborder=\"0\" allowfullscreen></iframe></div>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": []
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Mode [all]: There are 686 modules in unet to be set as requires_grad=True.\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 50/50 [00:04<00:00, 12.48it/s]\n"
     ]
    }
   ],
   "source": [
    "#### Launch the gradio demo\n",
    "\n",
    "import gradio as gr\n",
    "'''\n",
    "launch app\n",
    "'''\n",
    "title = \"DisCo Demo (Video Demo Comming Soon!)\"\n",
    "description = \"\"\"<p style='text-align: center'> <a href='https://disco-dance.github.io/' target='_blank'>Project Page</a> | <a href='https://arxiv.org/pdf/2212.11270.pdf' target='_blank'>Paper</a> | <a href='https://github.com/microsoft/X-Decoder' target='_blank'>Github Repo</a> | <a href='https://youtu.be/wYp6vmyolqE' target='_blank'>Video</a> </p>\n",
    "<p>Skip the queue by duplicating this space and upgrading to GPU in settings</p>\n",
    "<a href=\"https://huggingface.co/spaces/xdecoder/Demo?duplicate=true\"><img src=\"https://bit.ly/3gLdBN6\" alt=\"Duplicate Space\"></a>\n",
    "\"\"\"\n",
    "\n",
    "\n",
    "\n",
    "with gr.Blocks() as demo:\n",
    "    gr.Markdown(\n",
    "    \"\"\"\n",
    "    # DisCo Demo (Video Demo Comming Soon!)\n",
    "    Start edit the human with provided human foreground, background, pose. \n",
    "    \n",
    "    Note that for self-uploaded images, TikTok-Style human images are preferred.\n",
    "    \n",
    "    [Project Page](https://disco-dance.github.io/) | [Github](https://github.com/Wangt-CN/DisCo)\n",
    "    \"\"\")\n",
    "    #     gr.Markdown(\n",
    "#     \"\"\"\n",
    "#     ## DisCo Demo (Video Demo Comming Soon!)\n",
    "#     <p style='text-align: center'> <a href='https://disco-dance.github.io/' target='_blank'>Project Page</a> | <a href='https://arxiv.org/pdf/2212.11270.pdf' target='_blank'>Paper</a> | <a href='https://github.com/microsoft/X-Decoder' target='_blank'>Github Repo</a> | <a href='https://youtu.be/wYp6vmyolqE' target='_blank'>Video</a> </p>\n",
    "# <p>Skip the queue by duplicating this space and upgrading to GPU in settings</p>\n",
    "# <a href=\"https://huggingface.co/spaces/xdecoder/Demo?duplicate=true\"><img src=\"https://bit.ly/3gLdBN6\" alt=\"Duplicate Space\"></a>\n",
    "#     \"\"\")\n",
    "    \n",
    "    with gr.Row().style(equal_height=False):\n",
    "        with gr.Column(min_width=400, scale=2):\n",
    "            input_fg = gr.Image(type='pil',label=\"Foreground Image\")\n",
    "            gr.Examples(examples=[\"./demo_data/fg/masked_images/00035.png\", \"./demo_data/fg/masked_images/00335.png\", \"./demo_data/fg/masked_images/00147.png\", \"./demo_data/fg/masked_images/00072.png\", \"./demo_data/fg/masked_images/00115.png\"], inputs=input_fg)\n",
    "\n",
    "            input_bg = gr.Image(type='pil',label=\"Background Image\")\n",
    "            gr.Examples(examples=[\"./demo_data/bg/masked_images/00035.png\", \"./demo_data/bg/masked_images/00335.png\", \"./demo_data/bg/masked_images/00147.png\", \"./demo_data/bg/masked_images/00072.png\", \"./demo_data/bg/masked_images/00115.png\"], inputs=input_bg)\n",
    "\n",
    "            input_pose = gr.Image(type='pil',label=\"Target Pose\",scale=1)\n",
    "            gr.Examples(examples=[\"./demo_data/pose_img/0049.png\",\"./demo_data/pose_img/0198.png\",\"./demo_data/pose_img/0213.png\",\"./demo_data/pose_img/0264.png\",\"./demo_data/pose_img/0144.png\",\"./demo_data/pose_img/0054.png\"], inputs=input_pose)\n",
    "\n",
    "            btn = gr.Button(\"Generate\")\n",
    "            \n",
    "    \n",
    "        with gr.Column(min_width=150):\n",
    "            output_img = gr.Image(type='pil',label=\"Edited Human Image\")\n",
    "        \n",
    "    btn.click(inference_masked, inputs=[input_fg, input_bg, input_pose], outputs=[output_img])\n",
    "    \n",
    "demo.queue(concurrency_count=2)\n",
    "demo.launch(share=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "141af2b1",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python [conda env:disco]",
   "language": "python",
   "name": "conda-env-disco-py"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}