Update app.py

app.py

@@ -19,6 +19,7 @@ from torch import Tensor, nn
 from transformers import CLIPTextModel, CLIPTokenizer
 from transformers import T5EncoderModel, T5Tokenizer
 from transformers import MarianMTModel, MarianTokenizer, pipeline
+from huggingface_hub import snapshot_download
 
 class HFEmbedder(nn.Module):
     def __init__(self, version: str, max_length: int, **hf_kwargs):
@@ -777,20 +778,41 @@ TRANSLATORS = {
 
 translators_cache = {}
 
-from transformers import MarianMTModel, MarianTokenizer
+
+
+# 모델 캐시 디렉토리 설정
+os.environ['TRANSFORMERS_CACHE'] = '/tmp/transformers_cache'
+
+def download_model(model_name):
+    """모델을 미리 다운로드"""
+    try:
+        cache_dir = os.path.join('/tmp/transformers_cache', model_name.split('/')[-1])
+        snapshot_download(
+            repo_id=model_name,
+            cache_dir=cache_dir,
+            local_files_only=False
+        )
+        return cache_dir
+    except Exception as e:
+        print(f"Error downloading model {model_name}: {e}")
+        return None
 
 def get_translator(lang):
-    """단일 번역기를 초기화하고 반환하는 함수"""
+    """번역기 초기화 및 반환"""
     if lang == "English":
         return None
         
     if lang not in translators_cache:
         try:
             model_name = TRANSLATORS[lang]
-            tokenizer = MarianTokenizer.from_pretrained(model_name)
-            model = MarianMTModel.from_pretrained(model_name)
+            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)
             
-            # CPU에서 실행
             model = model.to("cpu")
             
             translators_cache[lang] = {
@@ -804,35 +826,26 @@ def get_translator(lang):
             
     return translators_cache[lang]
 
-def translate_prompt(prompt, source_lang):
-    """프롬프트를 번역하는 함수"""
-    if source_lang == "English":
-        return prompt
-    
-    translator_info = get_translator(source_lang)
+def translate_text(text, translator_info):
+    """번역 수행"""
     if translator_info is None:
-        print(f"No translator available for {source_lang}, using original prompt")
-        return prompt
-    
+        return text
+        
     try:
         tokenizer = translator_info["tokenizer"]
         model = translator_info["model"]
         
-        # 텍스트를 토큰화하고 모델 입력으로 변환
-        inputs = tokenizer(prompt, return_tensors="pt", padding=True, truncation=True, max_length=512)
-        
-        # 번역 수행
+        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)
         
-        # 번역된 텍스트 디코딩
-        translated_text = tokenizer.decode(translated[0], skip_special_tokens=True)
+        print(f"Original text: {text}")
+        print(f"Translated text: {result}")
         
-        print(f"Original ({source_lang}): {prompt}")
-        print(f"Translated: {translated_text}")
-        return translated_text
+        return result
     except Exception as e:
-        print(f"Translation error for {source_lang}: {e}")
-        return prompt
+        print(f"Translation error: {e}")
+        return text
 
 @spaces.GPU
 @torch.no_grad()
@@ -843,13 +856,15 @@ def generate_image(
 ):
     try:
         if source_lang != "English":
-            translated_prompt = translate_prompt(prompt, source_lang)
+            translator_info = get_translator(source_lang)
+            translated_prompt = translate_text(prompt, translator_info)
             print(f"Using translated prompt: {translated_prompt}")
         else:
             translated_prompt = prompt
     except Exception as e:
         print(f"Translation failed: {e}")
         translated_prompt = prompt
+