Spaces:
Running
Running
Ley_Fill7
commited on
Commit
•
0927a59
1
Parent(s):
f54eb4b
Add app file
Browse files
app.py
ADDED
@@ -0,0 +1,93 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import random
|
2 |
+
import os
|
3 |
+
import re
|
4 |
+
from langchain_nvidia_ai_endpoints import ChatNVIDIA
|
5 |
+
|
6 |
+
# Set environmental variables
|
7 |
+
nvidia_api_key = os.getenv("NVIDIA_KEY")
|
8 |
+
|
9 |
+
# Initialize ChatNVIDIA for text generation
|
10 |
+
client = ChatNVIDIA(
|
11 |
+
model="meta/llama-3.1-8b-instruct",
|
12 |
+
api_key=nvidia_api_key,
|
13 |
+
temperature=0.5,
|
14 |
+
top_p=0.85,
|
15 |
+
max_tokens=80
|
16 |
+
)
|
17 |
+
|
18 |
+
# Get the absolute path of the script's directory
|
19 |
+
script_dir = os.path.dirname(os.path.abspath(__file__))
|
20 |
+
phrases_file = os.path.join(script_dir, "phrases.txt")
|
21 |
+
|
22 |
+
# Read phrases from the text file
|
23 |
+
with open(phrases_file, "r") as f:
|
24 |
+
lines = f.readlines()
|
25 |
+
|
26 |
+
def clean_response(response_text):
|
27 |
+
# Join the response list into a single string
|
28 |
+
response_end = " ".join(response_text).strip()
|
29 |
+
|
30 |
+
# Replace multiple spaces with a single space
|
31 |
+
response_end = re.sub(r'\s+', ' ', response_end) # Collapses multiple spaces into one
|
32 |
+
|
33 |
+
# Attempt to fix improper word splitting only
|
34 |
+
# This regex finds single-letter space single-letter cases more precisely and joins them
|
35 |
+
response_end = re.sub(r'(\b\w)\s+(\w\b)', r'\1\2', response_end, flags=re.UNICODE)
|
36 |
+
|
37 |
+
response_end = re.sub(r'\s*([,.!?])\s*', r'\1 ', response_end)
|
38 |
+
|
39 |
+
# Remove any extraneous newlines or extra spaces left over
|
40 |
+
response_end = response_end.replace("\n", " ").replace(" ", " ")
|
41 |
+
|
42 |
+
# Trim any leading or trailing spaces
|
43 |
+
response_end = response_end.strip()
|
44 |
+
|
45 |
+
return response_text
|
46 |
+
|
47 |
+
def generate(starting_text):
|
48 |
+
for count in range(6):
|
49 |
+
seed = random.randint(100, 1000000)
|
50 |
+
random.seed(seed)
|
51 |
+
|
52 |
+
# If the text field is empty
|
53 |
+
if starting_text == "":
|
54 |
+
starting_text = lines[random.randrange(0, len(lines))].strip()
|
55 |
+
starting_text = re.sub(r"[,:\-–.!;?_]", '', starting_text)
|
56 |
+
|
57 |
+
# Format messages for the client
|
58 |
+
messages = [
|
59 |
+
{"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."},
|
60 |
+
{"role": "user", "content": starting_text}
|
61 |
+
]
|
62 |
+
|
63 |
+
# Stream response from LLaMA
|
64 |
+
response = client.stream(messages)
|
65 |
+
|
66 |
+
# Buffer to accumulate response chunks
|
67 |
+
buffer = []
|
68 |
+
for chunk in response:
|
69 |
+
if chunk.content:
|
70 |
+
resp = chunk.content.strip()
|
71 |
+
if resp != starting_text:
|
72 |
+
buffer.append(resp)
|
73 |
+
|
74 |
+
# Join the buffered chunks into a single string
|
75 |
+
response_text = " ".join(buffer).strip()
|
76 |
+
|
77 |
+
# Clean up the response text
|
78 |
+
response_end = clean_response(response_text)
|
79 |
+
|
80 |
+
# Print only the final response if it's not empty
|
81 |
+
if response_end != "":
|
82 |
+
return response_end
|
83 |
+
|
84 |
+
# Example usage
|
85 |
+
if __name__ == "__main__":
|
86 |
+
# Get user input for the starting text
|
87 |
+
user_input = input("Enter a starting text for the Midjourney prompt: ")
|
88 |
+
|
89 |
+
# Generate the prompt
|
90 |
+
generated_prompt = generate(user_input)
|
91 |
+
|
92 |
+
# Print the generated prompt
|
93 |
+
print(f"Generated Prompt:\n{generated_prompt}")
|