diff --git a/.gitignore b/.gitignore index 9ab52d7..6ab807a 100644 --- a/.gitignore +++ b/.gitignore @@ -14,3 +14,7 @@ frame_list.txt token.txt *.mp4 *.spec +*.mov +libnvrtc-builtins.so.11.1 +*.jpg +*.webm diff --git a/1.json b/1.json old mode 100644 new mode 100755 diff --git a/1.py b/1.py old mode 100644 new mode 100755 index 70f747a..f9179bf --- a/1.py +++ b/1.py @@ -56,12 +56,12 @@ def delete(): async def gpu(ctx): await ctx.channel.send(content=f"bot run on {torch.cuda.get_device_name(0)}\nvram: {int(torch.cuda.get_device_properties('cuda').total_memory/1024/1024/1024)}GB") @bot.command() -async def interpolate(ctx, arg1="--model",arg2="stable_e3", arg3="--discord", arg4="x"): +async def interpolate(ctx, arg1="--model",arg2="converted", arg3="--discord", arg4="x"): wrong=False ytdl=False gifuse=False ytdlurl="none" - model_name="stable_e3" + model_name="converted" ################ ################ if arg1=="--model": @@ -92,7 +92,7 @@ async def interpolate(ctx, arg1="--model",arg2="stable_e3", arg3="--discord", ar else: ytdlurl="https://www.youtube.com/watch?v=dQw4w9WgXcQ" await ctx.channel.send(embed=embedVar) - filename=int(random()*1000000000) + filename=int(random()*100000000000000) delete() ################ Downloading video ################ @@ -111,7 +111,7 @@ async def interpolate(ctx, arg1="--model",arg2="stable_e3", arg3="--discord", ar gifuse=True else: gifuse=False - # + ################ Using cv2 to get fps, width, height, frames number ################ @@ -123,8 +123,6 @@ async def interpolate(ctx, arg1="--model",arg2="stable_e3", arg3="--discord", ar height = video.get(cv2.CAP_PROP_FRAME_HEIGHT) # float `height` round(width) frames = video.get(cv2.CAP_PROP_FRAME_COUNT) length = int(frames/fps)+2 - if length>75: - length=75+7 bitrate = int(63200/length) await message.edit(content=f"bitrate {bitrate}K/s, fps: {fps*2}, frames: {frames}\n") except Exception as e: @@ -132,16 +130,12 @@ async def interpolate(ctx, arg1="--model",arg2="stable_e3", arg3="--discord", ar ################ Using ffmpeg to extract frames and audio ################ - - await message.edit(content=f"starting frame extraction🎞️\n") - if width 0 0 - 759 + 755 521 @@ -475,43 +475,11 @@ Interpolation factor - - - - 10 - 70 - 102 - 17 - - - - half precision - - - true - - - - - - 10 - 100 - 102 - 17 - - - - multi gpu mode - - - true - - 10 - 130 + 70 102 20 @@ -524,7 +492,7 @@ 80 - 130 + 70 141 20 @@ -540,7 +508,7 @@ 10 - 160 + 100 41 21 @@ -553,7 +521,7 @@ 80 - 160 + 100 141 20 @@ -563,7 +531,7 @@ 10 - 190 + 130 102 20 @@ -576,7 +544,7 @@ 80 - 190 + 130 139 20 @@ -671,13 +639,13 @@ 10 - 220 - 211 - 17 + 160 + 92 + 23 - resume (expeprimental work only for 2x) + TRT @@ -720,14 +688,14 @@ - 0 + -10 0 171 23 - save settings look at github repo + save settings @@ -740,7 +708,39 @@ - load + load (beta) + + + + + + 790 + 120 + 102 + 17 + + + + half precision + + + true + + + + + + 950 + 370 + 102 + 17 + + + + multi gpu mode + + + true diff --git a/generate.py b/generate.py old mode 100644 new mode 100755 index c640da9..f105283 --- a/generate.py +++ b/generate.py @@ -7,39 +7,34 @@ import threading import torch import platform +import video import os import numpy as np from tqdm import tqdm import utils -def interpolation(batch_size=5, img_fmt="png", torch_device="cuda", temp_img = "frameseq/", GPUid=1, GPUid2=2, fp16=True, modelp="1.pth"): + +def interpolation(batch_size=5, img_fmt="png", torch_device="cuda", temp_img = "frameseq/", GPUid=0, GPUid2=2, fp16=True, modelp="1.pth", TensorRT=True): #torch.cuda.set_device(GPUid) ossystem=platform.system() print(ossystem) - if fp16==True: - torch.set_default_tensor_type(torch.cuda.HalfTensor) + torch.cuda.device(GPUid) device = torch.device(torch_device) torch.backends.cudnn.enabled = True torch.backends.cudnn.benchmark = True torch.manual_seed(5325) torch.cuda.manual_seed(5325) - from cain import CAIN - model = CAIN(3) - if fp16==True: - model.half() - - - # model to device - if GPUid2==True: - model = torch.nn.DataParallel(model).to("cuda") + if TensorRT==True: + from torch2trt import TRTModule + model_trt = TRTModule() + model_trt.load_state_dict(torch.load(modelp)) else: - model = torch.nn.DataParallel(model, device_ids=[GPUid]) - - checkpoint = torch.load(modelp) - start_epoch = checkpoint['epoch'] + 1 - model.load_state_dict(checkpoint['state_dict']) - del checkpoint + from model.cain import CAIN + model = CAIN(depth=3) + checkpoint = torch.load(modelp) + model.load_state_dict(checkpoint) + model.cuda().half() if ossystem=='Linux': def save(): utils.save_image(out[b], temp_img+"/"+savepath) @@ -66,14 +61,13 @@ def test(): for i, (images, meta) in enumerate(tqdm(test_loader)): # Build input batch - if fp16==True: - im1, im2 = images[0].to(device).half(), images[1].to(device).half() - else: - im1, im2 = images[0].to(device), images[1].to(device) - + im1, im2 = images[0].to(device).half(), images[1].to(device).half() # Forward - out, _ = model(im1, im2) + if TensorRT==True: + out, _ = model_trt(im1, im2) + else: + out, _ = model(im1, im2) for b in range(images[0].size(0)): paths = meta['imgpath'][0][b].split('/') fp = temp_img @@ -99,7 +93,9 @@ def test(): tsave.start() # + test() + video.clean() diff --git a/input_and_output.py b/input_and_output.py old mode 100644 new mode 100755 index f651279..36bccae --- a/input_and_output.py +++ b/input_and_output.py @@ -68,6 +68,8 @@ def list_frame(dir="./frames", text_path="./frames"): print(os.path.join(dir, file)) string = (f"file 'frames/{os.path.join(file)}'\n") txt.write(string) + txt.close() + def ExportVideo(dir_path, proresmode, imtype, fps, factor, filetype, useprores, line): @@ -82,7 +84,7 @@ def ExportVideo(dir_path, proresmode, imtype, fps, factor, filetype, useprores, #list_frame(dir=f"frames") if ossystem=='Linux': if useprores==True: - + print("prores") os.system(f'ffmpeg -r {float(fpss)} -pattern_type glob -i "{dir_path}frames/*.{imtype}" -c:v prores_ks {line} -profile:v {proresmode} "{dir_path}/video.{filetype}"') torch.cuda.empty_cache() else: diff --git a/main.py b/main.py old mode 100644 new mode 100755 index db56b6c..5210c66 --- a/main.py +++ b/main.py @@ -18,6 +18,8 @@ global OutputPath OutputPath="" import platform +os.system("export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/cuda/lib64") + ossystem=platform.system() print(ossystem) ############## @@ -70,10 +72,10 @@ def extract(): input_and_output.ExtractFramesOrSplit(Type=self.comboBox_5.currentText(), chunksize="0", dir_path=OutputPath, Line=self.lineEdit.text(), input=video[2]) def interpolate(): if self.comboBox_2.currentText()=="2x": - generate.interpolation(batch_size=int(self.comboBox_6.currentText()), img_fmt=self.comboBox_5.currentText(), torch_device="cuda", temp_img = f"{OutputPath}/cain/frames", GPUid=self.spin_3.value(), GPUid2=self.checkBox_7.isChecked(), fp16=self.checkBox_4.isChecked(), modelp=self.comboBox_4.currentText()) + generate.interpolation(batch_size=int(self.comboBox_6.currentText()), img_fmt=self.comboBox_5.currentText(), torch_device="cuda", temp_img = f"{OutputPath}/cain/frames", GPUid=self.spin_3.value(), GPUid2=self.checkBox_7.isChecked(), fp16=self.checkBox_4.isChecked(), modelp=self.comboBox_4.currentText(), TensorRT=self.checkBox.isChecked()) else: - generate.interpolation(batch_size=int(self.comboBox_6.currentText()), img_fmt=self.comboBox_5.currentText(), torch_device="cuda", temp_img = f"{OutputPath}/cain/frames", GPUid=self.spin_3.value(), GPUid2=self.checkBox_7.isChecked(), fp16=self.checkBox_4.isChecked(), modelp=self.comboBox_4.currentText()) - generate.interpolation(batch_size=int(self.comboBox_6.currentText()), img_fmt=self.comboBox_5.currentText(), torch_device="cuda", temp_img = f"{OutputPath}/cain/frames", GPUid=self.spin_3.value(), GPUid2=self.checkBox_7.isChecked(), fp16=self.checkBox_4.isChecked(), modelp=self.comboBox_4.currentText()) + generate.interpolation(batch_size=int(self.comboBox_6.currentText()), img_fmt=self.comboBox_5.currentText(), torch_device="cuda", temp_img = f"{OutputPath}/cain/frames", GPUid=self.spin_3.value(), GPUid2=self.checkBox_7.isChecked(), fp16=self.checkBox_4.isChecked(), modelp=self.comboBox_4.currentText(), TensorRT=self.checkBox.isChecked()) + generate.interpolation(batch_size=int(self.comboBox_6.currentText()), img_fmt=self.comboBox_5.currentText(), torch_device="cuda", temp_img = f"{OutputPath}/cain/frames", GPUid=self.spin_3.value(), GPUid2=self.checkBox_7.isChecked(), fp16=self.checkBox_4.isChecked(), modelp=self.comboBox_4.currentText(), TensorRT=self.checkBox.isChecked()) def all(): extract() diff --git a/cain.py b/model/cain.py old mode 100644 new mode 100755 similarity index 100% rename from cain.py rename to model/cain.py diff --git a/partialconv2d.py b/partialconv2d.py old mode 100644 new mode 100755 diff --git a/save.txt b/save.txt old mode 100644 new mode 100755 index aabc650..b4f5e9d --- a/save.txt +++ b/save.txt @@ -7,6 +7,6 @@ lt 2x True 5 -stable_e3.pth +temp_conv.pth False 0 \ No newline at end of file diff --git a/save_loda.py b/save_loda.py old mode 100644 new mode 100755 diff --git a/to copy/1.json b/to copy/1.json old mode 100644 new mode 100755 diff --git a/to copy/avcodec-58.dll b/to copy/avcodec-58.dll old mode 100644 new mode 100755 diff --git a/to copy/avdevice-58.dll b/to copy/avdevice-58.dll old mode 100644 new mode 100755 diff --git a/to copy/avfilter-7.dll b/to copy/avfilter-7.dll old mode 100644 new mode 100755 diff --git a/to copy/avformat-58.dll b/to copy/avformat-58.dll old mode 100644 new mode 100755 diff --git a/to copy/avutil-56.dll b/to copy/avutil-56.dll old mode 100644 new mode 100755 diff --git a/to copy/ffmpeg.exe b/to copy/ffmpeg.exe old mode 100644 new mode 100755 diff --git a/to copy/ffplay.exe b/to copy/ffplay.exe old mode 100644 new mode 100755 diff --git a/to copy/ffprobe.exe b/to copy/ffprobe.exe old mode 100644 new mode 100755 diff --git a/to copy/license/cain.txt b/to copy/license/cain.txt old mode 100644 new mode 100755 diff --git a/to copy/license/gui and few other things.txt b/to copy/license/gui and few other things.txt old mode 100644 new mode 100755 diff --git a/to copy/license/if i forgot about something.txt b/to copy/license/if i forgot about something.txt old mode 100644 new mode 100755 diff --git a/to copy/license/partialconv2d.txt b/to copy/license/partialconv2d.txt old mode 100644 new mode 100755 diff --git a/to copy/license/pytorch_LICENSE b/to copy/license/pytorch_LICENSE old mode 100644 new mode 100755 diff --git a/to copy/license/tqdm.txt b/to copy/license/tqdm.txt old mode 100644 new mode 100755 diff --git a/to copy/postproc-55.dll b/to copy/postproc-55.dll old mode 100644 new mode 100755 diff --git a/to copy/swresample-3.dll b/to copy/swresample-3.dll old mode 100644 new mode 100755 diff --git a/to copy/swscale-5.dll b/to copy/swscale-5.dll old mode 100644 new mode 100755 diff --git a/trt.py b/trt.py new file mode 100644 index 0000000..184c44b --- /dev/null +++ b/trt.py @@ -0,0 +1,29 @@ +import torch +import numpy as np +import torchvision +from model.cain import CAIN +import time +fps=0 +print("Building model: CAIN") +model = CAIN(depth=3) +from torch2trt import torch2trt +data = torch.randn((1, 3, 256, 256)) +data1 = torch.randn((1, 3, 256, 256)) +checkpoint = torch.load("temp_conv.pth") +model.load_state_dict(checkpoint) + +model_trt = torch2trt(model.cuda().half(), [data.cuda().half(),data1.cuda().half()], fp16_mode=True) +startti=time.time() +model_trt.cuda() +torch.save(model_trt.state_dict(), 'converted.pth') +from torch2trt import TRTModule +del model_trt +model_trt = TRTModule() + +model_trt.load_state_dict(torch.load('converted.pth')) +while fps<1000: + x=model_trt(data.cuda(), data1.cuda()) + print(fps) + fps+=1 +endti=time.time() +print(endti-startti) diff --git a/utils.py b/utils.py old mode 100644 new mode 100755 diff --git a/ver.txt b/ver.txt old mode 100644 new mode 100755 diff --git a/video.py b/video.py old mode 100644 new mode 100755 index a5b2140..3cd875a --- a/video.py +++ b/video.py @@ -1,4 +1,5 @@ import os + import glob import numpy as np import torch @@ -6,6 +7,19 @@ from torchvision import transforms from PIL import Image +global img1 +global img2 +img2="nothing" +img1="nothing" +framenum=0 + +def clean(): + global img1 + global img2 + print(img1) + print(img2) + img1="nothing" + img2="nothing" class Video(Dataset): def __init__(self, data_root, fmt='png'): images = sorted(glob.glob(os.path.join(data_root, '*.%s' % fmt))) @@ -17,23 +31,36 @@ def __init__(self, data_root, fmt='png'): # re images = sorted(glob.glob(os.path.join(data_root, '*.%s' % fmt))) self.imglist = [[images[i], images[i+1]] for i in range(len(images)-1)] + img2="nothing" + img1="nothing" print('[%d] images ready to be loaded' % len(self.imglist)) + def __getitem__(self, index): + T = transforms.ToTensor() + global img1 + global img2 imgpaths = self.imglist[index] - - # Load images - img1 = Image.open(imgpaths[0]) + if img1=="nothing": + img1 = Image.open(imgpaths[0]) + else: + img1=img2 img2 = Image.open(imgpaths[1]) + #img2 = Image.open(imgpaths[1]) + #print(img1) + #print(img2) + - T = transforms.ToTensor() - img1 = T(img1) - img2 = T(img2) - # - # - # - # + if torch.is_tensor(img1): + somevar=False + else: + img1 = T(img1) + if torch.is_tensor(img2): + somevar=False + else: + img2 = T(img2) + #img2 = T(img2) imgs = [img1, img2] meta = {'imgpath': imgpaths} return imgs, meta diff --git a/wget.exe b/wget.exe old mode 100644 new mode 100755