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

2025-01-17 16:06:01 +09:00
parent c296204541
commit 6e5f990fc1
5 changed files with 121 additions and 38 deletions
--- a/custom_apps/imagen/custom_imagen.py
+++ b/custom_apps/imagen/custom_imagen.py
@@ -42,6 +42,10 @@ 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}_{_datetime}.png"), include_generation_parameters=False)
+        return 1
+    
    else:
        for i in range(len(images.images)):
            images[i].save(location=os.path.join(_folder,f"imagen_{_file_name}_{i+1}_{_datetime}.png"), include_generation_parameters=False)
+        
+        return len(images.images)