Update app.py

app.py

@@ -34,35 +34,46 @@ class GlobalVars:
 g = GlobalVars()
 
 def initialize_models(device):
-    # 3D 생성 파이프라인
-    g.trellis_pipeline = TrellisImageTo3DPipeline.from_pretrained(
-        "JeffreyXiang/TRELLIS-image-large"
-    )
-    if torch.cuda.is_available():
-        g.trellis_pipeline = g.trellis_pipeline.to("cuda")
-    
-    # 이미지 생성 파이프라인
-    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 if device != "cuda" else 0
-    )
+    try:
+        print("Initializing models...")
+        # 3D 생성 파이프라인
+        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...")
+        g.flux_pipe = FluxPipeline.from_pretrained(
+            "black-forest-labs/FLUX.1-dev",
+            torch_dtype=torch.bfloat16,
+            device_map="balanced"
+        )
+        
+        # Hyper-SD LoRA 로드
+        print("Loading LoRA weights...")
+        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)
+        
+        # 번역기 초기화
+        print("Initializing translator...")
+        g.translator = transformers_pipeline(
+            "translation", 
+            model="Helsinki-NLP/opus-mt-ko-en",
+            device=device if device != "cuda" else 0
+        )
+        print("Model initialization completed successfully")
+        
+    except Exception as e:
+        print(f"Error during model initialization: {str(e)}")
+        raise
     
 # CUDA 메모리 관리 설정
 torch.cuda.empty_cache()
@@ -113,6 +124,10 @@ def preprocess_image(image: Image.Image) -> Tuple[str, Image.Image]:
         return "", None
         
     try:
+        if g.trellis_pipeline is None:
+            print("Error: trellis_pipeline is not initialized")
+            return "", None
+            
         # webp 이미지를 RGB로 변환
         if isinstance(image, str) and image.endswith('.webp'):
             image = Image.open(image).convert('RGB')
@@ -370,7 +385,7 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
     trial_id = gr.Textbox(visible=False)
     output_buf = gr.State()
 
-    # Examples 갤러리를 맨 아래로 이동
+# Examples 갤러리를 맨 아래로 이동
     if example_images:
         gr.Markdown("""### Example Images""")
         with gr.Row():
@@ -379,11 +394,14 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
                 label="Click an image to use it",
                 show_label=True,
                 elem_id="gallery",
-                columns=5,
-                rows=5,
-                height=600,
-                allow_preview=True
+                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])
@@ -454,24 +472,33 @@ if __name__ == "__main__":
         device = "cuda" if torch.cuda.is_available() else "cpu"
         print(f"Using device: {device}")
         
+        # CUDA 메모리 초기화
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+            torch.cuda.synchronize()
+        
         # 모델 초기화
         initialize_models(device)
         
         # 초기 이미지 전처리 테스트
         try:
             test_image = Image.fromarray(np.zeros((512, 512, 3), dtype=np.uint8))
-            g.trellis_pipeline.preprocess_image(test_image)
+            if g.trellis_pipeline is not None:
+                g.trellis_pipeline.preprocess_image(test_image)
+            else:
+                print("Warning: trellis_pipeline is None")
         except Exception as e:
             print(f"Warning: Initial preprocessing test failed: {e}")
             
         # Gradio 인터페이스 실행
         demo.queue()  # 큐 기능 활성화
         demo.launch(
-            allowed_paths=[PERSISTENT_DIR],
+            allowed_paths=[PERSISTENT_DIR, TMP_DIR],  # TMP_DIR 추가
             server_name="0.0.0.0",
             server_port=7860,
             show_error=True,
-            share=True  # share를 True로 설정
+            share=True,  # share를 True로 설정
+            enable_queue=True  # 큐 활성화
         )
         
     except Exception as e: