edit : 이미지 생성시 생성 갯수 지정 가능

.gitignore

@@ -144,3 +144,4 @@ cython_debug/
 log
 sheet_counter.txt
 
+*output

custom_apps/dalle3/BingImageCreator.py

@@ -196,12 +196,13 @@ class ImageGenAsync:
             os.mkdir(output_dir)
         try:
             model = "dalle3"
-            jpeg_index = 0
+            jpeg_index = 1
             parsing_file_name = prompt_to_filenames(prompt)
+            _datetime = D.date_file_name()
 
             for link in links[:download_count]:
                 if download_count == 1:
-                    _path = os.path.join(output_dir, f"{model}_{parsing_file_name}_{D.date_file_name()}.jpg")
+                    _path = os.path.join(output_dir, f"{model}_{parsing_file_name}_{_datetime}.jpg")
                     if os.path.exists(_path):
                         raise Exception("파일 이미 존재함")
                     
@@ -216,7 +217,7 @@ class ImageGenAsync:
                     jpeg_index += 1
                 
                 else:
-                    _path = os.path.join(output_dir, f"{model}_{parsing_file_name}_{jpeg_index}_{D.date_file_name()}.jpg")
+                    _path = os.path.join(output_dir, f"{model}_{parsing_file_name}_{jpeg_index}_{_datetime}.jpg")
                     if os.path.exists(_path):
                         raise Exception("파일 이미 존재함")
                     

custom_apps/imagen/custom_imagen.py

@@ -22,6 +22,7 @@ def imagen_generate_image(prompt,download_count=1):
 
     _file_name = prompt_to_filenames(prompt)
     _folder = os.path.join(OUTPUT_FOLDER,"imagen")
+    _datetime = D.date_file_name()
     
     if not os.path.isdir(_folder):
         raise FileExistsError(f"FileExistsError: {_folder}")
@@ -40,7 +41,7 @@ def imagen_generate_image(prompt,download_count=1):
     )
     
     if len(images.images) <=1 :
-        images[0].save(location=os.path.join(_folder,f"imagen_{_file_name}_{D.date_file_name()}.png"), include_generation_parameters=False)
+        images[0].save(location=os.path.join(_folder,f"imagen_{_file_name}_{_datetime}.png"), include_generation_parameters=False)
     else:
         for i in range(len(images.images)):
-            images[i].save(location=os.path.join(_folder,f"imagen_{_file_name}_{i}_{D.date_file_name()}.png"), include_generation_parameters=False)
+            images[i].save(location=os.path.join(_folder,f"imagen_{_file_name}_{i+1}_{_datetime}.png"), include_generation_parameters=False)

rest/app/models.py

@@ -571,4 +571,4 @@ class ImageGenerateReq(BaseModel):
     ### [Request] image generate request
     """
     prompt : str = Field(description='프롬프트', example='검은색 안경')
-    # downloadCount : int = Field(1, description='이미지 생성 갯수', example=1)
+    downloadCount : int = Field(1, description='이미지 생성 갯수', example=1)

rest/app/routes/services.py

@@ -21,6 +21,7 @@ from custom_logger.custom_log import custom_logger as LOG
 
 from custom_apps.dalle3.custom_dalle import DallEArgument,dalle3_generate_image
 from custom_apps.imagen.custom_imagen import imagen_generate_image
+from rest.app.utils.parsing_utils import download_range
 
 
 router = APIRouter(prefix="/services")
@@ -39,9 +40,14 @@ async def dalle3(request: Request, request_body_info: M.ImageGenerateReq):
     """
     response = M.ResponseBase()
     try:
-        args = DallEArgument(prompt=request_body_info.prompt
-                    #   , download_count=request_body_info.downloadCount
-                    )
+        if not download_range(request_body_info.downloadCount):
+            raise Exception(f"downloadCount is 1~4 (current value = {request_body_info.downloadCount})")
+
+
+        args = DallEArgument(
+            prompt=request_body_info.prompt,
+            download_count=request_body_info.downloadCount
+            )       
         
         dalle3_generate_image(args)
         
@@ -66,7 +72,12 @@ async def imagen(request: Request, request_body_info: M.ImageGenerateReq):
     response = M.ResponseBase()
     try:
         
-        imagen_generate_image(request_body_info.prompt)
+        if not download_range(request_body_info.downloadCount):
+            raise Exception(f"downloadCount is 1~4 (current value = {request_body_info.downloadCount})")
+        
+        imagen_generate_image(prompt=request_body_info.prompt,
+                              download_count=request_body_info.downloadCount
+                              )
         
         return response.set_message()
     

rest/app/utils/parsing_utils.py

@@ -11,4 +11,15 @@ def prompt_to_filenames(prompt):
         else:
             filename += i
             
-    return filename
+    return filename
+
+
+def download_range(download_count:int):
+    _min = 1
+    _max = 4
+    
+    if _min <= download_count and download_count <= _max:
+        return True
+
+    return False
+