muse_train_maskgit1.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "MASK_GIT = 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import configparser\n",
    "\n",
    "# Loading configurations\n",
    "configParser = configparser.RawConfigParser()   \n",
    "configFilePath = r'configuration.txt'\n",
    "configParser.read(configFilePath)\n",
    "\n",
    "insert_amd_env_vars =  int(configParser.get('COMMON', 'insert_amd_env_vars'))\n",
    "HSA_OVERRIDE_GFX_VERSION =  configParser.get('COMMON', 'HSA_OVERRIDE_GFX_VERSION')\n",
    "ROCM_PATH =  configParser.get('COMMON', 'ROCM_PATH')\n",
    "\n",
    "if(insert_amd_env_vars != 0):\n",
    "    os.environ[\"HSA_OVERRIDE_GFX_VERSION\"] = HSA_OVERRIDE_GFX_VERSION\n",
    "    os.environ[\"ROCM_PATH\"] = ROCM_PATH\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "torch.cuda.is_available()\n",
    "import torch. multiprocessing as mp \n",
    "mp.set_start_method('spawn')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "import torch.multiprocessing as mp \n",
    "mp.set_start_method('spawn')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from textwrap import wrap\n",
    "import os\n",
    "import gc\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "\n",
    "#TODO document jupyter\n",
    "import pickle\n",
    "import configparser\n",
    "import sqlite3 as sl\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "from matplotlib import pyplot as plt\n",
    "from pathlib import Path\n",
    "\n",
    "configParser = configparser.RawConfigParser()   \n",
    "configFilePath = r'configuration.txt'\n",
    "configParser.read(configFilePath)\n",
    "datasetPathDatabase =  configParser.get('COMMON', 'datasetPathDatabase') + '/dataset.db'\n",
    "datasetPathDatabaseAdditional =  configParser.get('COMMON', 'datasetPathDatabase') + '/datasetAdditional.db'\n",
    "datasetPathDatabaseVgg =  configParser.get('COMMON', 'datasetPathDatabase') + '/datasetFacesBlurred.db'\n",
    "muse_image_size =  int(configParser.get('COMMON', 'muse_image_size'))\n",
    "muse_begin_with_image_size = int(configParser.get('COMMON', 'muse_begin_with_image_size'))\n",
    "unet1_dim =  int(configParser.get('COMMON', 'unet1_dim'))\n",
    "unet2_dim =  int(configParser.get('COMMON', 'unet2_dim'))\n",
    "mask_git1_image_size =  int(configParser.get('COMMON', 'mask_git1_image_size'))\n",
    "audio_embs =  str(configParser.get('COMMON', 'audio_embs'))\n",
    "audio_length_used =  configParser.get('train_muse', 'audio_length_used') \n",
    "model_filename =  configParser.get('train_muse', 'model_filename') + '_'  + audio_length_used +  's.pt'\n",
    "sub_epochs=  int(configParser.get('train_muse', 'sub_epochs') )\n",
    "inner_epochs=  int(configParser.get('train_muse', 'inner_epochs') )\n",
    "batch_size=  int(configParser.get('train_muse', 'batch_size') )\n",
    "timesteps= int(configParser.get('COMMON', 'timesteps') )\n",
    "sample_every=  int(configParser.get('train_muse', 'sample_every') ) - 1\n",
    "sample_probability = int(configParser.get('train_muse', 'sample_probability') )\n",
    "save_model_every=  int(configParser.get('train_muse', 'save_model_every') ) - 1\n",
    "sample_every_offset=  int(configParser.get('train_muse', 'sample_every_offset') ) - 1\n",
    "save_every_offset=  int(configParser.get('train_muse', 'save_every_offset') ) - 1\n",
    "imagen_samples_folder = configParser.get('train_muse', 'imagen_samples_folder') \n",
    "db_chunk = int(configParser.get('train_muse', 'db_chunk'))\n",
    "dask_chunk = int(configParser.get('train_muse', 'dask_chunk'))\n",
    "stop_at_no_of_samples = int(configParser.get('train_muse', 'stop_at_no_of_samples'))\n",
    "epochs = int(configParser.get('train_muse', 'epochs'))\n",
    "ignore_speaker_embedding = bool(int(configParser.get('train_muse', 'ignore_speaker_embedding') ))\n",
    "ignore_speech_brain = bool(int(configParser.get('train_muse', 'ignore_speech_brain') ))\n",
    "ignore_pyannote_titanet_speakernet = bool(int(configParser.get('train_muse', 'ignore_pyannote_titanet_speakernet') ))\n",
    "ignore_audio_features = bool(int(configParser.get('train_muse', 'ignore_audio_features') ))\n",
    "ignore_pyAudioAnalysis = bool(int(configParser.get('train_muse', 'ignore_pyAudioAnalysis') ))\n",
    "ignore_librosa = bool(int(configParser.get('train_muse', 'ignore_librosa') ))\n",
    "ignore_image_guide = bool(int(configParser.get('train_muse', 'ignore_image_guide') ))\n",
    "ignore_additional_attributes = bool(int(configParser.get('train_muse', 'ignore_additional_attributes') ))\n",
    "ignore_age = bool(int(configParser.get('train_muse', 'ignore_age') ))\n",
    "ignore_gender = bool(int(configParser.get('train_muse', 'ignore_gender') ))\n",
    "ignore_ethnicity = bool(int(configParser.get('train_muse', 'ignore_ethnicity') ))\n",
    "ignore_language_spoken = bool(int(configParser.get('train_muse', 'ignore_language_spoken') ))\n",
    "\n",
    "vae_file = configParser.get('train_muse', 'vae_file')\n",
    "\n",
    "\n",
    "pd.set_option('display.max_columns', None)\n",
    "pd.options.mode.chained_assignment = None\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def speaker_emb_preprocess(speaker_emb2):\n",
    "\n",
    "    if(ignore_speaker_embedding or ignore_speech_brain):\n",
    "        #print('ignore speech brain')\n",
    "        return np.zeros((1,768))\n",
    "    else:\n",
    "        speaker_emb2 = pickle.loads(speaker_emb2)\n",
    "        #print(speaker_emb2.print)\n",
    "        speaker_emb2 = speaker_emb2.squeeze()\n",
    "        speaker_emb2 = np.pad(speaker_emb2, (288), 'constant', constant_values=(0))\n",
    "        speaker_emb2 = np.tile(speaker_emb2, (1, 1))\n",
    "        speaker_emb2 = speaker_emb2 / 200.0\n",
    "        #print(speaker_emb2.shape)\n",
    "        #speaker_emb2 = np.array(speaker_emb2).tolist()\n",
    "        #print(type(speaker_emb2))\n",
    "        return speaker_emb2\n",
    "\n",
    "\n",
    "def audio_emb_preprocess2(speaker_emb2):\n",
    "    \n",
    "    if(ignore_pyannote_titanet_speakernet or ignore_speaker_embedding):\n",
    "        #print('ignore speech brain')\n",
    "        return np.zeros((28,768))\n",
    "    else:\n",
    "        #print(speaker_emb2)\n",
    "        speaker_emb2 = pickle.loads(speaker_emb2)\n",
    "        a = speaker_emb2 #np.zeros(shape=(24, 512))\n",
    "        a = a[-28:,:] # Actually only last 39 are the relevant ones\n",
    "        #print('au')\n",
    "        #print(a.shape)\n",
    "        b = np.zeros(shape=(a.shape[0], 768-a.shape[1]))\n",
    "        #c = np.zeros(shape=(28-a.shape[0], 768))\n",
    "        arr = np.concatenate((a, b), axis=1)\n",
    "        #arr = np.concatenate((arr, c), axis=0)\n",
    "        arr = arr / (1 if(audio_embs == 'wav2vec') else 10 if(audio_embs == 'openl3')  else 1 if(audio_embs == 'pyannoteTitaNet') else 1)\n",
    "        #print(str(arr.max()))\n",
    "        #print(arr.shape)\n",
    "        speaker_emb2 = np.array(arr)\n",
    "        return speaker_emb2\n",
    "\n",
    "def audio_features_preprocess(video_id):\n",
    "    #  79 belong to pyaudioanalysis\n",
    "    # 111 belogn to liborsa\n",
    "    if(ignore_audio_features):\n",
    "        #print('ignore speech brain')\n",
    "        return np.zeros((190,768))\n",
    "    else:\n",
    "        #print(video_id)\n",
    "        conAdditional = sl.connect(datasetPathDatabaseAdditional)\n",
    "        cur = conAdditional.cursor()\n",
    "        sql = '''SELECT AUDIO_FEATURES FROM AUDIO WHERE VIDEO_ID = ?'''\n",
    "        cur.execute(sql, [video_id])\n",
    "        audio_features = cur.fetchall()\n",
    "        #print(audio_features[0])\n",
    "        audio_features = pickle.loads(audio_features[0][0])\n",
    "        #print(audio_features.shape) # 190 x 128\n",
    "        audio_features = audio_features[0:190]\n",
    "        \n",
    "        #import sys\n",
    "        #np.set_printoptions(threshold=sys.maxsize)\n",
    "        #print(np.argwhere(audio_features == 0))\n",
    "\n",
    "        if(ignore_pyAudioAnalysis and not ignore_librosa):\n",
    "            zpa = np.zeros((190-110,128))\n",
    "            audio_features = audio_features[80:190]\n",
    "            audio_features = np.vstack((audio_features,zpa))\n",
    "            #print('ignore pyAudioAnalysis')\n",
    "            #print(audio_features.shape)\n",
    "        elif(ignore_librosa and not ignore_pyAudioAnalysis):\n",
    "            zpa = np.zeros((190-80,128))\n",
    "            audio_features = audio_features[0:80]\n",
    "            audio_features = np.vstack((audio_features,zpa))\n",
    "            #print('ignore librosa')\n",
    "            #print(audio_features.shape)\n",
    "        else:\n",
    "            audio_features = np.zeros((190,128))\n",
    "            #print('ignore pyaudioanalysis and librosa')\n",
    "\n",
    "\n",
    "        z1 = np.zeros((190,768-128))\n",
    "        audio_features = np.hstack((audio_features,z1))\n",
    "        audio_features = audio_features / 100.0\n",
    "        #print(audio_features.shape)\n",
    "    return audio_features\n",
    "\n",
    "def audio_transformer_features_preprocess(video_id):\n",
    "    #print(video_id)\n",
    "    conAdditional = sl.connect(datasetPathDatabaseAdditional)\n",
    "    cur = conAdditional.cursor()\n",
    "    sql = '''SELECT AUDIO_FEATURES FROM AUDIO_TRANSFORMER WHERE VIDEO_ID = ?'''\n",
    "    cur.execute(sql, [video_id])\n",
    "    audio_features = cur.fetchall()\n",
    "    #print(audio_features[0])\n",
    "    audio_features = pickle.loads(audio_features[0][0])\n",
    "    #print(audio_features.shape) # 514 x 768\n",
    "    audio_features = audio_features.squeeze()\n",
    "    audio_features = audio_features[0::3]#172*768\n",
    "    #z1 = np.zeros((161,768-128))\n",
    "    #audio_features = np.hstack((audio_features,z1))\n",
    "    #audio_features = audio_features / 100.0\n",
    "    return audio_features\n",
    "\n",
    "import random\n",
    "def image_guide_preprocess(face_id):\n",
    "    #print(face_id)\n",
    "\n",
    "    if(random.random() > 2):\n",
    "        image_guide = np.zeros((49,768))\n",
    "        #print(image_guide.shape)\n",
    "    else:\n",
    "        conAdditional = sl.connect(datasetPathDatabaseVgg)\n",
    "        cur = conAdditional.cursor()\n",
    "        sql = '''SELECT BLURRED_FACE_EMB FROM FACES_BLURRED WHERE FACE_ID = ?'''\n",
    "        cur.execute(sql, [face_id])\n",
    "        image_guide = cur.fetchall()\n",
    "        #print(image_guide)\n",
    "        image_guide = pickle.loads(image_guide[0][0])\n",
    "        image_guide = image_guide.squeeze() #197 x 768\n",
    "        image_guide = image_guide[1::4] #49 x 768\n",
    "    return image_guide\n",
    "\n",
    "\n",
    "#boxBlurMin =  int(configParser.get('extractVggBlurred', 'boxBlurMin'))\n",
    "#boxBlurMax =  int(configParser.get('extractVggBlurred', 'boxBlurMax'))\n",
    "\n",
    "#gaussianBlurMin =  int(configParser.get('extractVggBlurred', 'gaussianBlurMin'))\n",
    "#gaussianBlurMax =  int(configParser.get('extractVggBlurred', 'gaussianBlurMax'))\n",
    "\n",
    "\n",
    "from PIL import Image,ImageFilter\n",
    "import random\n",
    "#import cv2\n",
    "import numpy as np\n",
    "#from matplotlib import pyplot as plt\n",
    "  \n",
    "\n",
    "\n",
    "def image_guide_preprocess_low_res_dummy(path):\n",
    "    #print(face_id)\n",
    "\n",
    "    image_guide = np.zeros((1, 768))\n",
    "    return image_guide\n",
    "\n",
    "def image_guide_preprocess_low_res(path):\n",
    "    #print(face_id)\n",
    "\n",
    "\n",
    "    if(ignore_image_guide):\n",
    "        return np.nan\n",
    "    else:\n",
    "\n",
    "        image = Image.open(path).convert('RGB')\n",
    "\n",
    "        #print(image.size)\n",
    "\n",
    "        w_s = muse_image_size / (1+2 * 0.4)\n",
    "        h_s = muse_image_size / (1+2 * 0.4)\n",
    "\n",
    "        image = image.crop((0.2*w_s, 0.0*h_s, 1.6*w_s, 1.4*h_s))\n",
    "\n",
    "        image = image.resize((muse_image_size,muse_image_size))\n",
    "\n",
    "        image = image.resize((muse_begin_with_image_size,muse_begin_with_image_size))\n",
    "        im = image\n",
    "        \n",
    "\n",
    "        #print('saving')\n",
    "        #image.save('opop.png')\n",
    "\n",
    "        #print(np.array(image,np.float32).shape)\n",
    "\n",
    "        pix = np.array(image, np.float32)\n",
    "        pix = np.moveaxis(pix, -1, 0)\n",
    "\n",
    "        pix = pix / 255\n",
    "        image.close()\n",
    "        im.close()\n",
    "        return pix\n",
    "    return image_guide\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "\n",
    "\n",
    "def process_age(age):\n",
    "    if(random.random() < 0.2 or ignore_age or ignore_additional_attributes):\n",
    "        x = np.zeros(768)\n",
    "        x[767] = 1\n",
    "        return x\n",
    "    else:\n",
    "        x = np.ones(768) * (age / 100.0)\n",
    "        x[767] = 0\n",
    "        return x\n",
    "\n",
    "def process_gender(gender):\n",
    "    if(random.random() < 0.2 or ignore_gender or ignore_additional_attributes):\n",
    "        return np.zeros(768)\n",
    "    elif(gender == 'man'):\n",
    "        return np.ones(768)\n",
    "    elif(gender == \"woman\"):\n",
    "        return np.ones(768) * -1\n",
    "    else:\n",
    "        return np.zeros(768)\n",
    "    \n",
    "# TODO\n",
    "def process_ethnicity(eth):\n",
    "    x = np.zeros(768)\n",
    "    if(random.random() < 0.2 or ignore_ethnicity or ignore_additional_attributes):\n",
    "        x = x\n",
    "    elif(eth == \"indian\"):\n",
    "        x[0] = 1\n",
    "    elif(eth == \"asian\"):\n",
    "        x[16]=1\n",
    "    elif(eth == \"latino hispanic\"):\n",
    "        x[2]=1\n",
    "    elif(eth == \"black\"):\n",
    "        x[3]=1\n",
    "    elif(eth == \"middle eastern\"):\n",
    "        x[4]=1\n",
    "    elif(eth == \"white\"):\n",
    "        x[5]=1 \n",
    "    else:\n",
    "        x = x\n",
    "    return x     \n",
    "\n",
    "def process_language(lan):\n",
    "    x = np.zeros(768)\n",
    "    if(random.random() < 0.2 or ignore_language_spoken or ignore_additional_attributes):\n",
    "        x = x\n",
    "    elif(lan == \"Arabic\"):\n",
    "        x[0] = 1\n",
    "    elif(lan == \"Portuguese\"):\n",
    "        x[16]=1\n",
    "    elif(lan == \"Romansh_Sursilvan\"):\n",
    "        x[2]=1\n",
    "    elif(lan == \"Japanese\"):\n",
    "        x[3]=1\n",
    "    elif(lan == \"Ukranian\"):\n",
    "        x[4]=1\n",
    "    elif(lan == \"German\"):\n",
    "        x[5]=1   \n",
    "    elif(lan == \"Chinese_China\"):\n",
    "        x[6]=1   \n",
    "    elif(lan == \"Welsh\"):\n",
    "        x[7]=1  \n",
    "    elif(lan == \"English\"):\n",
    "        x[8]=1\n",
    "    elif(lan == \"Kabyle\"):\n",
    "        x[9]=1 \n",
    "    elif(lan == \"Kyrgyz\"):\n",
    "        x[10]=1\n",
    "    elif(lan == \"Georgian\"):\n",
    "        x[11]=1\n",
    "    elif(lan == \"Persian\"):\n",
    "        x[12]=1 \n",
    "    elif(lan == \"French\"):\n",
    "        x[13]=1\n",
    "    elif(lan == \"Interlingua\"):\n",
    "        x[14]=1\n",
    "    elif(lan == \"Swedish\"):\n",
    "        x[15]=1\n",
    "    elif(lan == \"Spanish\"):\n",
    "        x[16]=1 \n",
    "    elif(lan == \"Dhivehi\"):\n",
    "        x[17]=1\n",
    "    elif(lan == \"Kinyarwanda\"):\n",
    "        x[18]=1 \n",
    "    elif(lan == \"Tatar\"):\n",
    "        x[19]=1\n",
    "    elif(lan == \"Hakha_Chin\"):\n",
    "        x[20]=1 \n",
    "    elif(lan == \"Tamil\"):\n",
    "        x[21]=1 \n",
    "    elif(lan == \"Greek\"):\n",
    "        x[22]=1\n",
    "    elif(lan == \"Latvian\"):\n",
    "        x[23]=1 \n",
    "    elif(lan == \"Russian\"):\n",
    "        x[24]=1\n",
    "    elif(lan == \"Breton\"):\n",
    "        x[25]=1\n",
    "    elif(lan == \"Catalan\"):\n",
    "        x[26]=1    \n",
    "    elif(lan == \"Maltese\"):\n",
    "        x[27]=1 \n",
    "    elif(lan == \"Slovenian\"):\n",
    "        x[28]=1    \n",
    "    elif(lan == \"Indonesian\"):\n",
    "        x[29]=1    \n",
    "    elif(lan == \"Dutch\"):\n",
    "        x[30]=1\n",
    "    elif(lan == \"Chinese_Taiwan\"):\n",
    "        x[31]=1 \n",
    "    elif(lan == \"Sakha\"):\n",
    "        x[32]=1 \n",
    "    elif(lan == \"Polish\"):\n",
    "        x[33]=1 \n",
    "    elif(lan == \"Czech\"):\n",
    "        x[34]=1 \n",
    "    elif(lan == \"Romanian\"):\n",
    "        x[35]=1 \n",
    "    elif(lan == \"Mangolian\"):\n",
    "        x[36]=1 \n",
    "    elif(lan == \"Italian\"):\n",
    "        x[37]=1 \n",
    "    elif(lan == \"Chinese_Hongkong\"):\n",
    "        x[38]=1 \n",
    "    elif(lan == \"Estonian\"):\n",
    "        x[39]=1 \n",
    "    elif(lan == \"Basque\"):\n",
    "        x[40]=1 \n",
    "    elif(lan == \"Esperanto\"):\n",
    "        x[41]=1 \n",
    "    elif(lan == \"Frisian\"):\n",
    "        x[42]=1 \n",
    "    elif(lan == \"Turkish\"):\n",
    "        x[43]=1 \n",
    "    elif(lan == \"Chuvash\"):\n",
    "        x[44]=1 \n",
    "    else:\n",
    "        x = x\n",
    "    return x       "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "def process_input(row):\n",
    "    #print(row)\n",
    "    age = row[\"caption_a\"]\n",
    "    ethnicity=row[\"caption_e\"]\n",
    "    gender=row[\"caption_g\"]\n",
    "    language=row[\"caption_l\"]\n",
    "    speaker_emb=row[\"SPEAKER_EMB\"]\n",
    "    audio_emb=row[\"AUDIO_EMB\"]\n",
    "    audio_features=row[\"AUDIO_FEATURES\"]\n",
    "    image_guide=row[\"image_guide\"]\n",
    "    #print(image_guide)\n",
    "    #print(speaker_emb)\n",
    "    speaker_emb = np.asarray(speaker_emb, dtype=np.float32)\n",
    "    speaker_emb = speaker_emb.squeeze()\n",
    "    audio_emb = np.asarray(audio_emb, dtype=np.float32)\n",
    "    audio_emb = audio_emb.squeeze()\n",
    "    #print(speaker_emb.shape)\n",
    "    h = np.vstack((age, ethnicity))\n",
    "    h = np.vstack((h, gender))\n",
    "    h = np.vstack((h, language))\n",
    "    h = np.vstack((h, speaker_emb))\n",
    "    h = np.vstack((h, audio_emb))\n",
    "    h = np.vstack((h,audio_features))\n",
    "    h = np.vstack((h,image_guide))\n",
    "    #print('aaaaaaaaaaa')\n",
    "    #print(h.shape)\n",
    "    j = np.zeros(768)\n",
    "    j = np.tile(j,(256-h.shape[0],1))\n",
    "    h = np.vstack((h, j))\n",
    "    #print(h.shape)\n",
    "    \n",
    "    return h\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "def process_image_path(path):\n",
    "    \n",
    "    image = Image.open(path).convert('RGB')\n",
    "\n",
    "    w_s = muse_image_size / (1+2 * 0.4)\n",
    "    h_s = muse_image_size / (1+2 * 0.4)\n",
    "\n",
    "    #print(image.size)\n",
    "    image = image.crop((0.2*w_s,0.0*h_s,1.6*w_s,1.4*h_s))\n",
    "    \n",
    "\n",
    "    #print('saving')\n",
    "    #image.save(str(random.random()) + '.png')\n",
    "\n",
    "    if(MASK_GIT == 1):\n",
    "        image = image.resize((mask_git1_image_size,mask_git1_image_size))\n",
    "    else:\n",
    "        image = image.resize((muse_image_size,muse_image_size))\n",
    "    \n",
    "\n",
    "    #print(np.array(image,np.float32).shape)\n",
    "    pix = np.array(image,np.float32)\n",
    "    pix = np.moveaxis(pix, -1, 0)\n",
    "    \n",
    "    pix = pix / 255\n",
    "    image.close()\n",
    "    return pix"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import multiprocessing\n",
    "import muse_train_batch\n",
    "import dask.dataframe as dd\n",
    "import math\n",
    "\n",
    "procs = []\n",
    "\n",
    "for epoch in range(epochs):\n",
    "    print('-------------------- TRAINING IN EPOCH = ' + str(epoch) + ' ---------------------')\n",
    "    stop_now = False\n",
    "\n",
    "\n",
    "    while(not stop_now):\n",
    "\n",
    "        con = sl.connect(datasetPathDatabase)\n",
    "\n",
    "        print(\"Getting Data...\")\n",
    "        #import time\n",
    "        #start_time = time.time()\n",
    "        ## TODO: NEED TO CHECK F.ID CORRESPONDS TO CORRECT FACE_PATH\n",
    "\n",
    "        data = con.execute(\"SELECT V.ID,F.ID,V.VIDEO_PATH, V.AGE CAPTION_A, \" + \n",
    "                            \"V.ETHNICITY CAPTION_E, \" +\n",
    "                            \"lower(V.GENDER) CAPTION_G, \" +\n",
    "                                \"A.SPEAKER_EMB, \"+ (\"A.WAV_TO_VEC, \" if(audio_embs == 'wav2vec') else \"A.AUDIO_EMB2, \" if(audio_embs == 'openl3')  else \"A.PYANNOTE_TITANET, \" if(audio_embs == 'pyannoteTitaNet') else ', ') +\n",
    "                            \"A.AUDIO_FEATURES, \" +\n",
    "                            \"A.LANG CAPTION_L, \"+\n",
    "                            \"F.FACE_PATH \"+\n",
    "                            \"FROM VIDEO V \"+\n",
    "                            \"INNER JOIN AUDIO A ON V.ID = A.VIDEO_ID \" +\n",
    "                            \"INNER JOIN FACE F ON F.ID = (select ID from FACE f2 where f2.video_id = v.ID ORDER By ID limit 1 ) \" + \n",
    "                            \"WHERE AUDIO_LENGTH = \" + audio_length_used + ' ' +\n",
    "                            \"AND V.ID in (select V2.ID from VIDEO v2 WHERE V2.TRAINED IS NULL AND AUDIO_PRE IN (3,4) AND FACES_PRE = 2 ORDER BY ABS(RANDOM()) LIMIT \"+ str(db_chunk) + \")\" )\n",
    "        dataGotten = data.fetchall()\n",
    "\n",
    "        \n",
    "        print(\"Data Gotten\")\n",
    "        for ps in procs:\n",
    "            ps.join()\n",
    "        procs.clear()\n",
    "\n",
    "        if(len(dataGotten) == 0):\n",
    "            stop_now = True\n",
    "\n",
    "        \n",
    "\n",
    "        print(\"Preprocessing Data...\")\n",
    "        df = pd.DataFrame(dataGotten,columns = ['ID','FACE_ID','VIDEO_PATH','caption_a','caption_e','caption_g','SPEAKER_EMB','AUDIO_EMB','AUDIO_FEATURES','caption_l','image_path'])\n",
    "        df = df.sample(frac=1).reset_index(drop=True)\n",
    "        df = dd.from_pandas(df, npartitions=int(db_chunk/dask_chunk))\n",
    "        df[\"image_guide\"] = np.nan\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "        \n",
    "\n",
    "\n",
    "\n",
    "        #print(df.head(10))\n",
    "        data_frame = df[[\"ID\",\"FACE_ID\",\"image_path\",\"caption_a\",\"caption_e\",\"caption_g\",\"caption_l\"]]\n",
    "        data_frame['SPEAKER_EMB'] = df['SPEAKER_EMB']\n",
    "        data_frame['AUDIO_EMB'] = df['AUDIO_EMB']\n",
    "        data_frame['image_guide'] = df['image_guide']\n",
    "\n",
    "        \n",
    "        data_frame['SPEAKER_EMB'] = data_frame['SPEAKER_EMB'].apply(lambda x: speaker_emb_preprocess(x),meta=('1', 'object'))\n",
    "        data_frame['AUDIO_EMB'] = data_frame['AUDIO_EMB'].apply(lambda x: audio_emb_preprocess2(x),meta=('1', 'object'))\n",
    "        data_frame['AUDIO_FEATURES'] = data_frame['ID'].apply(lambda x: audio_features_preprocess(x),meta=('1', 'object'))\n",
    "        data_frame = data_frame.drop(['ID'], axis=1)\n",
    "        data_frame['image_guide'] = data_frame['image_path'].apply(lambda x: image_guide_preprocess_low_res_dummy(x),meta=('1', 'str'))\n",
    "        data_frame = data_frame.drop(['FACE_ID'], axis=1)\n",
    "        data_frame['caption_a'] = data_frame['caption_a'].apply(lambda x: process_age(x),meta=('1', 'object'))\n",
    "        data_frame['caption_g'] = data_frame['caption_g'].apply(lambda x: process_gender(x),meta=('1', 'object'))\n",
    "        data_frame['caption_l'] = data_frame['caption_l'].apply(lambda x: process_language(x),meta=('1', 'object'))\n",
    "        data_frame['caption_e'] = data_frame['caption_e'].apply(lambda x: process_ethnicity(x),meta=('1', 'object'))\n",
    "        data_frame['low_res_image'] = data_frame['image_path'].apply(lambda x: image_guide_preprocess_low_res(x),meta=('1', 'object'))\n",
    "\n",
    "        data_frame['INPUT'] = data_frame['SPEAKER_EMB']\n",
    "        #print(data_frame)\n",
    "\n",
    "        data_frame['INPUT'] = data_frame.apply(process_input,args=(),axis=1,meta=('1', 'object'))\n",
    "\n",
    "        #for index, row in data_frame.iterrows():\n",
    "        #    x = process_input(data_frame.loc[index,\"caption_a\"],data_frame.loc[index,\"caption_e\"]\n",
    "        #                        ,data_frame.loc[index,\"caption_g\"],\n",
    "        #                    data_frame.loc[index,\"caption_l\"],data_frame.loc[index,\"SPEAKER_EMB\"],\n",
    "        #                    data_frame.loc[index,\"AUDIO_EMB\"],data_frame.loc[index,\"AUDIO_FEATURES\"],\n",
    "        #                    data_frame.loc[index,\"image_guide\"])\n",
    "        #    x = [x]\n",
    "        #    #AADFS = AADFS\n",
    "        #    data_frame.loc[index,\"INPUT\"] = x\n",
    "\n",
    "        data_frame = data_frame.drop(['caption_e', 'caption_g','caption_l','SPEAKER_EMB','AUDIO_EMB'], axis=1)\n",
    "\n",
    "\n",
    "        data_frame['image_path'] = data_frame['image_path'].apply(lambda x: process_image_path(x),meta=('1', 'str'))\n",
    "        data_frame.compute()\n",
    "\n",
    "        #print(\"--- %s seconds ---\" % (time.time() - start_time))\n",
    "        #raise('aa')\n",
    "        print(\"Data Preprocessed...\")\n",
    "        \n",
    "\n",
    "        \n",
    "\n",
    "        parts = range(data_frame.npartitions)\n",
    "\n",
    "        for sub_epoch in range(sub_epochs):\n",
    "\n",
    "            print('Training in sub epoch ' + str(sub_epoch) + ' from ' + str(sub_epochs))\n",
    "\n",
    "            for part_no in parts:\n",
    "\n",
    "                print('Training in part ' + str(part_no))\n",
    "\n",
    "                part = data_frame[['INPUT','image_path','low_res_image']].partitions[part_no]\n",
    "\n",
    "                \n",
    "\n",
    "                input0 = np.array(part['INPUT'])\n",
    "                input0 = np.array([np.array(xi) for xi in input0])\n",
    "                input0[np.isnan(input0)] = 0\n",
    "                input0[input0 > 10] = 10\n",
    "                input0[input0 < -10] = -10\n",
    "                #print(sys.getsizeof(input))\n",
    "                \n",
    "                output = np.array(part['image_path'])\n",
    "                output = np.array([np.array(xi) for xi in output])\n",
    "                output.squeeze().shape\n",
    "                #print(sys.getsizeof(output))\n",
    "\n",
    "                input2 = np.array(part['low_res_image'])\n",
    "                input2 = np.array([np.array(xi) for xi in input2])\n",
    "                input2.squeeze()\n",
    "                #print(sys.getsizeof(input2))\n",
    "                \n",
    "\n",
    "\n",
    "            \n",
    "                \n",
    "                #print(len(procs))\n",
    "                #if(len(procs) > 0):\n",
    "                #    procs[0].join() # Wait for previous process to finish\n",
    "                #    print(\"Model trained using a batch of data...\")\n",
    "                #    procs = []\n",
    "                #print(len(procs))\n",
    "\n",
    "\n",
    "                print(\"Training model ...\")\n",
    "                \n",
    "                if(np.shape(input0)[0] > 0):\n",
    "\n",
    "                    if(MASK_GIT == 1):\n",
    "                        proc = multiprocessing.Process(target=muse_train_batch.train_batch_maskgit1, args=(input0,input2,output,model_filename,inner_epochs,batch_size,sample_every,save_model_every,muse_image_size,unet1_dim,unet2_dim,timesteps,muse_begin_with_image_size,mask_git1_image_size,imagen_samples_folder,sample_probability,dask_chunk,ignore_image_guide,vae_file,))\n",
    "                    else:\n",
    "                        proc = multiprocessing.Process(target=muse_train_batch.train_batch_maskgit2, args=(input0,input2,output,model_filename,inner_epochs,batch_size,sample_every,save_model_every,muse_image_size,unet1_dim,unet2_dim,timesteps,muse_begin_with_image_size,mask_git1_image_size,imagen_samples_folder,sample_probability,dask_chunk,ignore_image_guide,))\n",
    "                    procs.append(proc)\n",
    "                    proc.start()\n",
    "\n",
    "\n",
    "                    if(sub_epochs != 1):\n",
    "                        for ps in procs:\n",
    "                            ps.join()\n",
    "                        procs.clear()\n",
    "\n",
    "\n",
    "                \n",
    "\n",
    "                    #print(\"Model trained using this batch of data...\")\n",
    "\n",
    "                    #proc = multiprocessing.Process(target=train_muse_batch.train_batch_unet2, args=(input,output,model_filename,sub_epochs,batch_size,sample_every,save_model_every,image_size,unet1_dim,unet2_dim,))\n",
    "                    #proc.start()two_unets_pyannote_nemo\n",
    "                    #proc.join()\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "                \n",
    "                    data.close()\n",
    "\n",
    "                    data5 = np.array(df['ID']).astype('int')\n",
    "                    data6 = []\n",
    "                    for x in data5:\n",
    "                        data6.append(int(x))\n",
    "                    print(len(data6))\n",
    "                    #print(data6)\n",
    "                    sql5 = \"UPDATE VIDEO SET TRAINED = 1 WHERE ID IN ({0})\".format(', '.join('?' for _ in data6))\n",
    "                    #print(sql5)\n",
    "                    cur5 = con.cursor()\n",
    "                    cur5.execute(sql5, data6)\n",
    "                    con.commit()\n",
    "                    cur5.close()\n",
    "                else:\n",
    "                    stop_now == True\n",
    "\n",
    "\n",
    "        data10 = con.execute(\"SELECT COUNT(*) FROM VIDEO v WHERE V.TRAINED = 1\")\n",
    "        dataGotten10 = data10.fetchall()\n",
    "\n",
    "\n",
    "        print('-------------')\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "        #print(dataGotten10[0])\n",
    "        if(dataGotten10[0][0] > stop_at_no_of_samples or stop_now == True):\n",
    "            for ps in procs:\n",
    "                ps.join()\n",
    "            procs.clear()\n",
    "            print('--------------- FINISHED ------------------')\n",
    "            stop_now = True\n",
    "            sql50 = \"UPDATE VIDEO SET TRAINED = NULL\"\n",
    "            cur50 = con.cursor()\n",
    "            cur50.execute(sql50)\n",
    "            con.commit()\n",
    "            cur50.close()\n",
    "            print('--------------- TRAINED COLUMN RESET ------------------')\n",
    "        \n",
    "\n",
    "\n",
    "        con.close()\n",
    "        del dataGotten\n",
    "    \n",
    "        del df\n",
    "        del data_frame\n",
    "        del input0\n",
    "        del input2\n",
    "        del output\n",
    "        gc.collect()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "ds2f_m_i",
   "language": "python",
   "name": "python3"
  },
  "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.17"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "78dccc615e96ab04385280185a87a524fe0822daf5dbad2f4bf2e7d7b28366a7"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}