Update app.py

app.py

@@ -51,21 +51,28 @@ def load_gallery_images():
     
     return [(os.path.join("gallery", row[2]), f"{row[0]}: {row[1]}") for row in rows]
 
-# 번역 모델 초기화
-translator = pipeline("translation", model="Helsinki-NLP/opus-mt-ko-en")
-
-# 프롬프트 처리 함수 추가
-def process_prompt(prompt):
-    if any('\u3131' <= char <= '\u3163' or '\uac00' <= char <= '\ud7a3' for char in prompt):
-        translated = translator(prompt)[0]['translation_text']
-        return prompt, translated
-    return prompt, prompt
+# CPU에서 번역 모델 초기화
+@spaces.CPU
+def init_translator():
+    return pipeline("translation", model="Helsinki-NLP/opus-mt-ko-en", device="cpu")
+
+# CPU에서 번역 처리
+@spaces.CPU
+def translate_text(translator, text):
+    if any('\u3131' <= char <= '\u3163' or '\uac00' <= char <= '\ud7a3' for char in text):
+        translated = translator(text)[0]['translation_text']
+        return text, translated
+    return text, text
+
+# 프롬프트 처리 함수
+def process_prompt(translator, prompt):
+    return translate_text(translator, prompt)
 
 KEY_JSON = os.getenv("KEY_JSON")
 with open(KEY_JSON, 'r') as f:
     loras = json.load(f)
 
-# Initialize the base model
+# 기본 모델 초기화
 dtype = torch.bfloat16
 device = "cuda" if torch.cuda.is_available() else "cpu"
 base_model = "black-forest-labs/FLUX.1-dev"
@@ -136,11 +143,12 @@ def generate_image(prompt_mash, steps, seed, cfg_scale, width, height, lora_scal
         ):
             yield img
 
-def run_lora(prompt, cfg_scale, steps, selected_index, randomize_seed, seed, width, height, lora_scale, progress=gr.Progress(track_tqdm=True)):
+@spaces.GPU(duration=70)
+def run_lora(prompt, cfg_scale, steps, selected_index, randomize_seed, seed, width, height, lora_scale, translator, progress=gr.Progress(track_tqdm=True)):
     if selected_index is None:
         raise gr.Error("진행하기 전에 LoRA를 선택해야 합니다.")
     
-    original_prompt, english_prompt = process_prompt(prompt)
+    original_prompt, english_prompt = process_prompt(translator, prompt)
     
     selected_lora = loras[selected_index]
     lora_path = selected_lora["repo"]
