add secrete

model/model_manager.py

@@ -2,7 +2,7 @@ import concurrent.futures
 import random
 import gradio as gr
 import requests
-import io, base64, json
+import io, base64, json, os
 import spaces
 from PIL import Image
 from .models import IMAGE_GENERATION_MODELS, IMAGE_EDITION_MODELS, VIDEO_GENERATION_MODELS, MUSEUM_UNSUPPORTED_MODELS, DESIRED_APPEAR_MODEL, load_pipeline
@@ -33,8 +33,8 @@ class ModelManager:
         model_id = "meta-llama/Meta-Llama-Guard-2-8B"
         device = "cuda"
         dtype = torch.bfloat16
-        tokenizer = AutoTokenizer.from_pretrained(model_id)
-        model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=dtype, device_map=device)
+        tokenizer = AutoTokenizer.from_pretrained(model_id, token=os.environ['HF_GUARD'])
+        model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=dtype, device_map=device, token=os.environ['HF_GUARD'])
         chat = [{"role": "user", "content": prompt}]
         input_ids = tokenizer.apply_chat_template(chat, return_tensors="pt").to(device)
         output = model.generate(input_ids=input_ids, max_new_tokens=100, pad_token_id=0)