train_imagen_all_stable.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "UNET = 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": [
    "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",
    "image_size =  int(configParser.get('COMMON', 'resizeImageTo'))\n",
    "begin_with_image_size = int(configParser.get('COMMON', 'begin_with_image_size'))\n",
    "unet1_dim =  int(configParser.get('COMMON', 'unet1_dim'))\n",
    "unet2_dim =  int(configParser.get('COMMON', 'unet2_dim'))\n",
    "unet1_image_size =  int(configParser.get('COMMON', 'unet1_image_size'))\n",
    "audio_embs =  str(configParser.get('COMMON', 'audio_embs'))\n",
    "audio_length_used =  configParser.get('train_imagen', 'audio_length_used') \n",
    "model_filename =  configParser.get('train_imagen', 'model_filename')\n",
    "sub_epochs=  int(configParser.get('train_imagen', 'sub_epochs') )\n",
    "inner_epochs=  int(configParser.get('train_imagen', 'inner_epochs') )\n",
    "batch_size=  int(configParser.get('train_imagen', 'batch_size') )\n",
    "timesteps= int(configParser.get('COMMON', 'timesteps') )\n",
    "sample_every=  int(configParser.get('train_imagen', 'sample_every') ) - 1\n",
    "sample_probability = int(configParser.get('train_imagen', 'sample_probability') )\n",
    "save_model_every=  int(configParser.get('train_imagen', 'save_model_every') )\n",
    "sample_every_offset=  int(configParser.get('train_imagen', 'sample_every_offset') ) - 1\n",
    "save_every_offset=  int(configParser.get('train_imagen', 'save_every_offset') ) - 1\n",
    "imagen_samples_folder = configParser.get('train_imagen', 'imagen_samples_folder') \n",
    "db_chunk = int(configParser.get('train_imagen', 'db_chunk'))\n",
    "dask_chunk = int(configParser.get('train_imagen', 'dask_chunk'))\n",
    "stop_at_no_of_samples = int(configParser.get('train_imagen', 'stop_at_no_of_samples'))\n",
    "epochs = int(configParser.get('train_imagen', 'epochs'))\n",
    "ignore_speaker_embedding = bool(int(configParser.get('train_imagen', 'ignore_speaker_embedding') ))\n",
    "ignore_speech_brain = bool(int(configParser.get('train_imagen', 'ignore_speech_brain') ))\n",
    "ignore_pyannote_titanet_speakernet = bool(int(configParser.get('train_imagen', 'ignore_pyannote_titanet_speakernet') ))\n",
    "ignore_audio_features = bool(int(configParser.get('train_imagen', 'ignore_audio_features') ))\n",
    "ignore_pyAudioAnalysis = bool(int(configParser.get('train_imagen', 'ignore_pyAudioAnalysis') ))\n",
    "ignore_librosa = bool(int(configParser.get('train_imagen', 'ignore_librosa') ))\n",
    "ignore_image_guide = bool(int(configParser.get('train_imagen', 'ignore_image_guide') ))\n",
    "ignore_additional_attributes = bool(int(configParser.get('train_imagen', 'ignore_additional_attributes') ))\n",
    "ignore_age = bool(int(configParser.get('train_imagen', 'ignore_age') ))\n",
    "ignore_gender = bool(int(configParser.get('train_imagen', 'ignore_gender') ))\n",
    "ignore_ethnicity = bool(int(configParser.get('train_imagen', 'ignore_ethnicity') ))\n",
    "ignore_language_spoken = bool(int(configParser.get('train_imagen', 'ignore_language_spoken') ))\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 = image_size / (1+2 * 0.4)\n",
    "        h_s = 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((image_size,image_size))\n",
    "\n",
    "        image = image.resize((begin_with_image_size,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",
    "        try:\n",
    "            x = np.ones(768) * (age / 100.0)\n",
    "            x[767] = 0\n",
    "        except:\n",
    "            x = np.zeros(768)\n",
    "            x[767] = 1\n",
    "\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 = image_size / (1+2 * 0.4)\n",
    "    h_s = 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",
    "    image = image.resize((image_size,image_size))\n",
    "\n",
    "    #print('saving')\n",
    "    #image.save(str(random.random()) + '.png')\n",
    "\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": [
    "from scipy.signal import savgol_filter\n",
    "from torch.utils.data import TensorDataset, DataLoader\n",
    "import time\n",
    "import torch\n",
    "from torch.utils.data import Dataset\n",
    "from torch.utils.data import DataLoader\n",
    "from torchvision.transforms import ToTensor,Compose\n",
    "import dask.dataframe as dd\n",
    "from multiprocessing import Manager,RLock,Value,Queue\n",
    "import ctypes\n",
    "\n",
    "class CustomDataset(Dataset):\n",
    "\n",
    "    total_rows = None\n",
    "    con = None\n",
    "    lock = None\n",
    "    samples_seen =  Value('i', 0,lock=False)\n",
    "    \n",
    "    def no_of_samples_seen(self):\n",
    "        return self.samples_seen.value\n",
    "\n",
    "    def __init__(self,datasetPathDatabase_a,lock,last_sample_no):\n",
    "        self.lock = lock\n",
    "        self.samples_seen.value = last_sample_no\n",
    "        self.con = sl.connect(datasetPathDatabase_a)\n",
    "        pass\n",
    "\n",
    "\n",
    "\n",
    "    def __len__(self):\n",
    "        \n",
    "        if(self.total_rows == None):\n",
    "            self.lock.acquire(block=True)\n",
    "            data = self.con.execute(\"SELECT COUNT(*) FROM VIDEO v \")\n",
    "            count = data.fetchall()\n",
    "            self.total_rows = count[0][0]\n",
    "            data.close()\n",
    "            print(\"Found \" + str(self.total_rows) + \" data points.\")\n",
    "\n",
    "            #data = self.con.execute(\"SELECT COUNT(*) FROM VIDEO v WHERE V.TRAINED = 1\")\n",
    "            #count = data.fetchall()\n",
    "            #data.close()\n",
    "            #self.samples_seen.value = count[0][0]\n",
    "            #print(\"Starting from \" + str(self.samples_seen.value) + \" data points.\")\n",
    "            self.lock.release()\n",
    "\n",
    "        return self.total_rows\n",
    "\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        self.lock.acquire(block=True)\n",
    "        start_time = time.time()\n",
    "        sql = (\"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(1) + \")\")\n",
    "        data = self.con.execute(sql)\n",
    "        dataGotten = data.fetchall()\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",
    "        data.close()\n",
    "        \n",
    "        #print(\"aaaa\" + str(len(df)))\n",
    "        if(len(df) == 0):\n",
    "            print(\"Trying to get items\")\n",
    "            data = self.con.execute(sql)\n",
    "            dataGotten = data.fetchall()\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",
    "            data.close()\n",
    "            \n",
    "            if(len(df) == 0):\n",
    "                print(\"Trying to get items\")\n",
    "                data = self.con.execute(sql)\n",
    "                dataGotten = data.fetchall()\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",
    "                data.close()\n",
    "\n",
    "                if(len(df) == 0):\n",
    "                    print(\"Trying to get items\")\n",
    "                    data = self.con.execute(sql)\n",
    "                    dataGotten = data.fetchall()\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",
    "                    data.close()\n",
    "\n",
    "                    if(len(df) == 0):\n",
    "                        print(\"Reintializing dataset\")\n",
    "                        sql50 = \"UPDATE VIDEO SET TRAINED = NULL\"\n",
    "                        cur50 = self.con.cursor()\n",
    "                        cur50.execute(sql50)\n",
    "                        self.con.commit()\n",
    "                        cur50.close()\n",
    "\n",
    "                        data = self.con.execute(sql)\n",
    "                        dataGotten = data.fetchall()\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",
    "                        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 = self.con.cursor()\n",
    "        cur5.execute(sql5, data6)\n",
    "        self.con.commit()\n",
    "        cur5.close()\n",
    "        self.samples_seen.value +=  1\n",
    "        print (self.samples_seen.value,end='\\r')\n",
    "\n",
    "\n",
    "        self.lock.release()\n",
    "\n",
    "        \n",
    "        \n",
    "        df = df.sample(frac=1).reset_index(drop=True)\n",
    "        df = dd.from_pandas(df, npartitions=1)\n",
    "        df[\"image_guide\"] = np.nan\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",
    "        part = data_frame[['INPUT','image_path','low_res_image']].partitions[0]\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",
    "        #    prProcess,\n",
    "        input0 = torch.from_numpy(input0)\n",
    "        input0 = input0.to(torch.float)\n",
    "\n",
    "\n",
    "        input2 = torch.from_numpy(input2)\n",
    "        input2 = input2.to(torch.float)\n",
    "\n",
    "        output = torch.from_numpy(output)\n",
    "        output = output.to(torch.float)\n",
    "\n",
    "        input0 = input0.squeeze()\n",
    "        input2 = input2.squeeze()\n",
    "        output = output.squeeze()\n",
    "\n",
    "        return output,input0\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_loss(sample_no,loss_value):\n",
    "\n",
    "    if(UNET == 1):\n",
    "        my_file = Path(model_filename + '/loss_total_1.picke')\n",
    "        if my_file.is_file():\n",
    "            with open(model_filename + '/loss_total_1.picke', 'rb') as handle:\n",
    "                loss_total = pickle.load(handle)\n",
    "        else:\n",
    "            loss_total = []\n",
    "    else:\n",
    "        my_file = Path(model_filename + '/loss_total_2.picke')\n",
    "        if my_file.is_file():\n",
    "            with open(model_filename + '/loss_total_2.picke', 'rb') as handle:\n",
    "                loss_total = pickle.load(handle)\n",
    "        else:\n",
    "            loss_total = []\n",
    "    \n",
    "    #print(loss_list)\n",
    "    #print(loss_total)\n",
    "    loss_total.append([sample_no,loss_value])\n",
    "\n",
    "    if(UNET == 1):\n",
    "        with open(model_filename +'/loss_total_1.picke', 'wb') as handle:\n",
    "            pickle.dump(loss_total, handle)\n",
    "    else:\n",
    "        with open(model_filename +'/loss_total_2.picke', 'wb') as handle:\n",
    "            pickle.dump(loss_total, handle)\n",
    "    #print(loss_total)\n",
    "    fig = plt.figure()\n",
    "    x_val = [x[0] for x in loss_total]\n",
    "    y_val = [x[1] for x in loss_total]\n",
    "    plt.plot(x_val,y_val)\n",
    "    plt.title(\"Training Loss\")\n",
    "    plt.xlabel(\"Training Sample\")\n",
    "    plt.ylabel(\"MSE Loss\")\n",
    "    \n",
    "\n",
    "\n",
    "\n",
    "    if(UNET == 1):\n",
    "        if(os.path.isfile(model_filename + '/loss_1_plot.png')):\n",
    "            os.remove(model_filename + '/loss_1_plot.png')\n",
    "        fig.savefig(model_filename + '/loss_1_plot.png')\n",
    "    else:\n",
    "        if(os.path.isfile(model_filename + '/loss_2_plot.png')):\n",
    "            os.remove(model_filename + '/loss_2_plot.png')\n",
    "        fig.savefig(model_filename + '/loss_2_plot.png')\n",
    "    plt.close()\n",
    "\n",
    "    try:\n",
    "        fig = plt.figure()\n",
    "        plt.plot(x_val[1000::],y_val[1000::]\n",
    "                ,'.',zorder=-100)\n",
    "\n",
    "        #plt.axvline(x=1000,linestyle='--',color='green',label='1000 inner epochs')\n",
    "        #plt.axvline(x=17000,linestyle='--',color='purple',label='100 inner epochs - 100000 unique samples seen' )\n",
    "        #plt.axvline(x=23100,linestyle='-.',color='black',label='1 inner epoch - end of 1st epoch ')\n",
    "        #plt.legend(bbox_to_anchor = (1.0, 1), loc = 'upper right')\n",
    "        plt.grid()\n",
    "        yhat = savgol_filter(y_val, 1000, 3)\n",
    "        plt.plot(x_val[1000::],yhat[1000::],'r')\n",
    "        plt.title(\"Training Loss\")\n",
    "        plt.xlabel(\"Training Sample\")\n",
    "        plt.ylabel(\"MSE Loss\")\n",
    "        \n",
    "\n",
    "\n",
    "\n",
    "        if(UNET == 1):\n",
    "            if(os.path.isfile(model_filename + '/loss_zoomed_1_plot.png')):\n",
    "                os.remove(model_filename + '/loss_zoomed_1_plot.png')\n",
    "            fig.savefig(model_filename + '/loss_zoomed_1_plot.png')\n",
    "        else:\n",
    "            if(os.path.isfile(model_filename + '/loss_zoomed_2_plot.png')):\n",
    "                os.remove(model_filename + '/loss_zoomed_2_plot.png')\n",
    "            fig.savefig(model_filename + '/loss_zoomed_2_plot.png')\n",
    "        plt.close()\n",
    "\n",
    "\n",
    "    except:\n",
    "        try:\n",
    "            plt.close()\n",
    "        except:\n",
    "            pass\n",
    "\n",
    "\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "last_sample_no = 0\n",
    "if(UNET == 1):\n",
    "    my_file = Path(model_filename + '/last_sample_no1.picke')\n",
    "    if my_file.is_file():\n",
    "        with open(model_filename + '/last_sample_no1.picke', 'rb') as handle:\n",
    "            last_sample_no = pickle.load(handle)\n",
    "else:\n",
    "    my_file = Path(model_filename + '/last_sample_no2.picke')\n",
    "    if my_file.is_file():\n",
    "        with open(model_filename + '/last_sample_no2.picke', 'rb') as handle:\n",
    "            last_sample_no = pickle.load(handle)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "my_dataset = CustomDataset(datasetPathDatabase,RLock(),last_sample_no)\n",
    "print(len(my_dataset))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "my_dataloader = DataLoader(my_dataset,batch_size = batch_size * 16,shuffle=False,num_workers=16) # create your dataloader"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from imagen_pytorch import Unet, Imagen, ImagenTrainer,NullUnet\n",
    "from imagen_pytorch.data import Dataset\n",
    "\n",
    "unet0 = NullUnet()  # add a placeholder \"null\" unet for the base unet\n",
    "\n",
    "unet1 = Unet(\n",
    "    dim = unet1_dim,\n",
    "    cond_dim = 512,\n",
    "    dim_mults = (1, 2, 4, 8),\n",
    "    num_resnet_blocks = 3,\n",
    "    layer_attns = (False, True, True, True),\n",
    "    layer_cross_attns = (False, True, True, True)\n",
    ")\n",
    "\n",
    "unet2 = Unet(\n",
    "    dim = unet2_dim,\n",
    "    cond_dim = 512,\n",
    "    dim_mults = (1, 2, 4, 8),\n",
    "    num_resnet_blocks = (2, 4, 8, 8),\n",
    "    layer_attns = (False, False, False, True),\n",
    "    layer_cross_attns = (False, False, False, True)\n",
    ")\n",
    "\n",
    "if(ignore_image_guide):\n",
    "    unets = (unet1, unet2)\n",
    "    image_sizes= (unet1_image_size, image_size)\n",
    "    unet_to_train = UNET #1\n",
    "    #start_image_or_video = None\n",
    "else:\n",
    "    unets = (unet0,unet1, unet2)\n",
    "    image_sizes= (begin_with_image_size,unet1_image_size, image_size)\n",
    "    unet_to_train = UNET + 1 #2\n",
    "    #start_image_or_video = input2[:1,:]\n",
    "\n",
    "#print(input2[:1,:])\n",
    "#print(input2)\n",
    "\n",
    "imagen = Imagen(\n",
    "    unets = unets,\n",
    "    image_sizes = image_sizes,\n",
    "    timesteps = timesteps,\n",
    "    cond_drop_prob = 0.1\n",
    ").cuda()\n",
    "\n",
    "\n",
    "\n",
    "trainer = ImagenTrainer(\n",
    "    imagen = imagen,\n",
    "    split_valid_from_train = False, # whether to split the validation dataset from the training\n",
    "    checkpoint_every = save_model_every,\n",
    "    checkpoint_path = model_filename,\n",
    "    max_checkpoints_keep = 1,\n",
    "    only_train_unet_number = unet_to_train\n",
    ").cuda()\n",
    "\n",
    "\n",
    "trainer.add_train_dataloader(my_dataloader)\n",
    "\n",
    "#if not os.path.exists(imagen_samples_folder):\n",
    "#    os.makedirs(imagen_samples_folder)\n",
    "\n",
    "#my_file = Path(model_filename)\n",
    "#if my_file.is_file():\n",
    "#    print('Using model file ' + model_filename)\n",
    "#    trainer.load(model_filename)\n",
    "\n",
    "#trainer.load_from_checkpoint_folder()\n",
    "#print(epochs)\n",
    "for i in range(epochs):\n",
    "    \n",
    "    loss = trainer.train_step(unet_number = unet_to_train,max_batch_size = batch_size)\n",
    "    print(f'loss: {loss}')\n",
    "    #print(\"At sample seen \" + str(my_dataset.no_of_samples_seen()))\n",
    "    plot_loss(my_dataset.no_of_samples_seen(),loss)\n",
    "    #print(loss_list)\n",
    "    #trainer.save_to_checkpoint_folder()\n",
    "    print('epoch:' + str(i+1) + ' from ' + str(epochs))\n",
    "    last_sample_no = my_dataset.no_of_samples_seen()\n",
    "    if(UNET == 1):\n",
    "        with open(model_filename + '/last_sample_no1.picke', 'wb') as handle:\n",
    "            pickle.dump(last_sample_no, handle)\n",
    "    else:\n",
    "        with open(model_filename + '/last_sample_no2.picke', 'wb') as handle:\n",
    "            pickle.dump(last_sample_no, handle)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "trainer.save_to_checkpoint_folder()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from torchvision.utils import save_image\n",
    "from torchvision.transforms import Resize\n",
    "cond_scale = random.uniform(5.1, 9.9)\n",
    "now =time.time()\n",
    "seconds = now\n",
    "output, input0 = my_dataset.__getitem__(1)\n",
    "save_image(output,'ground_truth.png')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "resize_img_tensor = Resize((begin_with_image_size, begin_with_image_size))\n",
    "\n",
    "if(ignore_image_guide):\n",
    "    start_image_or_video = None\n",
    "else:\n",
    "    start_image_or_video = resize_img_tensor(output)\n",
    "\n",
    "start_image_or_video = torch.unsqueeze(start_image_or_video,0)\n",
    "input0 = torch.unsqueeze(input0,0)\n",
    "output = torch.unsqueeze(output,0)\n",
    "\n",
    "print(start_image_or_video.shape)\n",
    "print(input0.shape)\n",
    "print(output.shape)\n",
    "\n",
    "images = trainer.sample(text_embeds=input0,start_image_or_video = start_image_or_video,start_at_unet_number = unet_to_train\n",
    "            ,stop_at_unet_number=unet_to_train,batch_size = 1, return_pil_images = True,cond_scale=cond_scale) # returns List[Image]\n",
    "images[0].save('generated.png')"
   ]
  }
 ],
 "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
}