inference.py

#!/usr/bin/env python3
"""
OpenAI-compatible inference client for HF Playground.

Usage:
    python inference.py chat --model Qwen/Qwen2.5-3B-Instruct --message "Hello"
    python inference.py completion --model Qwen/Qwen2.5-3B-Instruct --prompt "Once upon"
    python inference.py image --model runwayml/stable-diffusion-v1-5 --prompt "A dragon"
    python inference.py tts --model suno/bark-small --text "Hello world"
    python inference.py models

Env vars:
    HF_SPACE_URL  — base URL (default: https://stanley-00-z-image-turbo.hf.space)
    HF_TOKEN      — optional API key
"""

DEFAULT_BASE_URL = "https://stanley-00-z-image-turbo.hf.space"
import argparse
import json
import os
import sys
import base64

try:
    from openai import OpenAI
except ImportError:
    print("Error: openai package not installed. Run: pip install openai")
    sys.exit(1)


def get_client(args):
    base_url = args.base_url or os.getenv("HF_SPACE_URL", DEFAULT_BASE_URL).rstrip("/")
    api_key = args.api_key or os.environ.get("HF_TOKEN", "local")
    return OpenAI(base_url=f"{base_url}/v1", api_key=api_key)


def cmd_models(args):
    client = get_client(args)
    models = client.models.list()
    print(f"{'ID':<55} {'Owner':<15}")
    print("-" * 70)
    for m in models.data:
        print(f"{m.id:<55} {getattr(m, 'owned_by', ''):<15}")


def cmd_chat(args):
    client = get_client(args)
    messages = [{"role": "user", "content": args.message}]
    if args.system:
        messages.insert(0, {"role": "system", "content": args.system})

    if args.stream:
        stream = client.chat.completions.create(
            model=args.model,
            messages=messages,
            max_tokens=args.max_tokens,
            temperature=args.temperature,
            stream=True,
        )
        for chunk in stream:
            delta = chunk.choices[0].delta
            if delta.content:
                print(delta.content, end="", flush=True)
        print()
    else:
        response = client.chat.completions.create(
            model=args.model,
            messages=messages,
            max_tokens=args.max_tokens,
            temperature=args.temperature,
        )
        print(response.choices[0].message.content)
        if response.usage:
            print(f"\n[tokens: {response.usage.prompt_tokens} prompt + {response.usage.completion_tokens} completion = {response.usage.total_tokens} total]")


def cmd_completion(args):
    client = get_client(args)
    if args.stream:
        stream = client.completions.create(
            model=args.model,
            prompt=args.prompt,
            max_tokens=args.max_tokens,
            temperature=args.temperature,
            stream=True,
        )
        for chunk in stream:
            text = chunk.choices[0].text
            if text:
                print(text, end="", flush=True)
        print()
    else:
        response = client.completions.create(
            model=args.model,
            prompt=args.prompt,
            max_tokens=args.max_tokens,
            temperature=args.temperature,
        )
        print(response.choices[0].text)
        if response.usage:
            print(f"\n[tokens: {response.usage.prompt_tokens} prompt + {response.usage.completion_tokens} completion = {response.usage.total_tokens} total]")


def cmd_image(args):
    client = get_client(args)
    response = client.images.generate(
        model=args.model,
        prompt=args.prompt,
        size=args.size,
    )
    if response.data and response.data[0].b64_json:
        img_bytes = base64.b64decode(response.data[0].b64_json)
        out_path = args.output or "generated_image.png"
        with open(out_path, "wb") as f:
            f.write(img_bytes)
        print(f"Saved to {out_path} ({len(img_bytes)} bytes)")
    elif response.data and response.data[0].url:
        print(f"Image URL: {response.data[0].url}")
    else:
        print("No image returned")


def cmd_tts(args):
    import requests as req

    base_url = args.base_url or os.environ.get("HF_SPACE_URL", DEFAULT_BASE_URL).rstrip("/")
    api_key = args.api_key or os.environ.get("HF_TOKEN", "local")

    resp = req.post(
        f"{base_url}/v1/audio/speech",
        headers={"Authorization": f"Bearer {api_key}", "Content-Type": "application/json"},
        json={"model": args.model, "input": args.text},
    )
    if resp.status_code == 200:
        out_path = args.output or "speech.wav"
        with open(out_path, "wb") as f:
            f.write(resp.content)
        print(f"Saved to {out_path} ({len(resp.content)} bytes)")
    else:
        print(f"Error {resp.status_code}: {resp.text}")


def main():
    parser = argparse.ArgumentParser(description="HF Playground — OpenAI-compatible inference client")
    parser.add_argument("--base-url", help=f"Server base URL (default: HF_SPACE_URL env or {DEFAULT_BASE_URL})")
    parser.add_argument("--api-key", help="API key (default: HF_TOKEN env or 'local')")
    sub = parser.add_subparsers(dest="command", help="Command to run")

    # models
    sub.add_parser("models", help="List available models")

    # chat
    p_chat = sub.add_parser("chat", help="Chat completion")
    p_chat.add_argument("--model", "-m", default="Qwen/Qwen2.5-3B-Instruct")
    p_chat.add_argument("--message", "-M", required=True, help="User message")
    p_chat.add_argument("--system", "-s", help="System prompt")
    p_chat.add_argument("--max-tokens", type=int, default=512)
    p_chat.add_argument("--temperature", type=float, default=0.7)
    p_chat.add_argument("--no-stream", dest="stream", action="store_false", default=True)

    # completion
    p_comp = sub.add_parser("completion", help="Text completion")
    p_comp.add_argument("--model", "-m", default="Qwen/Qwen2.5-3B-Instruct")
    p_comp.add_argument("--prompt", "-p", required=True, help="Prompt text")
    p_comp.add_argument("--max-tokens", type=int, default=512)
    p_comp.add_argument("--temperature", type=float, default=0.7)
    p_comp.add_argument("--no-stream", dest="stream", action="store_false", default=True)

    # image
    p_img = sub.add_parser("image", help="Image generation")
    p_img.add_argument("--model", "-m", default="runwayml/stable-diffusion-v1-5")
    p_img.add_argument("--prompt", "-p", required=True, help="Image prompt")
    p_img.add_argument("--size", default="512x512", help="Image size WxH")
    p_img.add_argument("--output", "-o", help="Output file path")

    # tts
    p_tts = sub.add_parser("tts", help="Text-to-speech")
    p_tts.add_argument("--model", "-m", default="suno/bark-small")
    p_tts.add_argument("--text", "-t", required=True, help="Text to speak")
    p_tts.add_argument("--output", "-o", help="Output WAV path")

    args = parser.parse_args()

    if not args.command:
        parser.print_help()
        sys.exit(1)

    cmds = {
        "models": cmd_models,
        "chat": cmd_chat,
        "completion": cmd_completion,
        "image": cmd_image,
        "tts": cmd_tts,
    }
    cmds[args.command](args)


if __name__ == "__main__":
    main()