@@ -159,21 +167,21 @@ def run_lora(prompt, cfg_scale, steps, selected_index, randomize_seed, seed, wid
     with calculateDuration("LoRA 언로드"):
         pipe.unload_lora_weights()
         
-    # Load LoRA weights
+    # LoRA 가중치 로드
     with calculateDuration(f"{selected_lora['title']}의 LoRA 가중치 로드"):
         if "weights" in selected_lora:
             pipe.load_lora_weights(lora_path, weight_name=selected_lora["weights"])
         else:
             pipe.load_lora_weights(lora_path)
 
-    # Set random seed for reproducibility
+    # 재현성을 위한 시드 설정
     with calculateDuration("시드 무작위화"):
         if randomize_seed:
             seed = random.randint(0, MAX_SEED)
 
     image_generator = generate_image(prompt_mash, steps, seed, cfg_scale, width, height, lora_scale, progress)
     
-    # Consume the generator to get the final image
+    # 최종 이미지를 얻기 위해 제너레이터 소비
     final_image = None
     step_counter = 0
     for image in image_generator:
@@ -187,32 +195,31 @@ def run_lora(prompt, cfg_scale, steps, selected_index, randomize_seed, seed, wid
     
     yield final_image, seed, gr.update(value=progress_bar, visible=False), original_prompt, english_prompt
 
-
 def get_huggingface_safetensors(link):
-  split_link = link.split("/")
-  if(len(split_link) == 2):
-            model_card = ModelCard.load(link)
-            base_model = model_card.data.get("base_model")
-            print(base_model)
-            if((base_model != "black-forest-labs/FLUX.1-dev") and (base_model != "black-forest-labs/FLUX.1-schnell")):
-                raise Exception("Not a FLUX LoRA!")
-            image_path = model_card.data.get("widget", [{}])[0].get("output", {}).get("url", None)
-            trigger_word = model_card.data.get("instance_prompt", "")
-            image_url = f"https://huggingface.co/{link}/resolve/main/{image_path}" if image_path else None
-            fs = HfFileSystem()
-            try:
-                list_of_files = fs.ls(link, detail=False)
-                for file in list_of_files:
-                    if(file.endswith(".safetensors")):
-                        safetensors_name = file.split("/")[-1]
-                    if (not image_url and file.lower().endswith((".jpg", ".jpeg", ".png", ".webp"))):
-                      image_elements = file.split("/")
-                      image_url = f"https://huggingface.co/{link}/resolve/main/{image_elements[-1]}"
-            except Exception as e:
-              print(e)
-              gr.Warning(f"You didn't include a link neither a valid Hugging Face repository with a *.safetensors LoRA")
-              raise Exception(f"You didn't include a link neither a valid Hugging Face repository with a *.safetensors LoRA")
-            return split_link[1], link, safetensors_name, trigger_word, image_url
+    split_link = link.split("/")
+    if(len(split_link) == 2):
+        model_card = ModelCard.load(link)
+        base_model = model_card.data.get("base_model")
+        print(base_model)
+        if((base_model != "black-forest-labs/FLUX.1-dev") and (base_model != "black-forest-labs/FLUX.1-schnell")):
+            raise Exception("Not a FLUX LoRA!")
+        image_path = model_card.data.get("widget", [{}])[0].get("output", {}).get("url", None)
+        trigger_word = model_card.data.get("instance_prompt", "")
+        image_url = f"https://huggingface.co/{link}/resolve/main/{image_path}" if image_path else None
+        fs = HfFileSystem()
+        try:
+            list_of_files = fs.ls(link, detail=False)
+            for file in list_of_files:
+                if(file.endswith(".safetensors")):
+                    safetensors_name = file.split("/")[-1]
+                if (not image_url and file.lower().endswith((".jpg", ".jpeg", ".png", ".webp"))):
+                    image_elements = file.split("/")
+                    image_url = f"https://huggingface.co/{link}/resolve/main/{image_elements[-1]}"
+        except Exception as e:
+            print(e)
+            gr.Warning(f"You didn't include a link neither a valid Hugging Face repository with a *.safetensors LoRA")
+            raise Exception(f"You didn't include a link neither a valid Hugging Face repository with a *.safetensors LoRA")
+        return split_link[1], link, safetensors_name, trigger_word, image_url
 
 def check_custom_model(link):
     if(link.startswith("https://")):
@@ -263,7 +270,6 @@ def add_custom_lora(custom_lora):
 def remove_custom_lora():
     return gr.update(visible=False), gr.update(visible=False), gr.update(), "", None, ""
 
-
 run_lora.zerogpu = True
 
 css = """
@@ -279,8 +285,12 @@ if not os.path.exists('gallery'):
 # 데이터베이스 초기화
 init_db()
 
-with gr.Blocks(theme="Nymbo/Nymbo_Theme", css=css) as app:
+# 번역기 초기화
+translator = init_translator()
 
+with gr.Blocks(theme="Nymbo/Nymbo_Theme", css=css) as app:
+    # 번역기를 상태로 추가
+    translator_state = gr.State(translator)
     selected_index = gr.State(None)
     
     with gr.Tabs():
@@ -342,11 +352,13 @@ with gr.Blocks(theme="Nymbo/Nymbo_Theme", css=css) as app:
         inputs=[width, height],
         outputs=[prompt, selected_info, selected_index, width, height]
     )
+    
     custom_lora.input(
         add_custom_lora,
         inputs=[custom_lora],
         outputs=[custom_lora_info, custom_lora_button, gallery, selected_info, selected_index, prompt]
     )
+    
     custom_lora_button.click(
         remove_custom_lora,
         outputs=[custom_lora_info, custom_lora_button, gallery, selected_info, selected_index, custom_lora]
@@ -355,7 +367,8 @@ with gr.Blocks(theme="Nymbo/Nymbo_Theme", css=css) as app:
     gr.on(
         triggers=[generate_button.click, prompt.submit],
         fn=run_lora,
-        inputs=[prompt, cfg_scale, steps, selected_index, randomize_seed, seed, width, height, lora_scale],
+        inputs=[prompt, cfg_scale, steps, selected_index, randomize_seed, 
+                seed, width, height, lora_scale, translator_state],
         outputs=[result, seed, progress_bar, original_prompt_display, english_prompt_display]
     )