feat: add MiniMax provider support for prompt refinement (M3 default)

opensora/utils/prompt_refine.py

@@ -1,9 +1,46 @@
 import base64
 import os
+import re
 from mimetypes import guess_type
 
 from openai import OpenAI
 
+MINIMAX_BASE_URL = os.environ.get("MINIMAX_BASE_URL", "https://api.minimax.io/v1")
+
+MINIMAX_MODELS = [
+    "MiniMax-M3",
+    "MiniMax-M2.7",
+    "MiniMax-M2.7-highspeed",
+]
+
+
+def _get_client(model: str) -> OpenAI:
+    """Return an OpenAI-compatible client configured for the given model."""
+    if model.startswith("MiniMax"):
+        api_key = os.environ.get("MINIMAX_API_KEY")
+        if not api_key:
+            raise ValueError(
+                "MINIMAX_API_KEY environment variable is not set. "
+                "Please set it to use MiniMax models."
+            )
+        return OpenAI(api_key=api_key, base_url=MINIMAX_BASE_URL)
+    return OpenAI(api_key=os.environ.get("OPENAI_API_KEY"))
+
+
+def has_minimax_key() -> bool:
+    """Check whether the MiniMax API key is configured."""
+    return bool(os.environ.get("MINIMAX_API_KEY"))
+
+
+def _strip_think_tags(text: str) -> str:
+    """Remove chain-of-thought <think>...</think> blocks from model output."""
+    return re.sub(r"<think>.*?</think>", "", text, flags=re.DOTALL).strip()
+
+
+def _extra_tokens(model: str) -> int:
+    """Extra token budget for MiniMax chain-of-thought <think> blocks."""
+    return 500 if model.startswith("MiniMax") else 0
+
 sys_prompt_t2v = """You are part of a team of bots that creates videos. The workflow is that you first create a caption of the video, and then the assistant bot will generate the video based on the caption. You work with an assistant bot that will draw anything you say.
 
 For example, outputting "a beautiful morning in the woods with the sun peaking through the trees" will trigger your partner bot to output an video of a forest morning, as described. You will be prompted by people looking to create detailed, amazing videos. The way to accomplish this is to take their short prompts and make them extremely detailed and descriptive.
@@ -72,12 +109,31 @@ def image_to_url(image_path):
     return f"data:{mime_type};base64,{base64_encoded_data}"
 
 
-def refine_prompt(prompt: str, retry_times: int = 3, type: str = "t2v", image_path: str = None):
+def refine_prompt(
+    prompt: str,
+    retry_times: int = 3,
+    type: str = "t2v",
+    image_path: str = None,
+    model: str = None,
+):
     """
-    Refine a prompt to a format that can be used by the model for inference
+    Refine a prompt to a format that can be used by the model for inference.
+
+    Args:
+        prompt: The input prompt text.
+        retry_times: Number of retry attempts on failure.
+        type: Prompt type - "t2v" (text-to-video), "t2i" (text-to-image),
+              "i2v" (image-to-video), or "motion_score".
+        image_path: Path to reference image (required when type="i2v").
+        model: LLM model to use for refinement. Defaults to the PROMPT_MODEL
+               environment variable, or "gpt-4o" if unset. Use MiniMax models
+               (e.g. "MiniMax-M3") together with MINIMAX_API_KEY.
     """
+    if model is None:
+        model = os.environ.get("PROMPT_MODEL", "gpt-4o")
 
-    client = OpenAI(api_key=os.environ.get("OPENAI_API_KEY"))
+    client = _get_client(model)
+    extra = _extra_tokens(model)
 
     text = prompt.strip()
     response = None
@@ -115,11 +171,11 @@ def refine_prompt(prompt: str, retry_times: int = 3, type: str = "t2v", image_pa
                         "content": f'Create an imaginative video descriptive caption or modify an earlier caption in ENGLISH for the user input: " {text} "',
                     },
                 ],
-                model="gpt-4o",  # glm-4-plus and gpt-4o have be tested
+                model=model,  # glm-4-plus, gpt-4o, and MiniMax-M3 have been tested
                 temperature=0.01,
                 top_p=0.7,
                 stream=False,
-                max_tokens=250,
+                max_tokens=250 + extra,
             )
         elif type == "t2i":
             response = client.chat.completions.create(
@@ -146,15 +202,15 @@ def refine_prompt(prompt: str, retry_times: int = 3, type: str = "t2v", image_pa
                         "content": f'Create an imaginative image descriptive caption or modify an earlier caption in ENGLISH for the user input: " {text} "',
                     },
                 ],
-                model="gpt-4o",  # glm-4-plus and gpt-4o have be tested
+                model=model,  # glm-4-plus, gpt-4o, and MiniMax-M3 have been tested
                 temperature=0.01,
                 top_p=0.7,
                 stream=False,
-                max_tokens=250,
+                max_tokens=250 + extra,
             )
         elif type == "i2v":
             response = client.chat.completions.create(
-                model="gpt-4o",
+                model=model,
                 messages=[
                     {"role": "system", "content": f"{sys_prompt_i2v}"},
                     {
@@ -200,7 +256,7 @@ def refine_prompt(prompt: str, retry_times: int = 3, type: str = "t2v", image_pa
                 temperature=0.01,
                 top_p=0.7,
                 stream=False,
-                max_tokens=250,
+                max_tokens=250 + extra,
             )
         elif type == "motion_score":
             response = client.chat.completions.create(
@@ -211,24 +267,43 @@ def refine_prompt(prompt: str, retry_times: int = 3, type: str = "t2v", image_pa
                         "content": f"{text}",
                     },
                 ],
-                model="gpt-4o",  # glm-4-plus and gpt-4o have be tested
+                model=model,  # glm-4-plus, gpt-4o, and MiniMax-M3 have been tested
                 temperature=0.01,
                 top_p=0.7,
                 stream=False,
-                max_tokens=100,
+                max_tokens=100 + extra,
             )
         if response is None:
             continue
         if response.choices:
-            return response.choices[0].message.content
+            return _strip_think_tags(response.choices[0].message.content)
     return prompt
 
 
-def refine_prompts(prompts: list[str], retry_times: int = 3, type: str = "t2v", image_paths: list[str] = None):
+def refine_prompts(
+    prompts: list[str],
+    retry_times: int = 3,
+    type: str = "t2v",
+    image_paths: list[str] = None,
+    model: str = None,
+):
     if image_paths is None:
         image_paths = [None] * len(prompts)
     refined_prompts = []
     for prompt, image_path in zip(prompts, image_paths):
-        refined_prompt = refine_prompt(prompt, retry_times=retry_times, type=type, image_path=image_path)
+        refined_prompt = refine_prompt(
+            prompt, retry_times=retry_times, type=type, image_path=image_path, model=model
+        )
         refined_prompts.append(refined_prompt)
     return refined_prompts
+
+
+def refine_prompts_by_minimax(
+    prompts: list[str],
+    retry_times: int = 3,
+    type: str = "t2v",
+    image_paths: list[str] = None,
+    model: str = "MiniMax-M3",
+):
+    """Refine prompts using the MiniMax API (requires MINIMAX_API_KEY)."""
+    return refine_prompts(prompts, retry_times=retry_times, type=type, image_paths=image_paths, model=model)

tests/test_minimax_prompt_refine.py

@@ -0,0 +1,190 @@
+"""Unit tests for MiniMax provider support in prompt_refine.py."""
+import os
+import sys
+import unittest
+from unittest.mock import MagicMock, patch
+
+sys.path.insert(0, os.path.join(os.path.dirname(__file__), ".."))
+
+from opensora.utils.prompt_refine import (
+    MINIMAX_BASE_URL,
+    MINIMAX_MODELS,
+    _get_client,
+    _strip_think_tags,
+    has_minimax_key,
+    refine_prompts_by_minimax,
+)
+
+
+class TestStripThinkTags(unittest.TestCase):
+    def test_strips_think_block(self):
+        text = "<think>Some reasoning here.</think>\nActual answer."
+        self.assertEqual(_strip_think_tags(text), "Actual answer.")
+
+    def test_strips_multiline_think_block(self):
+        text = "<think>\nline1\nline2\n</think>\nAnswer."
+        self.assertEqual(_strip_think_tags(text), "Answer.")
+
+    def test_no_think_block_unchanged(self):
+        text = "Just a plain answer."
+        self.assertEqual(_strip_think_tags(text), "Just a plain answer.")
+
+    def test_empty_string(self):
+        self.assertEqual(_strip_think_tags(""), "")
+
+
+class TestMiniMaxConfig(unittest.TestCase):
+    def test_minimax_models_list(self):
+        self.assertIn("MiniMax-M3", MINIMAX_MODELS)
+        self.assertIn("MiniMax-M2.7", MINIMAX_MODELS)
+        self.assertIn("MiniMax-M2.7-highspeed", MINIMAX_MODELS)
+        self.assertEqual(MINIMAX_MODELS[0], "MiniMax-M3")
+        self.assertEqual(len(MINIMAX_MODELS), 3)
+
+    def test_minimax_base_url_default(self):
+        self.assertTrue(MINIMAX_BASE_URL.startswith("https://api.minimax.io"))
+
+    def test_has_minimax_key_false_when_unset(self):
+        with patch.dict(os.environ, {}, clear=True):
+            os.environ.pop("MINIMAX_API_KEY", None)
+            self.assertFalse(has_minimax_key())
+
+    def test_has_minimax_key_true_when_set(self):
+        with patch.dict(os.environ, {"MINIMAX_API_KEY": "sk-test-key"}):
+            self.assertTrue(has_minimax_key())
+
+
+class TestGetClient(unittest.TestCase):
+    @patch("opensora.utils.prompt_refine.OpenAI")
+    def test_minimax_model_uses_minimax_base_url(self, mock_openai):
+        with patch.dict(os.environ, {"MINIMAX_API_KEY": "sk-test-key"}):
+            _get_client("MiniMax-M2.7")
+            mock_openai.assert_called_once_with(
+                api_key="sk-test-key",
+                base_url=MINIMAX_BASE_URL,
+            )
+
+    @patch("opensora.utils.prompt_refine.OpenAI")
+    def test_minimax_highspeed_uses_minimax_base_url(self, mock_openai):
+        with patch.dict(os.environ, {"MINIMAX_API_KEY": "sk-test-key"}):
+            _get_client("MiniMax-M2.7-highspeed")
+            mock_openai.assert_called_once_with(
+                api_key="sk-test-key",
+                base_url=MINIMAX_BASE_URL,
+            )
+
+    @patch("opensora.utils.prompt_refine.OpenAI")
+    def test_openai_model_uses_openai_key(self, mock_openai):
+        with patch.dict(os.environ, {"OPENAI_API_KEY": "sk-openai-key"}):
+            _get_client("gpt-4o")
+            mock_openai.assert_called_once_with(api_key="sk-openai-key")
+
+    def test_minimax_model_raises_without_api_key(self):
+        env = {k: v for k, v in os.environ.items() if k != "MINIMAX_API_KEY"}
+        with patch.dict(os.environ, env, clear=True):
+            os.environ.pop("MINIMAX_API_KEY", None)
+            with self.assertRaises(ValueError) as ctx:
+                _get_client("MiniMax-M2.7")
+            self.assertIn("MINIMAX_API_KEY", str(ctx.exception))
+
+
+class TestRefinePrompt(unittest.TestCase):
+    def _make_mock_response(self, content="A refined prompt."):
+        choice = MagicMock()
+        choice.message.content = content
+        response = MagicMock()
+        response.choices = [choice]
+        return response
+
+    @patch("opensora.utils.prompt_refine.OpenAI")
+    def test_refine_prompt_uses_minimax_model(self, mock_openai_cls):
+        mock_client = MagicMock()
+        mock_client.chat.completions.create.return_value = self._make_mock_response(
+            "A beautiful forest scene with sunlight streaming through the trees."
+        )
+        mock_openai_cls.return_value = mock_client
+
+        from opensora.utils.prompt_refine import refine_prompt
+
+        with patch.dict(os.environ, {"MINIMAX_API_KEY": "sk-test-key"}):
+            result = refine_prompt("a forest", type="t2v", model="MiniMax-M2.7")
+
+        self.assertEqual(result, "A beautiful forest scene with sunlight streaming through the trees.")
+        call_kwargs = mock_client.chat.completions.create.call_args[1]
+        self.assertEqual(call_kwargs["model"], "MiniMax-M2.7")
+        self.assertGreater(call_kwargs["temperature"], 0.0)
+
+    @patch("opensora.utils.prompt_refine.OpenAI")
+    def test_refine_prompt_default_model_is_gpt4o(self, mock_openai_cls):
+        mock_client = MagicMock()
+        mock_client.chat.completions.create.return_value = self._make_mock_response()
+        mock_openai_cls.return_value = mock_client
+
+        from opensora.utils.prompt_refine import refine_prompt
+
+        env = {k: v for k, v in os.environ.items() if k != "PROMPT_MODEL"}
+        with patch.dict(os.environ, env, clear=True):
+            os.environ.pop("PROMPT_MODEL", None)
+            refine_prompt("a forest", type="t2v")
+
+        call_kwargs = mock_client.chat.completions.create.call_args[1]
+        self.assertEqual(call_kwargs["model"], "gpt-4o")
+
+    @patch("opensora.utils.prompt_refine.OpenAI")
+    def test_prompt_model_env_overrides_default(self, mock_openai_cls):
+        mock_client = MagicMock()
+        mock_client.chat.completions.create.return_value = self._make_mock_response()
+        mock_openai_cls.return_value = mock_client
+
+        from opensora.utils.prompt_refine import refine_prompt
+
+        with patch.dict(os.environ, {"MINIMAX_API_KEY": "sk-key", "PROMPT_MODEL": "MiniMax-M2.7"}):
+            refine_prompt("a forest", type="t2v")
+
+        call_kwargs = mock_client.chat.completions.create.call_args[1]
+        self.assertEqual(call_kwargs["model"], "MiniMax-M2.7")
+
+    @patch("opensora.utils.prompt_refine.OpenAI")
+    def test_temperature_is_positive(self, mock_openai_cls):
+        mock_client = MagicMock()
+        mock_client.chat.completions.create.return_value = self._make_mock_response()
+        mock_openai_cls.return_value = mock_client
+
+        from opensora.utils.prompt_refine import refine_prompt
+
+        with patch.dict(os.environ, {"MINIMAX_API_KEY": "sk-test-key"}):
+            refine_prompt("a forest", type="t2v", model="MiniMax-M2.7")
+
+        call_kwargs = mock_client.chat.completions.create.call_args[1]
+        self.assertGreater(call_kwargs["temperature"], 0.0, "temperature must be > 0 for MiniMax")
+        self.assertLessEqual(call_kwargs["temperature"], 1.0)
+
+    @patch("opensora.utils.prompt_refine.OpenAI")
+    def test_refine_prompts_by_minimax_uses_m3(self, mock_openai_cls):
+        mock_client = MagicMock()
+        mock_client.chat.completions.create.return_value = self._make_mock_response("Refined.")
+        mock_openai_cls.return_value = mock_client
+
+        with patch.dict(os.environ, {"MINIMAX_API_KEY": "sk-test-key"}):
+            results = refine_prompts_by_minimax(["a dog"], type="t2v")
+
+        self.assertEqual(results, ["Refined."])
+        call_kwargs = mock_client.chat.completions.create.call_args[1]
+        self.assertEqual(call_kwargs["model"], "MiniMax-M3")
+
+    @patch("opensora.utils.prompt_refine.OpenAI")
+    def test_returns_original_on_failure(self, mock_openai_cls):
+        mock_client = MagicMock()
+        mock_client.chat.completions.create.return_value = MagicMock(choices=[])
+        mock_openai_cls.return_value = mock_client
+
+        from opensora.utils.prompt_refine import refine_prompt
+
+        with patch.dict(os.environ, {"MINIMAX_API_KEY": "sk-test-key"}):
+            result = refine_prompt("original prompt", type="t2v", model="MiniMax-M2.7", retry_times=1)
+
+        self.assertEqual(result, "original prompt")
+
+
+if __name__ == "__main__":
+    unittest.main()