Update app.py

app.py

@@ -24,9 +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):
         super().__init__()
@@ -749,8 +746,12 @@ 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"))
+
+
+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")
+
 
 
 @spaces.GPU
@@ -762,14 +763,17 @@ 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)
+    
+    # 한글, 일본어, 중국어가 있으면 번역
     if contains_korean(prompt):
         translated_prompt = ko_translator(prompt, max_length=512)[0]['translation_text']
         print(f"Translated Korean prompt: {translated_prompt}")
@@ -778,6 +782,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
+
+
+
+    
 
     if seed == 0:
         seed = int(random.random() * 1000000)