1

ai_pyppeteer.py

+import multiprocessing
+
 import requests as req
 import time
 from loguru import logger
@@ -464,7 +466,7 @@ def process_start(browser_item, lock):
 
 if __name__ == '__main__':
     try:
-        logger.info('__main__2')
+        multiprocessing.freeze_support()
         browser_list = [
             {
                 'num': 1,

download_img.py

@@ -8,13 +8,13 @@ from Helper import RedisCon
 import requests
 import os
 from custom_oss_upload import AliyunOss
+import multiprocessing
 
 
 class DownLoad:
-    def __init__(self, lock, task):
-        self.key = ''
+    def __init__(self, lock):
+        self.key = 'laravel_database_discord_img_download'
         self.lock = lock
-        self.task = task
         self.header = {
             'Accept': 'application/json',
             'Content-Type': 'application/json',
@@ -22,12 +22,12 @@ class DownLoad:
         }
 
     # 图片下载回调
-    async def task_callback(self, status, result_img, message=''):
-        url = self.task['callback']
+    async def task_callback(self, status, result_img, task, message=''):
+        url = task['callback']
         for i in range(5):
             try:
                 params = {
-                    'prompt_id': self.task['prompt_id'],
+                    'prompt_id': task['prompt_id'],
                 }
                 if status:
                     # 为真-成功
@@ -62,66 +62,62 @@ class DownLoad:
             except Exception as e:
                 logger.error(e)
 
-    async def upload_ali_oss(self, path, cur_uuid):
+    async def upload_ali_oss(self, path, cur_uuid, task):
         # path--本地图片路径
         # name- 阿里云存储位置
         name = f"files/download/admin/{cur_uuid}"
         result_img = await AliyunOss().put_object_from_file(name, path)
         if result_img:
             # 回调
-            await self.task_callback(True, result_img)
+            await self.task_callback(True, result_img, task)
+
+    async def get_cur_task(self):
+        task = RedisCon.get_task(self.key)
+        if task:
+            return json.loads(task)
+        return False
 
     # 开始
     async def on_start(self):
-        logger.info('开始执行')
-        cur_uuid = uuid.uuid4()
-        cur_uuid = str(cur_uuid) + '.png'
-        # 下载图片
-        path = await self.download_img(self.task['result_img'], cur_uuid)
-        if path:
-            logger.info('图片下载成功')
-            # 图片下载成功--准备上传阿里云
-            await self.upload_ali_oss(path, cur_uuid)
-        else:
-            # 图片下载失败-直接回调-失败状态
-            await self.task_callback(False, '', '转存discord图片失败')
+        while True:
+            # 获取任务
+            task = await self.get_cur_task()
+            if task:
+                cur_uuid = uuid.uuid4()
+                cur_uuid = str(cur_uuid) + '.png'
+                # 下载图片
+                path = await self.download_img(task['result_img'], cur_uuid)
+                if path:
+                    logger.info('图片下载成功')
+                    # 图片下载成功--准备上传阿里云
+                    await self.upload_ali_oss(path, cur_uuid, task)
+                else:
+                    # 图片下载失败-直接回调-失败状态
+                    await self.task_callback(False, '', task, '转存discord图片失败')
+            else:
+                time.sleep(3)
+                logger.info('不存在任务,等待3秒')
 
 
-def process_start(lock, task):
+def process_start(lock):
     loop = asyncio.get_event_loop()
-    keyword = DownLoad(lock, task)
+    keyword = DownLoad(lock)
     # 生产
     loop.run_until_complete(keyword.on_start())
 
 
 if __name__ == '__main__':
     # 获取任务并实例化
-    key = 'laravel_database_discord_img_download'
-    while True:
-        try:
-            tasks = list()
-            task_num = RedisCon.get_task_num(key)
-            logger.info(f'{task_num}个任务')
-            nums = 5
-            process = 0
-            if task_num and task_num > 0:
-                if task_num < nums:
-                    process = task_num
-                else:
-                    process = nums
-            lock = Lock()
-            l = []
-            for i in range(process):
-                task = RedisCon.get_task(key)
-                if task:
-                    task = json.loads(task)
-                    p = Process(target=process_start, args=(lock, task))
-                    p.start()
-                    l.append(p)  # 将进程加入列表中，这样 5 个子进程就会全部执行
+    multiprocessing.freeze_support()
+    try:
+        lock = Lock()
+        l = []
+        for i in range(5):
+            p = Process(target=process_start, args=(lock,))
+            p.start()
+            l.append(p)
 
-            for p in l:
-                p.join()  # 当5 个子进程全部执行完，才会执行主进程
-        except Exception as e:
-            logger.error(e)
-            logger.info('等待3秒')
-            time.sleep(3)
+        for p in l:
+            p.join()  # 当5 个子进程全部执行完，才会执行主进程
+    except Exception as e:
+        logger.error(e)