import gradio as gr
from inference import generate_rag_caption

def predict(image):
    """
    Wrapper for Gradio to generate a caption from an uploaded image.
    Input: PIL Image
    Output: Caption (str)
    """
    caption = generate_rag_caption(image)
    return caption

# Gradio interface
iface = gr.Interface(
    fn=predict,
    inputs=gr.Image(type="pil", label="Upload an Image"),
    outputs=gr.Textbox(label="Generated Caption"),
    title="RAG Image Captioning",
    description="Upload an image to generate a caption using a RAG-based pipeline with CLIP, T5, and SentenceTransformer."
)

if __name__ == "__main__":
    iface.launch(share=True,ssr_mode=False)
    
# FastAPI app
app = FastAPI()

# REST API endpoint for image captioning
@app.post("/caption-image")
async def caption_image(file: UploadFile = File(...)):
    contents = await file.read()
    image = Image.open(io.BytesIO(contents)).convert("RGB")
    caption = generate_rag_caption(image)
    return {"caption": caption}

# Mount Gradio at root "/"
app = gr.mount_gradio_app(app, iface, path="/")

# import gradio as gr
# from inference import generate_rag_caption
# from fastapi import FastAPI, File, UploadFile
# from PIL import Image
# import io

# # Function for Gradio
# def predict(image):
#     caption = generate_rag_caption(image)
#     return caption

# # Gradio interface
# iface = gr.Interface(
#     fn=predict,
#     inputs=gr.Image(type="pil", label="Upload an Image"),
#     outputs=gr.Textbox(label="Generated Caption"),
#     title="RAG Image Captioning",
#     description="Upload an image to generate a caption using a RAG-based pipeline with CLIP, T5, and SentenceTransformer."
# )

# # FastAPI app
# app = FastAPI()

# # REST API endpoint for image captioning
# @app.post("/caption-image")
# async def caption_image(file: UploadFile = File(...)):
#     contents = await file.read()
#     image = Image.open(io.BytesIO(contents)).convert("RGB")
#     caption = generate_rag_caption(image)
#     return {"caption": caption}

# # Mount Gradio at root "/"
# app = gr.mount_gradio_app(app, iface, path="/")

# import gradio as gr
# from inference import generate_rag_caption
# from fastapi import FastAPI, File, UploadFile
# from PIL import Image
# import io

# # Step 1: Create FastAPI app
# app = FastAPI()

# # Step 2: Add an API route that accepts image uploads
# @app.post("/caption-image")
# async def caption_image(file: UploadFile = File(...)):
#     contents = await file.read()
#     image = Image.open(io.BytesIO(contents)).convert("RGB")
#     caption = generate_rag_caption(image)
#     return {"caption": caption}

# # Step 3: Create the original Gradio interface
# gradio_interface = gr.Interface(
#     fn=generate_rag_caption,
#     inputs=gr.Image(type="pil", label="Upload an Image"),
#     outputs=gr.Textbox(label="Generated Caption"),
#     title="RAG Image Captioning",
#     description="Upload an image to generate a caption using a RAG-based pipeline with CLIP, T5, and SentenceTransformer."
# )

# # ✅ Step 4: Mount Gradio interface on root path `/`
# app = gr.mount_gradio_app(app, gradio_interface, path="/")