Update app.py

app.py

@@ -19,11 +19,50 @@ from trellis.representations import Gaussian, MeshExtractResult
 from trellis.utils import render_utils, postprocessing_utils
 from diffusers import FluxPipeline
 from typing import Tuple, Dict, Any  # Tuple import 추가
-
-# 파일 상단의 import 문 수정
+# 파일 상단의 import 문
 import transformers
 from transformers import pipeline as transformers_pipeline
-from transformers import Pipeline  # Pipeline 타입 힌트를 위해 추가
+from transformers import Pipeline
+
+# 전역 변수 초기화
+class GlobalVars:
+    def __init__(self):
+        self.translator = None
+        self.trellis_pipeline = None
+        self.flux_pipe = None
+
+g = GlobalVars()
+
+def initialize_models(device):
+    # 3D 생성 파이프라인
+    g.trellis_pipeline = TrellisImageTo3DPipeline.from_pretrained(
+        "JeffreyXiang/TRELLIS-image-large"
+    )
+    g.trellis_pipeline.to(device)
+    
+    # 이미지 생성 파이프라인
+    g.flux_pipe = FluxPipeline.from_pretrained(
+        "black-forest-labs/FLUX.1-dev",
+        torch_dtype=torch.bfloat16,
+        device_map="balanced"
+    )
+    
+    # Hyper-SD LoRA 로드
+    lora_path = hf_hub_download(
+        "ByteDance/Hyper-SD",
+        "Hyper-FLUX.1-dev-8steps-lora.safetensors",
+        use_auth_token=HF_TOKEN
+    )
+    g.flux_pipe.load_lora_weights(lora_path)
+    g.flux_pipe.fuse_lora(lora_scale=0.125)
+    
+    # 번역기 초기화
+    g.translator = transformers_pipeline(
+        "translation", 
+        model="Helsinki-NLP/opus-mt-ko-en",
+        device=device
+    )
+    
 # CUDA 메모리 관리 설정
 torch.cuda.empty_cache()
 torch.backends.cuda.matmul.allow_tf32 = True
@@ -71,7 +110,7 @@ class timer:
 
 def preprocess_image(image: Image.Image) -> Tuple[str, Image.Image]:
     trial_id = str(uuid.uuid4())
-    processed_image = pipeline.preprocess_image(image)
+    processed_image = g.trellis_pipeline.preprocess_image(image)
     processed_image.save(f"{TMP_DIR}/{trial_id}.png")
     return trial_id, processed_image
 
@@ -169,7 +208,7 @@ def text_to_image(prompt: str, height: int, width: int, steps: int, scales: floa
     
     # 프롬프트 전처리
     if contains_korean(prompt):
-        translated = translator(prompt)[0]['translation_text']
+        translated = g.translator(prompt)[0]['translation_text']
         prompt = translated
     
     # 프롬프트 형식 강제
@@ -177,7 +216,7 @@ def text_to_image(prompt: str, height: int, width: int, steps: int, scales: floa
     
     with torch.inference_mode(), torch.autocast("cuda", dtype=torch.bfloat16):
         try:
-            generated_image = flux_pipe(  # pipe를 flux_pipe로 변경
+            generated_image = g.flux_pipe(
                 prompt=[formatted_prompt],
                 generator=torch.Generator().manual_seed(int(seed)),
                 num_inference_steps=int(steps),
@@ -330,35 +369,8 @@ if __name__ == "__main__":
     print(f"Using device: {device}")
 
     try:
-        # 3D 생성 파이프라인
-        trellis_pipeline = TrellisImageTo3DPipeline.from_pretrained(
-            "JeffreyXiang/TRELLIS-image-large"
-        )
-        trellis_pipeline.to(device)
-        
-        # 이미지 생성 파이프라인
-        flux_pipe = FluxPipeline.from_pretrained(
-            "black-forest-labs/FLUX.1-dev",
-            torch_dtype=torch.bfloat16,
-            device_map="balanced"
-        )
-        
-        # Hyper-SD LoRA 로드
-        lora_path = hf_hub_download(
-            "ByteDance/Hyper-SD",
-            "Hyper-FLUX.1-dev-8steps-lora.safetensors",
-            use_auth_token=HF_TOKEN
-        )
-        flux_pipe.load_lora_weights(lora_path)
-        flux_pipe.fuse_lora(lora_scale=0.125)
-        
-        # 번역기 초기화
-        global translator
-        translator = transformers_pipeline(
-            "translation", 
-            model="Helsinki-NLP/opus-mt-ko-en",
-            device=device
-        )
+        # 모델 초기화
+        initialize_models(device)
         
         # CUDA 메모리 초기화
         if torch.cuda.is_available():
@@ -367,7 +379,7 @@ if __name__ == "__main__":
         # 초기 이미지 전처리 테스트
         try:
             test_image = Image.fromarray(np.zeros((512, 512, 3), dtype=np.uint8))
-            trellis_pipeline.preprocess_image(test_image)
+            g.trellis_pipeline.preprocess_image(test_image)
         except Exception as e:
             print(f"Warning: Initial preprocessing test failed: {e}")