80 lines
6.3 KiB
Python
80 lines
6.3 KiB
Python
"""Bounded, tiled Real-ESRGAN video inference. No CPU fallback or frame interpolation."""
|
|
import argparse, json, subprocess, sys, os, hashlib
|
|
from fractions import Fraction
|
|
from pathlib import Path
|
|
import numpy as np
|
|
import torch
|
|
from PIL import Image
|
|
sys.path.insert(0, str(Path(__file__).parent / 'vendor'))
|
|
from srvgg_arch import SRVGGNetCompact
|
|
|
|
def probe(file):
|
|
return json.loads(subprocess.check_output(['ffprobe','-v','error','-protocol_whitelist','file,pipe','-show_streams','-show_format','-of','json',file], timeout=30))
|
|
|
|
def main():
|
|
p=argparse.ArgumentParser()
|
|
for name in ['input','output','model','resolution']: p.add_argument('--'+name, required=True)
|
|
p.add_argument('--tile', type=int, default=192)
|
|
a=p.parse_args()
|
|
if not torch.cuda.is_available(): raise RuntimeError('CUDA unavailable; refusing CPU fallback')
|
|
torch.set_num_threads(4)
|
|
info=probe(a.input); v=next(s for s in info['streams'] if s['codec_type']=='video')
|
|
if not any(name in info['format']['format_name'].split(',') for name in ['mov','mp4','matroska','webm','avi','flv','mpegts']): raise ValueError('Unsupported video container')
|
|
rotation=next((int(side['rotation'])%360 for side in v.get('side_data_list',[]) if 'rotation' in side),0)
|
|
if rotation%90: raise ValueError('Unsupported rotation angle')
|
|
w,h=int(v['width']),int(v['height'])
|
|
if rotation in [90,270]: w,h=h,w
|
|
fps=Fraction(v['avg_frame_rate'])
|
|
duration=float(info['format']['duration'])
|
|
if not 0<duration<=120 or not 0<float(fps)<=60 or w*h>3840*2160: raise ValueError('Input exceeds 120 s / 60 fps / 4K limits')
|
|
# Reject variable frame rate rather than silently altering timing/audio sync.
|
|
frame_times=subprocess.check_output(['ffprobe','-v','error','-protocol_whitelist','file,pipe','-select_streams','v:0','-show_entries','frame=best_effort_timestamp_time','-of','csv=p=0',a.input],timeout=60).decode().splitlines()
|
|
timestamps=[float(s.split(',')[0]) for s in frame_times if s.split(',')[0].strip()]
|
|
if any(abs((b-c)-float(1/fps))>max(.002,float(1/fps)*.05) for c,b in zip(timestamps,timestamps[1:])): raise ValueError('Variable frame rate video must be converted to constant frame rate first')
|
|
long,short={'1080P':(1920,1080),'2K':(2560,1440),'4K':(3840,2160)}[a.resolution]
|
|
scale=min((long if w>=h else short)/w,(short if w>=h else long)/h)
|
|
if not 1<scale<=4: raise ValueError('Target must exceed source resolution by no more than 4x')
|
|
ow,oh=int(w*scale)//2*2,int(h*scale)//2*2
|
|
model=SRVGGNetCompact(num_in_ch=3,num_out_ch=3,num_feat=64,num_conv=32,upscale=4,act_type='prelu')
|
|
if hashlib.sha256(Path(a.model).read_bytes()).hexdigest() != '8dc7edb9ac80ccdc30c3a5dca6616509367f05fbc184ad95b731f05bece96292': raise RuntimeError('Model checksum mismatch')
|
|
weights=torch.load(a.model,map_location='cpu',weights_only=True)
|
|
model.load_state_dict(weights.get('params_ema',weights.get('params',weights)),strict=True)
|
|
model.eval().to('cuda')
|
|
temp=a.output+'.silent.mp4'
|
|
decode=subprocess.Popen(['ffmpeg','-v','error','-filter_threads','1','-filter_complex_threads','1','-threads','2','-protocol_whitelist','file,pipe','-i',a.input,'-map','0:v:0','-f','rawvideo','-pix_fmt','rgb24','-vsync','0','pipe:1'],stdout=subprocess.PIPE)
|
|
encode=subprocess.Popen(['ffmpeg','-v','error','-filter_threads','1','-filter_complex_threads','1','-y','-f','rawvideo','-pix_fmt','rgb24','-s',f'{ow}x{oh}','-r',str(fps),'-i','pipe:0','-an','-c:v','libx264','-threads','4','-preset','fast','-crf','18','-pix_fmt','yuv420p',temp],stdin=subprocess.PIPE)
|
|
count=0
|
|
try:
|
|
with torch.inference_mode():
|
|
while True:
|
|
raw=decode.stdout.read(w*h*3)
|
|
if not raw: break
|
|
if len(raw)!=w*h*3: raise RuntimeError('Truncated video frame')
|
|
source=np.frombuffer(raw,dtype=np.uint8).reshape(h,w,3)
|
|
# CPU output buffer, bounded GPU tiles with enough halo for 35 convolutions.
|
|
output=np.empty((h*4,w*4,3),dtype=np.uint8)
|
|
for y in range(0,h,a.tile):
|
|
for x in range(0,w,a.tile):
|
|
x0,y0=max(0,x-40),max(0,y-40); x1,y1=min(w,x+a.tile+40),min(h,y+a.tile+40)
|
|
tile=torch.from_numpy(source[y0:y1,x0:x1].copy()).permute(2,0,1).float().unsqueeze(0).to('cuda')/255
|
|
enhanced=model(tile).clamp_(0,1).squeeze(0).permute(1,2,0).mul_(255).round_().byte().cpu().numpy()
|
|
ex,ey=min(w,x+a.tile),min(h,y+a.tile)
|
|
output[y*4:ey*4,x*4:ex*4]=enhanced[(y-y0)*4:(ey-y0)*4,(x-x0)*4:(ex-x0)*4]
|
|
frame=Image.fromarray(output)
|
|
if (ow,oh)!=(w*4,h*4): frame=frame.resize((ow,oh),Image.Resampling.LANCZOS)
|
|
encode.stdin.write(frame.tobytes()); count+=1
|
|
if count%5==0: print(json.dumps({'stage':f'AI 超分 {count}/{len(timestamps)} 帧'}),flush=True)
|
|
encode.stdin.close()
|
|
if decode.wait(timeout=30)!=0 or encode.wait(timeout=120)!=0: raise RuntimeError('Video codec failed')
|
|
finally:
|
|
for process in [decode,encode]:
|
|
if process.poll() is None: process.kill(); process.wait()
|
|
if count!=len(timestamps) or count==0: raise RuntimeError('Output frame count mismatch')
|
|
subprocess.run(['ffmpeg','-v','error','-filter_threads','1','-filter_complex_threads','1','-y','-protocol_whitelist','file,pipe','-i',temp,'-protocol_whitelist','file,pipe','-i',a.input,'-map','0:v:0','-map','1:a?','-map_metadata','1','-metadata:s:v:0','rotate=0','-c:v','copy','-c:a','aac','-b:a','192k','-movflags','+faststart',a.output],check=True,timeout=120)
|
|
actual=probe(a.output); stream=next(s for s in actual['streams'] if s['codec_type']=='video')
|
|
if int(stream['nb_frames'])!=count or abs(float(actual['format']['duration'])-duration)>.2: raise RuntimeError('Output timing mismatch')
|
|
if any(s['codec_type']=='audio' for s in actual['streams']) != any(s['codec_type']=='audio' for s in info['streams']): raise RuntimeError('Output audio stream missing')
|
|
os.unlink(temp)
|
|
print(json.dumps({'result':{'width':ow,'height':oh,'durationSeconds':float(actual['format']['duration']),'fps':float(fps),'frames':count,'audio':any(s['codec_type']=='audio' for s in actual['streams']),'engine':'Real-ESRGAN','model':'realesr-general-x4v3'}}),flush=True)
|
|
if __name__=='__main__': main()
|