feat: add MiniMax as alternative LLM provider

README.md

@@ -84,6 +84,13 @@ Please refer to the [Installation Guide](installation/Installation.md) for detai
 
 **Prerequest:** Before running 4KAgent, please fill in the API key in [config file](config.yml)
 
+4KAgent supports multiple LLM providers for the reasoning agent. Use `--llm_provider` to select:
+- `gpt4` (default) — OpenAI GPT-4
+- `azure` — Azure OpenAI
+- `minimax` — [MiniMax](https://www.minimaxi.com) (MiniMax-M2.7 with 204K context)
+
+Set the corresponding API key in `config.yml` or via environment variable (`MINIMAX_API_KEY` for MiniMax).
+
 The inference of 4KAgent relies on profile, we present examples here:
 
 **Profiles use 'llama_vision' as the VLM in perception agent:**
@@ -123,10 +130,20 @@ CUDA_VISIBLE_DEVICES=1 python infer_4kagent.py \
   --tool_run_gpu_id 2
 ```
 
-We recommend the `FastGen4K_P` profile, which infers faster and has good perceptual quality. 
+We recommend the `FastGen4K_P` profile, which infers faster and has good perceptual quality.
 
 `tool_run_gpu_id` is used to specify the GPU to execute tools (restoration methods). For GPUs with larger VRAM, `tool_run_gpu_id` can be set as the same as `CUDA_VISIBLE_DEVICES`.
 
+**Using MiniMax as the LLM provider:**
+```bash
+CUDA_VISIBLE_DEVICES=1 python infer_4kagent.py \
+  --input_dir ./assets/profile_test_example/classicsr \
+  --output_dir ./outputs/4KAgent_test/classicsr \
+  --profile_name ExpSR_s4_F \
+  --tool_run_gpu_id 2 \
+  --llm_provider minimax
+```
+
 **Old Photo 4K SR**
 ```bash
 # Set up depictqa in portal A:

config.yml

@@ -16,4 +16,10 @@ AZUREGPT:
   MAX_TOKENS: 3000
   TEMPERATURE: 0.0
   ENDPOINT: ""
-  API_VERSION: ""
+  API_VERSION: ""
+
+MINIMAX:
+  API_KEY: ""
+  MODEL: "MiniMax-M2.7"
+  MAX_TOKENS: 3000
+  TEMPERATURE: 0.7

infer_4kagent.py

@@ -14,6 +14,9 @@ def parse_args():
     parser.add_argument("--output_dir", type=str, default="./outputs/LQ_results", help="Path to the output directory")
     parser.add_argument("--profile_name", type=str, default="", help="Profile Name for the 4KAgent")
     parser.add_argument("--tool_run_gpu_id", type=int, default=0, help="GPU ID to run tools the toolbox")
+    parser.add_argument("--llm_provider", type=str, default="gpt4",
+                        choices=["gpt4", "azure", "minimax"],
+                        help="LLM provider for the reasoning agent (default: gpt4)")
     return parser.parse_args()
 
 
@@ -23,6 +26,7 @@ def main():
     output_dir = Path(args.output_dir).resolve()
     profile_name = args.profile_name
     tool_run_gpu_id = args.tool_run_gpu_id
+    llm_provider = args.llm_provider
 
     output_dir.mkdir(parents=True, exist_ok=True)
 
@@ -53,7 +57,8 @@ def main():
             with_reflection=True,
             silent=False,
             tool_run_gpu_id=tool_run_gpu_id,
-            profile_name=profile_name
+            profile_name=profile_name,
+            llm_provider=llm_provider
         )
 
         agent.run()

llm/__init__.py

@@ -1,8 +1,9 @@
 from .gpt4 import GPT4
 from .depictqa import DepictQA
 from .azuregpt import AzureGPT
+from .minimax import MiniMaxLLM
 from .qwen_vl import PerceptionVLMAgent
 from .llama_vision import LlamaVisionAgent
 
 
-__all__ = ["GPT4", "AzureGPT", "DepictQA", "PerceptionVLMAgent", "LlamaVisionAgent"]
+__all__ = ["GPT4", "AzureGPT", "MiniMaxLLM", "DepictQA", "PerceptionVLMAgent", "LlamaVisionAgent"]

llm/minimax.py

