Update app.py

app.py

@@ -1,5 +1,6 @@
-import os
+# import os
 import spaces
+
 import time
 import gradio as gr
 import torch
@@ -8,9 +9,11 @@ from torchvision import transforms
 from dataclasses import dataclass
 import math
 from typing import Callable
+
 from tqdm import tqdm
 import bitsandbytes as bnb
 from bitsandbytes.nn.modules import Params4bit, QuantState
+
 import torch
 import random
 from einops import rearrange, repeat
@@ -18,8 +21,8 @@ from diffusers import AutoencoderKL
 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
+# from optimum.quanto import freeze, qfloat8, quantize
+from transformers import pipeline
 
 class HFEmbedder(nn.Module):
     def __init__(self, version: str, max_length: int, **hf_kwargs):
@@ -737,6 +740,7 @@ model_zero_init = False
 
 # 언어-모델 매핑 딕셔너리 추가
 TRANSLATORS = {
+    
     "Korean": "Helsinki-NLP/opus-mt-ko-en",
     "Japanese": "Helsinki-NLP/opus-mt-ja-en", 
     "Chinese": "Helsinki-NLP/opus-mt-zh-en",
@@ -778,6 +782,42 @@ TRANSLATORS = {
 # 번역기 캐시 딕셔너리
 translators_cache = {}
 
+def initialize_translators():
+    global translators_cache
+    translators_cache = {}
+    
+    for lang, model_name in TRANSLATORS.items():
+        try:
+            translator = pipeline(
+                task="translation",
+                model=model_name,
+                device="cpu"  # CPU에서 실행
+            )
+            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
+
+def translate_prompt(prompt, source_lang):
+    if source_lang == "English":
+        return prompt
+        
+    translator = translators_cache.get(source_lang)
+    if translator is None:
+        print(f"No translator available for {source_lang}, using original prompt")
+        return prompt
+    
+    try:
+        translation = translator(prompt, max_length=512)
+        translated_text = translation[0]['translation_text']
+        print(f"Original ({source_lang}): {prompt}")
+        print(f"Translated: {translated_text}")
+        return translated_text
+    except Exception as e:
+        print(f"Translation error for {source_lang}: {e}")
+        return prompt
+        
 def get_translator(lang):
     if lang == "English":
         return None
@@ -800,24 +840,6 @@ def get_translator(lang):
             
     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
@@ -846,19 +868,14 @@ def generate_image(
     progress=gr.Progress(track_tqdm=True),
 ):
     # 번역 처리
-    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:
+    try:
+        translated_prompt = translate_prompt(prompt, source_lang)
+        print(f"Final prompt to use: {translated_prompt}")
+    except Exception as e:
+        print(f"Translation failed: {e}")
         translated_prompt = prompt
-
-
-
     
+        
     if seed == 0:
         seed = int(random.random() * 1000000)
         
@@ -914,10 +931,10 @@ footer {
     visibility: hidden;
 }
 """
-
 def create_demo():
     with gr.Blocks(theme="Yntec/HaleyCH_Theme_Orange", css=css) as demo:
-        gr.Markdown("# Multilingual FLUXllama (36 Languages Support)")
+        gr.Markdown("# Multilingual FLUX")
+        gr.Markdown("Supported languages: " + ", ".join(["English"] + sorted(list(TRANSLATORS.keys()))))
         
         with gr.Row():
             with gr.Column():
@@ -927,12 +944,11 @@ def create_demo():
                     label="Source Language"
                 )
                 
-
-                
                 prompt = gr.Textbox(
                     label="Prompt",
-                    value="A cute and fluffy golden retriever puppy sitting upright..."
+                    value="A beautiful landscape"
                 )
+
                 
                 width = gr.Slider(minimum=128, maximum=2048, step=64, label="Width", value=768)
                 height = gr.Slider(minimum=128, maximum=2048, step=64, label="Height", value=768)
@@ -959,18 +975,21 @@ 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],
             # Korean
             ["벚꽃이 흩날리는 서울의 봄 풍경", "Korean", 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],
+            # Japanese
+            ["富士山と桜の美しい風景", "Japanese", 768, 768, 3.5, 30, 0, False, None, 0.8, True],
             # Chinese
-            ["美丽的中国长城日落景色", "Chinese", 768, 768, 3.5, 30, 0, False, None, 0.8, True]
+            ["长城日落的壮丽景色", "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]
         ]
 
+
+
         gr.Examples(
             examples=examples,
             inputs=[
@@ -1002,5 +1021,8 @@ def create_demo():
     return demo
 
 if __name__ == "__main__":
+    print("Initializing translators...")
+    initialize_translators()  # 번역기 초기화
+    print("Starting demo...")
     demo = create_demo()
-    demo.launch(share=True)  # share=True 추가
+    demo.launch(share=True)