main.py

from fastapi import FastAPI, UploadFile, File, HTTPException
from fastapi.responses import JSONResponse
from gradio_client import Client, handle_file
import base64

import os

app = FastAPI()

HF_TOKEN = os.getenv("HF_TOKEN")
# Initialize the Gradio client with the token
client = Client("Makhinur/Bringingoldphotoliveagain", hf_token=HF_TOKEN)


@app.post("/upload/")
async def upload_image(file: UploadFile = File(...)):
    try:
        # Save the uploaded file temporarily
        with open(file.filename, "wb") as buffer:
            buffer.write(await file.read())
        
        # Use the Gradio client to process the image
        result = gradio_client.predict(
            img=handle_file(file.filename),
            api_name="/predict"
        )

        # Read the output image and encode it in base64
        with open(result[0], "rb") as result_file:
            encoded_string = base64.b64encode(result_file.read()).decode('utf-8')

        # Create a JSON response with the base64 encoded image
        return JSONResponse(content={"sketch_image_base64": f"data:image/jpeg;base64,{encoded_string}"})

    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))