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

# Initialize the pipelines
device = "cuda" if torch.cuda.is_available() else "cpu"

text_generator = pipeline("text-generation", model="gpt2", device=device)
summarizer = pipeline("summarization", model="facebook/bart-large-cnn", device=device)
title_generator = pipeline("text2text-generation", model="fabiochiu/t5-small-medium-title-generation", device=device)

# Initialize the Stable Diffusion pipeline
stable_diffusion = StableDiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5", torch_dtype=torch.float16)
stable_diffusion.to(device)

# Function to generate blog post
def generate_blog_post(query):
    # Generate the article
    article = text_generator(query, max_length=500, num_return_sequences=1)[0]['generated_text']
    
    # Generate a title for the article
    title = title_generator(article, max_length=30, num_return_sequences=1)[0]['generated_text']
    
    # Generate a cover image using Stable Diffusion
    cover_image = stable_diffusion(query, num_inference_steps=50, guidance_scale=7.5).images[0]
    
    # Generate a summary of the article
    summary = summarizer(article, max_length=100, min_length=30, do_sample=False)[0]['summary_text']
    
    return title, article, cover_image, summary

# Create the Gradio interface
iface = gr.Interface(
    fn=generate_blog_post,
    inputs=gr.Textbox(lines=2, placeholder="Enter your blog post topic..."),
    outputs=[
        gr.Textbox(label="Generated Title"),
        gr.Textbox(label="Generated Article"),
        gr.Image(label="Cover Image"),
        gr.Textbox(label="Article Summary")
    ],
    title="Blog Post Generator",
    description="Enter a topic, and I'll generate a blog post with a title, cover image, and summary!"
)

# Launch the app
iface.launch()