main.py

import os
import requests
from fastapi import FastAPI, Request, Depends
from fastapi.responses import JSONResponse
from fastapi.staticfiles import StaticFiles
from pydantic import BaseModel
from typing import Optional, Any
from fastapi.responses import HTMLResponse

app = FastAPI()

# Montage du répertoire statique
app.mount("/static", StaticFiles(directory="static"), name="static")

# Route racine pour servir index.html
@app.get("/", include_in_schema=False)
async def root():
    with open("static/index.html", "r") as html_file:
        content = html_file.read()
    return HTMLResponse(content=content)

# Check whether we are executing inside a Hugging Face Space
SPACE_NAME = os.getenv("SPACE_NAME", default=None)
if SPACE_NAME is not None:
    print(f"Running inside {SPACE_NAME} Space.")
    try:
        # Try to auto-login using the Space's environment variables
        login(automatically=True)
    except Exception as e:
        print(f"Failed to auto-login ({str(e)}). Manually check the HF_ACCESS_TOKEN environment variable.")
        sys.exit(1)

try:
    HUGGINGFACE_TOKEN = os.environ['HF_ACCESS_TOKEN']
except KeyError:
    print('The environment variable "HF_ACCESS_TOKEN" is not found. Please configure it correctly in your Space.')
    sys.exit(1)

# Set up the API endpoint and headers
model_id = "152334H/miqu-1-70b-sf"
endpoint = f"https://api-inference.huggingface.co/models/{model_id}"
headers = {"Authorization": f"Bearer {HUGGINGFACE_TOKEN}"}

# Helper function to read raw request bodies
async def parse_raw(request: Request):
    return await request.body()

# Generate text using the Inference API
def generate_text(input_text: str) -> str:
    data = {"inputs": input_text}
    try:
        response = requests.post(endpoint, json=data, headers=headers)
        response.raise_for_status()  # Raise an exception for non-2xx status codes

        print(response.json())  # Print the entire response object

        return response.json()["generated_text"]
    except requests.exceptions.RequestException as e:
        print(f"RequestException: {e}")
        return ""
    except KeyError as e:
        print(f"KeyError: {e}. Check the response object for the correct key.")
        return ""

# Route for generating text
@app.post("/generate_text")
async def generate_text_route(data: BaseModel = Depends(parse_raw)):
    input_text = data.decode("utf-8")
    if not input_text or len(input_text) <= 0:
        return JSONResponse({"error": "Empty input received."}, status_code=400)

    return {"output": generate_text(input_text)}

# Mount static files
app.mount("/static", StaticFiles(directory="static"), name="static")