Update app.py

app.py

@@ -57,6 +57,7 @@ class Sam3Config(PretrainedConfig):
         self.input_modality = input_modality
         self.head_type = head_type
         self.version = version
+        
         self.hidden_size = self.d_model
         self.num_attention_heads = self.n_heads
 
@@ -71,11 +72,6 @@ try:
         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}")
@@ -83,14 +79,15 @@ except Exception as 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
+    # Set generation parameters within a GenerationConfig object
+    # We set use_cache=False here to bypass the onnx export issue
     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
+        use_cache=False, 
     )
 
     gen_pipeline = pipeline(
@@ -98,11 +95,12 @@ def generate_text(prompt, max_length=128, temperature=0.8, top_k=60, top_p=0.9):
         model=model,
         tokenizer=tokenizer,
         device=device,
-        generation_config=gen_config
     )
     
+    # Pass all generation parameters to the pipeline
     generated_text = gen_pipeline(
         prompt,
+        **gen_config.to_dict() 
     )
     return generated_text[0]["generated_text"]