File size: 3,311 Bytes
a4ba82a | 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 | from fractions import Fraction
import io
import json
from PIL import Image
from .bootstrap import production as a
def instruction_suffix(instruction):
value = instruction.strip()
if not value:
return ''
if len(value) > 2000:
raise ValueError('Keep the optional task instruction below 2000 characters.')
return '\n' + json.dumps({'episode_instruction': value}, ensure_ascii=False, sort_keys=True, separators=(',', ':'), allow_nan=False)
def human_transport(frames):
from flora_runtime.exact_frame_serving import encode_regularized_image
output = {}
for index, image in frames.items():
payload = encode_regularized_image(image, encoding='jpeg_q98_444')['payload']
with Image.open(io.BytesIO(payload)) as decoded:
output[index] = decoded.convert('RGB')
return output
def temporal_inputs(frames, window, pool, proc, roles, instruction=''):
from transformers.video_utils import VideoMetadata
from flora_runtime.exact_frame_serving import regularized_resize_dims
profile = window.window_profile
if len(frames) != profile.slot_count:
raise ValueError('Incorrect native temporal slot count')
regularized = []
for image in frames:
image = image.convert('RGB')
height, width = regularized_resize_dims(height=image.height, width=image.width, min_pixels=256, max_pixels=40000 if len(roles) == 1 else 160000)
if image.size != (width, height):
image = image.resize((width, height), Image.Resampling.BICUBIC)
regularized.append(image)
prompt = a.temporal_prompt(window, pool) + instruction_suffix(instruction)
text = proc.apply_chat_template([dict(role='user', content=prompt.replace('<video>', proc.video_token))], tokenize=False, add_generation_prompt=True, enable_thinking=False)
inputs = proc(text=[text], videos=[regularized], video_metadata=[VideoMetadata(total_num_frames=profile.slot_count, fps=profile.slot_fps, frames_indices=list(range(profile.slot_count)))], do_sample_frames=False, return_metadata=True, return_tensors='pt')
inputs.pop('video_metadata', None)
return (inputs, text)
def three_view_caption_inputs(plan, frames, proc, instruction=''):
from transformers.video_utils import VideoMetadata
from plans import HEADING, INSTRUCTION
time_base = Fraction(plan['time_base'])
content, videos, metadata = ([HEADING], [], [])
for block in plan['blocks']:
times = ', '.join((f'{float(t * time_base):.9f}' for t in block['pts_ticks']))
content.append(block['role'] + ':\nFrames are in this order; native source timestamps in seconds: [' + times + '].\n' + proc.video_token)
videos.append([frames[i] for i in block['indices']])
metadata.append(VideoMetadata(total_num_frames=len(block['indices']), fps=float(Fraction(plan['clock_fps'])), frames_indices=[i - block['source_start_frame'] for i in block['indices']]))
content.append(INSTRUCTION + instruction_suffix(instruction))
text = proc.apply_chat_template([dict(role='user', content=''.join(content))], tokenize=False, add_generation_prompt=True, enable_thinking=False)
result = proc(text=[text], videos=videos, video_metadata=metadata, do_sample_frames=False, return_tensors='pt', padding=True)
return (result, text)
|