Spaces:
Running
Running
| import requests | |
| import json | |
| import base64 | |
| import argparse | |
| VIDEO_WIDTH = 1080 | |
| VIDEO_HEIGHT = 1920 | |
| SPEAKER_ID = 0 | |
| class RestAPIApplication: | |
| def __init__(self, ip, port): | |
| if port < 0: | |
| self.post_request_addr = f"http://{ip}/register/" | |
| self.post_headers = {"Content-Type": "application/json"} | |
| self.generate_addr = (lambda id_: f'http://{ip}/generate/{id_}') | |
| else: | |
| self.post_request_addr = f"http://{ip}:{port}/register/" | |
| self.post_headers = {"Content-Type": "application/json"} | |
| self.generate_addr = (lambda id_: f'http://{ip}:{port}/generate/{id_}') | |
| def _get_json_request(text, lang, duration_rate, action, background_data=None, is_video_background=False): | |
| request_form = dict() | |
| request_form['text'] = text | |
| request_form['speaker'] = SPEAKER_ID | |
| request_form['width'] = VIDEO_WIDTH | |
| request_form['height'] = VIDEO_HEIGHT | |
| request_form['action'] = action | |
| if background_data is not None: | |
| background_base64 = base64.b64encode(background_data).decode("UTF-8") | |
| else: | |
| background_base64 = "" | |
| request_form['background'] = background_base64 | |
| request_form['durationRate'] = duration_rate | |
| request_form['isVideoBackground'] = is_video_background | |
| request_form['lang'] = lang | |
| request_as_json = json.dumps(request_form) | |
| return request_as_json | |
| def _get_video_id(results): | |
| return json.loads(bytes.decode(results.content))['id'] | |
| def get_video(self, text, lang, duration_rate, action, background_data=None, is_video_background=False): | |
| request_json = self._get_json_request(text, lang, duration_rate, action, background_data, is_video_background) | |
| # POST request with jsonified request | |
| results = requests.post(self.post_request_addr, headers=self.post_headers, data=request_json) | |
| # GET video with the given id | |
| video_id = self._get_video_id(results) | |
| video_results = requests.get(self.generate_addr(video_id)) | |
| return video_results.content | |
| def parse_args(): | |
| parser = argparse.ArgumentParser( | |
| description='REST API interface for talking face generation submitted to CVPR2022') | |
| parser.add_argument('-i', '--ip', dest='rest_ip', type=str, default="127.0.0.1", help="IP for REST API") | |
| parser.add_argument('-p', '--port', dest='rest_port', type=int, default=8080, help="Port for REST API") | |
| args = parser.parse_args() | |
| return args | |
| if __name__ == '__main__': | |
| args = parse_args() | |
| rest_api_application = RestAPIApplication(args.rest_ip, args.rest_port) | |