Coverage for wrapper/realesrgan/wrapper_realesrgan.py: 77%
217 statements
« prev ^ index » next coverage.py v7.15.4, created at 2026-08-09 04:47 +0000
1"""
2Wrapper for Real-ESRGAN image and video upscaling.
3"""
4import os
5import logging
6import tempfile
7import aiofiles
8import asyncio
10from PIL import Image
12from datetime import timedelta
14import torch
15from torch import distributed as dist
16from torch import inference_mode
18from typing import override
19from typing import List
20from typing import Optional
21from typing import Union
22from typing import Dict
23from typing import Any
25from model_timing import GenTimer
26from wrapper_model import ModelGeneration
28from console_utils import bytes_to_human
30from image_utils import base64_to_img
31from media_utils import base64_to_video_frames
32from file_utils import base64_to_binary
33from media_utils import get_video_fps
34from media_utils import save_video_frames
35from media_utils import get_video_size
37from RealESRGAN import RealESRGAN
40class RealESRGANGeneration(ModelGeneration):
41 """Handle image and video upscaling using Real-ESRGAN."""
43 def __init__(self) -> None:
44 super().__init__("realesrgan")
46 # Model components
47 self.SCALING_FACTORS = [2, 4, 8]
48 self.models: Dict[int, RealESRGAN] = {}
50 self.GPU: Optional[str] = None
51 if torch.cuda.is_available():
52 self.GPU = torch.cuda.get_device_name(0)
54 def __del__(self) -> None:
55 if self.models:
56 del self.models
57 super().__del__()
59 def init_parallelism(self) -> None:
60 self.load_timer.start("torch_dist")
62 self.rank = int(os.getenv("RANK", 0))
63 self.local_rank = int(os.getenv("LOCAL_RANK", 0))
64 self.world_size = int(os.getenv("WORLD_SIZE", 1))
66 self.device_id = self.local_rank
68 if not torch.cuda.is_available():
69 self.device_id = 0
70 self.device = torch.device("cpu")
71 logging.warning("CUDA is not available. Running on CPU.")
72 self.load_timer.end("torch_dist")
73 return # CPU mode, no parallelism needed
75 self.device = torch.device(f"cuda:{self.device_id}")
77 torch.cuda.set_device(self.local_rank)
79 if self.world_size <= 1:
80 self.load_timer.end("torch_dist")
81 return # Single GPU mode, no parallelism needed
83 if not dist.is_initialized():
84 dist.init_process_group(
85 backend="nccl",
86 init_method="env://",
87 rank=self.rank,
88 world_size=self.world_size,
89 timeout=timedelta(hours=24), # Prevent NCCL timeout
90 )
92 self.load_timer.end("torch_dist")
94 def load_model(self) -> None:
95 self.load_timer.start("realesrgan")
96 for scaling_factor in self.SCALING_FACTORS:
97 self.load_timer.start(f"realesrgan{scaling_factor}x")
98 self.models[scaling_factor] = RealESRGAN(self.device, scale=scaling_factor)
99 self.models[scaling_factor].load_weights(
100 f"ai-forever/Real-ESRGAN/RealESRGAN_x{scaling_factor}.pth",
101 download=False,
102 )
103 self.load_timer.end(f"realesrgan{scaling_factor}x")
104 logging.info("Loaded Real-ESRGAN.")
105 self.load_timer.end("realesrgan")
107 def init_model_parallelism(self) -> None:
108 if not dist.is_initialized() or self.world_size <= 1:
109 return
110 logging.info("Real-ESRGAN supports video parallelism.")
112 def model_compile(self) -> None:
113 if not self.torch_compile:
114 return
115 self.load_timer.start("compile")
116 logging.warning("torch.compile() not supported.")
117 # Error: accessing tensor output of CUDAGraphs that has been overwritten by a subsequent run.
118 # RealESRGAN/rrdbnet_arch.py", line 120, in forward
119 # out = self.conv_last(self.lrelu(self.conv_hr(feat))).
120 # To prevent overwriting, clone the tensor outside of torch.compile() or
121 # call torch.compiler.cudagraph_mark_step_begin() before each model invocation.
122 """
123 for scaling_factor in self.SCALING_FACTORS:
124 self.load_timer.start(f"compile{scaling_factor}x")
125 self.models[scaling_factor].model = torch.compile(
126 self.models[scaling_factor].model,
127 mode="reduce-overhead")
128 self.load_timer.end(f"compile{scaling_factor}x")
129 """
130 self.load_timer.end("compile")
132 def _assert_model_init(self) -> None:
133 super()._assert_model_init()
134 if not self.models:
135 raise ValueError("Real-ESRGAN not initialized.")
137 @inference_mode()
138 async def warmup(self) -> None:
139 logging.info(f"[{self.rank}] Warmup for Real-ESRGAN generation.")
140 video = [Image.new("RGB", (640, 480), color=(255, 255, 255))] * self.world_size * 2
141 await self.generate(
142 video=video,
143 width=1280,
144 height=960)
146 def _chunk_list_image(
147 self,
148 images: List[Image.Image]
149 ) -> List[Optional[Image.Image]]:
150 """
151 Chunk the list of images based on the world size.
152 This is used to distribute the workload across multiple ranks.
153 Each rank will process only its assigned images (not None).
154 """
155 if self.world_size == 1 or len(images) < 1:
156 return [img for img in images]
157 # Chunk one image per rank
158 ret: List[Optional[Image.Image]] = []
159 for it, image in enumerate(images):
160 if it % self.world_size == self.rank:
161 ret.append(image)
162 else:
163 ret.append(None)
164 return ret
166 def _gather_chunks(
167 self,
168 chunked_images: List[Optional[Image.Image]]
169 ) -> List[Image.Image]:
170 """
171 Gather the images from all ranks and put them into a sinle one.
172 Each rank will return its assigned images (not None).
173 This is used to collect the results after processing.
174 """
175 if self.world_size == 1:
176 return [img for img in chunked_images if img is not None]
178 gathered_lists = None
179 if self.rank == 0:
180 gathered_lists = [None] * self.world_size
181 dist.gather_object(chunked_images, gathered_lists, dst=0)
183 if self.rank != 0:
184 return []
186 assert gathered_lists is not None
187 ret = []
188 for position_images in zip(*gathered_lists):
189 for img in position_images:
190 if img is not None:
191 ret.append(img)
192 break
193 if len(ret) != len(chunked_images):
194 raise ValueError("Gathered images do not match the original chunked images length.")
195 return ret
197 @override
198 @inference_mode()
199 async def generate(
200 self,
201 job_id: Optional[str] = None,
202 image: Optional[Image.Image] = None,
203 video: Optional[List[Image.Image]] = None,
204 height: int = 960,
205 width: int = 1280,
206 batch_size: int = 4,
207 patches_size: int = 192,
208 padding: int = 24,
209 pad_size: int = 15,
210 video_fps: int = 30,
211 output_type: str = "pil", # "pil", "video_binary", "video_path"
212 ) -> Optional[Union[List[Image.Image], str, bytes]]:
213 gen_timer = self._new_gen_timer(job_id)
215 self._assert_model_init()
217 self.running = True # Mark running to avoid concurrent calls
219 try:
220 # Video upscaling
221 if video is not None:
222 if self.rank == 0:
223 video_len = get_video_size(video)
224 logging.info(
225 f"[{self.rank}] Upscaling video with {len(video)} frames and {bytes_to_human(video_len)}.")
226 ret: List[Optional[Image.Image]] = []
227 video_frames = video
228 chunked_video_frames = self._chunk_list_image(video_frames)
229 for it, frame in enumerate(chunked_video_frames):
230 if frame is None:
231 ret.append(None)
232 else:
233 gen_timer.start(f"frame_{it:03d}")
234 resized_frame = await asyncio.to_thread(
235 self.generate_image,
236 image=frame,
237 height=height,
238 width=width,
239 batch_size=batch_size,
240 patches_size=patches_size,
241 padding=padding,
242 pad_size=pad_size)
243 ret.append(resized_frame)
244 gen_timer.end(f"frame_{it:03d}")
245 gathered_frames = self._gather_chunks(ret)
246 if self.rank == 0:
247 logging.info(f"[{self.rank}] Generated {len(video)}->{len(gathered_frames)} upscaled video frames.")
248 else:
249 return [] # Skip non-rank 0 processes
250 return await self._output_video(
251 job_id,
252 gen_timer,
253 gathered_frames,
254 video_fps,
255 output_type)
257 # Image upscaling
258 if image is not None:
259 if self.rank != 0:
260 logging.debug(f"[{self.rank}] Skipping image upscaling, not rank 0.")
261 return []
262 out_image = await asyncio.to_thread(
263 self.generate_image,
264 image=image,
265 height=height,
266 width=width,
267 batch_size=batch_size,
268 patches_size=patches_size,
269 padding=padding,
270 pad_size=pad_size)
271 if self.rank == 0:
272 logging.info(f"[{self.rank}] Generated one upscaled image.")
273 return [out_image]
275 # Missing inputs
276 raise ValueError("Image or video required for Real-ESRGAN generation.")
277 finally:
278 self.running = False
279 gen_timer.end("total")
281 async def _output_video(
282 self,
283 job_id: Optional[str],
284 gen_timer: GenTimer,
285 video_frames: List[Image.Image],
286 video_fps: int = 30,
287 output_type: str = "pil", # "pil", "video_binary", "video_path"
288 ) -> Optional[Union[List[Image.Image], str, bytes]]:
289 gen_timer.start("output")
290 try:
291 if output_type == "pil":
292 return video_frames
294 if output_type in ("video_binary", "video_path"):
295 if not job_id:
296 video_path = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False).name
297 else:
298 video_path = f"/tmp/{job_id}.mp4"
299 video_path = await save_video_frames(
300 video_frames,
301 out_video_path=video_path,
302 fps=video_fps)
303 if output_type == "video_path":
304 return video_path
306 # video_binary
307 async with aiofiles.open(video_path, "rb") as f:
308 video_binary = await f.read()
309 return video_binary
311 logging.error(f"Unknown output type: {output_type}")
312 return None
313 finally:
314 gen_timer.end("output")
316 def get_model_scaling_factor(
317 self,
318 input_width: int,
319 input_height: int,
320 output_width: int,
321 output_height: int,
322 ) -> int:
323 if output_width < input_width or output_height < input_height:
324 raise ValueError(
325 f"Output size {output_width}x{output_height} must be larger than input {input_width}x{input_height}.")
326 max_scaling_factor = max(
327 output_width / input_width,
328 output_height / input_height)
329 if max_scaling_factor > self.SCALING_FACTORS[-1]:
330 raise ValueError(f"Scaling factor {max_scaling_factor}x > max {self.SCALING_FACTORS[-1]}x.")
331 # Get the next scaling factor in self.SCALING_FACTORS
332 model_scaling_factor = min(
333 [sf for sf in self.SCALING_FACTORS if sf >= max_scaling_factor],
334 default=self.SCALING_FACTORS[-1]
335 )
336 if model_scaling_factor not in self.models:
337 raise ValueError(f"Model for {model_scaling_factor}x not loaded.")
338 return model_scaling_factor
340 @torch.inference_mode()
341 def generate_image(
342 self,
343 image: Image.Image,
344 height: int = 768,
345 width: int = 1024,
346 batch_size: int = 4,
347 patches_size: int = 192,
348 padding: int = 24,
349 pad_size: int = 15,
350 ) -> Image.Image:
351 input_width, input_height = image.size
352 model_scaling_factor = self.get_model_scaling_factor(
353 input_width, input_height,
354 width, height)
355 logging.debug(
356 f"[{self.rank}] {input_width}x{input_height}->{width}x{height} Using scaling {model_scaling_factor}x.")
358 model = self.models.get(model_scaling_factor, None)
359 if model is None:
360 raise ValueError(f"Model for {model_scaling_factor}x not loaded.")
361 output_image = model.predict(
362 image,
363 batch_size=batch_size,
364 patches_size=patches_size,
365 padding=padding,
366 pad_size=pad_size)
368 logging.debug(f"[{self.rank}] Generated image with size {output_image.size}")
369 if output_image.size[0] != width or output_image.size[1] != height:
370 logging.debug(f"[{self.rank}] Downscaling now to {width}x{height}.")
371 output_image = output_image.resize((width, height), Image.Resampling.LANCZOS)
373 return output_image
375 def get_health(self) -> Dict[str, Any]:
376 ret = super().get_health()
377 ret.update({
378 "gpu": self.GPU,
379 "rank": self.rank,
380 "world_size": self.world_size,
381 "torch_compile": self.torch_compile,
382 })
383 return ret
385 async def get_rest_args(self, data_json: Dict[str, str]) -> Dict[str, Any]:
386 if data_json is None or not isinstance(data_json, dict):
387 raise ValueError("Missing JSON body")
389 job_id = data_json.get("job_id", None)
391 img_base64 = data_json.get("img", None)
392 img = None
393 if img_base64 is not None:
394 img = base64_to_img(img_base64)
396 video_base64 = data_json.get("video", None)
397 video_frames = None
398 video_fps: float = -1.0
399 if video_base64 is not None:
400 video_frames = base64_to_video_frames(video_base64)
401 video_binary = base64_to_binary(video_base64)
402 video_fps = get_video_fps(video_binary)
404 rest_args = {
405 "task": self.model_name,
406 "args": {
407 "job_id": job_id,
408 "image": img,
409 "video": video_frames,
410 "width": int(data_json.get("width", 640)),
411 "height": int(data_json.get("height", 480)),
412 "batch_size": int(data_json.get("batch_size", 4)),
413 "patches_size": int(data_json.get("patches_size", 192)),
414 "padding": int(data_json.get("padding", 24)),
415 "pad_size": int(data_json.get("pad_size", 15)),
416 "output_type": data_json.get("output_type", "pil"),
417 }
418 }
419 if video_fps > 0:
420 rest_args["args"]["video_fps"] = video_fps
421 return rest_args