Update app.py

app.py

@@ -64,10 +64,8 @@ except ModuleNotFoundError:
 # ─────────────────────────────────────────────────────────────
 caption_model = VisionEncoderDecoderModel.from_pretrained(
     "nlpconnect/vit-gpt2-image-captioning",
-    use_safetensors=True,
-    low_cpu_mem_usage=True          # 그대로 두어도 OK
-)
-caption_model.to("cpu")             # ★ 추가
+    use_safetensors=True      # OK
+).eval()                      # 평가 모드로
 
 feature_extractor = ViTImageProcessor.from_pretrained(
     "nlpconnect/vit-gpt2-image-captioning"
@@ -76,6 +74,7 @@ tokenizer = AutoTokenizer.from_pretrained(
     "nlpconnect/vit-gpt2-image-captioning"
 )
 
+
 # ─────────────────────────────────────────────────────────────
 # 5. MusicGen 모델
 # ─────────────────────────────────────────────────────────────
@@ -86,11 +85,15 @@ musicgen.set_generation_params(duration=10)
 # 6. 파이프라인 함수
 # ─────────────────────────────────────────────────────────────
 def generate_caption(image: Image.Image) -> str:
-    pixel_values = feature_extractor(images=image, return_tensors="pt").pixel_values
-    caption_model.to(pixel_values.device)       # ★ 안전 이동
-    with torch.no_grad():
-        ids = caption_model.generate(pixel_values, max_length=50)
-    return tokenizer.decode(ids[0], skip_special_tokens=True)
+    with torch.no_grad():                               # ★ 메모리 절약
+        pixel_values = feature_extractor(
+            images=image, return_tensors="pt"
+        ).pixel_values
+        output_ids = caption_model.generate(
+            pixel_values.to(caption_model.device),       # CPU 디바이스 통일
+            max_length=50
+        )
+    return tokenizer.decode(output_ids[0], skip_special_tokens=True)
 
 
 def generate_music(prompt: str) -> str: