Merge pull request #203 from Soulter/feat-custom-t2i-tmpl

自定义文转图 HTML 模板

util/t2i/renderer.py

@@ -12,7 +12,22 @@ class TextToImageRenderer:
         self.local_strategy = LocalRenderStrategy()
         self.context = RenderContext(self.network_strategy)
 
+    async def render_custom_template(self, tmpl_str: str, tmpl_data: dict, return_url: bool = False):
+        '''使用自定义文转图模板。该方法会通过网络调用 t2i 终结点图文渲染API。
+        @param tmpl_str: HTML Jinja2 模板。
+        @param tmpl_data: jinja2 模板数据。
+
+        @return: 图片 URL 或者文件路径，取决于 return_url 参数。
+
+        @example: 参见 https://astrbot.soulter.top 插件开发部分。
+        '''
+        local = locals()
+        local.pop('self')
+        return await self.network_strategy.render_custom_template(**local)
+
     async def render(self, text: str, use_network: bool = True, return_url: bool = False):
+        '''使用默认文转图模板。
+        '''
         if use_network:
             try:
                 return await self.context.render(text, return_url=return_url)

util/t2i/strategies/base_strategy.py

@@ -4,4 +4,7 @@ class RenderStrategy(ABC):
     @abstractmethod
     def render(self, text: str, return_url: bool) -> str:
         pass
-    
+    
+    @abstractmethod
+    def render_custom_template(self, tmpl_str: str, tmpl_data: dict, return_url: bool) -> str:
+        pass

util/t2i/strategies/local_strategy.py

@@ -6,6 +6,9 @@ from PIL import ImageFont, Image, ImageDraw
 from util.io import save_temp_img
 
 class LocalRenderStrategy(RenderStrategy):
+
+    async def render_custom_template(self, tmpl_str: str, tmpl_data: dict, return_url: bool=True) -> str:
+        raise NotImplementedError
     
     def get_font(self, size: int) -> ImageFont.FreeTypeFont:
         # common and default fonts on Windows, macOS and Linux

util/t2i/strategies/network_strategy.py

@@ -5,40 +5,40 @@ from .base_strategy import RenderStrategy
 from type.config import VERSION
 from util.io import download_image_by_url
 
+ASTRBOT_T2I_DEFAULT_ENDPOINT = "https://t2i.soulter.top/text2img"
+
 class NetworkRenderStrategy(RenderStrategy):
-    def __init__(self) -> None:
-        super().__init__()  
-        self.BASE_RENDER_URL = "https://t2i.soulter.top/text2img"
+    def __init__(self, base_url: str = ASTRBOT_T2I_DEFAULT_ENDPOINT) -> None:
+        super().__init__()
+        self.BASE_RENDER_URL = base_url
         self.TEMPLATE_PATH = os.path.join(os.path.dirname(__file__), "template")
 
+    async def render_custom_template(self, tmpl_str: str, tmpl_data: dict, return_url: bool=True) -> str:
+        '''使用自定义文转图模板'''
+        post_data = {
+            "tmpl": tmpl_str,
+            "json": return_url,
+            "tmpldata": tmpl_data,
+            "options": {
+                "full_page": True,
+                "type": "jpeg",
+                "quality": 40,
+            }
+        }
+        if return_url:
+            async with aiohttp.ClientSession() as session:
+                async with session.post(f"{self.BASE_RENDER_URL}/generate", json=post_data) as resp:
+                    ret = await resp.json()
+                    return f"{self.BASE_RENDER_URL}/{ret['data']['id']}"
+        return await download_image_by_url(f"{self.BASE_RENDER_URL}/generate", post=True, post_data=post_data)
+
+
     async def render(self, text: str, return_url: bool=False) -> str:
         '''
         返回图像的文件路径
         '''
         with open(os.path.join(self.TEMPLATE_PATH, "base.html"), "r", encoding='utf-8') as f:
             tmpl_str = f.read()
-
         assert(tmpl_str)
-
         text = text.replace("`", "\`")
-
-        post_data = {
-            "tmpl": tmpl_str,
-            "json": return_url,
-            "tmpldata": {
-                "text": text,
-                "version": f"v{VERSION}",
-            },
-            "options": {
-                "full_page": True,
-                "type": "jpeg",
-                "quality": 40,
-            }
-        }
-        
-        if return_url:
-            async with aiohttp.ClientSession() as session:
-                async with session.post(f"{self.BASE_RENDER_URL}/generate", json=post_data) as resp:
-                    ret = await resp.json()
-                    return f"{self.BASE_RENDER_URL}/{ret['data']['id']}"
-        return await download_image_by_url(f"{self.BASE_RENDER_URL}/generate", post=True, post_data=post_data)
+        return await self.render_custom_template(tmpl_str, {"text": text, "version": f"v{VERSION}"}, return_url)