Spaces:
Sleeping
Sleeping
| import os | |
| import tempfile | |
| import requests | |
| import streamlit as st | |
| from src.pipeline.main import LearnableAI | |
| from src.services.use_case_three.story_generator import StoryGenerator | |
| from src.services.image_generation.image_gen import image_generation | |
| def main(): | |
| st.title("LearnableAI") | |
| learnable_ai = LearnableAI() | |
| st.sidebar.header("Configuration") | |
| use_case = st.sidebar.radio( | |
| "Select Use Case", | |
| ["Use Case: 1(Word-Sentences)", "Use Case: 2(Image Describer)", "Use Case: 3(Story Generator From Image)"] | |
| ) | |
| difficulty = st.sidebar.selectbox( | |
| "Select Difficulty Level", | |
| ["Simple", "Easy", "Challenging", "Advanced"] | |
| ) | |
| if use_case == "Use Case: 1(Word-Sentences)": | |
| st.header("Word to Sentence Generation") | |
| st.write( | |
| "Enter a list of words and the child interest to generate the content for eg.Words: Apple, Banana interest: Ironman") | |
| words = st.text_area( | |
| "Words:", | |
| "", | |
| help="Enter multiple words separated by commas" | |
| ) | |
| interest = st.text_input("Child Interest:", "", | |
| help="Enter the interest area for the content ie. Superheroes, Animals, etc.") | |
| if st.button("Generate Content", key="text_process"): | |
| if words and interest: | |
| try: | |
| word_list = [w.strip() for w in words.split(",")] if "," in words else [words.strip()] | |
| result = learnable_ai.word_to_sentence( | |
| words=word_list, | |
| interest=interest, | |
| difficulties=difficulty | |
| ) | |
| st.success("Content Generated Successfully!") | |
| st.write("### Results") | |
| st.write(result.split("</think>")[1].split("Image Generation Prompt")[0]) | |
| image_prompt = result.split("Image Generation Prompt")[1] | |
| input = { | |
| "prompt": image_prompt, | |
| "go_fast": True, | |
| "megapixels": "1", | |
| "num_outputs": 1, | |
| "aspect_ratio": "1:1", | |
| "output_format": "webp", | |
| "output_quality": 80, | |
| "num_inference_steps": 4 | |
| } | |
| output_image_url = image_generation(input) | |
| image_bytes = requests.get(output_image_url).content | |
| st.write("### Generated Image") | |
| st.image(image_bytes, caption="Generated Image") | |
| except Exception as e: | |
| st.error(f"An error occurred: {str(e)}") | |
| else: | |
| st.warning("Please fill in both words and interest fields!") | |
| elif use_case == "Use Case: 2(Image Describer)": | |
| st.header("Image Description") | |
| st.write("Upload an image or capture from camera to generate to generate the description of the image") | |
| upload_method = st.radio( | |
| "Choose Upload Method", | |
| ["Upload File", "Capture from Camera"] | |
| ) | |
| if upload_method == "Upload File": | |
| image_file = st.file_uploader( | |
| "Upload an image", | |
| type=["jpg", "jpeg", "png"] | |
| ) | |
| else: | |
| image_file = st.camera_input("Take a picture") | |
| if image_file is not None: | |
| st.image(image_file, caption="Uploaded Image", use_column_width=True) | |
| if st.button("Generate Description", key="image_process"): | |
| try: # temp image file | |
| with tempfile.NamedTemporaryFile(delete=False, suffix='.jpg') as temp_file: | |
| temp_file.write(image_file.read()) | |
| temp_image_path = temp_file.name | |
| result = learnable_ai.image_describer( | |
| image_path=temp_image_path, | |
| difficulties=difficulty | |
| ) | |
| st.success("Image Processed Successfully!") | |
| st.write("### Image Description Results") | |
| st.write(result) | |
| os.remove(temp_image_path) # remove temp image file | |
| except Exception as e: | |
| st.error(f"An error occurred while processing the image: {str(e)}") | |
| elif use_case == "Use Case: 3(Story Generator From Image)": | |
| st.title("Story Generator from Images") | |
| st.write("Story writer from the images. Upload up to 3 images to generate a children's story.") | |
| st.subheader("Upload up to 3 images and generate a children's story") | |
| uploaded_files = st.file_uploader( | |
| "Upload up to 3 images", | |
| type=["jpg", "jpeg", "png"], | |
| accept_multiple_files=True | |
| ) | |
| if uploaded_files: | |
| if len(uploaded_files) > 3: | |
| st.error("You can upload a maximum of 3 images.") | |
| else: | |
| st.write("Uploaded Images:") | |
| for uploaded_file in uploaded_files: | |
| st.image(uploaded_file, caption=uploaded_file.name, use_column_width=True) | |
| if st.button("Generate Story"): | |
| story_generator = StoryGenerator() | |
| story_parts = [] | |
| try: | |
| for uploaded_file in uploaded_files: | |
| with tempfile.NamedTemporaryFile(delete=False, suffix='.jpg') as temp_file: | |
| temp_file.write(uploaded_file.read()) | |
| temp_image_path = temp_file.name | |
| story_part = story_generator.process_image(temp_image_path) | |
| story_parts.append(story_part) | |
| os.remove(temp_image_path) | |
| full_story = "\n\n".join(story_parts) | |
| st.success("Story Generated Successfully!") | |
| st.write("### Your Story") | |
| st.write(full_story) | |
| except Exception as e: | |
| st.error(f"An error occurred: {str(e)}") | |
| if __name__ == "__main__": | |
| main() | |