H / server.py
Maxenoz's picture
Create server.py
ad79e87 verified
import os
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
import openai
import torch
from transformers import pipeline
from diffusers import StableDiffusionPipeline
from PIL import Image
import io
import base64
import soundfile as sf
import numpy as np
# Initialize APIs
openai.api_key = os.getenv("OPENAI_API_KEY")
app = FastAPI(title="Maxenoz AI Server")
# Request models
class TextRequest(BaseModel):
input_text: str
class ImageRequest(BaseModel):
prompt: str
class AudioRequest(BaseModel):
seed: int = 42
# --- Health check ---
@app.get("/")
def health():
return {"status": "OK", "message": "Server is running"}
# --- GPT Text Endpoint ---
@app.post("/ai/text")
def ai_text(request: TextRequest):
try:
response = openai.ChatCompletion.create(
model="gpt-4-turbo",
messages=[{"role": "user", "content": request.input_text}]
)
return {"output_text": response.choices[0].message.content}
except Exception as e:
raise HTTPException(status_code=500, detail=str(e))
# --- Image Generation Endpoint ---
@app.post("/ai/image")
def ai_image(request: ImageRequest):
try:
# Minimal example using StableDiffusion
pipe = StableDiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5")
pipe = pipe.to("cpu") # Change to "cuda" if GPU available
image = pipe(request.prompt).images[0]
# Encode image to base64
buf = io.BytesIO()
image.save(buf, format="PNG")
img_bytes = buf.getvalue()
img_b64 = base64.b64encode(img_bytes).decode("utf-8")
return {"image_base64": img_b64}
except Exception as e:
raise HTTPException(status_code=500, detail=str(e))
# --- Audio Generation Endpoint (placeholder) ---
@app.post("/ai/audio")
def ai_audio(request: AudioRequest):
try:
# Example: generate 1-second silent audio
audio = np.zeros(16000) # 16kHz, 1 sec
buf = io.BytesIO()
sf.write(buf, audio, 16000, format="WAV")
audio_b64 = base64.b64encode(buf.getvalue()).decode("utf-8")
return {"audio_base64": audio_b64}
except Exception as e:
raise HTTPException(status_code=500, detail=str(e))
# --- Run server ---
if __name__ == "__main__":
import uvicorn
uvicorn.run(app, host="0.0.0.0", port=7860)