edit: faiss 전용 rest 서버 추가

custom_apps/imagen/custom_imagen.py

@@ -5,8 +5,8 @@ import os
 from vertexai.preview.vision_models import ImageGenerationModel
 
 from const import OUTPUT_FOLDER
-from rest.app.utils.parsing_utils import prompt_to_filenames
-from rest.app.utils.date_utils import D
+from main_rest.app.utils.parsing_utils import prompt_to_filenames
+from main_rest.app.utils.date_utils import D
 
 
 class ImagenConst:
@@ -48,4 +48,44 @@ def imagen_generate_image(prompt,download_count=1):
         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)
+        return len(images.images)
+    
+def imagen_generate_image_data(prompt,download_count=1):
+    vertexai.init(project=ImagenConst.project_id, location=ImagenConst.location)
+
+    model = ImageGenerationModel.from_pretrained(ImagenConst.model)
+    
+    images = model.generate_images(
+        prompt=prompt,
+        # Optional parameters
+        number_of_images=download_count,
+        language="ko",
+        # You can't use a seed value and watermark at the same time.
+        # add_watermark=False,
+        # seed=100,
+        aspect_ratio="1:1",
+        safety_filter_level="block_some",
+        person_generation="dont_allow",
+    )
+    
+    return images.images[0]._pil_image
+
+def imagen_generate_image_path(image_prompt):
+    
+    MODEL = "imagen"
+    QUERY = "query"
+    create_time = D.date_file_name()
+    
+    folder_name = os.path.join(OUTPUT_FOLDER,f"{MODEL}_{QUERY}_{image_prompt}_{create_time}")
+    if not os.path.exists(folder_name):
+        os.makedirs(folder_name)
+    
+    generate_img = imagen_generate_image_data(image_prompt)
+    
+    generate_img.save(os.path.join(folder_name,f"query.png"))
+       
+    return os.path.join(folder_name,f"query.png")
+
+if __name__ == '__main__':
+    pass
+    # imagen_generate_image_data("cat")