Create app.py

app.py

@@ -0,0 +1,161 @@
+import gradio as gr
+from gradio import Server
+import spaces
+import torch
+from transformers import AutoTokenizer, AutoModelForCausalLM
+
+# ─────────────────────────────────────────────
+# 1. MODEL SETUP
+# ─────────────────────────────────────────────
+MODEL_ID = "Smilyai-labs/Mira-1-large"
+
+print(f"Loading tokenizer: {MODEL_ID}")
+tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True)
+
+print(f"Loading model: {MODEL_ID}")
+model = AutoModelForCausalLM.from_pretrained(
+    MODEL_ID,
+    torch_dtype=torch.float16,
+    device_map="auto",       # ZeroGPU manages CUDA device
+    trust_remote_code=True,  # Needed for Qwen-based custom archs
+)
+model.eval()
+
+# ─────────────────────────────────────────────
+# 2. INFERENCE FUNCTION (ZeroGPU decorated)
+# ─────────────────────────────────────────────
+@spaces.GPU(duration=120)
+def generate(
+    prompt: str,
+    system_prompt: str = "You are a helpful assistant.",
+    max_new_tokens: int = 512,
+    temperature: float = 0.7,
+    top_p: float = 0.9,
+    do_sample: bool = True,
+) -> str:
+    """
+    Generate a text response from Mira-1-Large.
+
+    Args:
+        prompt: The user message / prompt to send to the model.
+        system_prompt: System-level instruction for the model.
+        max_new_tokens: Maximum number of tokens to generate.
+        temperature: Sampling temperature (higher = more creative).
+        top_p: Nucleus sampling probability mass.
+        do_sample: Whether to use sampling (True) or greedy decoding (False).
+
+    Returns:
+        The model's text response as a string.
+    """
+    # Build chat-style messages (Qwen uses apply_chat_template)
+    messages = [
+        {"role": "system", "content": system_prompt},
+        {"role": "user",   "content": prompt},
+    ]
+
+    # Qwen / Mira chat template
+    text = tokenizer.apply_chat_template(
+        messages,
+        tokenize=False,
+        add_generation_prompt=True,
+    )
+
+    inputs = tokenizer(text, return_tensors="pt").to(model.device)
+
+    with torch.no_grad():
+        output_ids = model.generate(
+            **inputs,
+            max_new_tokens=max_new_tokens,
+            temperature=temperature,
+            top_p=top_p,
+            do_sample=do_sample,
+            pad_token_id=tokenizer.eos_token_id,
+        )
+
+    # Decode only the newly generated tokens
+    new_tokens = output_ids[0][inputs["input_ids"].shape[1]:]
+    response = tokenizer.decode(new_tokens, skip_special_tokens=True)
+    return response
+
+
+# ─────────────────────────────────────────────
+# 3. STREAMING INFERENCE (SSE / token-by-token)
+# ─────────────────────────────────────────────
+@spaces.GPU(duration=120)
+def generate_stream(
+    prompt: str,
+    system_prompt: str = "You are a helpful assistant.",
+    max_new_tokens: int = 512,
+    temperature: float = 0.7,
+    top_p: float = 0.9,
+) -> str:
+    """
+    Stream a text response token-by-token from Mira-1-Large via SSE.
+
+    Args:
+        prompt: The user message / prompt.
+        system_prompt: System-level instruction for the model.
+        max_new_tokens: Maximum number of tokens to generate.
+        temperature: Sampling temperature.
+        top_p: Nucleus sampling probability mass.
+
+    Yields:
+        Partial response strings, growing with each new token.
+    """
+    from transformers import TextIteratorStreamer
+    from threading import Thread
+
+    messages = [
+        {"role": "system", "content": system_prompt},
+        {"role": "user",   "content": prompt},
+    ]
+    text = tokenizer.apply_chat_template(
+        messages, tokenize=False, add_generation_prompt=True
+    )
+    inputs = tokenizer(text, return_tensors="pt").to(model.device)
+
+    streamer = TextIteratorStreamer(
+        tokenizer, skip_prompt=True, skip_special_tokens=True
+    )
+
+    gen_kwargs = dict(
+        **inputs,
+        max_new_tokens=max_new_tokens,
+        temperature=temperature,
+        top_p=top_p,
+        do_sample=True,
+        streamer=streamer,
+        pad_token_id=tokenizer.eos_token_id,
+    )
+
+    thread = Thread(target=model.generate, kwargs=gen_kwargs)
+    thread.start()
+
+    partial = ""
+    for new_text in streamer:
+        partial += new_text
+        yield partial
+
+
+# ─────────────────────────────────────────────
+# 4. gr.Server — REST API + OPTIONAL SWAGGER UI
+# ───────────────────��─────────────────────────
+app = Server(
+    title="Mira-1-Large API",
+    summary="ZeroGPU-backed REST API for Smilyai-labs/Mira-1-large (Qwen arch)",
+    version="1.0.0",
+)
+
+# Register as Gradio API endpoints (queued, SSE-streaming capable)
+app.api(generate,        name="generate")        # POST /gradio_api/call/generate
+app.api(generate_stream, name="generate_stream") # POST /gradio_api/call/generate_stream
+
+# Optional: plain FastAPI GET health-check route
+@app.get("/health")
+def health():
+    return {"status": "ok", "model": MODEL_ID}
+
+# ─────────────────────────────────────────────
+# 5. LAUNCH
+# ─────────────────────────────────────────────
+app.launch()