Spaces:
Running
Running
| import random | |
| import os | |
| import re | |
| import streamlit as st | |
| from langchain_nvidia_ai_endpoints import ChatNVIDIA | |
| # Set environmental variables | |
| nvidia_api_key = os.getenv("NVIDIA_KEY") | |
| # Initialize ChatNVIDIA for text generation | |
| client = ChatNVIDIA( | |
| model="meta/llama-3.1-8b-instruct", | |
| api_key=nvidia_api_key, | |
| temperature=0.5, | |
| top_p=0.85, | |
| max_tokens=80 | |
| ) | |
| # Get the absolute path of the script's directory | |
| script_dir = os.path.dirname(os.path.abspath(__file__)) | |
| phrases_file = os.path.join(script_dir, "phrases.txt") | |
| # Read phrases from the text file | |
| with open(phrases_file, "r") as f: | |
| lines = f.readlines() | |
| def clean_response(response_text): | |
| # Join the response list into a single string | |
| response_end = " ".join(response_text).strip() | |
| # Replace multiple spaces with a single space | |
| response_end = re.sub(r'\s+', ' ', response_end) # Collapses multiple spaces into one | |
| # Attempt to fix improper word splitting only | |
| # This regex finds single-letter space single-letter cases more precisely and joins them | |
| response_end = re.sub(r'(\b\w)\s+(\w\b)', r'\1\2', response_end, flags=re.UNICODE) | |
| response_end = re.sub(r'\s*([,.!?])\s*', r'\1 ', response_end) | |
| # Remove any extraneous newlines or extra spaces left over | |
| response_end = response_end.replace("\n", " ").replace(" ", " ") | |
| # Trim any leading or trailing spaces | |
| response_end = response_end.strip() | |
| return response_text | |
| def generate(starting_text): | |
| for count in range(6): | |
| seed = random.randint(100, 1000000) | |
| random.seed(seed) | |
| # If the text field is empty | |
| if starting_text == "": | |
| starting_text = lines[random.randrange(0, len(lines))].strip() | |
| starting_text = re.sub(r"[,:\-–.!;?_]", '', starting_text) | |
| # Format messages for the client | |
| messages = [ | |
| {"role": "system", "content": "Create a vivid and imaginative MidJourney prompt for image generation. Focus on descriptive language, colors, and setting to evoke a specific visual scene. Avoid storytelling; focus on visual elements."}, | |
| {"role": "user", "content": starting_text} | |
| ] | |
| # Stream response from LLaMA | |
| response = client.stream(messages) | |
| # Buffer to accumulate response chunks | |
| buffer = [] | |
| for chunk in response: | |
| if chunk.content: | |
| resp = chunk.content.strip() | |
| if resp != starting_text: | |
| buffer.append(resp) | |
| # Join the buffered chunks into a single string | |
| response_text = " ".join(buffer).strip() | |
| # Clean up the response text | |
| response_end = clean_response(response_text) | |
| # Print only the final response if it's not empty | |
| if response_end != "": | |
| return response_end | |
| # Streamlit UI | |
| st.title("Generate Sci-fi/Space-themed MidJourney Prompts") | |
| user_input = st.text_input("Enter a short text or just press 'Generate Prompt' to use our space-themed RAG to get space-themed prompt:") | |
| if st.button("Generate Prompt"): | |
| generated_prompt = generate(user_input) | |
| st.write("**MidJourney Prompt:**") | |
| st.write(generated_prompt) |