import spaces
import gradio as gr
import torch
from transformers import (
    BlipProcessor,
    BlipForConditionalGeneration,
    AutoTokenizer,
    AutoModelForCausalLM,
)


APP_TITLE = "StoryLens AI"

STORY_STYLES = [
    "Fantasy",
    "Adventure",
    "Mystery",
    "Sci-Fi",
    "Inspirational",
    "Comedy",
    "Historical",
]

DEFAULT_STYLE = STORY_STYLES[0]

device = "cuda" if torch.cuda.is_available() else "cpu"
dtype = torch.float16 if device == "cuda" else torch.float32


blip_id = "Salesforce/blip-image-captioning-base"

caption_processor = BlipProcessor.from_pretrained(blip_id)

caption_model = BlipForConditionalGeneration.from_pretrained(
    blip_id,
    torch_dtype=dtype,
).to(device)

llm_id = "Qwen/Qwen2.5-1.5B-Instruct"

tokenizer = AutoTokenizer.from_pretrained(llm_id)

story_model = AutoModelForCausalLM.from_pretrained(
    llm_id,
    torch_dtype=dtype,
).to(device)

@spaces.GPU
def generate_caption(image):
    inputs = caption_processor(
        image,
        return_tensors="pt"
    ).to(device, dtype)

    output = caption_model.generate(
        **inputs,
        max_new_tokens=30
    )

    return caption_processor.decode(
        output[0],
        skip_special_tokens=True
    )


@spaces.GPU
def generate_story(caption, style, max_tokens):

    prompt = (
        f"Create an engaging story inspired by this image description: '{caption}'. "
        f"Use the '{style}' style. "
        "The story should be creative, descriptive, and between 3 and 5 sentences."
    )

    messages = [
        {
            "role": "user",
            "content": prompt
        }
    ]

    inputs = tokenizer.apply_chat_template(
        messages,
        add_generation_prompt=True,
        return_tensors="pt",
        return_dict=False,
    ).to(device)

    outputs = story_model.generate(
        inputs,
        max_new_tokens=max_tokens,
        do_sample=True,
        temperature=0.8,
        top_p=0.9,
        pad_token_id=tokenizer.eos_token_id,
    )

    return tokenizer.decode(
        outputs[0][inputs.shape[1]:],
        skip_special_tokens=True,
    )


def generate_output(image, style, length):

    if image is None:
        return "Please upload an image first.", ""

    image = image.convert("RGB")

    caption = generate_caption(image)

    story = generate_story(
        caption,
        style,
        length,
    )

    return caption, story


custom_theme = gr.themes.Default(
    primary_hue="blue",
    secondary_hue="blue",
    neutral_hue="slate",
)

css = """
body{
    background:#0f172a;
}

.gradio-container{
    background:#0f172a !important;
}

footer{
    display:none;
}
"""

demo = gr.Interface(
    fn=generate_output,

    inputs=[
        gr.Image(
            type="pil",
            label="Upload Image"
        ),

        gr.Dropdown(
            choices=STORY_STYLES,
            value=DEFAULT_STYLE,
            label="Story Style"
        ),

        gr.Slider(
            minimum=40,
            maximum=200,
            value=120,
            step=10,
            label="Story Length"
        ),
    ],

    outputs=[
        gr.Textbox(
            label="Image Caption"
        ),

        gr.Textbox(
            label="Generated Story",
            lines=8
        ),
    ],

    title=APP_TITLE,

    description="""
Upload any image and let AI first describe it, then transform that description
into a creative story based on your selected writing style.
""",

    theme=custom_theme,
    css=css,
)

demo.launch(server_name="0.0.0.0", server_port=7860)