Update app.py

app.py

@@ -22,7 +22,8 @@ from torch import Tensor, nn
 from transformers import CLIPTextModel, CLIPTokenizer
 from transformers import T5EncoderModel, T5Tokenizer
 # from optimum.quanto import freeze, qfloat8, quantize
-from transformers import pipeline
+
+from transformers import AutoModelForSeq2SeqGeneration, AutoTokenizer, pipeline
 
 class HFEmbedder(nn.Module):
     def __init__(self, version: str, max_length: int, **hf_kwargs):
@@ -779,31 +780,36 @@ TRANSLATORS = {
     "Austroasiatic": "Helsinki-NLP/opus-mt-aav-en"
 }
 
-# 번역기 캐시 딕셔너리
 translators_cache = {}
 
-def initialize_translators():
-    global translators_cache
-    translators_cache = {}
-    
-    for lang, model_name in TRANSLATORS.items():
+def get_translator(lang):
+    """단일 번역기를 초기화하고 반환하는 함수"""
+    if lang == "English":
+        return None
+        
+    if lang not in translators_cache:
         try:
+            model_name = TRANSLATORS[lang]
             translator = pipeline(
-                task="translation",
+                "translation",
                 model=model_name,
-                device="cpu"  # CPU에서 실행
+                device="cpu",
+                framework="pt"
             )
             translators_cache[lang] = translator
             print(f"Successfully loaded translator for {lang}")
         except Exception as e:
             print(f"Error loading translator for {lang}: {e}")
             translators_cache[lang] = None
+            
+    return translators_cache[lang]
 
 def translate_prompt(prompt, source_lang):
+    """프롬프트를 번역하는 함수"""
     if source_lang == "English":
         return prompt
-        
-    translator = translators_cache.get(source_lang)
+    
+    translator = get_translator(source_lang)
     if translator is None:
         print(f"No translator available for {source_lang}, using original prompt")
         return prompt
@@ -817,6 +823,8 @@ def translate_prompt(prompt, source_lang):
     except Exception as e:
         print(f"Translation error for {source_lang}: {e}")
         return prompt
+
+
         
 def get_translator(lang):
     if lang == "English":
@@ -869,12 +877,16 @@ def generate_image(
 ):
     # 번역 처리
     try:
-        translated_prompt = translate_prompt(prompt, source_lang)
-        print(f"Final prompt to use: {translated_prompt}")
+        if source_lang != "English":
+            translated_prompt = translate_prompt(prompt, source_lang)
+            print(f"Using translated prompt: {translated_prompt}")
+        else:
+            translated_prompt = prompt
     except Exception as e:
         print(f"Translation failed: {e}")
         translated_prompt = prompt
-    
+        
+
         
     if seed == 0:
         seed = int(random.random() * 1000000)
@@ -931,10 +943,11 @@ footer {
     visibility: hidden;
 }
 """
+
 def create_demo():
     with gr.Blocks(theme="Yntec/HaleyCH_Theme_Orange", css=css) as demo:
-        gr.Markdown("# Multilingual FLUXllma(36 languages support")
-        gr.Markdown("Supported languages: " + ", ".join(["English"] + sorted(list(TRANSLATORS.keys()))))
+        gr.Markdown("# Multilingual FLUX")
+        gr.Markdown("### Supported languages: " + ", ".join(["English"] + sorted(list(TRANSLATORS.keys()))))
         
         with gr.Row():
             with gr.Column():
@@ -948,6 +961,7 @@ def create_demo():
                     label="Prompt",
                     value="A beautiful landscape"
                 )
+                
 
                 
                 width = gr.Slider(minimum=128, maximum=2048, step=64, label="Width", value=768)
@@ -975,6 +989,7 @@ def create_demo():
                 output_seed = gr.Text(label="Used Seed")
                 translated_prompt = gr.Text(label="Translated Prompt")
 
+        # 다국어 예제
         examples = [
             # English
             ["A beautiful sunset over mountains", "English", 768, 768, 3.5, 30, 0, False, None, 0.8, True],
@@ -988,8 +1003,6 @@ def create_demo():
             ["Un hermoso atardecer en la playa", "Spanish", 768, 768, 3.5, 30, 0, False, None, 0.8, True]
         ]
 
-
-
         gr.Examples(
             examples=examples,
             inputs=[
@@ -998,7 +1011,7 @@ def create_demo():
             ],
             outputs=[output_image, output_seed, translated_prompt],
             fn=generate_image,
-            cache_examples=False
+            cache_examples=True
         )
 
 
@@ -1021,8 +1034,6 @@ def create_demo():
     return demo
 
 if __name__ == "__main__":
-    print("Initializing translators...")
-    initialize_translators()  # 번역기 초기화
     print("Starting demo...")
     demo = create_demo()
     demo.launch(share=True)