Update app.py

app.py

@@ -1,8 +1,12 @@
+
+
+
 #!/usr/bin/env python
 
 import os
 import re
 import tempfile
+import gc  # garbage collector 추가
 from collections.abc import Iterator
 from threading import Thread
 import json
@@ -20,6 +24,15 @@ import pandas as pd
 # PDF 텍스트 추출
 import PyPDF2
 
+##############################################################################
+# 메모리 정리 함수 추가
+##############################################################################
+def clear_cuda_cache():
+    """CUDA 캐시를 명시적으로 비웁니다."""
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
+        gc.collect()
+
 ##############################################################################
 # SERPHouse API key from environment variable
 ##############################################################################
@@ -122,13 +135,11 @@ def do_web_search(query: str) -> str:
         # 모델에게 명확한 지침 추가
         instructions = """
 # 웹 검색 결과
-
 아래는 검색 결과입니다. 질문에 답변할 때 이 정보를 활용하세요:
 1. 각 결과의 제목, 내용, 출처 링크를 참고하세요
 2. 답변에 관련 정보의 출처를 명시적으로 인용하세요 (예: "X 출처에 따르면...")
 3. 응답에 실제 출처 링크를 포함하세요
 4. 여러 출처의 정보를 종합하여 답변하세요
-
 """
         
         search_results = instructions + "\n".join(summary_lines)
@@ -144,14 +155,15 @@ def do_web_search(query: str) -> str:
 # 모델/프로세서 로딩
 ##############################################################################
 MAX_CONTENT_CHARS = 4000
-model_id = os.getenv("MODEL_ID", "VIDraft/Gemma3-R1945-27B")
+MAX_INPUT_LENGTH = 4096  # 최대 입력 토큰 수 제한 추가
+model_id = os.getenv("MODEL_ID", "mlabonne/gemma-3-27b-it-abliterated")
 
 processor = AutoProcessor.from_pretrained(model_id, padding_side="left")
 model = Gemma3ForConditionalGeneration.from_pretrained(
     model_id,
     device_map="auto",
     torch_dtype=torch.bfloat16,
-    attn_implementation="eager"
+    attn_implementation="eager"  # 가능하다면 "flash_attention_2"로 변경
 )
 MAX_NUM_IMAGES = int(os.getenv("MAX_NUM_IMAGES", "5"))
 
@@ -284,7 +296,7 @@ def validate_media_constraints(message: dict, history: list[dict]) -> bool:
 
 
 ##############################################################################
-# 비디오 처리
+# 비디오 처리 - 임시 파일 추적 코드 추가
 ##############################################################################
 def downsample_video(video_path: str) -> list[tuple[Image.Image, float]]:
     vidcap = cv2.VideoCapture(video_path)
@@ -298,6 +310,8 @@ def downsample_video(video_path: str) -> list[tuple[Image.Image, float]]:
         success, image = vidcap.read()
         if success:
             image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+            # 이미지 크기 줄이기 추가
+            image = cv2.resize(image, (0, 0), fx=0.5, fy=0.5)
             pil_image = Image.fromarray(image)
             timestamp = round(i / fps, 2)
             frames.append((pil_image, timestamp))
@@ -308,17 +322,20 @@ def downsample_video(video_path: str) -> list[tuple[Image.Image, float]]:
     return frames
 
 
-def process_video(video_path: str) -> list[dict]:
+def process_video(video_path: str) -> tuple[list[dict], list[str]]:
     content = []
+    temp_files = []  # 임시 파일 추적을 위한 리스트
+    
     frames = downsample_video(video_path)
     for frame in frames:
         pil_image, timestamp = frame
         with tempfile.NamedTemporaryFile(delete=False, suffix=".png") as temp_file:
             pil_image.save(temp_file.name)
+            temp_files.append(temp_file.name)  # 추적을 위해 경로 저장
             content.append({"type": "text", "text": f"Frame {timestamp}:"})
             content.append({"type": "image", "url": temp_file.name})
-    logger.debug(f"{content=}")
-    return content
+    
+    return content, temp_files
 
 
 ##############################################################################
@@ -360,9 +377,11 @@ def is_document_file(file_path: str) -> bool:
     )
 
 
-def process_new_user_message(message: dict) -> list[dict]:
+def process_new_user_message(message: dict) -> tuple[list[dict], list[str]]:
+    temp_files = []  # 임시 파일 추적용 리스트
+    
     if not message["files"]:
-        return [{"type": "text", "text": message["text"]}]
+        return [{"type": "text", "text": message["text"]}], temp_files
 
     video_files = [f for f in message["files"] if is_video_file(f)]
     image_files = [f for f in message["files"] if is_image_file(f)]
@@ -385,19 +404,21 @@ def process_new_user_message(message: dict) -> list[dict]:
         content_list.append({"type": "text", "text": pdf_markdown})
 
     if video_files:
-        content_list += process_video(video_files[0])
-        return content_list
+        video_content, video_temp_files = process_video(video_files[0])
+        content_list += video_content
+        temp_files.extend(video_temp_files)
+        return content_list, temp_files
 
     if "<image>" in message["text"] and image_files:
         interleaved_content = process_interleaved_images({"text": message["text"], "files": image_files})
         if content_list and content_list[0]["type"] == "text":
             content_list = content_list[1:]
-        return interleaved_content + content_list
+        return interleaved_content + content_list, temp_files
     else:
         for img_path in image_files:
             content_list.append({"type": "image", "url": img_path})
 
-    return content_list
+    return content_list, temp_files
 
 
 ##############################################################################
