import gradio as gr
import torch
from diffusers import StableDiffusionPipeline
from transformers import pipeline

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

text_generator = pipeline(
    "text-generation", model="openchat/openchat-3.5-0106", device=device
)
summarizer = pipeline(
    "summarization", model="sshleifer/distilbart-cnn-12-6", device=device
)
title_generator = pipeline(
    "text2text-generation",
    model="fabiochiu/t5-small-medium-title-generation",
    device=device,
)

stable_diffusion = StableDiffusionPipeline.from_pretrained("prompthero/openjourney-v4")
stable_diffusion.to(device)


def generate_blog_post(query):
    # Generate the article
    print("Generating article.")
    article = text_generator(query, max_length=500, num_return_sequences=1)[0][
        "generated_text"
    ]
    print(f"{article = }")

    # Generate a title for the article
    print("Generating the title.")
    title = title_generator(article, max_length=30, num_return_sequences=1)[0][
        "generated_text"
    ]
    print(f"{title = }")

    # Generate a cover image using Stable Diffusion
    print("Generating the cover.")
    cover = stable_diffusion(title, num_inference_steps=20, guidance_scale=7.5).images[
        0
    ]

    # Generate a summary of the article
    print("Generating the summary.")
    summary = summarizer(article, max_length=100, min_length=30, do_sample=False)[0][
        "summary_text"
    ]
    print(f"{summary = }")

    return title, cover, summary, article


with gr.Blocks() as iface:
    gr.Markdown("# Blog Post Generator")
    gr.Markdown(
        "Enter a topic, and I'll generate a blog post with a title, cover image, and summary!"
    )

    with gr.Row():
        topic_input = gr.Textbox(lines=2, placeholder="Enter your blog post topic...")

    generate_button = gr.Button("Generate Blog Post", size="sm")

    with gr.Row():
        with gr.Column(scale=2):
            title_output = gr.Textbox(label="Title")
            article_output = gr.Textbox(label="Article", lines=10)

        with gr.Column(scale=1):
            cover_output = gr.Image(label="Cover")
            summary_output = gr.Textbox(label="Summary", lines=5)

    generate_button.click(
        generate_blog_post,
        inputs=topic_input,
        outputs=[title_output, cover_output, summary_output, article_output],
    )

iface.launch()