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') | |