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)