Merge pull request #25 from MoDeep11/image_generate_api_finished

Image generate api finished

Author: kimwonshin

ai/image/image_generator.py

@@ -18,7 +18,6 @@
 from boto3 import client
 import io
 from PIL import Image
-import uuid
 
 dotenv.load_dotenv()  # .env 파일에서 환경 변수 로드
 AWS_ACCESS_KEY = os.getenv("AWS_ACCESS_KEY")
@@ -83,6 +82,7 @@ class ImageGenerationRequest(BaseModel):
     userId: int
     emotion: str
     content: str
+    taskId: str
 
 
 class ImageGenerator:
@@ -132,12 +132,10 @@ async def generate(self, request: ImageGenerationRequest) -> dict:
                 ),
             )
             os.makedirs(TEMP_DIR, exist_ok=True)
-            id1 = uuid.uuid4()
-            id2 = uuid.uuid4()
-            local_file = os.path.join(TEMP_DIR, f"generated_image_{id1}_{id2}.png")
+            local_file = os.path.join(TEMP_DIR, f"generated_image_{request.taskId}.png")
             await save_image(response, local_file)
 
-            s3_key = f"ai-gen/images/diary_{id1}_{id2}.png"
+            s3_key = f"ai-gen/images/{request.taskId}.png"
             s3_url = await s3_uploader.upload_file(local_file, s3_key)
             image_urls = [s3_url]
 

ai/image/sticker_generator.py

@@ -55,6 +55,7 @@ class StickerGenerationRequest(BaseModel):
     diaryId: int
     emotion: str
     content: str
+    taskId: str
 
 
 class StickerGenerator:
@@ -123,19 +124,22 @@ async def generate(self, request: StickerGenerationRequest) -> dict:
         )
 
         os.makedirs(TEMP_DIR, exist_ok=True)
-        local_file = os.path.join(TEMP_DIR, f"generated_sticker_{request.diaryId}_{request.userId}_1.png")
+        local_file = os.path.join(TEMP_DIR, f"generated_sticker_{request.taskId}.png")
         await save_image(image_response, local_file)
 
         image_text_parts = [part.text for part in image_response.parts if part.text]
         description = " ".join(image_text_parts)
 
-        s3_key = f"ai-gen/stickers/sticker_{request.diaryId}_{request.userId}_1.png"
+        s3_key = f"ai-gen/stickers/{request.taskId}.png"
         s3_url = await s3_uploader.upload_file(local_file, s3_key)
 
+        # S3 업로드 후 로컬 파일 삭제
+        if os.path.exists(local_file):
+            os.remove(local_file)
+
         stickers = [{
             "imageUrl": s3_url,
             "keyword": description,
-            "localPath": local_file,
         }]
 
         return {

server/app/router/diary_router.py

@@ -159,7 +159,7 @@ async def generate_sticker(request: StickerGenerateRequest, background_tasks: Ba
     - content: 일기 내용
     - count: 생성할 스티커 개수 (기본: 1, 최대: 5)
     """
-    task_id = f"sticker_task_{request.diaryId}_{request.userId}"
+    task_id = str(uuid.uuid4())
 
     # 백그라운드 작업으로 실행
     background_tasks.add_task(
@@ -194,7 +194,7 @@ async def generate_image(request: ImageGenerateRequest, background_tasks: Backgr
     - emotion: 감정 (HAPPY, SAD, ANGRY, ANXIOUS, NORMAL)
     - content: 일기 내용
     """
-    task_id = f"image_task_{request.diaryId}_{request.userId}"
+    task_id = str(uuid.uuid4())
 
     background_tasks.add_task(
         process_image_generation,
@@ -231,7 +231,7 @@ async def process_sticker_generation(
     """스티커 생성 백그라운드 작업"""
     print(f"[Router] process_sticker_generation started (task_id={task_id})")
     try:
-        result = await diary_service.generate_stickers(user_id, diary_id, emotion, content)
+        result = await diary_service.generate_stickers(user_id, diary_id, emotion, content, task_id)
         print(f"[Router] generate_stickers result: {result}")
         # 스티커 생성이 완료되면 Spring 서버에 콜백
         print(f"[Router] calling patch_sticker_callback for task_id={task_id}")
@@ -276,7 +276,7 @@ async def process_image_generation(
     """이미지 생성 백그라운드 작업"""
     print(f"[Router] process_image_generation started (task_id={task_id})")
     try:
-        result = await diary_service.generate_image(user_id, diary_id, emotion, content)
+        result = await diary_service.generate_image(user_id, diary_id, emotion, content, task_id)
         print(f"[Router] generate_image result: {result}")
         # 이미지 생성이 완료되면 Spring 서버에 콜백
         print(f"[Router] calling patch_image_callback for task_id={task_id}")

server/app/service/diary_service.py

@@ -130,23 +130,25 @@ async def create_diary(self, user_info, history, user_audio_urls=None, ai_audio_
                 "tags": []
             }        
 
-    async def generate_image(self, user_id: int, diary_id: int, emotion: str, content: str) -> Dict:
+    async def generate_image(self, user_id: int, diary_id: int, emotion: str, content: str, task_id: str) -> Dict:
         """이미지 생성"""
         request = ImageGenerationRequest(
             diaryId=diary_id,
             userId=user_id,
             emotion=emotion,
             content=content,
+            taskId=task_id,
         )
         return await image_generator.generate(request)
 
-    async def generate_stickers(self, user_id: int, diary_id: int, emotion: str, content: str) -> Dict:
+    async def generate_stickers(self, user_id: int, diary_id: int, emotion: str, content: str, task_id: str) -> Dict:
         """스티커 생성"""
         request = StickerGenerationRequest(
             userId=user_id,
             diaryId=diary_id,
             emotion=emotion,
             content=content,
+            taskId=task_id,
         )
         return await sticker_generator.generate(request)
 
@@ -191,12 +193,20 @@ async def patch_image_callback(self, diary_id: int, task_id: str, user_id: int,
 
     async def patch_sticker_callback(self, diary_id: int, task_id: str, user_id: int, stickers: list):
         """Spring 백엔드에 스티커 생성 완료를 알리는 콜백"""
-        url = f"{self.backend_url}/internal/diary/{diary_id}/ai-image"
+        url = f"{self.backend_url}/internal/v1/stickers"
+        if stickers:
+            image_url = stickers[0]["imageUrl"]
+            keyword = stickers[0]["keyword"]
+        else:
+            image_url = ""
+            keyword = ""
         payload = {
-            "diaryId": diary_id,
             "taskId": task_id,
             "userId": user_id,
-            "stickers": stickers,
+            "diaryId": diary_id,
+            "status": "success",
+            "imageUrl": image_url,
+            "keyword": keyword,
         }
 
         # 비동기 흐름 확인용 로그