feat: stable diffusion (#265)

* feat: stable diffusion
* docs: update docs

Author: timerring

README.md

@@ -22,6 +22,7 @@
   <img src="assets/siliconcloud-color.svg" alt="SiliconFlow" width="15" height="60" />
   <img src="assets/siliconcloud-text.svg" alt="SiliconFlow" width="100" height="60" />
   <img src="assets/wenxin-color.svg" alt="Baidu ERNIE" width="60" height="60" />
+  <img src="assets/stability-brand-color.svg" alt="Stability AI" width="80" height="60" />
 </div>
 
 ##  1. Introduction
@@ -50,7 +51,8 @@
   - `Minimax image-01`
   - `Kwai Kolors`
   - `Tencent Hunyuan`
-  - `Baidu ERNIE`
+  - `Baidu ERNIE irag-1.0`
+  - `Stable Diffusion 3.5 large turbo`
 
 项目架构流程如下：
 
@@ -223,6 +225,12 @@ MLLM 模型主要用于自动切片后的切片标题生成，此功能默认关
 
 请自行[注册账号](https://console.bce.baidu.com/iam/key/list)并申请 API Key，填写到 `bilive.toml` 文件中对应的 `BAIDU_API_KEY` 中。
 
+##### 3.3.5 Stability SD 3.5 large turbo 模型
+
+> 如需使用 Stability SD 3.5 large turbo  模型，请将 `IMAGE_GEN_MODEL` 参数设置为 `stability`。
+
+请自行[注册账号](https://platform.stability.ai/account/keys)并申请 API Key，填写到 `bilive.toml` 文件中对应的 `STABILITY_API_KEY` 中。
+
 #### 4. bilitool 登录
 
 > 由于一般日志打印不出二维码效果（docker 的日志不确定是否能打印，等发布新image时再修改，docker 版本请先参考文档 [bilive](https://bilive.timerring.com)，本 README 只针对源码部署），所以这步需要提前在机器上安装 [bilitool](https://github.com/timerring/bilitool):

assets/stability-brand-color.svg

@@ -37,9 +37,10 @@ qwen_api_key = "" # Apply for your own Qwen API key at https://bailian.console.a
 
 [cover]
 generate_cover = false # whether to generate cover
-image_gen_model = "minimax" # the image generation model, can be "minimax" or "siliconflow" or "tencent" or "baidu"
+image_gen_model = "minimax" # the image generation model, can be "minimax" or "siliconflow" or "tencent" or "baidu" or "stability"
 minimax_api_key = "" # Apply for your own Minimax API key at https://platform.minimaxi.com/user-center/basic-information/interface-key
 siliconflow_api_key = "" # Apply for your own SiliconFlow API key at https://cloud.siliconflow.cn/i/3Szr5BVg
 tencent_secret_id = "" # Apply for your own Tencent Cloud API key at https://console.cloud.tencent.com/cam/capi
 tencent_secret_key = "" # Apply for your own Tencent Cloud secret key as above
 baidu_api_key = "" # Apply for your own Baidu API key at https://console.bce.baidu.com/iam/key/list
+stability_api_key = "" # Apply for your own Stability API key at https://platform.stability.ai/account/keys

bilive.toml

@@ -78,4 +78,5 @@ def get_interface_config():
 SILICONFLOW_API_KEY = config.get('cover', {}).get('siliconflow_api_key')
 TENCENT_SECRET_ID = config.get('cover', {}).get('tencent_secret_id')
 TENCENT_SECRET_KEY = config.get('cover', {}).get('tencent_secret_key')
-BAIDU_API_KEY = config.get('cover', {}).get('baidu_api_key')
+BAIDU_API_KEY = config.get('cover', {}).get('baidu_api_key')
+STABILITY_API_KEY = config.get('cover', {}).get('stability_api_key')

src/config.py

@@ -67,6 +67,12 @@ def wrapper(video_path):
                 from .image_model_sdk.baidu_sdk import baidu_generate_cover
 
                 return baidu_generate_cover(cover_path)
+            elif model_type == "stability":
+                from .image_model_sdk.stability_sdk import (
+                    stable_diffusion_generate_cover,
+                )
+
+                return stable_diffusion_generate_cover(cover_path)
             else:
                 upload_log.error(f"Unsupported model type: {model_type}")
                 return None

src/cover/cover_generator.py

@@ -8,6 +8,7 @@
 from src.upload.bilitool.bilitool.model.model import Model
 from src.config import BAIDU_API_KEY
 
+
 def cover_up(img: str):
     """Upload the cover image
     Parameters
@@ -20,14 +21,15 @@ def cover_up(img: str):
     """
     from PIL import Image
     from io import BytesIO
+
     request = requests.Session()
     request.headers = {
-        'user-agent': "Mozilla/5.0 (X11; Linux x86_64) AppleWebKit/537.36 Chrome/63.0.3239.108",
-        'referer': "https://www.bilibili.com/",
-        'connection': 'keep-alive'
+        "user-agent": "Mozilla/5.0 (X11; Linux x86_64) AppleWebKit/537.36 Chrome/63.0.3239.108",
+        "referer": "https://www.bilibili.com/",
+        "connection": "keep-alive",
     }
     model = Model().get_config()
-    request.cookies.set('SESSDATA', model['cookies']['SESSDATA'])
+    request.cookies.set("SESSDATA", model["cookies"]["SESSDATA"])
     with Image.open(img) as im:
         # you should keep the image ratio 16:10
         xsize, ysize = im.size
@@ -44,7 +46,7 @@ def cover_up(img: str):
         data={
             "cover": b"data:image/jpeg;base64,"
             + (base64.b64encode(buffered.getvalue())),
-            "csrf": model['cookies']['bili_jct']
+            "csrf": model["cookies"]["bili_jct"],
         },
         timeout=30,
     )
@@ -55,24 +57,33 @@ def cover_up(img: str):
     print(res["data"]["url"], flush=True)
     return res["data"]["url"]
 
+
 def baidu_generate_cover(your_file_path):
+    """Generater cover image using baidu api
+    Args:
+        your_file_path: str, path to the image file
+    Returns:
+        str, local download path of the generated cover image file
+    """
     try:
         cover_url = cover_up(your_file_path)
 
         url = "https://qianfan.baidubce.com/v2/images/generations"
-        payload = json.dumps({
-        "model": "irag-1.0",
-        "prompt": "这是视频截图，请根据该图生成对应的动漫类型的封面",
-        "refer_image": cover_url
-    })
+        payload = json.dumps(
+            {
+                "model": "irag-1.0",
+                "prompt": "这是视频截图，请根据该图生成对应的动漫类型的封面",
+                "refer_image": cover_url,
+            }
+        )
         headers = {
-            'Content-Type': 'application/json',
-            'Authorization': f'Bearer {BAIDU_API_KEY}'
+            "Content-Type": "application/json",
+            "Authorization": f"Bearer {BAIDU_API_KEY}",
         }
-        
+
         response = requests.request("POST", url, headers=headers, data=payload)
         if response.status_code == 200:
-            image_url=response.json()['data'][0]['url']
+            image_url = response.json()["data"][0]["url"]
             img_data = requests.get(image_url).content
             cover_name = time.strftime("%Y%m%d%H%M%S") + ".png"
             temp_cover_path = os.path.join(os.path.dirname(your_file_path), cover_name)
@@ -87,5 +98,6 @@ def baidu_generate_cover(your_file_path):
         print(e, flush=True)
         return None
 
-if __name__ == '__main__':
-    print(baidu_generate_cover(""))
+
+if __name__ == "__main__":
+    print(baidu_generate_cover(""))

src/cover/image_model_sdk/baidu_sdk.py

@@ -0,0 +1,46 @@
+import requests
+from src.config import STABILITY_API_KEY
+import time
+import os
+
+
+def stable_diffusion_generate_cover(your_file_path):
+    """Generater cover image using stability api
+    Args:
+        image_path: str, path to the image file
+    Returns:
+        str, local download path of the generated cover image file
+    """
+
+    cover_name = time.strftime("%Y%m%d%H%M%S") + ".jpeg"
+    temp_cover_path = os.path.join(os.path.dirname(your_file_path), cover_name)
+
+    with open(your_file_path, "rb") as img_file:
+        response = requests.post(
+            f"https://api.stability.ai/v2beta/stable-image/generate/sd3",
+            headers={
+                "authorization": f"Bearer {STABILITY_API_KEY}",
+                "accept": "image/*",
+            },
+            files={"image": ("image.jpg", img_file, "image/jpeg")},
+            data={
+                "prompt": "This is a video screenshot, please generate a cover in the style of a manga",  # English only
+                "strength": 0.75,
+                "output_format": "jpeg",
+                "mode": "image-to-image",
+                "model": "sd3.5-large-turbo",
+            },
+        )
+
+    if response.status_code == 200:
+        with open(temp_cover_path, "wb") as file:
+            file.write(response.content)
+        os.remove(your_file_path)
+        return temp_cover_path
+    else:
+        raise Exception(str(response.json()))
+        return None
+
+
+if __name__ == "__main__":
+    print(stable_diffusion_generate_cover(""))