Update app.py

app.py

@@ -1,4 +1,4 @@
-import gradio as gr
+# import gradio as gr
 import torch
 from PIL import Image
 from transformers import BlipProcessor, BlipForConditionalGeneration
@@ -18,16 +18,34 @@ def generate_caption(image):
     # 고속 처리를 위한 리사이즈
     image = image.resize((384, 384))
     
-    # 설명 생성: 최소 20단어, 빔 서치 활용
     inputs = processor(images=image, return_tensors="pt").to(device)
+    # 충분히 긴 토큰 범위를 생성하고, 단어 수로 후처리
     output_ids = model.generate(
         **inputs,
-        max_length=100,
-        min_length=20,
+        max_length=200,      # 토큰 길이 최대
+        min_length=60,       # 토큰 길이 최소
         num_beams=5,
         early_stopping=True
     )
-    caption = processor.decode(output_ids[0], skip_special_tokens=True)
+    caption = processor.decode(output_ids[0], skip_special_tokens=True).strip()
+    
+    # 단어 수로 30~50 사이에 맞추기
+    words = caption.split()
+    if len(words) > 50:
+        words = words[:50]
+    elif len(words) < 30:
+        # 너무 짧으면 토큰 최소 길이 없이 재생성해보고
+        output_ids = model.generate(
+            **inputs,
+            max_length=200,
+            num_beams=5,
+            early_stopping=True
+        )
+        caption = processor.decode(output_ids[0], skip_special_tokens=True).strip()
+        words = caption.split()
+        if len(words) > 50:
+            words = words[:50]
+    caption = " ".join(words)
 
     # "Asian"이 포함된 경우 "Korean"으로 대체
     if "Asian" in caption:
@@ -45,7 +63,6 @@ with gr.Blocks(title="이미지 설명 생성기") as demo:
             image_input = gr.Image(label="입력 이미지", type="pil")
         with gr.Column():
             caption_output = gr.Textbox(label="생성된 설명", lines=3, show_copy_button=True)
-            # HTML로 버튼 생성
             gr.HTML("""
             <div style='margin-top: 10px; text-align: center;'>
                 <a href="https://huggingface.co/spaces/VIDraft/stable-diffusion-3.5-large-turboX" target="_blank">
@@ -56,13 +73,11 @@ with gr.Blocks(title="이미지 설명 생성기") as demo:
             </div>
             """)
 
-    # 업로드 → 설명 자동 생성 연결
     image_input.upload(fn=generate_caption, inputs=image_input, outputs=caption_output)
 
 # 5. 앱 실행
 demo.launch(debug=True)
 
-
 # import gradio as gr
 # import torch
 # from PIL import Image