Update app.py

app.py

@@ -1,9 +1,12 @@
 import gradio as gr
 import torch
 from dataclasses import dataclass
-from transformers import AutoTokenizer, PretrainedConfig, pipeline
+from transformers import AutoTokenizer, PretrainedConfig, pipeline, GenerationConfig
 from optimum.onnxruntime import ORTModelForCausalLM
 import onnx
+import logging
+
+logging.basicConfig(level=logging.INFO)
 
 # -----------------------------------------------------------------------------
 # Configuration and Special Tokens
@@ -40,7 +43,6 @@ class Sam3Config(PretrainedConfig):
     _attn_implementation_internal: str = "eager"
     is_encoder_decoder: bool = False
     
-    # These are the required attributes for ORTModelForCausalLM
     hidden_size: int = 384
     num_attention_heads: int = 6
 
@@ -55,8 +57,6 @@ class Sam3Config(PretrainedConfig):
         self.input_modality = input_modality
         self.head_type = head_type
         self.version = version
-        
-        # Ensure hidden_size and num_attention_heads are set correctly
         self.hidden_size = self.d_model
         self.num_attention_heads = self.n_heads
 
@@ -64,28 +64,45 @@ class Sam3Config(PretrainedConfig):
 model_config = Sam3Config()
 
 # Load the ONNX model by providing the configuration
-model = ORTModelForCausalLM.from_pretrained(
-    "Smilyai-labs/Sam-3.0-2-onnx",
-    config=model_config,
-    trust_remote_code=True
-)
+try:
+    model = ORTModelForCausalLM.from_pretrained(
+        "Smilyai-labs/Sam-3.0-2-onnx",
+        config=model_config,
+        trust_remote_code=True,
+    )
+    logging.info("ONNX model loaded successfully.")
+
+    # Fix the use_cache issue by setting it to False if the model doesn't support it
+    if not getattr(model, "_is_stateful", True):
+        logging.warning("Model does not support `_is_stateful`, setting `use_cache=False` for generation.")
+        model.generation_config.use_cache = False
+    
+except Exception as e:
+    logging.error(f"Failed to load ONNX model: {e}")
+    raise e
 
 # Define a function to generate text
 def generate_text(prompt, max_length=128, temperature=0.8, top_k=60, top_p=0.9):
+    # Set generation parameters
+    gen_config = GenerationConfig(
+        max_length=max_length,
+        temperature=temperature,
+        top_k=top_k,
+        top_p=top_p,
+        do_sample=True,
+        use_cache=False, # Explicitly disable cache to avoid the error
+    )
+
     gen_pipeline = pipeline(
         "text-generation",
         model=model,
         tokenizer=tokenizer,
-        device=device
+        device=device,
+        generation_config=gen_config
     )
     
     generated_text = gen_pipeline(
         prompt,
-        max_length=max_length,
-        temperature=temperature,
-        top_k=top_k,
-        top_p=top_p,
-        do_sample=True,
     )
     return generated_text[0]["generated_text"]