Update app.py

app.py

@@ -3,29 +3,39 @@ from transformers import pipeline, AutoModelForCausalLM, AutoTokenizer, BitsAndB
 import torch
 from diffusers import StableDiffusionPipeline
 
+
 # Load the text model
 @st.cache_resource
 def load_text_model():
     model_name = "meta-llama/Llama-2-7b-chat-hf"
+    
     bnb_config = BitsAndBytesConfig(
         load_in_4bit=True,  
         bnb_4bit_compute_dtype=torch.float16,  
         bnb_4bit_use_double_quant=True
     )
-    tokenizer = AutoTokenizer.from_pretrained(model_name)
+
+    tokenizer = AutoTokenizer.from_pretrained(model_name, use_auth_token=HF_TOKEN)
     model = AutoModelForCausalLM.from_pretrained(
         model_name, 
+        load_in_4bit=True,
         quantization_config=bnb_config,  
-        device_map="auto"
+        device_map="auto",
+        use_auth_token=HF_TOKEN
     )
+    
     return pipeline("text-generation", model=model, tokenizer=tokenizer)
 
 # Load the image model
 @st.cache_resource
 def load_image_model():
     model_id = "runwayml/stable-diffusion-v1-5"
-    pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16)
+    
+    pipe = StableDiffusionPipeline.from_pretrained(
+        model_id, torch_dtype=torch.float16, use_auth_token=HF_TOKEN
+    )
     pipe.to("cuda")  # Use GPU
+    
     return pipe
 
 llama_pipe = load_text_model()
@@ -50,4 +60,3 @@ if st.button("Generate Blog"):
             image = image_pipe(keywords).images[0]
             st.subheader("Generated Image")
             st.image(image, caption="AI-Generated Image", use_column_width=True)
-