绘图模型添加model=2

ai_pic.py

@@ -73,7 +73,7 @@ def start_paint(id,data):
     else:
         print(f"请求失败，状态码: {response.status_code}")
 
-def get_websocket_data(client_id):
+def get_websocket_data(client_id, end_node = "16"):
     """
     获取生图任务运行状态，根据实例id进行跟踪
     """
@@ -81,7 +81,7 @@ def get_websocket_data(client_id):
     def on_message(ws, message):
         nonlocal output_image_filenames
         message_dict = json.loads(message)
-        if message_dict["type"] == "executed" and message_dict["data"]["node"] == "16":
+        if message_dict["type"] == "executed" and message_dict["data"]["node"] == end_node:
             output_data = message_dict["data"]
             if "output" in output_data and "images" in output_data["output"]:
                 image_list = output_data["output"]["images"]
@@ -191,7 +191,7 @@ def story_start_p2p(qq, model=1):
         data["3"]["inputs"]["seed"] = random.randint(1, 2 ** 32 - 1)
         id = generate_instance_id()
         start_paint(id, data)
-        return get_images_from_filenames(get_websocket_data(id))
+        return get_images_from_filenames(get_websocket_data(id, "9"))
 
 
 def story_start_p2p_sese(qq, model=1):
@@ -217,5 +217,5 @@ def story_start_p2p_sese(qq, model=1):
         data['6']['inputs']["text"] = "nsfw,1girl,best quality,masterpiece,"
         id = generate_instance_id()
         start_paint(id, data)
-        return get_images_from_filenames(get_websocket_data(id))
+        return get_images_from_filenames(get_websocket_data(id, "9"))