FERMAT_TEST_REST_SERVER/GLASSES_AI_SERVER
4
0
Fork 0
Code Issues Pull Requests Projects Releases 1 Wiki Activity

edit : 원격 저장으로 변경

Browse Source
This commit is contained in:
김재웅
2025-04-28 14:23:46 +09:00
parent 6b212125a4
commit 4c2ea70289
8 changed files with 64 additions and 19 deletions
Download Patch File Download Diff File Expand all files Collapse all files

16
custom_apps/imagen/custom_imagen.py
View File

@@ -4,10 +4,10 @@ import vertexai
import os
from vertexai.preview.vision_models import ImageGenerationModel
from const import OUTPUT_FOLDER
from const import REMOTE_FOLDER, TEMP_FOLDER
from main_rest.app.utils.parsing_utils import prompt_to_filenames
from main_rest.app.utils.date_utils import D
from utils.custom_sftp import sftp_client
class ImagenConst:
project_id = "glasses-imagen"
@@ -21,11 +21,12 @@ def imagen_generate_image(prompt,download_count=1):
model = ImageGenerationModel.from_pretrained(ImagenConst.model)
_file_name = prompt_to_filenames(prompt)
_folder = os.path.join(OUTPUT_FOLDER,"imagen")
_folder = os.path.join(TEMP_FOLDER)
_remote_folder = os.path.join(REMOTE_FOLDER,"imagen")
_datetime = D.date_file_name()
if not os.path.isdir(_folder):
raise FileExistsError(f"FileExistsError: {_folder}")
# if not os.path.isdir(_folder):
# raise FileExistsError(f"FileExistsError: {_folder}")
images = model.generate_images(
prompt=prompt,
@@ -42,12 +43,13 @@ 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)
sftp_client.remote_copy_data(os.path.join(_folder,f"imagen_{_file_name}_{_datetime}.png"), os.path.join(_remote_folder,f"imagen_{_file_name}_{_datetime}.png"))
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)
sftp_client.remote_copy_data(os.path.join(_folder,f"imagen_{_file_name}_{i+1}_{_datetime}.png"), os.path.join(_remote_folder,f"imagen_{_file_name}_{i+1}_{_datetime}.png"))
return len(images.images)
def imagen_generate_image_data(prompt,download_count=1):
@@ -76,7 +78,7 @@ def imagen_generate_image_path(image_prompt):
QUERY = "query"
create_time = D.date_file_name()
folder_name = os.path.join(OUTPUT_FOLDER,f"{MODEL}_{QUERY}_{image_prompt}_{create_time}")
folder_name = os.path.join(TEMP_FOLDER,f"{MODEL}_{QUERY}_{image_prompt}_{create_time}")
if not os.path.exists(folder_name):
os.makedirs(folder_name)