Update app.py

app.py

@@ -16,20 +16,18 @@ from diffusers import FluxPipeline
 from typing import *
 
 
+MAX_SEED = np.iinfo(np.int32).max
+TMP_DIR = "/tmp/Trellis-demo"
+os.makedirs(TMP_DIR, exist_ok=True)
+
 # 메모리 관련 환경 변수
 os.environ['PYTORCH_CUDA_ALLOC_CONF'] = 'max_split_size_mb:128'
 os.environ['TRANSFORMERS_CACHE'] = '/tmp/transformers_cache'
 os.environ['TORCH_HOME'] = '/tmp/torch_home'
 os.environ['HF_HOME'] = '/tmp/huggingface'
 os.environ['XDG_CACHE_HOME'] = '/tmp/cache'
-
-# 환경 변수 설정
 os.environ['SPCONV_ALGO'] = 'native'
-os.environ['WARP_USE_CPU'] = '1'  # Warp를 CPU 모드로 강제
-
-MAX_SEED = np.iinfo(np.int32).max
-TMP_DIR = "/tmp/Trellis-demo"
-os.makedirs(TMP_DIR, exist_ok=True)
+os.environ['WARP_USE_CPU'] = '1'
 
 def initialize_models():
     global pipeline, translator, flux_pipe
@@ -39,22 +37,20 @@ def initialize_models():
         for dir_path in ['/tmp/transformers_cache', '/tmp/torch_home', '/tmp/huggingface', '/tmp/cache']:
             os.makedirs(dir_path, exist_ok=True)
         
-        # Trellis 파이프라인 초기화 (메모리 최적화)
+        # Trellis 파이프라인 초기화
         pipeline = TrellisImageTo3DPipeline.from_pretrained(
-            "JeffreyXiang/TRELLIS-image-large",
-            resume_download=True,
-            local_files_only=False
+            "JeffreyXiang/TRELLIS-image-large"
         )
         
-        # 번역기 초기화 (더 작은 모델 사용)
+        # 번역기 초기화
         translator = translation_pipeline(
             "translation",
             model="Helsinki-NLP/opus-mt-ko-en",
             device="cpu"
         )
         
-        # Flux 파이프라인 초기화 (메모리 최적화)
-        flux_pipe = None  # 필요할 때 로드
+        # Flux 파이프라인은 필요할 때 로드
+        flux_pipe = None
         
         print("Models initialized successfully")
         return True
@@ -89,6 +85,7 @@ def free_memory():
                 except Exception as e:
                     print(f'Error deleting {file_path}: {e}')
 
+
 @spaces.GPU
 def setup_gpu_model(model):
     """GPU 설정이 필요한 모델을 처리하는 함수"""
@@ -171,7 +168,7 @@ def image_to_3d(trial_id: str, seed: int, randomize_seed: bool, ss_guidance_stre
             
         input_image = Image.open(f"{TMP_DIR}/{trial_id}.png")
         
-        # 더 작은 이미지 크기 사용
+        # 이미지 크기 제한
         max_size = 512
         if max(input_image.size) > max_size:
             ratio = max_size / max(input_image.size)
@@ -200,6 +197,7 @@ def image_to_3d(trial_id: str, seed: int, randomize_seed: bool, ss_guidance_stre
                 }
             )
         
+        # 비디오 프레임 수 감소
         video = render_utils.render_video(outputs['gaussian'][0], num_frames=30)['color']
         video_geo = render_utils.render_video(outputs['mesh'][0], num_frames=30)['normal']
         video = [np.concatenate([video[i], video_geo[i]], axis=1) for i in range(len(video))]
@@ -225,7 +223,7 @@ def image_to_3d(trial_id: str, seed: int, randomize_seed: bool, ss_guidance_stre
 @spaces.GPU
 def generate_image_from_text(prompt, height, width, guidance_scale, num_steps):
     try:
-        # Flux 파이프라인 필요할 때 로드
+        # Flux 파이프라인 로드
         flux_pipe = load_flux_pipe()
         
         if torch.cuda.is_available():
@@ -416,6 +414,5 @@ if __name__ == "__main__":
         share=True,
         max_threads=1,
         show_error=True,
-        cache_examples=False,
-        enable_queue=True
-    )
+        cache_examples=False
+    )