Update app.py

app.py

@@ -2,18 +2,42 @@ import gradio as gr
 import torch
 import random
 import transformers
-from transformers import T5Tokenizer, T5ForConditionalGeneration
+from transformers import T5Tokenizer
+from optimum.onnxruntime import ORTModelForSeq2SeqLM
 
+# --- CUDA / Provider Setup ---
 if torch.cuda.is_available():
     device = "cuda"
-    print("Using GPU")
+    provider = "CUDAExecutionProvider"
+    print(f"Using GPU with {provider}")
 else:
     device = "cpu"
+    provider = "CPUExecutionProvider"
     print("Using CPU")
 
+# Load Model with ONNX Runtime for Execution Provider support
+# Note: This requires optimum installed: pip install optimum[onnxruntime-gpu]
+try:
+    model = ORTModelForSeq2SeqLM.from_pretrained(
+        "roborovski/superprompt-v1",
+        provider=provider,
+        export=False # Set True if you want to force generate ONNX files from pytorch
+    )
+    print(f"Model loaded successfully using {provider}")
+except Exception as e:
+    print(f"Failed to load ONNX model: {e}")
+    print("Falling back to standard PyTorch model...")
+    from transformers import T5ForConditionalGeneration
+    model = T5ForConditionalGeneration.from_pretrained(
+        "roborovski/superprompt-v1", legacy=False,
+        device_map="auto", 
+        torch_dtype="auto"
+    )
+    # Standard torch model doesn't use ExecutionProvider string, but we keep the logic intact
+    if device == "cuda":
+        model.to(device)
+
 tokenizer = T5Tokenizer.from_pretrained("roborovski/superprompt-v1")
-model = T5ForConditionalGeneration.from_pretrained("roborovski/superprompt-v1", device_map="auto", torch_dtype="auto")
-model.to(device)
 
 def generate(your_prompt, task_prefix, max_new_tokens, repetition_penalty, temperature, model_precision_type, top_p, top_k, seed):
     
@@ -21,18 +45,21 @@ def generate(your_prompt, task_prefix, max_new_tokens, repetition_penalty, tempe
         seed = random.randint(1, 2**32-1)
     transformers.set_seed(seed)
     
-    if model_precision_type == "fp16":
-        dtype = torch.float16
-    elif model_precision_type == "fp32":
-        dtype = torch.float32
-
-    model.to(dtype)
-
+    # ONNX Runtime models usually manage their own precision/quantization via the file loaded,
+    # but we can leave the UI option for users to switch logic if they were swapping models.
+    # For this specific implementation, the precision is largely determined by the loaded provider/weights.
+    
     repetition_penalty = float(repetition_penalty)
 
     input_text = f"{task_prefix}: {your_prompt}"
+    # ONNX models generally handle input tensors on the device they were initialized with automatically
     input_ids = tokenizer(input_text, return_tensors="pt").input_ids.to(device)
         
+    # ONNX Runtime generate function might differ slightly in arguments, but standard transformers args usually map over.
+    # We ensure we pass the device properly for PyTorch fallback.
+    if hasattr(model, 'device'):
+        input_ids = input_ids.to(model.device)
+        
     outputs = model.generate(
         input_ids,
         max_new_tokens=max_new_tokens,
@@ -74,7 +101,7 @@ gr.Interface(
     inputs=[your_prompt, task_prefix, max_new_tokens, repetition_penalty, temperature, model_precision_type, top_p, top_k, seed],
     outputs=gr.Textbox(label="Better Prompt"),
     title="SuperPrompt-v1",
-    description='Make your prompts more detailed! <br> <a href="https://github.com/Nick088Official/SuperPrompt-v1">Github Repository & Model used</a> <br> <a href="https://brianfitzgerald.xyz/prompt-augmentation/">Model Blog</a> <br> Hugging Face Space made by [Nick088](https://linktr.ee/Nick088)',
+    description='Make your prompts more detailed! <br> <br> Hugging Face Space made by Nick088 improved bu NeoPy/BF667',
     examples=examples,
     theme="NeoPy/Soft"
-).launch(share=True, mcp_mode=True)
+).launch(share=True)