@@ -429,6 +450,25 @@ def process_history(history: list[dict]) -> list[dict]:
     return messages
 
 
+##############################################################################
+# 모델 생성 함수에서 OOM 캐치
+##############################################################################
+def _model_gen_with_oom_catch(**kwargs):
+    """
+    별도 스레드에서 OutOfMemoryError를 잡아주기 위해
+    """
+    try:
+        model.generate(**kwargs)
+    except torch.cuda.OutOfMemoryError:
+        raise RuntimeError(
+            "[OutOfMemoryError] GPU 메모리가 부족합니다. "
+            "Max New Tokens을 줄이거나, 프롬프트 길이를 줄여주세요."
+        )
+    finally:
+        # 생성 완료 후 한번 더 캐시 비우기
+        clear_cuda_cache()
+
+
 ##############################################################################
 # 메인 추론 함수 (web search 체크 시 자동 키워드추출->검색->결과 system msg)
 ##############################################################################
@@ -446,6 +486,8 @@ def run(
         yield ""
         return
 
+    temp_files = []  # 임시 파일 추적용
+    
     try:
         combined_system_msg = ""
 
@@ -481,7 +523,9 @@ def run(
 
         messages.extend(process_history(history))
 
-        user_content = process_new_user_message(message)
+        user_content, user_temp_files = process_new_user_message(message)
+        temp_files.extend(user_temp_files)  # 임시 파일 추적
+        
         for item in user_content:
             if item["type"] == "text" and len(item["text"]) > MAX_CONTENT_CHARS:
                 item["text"] = item["text"][:MAX_CONTENT_CHARS] + "\n...(truncated)..."
@@ -494,7 +538,13 @@ def run(
             return_dict=True,
             return_tensors="pt",
         ).to(device=model.device, dtype=torch.bfloat16)
-
+        
+        # 입력 토큰 수 제한 추가
+        if inputs.input_ids.shape[1] > MAX_INPUT_LENGTH:
+            inputs.input_ids = inputs.input_ids[:, -MAX_INPUT_LENGTH:]
+            if 'attention_mask' in inputs:
+                inputs.attention_mask = inputs.attention_mask[:, -MAX_INPUT_LENGTH:]
+        
         streamer = TextIteratorStreamer(processor, timeout=30.0, skip_prompt=True, skip_special_tokens=True)
         gen_kwargs = dict(
             inputs,
@@ -513,22 +563,24 @@ def run(
     except Exception as e:
         logger.error(f"Error in run: {str(e)}")
         yield f"죄송합니다. 오류가 발생했습니다: {str(e)}"
-
-
-##############################################################################
-# [추가] 별도 함수에서 model.generate(...)를 호출, OOM 캐치
-##############################################################################
-def _model_gen_with_oom_catch(**kwargs):
-    """
-    별도 스레드에서 OutOfMemoryError를 잡아주기 위해
-    """
-    try:
-        model.generate(**kwargs)
-    except torch.cuda.OutOfMemoryError:
-        raise RuntimeError(
-            "[OutOfMemoryError] GPU 메모리가 부족합니다. "
-            "Max New Tokens을 줄이거나, 프롬프트 길이를 줄여주세요."
-        )
+    
+    finally:
+        # 임시 파일 삭제
+        for temp_file in temp_files:
+            try:
+                if os.path.exists(temp_file):
+                    os.unlink(temp_file)
+                    logger.info(f"Deleted temp file: {temp_file}")
+            except Exception as e:
+                logger.warning(f"Failed to delete temp file {temp_file}: {e}")
+        
+        # 명시적 메모리 정리
+        try:
+            del inputs, streamer
+        except:
+            pass
+        
+        clear_cuda_cache()
 
 
 ##############################################################################
@@ -632,12 +684,10 @@ css = """
     width: 100% !important;
     max-width: none !important; /* 1200px 제한 제거 */
 }
-
 .fillable {
     width: 100% !important; 
     max-width: 100% !important; 
 }
-
 /* 2) 배경을 연하고 투명한 파스텔 톤 그라디언트로 변경 */
 body {
     background: #f5f5f5; /* 그라디언트 대신 단색 사용 */
@@ -646,7 +696,6 @@ body {
     font-family: 'Helvetica Neue', Helvetica, Arial, sans-serif;
     color: #333;
 }
-
 /* 버튼 색상도 기존의 짙은 붉은-주황 → 파스텔 계열로 연���게 */
 button, .btn {
     background: #ffb6c1 !important; /* 그라디언트 대신 단색 사용 */
@@ -660,24 +709,21 @@ button, .btn {
     cursor: pointer;
     /* transition: transform 0.2s ease-in-out; - 혹시 모를 문제 부분 제거 */
 }
-
 button:hover, .btn:hover {
     /* transform: scale(1.03); - 혹시 모를 문제 부분 제거 */
     background: #ff69b4 !important;
 }
-
 #examples_container {
     margin: auto;
     width: 90%;
 }
-
 #examples_row {
     justify-content: center;
 }
 """
 
 title_html = """
-<h1 align="center" style="margin-bottom: 0.2em; font-size: 1.6em;"> 🤗 Gemma3-uncensored-R27B </h1>
+<h1 align="center" style="margin-bottom: 0.2em; font-size: 1.6em;"> 🤗 Gemma3-uncensored-R12B </h1>
 <p align="center" style="font-size:1.1em; color:#555;">
     ✅Agentic AI Platform ✅Reasoning & Uncensored ✅Multimodal & VLM ✅Deep-Research & RAG <br>
     Operates on an ✅'NVIDIA A100 GPU' as an independent local server, enhancing security and preventing information leakage.<br>