Update app.py

app.py

@@ -142,9 +142,23 @@ def get_model() -> tuple[MMAudio, FeaturesUtils, SequenceConfig]:
 
 net, feature_utils, seq_cfg = get_model()
 
+
+# search_videos 함수 수정
+@torch.no_grad()
+def search_videos(query):
+    try:
+        # CPU에서 번역 실행
+        query = translate_prompt(query)
+        return search_pixabay_videos(query, PIXABAY_API_KEY)
+    except Exception as e:
+        logging.error(f"Video search error: {e}")
+        return []
+
+# translate_prompt 함수도 수정
 def translate_prompt(text):
     try:
         if text and any(ord(char) >= 0x3131 and ord(char) <= 0xD7A3 for char in text):
+            # CPU에서 번역 실행
             with torch.no_grad():
                 translation = translator(text)[0]['translation_text']
             return translation
@@ -153,6 +167,22 @@ def translate_prompt(text):
         logging.error(f"Translation error: {e}")
         return text
 
+# 디바이스 설정 부분 수정
+if torch.cuda.is_available():
+    device = torch.device("cuda")
+    torch.backends.cuda.matmul.allow_tf32 = True
+    torch.backends.cudnn.allow_tf32 = True
+    torch.backends.cudnn.benchmark = True
+else:
+    device = torch.device("cpu")
+
+# 번역기 설정 수정
+translator = pipeline("translation", 
+                     model="Helsinki-NLP/opus-mt-ko-en", 
+                     device="cpu")  # 명시적으로 CPU 지정
+
+
+
 def search_pixabay_videos(query, api_key):
     try:
         base_url = "https://pixabay.com/api/videos/"
@@ -171,11 +201,6 @@ def search_pixabay_videos(query, api_key):
         logging.error(f"Pixabay API error: {e}")
         return []
 
-@torch.no_grad()
-def search_videos(query):
-    with torch.cuda.device("cpu"):
-        query = translate_prompt(query)
-        return search_pixabay_videos(query, PIXABAY_API_KEY)
 
 @spaces.GPU
 @torch.inference_mode()