app.py

from transformers import pipeline
import gradio as gr
import torch
import os

# ===== SMART DEVICE CONFIGURATION =====
def get_best_device():
    if torch.cuda.is_available():
        torch.backends.cudnn.benchmark = True
        return 0, torch.float16  # GPU with half precision
    return -1, torch.float32  # CPU fallback

device, dtype = get_best_device()
device_name = "GPU: " + torch.cuda.get_device_name(0) if device == 0 else "CPU"
print(f"⚡ Running on: {device_name}")

# ===== ERROR-PROOF MODEL LOADING =====
try:
    # Correct pipeline configuration (fixed trust_remote_code)
    model = pipeline(
        "text-generation",
        model="google/gemma-2b-it",
        device=device,
        torch_dtype=dtype,
        model_kwargs={
            "low_cpu_mem_usage": True,
            "trust_remote_code": True  # Correct placement
        }
    )
    
    # Pre-warm model
    model("Warmup", max_new_tokens=1)
    
except Exception as e:
    # Simplified fallback (removes duplicate trust_remote_code)
    model = pipeline(
        "text-generation",
        model="google/gemma-2b-it",
        device=device,
        torch_dtype=dtype
    )

# ===== OPTIMIZED GENERATION =====
def generate(prompt):
    try:
        return model(
            prompt,
            max_new_tokens=60,
            temperature=0.2,
            do_sample=False,
            pad_token_id=model.tokenizer.eos_token_id
        )[0]['generated_text']
    except Exception as e:
        return f"⚠️ Error: {str(e)}"

# ===== SIMPLE INTERFACE =====
with gr.Blocks() as demo:
    gr.Markdown("## Ask anything (1-2 second responses)")
    input = gr.Textbox(label="Your question")
    output = gr.Textbox(label="Answer")
    input.submit(generate, input, output)

demo.launch(server_name="0.0.0.0")