import os
from transformers import BlipProcessor, BlipForConditionalGeneration
import gradio as gr

# Load the token from the environment
HUGGINGFACE_TOKEN = os.getenv("HUGGINGFACE_TOKEN")

# Load the model and processor with the token
processor = BlipProcessor.from_pretrained(
    "quadranttechnologies/Imageclassification",
    use_auth_token=HUGGINGFACE_TOKEN
)
model = BlipForConditionalGeneration.from_pretrained(
    "quadranttechnologies/Imageclassification",
    use_auth_token=HUGGINGFACE_TOKEN
)

# Define your Gradio interface and logic as before
def generate_caption(image):
    try:
        inputs = processor(image, return_tensors="pt")
        outputs = model.generate(**inputs)
        caption = processor.decode(outputs[0], skip_special_tokens=True)
        return caption
    except Exception as e:
        return f"Error generating caption: {e}"

interface = gr.Interface(
    fn=generate_caption,
    inputs=gr.Image(type="pil"),
    outputs="text",
    title="Image Captioning Model",
    description="Upload an image to receive a caption generated by the model."
)

if __name__ == "__main__":
    interface.launch(share=True)