Update app.py

app.py

@@ -38,6 +38,8 @@ def initialize_models(device):
     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(
@@ -59,7 +61,7 @@ def initialize_models(device):
     g.translator = transformers_pipeline(
         "translation", 
         model="Helsinki-NLP/opus-mt-ko-en",
-        device=device
+        device=device if device != "cuda" else 0
     )
     
 # CUDA 메모리 관리 설정
@@ -171,32 +173,49 @@ def unpack_state(state: dict) -> Tuple[Gaussian, edict, str]:
 
 @spaces.GPU
 def image_to_3d(trial_id: str, seed: int, randomize_seed: bool, ss_guidance_strength: float, ss_sampling_steps: int, slat_guidance_strength: float, slat_sampling_steps: int) -> Tuple[dict, str]:
-    if randomize_seed:
-        seed = np.random.randint(0, MAX_SEED)
-    outputs = g.trellis_pipeline.run(  # pipeline을 g.trellis_pipeline으로 수정
-        Image.open(f"{TMP_DIR}/{trial_id}.png"),
-        seed=seed,
-        formats=["gaussian", "mesh"],
-        preprocess_image=False,
-        sparse_structure_sampler_params={
-            "steps": ss_sampling_steps,
-            "cfg_strength": ss_guidance_strength,
-        },
-        slat_sampler_params={
-            "steps": slat_sampling_steps,
-            "cfg_strength": slat_guidance_strength,
-        },
-    )
+    if not trial_id:  # trial_id 검증 추가
+        print("Error: No trial_id provided")
+        return None, None
+        
+    try:
+        if randomize_seed:
+            seed = np.random.randint(0, MAX_SEED)
+            
+        image_path = f"{TMP_DIR}/{trial_id}.png"
+        if not os.path.exists(image_path):
+            print(f"Error: Image file not found at {image_path}")
+            return None, None
+            
+        image = Image.open(image_path)
+        
+        outputs = g.trellis_pipeline.run(
+            image,
+            seed=seed,
+            formats=["gaussian", "mesh"],
+            preprocess_image=False,
+            sparse_structure_sampler_params={
+                "steps": ss_sampling_steps,
+                "cfg_strength": ss_guidance_strength,
+            },
+            slat_sampler_params={
+                "steps": slat_sampling_steps,
+                "cfg_strength": slat_guidance_strength,
+            },
+        )
 
-    video = render_utils.render_video(outputs['gaussian'][0], num_frames=120)['color']
-    video_geo = render_utils.render_video(outputs['mesh'][0], num_frames=120)['normal']
-    video = [np.concatenate([video[i], video_geo[i]], axis=1) for i in range(len(video))]
-    trial_id = uuid.uuid4()
-    video_path = f"{TMP_DIR}/{trial_id}.mp4"
-    os.makedirs(os.path.dirname(video_path), exist_ok=True)
-    imageio.mimsave(video_path, video, fps=15)
-    state = pack_state(outputs['gaussian'][0], outputs['mesh'][0], trial_id)
-    return state, video_path
+        video = render_utils.render_video(outputs['gaussian'][0], num_frames=120)['color']
+        video_geo = render_utils.render_video(outputs['mesh'][0], num_frames=120)['normal']
+        video = [np.concatenate([video[i], video_geo[i]], axis=1) for i in range(len(video))]
+        new_trial_id = str(uuid.uuid4())
+        video_path = f"{TMP_DIR}/{new_trial_id}.mp4"
+        os.makedirs(os.path.dirname(video_path), exist_ok=True)
+        imageio.mimsave(video_path, video, fps=15)
+        state = pack_state(outputs['gaussian'][0], outputs['mesh'][0], new_trial_id)
+        return state, video_path
+        
+    except Exception as e:
+        print(f"Error in image_to_3d: {str(e)}")
+        return None, None
 
 
 @spaces.GPU