Update app.py

app.py

@@ -2,9 +2,11 @@ import os
 from huggingface_hub import login
 import gradio as gr
 from PIL import Image
-from transformers import AutoProcessor, AutoModelForImageTextToText, pipeline
+from transformers import AutoProcessor, AutoModelForImageTextToText
 import torch
 import spaces
+from peft import PeftModel
+from transformers import AutoModelForCausalLM
 
 # Function to process vision information
 def process_vision_info(messages: list[dict]) -> list[Image.Image]:
@@ -33,11 +35,10 @@ def load_image_model():
 
 # Load text model on CPU
 def load_text_model():
-    return pipeline(
-        "text-generation",
-        model="tarotscientist/llama-2-7b-tarotreader",
-        device=-1  # Force CPU
-    )
+    base_model = AutoModelForCausalLM.from_pretrained("google/gemma-2-2b")
+    model = PeftModel.from_pretrained(base_model, "soonbob/gemma-2-2b-tarot")
+    tokenizer = AutoTokenizer.from_pretrained("google/gemma-2-2b")
+    return model, tokenizer
 
 # Generate card description with ZeroGPU
 @spaces.GPU
@@ -76,7 +77,7 @@ def generate_description(sample, model, processor):
 
 # Generate tarot interpretation with ZeroGPU
 @spaces.GPU
-def generate_interpretation(question, cards, model):
+def generate_interpretation(question, cards, model, tokenizer):
     prompt = f"""Analyze this tarot reading for the question: {question}
 
 Cards:
@@ -89,8 +90,9 @@ Provide a professional interpretation covering:
 - Practical advice
 - Potential outcomes"""
     # Use GPU for this inference call
-    response = model(prompt, max_length=1000, temperature=0.8, top_p=0.95)[0]['generated_text']
-    return response
+    input_ids = tokenizer(prompt, return_tensors="pt").to("cuda")
+    outputs = model.generate(**input_ids, max_new_tokens=32)
+    return (tokenizer.decode(outputs[0]))
 
 def main():
     """
@@ -105,7 +107,7 @@ def main():
 
     # Load models on CPU
     image_processor, image_model = load_image_model()
-    text_model = load_text_model()
+    text_model, text_tokenizer = load_text_model()
 
     # Define the tarot processing function
     def process_tarot(question, reason_img, result_img, recommendation_img):
@@ -130,7 +132,7 @@ def main():
                 output = "### Identifing Card Name...\n"
 
             # Generate the full interpretation using GPU
-            interpretation = generate_interpretation(question, cards, text_model)
+            interpretation = generate_interpretation(question, cards, text_model, text_tokenizer)
 
             # Format the output
             output += "### Card Analysis\n"