@@ -0,0 +1,207 @@
+from pathlib import Path
+import requests
+import logging
+from typing import Callable, Optional
+from time import sleep
+import random
+import re
+
+from .base_llm import BaseLLM
+from utils.misc import encode_img
+
+
+class MiniMaxLLM(BaseLLM):
+    """MiniMax LLM provider using the OpenAI-compatible API.
+
+    Parameters when called: img_path_lst, prompt, format_check.
+    """
+
+    MINIMAX_API_URL = "https://api.minimax.io/v1/chat/completions"
+
+    def __init__(self,
+                 config_path: Path = Path("config.yml"),
+                 log_path: Optional[Path | str] = None,
+                 logger: Optional[logging.Logger] = None,
+                 silent: bool = False,
+                 system_message: Optional[str] = None,
+                 model: Optional[str] = None
+                 ):
+        super().__init__(
+            config_path=config_path,
+            log_path=log_path,
+            logger=logger,
+            silent=silent
+        )
+
+        self.api_key = self.cfg["MINIMAX"]["API_KEY"]
+        if model is None:
+            self.model = self.cfg["MINIMAX"]["MODEL"]
+        else:
+            self.model = model
+        self.max_tokens = self.cfg["MINIMAX"]["MAX_TOKENS"]
+        # MiniMax requires temperature in (0.0, 1.0]
+        raw_temp = self.cfg["MINIMAX"]["TEMPERATURE"]
+        self.temperature = max(0.01, min(float(raw_temp), 1.0))
+
+        self.prompt_tokens = 0
+        self.completion_tokens = 0
+
+        self.system_message = system_message
+        if self.system_message is not None:
+            self._log("_Note: These user-assistant interactions are independent "
+                      "and the system message is always attached in each turn for MiniMax._")
+            self._log("**System message for MiniMax**")
+            self._log(self.system_message)
+
+    def query(self,
+              img_path_lst: Optional[list[Path]] = None,
+              prompt: str = "",
+              format_check: Optional[Callable[[object], None]] = None,
+              ) -> tuple[str, str]:
+        headers, payload = self._prepare_for_request(
+            prompt, img_path_lst)
+        while True:
+            response = self._send_request(headers, payload)
+
+            usage = response.json()["usage"]
+            self.prompt_tokens += usage["prompt_tokens"]
+            self.completion_tokens += usage["completion_tokens"]
+
+            rsp_text: str = response.json()['choices'][0]['message']['content']
+            # Strip MiniMax thinking tags if present
+            rsp_text = re.sub(r'<think>.*?</think>\s*', '', rsp_text, flags=re.DOTALL).strip()
+            if format_check is not None:
+                valid, rsp_text = self._check_syntax(rsp_text, format_check)
+                if not valid:
+                    continue
+            return prompt, rsp_text
+
+    def _prepare_for_request(self, prompt: str,
+                             img_path_lst: Optional[list[Path]] = None
+                             ) -> tuple[dict, dict]:
+        content = [{
+            "type": "text",
+            "text": prompt
+        }]
+        if img_path_lst is not None:
+            for img_path in img_path_lst:
+                img_base64 = encode_img(img_path)
+                content.append({
+                    "type": "image_url",
+                    "image_url": {
+                        "url": img_base64,
+                        "detail": "auto"
+                    }
+                })
+
+        messages = []
+        if self.system_message is not None:
+            messages.append({
+                "role": "system",
+                "content": self.system_message
+            })
+        messages.append({
+            "role": "user",
+            "content": content
+        })
+
+        headers = {
+            "Content-Type": "application/json",
+            "Authorization": f"Bearer {self.api_key}"
+        }
+        payload = {
+            "model": self.model,
+            "messages": messages,
+            "max_tokens": self.max_tokens,
+            "temperature": self.temperature
+        }
+
+        return headers, payload
+
+    def _send_request(self, headers: dict, payload: dict,
+                      max_retries: int = 5,
+                      initial_delay: int = 3,
+                      exp_base: int = 2,
+                      jitter: bool = True) -> requests.Response:
+        """Sends a request to the MiniMax API and handles errors with exponential backoff."""
+
+        n_retries = 0
+        backoff_delay = initial_delay
+        while True:
+            try:
+                response = requests.post(self.MINIMAX_API_URL,
+                                         headers=headers, json=payload)
+                is_valid, recommended_delay = self._check_response(response)
+                if is_valid:
+                    return response
+            except Exception as e:
+                self._log("An error occurred when sending a request: "
+                          f"{type(e).__name__}: {e}",
+                          level='warning')
+                recommended_delay = None
+
+            n_retries += 1
+            if n_retries > max_retries:
+                raise RuntimeError(
+                    "Too many errors occurred when querying MiniMax LLM.")
+            if recommended_delay is not None:
+                delay = recommended_delay
+            else:
+                backoff_delay *= exp_base * (1 + jitter*random.random())
+                delay = backoff_delay
+            self._log(
+                f"Retrying in {delay:.3f} seconds...", level='warning')
+            sleep(delay)
+
+    def _check_response(self, response: requests.Response) -> tuple[bool, Optional[float]]:
+        """Checks if the response is valid. If error occurs, gets the recommended delay if any."""
+
+        if "error" in response.json():
+            err_msg: str = response.json()['error']['message']
+            self._log(f"An error occurred when querying MiniMax LLM: {err_msg}",
+                      level='warning')
+
+            recommended_delay = None
+            if response.json()['error'].get('code') == 'rate_limit_exceeded':
+                match = re.search(
+                    R"(?<=Please try again in )(\d+m)?\d+\.?\d*(?=s)", err_msg)
+                if match is not None:
+                    t = match.group().split('m')
+                    m = t[0] if len(t) > 1 else 0
+                    s = t[-1]
+                    recommended_delay = 60*int(m) + float(s)
+
+            return False, recommended_delay
+
+        if (finish_reason := response.json()['choices'][0]['finish_reason']) != 'stop':
+            self._log(f"finish_reason is {finish_reason}", level='warning')
+
+        return True, None
+
+    def _check_syntax(self, rsp_text: str, format_check: Callable[[object], None]
+                      ) -> tuple[bool, str]:
+        """Checks whether the response is a valid Python object and follows the specified format."""
+        try:
+            obj = eval(rsp_text)
+        except:
+            inner_rsp_text = rsp_text.strip("```").lstrip("json").strip()
+            try:
+                obj = eval(inner_rsp_text)
+                rsp_text = inner_rsp_text
+            except:
+                self._log("Failed to parse the response:", level='warning')
+                self._log(rsp_text, level='warning')
+                return False, ""
+        try:
+            format_check(obj)
+        except AssertionError as e:
+            self._log(f"Failed to pass the format check: {e}", level='warning')
+            self._log(f"Response: {obj}", level='warning')
+            return False, ""
+        return True, rsp_text
+
+    def _post_process(self):
+        """Logs the token usage."""
+        self._log("Token usage so far: "
+                  f"{self.prompt_tokens} prompt tokens, "
+                  f"{self.completion_tokens} completion tokens")

