Coverage for wrapper/4kagent/wrapper_4kagent.py: 94%
124 statements
« prev ^ index » next coverage.py v7.15.4, created at 2026-08-09 04:47 +0000
1"""
2Wrapper for 4KAgent - agentic any-image-to-4K super-resolution.
3https://github.com/taco-group/4KAgent
4"""
6import asyncio
7import logging
8import os
9import shutil
10import sys
11import tempfile
13from pathlib import Path
14from typing import override
15from typing import Dict
16from typing import Optional
17from typing import Union
18from typing import Any
20import torch
21import yaml
23from PIL import Image
25from wrapper_model import ModelGeneration
27from image_utils import base64_to_img
30class Upscale4KAgent(ModelGeneration):
31 """
32 Wrapper for 4KAgent image super-resolution.
34 4KAgent is an agentic framework that upscales any image to 4K resolution.
35 The 4KAgent code is imported directly from the cloned repository at
36 FOURK_AGENT_DIR, which is added to sys.path during load_model().
38 Required environment variables (at least one LLM key is needed):
39 LLAMA_API_KEY – Meta Llama API key (used by llama_vision profiles)
40 OPENAI_API_KEY – OpenAI API key (used by GPT-based profiles)
41 AZURE_OPENAI_API_KEY / AZURE_OPENAI_ENDPOINT / AZURE_OPENAI_MODEL /
42 AZURE_OPENAI_API_VERSION – Azure OpenAI credentials
43 """
45 FOURK_AGENT_DIR: str = "/4kagent/4KAgent"
46 DEFAULT_PROFILE: str = "ExpSR_s4_P"
48 def __init__(self, model_name: str = "4kagent") -> None:
49 super().__init__(model_name)
50 self.fourk_agent_dir: Optional[str] = None
52 def __del__(self) -> None:
53 super().__del__()
55 def init_parallelism(self) -> None:
56 self.load_timer.start("torch_dist")
58 self.rank = int(os.getenv("RANK", 0))
59 self.local_rank = int(os.getenv("LOCAL_RANK", 0))
60 # 4KAgent manages its own GPU assignment internally.
61 self.world_size = 1
63 self.device_id = self.local_rank
64 if torch.cuda.is_available():
65 self.device = torch.device(f"cuda:{self.device_id}")
66 else:
67 self.device = torch.device("cpu")
69 self.load_timer.end("torch_dist")
71 def load_model(self) -> None:
72 if not os.path.isdir(self.FOURK_AGENT_DIR):
73 raise RuntimeError(
74 f"4KAgent not found at {self.FOURK_AGENT_DIR}. "
75 "Ensure the Docker image has cloned the repository."
76 )
77 self.fourk_agent_dir = self.FOURK_AGENT_DIR
78 # Add the 4KAgent repo to sys.path so its packages can be imported
79 # directly. This must happen before the first call to _run_4kagent().
80 if self.fourk_agent_dir not in sys.path:
81 sys.path.insert(0, self.fourk_agent_dir)
82 self._write_config()
84 # Warn if no LLM API key is configured. Most llama_vision-based profiles
85 # require LLAMA_API_KEY; GPT-based profiles need OPENAI_API_KEY or Azure
86 # credentials. A missing key will only cause a runtime failure when the
87 # VLM is invoked, not during loading.
88 llm_keys = [
89 os.getenv("LLAMA_API_KEY", ""),
90 os.getenv("OPENAI_API_KEY", ""),
91 os.getenv("AZURE_OPENAI_API_KEY", ""),
92 ]
93 if not any(llm_keys):
94 logging.warning(
95 "No LLM API key found (LLAMA_API_KEY / OPENAI_API_KEY / "
96 "AZURE_OPENAI_API_KEY). Generation will fail unless the chosen "
97 "profile does not require a VLM."
98 )
100 logging.info(f"Loaded 4KAgent from {self.fourk_agent_dir}.")
102 def _write_config(self) -> None:
103 """Write config.yml populated with API keys from environment variables."""
104 assert self.fourk_agent_dir is not None
105 config: Dict[str, Any] = {
106 "GPT": {
107 "API_KEY": os.getenv("OPENAI_API_KEY", ""),
108 "MODEL": os.getenv("OPENAI_MODEL", "gpt-4-turbo"),
109 "MAX_TOKENS": 3000,
110 "TEMPERATURE": 0.0,
111 },
112 "LLAMA": {
113 "API_KEY": os.getenv("LLAMA_API_KEY", ""),
114 "MODEL": os.getenv("LLAMA_MODEL", "llama3.1-405b"),
115 "MAX_TOKENS": 3000,
116 "TEMPERATURE": 0.0,
117 },
118 "AZUREGPT": {
119 "API_KEY": os.getenv("AZURE_OPENAI_API_KEY", ""),
120 "MODEL": os.getenv("AZURE_OPENAI_MODEL", ""),
121 "MAX_TOKENS": 3000,
122 "TEMPERATURE": 0.0,
123 "ENDPOINT": os.getenv("AZURE_OPENAI_ENDPOINT", ""),
124 "API_VERSION": os.getenv("AZURE_OPENAI_API_VERSION", ""),
125 },
126 }
127 config_path = os.path.join(self.fourk_agent_dir, "config.yml")
128 with open(config_path, "w") as fh:
129 yaml.dump(config, fh, default_flow_style=False)
130 logging.info(f"Wrote 4KAgent config to {config_path}.")
132 def init_model_parallelism(self) -> None:
133 if self.world_size > 1:
134 logging.warning("4KAgent does not support model parallelism.")
136 def model_compile(self) -> None:
137 # torch.compile is not applicable for direct-import execution.
138 pass
140 def _assert_model_init(self) -> None:
141 super()._assert_model_init()
142 if self.fourk_agent_dir is None:
143 raise ValueError("4KAgent directory not initialised.")
145 @torch.inference_mode()
146 async def warmup(self) -> None:
147 logging.info(f"[{self.rank}] Warmup for 4KAgent.")
148 warmup_image = Image.new("RGB", (256, 192), color=(128, 128, 128))
149 await self.generate(image=warmup_image)
151 @override
152 @torch.inference_mode()
153 async def generate(
154 self,
155 image: Optional[Image.Image] = None,
156 profile_name: str = DEFAULT_PROFILE,
157 tool_run_gpu_id: int = 0,
158 job_id: Optional[str] = None,
159 ) -> Image.Image:
160 """
161 Run 4KAgent super-resolution on *image* and return the upscaled result.
163 Args:
164 image: Input PIL image to upscale.
165 profile_name: 4KAgent profile (e.g. "ExpSR_s4_P", "FastGen4K_P").
166 Profiles using ``llama_vision`` work without a separate
167 DepictQA server; profiles using ``depictqa`` require one.
168 tool_run_gpu_id: GPU index for the restoration tool subprocesses.
169 job_id: Optional job identifier for timing telemetry.
171 Returns:
172 Upscaled PIL image.
173 """
174 gen_timer = self._new_gen_timer(job_id)
175 self._assert_model_init()
177 if image is None:
178 raise ValueError("An input image is required for 4KAgent generation.")
180 self.running = True
181 input_path: Optional[str] = None
182 output_dir: Optional[str] = None
183 try:
184 gen_timer.start("save_input")
185 if job_id:
186 input_path = f"/tmp/{job_id}.png"
187 else:
188 with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmp:
189 input_path = tmp.name
190 await asyncio.to_thread(image.save, input_path)
191 gen_timer.end("save_input")
193 output_dir = tempfile.mkdtemp(prefix="4kagent_out_")
195 gen_timer.start("inference")
196 await asyncio.to_thread(
197 self._run_4kagent,
198 Path(input_path),
199 Path(output_dir),
200 profile_name,
201 tool_run_gpu_id,
202 )
203 gen_timer.end("inference")
205 gen_timer.start("load_result")
206 result = await asyncio.to_thread(self._load_result, output_dir)
207 gen_timer.end("load_result")
209 return result
210 finally:
211 self.running = False
212 gen_timer.end("total")
213 if input_path and os.path.exists(input_path):
214 os.unlink(input_path)
215 if output_dir and os.path.exists(output_dir):
216 shutil.rmtree(output_dir, ignore_errors=True)
218 def _run_4kagent(
219 self,
220 input_path: Path,
221 output_dir: Path,
222 profile_name: str,
223 tool_run_gpu_id: int,
224 ) -> None:
225 """Import and run The4KAgent pipeline directly.
227 Called via asyncio.to_thread; 4KAgent handles one image at a time so
228 concurrent requests are serialised by the caller's semaphore.
229 The lazy import (after sys.path is set in load_model) is intentional:
230 Python's import system caches modules in sys.modules so subsequent
231 calls incur only a dict lookup.
232 """
233 assert self.fourk_agent_dir is not None
234 from pipeline.the4kagent_pipeline import The4KAgent # noqa: PLC0415
235 logging.info(
236 f"[{self.rank}] Running 4KAgent: profile={profile_name}, tool_gpu={tool_run_gpu_id}.",
237 )
238 agent = The4KAgent(
239 input_path=input_path,
240 output_dir=output_dir,
241 llm_config_path=Path(self.fourk_agent_dir) / "config.yml",
242 with_retrieval=True,
243 with_reflection=True,
244 silent=False,
245 tool_run_gpu_id=tool_run_gpu_id,
246 profile_name=profile_name,
247 )
248 agent.run()
250 def _load_result(self, output_dir: str) -> Image.Image:
251 """
252 Locate and load the result image produced by 4KAgent.
254 4KAgent writes to ``<output_dir>/<image_stem>/<step>/result.png``.
255 """
256 result_candidates = sorted(Path(output_dir).glob("*/*/result.png"))
257 if not result_candidates:
258 raise ValueError(
259 f"No result.png found under {output_dir}. "
260 "4KAgent may have failed or the output structure changed."
261 )
262 result_path = result_candidates[-1]
263 logging.info(f"[{self.rank}] Loading result from {result_path}.")
264 # .copy() ensures the image is fully loaded before the temp dir is removed.
265 return Image.open(result_path).copy()
267 def get_health(self) -> Dict[str, Any]:
268 ret = super().get_health()
269 ret.update({
270 "rank": self.rank,
271 "world_size": self.world_size,
272 "fourk_agent_dir": self.fourk_agent_dir,
273 })
274 return ret
276 async def get_rest_args(
277 self,
278 data_json: Dict[str, Union[str, int, float]],
279 ) -> Dict[str, Any]:
280 if data_json is None or not isinstance(data_json, dict):
281 raise ValueError("Missing JSON body")
283 img_base64 = data_json.get("img", None)
284 if img_base64 is None:
285 raise ValueError("Missing 'img' parameter")
287 image = base64_to_img(str(img_base64))
289 return {
290 "task": self.model_name,
291 "args": {
292 "job_id": data_json.get("job_id", None),
293 "image": image,
294 "profile_name": str(data_json.get("profile_name", self.DEFAULT_PROFILE)),
295 "tool_run_gpu_id": int(data_json.get("tool_run_gpu_id", 0)),
296 },
297 }