Spaces:
Running
Running
| import streamlit as st | |
| import requests | |
| from PIL import Image | |
| from io import BytesIO | |
| from gradio_client import Client | |
| import json | |
| def generate(outline, characters, settings, num_images): | |
| # Create a dynamic prompt with multiple image prompts based on the number of images | |
| prompt = f"""Hello! I would like to request a 4-paragraph story (700 words per paragraph) and a corresponding image prompt for each paragraph in JSON format described later in the prompt with the following detailed outline:\n\n{outline}\n\nCharacters: {characters}\n\nSettings: {settings}\n\nPlease generate the story with the following detailed JSON format:\n\np1, p2, p3, p4: Keys for story paragraphs; title: Key for story title; {', '.join([f'prompt{i+1}' for i in range(num_images)])}: Keys for image prompts for each paragraph (or more if additional images are needed). Please do not include any other text in the output. Only JSON is needed or it will break the system and cause loss. Please don't say 'Here is the requested output in JSON format' or anything similar. Thank you.""" | |
| client = Client("Be-Bo/llama-3-chatbot_70b") | |
| hikaye = client.predict( | |
| message=prompt, | |
| api_name="/chat" | |
| ) | |
| print("Debug: Generated hikaye:", hikaye) | |
| return hikaye | |
| def cover(prompts): | |
| api_key = st.secrets['apikey'] | |
| model = "mann-e/Mann-E_Turbo" | |
| headers = {"Authorization": f"Bearer {api_key}"} | |
| api_url = f"https://api-inference.huggingface.co/models/{model}" | |
| images = [] | |
| for prompt in prompts: | |
| data = {"inputs": prompt} | |
| response = requests.post(api_url, headers=headers, json=data) | |
| if 'image' in response.headers.get('content-type', '').lower(): | |
| image = Image.open(BytesIO(response.content)) | |
| images.append(image) | |
| else: | |
| st.error(f"Failed to fetch image for prompt: {prompt}") | |
| images.append(None) | |
| return images | |
| def parse_story_response(response, num_images): | |
| print("Debug: Raw response from API:", response) | |
| if not response: | |
| print("Debug: Response is empty or None.") | |
| return None, None, None, None, [] | |
| title = response.get('title', '') | |
| p1 = response.get('p1', '') | |
| p2 = response.get('p2', '') | |
| p3 = response.get('p3', '') | |
| p4 = response.get('p4', '') | |
| prompts = [response.get(f'prompt{i+1}', '') for i in range(num_images)] | |
| return title, p1, p2, p3, p4, prompts | |
| # Streamlit UI | |
| st.title('Story AI By Ozi') | |
| characters = st.text_area(label="Characters") | |
| outline = st.text_area(label="Story Outline") | |
| settings = st.text_area(label="Setting") | |
| # Number chooser for images (up to 5) | |
| num_images = st.slider('Number of Images to Generate', 1, 5, 4) | |
| if st.button(label="Generate"): | |
| with st.spinner('Generating story and cover images...'): | |
| hikaye = generate(outline, characters, settings, num_images) | |
| print("Debug: Story generation response:", hikaye) | |
| if hikaye: | |
| try: | |
| hikaye_json = json.loads(hikaye) | |
| except json.JSONDecodeError as e: | |
| st.error(f"Failed to parse JSON response: {e}") | |
| st.stop() | |
| title, p1, p2, p3, p4, prompts = parse_story_response(hikaye_json, num_images) | |
| if title and p1 and p2 and p3 and p4: | |
| st.markdown(f'### {title}') | |
| # Generate and display images | |
| images = cover(prompts) | |
| for i, image in enumerate(images): | |
| if image: | |
| st.image(image, caption=f"Image {i+1}: {prompts[i]}") | |
| else: | |
| st.error(f"Failed to generate image {i+1}.") | |
| # Display paragraphs | |
| st.markdown(f''' | |
| {p1} | |
| {p2} | |
| {p3} | |
| {p4} | |
| ''') | |
| else: | |
| st.error("Failed to generate or parse story.") | |
| else: | |
| st.error("No story data received.") | |