Skip to content
This repository was archived by the owner on Jul 30, 2024. It is now read-only.

Commit ac792b8

Browse files
harmindersinghnijjarharmindersinghnijjar
authored
Add files via upload
1 parent 51da334 commit ac792b8

File tree

3 files changed

+335
-0
lines changed

3 files changed

+335
-0
lines changed

‎requirements.txt‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,6 @@
1+
numpy==1.21.6
2+
openai==0.26.1
3+
Pillow==9.5.0
4+
playwright==1.33.0
5+
requests==2.28.1
6+

‎set_openai_api_key.py‎

Lines changed: 66 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,66 @@
1+
2+
# Import modules.
3+
import os
4+
5+
# Define functions to get OpenAI API key from file.
6+
def get_openai_api_key_from_file(filepath):
7+
with open(filepath, 'r') as infile:
8+
return infile.read().strip()
9+
10+
# Define function to get OpenAI API key from environment.
11+
def get_openai_api_key_from_env():
12+
return os.environ['OPENAI_API_KEY']
13+
14+
# Define function to get OpenAI API key from user input.
15+
def get_openai_api_key_from_user():
16+
return input("Please enter your OpenAI API key: ")
17+
18+
# Define function to set OpenAI API key in environment.
19+
def set_openai_api_key(api_key):
20+
# Set the OpenAI API key in the environment variable by using the command line.
21+
os.environ['OPENAI_API_KEY'] = api_key
22+
23+
# Ask user how they want to enter their OpenAI API key.
24+
def get_openai_api_key():
25+
# Run an infinite loop until the user enters a valid choice.
26+
while True:
27+
# Print the menu.
28+
print("Welcome to the Quora Automation OpenAI API key setup script.")
29+
print("How would you like to enter your OpenAI API key?")
30+
print("1. From a file")
31+
print("2. From the environment")
32+
print("3. From the console")
33+
print("4. Exit")
34+
choice = input("Please enter a number: ")
35+
# If the user enters 1, ask them for the filepath to the API key file.
36+
if choice == '1':
37+
filepath = input("Please enter the filepath to your API key file: ")
38+
return get_openai_api_key_from_file(filepath)
39+
# If the user enters 2, get the API key from the environment.
40+
elif choice == '2':
41+
return get_openai_api_key_from_env()
42+
# If the user enters 3, ask them to enter the API key.
43+
elif choice == '3':
44+
return get_openai_api_key_from_user()
45+
# If the user enters 4, exit the script.
46+
elif choice == '4':
47+
exit()
48+
# If the user enters anything else, ask them to try again.
49+
else:
50+
print("Invalid choice. Please try again.")
51+
52+
# Main program of the script.
53+
if __name__ == '__main__':
54+
# Check if the environment variable is already set.
55+
if "OPENAI_API_KEY" in os.environ:
56+
print("Your OpenAI API key is already set in the environment.")
57+
print("You can now run the Upscale Midjourney Downloader.")
58+
exit()
59+
# If the environment variable is not set, ask the user how they want to enter their OpenAI API key.
60+
else:
61+
print("Your OpenAI API key is not set in the environment.")
62+
print("Let's set it now.")
63+
api_key = get_openai_api_key()
64+
print("Your OpenAI API key is: %s" % api_key)
65+
set_openai_api_key(api_key)
66+
print("Your OpenAI API key has been set in the environment.")

‎upscale_midjourney_downloader.py‎

