Update app.py

app.py

@@ -805,21 +805,20 @@ def get_translator(lang):
     if lang not in translators_cache:
         try:
             model_name = TRANSLATORS[lang]
-            cache_dir = download_model(model_name)
             
-            if cache_dir is None:
-                return None
-                
-            tokenizer = MarianTokenizer.from_pretrained(model_name, cache_dir=cache_dir)
-            model = MarianMTModel.from_pretrained(model_name, cache_dir=cache_dir)
+            # pipeline 사용 대신 직접 모델 로드
+            tokenizer = MarianTokenizer.from_pretrained(model_name)
+            model = MarianMTModel.from_pretrained(model_name)
             
-            model = model.to("cpu")
+            # CPU에서 실행
+            model = model.to("cpu").eval()
             
             translators_cache[lang] = {
                 "model": model,
                 "tokenizer": tokenizer
             }
             print(f"Successfully loaded translator for {lang}")
+            
         except Exception as e:
             print(f"Error loading translator for {lang}: {e}")
             return None
@@ -835,18 +834,26 @@ def translate_text(text, translator_info):
         tokenizer = translator_info["tokenizer"]
         model = translator_info["model"]
         
+        # 입력 텍스트 전처리
         inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True, max_length=512)
-        translated = model.generate(**inputs)
-        result = tokenizer.decode(translated[0], skip_special_tokens=True)
+        
+        # 번역 수행
+        with torch.no_grad():
+            outputs = model.generate(**inputs)
+            
+        # 번역 결과 디코딩
+        translated = tokenizer.decode(outputs[0], skip_special_tokens=True)
         
         print(f"Original text: {text}")
-        print(f"Translated text: {result}")
+        print(f"Translated text: {translated}")
+        
+        return translated
         
-        return result
     except Exception as e:
         print(f"Translation error: {e}")
         return text
 
+
 @spaces.GPU
 @torch.no_grad()
 def generate_image(
@@ -854,17 +861,22 @@ def generate_image(
     do_img2img, init_image, image2image_strength, resize_img,
     progress=gr.Progress(track_tqdm=True),
 ):
+    # 번역 처리
     try:
         if source_lang != "English":
             translator_info = get_translator(source_lang)
-            translated_prompt = translate_text(prompt, translator_info)
-            print(f"Using translated prompt: {translated_prompt}")
+            if translator_info is not None:
+                translated_prompt = translate_text(prompt, translator_info)
+                print(f"Using translated prompt: {translated_prompt}")
+            else:
+                print(f"No translator available for {source_lang}, using original prompt")
+                translated_prompt = prompt
         else:
             translated_prompt = prompt
     except Exception as e:
         print(f"Translation failed: {e}")
         translated_prompt = prompt
-        
+