pipeline/the4kagent_pipeline.py

@@ -14,7 +14,7 @@
 
 from . import prompts
 from executor import executor, Tool
-from llm import GPT4, AzureGPT, DepictQA, PerceptionVLMAgent, LlamaVisionAgent
+from llm import GPT4, AzureGPT, MiniMaxLLM, DepictQA, PerceptionVLMAgent, LlamaVisionAgent
 
 from utils.img_tree import ImgTree
 from utils.logger import get_logger
@@ -62,6 +62,7 @@ def __init__(
         tool_run_gpu_id: Optional[int] = None,
         perception_agent_run_gpu_id: Optional[int] = None,
         profile_name: Optional[str] = None,
+        llm_provider: Optional[str] = None,
     ) -> None:
         # paths
         self._prepare_dir(input_path, output_dir)
@@ -73,7 +74,8 @@ def __init__(
             with_reflection,
             # with_rollback,
             tool_run_gpu_id,
-            profile_name
+            profile_name,
+            llm_provider
         )
         # components
         self._create_components(llm_config_path, schedule_experience_path, silent)
@@ -124,6 +126,7 @@ def _config(
         # with_rollback: bool,
         tool_run_gpu_id: Optional[int],
         profile_name: Optional[str] = None,
+        llm_provider: Optional[str] = None,
     ) -> None:
         # extract profile
         self.profile_name = profile_name or "FastGen4K_P"
@@ -167,7 +170,9 @@ def _config(
 
         self.fast_4k = self.profile.get("Fast4K", False)
         self.fast4k_side_thres = self.profile.get("Fast4kSideThres", 1024)
-
+
+        self.llm_provider = llm_provider or "gpt4"
+
         self.project_root = Path(__file__).resolve().parent.parent # 4kagent dir path
 
 
@@ -212,12 +217,27 @@ def _create_components(
             )
 
         # language models
-        self.gpt4 = GPT4(
-            config_path=llm_config_path,
-            logger=self.qa_logger,
-            silent=silent,
-            system_message=prompts.system_message,
-        )
+        if self.llm_provider == "minimax":
+            self.gpt4 = MiniMaxLLM(
+                config_path=llm_config_path,
+                logger=self.qa_logger,
+                silent=silent,
+                system_message=prompts.system_message,
+            )
+        elif self.llm_provider == "azure":
+            self.gpt4 = AzureGPT(
+                config_path=llm_config_path,
+                logger=self.qa_logger,
+                silent=silent,
+                system_message=prompts.system_message,
+            )
+        else:
+            self.gpt4 = GPT4(
+                config_path=llm_config_path,
+                logger=self.qa_logger,
+                silent=silent,
+                system_message=prompts.system_message,
+            )
         # self.gpt4 = AzureGPT(
         #     config_path=llm_config_path,
         #     logger=self.qa_logger,