Lines changed: 263 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,263 @@
1+
from playwright.sync_api import expect
2+
from playwright.sync_api import sync_playwright
3+
import openai
4+
import os
5+
import random
6+
import re
7+
import requests
8+
import shutil
9+
import time
10+
import uuid
11+
12+
# Set OpenAI API key.
13+
openai.api_key = os.getenv("OPENAI_API_KEY")
14+
15+
# Constants.
16+
ART_TYPE = ""
17+
BOT_COMMAND = "/imagine"
18+
CHANNEL_URL = ""
19+
DESCRIPTORS = ""
20+
TOPIC = ""
21+
PROMPT = f"Generate a Midjourney prompt to result in an {ART_TYPE} image about {TOPIC} include {DESCRIPTORS}"
22+
23+
24+
25+
# If any of the constants are not set, ask it as user input and exit.
26+
if ART_TYPE == "":
27+
print("Enter art type in line 18 of the source code.")
28+
exit()
29+
if BOT_COMMAND == "":
30+
print("Enter bot command in line 19 of the source code.")
31+
exit()
32+
if CHANNEL_URL == "":
33+
print("Enter channel URL in line 20 of the source code.")
34+
exit()
35+
if DESCRIPTORS == "":
36+
print("Enter descriptors in line 21 of the source code.")
37+
exit()
38+
if TOPIC == "":
39+
print("Enter topic in line 22 of the source code.")
40+
exit()
41+
if PROMPT == "":
42+
print("Enter prompt to send to OpenAI's API in line 23 of the source code.")
43+
exit()
44+
45+
# Function to login to Discord.
46+
def open_discord_channel(page):
47+
# Go to appropriate channel.
48+
page.goto(f"{CHANNEL_URL}")
49+
time.sleep(random.randint(1, 5))
50+
page.wait_for_load_state("networkidle") # This waits for the "networkidle"
51+
print("Opened appropriate channel.")
52+
print("Entering in the specified bot command.")
53+
bot_command(page, BOT_COMMAND)
54+
return page
55+
56+
# Function to ask GPT-3 for the image prompt.
57+
def gpt3_midjourney_prompt(prompt, engine='text-davinci-003', temp=0.7, top_p=1.0, tokens=400, freq_pen=0.0, pres_pen=0.0):
58+
try:
59+
if not prompt:
60+
raise ValueError("Prompt cannot be empty.")
61+
prompt = prompt.encode(encoding='ASCII',errors='ignore').decode()
62+
response = openai.Completion.create(
63+
engine=engine,
64+
prompt=prompt,
65+
temperature=temp,
66+
max_tokens=tokens,
67+
top_p=top_p,
68+
frequency_penalty=freq_pen,
69+
presence_penalty=pres_pen
70+
)
71+
if not response.choices:
72+
raise ValueError("No response from OpenAI API.")
73+
text = response.choices[0].text.strip()
74+
if not text:
75+
raise ValueError("Response text cannot be empty.")
76+
return text
77+
except Exception as e:
78+
print(f"Error occurred: {e} while generating prompt.")
79+
return None
80+
81+
# Function to click all upscale buttons.
82+
def select_upscale_option(page, option_text):
83+
# Click the last button on the page that contains the option text.
84+
page.locator(f"button:has-text('{option_text}')").locator("nth=-1").click()
85+
print(f"Clicked {option_text} upscale option.")
86+
87+
# Function to download upscaled images.
88+
def download_upscaled_images(page, prompt_text):
89+
# Wait for all four images to complete rendering by checking the contents of the last 4 messages to see if it
90+
# contains the phrase 'Make Variations', and 'Web'.
91+
messages = page.query_selector_all(".messageListItem-ZZ7v6g")
92+
last_four_messages = messages[-4:]
93+
# Get the inner text of the last four messages by evaluating the innerText property of the node.
94+
for message in last_four_messages:
95+
message = message.evaluate_handle('(node) => node.innerText')
96+
message = str(message)
97+
print(message)
98+
# Check to see if string contains the 'Make Variation' and 'Web'.
99+
if 'Make Variations' and 'Web' in message:
100+
try:
101+
print('Downloading upscaled images.')
102+
try:
103+
# Download last 4 images.
104+
last_image = page.query_selector_all('.originalLink-Azwuo9')[-1]
105+
second_last_image = page.query_selector_all('.originalLink-Azwuo9')[-2]
106+
third_last_image = page.query_selector_all('.originalLink-Azwuo9')[-3]
107+
fourth_last_image = page.query_selector_all('.originalLink-Azwuo9')[-4]
108+
109+
# Loop and download all 4 images with the same name as the prompt, along with last_image, second_last_image, etc.
110+
for image in [last_image, second_last_image, third_last_image, fourth_last_image]:
111+
src = image.get_attribute('href')
112+
url = src
113+
# # Only keep name until first period.
114+
# response = re.sub(r'\w+\.', '', response)
115+
# Remove all special characters from the response using regex.
116+
response = re.sub(r'[^a-zA-Z0-9\s]', '', prompt_text)
117+
# Replace all commas and spaces with underscores.
118+
response = response.replace(',', '_').replace(' ', '_')
119+
# Get first 200 characters of response.
120+
response = response[:200]
121+
r = requests.get(url, stream=True)
122+
with open(f'{str(response) + str(uuid.uuid1())}.png', 'wb') as out_file:
123+
shutil.copyfileobj(r.raw, out_file)
124+
del r
125+
126+
except Exception as e:
127+
print(f"An error, {e}, occured while downloading the images.")
128+
except Exception as e:
129+
print(e)
130+
else:
131+
# Call the function again.
132+
download_upscaled_images(page, prompt_text)
133+
134+
# Function to get the last message.
135+
def get_last_message(page):
136+
# Obtain the list of all messages.
137+
# messageListItem-ZZ7v6g
138+
messages = page.query_selector_all(".messageListItem-ZZ7v6g")
139+
# Select the last message.
140+
last_message = messages[-1]
141+
# Get the text of the last message.
142+
last_message = last_message.evaluate_handle('(node) => node.innerText')
143+
last_message = str(last_message)
144+
print(last_message)
145+
return last_message
146+
147+
# Function to wait for page to fully load and select all upscale options.
148+
def wait_and_select_upscale_options(page, prompt_text):
149+
prompt_text = prompt_text.lower()
150+
try:
151+
last_message = get_last_message(page)
152+
# Check to see if string contains the 'U1'.
153+
if 'U1' in last_message:
154+
print("Found upscale options.")
155+
print("Attempting to upscale all generated images.")
156+
try:
157+
# Select the 'U1' upscale option
158+
select_upscale_option(page, 'U1')
159+
time.sleep(random.randint(3, 5))
160+
# Select the 'U2' upscale option
161+
select_upscale_option(page, 'U2')
162+
time.sleep(random.randint(3, 5))
163+
# Select the 'U3' upscale option
164+
select_upscale_option(page, 'U3')
165+
time.sleep(random.randint(3, 5))
166+
# Select the 'U4' upscale option
167+
select_upscale_option(page, 'U4')
168+
time.sleep(random.randint(3, 5))
169+
except Exception as e:
170+
print("An error occured while selecting upscale options:", e)
171+
download_upscaled_images(page, prompt_text)
172+
else:
173+
print("Photo(s) not fully loaded.")
174+
# KNOWN ISSUE: Script will sometimes start monitoring second last message instead of the last one.
175+
# Keep looping until upscale options are found.
176+
wait_and_select_upscale_options(page, prompt_text)
177+
except Exception as e:
178+
print("An error occurred while finding the last message:", e)
179+
180+
181+
# Function to generate prompt and sumbit command.
182+
def generate_prompt_and_submit_command(page, prompt):
183+
try:
184+
# Generate prompt.
185+
prompt_text = gpt3_midjourney_prompt(prompt)
186+
time.sleep(random.randint(1, 5))
187+
pill_value = page.locator('xpath=//*[@id="app-mount"]/div[2]/div[1]/div[1]/div/div[2]/div/div/div/div/div[3]/div/main/form/div/div[2]/div/div[2]/div/div/div/span[2]/span[2]')
188+
pill_value.fill(prompt_text)
189+
# Submit prompt.
190+
time.sleep(random.randint(1, 5))
191+
# Press the Enter key.
192+
page.keyboard.press("Enter")
193+
print(f'Successfully submitted prompt: {prompt_text}')
194+
wait_and_select_upscale_options(page, prompt_text)
195+
except Exception as e:
196+
print("An error occurred while submitting the prompt:", e)
197+
198+
# Function to enter in bot command.
199+
def bot_command(page, command):
200+
try:
201+
print("Clicking on chat bar.")
202+
chat_bar = page.locator('xpath=//*[@id="app-mount"]/div[2]/div[1]/div[1]/div/div[2]/div/div/div/div/div[3]/div[2]/main/form/div/div[1]/div/div[3]/div/div[2]/div')
203+
time.sleep(random.randint(1, 5))
204+
print("Typing in bot command")
205+
chat_bar.fill(command)
206+
time.sleep(random.randint(1, 5))
207+
print("Selecting the prompt option in the suggestions menu")
208+
# Select the first option in the pop-up menu.
209+
prompt_option = page.locator('xpath=/html/body/div[1]/div[2]/div[1]/div[1]/div/div[2]/div/div/div/div/div[3]/div[2]/main/form/div/div[2]/div/div/div[2]/div[1]/div/div/div')
210+
time.sleep(random.randint(1, 5))
211+
# Click on the prompt option.
212+
prompt_option.click()
213+
print("Generating prompt using OpenAI's API.")
214+
generate_prompt_and_submit_command(page, PROMPT)
215+
except Exception as e:
216+
print("An error occurred while entering in the prompt:", e)
217+
218+
# Main function to login to Discord and run the bot.
219+
if __name__ == "__main__":
220+
with sync_playwright() as p:
221+
browser = p.chromium.launch(headless=False)
222+
# Create a new incognito browser context.
223+
page = browser.new_page()
224+
# Go to Discord login page.
225+
page.goto("https://www.discord.com/login")
226+
# Open credentials file and read in email and password.
227+
with open("credentials.txt", "r") as f:
228+
email = f.readline()
229+
password = f.readline()
230+
# Fill in email and password fields.
231+
print("Entering email")
232+
page.fill("input[name='email']", email)
233+
time.sleep(random.randint(1, 5))
234+
print("Entering password")
235+
page.fill("input[name='password']", password)
236+
time.sleep(random.randint(1, 5))
237+
# Click login button.
238+
print("Logging into Discord")
239+
page.click("button[type='submit']")
240+
time.sleep(random.randint(5, 10))
241+
# Wait for page URL to change for 15 seconds.
242+
page.wait_for_url("https://discord.com/channels/@me", timeout=15000)
243+
print("Successfully logged into Discord.")
244+
time.sleep(random.randint(1, 5))
245+
# Run the bot for 10 iterations.
246+
for i in range(10):
247+
open_discord_channel(page)
248+
# Go to appropriate channel.
249+
page.goto(f"{CHANNEL_URL}")
250+
time.sleep(random.randint(1, 5))
251+
page.wait_for_load_state("networkidle") # This waits for the "networkidle"
252+
print("Opened appropriate channel.")
253+
print("Entering in the specified bot command.")
254+
bot_command(page, BOT_COMMAND)
255+
# Print the number of iterations completed.
256+
print(f"Iteration {i+1} completed.")
257+
258+
259+
260+
261+
262+
263+

0 commit comments

Comments
 (0)