Update app.py

app.py

@@ -24,8 +24,6 @@ from transformers import T5EncoderModel, T5Tokenizer
 # from optimum.quanto import freeze, qfloat8, quantize
 from transformers import pipeline
 
-ko_translator = pipeline("translation", model="Helsinki-NLP/opus-mt-ko-en")
-ja_translator = pipeline("translation", model="Helsinki-NLP/opus-mt-ja-en")
 
 class HFEmbedder(nn.Module):
     def __init__(self, version: str, max_length: int, **hf_kwargs):
@@ -749,9 +747,6 @@ model = Flux().to(dtype=torch.bfloat16, device="cuda")
 result = model.load_state_dict(sd)
 model_zero_init = False
 
-# model = Flux().to(dtype=torch.bfloat16, device="cuda")
-# result = model.load_state_dict(load_file("/storage/dev/nyanko/flux-dev/flux1-dev.sft"))
-
 
 @spaces.GPU
 @torch.no_grad()
@@ -762,14 +757,28 @@ def generate_image(
 ):
     translated_prompt = prompt
     
-    # 한글 또는 일본어 문자 감지
+    # 한글, 일본어, 중국어, 스페인어 문자 감지
     def contains_korean(text):
         return any('\u3131' <= c <= '\u318E' or '\uAC00' <= c <= '\uD7A3' for c in text)
         
     def contains_japanese(text):
         return any('\u3040' <= c <= '\u309F' or '\u30A0' <= c <= '\u30FF' or '\u4E00' <= c <= '\u9FFF' for c in text)
     
-    # 한글이나 일본어가 있으면 번역
+    def contains_chinese(text):
+        return any('\u4e00' <= c <= '\u9fff' for c in text)
+    
+    def contains_spanish(text):
+        # 스페인어 특수 문자 포함 확인
+        spanish_chars = set('áéíóúüñ¿¡')
+        return any(c in spanish_chars for c in text.lower())
+
+    # 번역기 추가
+    ko_translator = pipeline("translation", model="Helsinki-NLP/opus-mt-ko-en")
+    ja_translator = pipeline("translation", model="Helsinki-NLP/opus-mt-ja-en")    
+    zh_translator = pipeline("translation", model="Helsinki-NLP/opus-mt-zh-en")
+    es_translator = pipeline("translation", model="Helsinki-NLP/opus-mt-es-en")
+    
+    # 각 언어 감지 후 번역
     if contains_korean(prompt):
         translated_prompt = ko_translator(prompt, max_length=512)[0]['translation_text']
         print(f"Translated Korean prompt: {translated_prompt}")
@@ -778,6 +787,14 @@ def generate_image(
         translated_prompt = ja_translator(prompt, max_length=512)[0]['translation_text'] 
         print(f"Translated Japanese prompt: {translated_prompt}")
         prompt = translated_prompt
+    elif contains_chinese(prompt):
+        translated_prompt = zh_translator(prompt, max_length=512)[0]['translation_text']
+        print(f"Translated Chinese prompt: {translated_prompt}")
+        prompt = translated_prompt
+    elif contains_spanish(prompt):
+        translated_prompt = es_translator(prompt, max_length=512)[0]['translation_text']
+        print(f"Translated Spanish prompt: {translated_prompt}")
+        prompt = translated_prompt
 
     if seed == 0:
         seed = int(random.random() * 1000000)
@@ -785,7 +802,6 @@ def generate_image(
     device = "cuda" if torch.cuda.is_available() else "cpu"
     torch_device = torch.device(device)
 
-
     
     global model, model_zero_init
     if not model_zero_init: