Gajendra5490's picture
Update main.py
edcdcf8 verified
from fastapi import FastAPI, File, UploadFile
from gradio_client import Client
import shutil
import os
import gradio as gr
# Initialize FastAPI app
app = FastAPI()
# Use `/tmp/` directory (writable on Hugging Face Spaces)
UPLOAD_DIR = "/tmp/temp_uploads"
os.makedirs(UPLOAD_DIR, exist_ok=True)
# Hugging Face Gradio API
client = Client("https://vila-m3-demo.monai.ngc.nvidia.com/")
@app.post("/")
async def process_image(file: UploadFile = File(...)):
"""Receives an image, saves it, and sends it to the Hugging Face model"""
try:
# Save the uploaded file
file_path = os.path.join(UPLOAD_DIR, file.filename)
with open(file_path, "wb") as buffer:
shutil.copyfileobj(file.file, buffer)
# Send image file path to Hugging Face API
result = client.predict(
image=handle_file('https://raw.githubusercontent.com/gradio-app/gradio/main/test/test_files/bus.png'),
#image = file_path, # File path of uploaded image
api_name="/input_image"
)
# Delete the file after processing
os.remove(file_path)
return {"result": result}
except Exception as e:
return {"error": str(e)}
# Create Gradio Interface
def upload_image(image):
file_path = os.path.join(UPLOAD_DIR, "uploaded_image.png")
image.save(file_path)
result = client.predict(file_path, api_name="/input_image")
os.remove(file_path)
return result
interface = gr.Interface(
fn=upload_image,
inputs=gr.Image(type="pil"),
outputs="text",
title="Image Processing App",
description="Upload an image to process it using the Hugging Face model."
)
# Run the Gradio interface in a separate thread
import threading
threading.Thread(target=interface.launch, daemon=True).start()
# 🚀 Ensure app is properly loaded for Hugging Face
if __name__ == "__main__":
import uvicorn
uvicorn.run("main:app", host="0.0.0.0", port=7860, reload=True)