186 lines
8.3 KiB
Python
186 lines
8.3 KiB
Python
import random
|
|
import string
|
|
import requests
|
|
import os
|
|
import undetected_chromedriver as uc
|
|
from selenium.webdriver.common.by import By
|
|
from selenium.webdriver.support.ui import WebDriverWait
|
|
from selenium.webdriver.support import expected_conditions as EC
|
|
import time
|
|
import traceback
|
|
from urllib.parse import urlparse
|
|
|
|
MIN_FILE_SIZE_BYTES = 10 * 1024
|
|
|
|
class MJScraperNode:
|
|
def __init__(self):
|
|
pass
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"search_url": ("STRING", {"default": "https://midjourney.com/showcase/top/"}),
|
|
"output_directory": ("STRING", {"default": "./Midjourney_Scraped_Images/"}),
|
|
"num_images_to_scrape": ("INT", {"default": 10, "min": 0, "max": 1000, "step": 1}),
|
|
"timeout": ("INT", {"default": 5, "min": 1, "max": 60, "step": 1}),
|
|
"headless": ("BOOLEAN", {"default": True}),
|
|
"cookies": ("STRING", {"default": ""}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("STRING",)
|
|
RETURN_NAMES = ("image_paths",)
|
|
FUNCTION = "scrape_images"
|
|
CATEGORY = "Image Scrapers"
|
|
|
|
def scrape_images(self, search_url, output_directory, num_images_to_scrape, timeout, headless, cookies):
|
|
img_links = []
|
|
driver = None
|
|
print("Starting MJ Scraper Node...")
|
|
try:
|
|
options = uc.ChromeOptions()
|
|
if headless:
|
|
options.add_argument("--headless")
|
|
options.add_argument("--no-sandbox")
|
|
options.add_argument("--disable-dev-shm-usage")
|
|
|
|
print("Initializing undetected_chromedriver...")
|
|
driver = uc.Chrome(options=options)
|
|
print("undetected_chromedriver initialized successfully.")
|
|
|
|
driver.implicitly_wait(timeout)
|
|
print(f"Navigating to: {search_url}")
|
|
driver.get(search_url)
|
|
print(f"Navigation complete. Current URL: {driver.current_url}")
|
|
|
|
if cookies:
|
|
print("Injecting provided cookies...")
|
|
for cookie_pair in cookies.split(';'):
|
|
if '=' in cookie_pair:
|
|
name, value = cookie_pair.strip().split('=', 1)
|
|
try:
|
|
driver.add_cookie({'name': name, 'value': value, 'domain': '.midjourney.com'})
|
|
print(f"Added cookie: {name}")
|
|
except Exception as e:
|
|
print(f"Error adding cookie {name}: {e}")
|
|
driver.get(search_url)
|
|
print(f"Re-navigated after cookie injection. Current URL: {driver.current_url}")
|
|
|
|
wait = WebDriverWait(driver, timeout)
|
|
|
|
last_height = driver.execute_script("return document.body.scrollHeight")
|
|
while True:
|
|
driver.execute_script("window.scrollTo(0, document.body.scrollHeight);")
|
|
try:
|
|
wait.until(EC.presence_of_element_located((By.CSS_SELECTOR, "img[src*='cdn.midjourney.com'], video[src*='cdn.midjourney.com']")))
|
|
except:
|
|
pass
|
|
|
|
time.sleep(3)
|
|
new_height = driver.execute_script("return document.body.scrollHeight")
|
|
if new_height == last_height:
|
|
break
|
|
last_height = new_height
|
|
|
|
image_elements = driver.find_elements(By.CSS_SELECTOR, "img[src*='cdn.midjourney.com']")
|
|
video_elements = driver.find_elements(By.CSS_SELECTOR, "video[src*='cdn.midjourney.com']")
|
|
|
|
all_elements = image_elements + video_elements
|
|
print(f"Found {len(all_elements)} potential links (images and videos).")
|
|
|
|
for elem in all_elements:
|
|
link = elem.get_attribute('src')
|
|
if not link:
|
|
link = elem.get_attribute('href')
|
|
|
|
if link:
|
|
print(f"Scraped link: {link}")
|
|
|
|
if "/video/" in link and link.endswith(".webp"):
|
|
link = link.replace("_640_N.webp", ".mp4")
|
|
print(f"Transformed video link to: {link}")
|
|
elif ".webp" in link:
|
|
link = link.replace("0_0_32_N.webp", "0_0.png")
|
|
link = link.replace("0_1_32_N.webp", "0_0.png")
|
|
link = link.replace("0_2_32_N.webp", "0_0.png")
|
|
link = link.replace("0_3_32_N.webp", "0_0.png")
|
|
link = link.replace("grid_0_32_N.webp", "0_0.png")
|
|
|
|
img_links.append(link)
|
|
if num_images_to_scrape > 0 and len(img_links) >= num_images_to_scrape:
|
|
print(f"Reached desired number of items ({num_images_to_scrape}). Stopping scraping.")
|
|
break
|
|
else:
|
|
print("Warning: Found an element but could not extract 'src' or 'href' attribute. Skipping.")
|
|
|
|
os.makedirs(output_directory, exist_ok=True)
|
|
|
|
downloaded_image_paths = []
|
|
print(f"Attempting to download {len(img_links)} items.")
|
|
for i, item_link in enumerate(img_links):
|
|
retries = 9
|
|
for attempt in range(retries):
|
|
try:
|
|
print(f"Downloading item {i+1}/{len(img_links)}: {item_link} (Attempt {attempt + 1}/{retries})")
|
|
time.sleep(0.8)
|
|
download_headers = {
|
|
'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.124 Safari/537.36',
|
|
'Accept': 'image/webp,image/apng,image/*,*/*;q=0.8,video/*',
|
|
'Accept-Language': 'en-US,en;q=0.9',
|
|
'Connection': 'keep-alive',
|
|
'Referer': search_url
|
|
}
|
|
item_data = requests.get(item_link, headers=download_headers, timeout=10).content
|
|
|
|
parsed_url = urlparse(item_link)
|
|
path = parsed_url.path
|
|
ext = os.path.splitext(path)[1]
|
|
if not ext:
|
|
ext = ".png"
|
|
|
|
item_name = ''.join(random.choices(string.ascii_letters + string.digits, k=8)) + ext
|
|
file_path = os.path.join(output_directory, item_name)
|
|
with open(file_path, 'wb') as handler:
|
|
handler.write(item_data)
|
|
|
|
file_size = os.path.getsize(file_path)
|
|
if file_size < MIN_FILE_SIZE_BYTES:
|
|
print(f"Warning: Downloaded file {item_name} is only {file_size} bytes, which is less than {MIN_FILE_SIZE_BYTES} bytes. Considering it corrupted.")
|
|
os.remove(file_path)
|
|
raise requests.exceptions.RequestException("File too small, likely corrupted.")
|
|
|
|
downloaded_image_paths.append(file_path)
|
|
print(f"Downloaded: {item_name} ({file_size} bytes)")
|
|
break
|
|
except requests.exceptions.RequestException as e:
|
|
print(f"Error downloading {item_link}: {e}")
|
|
print(f"RequestException details: {e}")
|
|
if attempt < retries - 1:
|
|
print(f"Retrying download in 2 seconds...")
|
|
time.sleep(2)
|
|
else:
|
|
print(f"Failed to download {item_link} after {retries} attempts.")
|
|
except Exception as e:
|
|
print(f"Error saving item: {e}")
|
|
print(f"General Exception details: {e}")
|
|
break
|
|
|
|
return (",".join(downloaded_image_paths),)
|
|
|
|
except Exception as e:
|
|
print(f"An unhandled error occurred during scraping: {e}")
|
|
traceback.print_exc()
|
|
return ("",)
|
|
finally:
|
|
if driver:
|
|
driver.quit()
|
|
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"MJScraper": MJScraperNode
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"MJScraper": "Midjourney Scraper"
|
|
} |