from transformers import pipeline
from IPython.display import display, Markdown
from PIL import Image
import requests
from io import BytesIO

# Load model
pipe = pipeline(
    "image-text-to-text",
    model="rahul7star/LFM2.5-VL-3B-tool-sft-merged",
    trust_remote_code=True,
)

# Conversation history
messages = []

def chat(user_text, image_url=None):
    global messages

    content = []

    if image_url:
        content.append({
            "type": "image",
            "url": image_url,
        })

    content.append({
        "type": "text",
        "text": user_text,
    })

    messages.append({
        "role": "user",
        "content": content,
    })

    output = pipe(text=messages)

    # Pipeline output is normally a list containing the generated conversation
    result = output[0]

    if isinstance(result, dict) and "generated_text" in result:
        generated = result["generated_text"]
    else:
        generated = result

    # If the pipeline returns the complete conversation,
    # extract the assistant's latest message.
    if isinstance(generated, list):
        assistant_message = generated[-1]

        if isinstance(assistant_message, dict):
            assistant_content = assistant_message.get("content", "")
        else:
            assistant_content = str(assistant_message)
    else:
        assistant_content = str(generated)

    messages.append({
        "role": "assistant",
        "content": [
            {
                "type": "text",
                "text": assistant_content,
            }
        ],
    })

    print("\nAssistant:")
    print(assistant_content)

    return assistant_content


# Test with image
image_url = "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/p-blog/candy.JPG"

chat(
    "What animal is on the candy?",
    image_url=image_url,
)