Update app.py

app.py

@@ -40,9 +40,6 @@ def initialize_models(device):
         g.trellis_pipeline = TrellisImageTo3DPipeline.from_pretrained(
             "JeffreyXiang/TRELLIS-image-large"
         )
-        if torch.cuda.is_available():
-            print("Moving trellis_pipeline to CUDA")
-            g.trellis_pipeline = g.trellis_pipeline.to("cuda")
         
         # 이미지 생성 파이프라인
         print("Loading flux_pipe...")
@@ -67,7 +64,7 @@ def initialize_models(device):
         g.translator = transformers_pipeline(
             "translation", 
             model="Helsinki-NLP/opus-mt-ko-en",
-            device=device if device != "cuda" else 0
+            device=device
         )
         print("Model initialization completed successfully")
         
@@ -79,11 +76,12 @@ def initialize_models(device):
 torch.cuda.empty_cache()
 torch.backends.cuda.matmul.allow_tf32 = True
 torch.backends.cudnn.benchmark = True
-
 # 환경 변수 설정
 os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "max_split_size_mb:512"
 os.environ['SPCONV_ALGO'] = 'native'
 os.environ['SPARSE_BACKEND'] = 'native'
+os.environ['CUDA_LAUNCH_BLOCKING'] = '1'  # CUDA 작업 동기화
+
 
 # Hugging Face 토큰 설정
 HF_TOKEN = os.getenv("HF_TOKEN")
@@ -385,24 +383,6 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
     trial_id = gr.Textbox(visible=False)
     output_buf = gr.State()
 
-# Examples 갤러리를 맨 아래로 이동
-    if example_images:
-        gr.Markdown("""### Example Images""")
-        with gr.Row():
-            gallery = gr.Gallery(
-                value=example_images,
-                label="Click an image to use it",
-                show_label=True,
-                elem_id="gallery",
-                columns=12,  # 한 줄에 12개
-                rows=2,      # 2줄
-                height=300,  # 높이 조정
-                allow_preview=True,
-                object_fit="contain"  # 이미지 비율 유지
-            )
-    
- 
-            
             def load_example(evt: gr.SelectData):
                 selected_image = Image.open(example_images[evt.index])
                 trial_id_val, processed_image = preprocess_image(selected_image)
@@ -468,15 +448,10 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
 
 if __name__ == "__main__":
     try:
-        # CUDA 사용 가능 여부 확인
-        device = "cuda" if torch.cuda.is_available() else "cpu"
+        # CPU로 초기화
+        device = "cpu"
         print(f"Using device: {device}")
         
-        # CUDA 메모리 초기화
-        if torch.cuda.is_available():
-            torch.cuda.empty_cache()
-            torch.cuda.synchronize()
-        
         # 모델 초기화
         initialize_models(device)
         
@@ -490,15 +465,30 @@ if __name__ == "__main__":
         except Exception as e:
             print(f"Warning: Initial preprocessing test failed: {e}")
             
+        # Examples 갤러리 설정
+        if example_images:
+            gr.Markdown("""### Example Images""")
+            with gr.Row():
+                gallery = gr.Gallery(
+                    value=example_images,
+                    label="Click an image to use it",
+                    show_label=True,
+                    elem_id="gallery",
+                    columns=12,  # 한 줄에 12개
+                    rows=2,      # 2줄
+                    height=300,  # 높이 조정
+                    allow_preview=True,
+                    object_fit="contain"  # 이미지 비율 유지
+                )
+        
         # Gradio 인터페이스 실행
         demo.queue()  # 큐 기능 활성화
         demo.launch(
-            allowed_paths=[PERSISTENT_DIR, TMP_DIR],  # TMP_DIR 추가
+            allowed_paths=[PERSISTENT_DIR, TMP_DIR],
             server_name="0.0.0.0",
             server_port=7860,
             show_error=True,
-            share=True,  # share를 True로 설정
-            enable_queue=True  # 큐 활성화
+            share=True
         )
         
     except Exception as e: