Update app.py

app.py

@@ -108,11 +108,23 @@ class timer:
         print(f"{self.method} took {str(round(end - self.start, 2))}s")
 
 def preprocess_image(image: Image.Image) -> Tuple[str, Image.Image]:
-    trial_id = str(uuid.uuid4())
-    processed_image = g.trellis_pipeline.preprocess_image(image)
-    processed_image.save(f"{TMP_DIR}/{trial_id}.png")
-    return trial_id, processed_image
-
+    if image is None:
+        print("Error: Input image is None")
+        return "", None
+        
+    try:
+        trial_id = str(uuid.uuid4())
+        processed_image = g.trellis_pipeline.preprocess_image(image)
+        if processed_image is not None:
+            processed_image.save(f"{TMP_DIR}/{trial_id}.png")
+            return trial_id, processed_image
+        else:
+            print("Error: Processed image is None")
+            return "", None
+    except Exception as e:
+        print(f"Error in image preprocessing: {str(e)}")
+        return "", None
+        
 def pack_state(gs: Gaussian, mesh: MeshExtractResult, trial_id: str) -> dict:
     return {
         'gaussian': {
@@ -201,20 +213,24 @@ def deactivate_button() -> gr.Button:
 
 @spaces.GPU
 def text_to_image(prompt: str, height: int, width: int, steps: int, scales: float, seed: int) -> Image.Image:
-    # 한글 감지 및 번역
-    def contains_korean(text):
-        return any(ord('가') <= ord(c) <= ord('힣') for c in text)
-    
-    # 프롬프트 전처리
-    if contains_korean(prompt):
-        translated = g.translator(prompt)[0]['translation_text']
-        prompt = translated
-    
-    # 프롬프트 형식 강제
-    formatted_prompt = f"wbgmsst, 3D, {prompt}, white background"
-    
-    with torch.inference_mode(), torch.autocast("cuda", dtype=torch.bfloat16):
-        try:
+    try:
+        # CUDA 메모리 정리
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+            
+        # 한글 감지 및 번역
+        def contains_korean(text):
+            return any(ord('가') <= ord(c) <= ord('힣') for c in text)
+        
+        # 프롬프트 전처리
+        if contains_korean(prompt):
+            translated = g.translator(prompt)[0]['translation_text']
+            prompt = translated
+        
+        # 프롬프트 형식 강제
+        formatted_prompt = f"wbgmsst, 3D, {prompt}, white background"
+        
+        with torch.inference_mode(), torch.autocast("cuda", dtype=torch.bfloat16):
             generated_image = g.flux_pipe(
                 prompt=[formatted_prompt],
                 generator=torch.Generator().manual_seed(int(seed)),
@@ -225,13 +241,17 @@ def text_to_image(prompt: str, height: int, width: int, steps: int, scales: floa
                 max_sequence_length=256
             ).images[0]
             
-            trial_id = str(uuid.uuid4())
-            generated_image.save(f"{TMP_DIR}/{trial_id}.png")
-            return generated_image
-            
-        except Exception as e:
-            print(f"Error in image generation: {str(e)}")
-            return None
+            if generated_image is not None:
+                trial_id = str(uuid.uuid4())
+                generated_image.save(f"{TMP_DIR}/{trial_id}.png")
+                return generated_image
+            else:
+                print("Error: Generated image is None")
+                return None
+                
+    except Exception as e:
+        print(f"Error in image generation: {str(e)}")
+        return None
 
 # Gradio Interface
 with gr.Blocks(theme=gr.themes.Soft()) as demo: