app.py

import gradio as gr
from transformers import BlipProcessor, BlipForConditionalGeneration, pipeline
from fpdf import FPDF
from PIL import Image
import os

# Load BLIP for image captioning (lightweight model)
processor = BlipProcessor.from_pretrained("Salesforce/blip-image-captioning-base")
model = BlipForConditionalGeneration.from_pretrained("Salesforce/blip-image-captioning-base")

# Load very lightweight LLM for story generation (CPU friendly)
story_gen = pipeline("text-generation", model="distilgpt2")

def generate_caption(image):
    image = Image.open(image).convert("RGB")
    inputs = processor(image, return_tensors="pt")
    out = model.generate(**inputs, max_length=50)
    caption = processor.decode(out[0], skip_special_tokens=True)
    return caption

def generate_story(captions, theme="fantasy"):
    prompt = f"Write a short children's story in a {theme} style, connecting these images:\n"
    for i, cap in enumerate(captions, 1):
        prompt += f"Scene {i}: {cap}\n"
    output = story_gen(prompt, max_length=200, do_sample=True, temperature=0.7)
    story = output[0]["generated_text"]
    return story

def create_pdf(images, story_text):
    pdf = FPDF()
    pdf.set_auto_page_break(auto=True, margin=15)

    story_pages = story_text.split("\n")
    for i, img in enumerate(images):
        pdf.add_page()
        img_path = f"temp_{i}.png"
        Image.open(img).save(img_path)
        pdf.image(img_path, x=10, y=20, w=90)
        os.remove(img_path)

        pdf.set_font("Arial", size=12)
        pdf.set_xy(110, 20)
        pdf.multi_cell(90, 10, story_pages[i] if i < len(story_pages) else "")

    pdf_path = "storybook.pdf"
    pdf.output(pdf_path)
    return pdf_path

def main(images, theme):
    captions = [generate_caption(img) for img in images]
    story = generate_story(captions, theme)
    pdf_file = create_pdf(images, story)
    return story, pdf_file

with gr.Blocks() as demo:
    gr.Markdown("# 📚 StoryCraft AI – AI Storybook Generator")
    gr.Markdown("Turn messy kids' doodles, drawings, or photos into a fun AI-generated storybook!")

    with gr.Row():
        image_input = gr.File(type="file", file_types=["image"], file_count="multiple", label="Upload Images")
        theme_input = gr.Dropdown(["fantasy", "adventure", "mystery", "comedy"], value="fantasy", label="Choose Theme")

    story_output = gr.Textbox(label="Generated Story", lines=15)
    pdf_output = gr.File(label="Download Storybook (PDF)")

    submit_btn = gr.Button("Generate Storybook")
    submit_btn.click(main, inputs=[image_input, theme_input], outputs=[story_output, pdf_output])

demo.launch()