Update app.py

app.py

@@ -8,7 +8,6 @@ import torch
 from PIL import Image
 import os
 import re
-import openai  # OpenAI 추가
 
 # 경로 및 설정
 CLIP_PATH = "google/siglip-so400m-patch14-384"
@@ -18,7 +17,6 @@ CHECKPOINT_PATH = Path("wpkklhc6")
 TITLE = "<h1><center>JoyCaption Pre-Alpha (2024-07-30a)</center></h1>"
 
 HF_TOKEN = os.environ.get("HF_TOKEN", None)
-openai.api_key = os.getenv("OPENAI_API_KEY")  # OpenAI API 키 설정
 
 # 이미지 어댑터 정의
 class ImageAdapter(nn.Module):
@@ -83,8 +81,19 @@ def stream_chat(input_image: Image.Image):
     embedded_bos = text_model.model.embed_tokens(torch.tensor([[tokenizer.bos_token_id]], device='cuda', dtype=torch.int64))
 
     # 프롬프트 구성
-    inputs_embeds = torch.cat([embedded_bos.expand(embedded_images.shape[0], -1, -1), embedded_images, prompt_embeds], dim=1)
-    input_ids = torch.cat([torch.tensor([[tokenizer.bos_token_id]], dtype=torch.long), torch.zeros((1, embedded_images.shape[1]), dtype=torch.long), prompt], dim=1).to('cuda')
+    inputs_embeds = torch.cat([
+        embedded_bos.expand(embedded_images.shape[0], -1, -1),
+        embedded_images,
+        prompt_embeds
+    ], dim=1)
+
+    # CPU에 있는 텐서를 GPU로 이동
+    input_ids = torch.cat([
+        torch.tensor([[tokenizer.bos_token_id]], dtype=torch.long).to('cuda'),
+        torch.zeros((1, embedded_images.shape[1]), dtype=torch.long).to('cuda'),
+        prompt.to('cuda')
+    ], dim=1)
+
     attention_mask = torch.ones_like(input_ids)
 
     # 텍스트 생성
@@ -140,10 +149,10 @@ def replace_gender_specific_words(caption, gender_prefix):
     if gender_prefix == "Korean man":
         caption = re.sub(r'\bwoman\b', "man", caption, flags=re.IGNORECASE)
         caption = re.sub(r'\bgirl\b', "boy", caption, flags=re.IGNORECASE)
-        caption = re.sub(r'\blady\b', "gentleman", flags=re.IGNORECASE)
-        caption = re.sub(r'\bshe\b', "he", flags=re.IGNORECASE)
-        caption = re.sub(r'\bher\b', "his", flags=re.IGNORECASE)
-        caption = re.sub(r'\bherself\b', "himself", flags=re.IGNORECASE)
+        caption = re.sub(r'\blady\b', "gentleman", caption, flags=re.IGNORECASE)
+        caption = re.sub(r'\bshe\b', "he", caption, flags=re.IGNORECASE)
+        caption = re.sub(r'\bher\b', "his", caption, flags=re.IGNORECASE)
+        caption = re.sub(r'\bherself\b', "himself", caption, flags=re.IGNORECASE)
     elif gender_prefix == "Korean woman":
         caption = re.sub(r'\bman\b', "woman", caption, flags=re.IGNORECASE)
         caption = re.sub(r'\bboy\b', "girl", caption, flags=re.IGNORECASE)
@@ -167,20 +176,6 @@ def replace_gender_words(caption, gender, age, hair_length, hair_style, hair_col
     caption = replace_gender_specific_words(caption, gender_prefix)
     return f"{gender_prefix}, age {age}, {hair_description}: {caption}"
 
-# GPT-4o mini API 호출 함수 추가
-def call_gpt4o_api(content, system_message, max_tokens, temperature, top_p):
-    response = openai.ChatCompletion.create(
-        model="gpt-4o-mini",  # OpenAI GPT-4o Mini 모델 사용
-        messages=[
-            {"role": "system", "content": system_message},
-            {"role": "user", "content": content},
-        ],
-        max_tokens=max_tokens,
-        temperature=temperature,
-        top_p=top_p,
-    )
-    return response.choices[0].message['content']
-
 # Recaption 함수
 def recaption(input_image: Image.Image, prefix: str, age: int, hair_length: str, hair_style: str, hair_color: str, hair_accessory: str):
     original_caption = stream_chat(input_image)