edit : 실제 이미지 생성한 갯수 response 에 추가

custom_apps/dalle3/BingImageCreator.py

@@ -181,6 +181,7 @@ class ImageGenAsync:
         links: list,
         output_dir: str,
         download_count: int,
+        counter,
         file_name: str = None,
         prompt: str = None
     ) -> None:
@@ -202,7 +203,8 @@ class ImageGenAsync:
 
             for link in links[:download_count]:
                 if download_count == 1:
-                    _path = os.path.join(output_dir, f"{model}_{parsing_file_name}_{_datetime}.jpg")
+                    _file_name = f"{model}_{parsing_file_name}_{_datetime}.png"
+                    _path = os.path.join(output_dir, _file_name)
                     if os.path.exists(_path):
                         raise Exception("파일 이미 존재함")
                     
@@ -211,26 +213,43 @@ class ImageGenAsync:
                     response = await self.session.get(link)
                     if response.status_code != 200:
                         raise Exception("Could not download image")
-                    # save response to file
-                    with open(_path,"wb") as output_file:
-                        output_file.write(response.content)
+                    
+                    # 이미지 파일 생성 x(error)
+                    if 'svg viewBox' in str(response.content):
+                        counter.add_error_messages(_file_name)
+                        
+                    # 이미지 파일 생성
+                    else:
+                        # save response to file
+                        with open(_path,"wb") as output_file:
+                            output_file.write(response.content)
+                            counter.count()
+                    
                     jpeg_index += 1
                 
                 else:
-                    _path = os.path.join(output_dir, f"{model}_{parsing_file_name}_{jpeg_index}_{_datetime}.jpg")
+                    _file_name = f"{model}_{parsing_file_name}_{jpeg_index}_{_datetime}.png"
+                    _path = os.path.join(output_dir, _file_name)
                     if os.path.exists(_path):
                         raise Exception("파일 이미 존재함")
-                    
-                    while os.path.exists(_path):
-                        jpeg_index += 1
+                                           
                     response = await self.session.get(link)
                     if response.status_code != 200:
                         raise Exception("Could not download image")
-                    # save response to file
-                    with open(_path,"wb") as output_file:
-                        output_file.write(response.content)
-                    jpeg_index += 1
                     
+                    # 이미지 파일 생성 x(error)
+                    if 'svg viewBox' in str(response.content) or 'var' in str(response.content):
+                        counter.add_error_messages(_file_name)
+                        
+                    # 이미지 파일 생성
+                    else:
+                        # save response to file
+                        with open(_path,"wb") as output_file:
+                            output_file.write(response.content)
+                            counter.count()
+                    
+                    jpeg_index += 1
+                                
         except httpx.InvalidURL as url_exception:
             raise Exception(
                 "Inappropriate contents found in the generated images. Please try again or try another prompt.",
@@ -241,18 +260,12 @@ async def async_image_gen(
     prompt: str,
     download_count: int,
     output_dir: str,
+    counter,
     u_cookie=None,
     debug_file=None,
     quiet=False,
     all_cookies=None,
 ):
-    async with ImageGenAsync(
-        u_cookie,
-        debug_file=debug_file,
-        quiet=quiet,
-        all_cookies=all_cookies,
-    ) as image_generator:
+    async with ImageGenAsync(u_cookie,debug_file=debug_file,quiet=quiet,all_cookies=all_cookies,) as image_generator:
         images = await image_generator.get_images(prompt)
-        await image_generator.save_images(
-            images, output_dir=output_dir, download_count=download_count,prompt=prompt
-        )
+        await image_generator.save_images(links=images, counter=counter, output_dir=output_dir, download_count=download_count,prompt=prompt)

custom_apps/dalle3/custom_dalle.py

@@ -47,12 +47,40 @@ class DallEArgument:
     quiet: bool = False
     asyncio: bool = True
     version: bool = False
+    
 
+class Counter:
+    """
+    생성된 이미지 카운트
+    """
+    
+    def __init__(self):
+        self.counter = 0
+        self.error_messages = []
+        
+    def reset(self):
+        self.counter = 0
+        self.error_messages = []
+        
+    def count(self):
+        self.counter += 1
+        
+    def get_counter(self):
+        return self.counter
+
+    def add_error_messages(self,error):
+        self.error_messages.append(error)
+        
+    def get_error_messages(self):
+        return self.error_messages
+    
 
 def dalle3_generate_image(args):   
     
     nest_asyncio.apply()
     
+    counter = Counter()
+    
     if not os.path.isdir(args.output_dir):
         raise FileExistsError(f"FileExistsError: {args.output_dir}")
     
@@ -76,13 +104,16 @@ def dalle3_generate_image(args):
 
     asyncio.run(
         async_image_gen(
-            args.prompt,
-            args.download_count,
-            args.output_dir,
-            args.U,
-            args.debug_file,
-            args.quiet,
+            prompt=args.prompt,
+            download_count=args.download_count,
+            output_dir=args.output_dir,
+            counter=counter,
+            u_cookie=args.U,
+            debug_file=args.debug_file,
+            quiet=args.quiet,
             all_cookies=cookie_json,
         ),
     )
+    
+    return counter