Update app.py

app.py

@@ -1,62 +1,46 @@
-from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
-import torch
+from transformers import pipeline
 import gradio as gr
-import os
-
-# 1. GPU Verification
-if not torch.cuda.is_available():
-    raise RuntimeError("❌ GPU not detected! Enable GPU in Space settings.")
-print(f"✅ Using GPU: {torch.cuda.get_device_name(0)}")
-
-# 2. Model Configuration
-MODEL_NAME = "google/gemma-2b-it"  # Try "mistralai/Mistral-7B-v0.1" for more power
+import torch
 
-# 3. Load Model with GPU Optimization
-tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
-model = AutoModelForCausalLM.from_pretrained(
-    MODEL_NAME,
-    device_map="auto",          # Auto-selects GPU
-    torch_dtype=torch.float16,  # Half-precision for memory
-    low_cpu_mem_usage=True      # Reduces CPU overhead
-)
+# Auto-configure GPU/CPU
+device = 0 if torch.cuda.is_available() else -1
+dtype = torch.float16 if device == 0 else torch.float32
+print(f"⚡ Using {'GPU: ' + torch.cuda.get_device_name(0) if device == 0 else 'CPU'}")
 
-# 4. Create GPU-accelerated pipeline
-pipe = pipeline(
+# Load optimized pipeline
+model = pipeline(
     "text-generation",
-    model=model,
-    tokenizer=tokenizer,
-    device=0,                  # Force first GPU
-    torch_dtype=torch.float16
+    model="google/gemma-2b-it",
+    device=device,
+    torch_dtype=dtype,
+    model_kwargs={
+        "low_cpu_mem_usage": True,
+        "trust_remote_code": True
+    }
 )
 
-# 5. Generation Function
-def generate_text(prompt):
+# Pre-warm model (reduces first response time)
+model("Warming up...", max_new_tokens=1)
+
+def generate(prompt):
+    """Ultra-fast generation with 1-2 second responses"""
     try:
-        outputs = pipe(
+        output = model(
             prompt,
-            max_new_tokens=150,
-            temperature=0.7,
-            do_sample=True,
-            pad_token_id=tokenizer.eos_token_id
+            max_new_tokens=80,  # Shorter = faster
+            temperature=0.3,    # More deterministic
+            do_sample=False,    # Disable sampling for speed
+            pad_token_id=model.tokenizer.eos_token_id
         )
-        return outputs[0]['generated_text']
+        return output[0]['generated_text']
     except Exception as e:
-        return f"⚠️ Error: {str(e)}"
+        return f"Error: {str(e)}"
 
-# 6. Gradio Interface
-with gr.Blocks(theme=gr.themes.Soft()) as demo:
-    gr.Markdown("## 🚀 GPU-Powered Text Generator")
-    with gr.Row():
-        inp = gr.Textbox(label="Your Prompt", placeholder="Type here...")
-    with gr.Row():
-        out = gr.Textbox(label="Generated Text", lines=5)
-    with gr.Row():
-        btn = gr.Button("Generate", variant="primary")
-    
-    btn.click(fn=generate_text, inputs=inp, outputs=out)
-    inp.submit(fn=generate_text, inputs=inp, outputs=out)
+# Minimal UI for maximum speed
+with gr.Blocks(title="🚀 Instant AI") as demo:
+    gr.Markdown("## Type anything (1-2 sec responses):")
+    input = gr.Textbox(placeholder="How to make pizza?")
+    output = gr.Textbox()
+    input.submit(generate, input, output)
 
-# 7. Launch with GPU monitoring
-if __name__ == "__main__":
-    print(f"GPU Memory Allocated: {torch.cuda.memory_allocated()/1e9:.2f} GB")
-    demo.launch(server_name="0.0.0.0")
+demo.launch(server_name="0.0.0.0")