Merge pull request #74 from 2Swon/main

add generate img

app/api/api_v1/endpoints/gree.py

@@ -16,6 +16,7 @@
 import uuid
 
 from app.schemas.greeFileDto import GreeFileSchema
+from app.services.image_service import create_image, check_image_status, upload_images_to_azure
 from app.services.upload_service import upload_file_to_azure, upload_greefile_to_azure, \
     upload_yaml_to_azure_blob, upload_gif_to_azure_blob
 from app.api.api_v1.endpoints.user import get_current_user
@@ -56,6 +57,39 @@ class SuccessMessage(BaseModel):
     message: str
 
 
+@router.post("/generate-and-upload-image/{gree_id}")
+async def generate_and_upload_image(
+        gree_id: int,
+        promptSelect: int,
+        current_user: Member = Depends(get_current_user),
+        db: AsyncSession = Depends(get_db)):
+
+    gree = await crud_get_gree_by_id(db, gree_id=gree_id, user_id=current_user.id)
+    if not gree:
+        raise HTTPException(status_code=404, detail="Gree not found")
+
+    try:
+        creation_response = await create_image(promptSelect, gree.raw_img)
+        if 'data' not in creation_response or 'id' not in creation_response['data']:
+            raise HTTPException(status_code=500, detail="Failed to initiate image creation.")
+    except Exception as e:
+        # 여기서 발생하는 예외는 create_image 함수 내부에서 발생한 예외를 처리합니다.
+        raise HTTPException(status_code=500, detail=f"Image creation request failed: {str(e)}")
+
+    image_data = await check_image_status(creation_response['data']['id'])
+
+    local_original_image_path_list = []
+    for image_url in image_data:
+        unique_filename = f"{uuid.uuid4()}.png"
+        local_path = f"temp/{unique_filename}"
+        await download_image_async(image_url, local_path)
+        local_original_image_path_list.append(local_path)
+
+    uploaded_urls = await upload_images_to_azure(local_original_image_path_list)
+
+    return {"uploaded_image_urls": uploaded_urls, "message": "Images uploaded successfully."}
+
+
 @router.put('/update/{gree_id}', response_model=SuccessMessage)
 async def update_gree(gree_id: int, gree_update: GreeUpdate, current_user: Member = Depends(get_current_user),
                       db: AsyncSession = Depends(get_db)):

app/services/image_service.py

@@ -0,0 +1,85 @@
+import uuid
+
+import aiofiles
+import httpx
+import asyncio
+
+from azure.storage.blob import ContentSettings
+from azure.storage.blob.aio import BlobServiceClient
+from typing import List
+import os
+
+from fastapi import HTTPException
+
+MID_API_KEY = os.getenv('MID_API_KEY')
+container_name = "greefile"
+AZURE_ACCOUNT_KEY = os.getenv("AZURE_ACCOUNT_KEY")
+connection_string = f'DefaultEndpointsProtocol=https;AccountName=greedotstorage;AccountKey={AZURE_ACCOUNT_KEY};EndpointSuffix=core.windows.net'
+
+promptDict = {
+    1: """
+    please make cute
+    This picture was painted by a child.
+    Please change this picture a little bit more cute.
+    I want to keep most of the original.
+    Arms and legs must be in the form of characters.
+    Please don't put your face and body in.
+    I just wish I had one character.
+    And the parts other than the character outline must be a white background, so please make this important.
+    """,
+    2: """
+    please make cool
+    This picture was painted by a child.
+    Please change this picture a little bit more cute.
+    I want to keep most of the original.
+    Arms and legs must be in the form of characters.
+    Please don't put your face and body in.
+    I just wish I had one character.
+    And the parts other than the character outline must be a white background, so please make this important.
+    """
+}
+
+async def create_image(promptSelect: int, raw_img_url: str) -> dict:
+    prompt = promptDict[promptSelect] + "\n" + raw_img_url
+    headers = {
+        'Authorization': f'Bearer {MID_API_KEY}',
+        'Content-Type': 'application/json'
+    }
+    data = {"prompt": prompt}
+    async with httpx.AsyncClient() as client:
+        response = await client.post('https://cl.imagineapi.dev/items/images/', json=data, headers=headers)
+        return response.json()
+
+async def check_image_status(image_id: str) -> list:
+    headers = {'Authorization': f'Bearer {MID_API_KEY}'}
+    async with httpx.AsyncClient() as client:
+        while True:
+            response = await client.get(f'https://cl.imagineapi.dev/items/images/{image_id}', headers=headers)
+            data = response.json()
+            if data['data']['status'] == 'completed':
+                if 'upscaled_urls' in data['data']:
+                    return data['data']['upscaled_urls']
+                else:
+                    # Handle case where 'upscaled_urls' is missing even though status is 'completed'
+                    raise HTTPException(status_code=500, detail="Image creation completed but did not return URLs.")
+            elif data['data']['status'] == 'failed':
+                # Handle failed image creation
+                raise HTTPException(status_code=500, detail="Image creation failed.")
+            await asyncio.sleep(5)  # Polling interval
+
+
+async def upload_images_to_azure(image_paths: List[str]) -> List[str]:
+    blob_service_client = BlobServiceClient.from_connection_string(connection_string)
+    uploaded_urls = []
+    async with blob_service_client:
+        for image_path in image_paths:
+            filename = os.path.basename(image_path) or str(uuid.uuid4())
+            blob_name = f"upload/{filename}.png"
+            blob_client = blob_service_client.get_blob_client(container=container_name, blob=blob_name)
+
+            async with aiofiles.open(image_path, 'rb') as file:
+                data = await file.read()
+
+            await blob_client.upload_blob(data, overwrite=True, content_settings=ContentSettings(content_type='image/png'))
+            uploaded_urls.append(blob_client.url)
+    return uploaded_urls