Update app.py

app.py

@@ -1,6 +1,8 @@
 import gradio as gr
 import spaces
 from transformers import AutoModelForCausalLM, AutoTokenizer
+from transformers import AutoProcessor, Llama4ForConditionalGeneration
+
 import torch
 
 #Qwen/Qwen2.5-14B-Instruct-1M
@@ -14,14 +16,23 @@ filename = "Llama-4-Scout-17B-16E-Instruct-UD-IQ2_XXS.gguf"
 torch_dtype = torch.bfloat16 # could be torch.float16 or torch.bfloat16 torch.float32 too
 cache_dir = "/data"
 
-model = AutoModelForCausalLM.from_pretrained(
+# model = AutoModelForCausalLM.from_pretrained(
+#     model_name,
+#     # subfolder=subfolder,
+#     gguf_file=filename,
+#     torch_dtype=torch_dtype,
+#     device_map="auto",
+#     cache_dir = cache_dir,
+# )
+model = Llama4ForConditionalGeneration.from_pretrained(
     model_name,
-    # subfolder=subfolder,
+    attn_implementation="flex_attention",
     gguf_file=filename,
     torch_dtype=torch_dtype,
     device_map="auto",
     cache_dir = cache_dir,
 )
+
 tokenizer = AutoTokenizer.from_pretrained(model_name
     , gguf_file=filename
     # , subfolder=subfolder
@@ -44,7 +55,8 @@ def generate(prompt, history):
     ]
     text = tokenizer.apply_chat_template(
         messages,
-        tokenize=False,
+        # tokenize=False,
+        tokenize=True,
         add_generation_prompt=True
     )
     model_inputs = tokenizer([text], return_tensors="pt").to(model.device)