From 6b8299eac86fce357d7ea8cef2614461d47e7ece Mon Sep 17 00:00:00 2001 From: Vladimir Mandic Date: Mon, 31 Mar 2025 16:21:20 -0400 Subject: [PATCH] lora mpt4 preview support Signed-off-by: Vladimir Mandic --- CHANGELOG.md | 2 ++ modules/modelloader.py | 42 +++++++++++++++++++++++++++++------ modules/ui_control_helpers.py | 26 ++-------------------- modules/ui_models.py | 3 +++ modules/video.py | 23 +++++++++++++++++++ 5 files changed, 65 insertions(+), 31 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index e7230fbf7..9af36d7bd 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -36,6 +36,8 @@ Plus... - **LoRA** - support for *Hunyuan, LTX, WAN, Mochi, Cog* - add option to apply LoRA directly on GPU or use CPU first in low-memory scenarios + - improve metadata and preview parallel fetch + - support for mp4 so first frame is extracted as used as lora preview - additional key points: - all models are auto-downloaded upon first use uses *system paths -> huggingface* folder diff --git a/modules/modelloader.py b/modules/modelloader.py index 43325b7c3..412ef0aa1 100644 --- a/modules/modelloader.py +++ b/modules/modelloader.py @@ -18,6 +18,7 @@ from modules.paths import script_path, models_path loggedin = None diffuser_repos = [] debug = shared.log.trace if os.environ.get('SD_DOWNLOAD_DEBUG', None) is not None else lambda *args, **kwargs: None +pbar = None def hf_login(token=None): @@ -61,12 +62,34 @@ def download_civit_meta(model_path: str, model_id): return f'CivitAI download error: id={model_id} url={url} code={r.status_code}' +def save_video_frame(filepath: str): + from modules import video + try: + frames, fps, duration, w, h, codec, frame = video.get_video_params(filepath, capture=True) + except Exception as e: + shared.log.error(f'Video: file={filepath} {e}') + return None + if frame is not None: + basename = os.path.splitext(filepath) + thumb = f'{basename[0]}.thumb.jpg' + shared.log.debug(f'Video: file={filepath} frames={frames} fps={fps} size={w}x{h} codec={codec} duration={duration} thumb={thumb}') + frame.save(thumb) + else: + shared.log.error(f'Video: file={filepath} no frames found') + return frame + + def download_civit_preview(model_path: str, preview_url: str): + global pbar # pylint: disable=global-statement + if model_path is None: + pbar = None + return '' ext = os.path.splitext(preview_url)[1] preview_file = os.path.splitext(model_path)[0] + ext is_video = preview_file.lower().endswith('.mp4') - if is_video: - shared.log.warning(f'CivitAI download: url="{preview_url}" skip video') + is_json = preview_file.lower().endswith('.json') + if is_json: + shared.log.warning(f'CivitAI download: url="{preview_url}" skip json') return '' if os.path.exists(preview_file): return '' @@ -77,20 +100,25 @@ def download_civit_preview(model_path: str, preview_url: str): written = 0 img = None shared.state.begin('CivitAI') + if pbar is None: + pbar = p.Progress(p.TextColumn('[cyan]{task.description}'), p.DownloadColumn(), p.BarColumn(), p.TaskProgressColumn(), p.TimeRemainingColumn(), p.TimeElapsedColumn(), p.TransferSpeedColumn(), console=shared.console) try: with open(preview_file, 'wb') as f: - with p.Progress(p.TextColumn('[cyan]{task.description}'), p.DownloadColumn(), p.BarColumn(), p.TaskProgressColumn(), p.TimeRemainingColumn(), p.TimeElapsedColumn(), p.TransferSpeedColumn(), console=shared.console) as progress: - task = progress.add_task(description="Download starting", total=total_size) + with pbar: + task = pbar.add_task(description="Download starting", total=total_size) for data in r.iter_content(block_size): written = written + len(data) f.write(data) - progress.update(task, advance=block_size, description="Downloading") + pbar.update(task, advance=block_size, description="Downloading") if written < 1024: # min threshold os.remove(preview_file) raise ValueError(f'removed invalid download: bytes={written}') - img = Image.open(preview_file) + if is_video: + img = save_video_frame(preview_file) + else: + img = Image.open(preview_file) except Exception as e: - os.remove(preview_file) + # os.remove(preview_file) res += f' error={e}' shared.log.error(f'CivitAI download error: url={preview_url} file="{preview_file}" written={written} {e}') shared.state.end() diff --git a/modules/ui_control_helpers.py b/modules/ui_control_helpers.py index 2b10e21d3..de8983be2 100644 --- a/modules/ui_control_helpers.py +++ b/modules/ui_control_helpers.py @@ -1,7 +1,7 @@ import os import gradio as gr from PIL import Image -from modules import shared, scripts, masking # pylint: disable=ungrouped-imports +from modules import shared, scripts, masking, video # pylint: disable=ungrouped-imports gr_height = None @@ -82,31 +82,9 @@ def display_units(num_units): return (num_units * [gr.update(visible=True)]) + ((max_units - num_units) * [gr.update(visible=False)]) -def get_video_params(filepath: str, capture: bool = False): - import cv2 - from modules.control.util import decode_fourcc - video = cv2.VideoCapture(filepath) - if not video.isOpened(): - msg = f'Control: video open failed: path="{filepath}"' - shared.log.error(msg) - raise RuntimeError(msg) - frames = int(video.get(cv2.CAP_PROP_FRAME_COUNT)) - fps = video.get(cv2.CAP_PROP_FPS) - duration = float(frames) / fps - w, h = int(video.get(cv2.CAP_PROP_FRAME_WIDTH)), int(video.get(cv2.CAP_PROP_FRAME_HEIGHT)) - codec = decode_fourcc(video.get(cv2.CAP_PROP_FOURCC)) - frame = None - if capture: - _status, frame = video.read() - frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) - frame = Image.fromarray(frame) - video.release() - return frames, fps, duration, w, h, codec, frame - - def get_video(filepath: str): try: - frames, fps, duration, w, h, codec, _cap = get_video_params(filepath) + frames, fps, duration, w, h, codec, _cap = video.get_video_params(filepath) shared.log.debug(f'Control: input video: path={filepath} frames={frames} fps={fps} size={w}x{h} codec={codec}') msg = f'Control input | Video | Size {w}x{h} | Frames {frames} | FPS {fps:.2f} | Duration {duration:.2f} | Codec {codec}' return msg diff --git a/modules/ui_models.py b/modules/ui_models.py index 4f6355a9a..59d8a9196 100644 --- a/modules/ui_models.py +++ b/modules/ui_models.py @@ -573,6 +573,8 @@ def create_ui(): def atomic_civit_search_metadata(item, res, rehash): from modules.modelloader import download_civit_preview, download_civit_meta + if item is None: + return meta = os.path.splitext(item['filename'])[0] + '.json' has_meta = os.path.isfile(meta) and os.stat(meta).st_size > 0 if ('card-no-preview.png' in item['preview'] or not has_meta) and os.path.isfile(item['filename']): @@ -629,6 +631,7 @@ def create_ui(): with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as executor: for fn in candidates: executor.submit(atomic_civit_search_metadata, fn, res, rehash) + atomic_civit_search_metadata(None, res, rehash) t1 = time.time() log.debug(f'CivitAI search metadata: items={i} time={t1-t0:.2f}') txt = '
'.join([r for r in res if len(r) > 0]) diff --git a/modules/video.py b/modules/video.py index d9e40a27f..bd26aac04 100644 --- a/modules/video.py +++ b/modules/video.py @@ -1,6 +1,7 @@ import os import threading import numpy as np +from PIL import Image from modules import shared, errors from modules.images_namegen import FilenameGenerator # pylint: disable=unused-import @@ -83,3 +84,25 @@ def save_video(p, images, filename = None, video_type: str = 'none', duration: f else: save_video_atomic(images, filename, video_type, duration, loop, interpolate, scale, pad, change) return filename + + +def get_video_params(filepath: str, capture: bool = False): + import cv2 + from modules.control.util import decode_fourcc + video = cv2.VideoCapture(filepath) + if not video.isOpened(): + msg = f'Video open failed: path="{filepath}"' + shared.log.error(msg) + raise RuntimeError(msg) + frames = int(video.get(cv2.CAP_PROP_FRAME_COUNT)) + fps = round(video.get(cv2.CAP_PROP_FPS), 2) + duration = round(float(frames) / fps, 2) + w, h = int(video.get(cv2.CAP_PROP_FRAME_WIDTH)), int(video.get(cv2.CAP_PROP_FRAME_HEIGHT)) + codec = decode_fourcc(video.get(cv2.CAP_PROP_FOURCC)) + frame = None + if capture: + _status, frame = video.read() + frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) + frame = Image.fromarray(frame) + video.release() + return frames, fps, duration, w, h, codec, frame