Update app.py

app.py

@@ -19,7 +19,7 @@ from torch import Tensor, nn
 from transformers import CLIPTextModel, CLIPTokenizer
 from transformers import T5EncoderModel, T5Tokenizer
 from transformers import pipeline, AutoTokenizer, MarianMTModel
-
+from transformers import MarianMTModel, MarianTokenizer
 
 class HFEmbedder(nn.Module):
     def __init__(self, version: str, max_length: int, **hf_kwargs):
@@ -778,7 +778,6 @@ TRANSLATORS = {
 # 번역기 캐시 딕셔너리
 translators_cache = {}
 
-# 번역기 초기화 부분 수정
 def get_translator(lang):
     if lang == "English":
         return None
@@ -786,23 +785,39 @@ def get_translator(lang):
     if lang not in translators_cache:
         try:
             model_name = TRANSLATORS[lang]
-            tokenizer = AutoTokenizer.from_pretrained(model_name)
-            model = MarianMTModel.from_pretrained(
-                model_name,
-                torch_dtype=torch.float32,  # float16 대신 float32 사용
-                low_cpu_mem_usage=True
-            ).to("cpu")  # 명시적으로 CPU 지정
+            tokenizer = MarianTokenizer.from_pretrained(model_name)
+            model = MarianMTModel.from_pretrained(model_name).to("cpu")
             
-            translators_cache[lang] = {
-                "model": model,
-                "tokenizer": tokenizer
-            }
+            translators_cache[lang] = pipeline(
+                "translation",
+                model=model,
+                tokenizer=tokenizer,
+                device="cpu"
+            )
         except Exception as e:
             print(f"Error loading translator for {lang}: {e}")
             return None
             
     return translators_cache[lang]
 
+def translate_prompt(prompt, source_lang):
+    if source_lang == "English":
+        return prompt
+    
+    translator = get_translator(source_lang)
+    if translator is None:
+        print(f"No translator available for {source_lang}, using original prompt")
+        return prompt
+        
+    try:
+        translated = translator(prompt, max_length=512)[0]['translation_text']
+        print(f"Original: {prompt}")
+        print(f"Translated from {source_lang}: {translated}")
+        return translated
+    except Exception as e:
+        print(f"Translation error: {e}")
+        return prompt
+
 def translate_text(text, translator_info):
     if translator_info is None:
         return text
@@ -823,23 +838,6 @@ def translate_text(text, translator_info):
         print(f"Translation error: {e}")
         return text
 
-def translate_prompt(prompt, source_lang):
-    if source_lang == "English":
-        return prompt
-    
-    translator = get_translator(source_lang)
-    if translator is None:
-        print(f"No translator available for {source_lang}, using original prompt")
-        return prompt
-        
-    try:
-        translated = translate_text(prompt, translator)
-        print(f"Translated from {source_lang}: {translated}")
-        return translated
-    except Exception as e:
-        print(f"Translation error: {e}")
-        return prompt
-
 @spaces.GPU
 @torch.no_grad()
 def generate_image(
@@ -847,13 +845,19 @@ def generate_image(
     do_img2img, init_image, image2image_strength, resize_img,
     progress=gr.Progress(track_tqdm=True),
 ):
-    try:
-        translated_prompt = translate_prompt(prompt, source_lang)
-        print(f"Using prompt: {translated_prompt}")
-    except Exception as e:
-        print(f"Translation failed: {e}")
+    # 번역 처리
+    if source_lang != "English":
+        try:
+            translated_prompt = translate_prompt(prompt, source_lang)
+            print(f"Using translated prompt: {translated_prompt}")
+        except Exception as e:
+            print(f"Translation failed: {e}")
+            translated_prompt = prompt
+    else:
         translated_prompt = prompt
 
+
+
     
     if seed == 0:
         seed = int(random.random() * 1000000)
@@ -910,11 +914,10 @@ footer {
     visibility: hidden;
 }
 """
-
-def create_demo():
+ef create_demo():
     with gr.Blocks(theme="Yntec/HaleyCH_Theme_Orange", css=css) as demo:
-    gr.Markdown("# FLUXllama Multilingual(36 Languages")  
-    
+        gr.Markdown("# Multilingual FLUX (36 Languages Support)")
+        
         with gr.Row():
             with gr.Column():
                 source_lang = gr.Dropdown(
@@ -923,6 +926,8 @@ def create_demo():
                     label="Source Language"
                 )
                 
+
+                
                 prompt = gr.Textbox(
                     label="Prompt",
                     value="A cute and fluffy golden retriever puppy sitting upright..."
@@ -953,16 +958,16 @@ def create_demo():
                 output_seed = gr.Text(label="Used Seed")
                 translated_prompt = gr.Text(label="Translated Prompt")
 
-        # Examples 섹션 수정
+                
         examples = [
-            # English examples
-            ["A majestic dragon soaring through clouds at sunset", "English", 768, 768, 3.5, 30, 0, False, None, 0.8, True],
-            # Korean examples
+            # English
+            ["A beautiful sunset over mountains", "English", 768, 768, 3.5, 30, 0, False, None, 0.8, True],
+            # Korean
             ["벚꽃이 흩날리는 서울의 봄 풍경", "Korean", 768, 768, 3.5, 30, 0, False, None, 0.8, True],
-            # Spanish examples
-            ["Una colorida fiesta mexicana con piñatas", "Spanish", 768, 768, 3.5, 30, 0, False, None, 0.8, True],
-            # Chinese examples
-            ["古老的江南水乡，小桥流水人家", "Chinese", 768, 768, 3.5, 30, 0, False, None, 0.8, True]
+            # Spanish
+            ["Un hermoso atardecer en la playa", "Spanish", 768, 768, 3.5, 30, 0, False, None, 0.8, True],
+            # Chinese
+            ["美丽的中国长城日落景色", "Chinese", 768, 768, 3.5, 30, 0, False, None, 0.8, True]
         ]
 
         gr.Examples(
@@ -973,9 +978,11 @@ def create_demo():
             ],
             outputs=[output_image, output_seed, translated_prompt],
             fn=generate_image,
-            cache_examples=True
+            cache_examples=False
         )
 
+
+
         do_img2img.change(
             fn=lambda x: [gr.update(visible=x), gr.update(visible=x), gr.update(visible=x)],
             inputs=[do_img2img],