# app.py
import os
import gradio as gr
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM
import spaces  # Mandatory library for Hugging Face ZeroGPU [1]

MODEL_ID = "DevStudio-AI/Devstudio-Coder-1.5B"

print("Loading tokenizer and base model...")
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)

# We load the model in 16-bit on CPU first
# ZeroGPU will automatically move the model to the GPU when the decorated function runs [1]
model = AutoModelForCausalLM.from_pretrained(
    MODEL_ID,
    torch_dtype=torch.float16,
    device_map="cpu"
)
print("Model successfully loaded on CPU. Awaiting ZeroGPU allocation...")

# The @spaces.GPU decorator dynamically requests Nvidia A100 resources for this call [1]
@spaces.GPU
def generate_code(prompt, temperature, max_tokens):
    try:
        # Move model to CUDA dynamically inside the GPU context [1]
        model.to("cuda")
        
        inputs = tokenizer(prompt, return_tensors="pt").to("cuda")
        
        outputs = model.generate(
            **inputs,
            max_new_tokens=int(max_tokens),
            temperature=float(temperature),
            do_sample=True,
            eos_token_id=tokenizer.eos_token_id
        )
        
        # Isolate and decode newly generated tokens
        generated_ids = outputs[0][inputs["input_ids"].shape[1]:]
        return tokenizer.decode(generated_ids, skip_special_tokens=True)
    except Exception as e:
        return f"Error during generation: {str(e)}"

# Define the Gradio web interface
demo = gr.Interface(
    fn=generate_code,
    inputs=[
        gr.Textbox(label="Prompt", placeholder="Enter your prompt here..."),
        gr.Slider(minimum=0.1, maximum=1.0, value=0.3, label="Temperature"),
        gr.Slider(minimum=64, maximum=2048, value=1024, step=64, label="Max Tokens")
    ],
    outputs=gr.Textbox(label="Generated Code"),
    title="DevStudio-1.5B API Engine",
    description="Static HTML + Tailwind CSS specialized completion endpoint running on ZeroGPU."
)

demo.launch()