Update app.py

app.py

@@ -10,7 +10,6 @@ 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
@@ -25,6 +24,9 @@ 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__()
@@ -747,48 +749,10 @@ 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")
-
-
-from transformers import M2M100ForConditionalGeneration, M2M100Tokenizer
-
-def translate_text(text):
-    try:
-        # M2M100은 다국어 번역을 한 모델로 처리할 수 있습니다
-        model_name = "facebook/m2m100_418M"
-        tokenizer = M2M100Tokenizer.from_pretrained(model_name)
-        model = M2M100ForConditionalGeneration.from_pretrained(model_name).to(device)
-        
-        # 언어 감지
-        def detect_language(text):
-            if any('\u3131' <= c <= '\u318E' or '\uAC00' <= c <= '\uD7A3' for c in text):
-                return 'ko'
-            elif any('\u3040' <= c <= '\u309F' or '\u30A0' <= c <= '\u30FF' for c in text):
-                return 'ja'
-            elif any('\u4e00' <= c <= '\u9fff' for c in text):
-                return 'zh'
-            return None
-        
-        src_lang = detect_language(text)
-        if src_lang is None:
-            return text
-            
-        tokenizer.src_lang = src_lang
-        encoded = tokenizer(text, return_tensors="pt").to(device)
-        generated_tokens = model.generate(
-            **encoded,
-            forced_bos_token_id=tokenizer.get_lang_id("en"),
-            max_length=128
-        )
-        return tokenizer.batch_decode(generated_tokens, skip_special_tokens=True)[0]
-    except Exception as e:
-        print(f"Translation error: {e}")
-        return text
-
 @spaces.GPU
 @torch.no_grad()
 def generate_image(
@@ -796,36 +760,24 @@ def generate_image(
     do_img2img, init_image, image2image_strength, resize_img,
     progress=gr.Progress(track_tqdm=True),
 ):
-    translated_prompt = translate_text(prompt)
-    if translated_prompt != prompt:
-        print(f"Translated prompt: {translated_prompt}")
-    prompt = translated_prompt
-
-    if seed == 0:
-        seed = int(random.random() * 1000000)
-def translate_text(text, src_lang, model_name):
-    try:
-        tokenizer = MarianTokenizer.from_pretrained(model_name)
-        model = MarianMTModel.from_pretrained(model_name)
-        model = model.to(device)
+    translated_prompt = prompt
+    
+    # 한글 또는 일본어 문자 감지
+    def contains_korean(text):
+        return any('\u3131' <= c <= '\u318E' or '\uAC00' <= c <= '\uD7A3' for c in text)
         
-        inputs = tokenizer(text, return_tensors="pt", padding=True).to(device)
-        translated = model.generate(**inputs)
-        translated_text = tokenizer.batch_decode(translated, skip_special_tokens=True)[0]
-        return translated_text
-    except:
-        return text  # 번역 실패시 원본 텍스트 반환
-
-# 기존의 translator 정의 부분을 삭제하고 아래 코드로 대체
-def translate_if_needed(prompt):
-    if any('\u3131' <= c <= '\u318E' or '\uAC00' <= c <= '\uD7A3' for c in prompt):  # Korean
-        return translate_text(prompt, 'ko', 'Helsinki-NLP/opus-mt-ko-en')
-    elif any('\u3040' <= c <= '\u309F' or '\u30A0' <= c <= '\u30FF' for c in prompt):  # Japanese
-        return translate_text(prompt, 'ja', 'Helsinki-NLP/opus-mt-ja-en')
-    elif any('\u4e00' <= c <= '\u9fff' for c in prompt):  # Chinese
-        return translate_text(prompt, 'zh', 'Helsinki-NLP/opus-mt-zh-en')
-    return prompt
-
+    def contains_japanese(text):
+        return any('\u3040' <= c <= '\u309F' or '\u30A0' <= c <= '\u30FF' or '\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}")
+        prompt = translated_prompt
+    elif contains_japanese(prompt):
+        translated_prompt = ja_translator(prompt, max_length=512)[0]['translation_text'] 
+        print(f"Translated Japanese prompt: {translated_prompt}")
+        prompt = translated_prompt
 
     if seed == 0:
         seed = int(random.random() * 1000000)
@@ -888,12 +840,13 @@ footer {
 }
 """
 
+
 def create_demo():
     with gr.Blocks(theme="Yntec/HaleyCH_Theme_Orange", css=css) as demo:
 
         with gr.Row():
             with gr.Column():
-                prompt = gr.Textbox(label="Prompt(한글/일본어/중국어 가능)", value="A cute and fluffy golden retriever puppy sitting upright, holding a neatly designed white sign with bold, colorful lettering that reads 'Have a Happy Day!' in cheerful fonts. The puppy has expressive, sparkling eyes, a happy smile, and fluffy ears slightly flopped. The background is a vibrant and sunny meadow with soft-focus flowers, glowing sunlight filtering through the trees, and a warm golden glow that enhances the joyful atmosphere. The sign is framed with small decorative flowers, adding a charming and wholesome touch. Ensure the text on the sign is clear and legible.")
+                prompt = gr.Textbox(label="Prompt(한글 가능)", value="A cute and fluffy golden retriever puppy sitting upright, holding a neatly designed white sign with bold, colorful lettering that reads 'Have a Happy Day!' in cheerful fonts. The puppy has expressive, sparkling eyes, a happy smile, and fluffy ears slightly flopped. The background is a vibrant and sunny meadow with soft-focus flowers, glowing sunlight filtering through the trees, and a warm golden glow that enhances the joyful atmosphere. The sign is framed with small decorative flowers, adding a charming and wholesome touch. Ensure the text on the sign is clear and legible.")
                 
                 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)
@@ -922,27 +875,19 @@ def create_demo():
             outputs=[init_image, image2image_strength, resize_img]
         )
 
-        gr.Examples(
-            examples=[
-                ["A magical fairy garden with glowing mushrooms and floating lanterns", 768, 768, 3.5, 30, 0, False, None, 0.8, True],  # English
-                ["아름다운 벚꽃이 흩날리는 한옥 정원에서 한복을 입은 소녀", 768, 768, 3.5, 30, 0, False, None, 0.8, True],  # Korean
-                ["夕暮れの富士山と桜の木の下で休んでいる可愛い柴犬", 768, 768, 3.5, 30, 0, False, None, 0.8, True],  # Japanese
-                ["古老的中国庭园里，一只熊猫正在竹林中悠闲地吃着竹子", 768, 768, 3.5, 30, 0, False, None, 0.8, True]   # Chinese
-            ],
-            inputs=[prompt, width, height, guidance, inference_steps, seed, do_img2img, init_image, image2image_strength, resize_img],
-            outputs=[output_image, output_seed],
-            fn=generate_image,
-            cache_examples=True
-        )
-
         generate_button.click(
             fn=generate_image,
             inputs=[prompt, width, height, guidance, inference_steps, seed, do_img2img, init_image, image2image_strength, resize_img],
             outputs=[output_image, output_seed]
         )
+        
+        examples = [
+            "a tiny astronaut hatching from an egg on the moon",
+            "a cat holding a sign that says hello world",
+            "an anime illustration of a wiener schnitzel",
+        ]
 
     return demo
 
 if __name__ == "__main__":
-    demo = create_demo()
-    demo.launch()
+    demo = create_demo()