Spaces:
Running
on
Zero
Running
on
Zero
File size: 3,475 Bytes
113884e |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 |
import os
import json
import decord
decord.bridge.set_bridge('torch')
import torch
from torch.utils.data import Dataset
import torchvision
import torchvision.transforms as T
from itertools import islice
from glob import glob
from PIL import Image
from einops import rearrange, repeat
def read_caption_file(caption_file):
with open(caption_file, 'r', encoding="utf8") as t:
return t.read()
def get_text_prompt(
text_prompt: str = '',
fallback_prompt: str= '',
file_path:str = '',
ext_types=['.mp4'],
use_caption=False
):
try:
if use_caption:
if len(text_prompt) > 1: return text_prompt
caption_file = ''
# Use caption on per-video basis (One caption PER video)
for ext in ext_types:
maybe_file = file_path.replace(ext, '.txt')
if maybe_file.endswith(ext_types): continue
if os.path.exists(maybe_file):
caption_file = maybe_file
break
if os.path.exists(caption_file):
return read_caption_file(caption_file)
# Return fallback prompt if no conditions are met.
return fallback_prompt
return text_prompt
except:
print(f"Couldn't read prompt caption for {file_path}. Using fallback.")
return fallback_prompt
def get_video_frames(vr, start_idx, sample_rate=1, max_frames=24):
max_range = len(vr)
frame_number = sorted((0, start_idx, max_range))[1]
frame_range = range(frame_number, max_range, sample_rate)
frame_range_indices = list(frame_range)[:max_frames]
return frame_range_indices
def get_prompt_ids(prompt, tokenizer):
prompt_ids = tokenizer(
prompt,
truncation=True,
padding="max_length",
max_length=tokenizer.model_max_length,
return_tensors="pt",
).input_ids
return prompt_ids
def process_video(vid_path, use_bucketing, w, h, get_frame_buckets, get_frame_batch):
if use_bucketing:
vr = decord.VideoReader(vid_path)
resize = get_frame_buckets(vr)
video = get_frame_batch(vr, resize=resize)
else:
vr = decord.VideoReader(vid_path, width=w, height=h)
video = get_frame_batch(vr)
return video, vr
def min_res(size, min_size): return 192 if size < 192 else size
def up_down_bucket(m_size, in_size, direction):
if direction == 'down': return abs(int(m_size - in_size))
if direction == 'up': return abs(int(m_size + in_size))
def get_bucket_sizes(size, direction: 'down', min_size):
multipliers = [64, 128]
for i, m in enumerate(multipliers):
res = up_down_bucket(m, size, direction)
multipliers[i] = min_res(res, min_size=min_size)
return multipliers
def closest_bucket(m_size, size, direction, min_size):
lst = get_bucket_sizes(m_size, direction, min_size)
return lst[min(range(len(lst)), key=lambda i: abs(lst[i]-size))]
def resolve_bucket(i,h,w): return (i / (h / w))
def sensible_buckets(m_width, m_height, w, h, min_size=192):
if h > w:
w = resolve_bucket(m_width, h, w)
w = closest_bucket(m_width, w, 'down', min_size=min_size)
return w, m_height
if h < w:
h = resolve_bucket(m_height, w, h)
h = closest_bucket(m_height, h, 'down', min_size=min_size)
return m_width, h
return m_width, m_height |