Update app.py

app.py

@@ -9,59 +9,58 @@ from diffusers import DiffusionPipeline
 from custom_pipeline import FLUXPipelineWithIntermediateOutputs
 from transformers import pipeline
 
-# Translation model loading with device specification
+# 번역 모델 설정 (CPU 사용)
 translator = pipeline("translation", model="Helsinki-NLP/opus-mt-ko-en", device="cpu")
 
-# Constants
+# 상수 정의
 MAX_SEED = np.iinfo(np.int32).max
 MAX_IMAGE_SIZE = 2048
 DEFAULT_WIDTH = 1024
 DEFAULT_HEIGHT = 1024
 DEFAULT_INFERENCE_STEPS = 1
-GPU_DURATION = 15  # Reduced from 25 to stay within quota
+GPU_DURATION = 15  # GPU 할당 시간 축소
 
-# Device and model setup with memory optimization
+# 모델 설정
 def setup_model():
     dtype = torch.float16
     pipe = FLUXPipelineWithIntermediateOutputs.from_pretrained(
         "black-forest-labs/FLUX.1-schnell", 
-        torch_dtype=dtype,
-        device_map="auto"  # Enable model parallelism
-    )
+        torch_dtype=dtype
+    ).to("cuda")
     return pipe
 
 pipe = setup_model()
 
