Update app.py

app.py

@@ -1,12 +1,14 @@
 import re
 import threading
 import gc
+import os
 import torch
 
 import gradio as gr
 import spaces
 import transformers
 from transformers import pipeline, AutoModelForCausalLM, AutoTokenizer
+from huggingface_hub import login
 
 # 모델 메모리 관리 및 최적화를 위한 설정
 DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
@@ -28,6 +30,17 @@ available_models = {
 pipe = None
 current_model_name = None
 
+# Hugging Face 토큰으로 로그인 시도
+try:
+    hf_token = os.getenv("HF_TOKEN")
+    if hf_token:
+        login(token=hf_token)
+        print("Hugging Face에 성공적으로 로그인했습니다.")
+    else:
+        print("경고: HF_TOKEN 환경 변수가 설정되지 않았습니다.")
+except Exception as e:
+    print(f"Hugging Face 로그인 에러: {str(e)}")
+
 # 최종 답변을 감지하기 위한 마커
 ANSWER_MARKER = "**답변**"
 
@@ -143,6 +156,14 @@ def load_model(model_names):
     
     # 모델 로드 (크기에 따라 최적화된 설정 적용)
     try:
+        # HF_TOKEN 환경 변수 확인
+        hf_token = os.getenv("HF_TOKEN")
+        # 공통 매개변수
+        common_params = {
+            "token": hf_token,  # 접근 제한 모델을 위한 토큰
+            "trust_remote_code": True,
+        }
+        
         # BF16 정밀도 사용 (A100에 최적화)
         if config["quantization"]:
             # 양자화 적용
@@ -159,9 +180,9 @@ def load_model(model_names):
                 torch_dtype=DTYPE,
                 quantization_config=quantization_config if config["quantization"] else None,
                 offload_folder="offload" if config["offload"] else None,
-                trust_remote_code=True
+                **common_params
             )
-            tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
+            tokenizer = AutoTokenizer.from_pretrained(model_name, **common_params)
             
             pipe = pipeline(
                 "text-generation",
@@ -177,7 +198,7 @@ def load_model(model_names):
                 model=model_name,
                 device_map="auto",
                 torch_dtype=DTYPE,
-                trust_remote_code=True
+                **common_params
             )
         
         current_model_name = model_name
@@ -414,5 +435,12 @@ if __name__ == "__main__":
         print(f"현재 GPU: {torch.cuda.current_device()}")
         print(f"GPU 이름: {torch.cuda.get_device_name(0)}")
     
+    # HF_TOKEN 환경 변수 확인
+    hf_token = os.getenv("HF_TOKEN")
+    if hf_token:
+        print("HF_TOKEN 환경 변수가 설정되어 있습니다.")
+    else:
+        print("경고: HF_TOKEN 환경 변수가 설정되지 않았습니다. 제한된 모델에 접근할 수 없습니다.")
+    
     # 큐 사용 및 앱 실행
     demo.queue(max_size=10).launch()