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

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()