Spaces:
Sleeping
Sleeping
import streamlit as st | |
from transformers import pipeline | |
from diffusers import StableDiffusionPipeline | |
from PIL import Image | |
import torch | |
# Initialize the text generation model | |
def initialize_text_generator(): | |
try: | |
text_generator = pipeline('text-generation', model='gpt2') | |
except Exception as e: | |
st.error(f"Error loading text generation model: {e}") | |
return None | |
return text_generator | |
# Initialize the image generation model | |
def initialize_image_generator(): | |
try: | |
device = "cuda" if torch.cuda.is_available() else "cpu" | |
image_generator = StableDiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v1-4").to(device) | |
except Exception as e: | |
st.error(f"Error loading image generation model: {e}") | |
return None | |
return image_generator | |
text_generator = initialize_text_generator() | |
image_generator = initialize_image_generator() | |
def generate_blog(title): | |
if text_generator is None or image_generator is None: | |
return "Failed to load models", None | |
# Generate blog content | |
blog_content = text_generator(title, max_length=500, num_return_sequences=1)[0]['generated_text'] | |
# Generate an image | |
image_prompt = f"An image representing {title}" | |
image = image_generator(image_prompt).images[0] | |
return blog_content, image | |
# Streamlit app | |
st.title('Blog Generator') | |
title = st.text_input('Enter the title of your blog:') | |
if title: | |
with st.spinner('Generating blog content and image...'): | |
blog_content, image = generate_blog(title) | |
if blog_content == "Failed to load models": | |
st.error(blog_content) | |
else: | |
st.success('Blog generated successfully!') | |
st.subheader('Blog Content') | |
st.write(blog_content) | |
st.subheader('Generated Image') | |
st.image(image, caption='Generated Image') | |