Update app.py

app.py

@@ -34,7 +34,8 @@ cache_dir = "/data"
 # )
 model = Llama4ForConditionalGeneration.from_pretrained(
     model_name,
-    attn_implementation="flex_attention",
+    # flex_attention is only needed for image
+    # attn_implementation="flex_attention", 
     # gguf_file=filename,
     cache_dir = cache_dir,
     torch_dtype=torch_dtype,
@@ -83,11 +84,11 @@ def generate(prompt, history):
     inputs = processor.apply_chat_template(
         messages,
         add_generation_prompt=True,
-        tokenize=True,
+        # tokenize=True,
         return_dict=True,
         return_tensors="pt",
     ).to(gpu_model.device)
-    outputs = gpu_model.generate(
+    outputs = model.generate(
         **inputs,
         max_new_tokens=512,
     )