-# Menu labels dictionary
-english_labels = {
-    "Generated Image": "Generated Image",
-    "Prompt": "Prompt",
-    "Enhance Image": "Enhance Image",
-    "Advanced Options": "Advanced Options",
-    "Seed": "Seed",
-    "Randomize Seed": "Randomize Seed",
-    "Width": "Width",
-    "Height": "Height",
-    "Inference Steps": "Inference Steps",
-    "Inspiration Gallery": "Inspiration Gallery"
+# 메뉴 레이블
+labels = {
+    "Generated Image": "생성된 이미지",
+    "Prompt": "프롬프트",
+    "Enhance Image": "이미지 향상",
+    "Advanced Options": "고급 설정",
+    "Seed": "시드",
+    "Randomize Seed": "랜덤 시드",
+    "Width": "너비",
+    "Height": "높이",
+    "Inference Steps": "추론 단계",
+    "Inspiration Gallery": "영감 갤러리"
 }
 
 def translate_if_korean(text):
-    """Safely translate Korean text to English."""
+    """한글 텍스트를 영어로 안전하게 번역"""
     try:
         if any('\u3131' <= char <= '\u3163' or '\uac00' <= char <= '\ud7a3' for char in text):
             return translator(text)[0]['translation_text']
         return text
     except Exception as e:
-        print(f"Translation error: {e}")
+        print(f"번역 오류: {e}")
         return text
 
-# Modified inference function with error handling and memory management
+# 이미지 생성 함수
 @spaces.GPU(duration=GPU_DURATION)
 def generate_image(prompt, seed=None, width=DEFAULT_WIDTH, height=DEFAULT_HEIGHT, 
                   randomize_seed=True, num_inference_steps=DEFAULT_INFERENCE_STEPS):
     try:
-        # Input validation
+        # 입력값 검증
         if not isinstance(seed, (int, type(None))):
             seed = None
             randomize_seed = True
@@ -71,7 +70,7 @@ def generate_image(prompt, seed=None, width=DEFAULT_WIDTH, height=DEFAULT_HEIGHT
         if seed is None or randomize_seed:
             seed = random.randint(0, MAX_SEED)
         
-        # Ensure valid dimensions
+        # 크기 유효성 검사
         width = min(max(256, width), MAX_IMAGE_SIZE)
         height = min(max(256, height), MAX_IMAGE_SIZE)
         
@@ -79,7 +78,7 @@ def generate_image(prompt, seed=None, width=DEFAULT_WIDTH, height=DEFAULT_HEIGHT
         
         start_time = time.time()
         
-        with torch.cuda.amp.autocast():  # Enable automatic mixed precision
+        with torch.cuda.amp.autocast():
             for img in pipe.generate_images(
                 prompt=prompt,
                 guidance_scale=0,
@@ -88,26 +87,25 @@ def generate_image(prompt, seed=None, width=DEFAULT_WIDTH, height=DEFAULT_HEIGHT
                 height=height,
                 generator=generator
             ):
-                latency = f"Processing Time: {(time.time()-start_time):.2f} seconds"
+                latency = f"처리 시간: {(time.time()-start_time):.2f} 초"
                 
-                # Clear CUDA cache after generation
+                # CUDA 캐시 정리
                 if torch.cuda.is_available():
                     torch.cuda.empty_cache()
                     
                 yield img, seed, latency
                 
     except Exception as e:
-        print(f"Error in generate_image: {e}")
-        # Return a blank image or error message
-        yield None, seed, f"Error: {str(e)}"
+        print(f"이미지 생성 오류: {e}")
+        yield None, seed, f"오류: {str(e)}"
 
-# Example generator with error handling
+# 예제 이미지 생성
 def generate_example_image(prompt):
     try:
         return next(generate_image(prompt, randomize_seed=True))
     except Exception as e:
-        print(f"Error in example generation: {e}")
-        return None, None, f"Error: {str(e)}"
+        print(f"예제 생성 오류: {e}")
+        return None, None, f"오류: {str(e)}"
 
 # Example prompts
 examples = [
@@ -119,63 +117,63 @@ examples = [
     "A cosmic coffee shop where baristas are constellations serving drinks made of stardust"
 ]
 
+
 css = """
 footer {
     visibility: hidden;
 }
 """
 
-# --- Gradio UI with improved error handling ---
+# Gradio UI 구성
 with gr.Blocks(theme="Yntec/HaleyCH_Theme_Orange", css=css) as demo:
     with gr.Column(elem_id="app-container"):
         with gr.Row():
             with gr.Column(scale=3):
-                result = gr.Image(label=english_labels["Generated Image"], 
+                result = gr.Image(label=labels["Generated Image"], 
                                 show_label=False, 
                                 interactive=False)
             with gr.Column(scale=1):
                 prompt = gr.Text(
-                    label=english_labels["Prompt"],
-                    placeholder="Describe the image you want to generate...",
+                    label=labels["Prompt"],
+                    placeholder="생성하고 싶은 이미지를 설명해주세요...",
                     lines=3,
                     show_label=False,
                     container=False,
                 )
-                enhanceBtn = gr.Button(f"🚀 {english_labels['Enhance Image']}")
+                enhanceBtn = gr.Button(f"🚀 {labels['Enhance Image']}")
 
-                with gr.Column(english_labels["Advanced Options"]):
+                with gr.Column(labels["Advanced Options"]):
                     with gr.Row():
                         latency = gr.Text(show_label=False)
                     with gr.Row():
-                        # Modified Number component with proper validation
                         seed = gr.Number(
-                            label=english_labels["Seed"], 
+                            label=labels["Seed"], 
                             value=42, 
                             precision=0,
                             minimum=0,
                             maximum=MAX_SEED
                         )
                         randomize_seed = gr.Checkbox(
-                            label=english_labels["Randomize Seed"], 
+                            label=labels["Randomize Seed"], 
                             value=True
                         )
                     with gr.Row():
                         width = gr.Slider(
-                            label=english_labels["Width"], 
+                            label=labels["Width"], 
                             minimum=256, 
                             maximum=MAX_IMAGE_SIZE, 
                             step=32, 
                             value=DEFAULT_WIDTH
                         )
                         height = gr.Slider(
-                            label=english_labels["Height"], 
+                            label=labels["Height"], 
                             minimum=256, 
                             maximum=MAX_IMAGE_SIZE, 
                             step=32, 
                             value=DEFAULT_HEIGHT
                         )
                         num_inference_steps = gr.Slider(
-                            label=english_labels["Inference Steps"], 
+                            label=labels["Inference Steps"], 
                             minimum=1, 
                             maximum=4, 
                             step=1, 
@@ -183,7 +181,7 @@ with gr.Blocks(theme="Yntec/HaleyCH_Theme_Orange", css=css) as demo:
                         )
 
         with gr.Row():
-            gr.Markdown(f"### 🌟 {english_labels['Inspiration Gallery']}")
+            gr.Markdown(f"### 🌟 {labels['Inspiration Gallery']}")
         with gr.Row():
             gr.Examples(
                 examples=examples,
@@ -193,7 +191,14 @@ with gr.Blocks(theme="Yntec/HaleyCH_Theme_Orange", css=css) as demo:
                 cache_examples=False
             )
 
-    # Event handling with improved error handling
+    # 이벤트 처리
+    def validated_generate(*args):
+        try:
+            return next(generate_image(*args))
+        except Exception as e:
+            print(f"검증 생성 오류: {e}")
+            return None, args[1], f"오류: {str(e)}"
+
     enhanceBtn.click(
         fn=generate_image,
         inputs=[prompt, seed, width, height],
@@ -203,14 +208,6 @@ with gr.Blocks(theme="Yntec/HaleyCH_Theme_Orange", css=css) as demo:
         queue=False
     )
 
-    # Modified event handler with proper input validation
-    def validated_generate(*args):
-        try:
-            return next(generate_image(*args))
-        except Exception as e:
-            print(f"Error in validated_generate: {e}")
-            return None, args[1], f"Error: {str(e)}"
-
     gr.on(
         triggers=[prompt.input, width.input, height.input, num_inference_steps.input],
         fn=validated_generate,
@@ -223,4 +220,5 @@ with gr.Blocks(theme="Yntec/HaleyCH_Theme_Orange", css=css) as demo:
     )
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch()
+