-
Notifications
You must be signed in to change notification settings - Fork 63
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
- Loading branch information
1 parent
674da5b
commit 054bd70
Showing
6 changed files
with
244 additions
and
241 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -1,177 +1,56 @@ | ||
r""" | ||
_____ | ||
_____/ ____\______ ________________ ____ ___________ | ||
/ _ \ __\/ ___// ___\_ __ \__ \ / _ \_/ __ \_ __ \ | ||
( <_> ) | \___ \\ \___| | \// __ \( <_> ) ___/| | \/ | ||
\____/|__| /____ >\___ >__| (____ /\____/ \___ >__| | ||
\/ \/ \/ \/ | ||
""" | ||
import asyncio | ||
import logging | ||
import pathlib | ||
import traceback | ||
|
||
from rich.live import Live | ||
from rich.table import Column | ||
|
||
try: | ||
from win32_setctime import setctime # pylint: disable=import-error | ||
except ModuleNotFoundError: | ||
pass | ||
import ofscraper.classes.placeholder as placeholder | ||
import ofscraper.classes.sessionbuilder as sessionbuilder | ||
import ofscraper.constants as constants | ||
import ofscraper.download.common as common | ||
import ofscraper.utils.console as console | ||
import ofscraper.utils.exit as exit | ||
import ofscraper.utils.paths as paths | ||
import ofscraper.utils.stdout as stdout | ||
from ofscraper.classes.semaphoreDelayed import semaphoreDelayed | ||
from ofscraper.download.alt_download import alt_download | ||
from ofscraper.download.common import ( | ||
convert_num_bytes, | ||
get_medialog, | ||
reset_globals, | ||
setDirectoriesDate, | ||
setupProgressBar, | ||
) | ||
from ofscraper.download.main_download import main_download | ||
from ofscraper.utils.run_async import run | ||
|
||
|
||
@run | ||
async def process_dicts(username, model_id, medialist): | ||
with stdout.lowstdout(): | ||
progress_group, overall_progress, job_progress = setupProgressBar() | ||
# This need to be here: https://stackoverflow.com/questions/73599594/asyncio-works-in-python-3-10-but-not-in-python-3-8 | ||
reset_globals() | ||
common.log = logging.getLogger("ofscraper-download") | ||
|
||
try: | ||
with Live( | ||
progress_group, | ||
refresh_per_second=constants.refreshScreen, | ||
console=console.shared_console, | ||
): | ||
aws = [] | ||
photo_count = 0 | ||
video_count = 0 | ||
audio_count = 0 | ||
skipped = 0 | ||
forced_skipped = 0 | ||
total_downloaded = 0 | ||
sum = 0 | ||
desc = "Progress: ({p_count} photos, {v_count} videos, {a_count} audios, {forced_skipped} skipped, {skipped} failed || {sumcount}/{mediacount}||{total_bytes_download}/{total_bytes})" | ||
|
||
async with sessionbuilder.sessionBuilder() as c: | ||
for ele in medialist: | ||
aws.append( | ||
asyncio.create_task( | ||
download(c, ele, model_id, username, job_progress) | ||
) | ||
) | ||
task1 = overall_progress.add_task( | ||
desc.format( | ||
p_count=photo_count, | ||
v_count=video_count, | ||
a_count=audio_count, | ||
skipped=skipped, | ||
mediacount=len(medialist), | ||
forced_skipped=forced_skipped, | ||
sumcount=0, | ||
total_bytes_download=0, | ||
total_bytes=0, | ||
), | ||
total=len(aws), | ||
visible=True, | ||
) | ||
for coro in asyncio.as_completed(aws): | ||
try: | ||
pack = await coro | ||
common.log.debug(f"unpack {pack} count {len(pack)}") | ||
media_type, num_bytes_downloaded = pack | ||
except Exception as e: | ||
common.log.traceback_(e) | ||
common.log.traceback_(traceback.format_exc()) | ||
media_type = "skipped" | ||
num_bytes_downloaded = 0 | ||
|
||
total_downloaded += num_bytes_downloaded | ||
total_bytes_downloaded = convert_num_bytes(total_downloaded) | ||
total_bytes = convert_num_bytes(common.total_data) | ||
if media_type == "images": | ||
photo_count += 1 | ||
|
||
elif media_type == "videos": | ||
video_count += 1 | ||
elif media_type == "audios": | ||
audio_count += 1 | ||
elif media_type == "skipped": | ||
skipped += 1 | ||
elif media_type == "forced_skipped": | ||
forced_skipped += 1 | ||
sum += 1 | ||
overall_progress.update( | ||
task1, | ||
description=desc.format( | ||
p_count=photo_count, | ||
v_count=video_count, | ||
a_count=audio_count, | ||
skipped=skipped, | ||
forced_skipped=forced_skipped, | ||
mediacount=len(medialist), | ||
sumcount=sum, | ||
total_bytes=total_bytes, | ||
total_bytes_download=total_bytes_downloaded, | ||
), | ||
refresh=True, | ||
advance=1, | ||
) | ||
overall_progress.remove_task(task1) | ||
setDirectoriesDate() | ||
common.log.warning( | ||
f"[bold]{username}[/bold] ({photo_count+audio_count+video_count} downloads total [{video_count} videos, {audio_count} audios, {photo_count} photos] {forced_skipped} skipped, {skipped} failed)" | ||
) | ||
return photo_count, video_count, audio_count, forced_skipped, skipped | ||
|
||
except Exception as E: | ||
with exit.DelayedKeyboardInterrupt(): | ||
raise E | ||
finally: | ||
await asyncio.get_event_loop().run_in_executor( | ||
common.cache_thread, common.cache.close | ||
) | ||
common.cache_thread.shutdown() | ||
|
||
|
||
async def download(c, ele, model_id, username, progress): | ||
async with common.maxfile_sem: | ||
with paths.set_directory( | ||
placeholder.Placeholders().getmediadir(ele, username, model_id) | ||
): | ||
try: | ||
if ele.url: | ||
return await main_download( | ||
c, | ||
ele, | ||
pathlib.Path(".").absolute(), | ||
username, | ||
model_id, | ||
progress, | ||
) | ||
elif ele.mpd: | ||
return await alt_download( | ||
c, | ||
ele, | ||
pathlib.Path(".").absolute(), | ||
username, | ||
model_id, | ||
progress, | ||
) | ||
except Exception as E: | ||
common.log.debug(f"{get_medialog(ele)} exception {E}") | ||
common.log.debug( | ||
f"{get_medialog(ele)} exception {traceback.format_exc()}" | ||
) | ||
return "skipped", 0 | ||
import time | ||
|
||
import httpx | ||
|
||
import ofscraper.db.operations as operations | ||
import ofscraper.download.downloadbatch as batchdownloader | ||
import ofscraper.download.downloadnormal as normaldownloader | ||
import ofscraper.utils.args as args_ | ||
import ofscraper.utils.config as config_ | ||
import ofscraper.utils.separate as seperate | ||
import ofscraper.utils.system as system | ||
|
||
|
||
def medialist_filter(medialist, model_id, username): | ||
log = logging.getLogger("shared") | ||
if not args_.getargs().dupe: | ||
media_ids = set( | ||
operations.get_media_ids_downloaded(model_id=model_id, username=username) | ||
) | ||
log.debug( | ||
f"Number of unique media ids in database for {username}: {len(media_ids)}" | ||
) | ||
medialist = seperate.separate_by_id(medialist, media_ids) | ||
log.debug(f"Number of new mediaids with dupe ids removed: {len(medialist)}") | ||
medialist = seperate.seperate_avatars(medialist) | ||
log.debug(f"Removed previously downloaded avatars/headers") | ||
log.debug(f"Final Number of media to download {len(medialist)}") | ||
|
||
else: | ||
log.info(f"forcing all downloads media count {len(medialist)}") | ||
return medialist | ||
|
||
|
||
def download_picker(username, model_id, medialist): | ||
medialist = medialist_filter(medialist, model_id, username) | ||
|
||
if len(medialist) == 0: | ||
logging.getLogger("shared").error( | ||
f"[bold]{username}[/bold] ({0} photos, {0} videos, {0} audios, {0} skipped, {0} failed)" | ||
) | ||
return 0, 0, 0, 0, 0 | ||
elif ( | ||
system.getcpu_count() > 1 | ||
and ( | ||
len(medialist) >= config_.get_download_semaphores(config_.read_config()) * 5 | ||
) | ||
and ( | ||
args_.getargs().downloadthreads | ||
or config_.get_threads(config_.read_config()) | ||
) | ||
> 0 | ||
): | ||
return batchdownloader.process_dicts(username, model_id, medialist) | ||
else: | ||
return normaldownloader.process_dicts(username, model_id, medialist) |
Oops, something went wrong.