| import aiohttp |
| import asyncio |
| from itertools import chain |
|
|
|
|
| class VideoGenerator: |
| def __init__(self): |
| self.base_urls = [f"https://yakova-depthflow-{i}.hf.space" for i in range(10)] |
| self.headers = {"accept": "application/json"} |
| self.default_params = { |
| "frame_rate": 30, |
| "duration": 3, |
| "quality": 1, |
| "ssaa": 0.8, |
| "raw": "false", |
| } |
|
|
| async def generate_video(self, base_url, params): |
| url = f"{base_url}/generate_video" |
|
|
| async with aiohttp.ClientSession() as session: |
| async with session.post( |
| url, params=params, headers=self.headers |
| ) as response: |
| if response.status == 200: |
| data = await response.json() |
| output_file = data.get("output_file") |
| return output_file |
| else: |
| print(f"Request to {url} failed with status: {response.status}") |
| return None |
|
|
| async def check_video_ready(self, base_url, output_file): |
| url = f"{base_url}/download/{output_file}" |
|
|
| async with aiohttp.ClientSession() as session: |
| while True: |
| async with session.get(url, headers=self.headers) as response: |
| if response.status == 200: |
| video_content = await response.read() |
| if len(video_content) > 0: |
| return url |
| else: |
| print( |
| f"Video {output_file} is ready but the file size is zero, retrying in 10 seconds..." |
| ) |
| await asyncio.sleep(10) |
| elif response.status == 404: |
| data = await response.json() |
| if data.get("detail") == "Video not found": |
| print( |
| f"Video {output_file} not ready yet, retrying in 10 seconds..." |
| ) |
| await asyncio.sleep(180) |
| else: |
| print(f"Unexpected response for {output_file}: {data}") |
| return None |
| else: |
| print(f"Request to {url} failed with status: {response.status}") |
| return None |
|
|
| async def process_image(self, base_url, image_link): |
| params = self.default_params.copy() |
| params["image_link"] = image_link |
|
|
| output_file = await self.generate_video(base_url, params) |
| if output_file: |
| print(f"Generated video file id: {output_file} for image {image_link}") |
| video_url = await self.check_video_ready(base_url, output_file) |
| if video_url: |
| print( |
| f"Video for {image_link} is ready and can be downloaded from: {video_url}" |
| ) |
| return video_url |
| else: |
| print(f"Failed to get the video URL for {image_link}") |
| return None |
| else: |
| print(f"Failed to generate the video for {image_link}") |
| return None |
|
|
| def flatten(self, nested_list): |
| return list(chain.from_iterable(nested_list)) |
|
|
| def nest(self, flat_list, nested_dims): |
| it = iter(flat_list) |
| return [[next(it) for _ in inner_list] for inner_list in nested_dims] |
|
|
| async def run(self, nested_image_links): |
| flat_image_links = self.flatten(nested_image_links) |
| tasks = [] |
| base_index = 0 |
|
|
| for image_link in flat_image_links: |
| base_url = self.base_urls[base_index % len(self.base_urls)] |
| tasks.append(self.process_image(base_url, image_link)) |
| base_index += 1 |
|
|
| flat_video_urls = await asyncio.gather(*tasks) |
| nested_video_urls = self.nest(flat_video_urls, nested_image_links) |
| return nested_video_urls |
|
|
|
|
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
|
|
| |
| |
| |
|
|
| |
|
|
