super_hwan.py

import gradio as gr
from transformers import (
    Blip2Processor,
    Blip2ForConditionalGeneration,
    AutoTokenizer,
    AutoModelForCausalLM,
)
from PIL import Image
import torch

# Set device
device = "cuda" if torch.cuda.is_available() else "cpu"

# Load image captioning model (BLIP-2)
processor = Blip2Processor.from_pretrained("Salesforce/blip2-flan-t5-xl")
blip_model = Blip2ForConditionalGeneration.from_pretrained("Salesforce/blip2-flan-t5-xl").to(device)

# Load text generation model (LLM)
llm_tokenizer = AutoTokenizer.from_pretrained("google/gemma-7b-it")
llm_model = AutoModelForCausalLM.from_pretrained("google/gemma-7b-it").to(device)

# Step 1: Generate image caption
def extract_caption(image):
    inputs = processor(images=image, return_tensors="pt").to(device)
    outputs = blip_model.generate(**inputs, max_new_tokens=50)
    caption = processor.tokenizer.decode(outputs[0], skip_special_tokens=True)
    return caption

# Step 2: Build fairytale prompt
def build_prompt(caption):
    return (
        f"Based on the image description: \"{caption}\", write a children's fairytale.\n"
        "The story must:\n"
        "- Start with 'Once upon a time'\n"
        "- Be at least 10 full sentences long\n"
        "- Include named characters, a clear setting, emotions, a challenge, and a resolution\n"
        "- Avoid mentions of babies or unrelated royalty unless relevant\n"
        "Here is the story:\nOnce upon a time"
    )

# Step 3: Generate story
def generate_fairytale(image):
    caption = extract_caption(image)
    prompt = build_prompt(caption)
    inputs = llm_tokenizer(prompt, return_tensors="pt").to(device)

    output = llm_model.generate(
        **inputs,
        max_new_tokens=500,
        do_sample=True,
        temperature=0.9,
        top_p=0.95,
        pad_token_id=llm_tokenizer.eos_token_id
    )

    result = llm_tokenizer.decode(output[0], skip_special_tokens=True)

    # Trim to only the story
    if "Once upon a time" in result:
        return "Once upon a time" + result.split("Once upon a time", 1)[-1].strip()
    else:
        return f"⚠️ Failed to generate story.\n\n[Prompt]\n{prompt}\n\n[Output]\n{result}"

# Gradio interface
with gr.Blocks() as demo:
    gr.Markdown("## 📖 AI Fairytale Generator\nUpload an image and get a magical story!")

    with gr.Row():
        image_input = gr.Image(type="pil", label="Upload an image")

    with gr.Row():
        generate_button = gr.Button("✨ Generate Fairytale")

    with gr.Row():
        output_text = gr.Textbox(label="Generated Story", lines=20)

    generate_button.click(fn=generate_fairytale, inputs=[image_input], outputs=[output_text])

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