Update app.py

app.py

@@ -65,8 +65,10 @@ except ModuleNotFoundError:
 caption_model = VisionEncoderDecoderModel.from_pretrained(
     "nlpconnect/vit-gpt2-image-captioning",
     use_safetensors=True,
-    low_cpu_mem_usage=True
+    low_cpu_mem_usage=True          # 그대로 두어도 OK
 )
+caption_model.to("cpu")             # ★ 추가
+
 feature_extractor = ViTImageProcessor.from_pretrained(
     "nlpconnect/vit-gpt2-image-captioning"
 )
@@ -85,9 +87,12 @@ musicgen.set_generation_params(duration=10)
 # ─────────────────────────────────────────────────────────────
 def generate_caption(image: Image.Image) -> str:
     pixel_values = feature_extractor(images=image, return_tensors="pt").pixel_values
-    ids = caption_model.generate(pixel_values, max_length=50)
+    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)
 
+
 def generate_music(prompt: str) -> str:
     wav = musicgen.generate([prompt])           # batch size = 1
     tmpdir = tempfile.mkdtemp()