diff --git a/THIRD_PARTY_NOTICES.md b/THIRD_PARTY_NOTICES.md index 38c85761..e90529bf 100644 --- a/THIRD_PARTY_NOTICES.md +++ b/THIRD_PARTY_NOTICES.md @@ -13,6 +13,7 @@ FreeVideo code uses [Apache-2.0](LICENSE). Its dependencies and models retain th | [VDN-H3 and H3 model weights](https://huggingface.co/OpenVDN/vdn-minimax-h3-edge/blob/main/LICENSE) | MiniMax H3 Community License | | [MiniMax H3 latent upscaler](https://huggingface.co/LBH-123-AI/Minimax_h3_latent_Upscaler) | [MIT](freevideo_engine/licenses/latent-upscaler-MIT.txt) | | [FreeToken](https://github.com/FlashML-org/FreeToken) prompt VLM operators, [Qwen3-VL-4B-Instruct](https://huggingface.co/Qwen/Qwen3-VL-4B-Instruct) optional weights | [Apache-2.0](freevideo_engine/licenses/FreeToken-Apache-2.0.txt) | +| [Qwen-Image-2.1](https://huggingface.co/Comfy-Org/Qwen-Image-2.1) Turbo image generation models, including Qwen3-VL-8B and VAE (optional download; [notice](freevideo_engine/image_gen/licenses/NOTICE.txt) installed with the models), © 2026 Hangzhou Tongyi Laboratory Technology Co., Ltd. All Rights Reserved. | [Qwen RESEARCH LICENSE AGREEMENT](https://huggingface.co/Comfy-Org/Qwen-Image-2.1/blob/main/LICENSE) | | [sherpa-onnx](https://github.com/k2-fsa/sherpa-onnx) 1.13.8 speech runtime, built by FreeVideo without text-to-speech, ALSA and PortAudio (modified build; `__init__.py` edited; source snapshot 11afbd00 kept by FreeVideo; optional download for automatic subtitles, [build note](freevideo_engine/licenses/sherpa-onnx-build-NOTE.txt)). Includes statically linked kaldi-native-fbank, kaldi-decoder, kaldifst, OpenFst, simple-sentencepiece (Apache-2.0), KissFFT (BSD-3-Clause), pybind11 (BSD-3-Clause), nlohmann/json (MIT), hclust-cpp and Darts-clone (BSD-2-Clause), ThreadPool (Zlib). The Windows build also contains msvcp140.dll and msvcp140_1.dll (Microsoft Visual C++ 2022 runtime 14.44.35211, © Microsoft Corporation), unmodified and redistributed as permitted by Microsoft for its redistributable files. | Apache-2.0 | | [ONNX Runtime](https://github.com/microsoft/onnxruntime) 1.28.2 (CPU), © Microsoft Corporation; its third-party notices are shipped inside the sherpa-onnx wheel | MIT | | [Eigen](https://gitlab.com/libeigen/eigen) 5.0.1, compiled into sherpa-onnx and ONNX Runtime; source: https://gitlab.com/libeigen/eigen/-/tree/5.0.1 | MPL-2.0 | diff --git a/__init__.py b/__init__.py index 89766033..3f20b504 100644 --- a/__init__.py +++ b/__init__.py @@ -17,6 +17,12 @@ async def comfy_entrypoint(): setup_routes() from .freevideo_engine.prompt_vlm.service import register as prompt_routes prompt_routes() + try: + from .freevideo_engine.image_gen.service import register as image_routes + image_routes() + except Exception as error: + import logging + logging.warning('[FreeVideo] Image generation routes not registered: %s', error) from .freevideo_engine.comfy_launcher_api import register as launcher_routes launcher_routes() from .freevideo_engine.comfy_updates import register as update_routes diff --git a/freevideo_engine/agent/cast.py b/freevideo_engine/agent/cast.py index 767f5495..3178bde2 100644 --- a/freevideo_engine/agent/cast.py +++ b/freevideo_engine/agent/cast.py @@ -13,6 +13,7 @@ from pathlib import Path import re import secrets +import shutil import threading import uuid @@ -537,7 +538,7 @@ def used(self, identities, film): # Reference pictures ------------------------------------------------------------------------------- def generate(self, identity, value=None, *, by='builtin', ref=None, revision=None, client=None): - """Submit a reference clip (and a face close-up) for this card; the scheduler lands the frames.""" + """Submit reference pictures or clips for this card; the scheduler lands the pictures.""" value = value or {} shape = schema.cast_schema('generate') schema.validate(shape, value) @@ -560,26 +561,79 @@ def generate(self, identity, value=None, *, by='builtin', ref=None, revision=Non card=view(card), job=job) if value['face'] and card['kind'] != 'character': raise spec.SpecError([schema.diagnostic('/face', 'false for a setting card', True)]) + ready = value.get('source') != 'video' and getattr(self.jobs.port, 'image_ready', lambda: False)() + source = value.get('source', 'image' if ready else 'video') + if source == 'image' and not ready: + raise CastError('pack_missing', errors.IMAGE_ERRORS['pack_missing'], + 'omit source, or use source video') seed = secrets.randbelow(spec.MAX_SEED + 1) token = uuid.uuid4().hex[:12] - parts = [('turn', SCENE if card['kind'] == 'setting' else TURN, False)] - if value['face']: - parts.append(('face', CLOSE_UP, True)) - items = [dict(kind='generate', prompt=reference_prompt(card, face=face), seed=seed, - label=('Cast · %s · %s' % (card['name'], part))[:120], - metadata=dict(cast_card=identity, cast_part=part), **shot) - for part, shot, face in parts] + if source == 'image': + from ..image_gen.cast_views import VIEWS, SETTING_VIEWS + # A picture takes seconds, so a character always gets every view, the face close-up included. + views = list(SETTING_VIEWS if card['kind'] == 'setting' else VIEWS) + items = self._image_items(card, views[:1], seed) + else: + parts = [('turn', SCENE if card['kind'] == 'setting' else TURN, False)] + if value['face']: + parts.append(('face', CLOSE_UP, True)) + items = [dict(kind='generate', prompt=reference_prompt(card, face=face), seed=seed, + label=('Cast · %s · %s' % (card['name'], part))[:120], + metadata=dict(cast_card=identity, cast_part=part), **shot) + for part, shot, face in parts] rows, _ = self.jobs.submit_request(dict(jobs=items), client='cast:' + identity, parent=dict(cast=identity, ref=token)) jobs = [row['id'] for row in rows] card['generating'] = dict(job=jobs[0], face_job=jobs[1] if len(jobs) > 1 else None, - replace=value['replace'], at=self._now(), by=by) + jobs=list(jobs), source=source, replace=value['replace'], at=self._now(), by=by) card['_generating'] = dict(ref=token, seed=seed, kind=card['kind'], user_images=card.get('_user_images', 0)) + if source == 'image': + card['_generating'].update(source=source, stage='front', views=views[1:], + job_views={jobs[0]: views[0]}, name=card['name'], + description=card['description']) result = dict(job=jobs[0], jobs=jobs, card=view(card)) self._remember(record, ref, digest, result, client) self._save(record) return result + def _image_items(self, card, views, seed, front=None, from_photo=False): + from ..image_gen.cast_views import size, uses_front, view_prompt, photo_prompt + # A picture somebody put on the card (the user, or an agent passing on the user's photo) is the person. + photo = next((row['media'] for row in card.get('images', []) if not row.get('generated')), None) + items = [] + for part in views: + width, height = size(part) + images = [front] if uses_front(part) else [photo] if part == 'front' and photo else [] + # When the front was made from a photo, the photo and not the text says what the person looks like. + prompt = (photo_prompt() if part == 'front' and photo else + view_prompt(dict(card, photo=True) if from_photo else card, part)) + items.append(dict(kind='image', prompt=prompt, width=width, height=height, seed=seed, + images=images, label=('Cast · %s · %s' % (card['name'], part))[:120], + metadata=dict(cast_card=card['id'], cast_part=part))) + return items + + def _image_views(self, identity, generating, hidden, rows): + """Submit edits of the finished front once, while this tick owns the card's claim.""" + front = rows[generating['job']] + if hidden['stage'] != 'front' or not hidden['views'] or not front or front['status'] != 'succeeded': + return False + with self.changed: + record = self.records.get(identity) + if record is None or record['card'].get('_generating') != hidden: + return True + record = deepcopy(record) + card = record['card'] + basis = dict(id=identity, kind=hidden['kind'], name=hidden['name'], description=hidden['description']) + items = self._image_items(basis, hidden['views'], hidden['seed'], dict(file=front['result']['file']), + from_photo=bool(front['spec'].get('images'))) + jobs, _ = self.jobs.submit_request(dict(jobs=items, ref='cast-' + hidden['ref'] + '-views'), + client='cast:' + identity, parent=dict(cast=identity, ref=hidden['ref'])) + card['generating']['jobs'].extend(row['id'] for row in jobs) + card['_generating']['job_views'].update({row['id']: part for row, part in zip(jobs, hidden['views'])}) + card['_generating'].update(stage='views', views=[]) + self._save(record) + return True + def allows(self, parent): """May Jobs still start this generation? Only while its card waits for it.""" with self.changed: @@ -588,7 +642,7 @@ def allows(self, parent): return bool(hidden) and hidden['ref'] == parent.get('ref') def _finished(self): - """Cards whose generation jobs have all ended, with those jobs as they ended.""" + """Cards ready to submit more views or land, with their current jobs as they ended.""" ready = [] with self.changed: for identity, record in self.records.items(): @@ -618,7 +672,11 @@ def tick(self): if not self._claim(identity, hidden): continue # Landed already, or another tick is landing it now. try: - self._land(identity, generating, hidden, rows) + if hidden.get('source') == 'image': + if not self._image_views(identity, generating, hidden, rows): + self._land_image(identity, generating, hidden, rows) + else: + self._land(identity, generating, hidden, rows) except Exception as error: # noqa: BLE001 - one card must never stop the others logging.getLogger(__name__).warning('Reference pictures could not be saved (%s)', type(error).__name__) self._fail(identity, hidden, dict(code='internal_error', job=generating['job'], @@ -632,11 +690,44 @@ def _claim(self, identity, hidden): with self.changed: record = self.records.get(identity) if (identity in self.landing or record is None - or (record['card'].get('_generating') or {}).get('ref') != hidden['ref']): + or record['card'].get('_generating') != hidden): return False self.landing.add(identity) return True + def _land_image(self, identity, generating, hidden, rows): + pictures, failure, staged = [], None, [] + try: + for job in _generation_jobs(generating): + row = rows[job] + if row is None or row['status'] != 'succeeded': + code = (row.get('error') or {}).get('code') if row else 'not_found' + failure = failure or dict(code=code or row['status'], job=job, message=( + 'The generation was cancelled.' if row and row['status'] == 'cancelled' + else 'The reference picture did not finish.')) + continue + folder = 'freevideo/agent/%s/ref-%s' % (identity, job) + file = folder + '-0.png' + temporary = self.jobs.port.input_path(folder + '-0.landing-' + uuid.uuid4().hex[:8] + '.png', create=True) + for left in temporary.parent.glob('ref-%s-*.landing-*.png' % job): + if left != temporary: + left.unlink(missing_ok=True) + try: + shutil.copyfile(self.jobs.port.input_path(row['result']['file']), temporary) + except Exception: # noqa: BLE001 - the other views still land + temporary.unlink(missing_ok=True) + failure = failure or dict(code='invalid_media', job=job, + message='The reference picture could not be read.') + continue + staged.append((temporary, file)) + part = hidden['job_views'][job] + pictures.append(dict(media=dict(file=file), by=generating['by'], from_job=job, + generated=True, label='other' if hidden['kind'] == 'setting' else part)) + self._settle(identity, hidden, pictures, failure, staged) + finally: + for out, _ in staged: + out.unlink(missing_ok=True) + def _land(self, identity, generating, hidden, rows): from . import media pictures, failure, staged = [], None, [] @@ -733,6 +824,8 @@ def _settle(self, identity, hidden, pictures, failure, staged=()): generate_error=failure, revision=card['revision'] + 1, updated_at=self._now(), by=by, generated=dict(job=generating['job'], face_job=generating.get('face_job'), seed=hidden['seed'], prompt_version=PROMPT_VERSION)) + if hidden.get('source') == 'image': + card['generated']['source'] = 'image' card.pop('_generating', None) self._save(record) @@ -772,4 +865,4 @@ def _unreadable(error): def _generation_jobs(generating): if not generating: return [] - return [generating['job']] + ([generating['face_job']] if generating.get('face_job') else []) + return generating.get('jobs') or [generating['job']] + ([generating['face_job']] if generating.get('face_job') else []) diff --git a/freevideo_engine/agent/cli.py b/freevideo_engine/agent/cli.py index ca0eadb5..3d91a67e 100644 --- a/freevideo_engine/agent/cli.py +++ b/freevideo_engine/agent/cli.py @@ -76,6 +76,9 @@ def build_parser(): schema.cli_arguments(command, schema.contracts()['GET', '/jobs']['query']) command = commands.add_parser('upload', help='Copy a local media file into FreeVideo inputs') command.add_argument('path') + command = commands.add_parser('generate-image', help='Generate or change a picture with the local image model') + schema.cli_arguments(command, schema.image(), exclude=('kind',)) + command.add_argument('--wait', action='store_true', help='Wait for the picture and print the result') command = commands.add_parser('drama', help='Plan, start, review and finish a durable film') drama = command.add_subparsers(dest='operation', required=True) for name, (method, path, summary) in schema.DRAMA_OPERATIONS.items(): @@ -196,6 +199,12 @@ def source(value): if args.wait: jobs = [client.wait(job['id'], on_progress=progress) for job in jobs] result = jobs[0] if len(jobs) == 1 and args.command == 'submit' else {'jobs': jobs} + elif args.command == 'generate-image': + value = dict(kind='image', **{k: values[k] for k in schema.image()['properties'] if k in values}) + schema.validate_job(value) + result = client.submit(value) + if args.wait: + result = client.wait(result['id'], on_progress=progress) elif args.command == 'job': result = client.job(args.id, **({'wait': values['wait']} if 'wait' in values else {})) elif args.command == 'wait': diff --git a/freevideo_engine/agent/comfy.py b/freevideo_engine/agent/comfy.py index 3fafea3e..179bf9d2 100644 --- a/freevideo_engine/agent/comfy.py +++ b/freevideo_engine/agent/comfy.py @@ -1,4 +1,4 @@ -"""The running ComfyUI server as seen by agent jobs: its queue, history, folders and the local VLM.""" +"""The running ComfyUI server as seen by agent jobs: queue, history, folders and local services.""" import json import os from pathlib import Path, PurePosixPath @@ -304,6 +304,46 @@ def vlm_cancel(self, key): except ValueError: pass + # Local image generation ---------------------------------------------------------------------------- + def _image(self): + from ..image_gen.service import get_service + return get_service(self.server) + + @property + def image(self): + service = self._image() + if service is None: + raise JobError('image_unavailable', 'The image service is not loaded') + return service + + def image_busy(self): + service = self._image() + return service is not None and service.busy() + + def image_ready(self): + try: + service = self._image() + return service is not None and bool(service.info()['ready']) + except Exception: # noqa: BLE001 - an optional service must not prevent video generation + return False + + def image_start(self, action, value): + try: + return self.image.start(action, value)['job'] + except ValueError as error: + from .jobs import IMAGE_ERRORS, IMAGE_WAIT_ERRORS + code = str(error) if str(error) in IMAGE_ERRORS else 'invalid_request' + raise JobError(code, IMAGE_ERRORS[code], retryable=code in IMAGE_WAIT_ERRORS) from None + + def image_status(self, key): + return self.image.status(key) + + def image_cancel(self, key): + try: + self.image.cancel(key) + except ValueError: + pass + # Media tools --------------------------------------------------------------------------------------- def import_file(self, source, target): from ..comfy_assets import media_kind diff --git a/freevideo_engine/agent/drama.py b/freevideo_engine/agent/drama.py index e5631d76..8d1e1154 100644 --- a/freevideo_engine/agent/drama.py +++ b/freevideo_engine/agent/drama.py @@ -20,6 +20,7 @@ import time from . import spec as specs +from ..chain_plan import MIN_PART_SECONDS EXAMPLE = { 'title': '末班车', 'name': 'the-last-train', 'aspect': '9:16', 'quality': 'light', @@ -33,7 +34,7 @@ 'looking up at the departure board.', 'camera': 'The camera tracks her with small amplitude, then ' 'holds still.', 'lines': [{'who': 'lin', 'line': '还来得及吗？', 'how': 'says breathlessly'}], 'sound': 'Heavy rain on the platform roof, splashing footsteps, the buzz of the lamp.'}, - {'id': 'arrival', 'seconds': 6, 'continue': True, 'framing': 'vertical medium close-up', + {'id': 'arrival', 'seconds': 10, 'continue': True, 'framing': 'vertical medium close-up', 'action': 'A train headlight sweeps across Lin\'s face; the wind of the arriving train lifts her wet hair and ' 'her worry turns into a relieved smile.', 'camera': 'The camera pushes in slowly toward her face.', 'lines': [{'who': 'lin', 'line': '来了。', 'how': 'says quietly'}], @@ -406,10 +407,16 @@ def intent(play, shot): shot.get('camera'), lines, looks) if x and x.strip()) +def shot_seconds(play, shot): + """A shot says its length; else the film does; else 6 s, and 10 s for a shot that continues the one before + (a part that continues another video is at least 10 s).""" + return shot.get('seconds', play.get('seconds', MIN_PART_SECONDS if shot.get('continue') else 6)) + + def shot_job(play, shot, index, ref): job = dict(ref=ref, label='%s · %s' % (play.get('title') or 'drama', shot['id']), project=play.get('name'), aspect=play.get('aspect', '9:16'), quality=play.get('quality', 'light'), - seconds=shot.get('seconds', play.get('seconds', 6)), metadata={'drama_shot': shot['id']}) + seconds=shot_seconds(play, shot), metadata={'drama_shot': shot['id']}) if play.get('megapixels'): job['megapixels'] = play['megapixels'] if shot.get('seed') is not None: @@ -448,7 +455,7 @@ def batch(play): for i, shot in enumerate(play['shots']): if shot.get('idea') and not shot.get('action') and not shot.get('prompt'): rows.append(dict(ref='s%d-idea' % i, kind='enhance', text=shot['idea'], - seconds=shot.get('seconds', play.get('seconds', 6)), project=play.get('name'), + seconds=shot_seconds(play, shot), project=play.get('name'), label='%s · %s idea' % (play.get('title') or 'drama', shot['id']))) rows.append(shot_job(play, shot, i, 's%d' % i)) rows.append(review_job(play, [(shot, '@s%d' % i) for i, shot in enumerate(play['shots'])])) diff --git a/freevideo_engine/agent/errors.py b/freevideo_engine/agent/errors.py index 0ca9a443..a0b44920 100644 --- a/freevideo_engine/agent/errors.py +++ b/freevideo_engine/agent/errors.py @@ -117,10 +117,37 @@ # Recovery is a function of the committed task and phase, not just an error code. +IMAGE_ERRORS = { + 'invalid_request': 'The picture request contains invalid fields.', + 'invalid_prompt': 'The picture prompt must contain text and be at most 2000 characters.', + 'invalid_dimensions': 'Picture width and height must be integers from 256 to 2048.', + 'invalid_seed': 'The picture seed must be an integer from -9223372036854775808 to 18446744073709551615.', + 'invalid_media': 'A reference picture must be a regular file inside the input directory without symlinks.', + 'too_many_images': 'Use at most four reference pictures.', + 'busy': 'Another picture is being generated; this picture is waiting for it.', + 'busy_video': 'A video is being generated; this picture is waiting for it.', + 'pack_missing': 'The image model is not installed on this computer; tell the user to install it in FreeVideo first and do not retry.', + 'unknown_job': 'The image service no longer has this picture job.', + 'closed': 'The image service has shut down.', + 'image_unavailable': 'The image service is not loaded on this computer.', + 'no_memory': 'There is not enough GPU or system memory for this picture.', + 'generation_failed': 'The image model could not finish this picture.', + 'worker_exit': 'The image worker stopped before finishing this picture.', + 'cancelled': 'The picture job was cancelled.', +} +for _code, _message in IMAGE_ERRORS.items(): + if _code not in ERRORS: + ERRORS[_code] = (503 if _code in ('pack_missing', 'closed', 'image_unavailable') else 422, + _code in ('busy_video', 'unknown_job', 'closed', 'image_unavailable', 'worker_exit'), + _message, 'Read the picture job before submitting a corrected request with a new ref.') +ERRORS['pack_missing'] = (503, False, IMAGE_ERRORS['pack_missing'], + 'Tell the user to install the image model in FreeVideo first; do not retry.') + RECOVERY = { 'generate': 'POST /freevideo/agent/v1/jobs/{id}/resume to retry unfinished generation blocks.', 'vlm': 'Submit the question again with a new ref.', 'enhance': 'Submit the prompt enhancement again with a new ref.', + 'image': 'Submit the picture request again with a new ref.', 'assemble': 'Submit the clip assembly again with a new ref.', 'refine': 'Submit the second-pass generation again with a new ref.', 'submission_unknown': 'Read the job list, then retry the original submission with the same ref.', @@ -138,6 +165,10 @@ def recovery(code, *, kind=None, stage='request', resumable=True): + if kind == 'image' and code == 'pack_missing': + return ERRORS[code][3] + if kind == 'image' and code == 'no_memory': + return 'Lower the picture dimensions or use fewer reference pictures, then submit with a new ref.' if kind in RECOVERY and stage == 'terminal': return REPAIR.get(code, '') + RECOVERY['refine' if kind == 'generate' and not resumable else kind] if stage == 'submission_unknown' and code in ('internal_error', 'storage_unavailable', 'unreachable'): @@ -177,6 +208,8 @@ def build(code, *, details=None, request=None, validation=False, message=None, n code = code if code in ERRORS else 'internal_error' _, retryable, default_message, _ = ERRORS[code] action = recovery(code, kind=kind, stage=stage, resumable=resumable) + if kind == 'image': + default_message = IMAGE_ERRORS.get(code, default_message) result = dict(code=code, message=message if validation and message else default_message, details=(details or []) if validation else safe_details(details or []), retryable=retryable, request_id=request or request_id(), diff --git a/freevideo_engine/agent/guide.py b/freevideo_engine/agent/guide.py index 3ac4ba2a..29633761 100644 --- a/freevideo_engine/agent/guide.py +++ b/freevideo_engine/agent/guide.py @@ -73,7 +73,7 @@ 'timing': [ 'Single-block frames are 24 fps on a 17n+5 grid (6 s -> 6.583 s, 10 s -> 10.125 s). ' 'Under 30 s defaults to one generation when the device supports this size and quality; unknown limits use 15 s. ' - 'Otherwise it uses parts of up to 15 s (30 s is 15 + 15), continued and joined automatically. ' + 'Otherwise it uses as few parts as possible, each at least 10 s; how many depends on how long this computer can generate in one run, and the parts continue and join automatically. ' 'Multiple blocks or continue_from need two-pass text generation: quality "light" ' 'or sampling.two_pass=true, no preview and no media. ' 'Multi-block output is cut to the planned duration. POST /plan shows the chain plan. ' @@ -189,6 +189,10 @@ 'otherwise make its reference pictures with POST /cast/{id}/generate. ' 'Before shooting, tell the user the casting in one sentence. Later shots use the card\'s current pictures. ' 'Keep cards you made so later films can reuse them.', + 'Pictures. When the user wants a picture, or a first frame or reference picture that does not exist yet, make ' + 'it with freevideo_generate_image; to change a picture, pass it in images. The result is an input file: use ' + 'result.file as a first frame, as a reference, or in images again. When where people stand in the shot matters, ' + 'make the first frame with people (cast pictures and places) and generate the video from it.', 'Cut by default. Wherever the film can change shot, make the new shot from the reference pictures it should ' 'have (cast cards, scene pictures) instead of continuing the previous video. Continue a video only for a ' 'long take that cannot be cut, an action that must stay unbroken, or sound that must stay unbroken, and ' @@ -218,6 +222,9 @@ '先选角，再拍摄。先搜索角色库，使用用户指定的卡片或明确符合要求的卡片；拿不准时询问用户。' '只有找不到合适卡片时才新建：用户提供了照片就用照片，否则通过 POST /cast/{id}/generate 生成参考图。' '开拍前用一句话告诉用户选角结果。后续镜头使用卡片当前的图片。保留你创建的卡片，方便以后拍片复用。', + '图片。用户要一张图，或者缺一张首帧、参考图时，用 freevideo_generate_image 生成；要改一张图，把它放进 images。' + '结果是一个输入文件：result.file 可以当首帧、当参考图，也可以再放进 images。镜头里人物站位要紧时，' + '先用 people（角色图和各自的位置）做出首帧，再用它生成视频。', '默认切镜头。能切镜头的地方就切：新镜头直接用对应的参考图（人物卡、场景图）生成，不接上一段视频。' '只有不能切镜的长镜头、必须连贯的动作、必须连贯的声音，才接着上一段生成，而且不建议一段接一段地连续接。' '一条视频最多 5 段（第一段加 4 次接着生成），接着生成的每一段至少 10 秒，总长约 70 秒；不建议接满，超过就切镜头。', @@ -523,7 +530,8 @@ def connection_text(value): 'continue the previous one carry no references (cast_reference inherited); POST /drama/plan shows each shot\'s cast. ' 'A shot with a carded character is at most 15 s; continue it in the next shot to go longer.', 'POST /cast/{id}/generate makes reference pictures from the description as a normal job (face=true adds a close-up). ' - 'When the job ends the card gets the sharpest frames (labels guessed, label_guess) and a new revision; a failure ' + 'It defaults to the installed image model, otherwise video frames (labels guessed, label_guess). ' + 'When generation ends the card gets the pictures and a new revision; a failure ' 'sets generate_error and keeps the pictures. If the user changed the pictures (meanwhile, or before an agent asked), ' 'they go to candidates.', ] diff --git a/freevideo_engine/agent/jobs.py b/freevideo_engine/agent/jobs.py index 0d7cfce8..8d29db26 100644 --- a/freevideo_engine/agent/jobs.py +++ b/freevideo_engine/agent/jobs.py @@ -1,8 +1,8 @@ """Agent jobs: one durable record per request, dependencies, GPU turn-taking and status. Generation uses the studio's durable chain runner. Jobs keeps Agent chains in -FIFO order until the runner is free. The local VLM (prompt enhancement and questions) only -starts when that queue is empty, and generation waits while the VLM runs. Jobs +FIFO order until the runner is free. Local image and VLM jobs (prompt enhancement and questions) only +start when that queue is empty, and video generation waits while either runs. Jobs may read other jobs' results (a frame, a clip, an enhanced prompt); such a job waits until those finish and fails with dependency_failed if one does not. """ @@ -15,6 +15,7 @@ import math import os import secrets +import shutil import stat from pathlib import Path import threading @@ -30,7 +31,40 @@ ACTIVE = ('waiting', 'queued', 'running', 'cancelling') CAST_AGAIN = 'POST /freevideo/agent/v1/cast/%s/generate to make the reference pictures again.' DONE = ('succeeded', 'failed', 'cancelled', 'interrupted') -GPU_KINDS = ('vlm', 'enhance') +GPU_KINDS = ('vlm', 'enhance', 'image') +IMAGE_ERRORS = errors.IMAGE_ERRORS +IMAGE_WAIT_ERRORS = ('busy', 'busy_video') +IMAGE_WAIT_NOTICES = { + 'video': 'A video is being generated; this picture starts when it finishes. Tell the user that once, then wait.', + 'image': 'Another picture is being generated; this picture starts when it finishes. Tell the user that once, then wait.', + 'gpu': 'The GPU is in use; this picture starts when it is free. Tell the user that once, then wait.', +} + + +def image_reference(port, relative): + """The name the image service reads an input picture by. + + It takes files in the input folder itself and its own results as they are, and checks them itself; + a picture anywhere else under input is copied next to them once, named by its contents.""" + folder, _, leaf = relative.rpartition('/') + if folder in ('', 'freevideo-image'): + return relative + try: + source = port.input_path(relative) + with open(source, 'rb') as stream: + digest = hashlib.file_digest(stream, 'sha256').hexdigest() + name = 'agent-image-' + digest[:16] + source.suffix.lower() + target = port.input_root / name + if not target.exists(): + part = target.with_name('%s.%s.part' % (name, uuid.uuid4().hex[:8])) + try: + shutil.copyfile(source, part) + os.replace(part, target) + finally: + part.unlink(missing_ok=True) + return name + except OSError: + raise JobError('invalid_media', IMAGE_ERRORS['invalid_media']) from None class JobError(Exception): @@ -107,6 +141,40 @@ def cut_problems(spec, seconds, prefix='', *, editor=None): return [dict(problem, got=got, message=problem['message'].replace('received None', 'received ' + got))] +def _continued_part_problem(seconds, chain): + """True when a video that continues another has a part the backend refuses, or is under 10 s. + + The backend compares each part's node input frames (overlap included); a part that + continues is also at least MIN_PART_SECONDS long. + """ + from ..chain_plan import MIN_PART_FRAMES, MIN_PART_SECONDS + return seconds < MIN_PART_SECONDS or any(block['frames'] < MIN_PART_FRAMES for block in chain['blocks']) + + +def length_problem(seconds, continued, prefix='', *, minimum=None): + """Suggest nearby whole-second lengths the same planner and minimum accept.""" + from ..chain_plan import PlanError, plan + available = [] + for candidate in range(max(1, int(seconds) - 30), int(seconds) + 31): + try: + chain = plan(candidate, continued=continued) + except PlanError: + continue + if not continued or not _continued_part_problem(candidate, chain): + available.append(candidate) + if minimum is not None and seconds >= minimum: + minimum = None # long enough on its own: the length only cannot be cut into parts that long + nearby = [s for s in available if s < seconds][-2:] + [s for s in available if s > seconds][:2] + rule = ('at least %d s per continued part' % minimum if minimum is not None + else 'a length that can be split into supported parts') + allowed = rule + '; available lengths: ' + ', '.join('%d s' % s for s in nearby) + problem = schema.diagnostic(prefix + '/seconds', allowed, seconds, nearest=nearby) + if minimum is not None: + problem['message'] = ('Each continued part must be at least %d seconds. Make it longer. ' % minimum + + problem['message']) + return problem + + def parts_problems(spec, runner, prefix='', *, editor=None, work=None, planned=None): """One video continues through a limited number of parts; say so before anything is generated. @@ -114,13 +182,24 @@ def parts_problems(spec, runner, prefix='', *, editor=None, work=None, planned=N `work` is the saved video continue_from names (output-relative), or None while it cannot be told. `planned` is the number of parts an unfinished job it continues will have, when that is known instead. """ - if runner is None or spec['kind'] != 'generate' or spec.get('refine_from'): + if spec['kind'] != 'generate' or spec.get('refine_from'): return [] - from ..chain_plan import plan + from ..chain_plan import MIN_PART_SECONDS, PlanError, plan from ..chain_runner import ChainTooLong from ..chain_store import reference - limit, continued = runner.max_parts, bool(spec.get('continue_from')) - count = lambda seconds: len(plan(seconds, max_frames=single_frame_limit(spec), continued=continued)['blocks']) + continued, frames = bool(spec.get('continue_from')), single_frame_limit(spec) + if editor is None: + # Planning errors must not be swallowed by the source-depth fallback below. + try: + chain = plan(spec['seconds'], max_frames=frames, continued=continued) + except PlanError: + return [length_problem(spec['seconds'], continued, prefix)] + if continued and _continued_part_problem(spec['seconds'], chain): + return [length_problem(spec['seconds'], continued, prefix, minimum=MIN_PART_SECONDS)] + if runner is None: + return [] + limit = runner.max_parts + count = lambda seconds: len(plan(seconds, max_frames=frames, continued=continued)['blocks']) try: with runner.lock: if editor is not None: @@ -130,7 +209,7 @@ def parts_problems(spec, runner, prefix='', *, editor=None, work=None, planned=N else: # A source the backend has no record of counts as one part, as it does there. earlier = (runner._work_parts(work, ()) if work else planned or 1) if continued else 0 - needed = earlier + count(spec['seconds']) + needed = earlier + len(chain['blocks']) except (ValueError, KeyError, TypeError): return [] # Not a chain the backend accepts; it says why when the chain is created. if needed <= limit: @@ -138,7 +217,14 @@ def parts_problems(spec, runner, prefix='', *, editor=None, work=None, planned=N if editor is not None: problem, got = schema.diagnostic(prefix + '/target', 'a range that takes at most %d parts' % limit, None), '%d parts' % needed else: - fit = next((s for s in range(int(spec['seconds']), 1, -1) if earlier + count(s) <= limit), None) + # A length the planner refuses does not fit; a part that continues a video is at least MIN_PART_SECONDS. + def parts(seconds): + try: + return count(seconds) + except ValueError: + return limit + 1 + shortest = MIN_PART_SECONDS if continued else 2 + fit = next((s for s in range(int(spec['seconds']), shortest - 1, -1) if earlier + parts(s) <= limit), None) if fit is None: problem = schema.diagnostic(prefix + '/continue_from', 'a video of fewer than %d parts' % limit, None) got = 'a video of %d parts' % earlier @@ -475,9 +561,9 @@ def touch(self, row, **changes): def finish(self, row, status, *, result=None, error=None, save=True): if error is not None: code = error.get('code', 'internal_error') - if code == 'invalid_json' and row['kind'] in GPU_KINDS: + if code == 'invalid_json' and row['kind'] in ('vlm', 'enhance'): code = 'vlm_invalid_json' - if code != error.get('code') or not error.get('request_id'): + if code != error.get('code') or not error.get('request_id') or row['kind'] == 'image' and code in IMAGE_ERRORS: error = errors.build(code, details=error.get('details'), request=row.get('request_id'), action_required=code != 'cancelled', kind=row['kind'], stage='terminal', resumable=not row['spec'].get('refine_from')) @@ -688,7 +774,9 @@ def cancel(self, identity): state = self.port.cancel(row['comfy']['prompt_id'], row['comfy']['node']) if state == 'cancelled': self.finish(row, 'cancelled', error=dict(code='cancelled', message='Removed from the queue')) - elif row['kind'] in GPU_KINDS and row.get('_vlm'): + elif row['kind'] == 'image' and row.get('_image'): + self.port.image_cancel(row['_image']) + elif row['kind'] in ('vlm', 'enhance') and row.get('_vlm'): self.port.vlm_cancel(row['_vlm']) return row @@ -903,7 +991,10 @@ def view(self, row, *, detail=True): result.update(cancel_requested=row.get('cancel_requested', False), attempt=row.get('attempt', 1)) if row['status'] == 'waiting': pending = [d for d in row.get('depends_on', ()) if self.jobs.get(d, {}).get('status') not in DONE] - result['waiting_for'] = pending or (['gpu'] if row['kind'] in GPU_KINDS + ('generate',) else []) + result['waiting_for'] = pending or row.get('waiting_for') or (['gpu'] if row['kind'] in GPU_KINDS + ('generate',) else []) + if row['kind'] == 'image' and not pending and row.get('waiting_for'): + # Only once it was really held back; a picture just submitted has not been looked at yet. + result['notice'] = IMAGE_WAIT_NOTICES[row['waiting_for'][0]] if row['status'] in ('queued', 'running', 'cancelling'): result['progress'] = progress if row.get('chain_id'): @@ -1042,8 +1133,8 @@ def advance(): def _tick(self): """Advance every unfinished job once, oldest first; safe to call often. - The GPU takes turns: a VLM job starts only with no generation ahead of it - or in ComfyUI's queue, and a generation never overtakes an earlier VLM job. + The GPU takes turns: local image and VLM jobs start only with no video ahead + of them or in ComfyUI's queue, and a video never overtakes an earlier GPU job. """ dispatches = [] generation_ahead = vlm_ahead = False @@ -1056,7 +1147,9 @@ def _tick(self): self.poll_chain(row) elif kind == 'generate' and row.get('comfy'): self.poll_generation(row) - elif kind in GPU_KINDS and row.get('_vlm'): + elif kind == 'image' and row.get('_image'): + self.poll_image(row) + elif kind in ('vlm', 'enhance') and row.get('_vlm'): self.poll_vlm(row) if row['status'] in ACTIVE: generation_ahead |= kind == 'generate' @@ -1067,7 +1160,7 @@ def _tick(self): if self.ready(row) != 'ready': continue if kind == 'generate': - if generation_ahead or vlm_ahead or self.port.vlm_busy(): + if generation_ahead or vlm_ahead or self.port.vlm_busy() or self.port.image_busy(): generation_ahead = True continue generation_ahead = True @@ -1075,9 +1168,16 @@ def _tick(self): continue elif kind in GPU_KINDS: chain_busy = self.runner_busy() - busy = generation_ahead or vlm_ahead or chain_busy or self.port.queue_busy() or self.port.vlm_busy() + video_busy = generation_ahead or chain_busy or self.port.queue_busy() + image_busy = self.port.image_busy() + busy = video_busy or vlm_ahead or image_busy or self.port.vlm_busy() vlm_ahead = True if busy: + if kind == 'image': + code = 'busy_video' if video_busy else 'busy' + if not self.wait_image(row, code, + waiting_for='video' if video_busy else 'image' if image_busy else 'gpu'): + self.finish(row, 'failed', error=dict(code=code)) continue row['_dispatching'] = True dispatches.append(row) @@ -1142,13 +1242,16 @@ def run(): await asyncio.to_thread(check) if row['kind'] == 'generate': await self.start_generation(row) - elif row['kind'] in GPU_KINDS: + elif row['kind'] == 'image': + await asyncio.to_thread(self.start_image, row) + elif row['kind'] in ('vlm', 'enhance'): await asyncio.to_thread(self.start_vlm, row) else: result = await asyncio.to_thread(run) succeeded = True except JobError as error: - failure = ('cancelled' if error.code == 'cancelled' else 'failed', error.view()) + if row['kind'] != 'image' or not self.wait_image(row, error.code): + failure = ('cancelled' if error.code == 'cancelled' else 'failed', error.view()) except specs.SpecError as error: failure = ('failed', dict(code='invalid_request', message=str(error), details=error.errors)) except Exception as error: # noqa: BLE001 - a job must always end with a reason @@ -1316,7 +1419,7 @@ def start_chain(self, row, *, payload=None, resume=False): if row.get('_cancel') or row['status'] == 'cancelled': return if (any(r['state'] in ('queued', 'running', 'cancelling') or r['output_state'] == 'joining' - for r in runner.records.values()) or self.port.queue_busy() or self.port.vlm_busy()): + for r in runner.records.values()) or self.port.queue_busy() or self.port.vlm_busy() or self.port.image_busy()): return # Submit and plan refuse this up front; the backend can still change before a queued or resumed job starts. if (any(b['state'] != 'done' and (b['after'] or b['before']) @@ -1325,7 +1428,8 @@ def start_chain(self, row, *, payload=None, resume=False): else bool(row['spec'].get('continue_from')) or not single_block(row['spec'])) and not self.continuation_ready(): raise continuation_unavailable() self.touch(row, chain_submission='creating', status='queued', queued_at=self.clock()) - from ..chain_runner import ChainBusy, ChainTooLong + from ..chain_plan import PlanError + from ..chain_runner import ChainBusy, ChainPartTooShort, ChainTooLong try: if resume: chain = runner.resume(row['chain_id'], client_id='agent:' + row['id']) @@ -1341,6 +1445,12 @@ def start_chain(self, row, *, payload=None, resume=False): row['resume_pending'] = False self.finish(row, 'failed', error=errors.build('chain_too_long', validation=True, message=str(error))) return + except (PlanError, ChainPartTooShort) as error: + row['resume_pending'] = False + problem = length_problem(row['spec']['seconds'], bool(row['spec'].get('continue_from')), + minimum=error.limit if isinstance(error, ChainPartTooShort) else None) + self.finish(row, 'failed', error=errors.validation([problem])) + return except Exception: row['resume_pending'] = False if resume: @@ -1542,6 +1652,70 @@ def generation_failure(self, row, chain): return chain_error(code, errors.ERRORS[code][2], next_action, details=[detail], retryable=code != 'insufficient_memory').view() + # Local image generation ---------------------------------------------------------------------------- + def wait_image(self, row, code, *, waiting_for=None): + if code not in IMAGE_WAIT_ERRORS: + return False + if row.get('_cancel'): + self.finish(row, 'cancelled', error=dict(code='cancelled')) + return True + waiting_for = [waiting_for or ('video' if code == 'busy_video' else 'image')] + row.pop('_image', None) + if row['status'] != 'waiting' or row.get('waiting_for') != waiting_for: + self.touch(row, status='waiting', started_at=None, waiting_for=waiting_for) + return True + + def start_image(self, row): + spec = row['spec'] + row['_phase'] = 'preparing images' + value = {k: spec[k] for k in ('width', 'height', 'seed') if k in spec} + if spec.get('people'): + # A first frame: the people's pictures at the asked places, in the scene shown or described. + action = 'compose' + value['people'] = [dict(person, image=image_reference(self.port, person['image'])) for person in spec['people']] + if spec.get('scene_image'): + value['scene_image'] = image_reference(self.port, spec['scene_image']) + else: + value['scene'] = spec['prompt'] + else: + action = 'generate' + value.update(prompt=spec['prompt'], images=[image_reference(self.port, name) for name in spec['images']]) + if row.get('_cancel'): + raise JobError('cancelled', IMAGE_ERRORS['cancelled']) + row['_image'] = self.port.image_start(action, value) + row['_phase'] = 'starting' + row.pop('waiting_for', None) + self.touch(row, status='running', started_at=self.clock()) + if row.get('_cancel'): + self.port.image_cancel(row['_image']) + + def poll_image(self, row): + try: + state = self.port.image_status(row['_image']) + except ValueError: + self.finish(row, 'failed', error=dict(code='unknown_job')) + return + phase = state.get('phase') + if phase in ('starting', 'progress', 'cancelling'): + row['_phase'] = state.get('stage') if phase == 'progress' else phase + elif phase == 'complete': + from PIL import Image + relative = state['file'] + with Image.open(self.port.input_path(relative)) as picture: + width, height = picture.size + result = dict(assets.view(relative, 'input'), width=width, height=height) + result['thumbnail'] = result['url'] + seed = state.get('seed', row['spec'].get('seed')) + if seed is not None: + result['seed'] = seed + self.finish(row, 'succeeded', result=result) + elif phase == 'cancelled': + self.finish(row, 'cancelled', error=dict(code='cancelled')) + else: + code = state.get('error') or 'generation_failed' + if not self.wait_image(row, code): + self.finish(row, 'failed', error=dict(code=code)) + # Local VLM ----------------------------------------------------------------------------------------- def start_vlm(self, row): spec = row['spec'] diff --git a/freevideo_engine/agent/mcp.py b/freevideo_engine/agent/mcp.py index 2b683e7a..cbe9969d 100644 --- a/freevideo_engine/agent/mcp.py +++ b/freevideo_engine/agent/mcp.py @@ -44,8 +44,9 @@ def tool_list(): dict(name='freevideo_prompt_guide', title='H3 prompt guide', description='How to write prompts: sections, shots and cuts, camera motion, dialogue, references, short-drama tips, and a full example.', annotations=read), dict(name='freevideo_schema', title='Job JSON Schemas', description='JSON Schemas of every job kind and of image edit recipes.', annotations=read), dict(name='freevideo_plan', title='Plan a job', description="Validate a job without running it: actual frames/seconds, sampling passes, warnings and a time estimate from this computer's history. Same fields as freevideo_generate.", annotations=read), - dict(name='freevideo_generate', title='Generate a video', description='Generate one video from an H3 prompt, up to 60 seconds. Under 30 s defaults to one generation if the device supports this size and quality; unknown limits use 15 s. Otherwise parts continue and join automatically (30 s is 15 + 15). Prefer a new shot with its reference pictures wherever the film can cut; continue only a take, action or sound that must stay unbroken. Use this for any single video, also after freevideo_enhance_prompt; the drama tools are only for a film of several separately chosen shots. Optionally from first/last frames or references (files, local paths, or frames of earlier jobs). Returns the job; with wait, its result (video path and URL).', annotations=write), - dict(name='freevideo_submit_batch', title='Submit a storyboard', description='Submit several jobs at once. Each row is a job object with "kind" (generate default, assemble, vlm, enhance) and an optional "ref"; later rows use "@ref" wherever a job id is expected (from_job, after, clips, refine_from). FreeVideo runs them in dependency order.', annotations=write), + dict(name='freevideo_generate', title='Generate a video', description='Generate one video from an H3 prompt, up to 60 seconds. Under 30 s defaults to one generation if the device supports this size and quality; unknown limits use 15 s. Otherwise the video is cut into as few parts as possible, each at least 10 s, and how many parts depends on how long this computer can generate in one run; the parts continue and join automatically. Prefer a new shot with its reference pictures wherever the film can cut; continue only a take, action or sound that must stay unbroken. Use this for any single video, also after freevideo_enhance_prompt; the drama tools are only for a film of several separately chosen shots. Optionally from first/last frames or references (files, local paths, or frames of earlier jobs). Returns the job; with wait, its result (video path and URL).', annotations=write), + dict(name='freevideo_generate_image', title='Generate a picture', description='Generate a picture from a prompt; with images, change or combine input pictures; with people, make a first frame with cast members standing where you say. The result is an input file to use as a first frame, as a reference, or in images again.', annotations=write), + dict(name='freevideo_submit_batch', title='Submit a storyboard', description='Submit several jobs at once. Each row is a job object with "kind" (generate default, assemble, vlm, enhance, image) and an optional "ref"; later rows use "@ref" wherever a job id is expected (from_job, after, clips, refine_from). FreeVideo runs them in dependency order.', annotations=write), dict(name='freevideo_job', title='Job status', description='Status, progress, result or error of a job.', annotations=read), dict(name='freevideo_wait', title='Wait for a job or the user', description="Wait at most 25 seconds for a job (with id) or, without id, at most 50 seconds for what the user says in FreeVideo (inbox). again=true means nothing new yet: call it again. A job wait returns early when the user says something. If the job is still active, call it again with the same ID; do not resubmit.", annotations=read), dict(name='freevideo_jobs', title='List jobs', description='Recent agent jobs, optionally by status (comma-separated), project or kind.', annotations=read), @@ -271,9 +272,9 @@ def tool(self, identity, name, args, token): return c.schema(), [] if name == 'freevideo_plan': return c.plan(dict(args, kind='generate')), [] - if name in ('freevideo_generate', 'freevideo_assemble', 'freevideo_ask', 'freevideo_enhance_prompt'): - kind = dict(freevideo_generate='generate', freevideo_assemble='assemble', freevideo_ask='vlm', - freevideo_enhance_prompt='enhance')[name] + from .schema import TOOL_KINDS + if name in TOOL_KINDS: + kind = TOOL_KINDS[name] takes = args.pop('takes', 1) response = c.post('/jobs', dict(args, kind=kind, **({'takes': takes} if takes != 1 else {}))) jobs = response['jobs'] diff --git a/freevideo_engine/agent/productions.py b/freevideo_engine/agent/productions.py index 809b2213..d23838c3 100644 --- a/freevideo_engine/agent/productions.py +++ b/freevideo_engine/agent/productions.py @@ -22,7 +22,7 @@ import uuid from . import drama, production, schema, spec, errors, operation_receipts -from .jobs import DONE, JobError, iso +from .jobs import DONE, JobError, iso, parts_problems from .forecast import generation_estimate MAX_GENERATION = 4 # Generations are zero-based: original 0, reshoots 1 through 4. @@ -306,6 +306,16 @@ def _material(self, state, *, local=True): problems.append(dict(row, path=new, field=prefix + row.get('field', old), message=row['message'].replace(old, new, 1), next=row['next'].replace(old, new, 1))) continue + if continued or shot['continues']: + # The same length rule as at submission, with this computer's single-run limit rather than the + # 15 s default: a shot that cannot be cut into parts of at least 10 s is refused here instead of + # starting a film in which that shot never gets a job. A shot not yet submitted carries no parent. + try: + pinned = spec.normalize(dict(item, continue_from=dict(from_job=placeholder)), port=self.jobs.port) + except spec.SpecError: + pinned = None # Its other problems are reported where they are found. + if pinned is not None: + problems.extend(parts_problems(pinned, None, prefix)) for dependency in normalized['depends_on']: if not (continued and dependency == placeholder) and dependency not in self.jobs.jobs: problems.append(schema.diagnostic(prefix + '/first', 'a job that exists on this server', dependency)) diff --git a/freevideo_engine/agent/routes.py b/freevideo_engine/agent/routes.py index 96400e7e..7eabc124 100644 --- a/freevideo_engine/agent/routes.py +++ b/freevideo_engine/agent/routes.py @@ -203,6 +203,7 @@ def capabilities(port, jobs): kinds={'generate': 'Video from a prompt and optional keyframes or references', 'assemble': 'Join finished clips into one video (cut or crossfade)', 'vlm': 'Ask the local vision-language model about images or sampled video frames', + 'image': 'Generate a picture from text or change input pictures with the local image model', 'enhance': 'Rewrite a rough idea into a full H3 prompt with the local VLM', 'export': 'Render a saved revision of an editing project to a video', 'listen': 'Listen to the sound of an editing project for subtitle lines and times', @@ -233,6 +234,8 @@ def plan(port, value, jobs=None): if jobs is not None: problems += parts_problems(spec, jobs.port.chain_runner(), work=jobs.continued_work(spec), planned=jobs.planned_parts(spec)) + else: + problems += parts_problems(spec, None) if problems: raise specs.SpecError(problems) if jobs is not None: @@ -724,6 +727,12 @@ def activity_item(row, *, card=False): width=size[0], height=size[1], thumbnail='/freevideo/library/thumbnail?id=' + quote(work, safe='') if work else None, play='/freevideo/agent/chat/video?id=' + quote(work, safe='') if work else None) + picture = view.get('result') or {} + if row['kind'] == 'image' and picture.get('file'): + # ComfyUI's own view of the input file, which the page can show without the agent's token. + folder, _, leaf = picture['file'].rpartition('/') + value.update(width=picture.get('width'), height=picture.get('height'), + thumbnail='/view?filename=%s&subfolder=%s&type=input' % (quote(leaf, safe=''), quote(folder, safe=''))) return value def agent_state(): diff --git a/freevideo_engine/agent/schema.py b/freevideo_engine/agent/schema.py index 2cdd74e8..5f38b8b3 100644 --- a/freevideo_engine/agent/schema.py +++ b/freevideo_engine/agent/schema.py @@ -15,6 +15,8 @@ QUALITY = {'light': (True, 8, 3), 'medium': (False, 12, 3), 'high': (False, 16, 3), 'max': (False, 20, 3)} ASPECTS = {'16:9': (16, 9), '9:16': (9, 16), '1:1': (1, 1), '4:3': (4, 3), '3:4': (3, 4), '21:9': (21, 9), '2:3': (2, 3), '3:2': (3, 2)} +IMAGE_ASPECTS = {'1:1': (1024, 1024), '16:9': (1344, 768), '9:16': (768, 1344), + '4:3': (1152, 896), '3:4': (896, 1152)} MAX_UPLOAD_BYTES = 2 * 1024 ** 3 MAX_JSON_UPLOAD_BYTES = 96 * 1024 ** 2 MAX_BODY_BYTES = 1024 * 1024 @@ -194,6 +196,34 @@ def vlm(*, local=False): ('kind', 'questions'), description='x-job-kind also checks the sum of images and sampled video frames across questions (at most 24).', **{'x-job-kind': 'vlm'}) +def image(*, local=False): + return obj(dict(kind={'const': 'image', 'type': 'string'}, + prompt=dict(type='string', minLength=1, maxLength=2000), + width=number(256, 2048, integer=True), height=number(256, 2048, integer=True), + aspect=enum(IMAGE_ASPECTS, description='Instead of width and height; 1:1 when all three are omitted.'), + seed=nullable(number(-2**63, 2**64 - 1, integer=True)), + images=array({'oneOf': [FILE, obj(dict(file=FILE), ('file',))]}, high=4, description='Pictures to change ' + 'or combine: input files, such as result.file of an earlier picture job. The prompt calls them image 1, image 2.'), + people=array(obj(dict(image=FILE, name=string(40), x=number(0, 1), height=number(.05, 1), feet=number(0, 1)), + ('image',)), 1, 3, description='Makes a first frame with people standing where you say. image is ' + "a full-body picture on a plain backdrop, such as a cast card's front picture. From 0 to 1: x is the person's " + 'centre across the width (default 0.5), height their size (0.62) and feet where they stand down the height ' + '(0.9). The prompt describes the scene, or scene_image shows it. Not together with images.'), + scene_image=FILE, + **common()), + ('kind', 'prompt'), allOf=[{'if': {'required': ['width']}, 'then': {'required': ['height']}}, + {'if': {'required': ['height']}, 'then': {'required': ['width']}}, + {'if': {'required': ['aspect']}, + 'then': {'not': {'anyOf': [{'required': ['width']}, {'required': ['height']}]}}}], + **{'x-job-kind': 'image'}) + + +def image_result(): + return obj(dict(file=FILE, url=string(), width=number(256, 2048, integer=True), + height=number(256, 2048, integer=True), seed=number(-2**63, 2**64 - 1, integer=True), + thumbnail=string()), ('file', 'url', 'width', 'height', 'thumbnail')) + + def enhance(*, local=False): return obj(dict(kind={'const': 'enhance', 'type': 'string'}, text=string(12000), seconds=number(1, 60, default=10), media=nullable(generation_media(local=local, enhance=True)), **common()), @@ -235,7 +265,8 @@ def image_edit(*, local=False): **{'x-job-kind': 'image_edit'}) -KINDS = dict(generate=generate, assemble=assemble, vlm=vlm, enhance=enhance, export=export, listen=listen, image_edit=image_edit) +KINDS = dict(generate=generate, assemble=assemble, vlm=vlm, enhance=enhance, image=image, + export=export, listen=listen, image_edit=image_edit) def job_schema(*, local=False, refs=False): @@ -444,7 +475,8 @@ def cast_schema(name, *, local=False): if name == 'delete': return obj(dict(revision=number(1, integer=True), ref=REF)) if name == 'generate': - return obj(dict(face=dict(type='boolean', default=False, description='Also make one face close-up.'), + return obj(dict(face=dict(type='boolean', default=False, description='Also make one face close-up (video; the image model always does).'), + source=enum(('image', 'video'), description='Defaults to image when the local image model is ready, otherwise video.'), replace=dict(type='boolean', default=False, description='Replace pictures the user did not add; default fills empty slots.'), revision=number(1, integer=True), ref=REF)) @@ -454,6 +486,24 @@ def cast_schema(name, *, local=False): return EMPTY +def cast_view(): + picture = obj(dict(media=obj(dict(file=FILE, url=string(2048)), ('file', 'url')), + by=enum(('user', 'agent', 'builtin')), from_job=JOB, generated={'type': 'boolean'}, + label=enum(LABELS), label_guess={'type': 'boolean'}), ('media', 'by', 'label')) + generating = obj(dict(job=JOB, face_job=nullable(JOB), jobs=array(JOB, 1, 5), + source=enum(('image', 'video')), replace={'type': 'boolean'}, at=string(), + by=enum(('user', 'agent', 'builtin'))), ('job', 'face_job', 'replace', 'at', 'by')) + generated = obj(dict(job=JOB, face_job=nullable(JOB), source=enum(('image', 'video')), + seed=number(0, 2**53 - 1, integer=True), prompt_version=number(1, integer=True)), + ('job', 'face_job', 'seed', 'prompt_version')) + return dict(obj(dict(id=CAST, kind=enum(('character', 'setting')), name=string(120), + description=nullable(string(1000)), revision=number(1, integer=True), + images=array(picture, 0, 6), candidates=array(picture, 0, 6), + generating=nullable(generating), generated=nullable(generated), + generate_error=nullable(obj(dict(code=string(), job=JOB, message=string()), ('code', 'job', 'message')))), + ('id', 'kind', 'name', 'revision', 'images', 'candidates', 'generating')), additionalProperties=True) + + # Image editing sessions: the image editor's version 2 document (web/image_editor_engine.js), # changed with the page's commands (web/image_commands.js). Shapes follow the page's # own fields (lineHeight, shape width), not the /media/edit recipe's. @@ -1017,6 +1067,11 @@ def bad(field, allowed, got=None): if flags['clip'] and item.get('end') is not None and item['end'] <= item.get('start', 0): bad('end', 'later than start', item['end']) kind = schema.get('x-job-kind') + if kind == 'image': + if value.get('people') and value.get('images'): + bad('images', 'omitted when people is supplied', value['images']) + if value.get('scene_image') and not value.get('people'): + bad('scene_image', 'omitted unless people is supplied', value['scene_image']) if kind == 'generate': if value.get('continue_from') and value.get('media'): bad('media', 'omitted when continue_from is supplied', value['media']) @@ -1159,12 +1214,12 @@ def schemas(): return dict(api_version=API_VERSION, job=dict(job_schema(), **{'$schema': S, '$defs': { k: f() for k, f in KINDS.items()}}), batch=batch_schema(), production=drama_schema('start'), drama={name: drama_schema(name) for name in DRAMA_OPERATIONS}, production_view=production_response(), - cast={name: cast_schema(name) for name in CAST_OPERATIONS}, + cast={name: cast_schema(name) for name in CAST_OPERATIONS}, cast_view=cast_view(), projects={name: project_schema(name) for name in PROJECT_OPERATIONS}, images={name: image_schema(name) for name in IMAGE_OPERATIONS}, sessions=dict(event=session_event(stored=True), meta=session_view(), **{name: session_schema(name) for name in SESSION_OPERATIONS}), - confirmations=confirm_settings()) + confirmations=confirm_settings(), image_result=image_result()) def openapi(): @@ -1208,6 +1263,12 @@ def openapi(): response_schema = production_response() elif suffix == '/confirmations': response_schema = obj(confirm_settings()['properties'], ('start', 'reshoot_seconds')) + elif suffix in ('/cast', '/freevideo/agent/cast') and method == 'POST' or ( + suffix in ('/cast/{id}', '/freevideo/agent/cast/{id}') and method in ('GET', 'PATCH')): + response_schema = {'$ref': '#/components/schemas/CastCard'} + elif suffix in ('/cast/{id}/generate', '/freevideo/agent/cast/{id}/generate'): + response_schema = obj(dict(job=JOB, jobs=array(JOB, 1, 5), + card={'$ref': '#/components/schemas/CastCard'}), ('job', 'jobs', 'card')) elif suffix.startswith('/sessions'): response_schema = (obj(dict(events=array(session_event(stored=True), 0, 500), last_seq=number(0, integer=True), stopped={'type': 'boolean'}), ('events', 'last_seq', 'stopped')) @@ -1237,12 +1298,16 @@ def openapi(): status=enum(('waiting', 'queued', 'running', 'cancelling', 'succeeded', 'failed', 'cancelled', 'interrupted')), chain_id=string(), last_heartbeat_at={'type': ['string', 'null']}, poll_after_ms={'type': ['integer', 'null']}, spec={'type': 'object'}, progress={'type': 'object'}, result={'type': 'object'}, + waiting_for=array(string()), notice=string(), error={'$ref': '#/components/schemas/Error'})} + job_view['allOf'] = [{'if': {'required': ['kind'], 'properties': {'kind': {'const': 'image'}}}, + 'then': {'properties': {'result': image_result()}}}] return dict(openapi='3.1.0', info=dict(title='FreeVideo Agent API', version=str(API_VERSION)), paths=paths, components=dict(securitySchemes={'bearer': {'type': 'http', 'scheme': 'bearer'}, 'page': {'type': 'apiKey', 'in': 'cookie', 'name': 'freevideo_agent_page'}}, schemas={ 'Session': session_view(), 'SessionEvent': session_event(stored=True), - 'JobSpec': schemas()['job'], 'Batch': batch_schema(), 'Job': job_view, + 'JobSpec': schemas()['job'], 'ImageResult': image_result(), 'CastCard': cast_view(), + 'Batch': batch_schema(), 'Job': job_view, # One job answers the way a submission does: the same job is also under 'job'. 'JobAnswer': dict(job_view, properties=dict(job_view['properties'], job={'$ref': '#/components/schemas/Job'})), 'Submitted': obj(dict(job={'$ref': '#/components/schemas/Job'}, @@ -1274,7 +1339,7 @@ def openapi(): TOOL_ENDPOINTS.update({SESSION_TOOLS[name]: SESSION_OPERATIONS[name][:2] for name in SESSION_TOOLS if name != 'post'}) TOOL_KINDS = dict(freevideo_generate='generate', freevideo_assemble='assemble', - freevideo_ask='vlm', freevideo_enhance_prompt='enhance') + freevideo_ask='vlm', freevideo_enhance_prompt='enhance', freevideo_generate_image='image') SESSION_ARGUMENT = dict(type='string', pattern='^s[0-9a-f]{12}$', @@ -1381,6 +1446,8 @@ def _tool_schema(name): kind = TOOL_KINDS.get(name, 'generate') value = KINDS[kind](local=True) value['properties'].pop('kind') + if kind == 'image': + value.pop('allOf') # Said in the field descriptions instead; the job is still checked when submitted. value['required'] = [k for k in value.get('required', []) if k != 'kind'] if name != 'freevideo_plan': value['properties']['wait'] = dict(CLIENT_WAIT, default=0 if kind == 'generate' else 25) diff --git a/freevideo_engine/agent/spec.py b/freevideo_engine/agent/spec.py index 8da22005..d44e60de 100644 --- a/freevideo_engine/agent/spec.py +++ b/freevideo_engine/agent/spec.py @@ -215,6 +215,19 @@ def normalize(value, *, port=None): item['video'] = media_item({k: v for k, v in video.items() if k != 'frames'} if isinstance(video, dict) else video, '/questions/video') result['questions'].append(item) result.update({k: value[k] for k in ('format', 'max_tokens', 'system') if k in value}) + elif kind == 'image': + width, height = ((value['width'], value['height']) if 'width' in value else + contract.IMAGE_ASPECTS[value.get('aspect', '1:1')]) + result.update(prompt=value['prompt'], width=width, height=height, + images=[item if isinstance(item, str) else item['file'] for item in value.get('images', [])]) + if value.get('seed') is not None: + result['seed'] = value['seed'] + if value.get('people'): + result['people'] = [dict(image=person['image'], x=person.get('x', .5), height=person.get('height', .62), + feet=person.get('feet', .9), **({'name': person['name']} if person.get('name') else {})) + for person in value['people']] + if value.get('scene_image'): + result['scene_image'] = value['scene_image'] elif kind == 'export': result.update(editor_project=value['editor_project'], revision=value['revision']) elif kind == 'listen': diff --git a/freevideo_engine/bootstrap.py b/freevideo_engine/bootstrap.py index ead734a7..5b990d94 100644 --- a/freevideo_engine/bootstrap.py +++ b/freevideo_engine/bootstrap.py @@ -328,6 +328,34 @@ def usable_git(path=None, env=None): return executable if version and version >= git_floor() else None +def gcc_major(executable, env=None): + """Major version of a GCC driver, or None when it does not run or does not answer like GCC.""" + try: + text = subprocess_output.output_external([str(executable), '-dumpfullversion', '-dumpversion'], env=env, + stderr=subprocess.DEVNULL, timeout=30) + return int(text.split()[0].split('.')[0]) + except (OSError, subprocess.SubprocessError, ValueError, IndexError): + return None + + +def host_compiler_choice(below, path=None, env=None, find=shutil.which, major=gcc_major): + """Which C/C++ compiler builds GPU code when the toolkit accepts GCC only below major version `below`. + + ({}, ...) when the system compiler is accepted, or when its version cannot be told: the build then + behaves as it always did. ({'CC': ..., 'CXX': ...}) when it is too new and a versioned pair + (g++-15 and gcc-15 ...) is installed beside it. None when no installed compiler is accepted. + """ + system = find('g++', path=path) + version = major(system, env) if system else None + if version is None or version < below: + return {} + for candidate in range(below - 1, 8, -1): + cxx, cc = find('g++-%d' % candidate, path=path), find('gcc-%d' % candidate, path=path) + if cxx and cc and major(cxx, env) == candidate and major(cc, env) == candidate: + return {'CC': cc, 'CXX': cxx} + return None + + def existing_parent(path): while not path.exists(): path = path.parent @@ -340,6 +368,10 @@ def prepared_precision(prepared): def model_target(row, model_dir, encoder_dir, prepared_dir=None): + from .model_status import family + if family(row) == 'image': + from .image_gen.image_pack import target + return target(row, model_dir) if row.get('sampling_file'): from .adaln_assets import asset_path from .sampling_assets import cache_root @@ -1060,6 +1092,8 @@ def setup_task(label): 'windows-runtime-check': ('Check GPU runtime', 'Verify CUDA and the Windows runtime libraries'), 'portable-git': ('Prepare Git', 'Install the verified Git packages'), 'portable-git-check': ('Prepare Git', 'Check the installed Git'), + 'portable-gcc': ('Prepare compiler', 'Install the verified compiler packages'), + 'portable-gcc-check': ('Prepare compiler', 'Check the installed compiler'), } title, detail = tasks.get(key, ('Prepare dependencies', 'Prepare verified source files')) for prefix, name in (('h3-text-encoder-', 'text encoder'), ('SageAttention-', 'SageAttention 2')): @@ -1160,6 +1194,11 @@ def initialize(self, value, ui): if value.get('disk_mode') == 'extreme': from .install_disk import environment self.env = environment(self.root, self.env) + # A folder the system code page cannot spell cannot hold what compiles kernels: + # every step runs from the short folder, and files and receipts keep the real one. + from . import launch_view + if launch_view.current(self.env) is None: + self.env = launch_view.environment_view(self.env, launch_view.for_engine(self.root)) self.network = dict(value.get('network', {}), events_path=str(self.run_dir / 'network.jsonl'), quiet=True) self.state = {'status': 'running', 'steps': [], 'plan': value} if self.damaged_before: @@ -1640,9 +1679,9 @@ def install_sage(self, uv, python): 'torch=str(torch.__version__),cuda=torch.version.cuda,cxx11_abi=torch.compiled_with_cxx11_abi(), ' 'machine=platform.machine(),glibc=platform.libc_ver())))') identity = json.loads(subprocess_output.output_python([str(python), '-c', code], env=self.env)) + compiler = self.host_compiler(identity['cuda']) identity.update(arch='.'.join(map(str, self.plan['inventory']['hardware']['capability'])), - source=self.spec['sageattention']['commit'], compiler=subprocess_output.output_external( - ['g++', '--version']).splitlines()[0]) + source=self.spec['sageattention']['commit'], compiler=compiler['identity']) wheel_dir = Path(self.plan['wheel_cache']) / wheel_key(identity) manifest = wheel_dir / 'manifest.json' record = json.loads(manifest.read_text(encoding='utf-8')) if manifest.is_file() else None @@ -1672,6 +1711,7 @@ def install_sage(self, uv, python): build_env = dict(self.env, CUDA_HOME=str(toolkit), PATH=str(toolkit / 'bin') + os.pathsep + self.env['PATH'], CPATH=headers, TORCH_CUDA_ARCH_LIST=identity['arch'], MAX_JOBS=str(build['jobs']), EXT_PARALLEL='1', NVCC_APPEND_FLAGS='--threads=' + str(build['nvcc_threads'])) + build_env.update(compiler['env'] if compiler['env'] is not None else self.install_compiler_linux()) self.command('sage2-toolchain-check', [python, '-c', 'from torch.utils.cpp_extension import is_ninja_available; ' 'assert is_ninja_available(), "Ninja must be visible in the build PATH; serial fallback is refused"'], env=build_env) @@ -1690,6 +1730,47 @@ def install_sage(self, uv, python): save(self.run_dir / 'sage-wheel.json', dict(record, reused=cached, cache=str(wheel_dir))) return manifest + def host_compiler(self, cuda): + """The compiler pair Sage is built with: the system's, a versioned one beside it, or FreeVideo's own. + + `env` is empty when the system compiler is used; `identity` names the compiler in the wheel cache key. + A portable compiler is only installed when a build is needed (see install_compiler_linux). + """ + spec = json.loads((PACKAGE / 'gcc_linux.json').read_text(encoding='utf-8')) + below = spec['host_compiler_below'].get(cuda) + choice = host_compiler_choice(below, self.env['PATH'], self.env) if below else {} + if choice is None: + return {'env': None, 'identity': 'conda-forge gcc ' + spec['version']} + cxx = choice.get('CXX', 'g++') + return {'env': choice, 'identity': subprocess_output.output_external([cxx, '--version']).splitlines()[0]} + + def install_compiler_linux(self): + """GCC from conda-forge, pinned by package hashes, for a system compiler newer than the toolkit accepts.""" + spec = json.loads((PACKAGE / 'gcc_linux.json').read_text(encoding='utf-8')) + prefix = self.root / 'tools' / ('gcc-' + spec['version']) + triple = 'x86_64-conda-linux-gnu' + compilers = {'CC': prefix / 'bin' / (triple + '-gcc'), 'CXX': prefix / 'bin' / (triple + '-g++')} + receipt, content = prefix / '.freevideo-gcc.json', {k: spec[k] for k in ('version', 'packages')} + if not (receipt.is_file() and json.loads(receipt.read_text(encoding='utf-8')) == content + and all(path.is_file() for path in compilers.values())): + micromamba, packages = self.fetch_conda_lock(spec, 'gcc-linux', 'compiler') + if prefix.is_symlink(): + prefix.unlink() + elif prefix.exists(): + shutil.rmtree(prefix) + explicit = self.root / 'downloads' / 'gcc-linux' / 'explicit.txt' + explicit.write_text('@EXPLICIT\n' + ''.join(path.resolve().as_uri() + '\n' for path in packages), encoding='utf-8') + cache = self.root / 'tools' / 'micromamba-root' + env = {k: v for k, v in self.env.items() if not k.upper().startswith(('CONDA', 'MAMBA'))} + self.command('portable-gcc', [micromamba, 'create', '--yes', '--no-rc', '--offline', + '--root-prefix', cache, '--prefix', prefix, '--file', explicit], env=env) + save(receipt, content) + shutil.rmtree(cache, ignore_errors=True) + shutil.rmtree(explicit.parent, ignore_errors=True) + micromamba.unlink(missing_ok=True) + self.command('portable-gcc-check', [compilers['CXX'], '--version']) + return {key: str(path) for key, path in compilers.items()} + def install_sage_windows(self, uv, python): spec = self.versions['windows']['sageattention'] directory = Path(self.plan['wheel_cache']) / 'windows-sage2' / spec['sha256'][:16] @@ -1750,7 +1831,10 @@ def install_git_linux(self): def fetch_git_linux(self, spec): """micromamba and the locked packages, verified by SHA-256, as one step.""" - directory = self.root / 'downloads' / 'git-linux' + return self.fetch_conda_lock(spec, 'git-linux', 'Git') + + def fetch_conda_lock(self, spec, folder, label): + directory = self.root / 'downloads' / folder tool = spec['micromamba'] micromamba = self.root / 'tools' / ('micromamba-' + tool['version']) channels = spec['channels'] @@ -1762,21 +1846,21 @@ def fetch_git_linux(self, spec): rows += [(directory / row['path'].split('/')[-1], [(n, channels[n] + '/' + row['path']) for n in names], row) for row in spec['packages']] total = sum(row['bytes'] for _, _, row in rows) - key = 'download-git-linux' - self.ui.begin(key, 'Prepare Git', detail='Download and verify Git %s' % spec['version']) + key = 'download-' + folder + self.ui.begin(key, 'Prepare ' + label, detail='Download and verify %s %s' % (label, spec['version'])) done = 0 try: for path, candidates, row in rows: def progress(current, size, speed, before=done): if self.cancel.is_set(): raise RuntimeError('Installation cancelled; partial download retained.') - # One step for micromamba and the packages: the line names Git with the step's total. - self.ui.update(key, done=before + current, total=total, rate=speed, unit='bytes', scope='git', - activity='download', name='Git', + # One step for micromamba and the packages: the line names the tool with the step's total. + self.ui.update(key, done=before + current, total=total, rate=speed, unit='bytes', scope=label.lower(), + activity='download', name=label, detail='%.1f MiB / %.1f MiB' % ((before + current) / 2**20, total / 2**20)) def feedback(event): if event['action'] in ('retry', 'route-retry'): - self.ui.update(key, activity='resume', name='Git') + self.ui.update(key, activity='resume', name=label) network.download(candidates, path, row['sha256'], progress, network=dict(self.network, quiet=True, event_callback=feedback), env=self.env, size=row['bytes'], category='github' if path == micromamba else 'download') @@ -2066,6 +2150,26 @@ def main(argv=None): prepare_end(ui, False) +def install_image_model(engine_root, model_root): + """The optional image model, after the video installation is complete; its failure never fails setup.""" + from .image_gen import image_pack + last = [-1] + def progress(value): + total, done = value.get('total') or 0, value.get('done') or 0 + if total and done * 100 // total != last[0]: + last[0] = done * 100 // total + print('Image model: %d%%' % last[0], flush=True) + print('Downloading the image model (%.1f GiB)…' % (sum(r['bytes'] for r in image_pack.rows()) / 2**30), flush=True) + try: + image_pack.install(Path(model_root), progress, engine_root=Path(engine_root)) + except image_pack.PackError as error: + print('The image model was not downloaded (%s). Download it later from the launcher.' % error.reason, file=sys.stderr) + except (OSError, ValueError, RuntimeError) as error: + print('The image model was not downloaded (%s). Download it later from the launcher.' % type(error).__name__, file=sys.stderr) + else: + print('Image model ready.') + + def run_main(argv, shown): parser = argparse.ArgumentParser(description=__doc__) parser.add_argument('--root', type=Path, default=Path(os.environ.get('FREEVIDEO_HOME', DEFAULT_ROOT))) @@ -2076,6 +2180,9 @@ def run_main(argv, shown): parser.add_argument('--sampling-caches', action=argparse.BooleanOptionalAction, default=None, help='Install every quality level in advance (default for new installations; existing ones keep ' 'their choice). The default 8 + 3 refinement tables are always installed.') + parser.add_argument('--image-model', action='store_true', + help='After setup, also download the optional image generation model (about 16 GiB). ' + 'A failure here does not fail setup; the launcher can download it later') parser.add_argument('--prepared-format', choices=('auto', 'int8', 'fp8'), default='auto', help='auto: an installation keeps its recorded format and new installations use int8 on GeForce ' 'and Ampere cards, FP8 elsewhere. int8/fp8 switch explicitly; files of the other format are kept') @@ -2224,6 +2331,8 @@ def interrupted(signum, frame): if installer.state['status'] == 'complete': if args.verbose: print('Configuration: %s' % (installer.root / 'machine.json')) + if args.image_model and platform.system() != 'Darwin': + install_image_model(installer.root, value['model_dir']) from .comfy_service import after_setup, continues_after_setup if continues_after_setup(no_launch=args.no_launch, approved_plan=args.approved_plan): print('\nSetup complete.') diff --git a/freevideo_engine/chain_plan.py b/freevideo_engine/chain_plan.py index d75a42ae..d7cf593f 100644 --- a/freevideo_engine/chain_plan.py +++ b/freevideo_engine/chain_plan.py @@ -11,6 +11,10 @@ FPS, UNIT, OVERLAP = 24, 17, 22 SINGLE_SECONDS, BLOCK_SECONDS, SHORTEST_BLOCK = 30, 15, 5 +# A part that continues another is at least this long (the first part is not held to it). +MIN_PART_SECONDS = 10 +MIN_PART_FRAMES = geometry(seconds=MIN_PART_SECONDS)['frames'] +SHORTEST_FIRST = 2 # whole seconds: what the first part keeps when the continuing parts take the rest MAX_SECONDS = 7200 MAX_BLOCKS = MAX_SECONDS // SHORTEST_BLOCK @@ -72,36 +76,51 @@ def plan(seconds, single=False, max_frames=None, single_limit_seconds=None, bloc fits = frames <= limit alone = dict(seconds=seconds, frames=new, single=True, hide_frames=0, blocks=[dict(seconds=seconds, frames=frames, new_frames=new)]) + head = following if continued else first # the longest part this video can start with # Under 30 s stays one generation unless a block length was asked for; 30 s is two blocks of 15 s. - if fits and (single or block_seconds is None and seconds < SINGLE_SECONDS - or seconds <= min(block, following if continued else first)): + if fits and (single or block_seconds is None and seconds < SINGLE_SECONDS or seconds <= min(block, head)): return alone - first, following = min(block, following if continued else first), min(block, following) - if first < SHORTEST_BLOCK or following < SHORTEST_BLOCK: + if block < MIN_PART_SECONDS: + if fits and seconds < SINGLE_SECONDS: + return alone # a block this short cannot continue another: a short video stays one generation + raise PlanError('block_too_short', 'A block that continues another is at least %d s long' % MIN_PART_SECONDS) + first, following = min(block, head), min(block, following) + shortest = MIN_PART_SECONDS if continued else SHORTEST_BLOCK + if first < shortest or following < MIN_PART_SECONDS: + if fits and seconds < SINGLE_SECONDS: + return alone # a requested block length never makes a part this device cannot do raise PlanError('device_limit', 'This device cannot generate this length at this size') - total = seconds = math.floor(seconds + .5) - lengths = [min(first, total)] - while sum(lengths) < total: - if len(lengths) >= MAX_BLOCKS: - raise PlanError('too_long', 'This length needs too many blocks') - lengths.append(min(following, total - sum(lengths))) - if len(lengths) > 1 and lengths[-1] < SHORTEST_BLOCK: - # No short tail: the last two blocks share their length, the earlier one gets the odd second. - pair = lengths[-2] + lengths[-1] - lengths[-2:] = [pair - pair // 2, pair // 2] - if fits and block_seconds is not None and total < SINGLE_SECONDS and min(lengths) < SHORTEST_BLOCK: - return alone # a requested block length never makes a part under 5 s + # Under 30 s the length stays as asked (it may be fractional); from 30 s it is whole seconds already. + total = seconds + # Every part that continues another is at least MIN_PART_SECONDS and at most `following`. The part that + # starts a new video is not held to it; a rest too short for the continuing parts is taken from that part. + rest = total if continued else total - min(first, total) + count = math.ceil(round(rest, 6) / following) if rest > 0 else 0 + rest = max(rest, MIN_PART_SECONDS * count) + whole = math.ceil(round(rest, 6)) + if continued and whole > math.ceil(round(total, 6)) or not continued and total - whole < SHORTEST_FIRST: + if fits and seconds < SINGLE_SECONDS: + return alone + raise PlanError('length_not_splittable', 'This length cannot be cut into parts of the requested length') + if count > MAX_BLOCKS: + raise PlanError('too_long', 'This length needs too many blocks') + lengths = [whole // count + (index < whole % count) for index in range(count)] + if continued: + lengths[0] = round(lengths[0] - (whole - total), 6) # a fractional request keeps its length in the first part + else: + lengths.insert(0, round(total - whole, 6)) + frames_all = int(round(FPS * total)) blocks, shown = [], 0 for index, length in enumerate(lengths): - target = FPS * sum(lengths[:index + 1]) + target = int(round(FPS * sum(lengths[:index + 1]))) if index == 0 and not continued: new = requested = _first_frames(length) else: # Each block takes the shortest step that reaches its whole second; nothing is cut in between. new = max(1, math.ceil((target - shown) / UNIT)) * UNIT requested = new + OVERLAP - if requested > limit or index and length > following or not index and length > first: + if requested > limit or length > (first if not index else following): raise PlanError('device_limit', 'This device cannot generate this length at this size') shown += new blocks.append(dict(seconds=length, frames=requested, new_frames=new)) - return dict(seconds=total, frames=FPS * total, single=False, hide_frames=shown - FPS * total, blocks=blocks) + return dict(seconds=total, frames=frames_all, single=False, hide_frames=shown - frames_all, blocks=blocks) diff --git a/freevideo_engine/chain_runner.py b/freevideo_engine/chain_runner.py index ffa5b35d..469353cf 100644 --- a/freevideo_engine/chain_runner.py +++ b/freevideo_engine/chain_runner.py @@ -8,7 +8,7 @@ import time import uuid -from . import chain_leads +from . import chain_leads, chain_plan from . import chain_store as store @@ -26,6 +26,16 @@ def __init__(self, parts, limit): self.message_zh = '这条视频最多 %d 段，请缩短一些。' % limit +class ChainPartTooShort(ValueError): + """A continuing part shorter than a continued video may have.""" + code = 'chain_part_too_short' + + def __init__(self, seconds, limit): + super().__init__('Each continued part must be at least %d seconds. Make it longer.' % limit) + self.seconds, self.limit = seconds, limit + self.message_zh = '接着生成的每一段至少 %d 秒，请加长一些。' % limit + + def _work_key(path): return PurePosixPath(str(path).replace('\\', '/')).as_posix() @@ -58,9 +68,10 @@ def node_prompt(prompt, node): class ChainRunner: def __init__(self, store_dir, host, clock=time.monotonic, wall=time.time, estimate=None, poll=0.5, inputs=None, - max_parts=None): + max_parts=None, min_part_frames=None): self.store_dir, self.host = store_dir, host self.max_parts = max_parts or store.MAX_PARTS + self.min_part_frames = chain_plan.MIN_PART_FRAMES if min_part_frames is None else min_part_frames self.clock, self.wall, self.estimate, self.poll = clock, wall, estimate, poll self.inputs = inputs or host.inputs self.lock = threading.RLock() @@ -175,6 +186,11 @@ def create(self, payload): record = store.validate(payload) with self.lock: # Before anything is written or queued: a refused chain leaves every record and file as it was. + short = [row['inputs']['frames'] for row in record['blocks'] + if (row['after'] is not None or row['before'] is not None) + and row['inputs']['frames'] < self.min_part_frames] + if short: + raise ChainPartTooShort(min(short) / chain_plan.FPS, chain_plan.MIN_PART_SECONDS) parts = max(self._depths(record).values()) if parts > self.max_parts: raise ChainTooLong(parts, self.max_parts) diff --git a/freevideo_engine/comfy_bridge.py b/freevideo_engine/comfy_bridge.py index f7404d67..4598e5d7 100644 --- a/freevideo_engine/comfy_bridge.py +++ b/freevideo_engine/comfy_bridge.py @@ -1185,10 +1185,7 @@ def generate(prompt, width, height, seconds, seed, output_directory, *, continuation = None def check_ending(): - try: - from .continuation_tail import check - except ImportError: - return + from .continuation_tail import check from .two_pass import sampling_canvas check(continue_from, sampling_canvas(canvas)[0]) diff --git a/freevideo_engine/comfy_chain.py b/freevideo_engine/comfy_chain.py index ffd73b6d..b28625de 100644 --- a/freevideo_engine/comfy_chain.py +++ b/freevideo_engine/comfy_chain.py @@ -7,7 +7,7 @@ import ssl from urllib import error, request -from .chain_runner import ChainBusy, ChainRunner, ChainTooLong +from .chain_runner import ChainBusy, ChainPartTooShort, ChainRunner, ChainTooLong from . import chain_plan, chain_store RUNNER = None @@ -225,9 +225,10 @@ def register(host_factory=None): host = host_factory(server) if fake else ComfyHost(server) if fake: logging.warning('[FreeVideo] Test background chain host is enabled') - # Only the test host may lift the limit on parts, never the package. + # Only the test host may lift the limits on parts, never the package. runner = ChainRunner(Path(root) / 'chains', host, estimate=None, - max_parts=getattr(host, 'max_parts', None) if fake else None) + max_parts=getattr(host, 'max_parts', None) if fake else None, + min_part_frames=getattr(host, 'min_part_frames', None) if fake else None) if fake: host.runner = runner RUNNER = runner @@ -275,6 +276,9 @@ async def create(request): except ChainTooLong as failure: return response({'error': str(failure), 'error_zh': failure.message_zh, 'code': failure.code, 'parts': failure.parts, 'max_parts': failure.limit}, 400) + except ChainPartTooShort as failure: + return response({'error': str(failure), 'error_zh': failure.message_zh, 'code': failure.code, + 'min_seconds': failure.limit}, 400) except chain_plan.PlanError as failure: return response({'error': str(failure), 'code': failure.code}, 400) except ValueError as failure: diff --git a/freevideo_engine/comfy_nodes.py b/freevideo_engine/comfy_nodes.py index 10f5181b..e11d77fa 100644 --- a/freevideo_engine/comfy_nodes.py +++ b/freevideo_engine/comfy_nodes.py @@ -130,6 +130,11 @@ def execute(cls, text, width, height, seconds, seed, first=None, last=None, from .resident_process import OWNER from .prompt_vlm.service import stop_prestart_for_generation stop_prestart_for_generation(PromptServer.instance) + try: # the image option must never keep a video from starting + from .image_gen.service import stop_prestart_for_generation as stop_image_for_generation + stop_image_for_generation(PromptServer.instance) + except Exception as error: + logging.warning('[FreeVideo] Could not stop image generation before the video: %s', error) prewarm = {} result_reused = [False] result_report = [None] diff --git a/freevideo_engine/compute_budget.py b/freevideo_engine/compute_budget.py index 449c1ed6..8a0764a7 100644 --- a/freevideo_engine/compute_budget.py +++ b/freevideo_engine/compute_budget.py @@ -52,6 +52,9 @@ def spare_compute_trials(row, engine, peak_bytes, headroom_bytes): workspace_per_lane_bytes=int(workspace), reason='Overlap unchanged head groups; extra live workspace uses measured spare VRAM'))) canvas = row.get('geometry', {}) + if canvas.get('early_reference') is not None: + # Measured spare memory has no stage identity in this proposal schema. + return result rows = canvas.get('video_tokens') if type(rows) is not int or rows <= 0: return result diff --git a/freevideo_engine/download_settings.py b/freevideo_engine/download_settings.py index 6a8a1e58..4afc98fa 100644 --- a/freevideo_engine/download_settings.py +++ b/freevideo_engine/download_settings.py @@ -10,7 +10,7 @@ from .proxy import MODES CHOICES = ('auto', 'official', 'hf-mirror', 'modelscope') -FAMILIES = ('edge-models', 'vdn-models', 'models') +FAMILIES = ('edge-models', 'vdn-models', 'models', 'image-models') def speed_text(row, zh=False): diff --git a/freevideo_engine/effort_forecast.py b/freevideo_engine/effort_forecast.py index 71d959b8..31839066 100644 --- a/freevideo_engine/effort_forecast.py +++ b/freevideo_engine/effort_forecast.py @@ -38,6 +38,10 @@ def timing_record(report, gpu_uuid, system, *, stage_details=True): if not same_device(hardware, gpu_uuid) or hardware.get('system') != system: return None shape = report.get('geometry', {}) + # The pooled ETA/capability schema has no per-step reference schedule. + # An early-only success cannot calibrate a sustained-reference request. + if shape.get('early_reference') is not None: + return None plan = report.get('sampling_plan') or shape.get('sampling_plan') or {} if not all(type(shape.get(k)) is int and shape[k] > 0 for k in ('width', 'height', 'frames')): return None diff --git a/freevideo_engine/encode_worker.py b/freevideo_engine/encode_worker.py index 464453c5..d5882626 100644 --- a/freevideo_engine/encode_worker.py +++ b/freevideo_engine/encode_worker.py @@ -321,12 +321,18 @@ def phase(name, **metrics): normalized = None media_kwargs = {} task = 't2va' + early_reference = None if request.get('media'): - from .media_request import task_for + from .media_request import task_for, reference_stages from .media_encoding import prepare - task = task_for(request['media']) + stages = reference_stages(request['media']) + early = [dict(reference_index=i, early_steps=row['early_steps']) + for i, row in enumerate(request['media'].get('references', [])) if 'early_steps' in row] + if early: + early_reference = dict(version=1, early_steps=early, conditions=[]) + task = task_for(stages[0][1]) if task != 't2va': - normalized, media_kwargs = prepare(request['media'], request['geometry'], Path(request['output']).parent / 'media') + normalized, media_kwargs = prepare(stages[0][1], request['geometry'], Path(request['output']).parent / 'media') # Saved with the conditioning, so a reused input cache still reports them. loading['reference_trims'] = [row['trimmed'] for row in normalized if row.get('trimmed')] if keyframes is not None: @@ -516,6 +522,26 @@ def attempt(index, failure): encoded, attempts = encode_with_recovery(clip, tokens, memory, torch, phase, max_attempts=4, attempt=attempt) loading['encoder_attempts'] = attempts + if early_reference is not None: + # Reuse this loaded Qwen instance, its precision and recovery policy. + # Re-tokenize: removing visual rows alone does not remove VL features. + from .media_encoding import reference_options + extra_started = time.perf_counter() + for step, stage_media in stages[1:]: + tokens = clip.tokenize(prompt, **reference_options(normalized, media_kwargs, step)) + size = input_size(tokens) + current.clear() + with readonly_pinning(torch, memory): + stage_encoded, stage_attempts = encode_with_recovery( + clip, tokens, memory, torch, phase, max_attempts=4, attempt=attempt) + stage_task = task_for(stage_media) + stage_value = to_cache(stage_encoded, prompt, task=stage_task) + stage_value['task'] = stage_task + early_reference['conditions'].append(dict(start_step=step, value=stage_value)) + loading.setdefault('early_reference_attempts', []).append(dict(start_step=step, attempts=stage_attempts)) + del stage_encoded + torch.cuda.synchronize() + loading['early_reference_extra_encode_seconds'] = time.perf_counter() - extra_started if stages[1:] else 0. # Finish async forward work at the existing conditioning boundary, so # it is not attributed to packing/saving on the host. torch.cuda.synchronize() @@ -532,6 +558,8 @@ def attempt(index, failure): task = 'fl2va' if images is not None else task value = to_cache(encoded, prompt, task=task) value['task'] = task + if early_reference is not None: + value['early_reference'] = early_reference torch.cuda.synchronize() prompt_encode_seconds = time.perf_counter() - encode_started # Keyframe and reference-media encoding read the drive themselves. @@ -586,6 +614,10 @@ def attempt(index, failure): info = encode_latents(normalized, value, request['base'], request['geometry'], resident=resident) keyframe_metrics.update(media_vae_encode_seconds=time.perf_counter() - tick, conditioning_info=info) task = info['task'] + elif early_reference is not None: + from .media_conditioning import describe + keyframe_metrics['conditioning_info'] = describe(value, **{ + k: request['geometry'][k] for k in ('width', 'height', 'frames')}) phase('encoder_save') output = Path(request['output']) output.parent.mkdir(parents=True, exist_ok=True) diff --git a/freevideo_engine/failure_details.py b/freevideo_engine/failure_details.py index ef0772d3..9d6a716c 100644 --- a/freevideo_engine/failure_details.py +++ b/freevideo_engine/failure_details.py @@ -508,6 +508,11 @@ def launcher_failure(value, *, zh=False, other_disk=False, disk_name='', can_rep ('It may have been quarantined by security software. Restore the file or allow the FreeVideo folder, ' 'then click Retry installation.', '可能被安全软件隔离了。恢复这个文件或允许 FreeVideo 所在的文件夹后，点“重试安装”。')), + # Windows refuses PyTorch's unsigned DLLs under Smart App Control; reinstalling changes nothing. + (r'WinError 4551|Application Control policy has blocked', 'app-control', + ('Windows blocked the PyTorch files', 'Windows 拦截了 PyTorch 的文件'), + ('Smart App Control in Windows blocked these files. Search Windows for "Smart App Control", or open Windows Security → App & browser control → Smart App Control settings.', + 'Windows 的“智能应用控制”拦截了这些文件。在 Windows 中搜索“智能应用控制”，或打开“Windows 安全中心 → 应用和浏览器控制 → 智能应用控制设置”。')), (r'PermissionError|Access is denied|Permission denied|WinError 5\b', 'permission', ('This folder is not writable', '无法写入这个目录'), ('Choose an installation folder you can write to, then retry.', '选择有写入权限的安装目录后重试。')), diff --git a/freevideo_engine/gcc_linux.json b/freevideo_engine/gcc_linux.json new file mode 100644 index 00000000..b832a479 --- /dev/null +++ b/freevideo_engine/gcc_linux.json @@ -0,0 +1,115 @@ +{ + "version": "14.4.0", + "host_compiler_below": { + "12.9": 15, + "13.0": 16 + }, + "micromamba": { + "version": "2.9.0-0", + "url": "https://github.com/mamba-org/micromamba-releases/releases/download/2.9.0-0/micromamba-linux-64", + "sha256": "366cd9cd8be14df1ab8ed50352a82111082a36686b2d389fdb79a92c3fafb3e3", + "bytes": 18292808 + }, + "channels": { + "official": "https://conda.anaconda.org/conda-forge", + "tuna": "https://mirrors.tuna.tsinghua.edu.cn/anaconda/cloud/conda-forge", + "nju": "https://mirror.nju.edu.cn/anaconda/cloud/conda-forge" + }, + "packages": [ + { + "path": "linux-64/_openmp_mutex-4.5-20_gnu.conda", + "sha256": "1dd3fffd892081df9726d7eb7e0dea6198962ba775bd88842135a4ddb4deb3c9", + "bytes": 28948 + }, + { + "path": "linux-64/binutils_impl_linux-64-2.46.1-default_hfdba357_103.conda", + "sha256": "98f9ddc6dc5d56a824b7eaea5bab6ad4bb47642288fc70d6b1b4024ac39d5595", + "bytes": 3706604 + }, + { + "path": "linux-64/binutils_linux-64-2.46.1-default_h4852527_103.conda", + "sha256": "606a42c754752fb905fe2012d52ff370d47b5b6adefd5f7befc4ec30ce758f23", + "bytes": 36335 + }, + { + "path": "linux-64/gcc_impl_linux-64-14.4.0-h4b2a866_7.conda", + "sha256": "fe3d0c2cbd59755ddd8364c79b8e1a09b3125581d9df07170f102bc150740632", + "bytes": 78647149 + }, + { + "path": "linux-64/gcc_linux-64-14.4.0-hc469d41_2.conda", + "sha256": "082fb68319027d485cc9b4d8effc3faa49d2edb13bc530f849acd321c8cfb6c2", + "bytes": 30385 + }, + { + "path": "linux-64/gxx_impl_linux-64-14.4.0-hc436dd5_7.conda", + "sha256": "8082bf31bf210f55ab3e69bbc301e5131b8fe2986c743b4d4486d469dc737ac3", + "bytes": 15847507 + }, + { + "path": "linux-64/gxx_linux-64-14.4.0-h0047a55_2.conda", + "sha256": "9801d3b03494c5530e6574eae1055d8b5bbb99d0ac75d3bca8757c17a2f3a001", + "bytes": 28787 + }, + { + "path": "linux-64/ld_impl_linux-64-2.46.1-default_hbd61a6d_103.conda", + "sha256": "901649b1edbe841fd4ee7e2d491a05184ed540eb01cdae506872f344babc1d28", + "bytes": 742109 + }, + { + "path": "linux-64/libgcc-16.2.0-ha9f2e26_7.conda", + "sha256": "e031634c3a928f9594eba868bc46ef5230d6571a364a0c63215b06c4563bfcea", + "bytes": 1058444 + }, + { + "path": "linux-64/libgomp-16.2.0-he0feb66_7.conda", + "sha256": "92a0bb350c7be6a4eedcb4e111c60aefd176018e6b1c009a6e025dbe9bbcd758", + "bytes": 640566 + }, + { + "path": "linux-64/libsanitizer-14.4.0-hf39dbba_7.conda", + "sha256": "7a4c47d699e6e9640aa8061f461ab31e71ba7d41994a451ea1e79a3801e8f6c3", + "bytes": 7211298 + }, + { + "path": "linux-64/libstdcxx-16.2.0-h934c35e_7.conda", + "sha256": "fa7018298629fc429971fffd6ecba3d27a7b02e4e890d6a93ea01ffebe8cb745", + "bytes": 6615436 + }, + { + "path": "linux-64/libzlib-1.3.2-h25fd6f3_3.conda", + "sha256": "eb8a0db0aa570124f7d2a93d7c7f596e3390df5e047818d873baad32985fc736", + "bytes": 63713 + }, + { + "path": "linux-64/zstd-1.5.7-hb78ec9c_7.conda", + "sha256": "47d682b9f6d6ec9eb1a6e6c3e75ea6273e899e78fb7fc59f81d39745009fbc60", + "bytes": 601301 + }, + { + "path": "noarch/kernel-headers_linux-64-5.14.0-he073ed8_3.conda", + "sha256": "c4d973dd80ab10028b3e4c58c1c6d8f437bfeaf6069dafb1130dabc550e9caf4", + "bytes": 1410267 + }, + { + "path": "noarch/libgcc-devel_linux-64-14.4.0-hd9a9cd0_107.conda", + "sha256": "fb09fdc5d12cc929ff0af73fbd7cf33be7b473a46fe11353a21ec1ffa468c657", + "bytes": 3090192 + }, + { + "path": "noarch/libstdcxx-devel_linux-64-14.4.0-ha5b54cb_107.conda", + "sha256": "a518b2878cecd324082d334d36907156b68074b18736c111b15ff34efcb22aaf", + "bytes": 19512223 + }, + { + "path": "noarch/sysroot_linux-64-2.34-h087de78_3.conda", + "sha256": "08c50c314b331730eeb8f87c39a45e5426a5463e078392ab283e2552d753443e", + "bytes": 40779381 + }, + { + "path": "noarch/tzdata-2026c-h151e31d_0.conda", + "sha256": "b928c30ddcb0e3f544c6eade8352737e6e610e263276b90232db6a578ef899d8", + "bytes": 118849 + } + ] +} diff --git a/freevideo_engine/generate.py b/freevideo_engine/generate.py index 1603fc77..47f7275e 100644 --- a/freevideo_engine/generate.py +++ b/freevideo_engine/generate.py @@ -718,6 +718,8 @@ def interrupted(signum, frame): if type(count) is not int or count < 0: raise ValueError('Invalid reference token count') canvas[name] = count + if 'early_reference' in info: + canvas['early_reference'] = info['early_reference'] profile['engine']['task'] = info['task'] if not condition.is_file(): raise FileNotFoundError(condition) diff --git a/freevideo_engine/image_gen/__init__.py b/freevideo_engine/image_gen/__init__.py new file mode 100644 index 00000000..fcabd633 --- /dev/null +++ b/freevideo_engine/image_gen/__init__.py @@ -0,0 +1 @@ +"""Image generation (Qwen-Image 2.1 Turbo): model parts, planning and the on-demand pack; importing this package loads nothing.""" diff --git a/freevideo_engine/image_gen/_model_utils.py b/freevideo_engine/image_gen/_model_utils.py new file mode 100644 index 00000000..72aab9b1 --- /dev/null +++ b/freevideo_engine/image_gen/_model_utils.py @@ -0,0 +1,122 @@ +# Copyright 2025 HuggingFace Inc. +# Copyright 2026 The HuggingFace Team. All rights reserved. +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy at http://www.apache.org/licenses/LICENSE-2.0 +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Small local equivalents of the inference helpers used by the upstream models.""" +from collections import OrderedDict +from copy import deepcopy +from functools import wraps +import inspect +import math + +import torch +from torch import nn + + +class ModelConfig(dict): + def __getattr__(self, name): + try: + return self[name] + except KeyError as exc: + raise AttributeError(name) from exc + + +def model_config(init): + """Record explicit/default constructor values, with no serialization framework.""" + signature = inspect.signature(init) + + @wraps(init) + def wrapped(self, *args, **kwargs): + if len(args) == 1 and isinstance(args[0], dict): + kwargs = {**args[0], **kwargs} + args = () + bound = signature.bind(self, *args, **kwargs) + bound.apply_defaults() + self.config = ModelConfig({key: deepcopy(value) for key, value in bound.arguments.items() if key != "self"}) + init(self, *args, **kwargs) + + return wrapped + + +class ModelOutput(OrderedDict): + """Attribute, mapping and tuple-index access used by Diffusers callers.""" + def __init__(self, **values): + super().__init__(values) + self.__dict__.update(values) + + def __getitem__(self, key): + return super().__getitem__(key) if isinstance(key, str) else self.to_tuple()[key] + + def to_tuple(self): + return tuple(self.values()) + + +AutoencoderKLOutput = ModelOutput + + +class RMSNorm(nn.Module): + """Upstream FP32 variance/rsqrt arithmetic and checkpoint-compatible weight.""" + def __init__(self, dim, eps): + super().__init__() + self.eps = eps + self.weight = nn.Parameter(torch.ones(dim)) + + def forward(self, hidden_states): + variance = hidden_states.float().pow(2).mean(-1, keepdim=True) + hidden_states = hidden_states * torch.rsqrt(variance + self.eps) + if self.weight.dtype in (torch.float16, torch.bfloat16): + hidden_states = hidden_states.to(self.weight.dtype) + return hidden_states * self.weight + + +class TimestepEmbedding(nn.Module): + def __init__(self, in_channels, time_embed_dim, sample_proj_bias=False): + super().__init__() + self.linear_1 = nn.Linear(in_channels, time_embed_dim, bias=sample_proj_bias) + self.act = nn.SiLU() + self.linear_2 = nn.Linear(time_embed_dim, time_embed_dim, bias=sample_proj_bias) + + def forward(self, sample): + return self.linear_2(self.act(self.linear_1(sample))) + + +class DiagonalGaussianDistribution: + """Inference posterior; arithmetic follows Diffusers' diagonal Gaussian.""" + def __init__(self, parameters, deterministic=False): + self.parameters = parameters + self.mean, self.logvar = parameters.chunk(2, dim=1) + self.logvar = self.logvar.clamp(-30.0, 20.0) + self.deterministic = deterministic + self.std = torch.exp(0.5 * self.logvar) + self.var = torch.exp(self.logvar) + if deterministic: + self.var = self.std = torch.zeros_like(self.mean) + + def sample(self, generator=None): + # Match randn_tensor's CPU-generator support for accelerator posteriors. + noise_device = generator.device if generator is not None else self.mean.device + noise = torch.randn(self.mean.shape, generator=generator, device=noise_device, dtype=self.mean.dtype) + return self.mean + self.std * noise.to(self.mean.device) + + def mode(self): + return self.mean + + def kl(self, other=None): + if self.deterministic: + return torch.tensor([0.0]) + if other is None: + value = self.mean.pow(2) + self.var - 1.0 - self.logvar + else: + value = (self.mean - other.mean).pow(2) / other.var + self.var / other.var - 1.0 - self.logvar + other.logvar + return 0.5 * value.sum(dim=(1, 2, 3)) + + def nll(self, sample, dims=(1, 2, 3)): + if self.deterministic: + return torch.tensor([0.0]) + return 0.5 * (math.log(2.0 * math.pi) + self.logvar + (sample - self.mean).pow(2) / self.var).sum(dim=dims) diff --git a/freevideo_engine/image_gen/cast_views.py b/freevideo_engine/image_gen/cast_views.py new file mode 100644 index 00000000..9c985ad8 --- /dev/null +++ b/freevideo_engine/image_gen/cast_views.py @@ -0,0 +1,98 @@ +"""Pictures of a cast card made by the image model: one place for the view prompts and sizes. + +The front view comes from the description (or from a user's photo, edited into a clean card); every +other view of a character is an edit of the front picture, so they show the same person. A setting is +a single wide picture from text. +""" + +import re + +VIEWS = ('front', 'three_quarter', 'side', 'back', 'face') +SETTING_VIEWS = ('scene',) + +SIZES = { + 'front': (832, 1472), + 'three_quarter': (832, 1472), + 'side': (832, 1472), + 'back': (832, 1472), + 'face': (1024, 1024), + 'scene': (1344, 768), +} + +BACKDROP = 'in front of a plain light grey studio backdrop, even soft studio light, sharp focus, no text' +SAME = (' Keep the same person, facial features, hairstyle, clothes and accessories. Full body standing, ' + 'feet visible, plain light grey background, soft even light. One photograph of one person.') + +_VIEW_TEXT = { + 'three_quarter': ('Turn the person in image 1 toward the left side of the picture by 45 degrees. ' + 'Turn the head and body together. Show the front-left three-quarter view, ' + 'with both eyes visible.'), + 'side': ('将图1中的同一个人完全转向画面左边，呈现90度正侧面。鼻尖、胸口、脚尖朝向画面左边，' + '只能看见一只眼睛，近侧肩膀遮住远侧肩膀。'), + 'back': ('Turn the person in image 1 all the way around. Show the back of the head and body. ' + 'The face is hidden. Show the back of the clothes, without front buttons or a shirt opening.'), + 'face': 'front-facing face close-up, framed from shoulders upward, centered head, eyes looking at the camera, neutral expression', +} + +_HAS_GLASSES = re.compile(r'\b(?:eye|sun)?glasses\b|眼镜|太阳镜', re.IGNORECASE) +_NO_GLASSES = re.compile( + r'\b(?:no|without|not wearing|does not wear)\s+(?:[\w-]+\s+){0,3}(?:eye|sun)?glasses\b' + r'|(?:不戴|不佩戴|未戴|没有|无)[^，。；,.;]{0,8}(?:眼镜|太阳镜)', re.IGNORECASE) + + +def _eyewear(card): + description = card.get('description') or '' + if _HAS_GLASSES.search(description) and not _NO_GLASSES.search(description): + return 'wearing the same glasses' + if card.get('photo') and not _NO_GLASSES.search(description): + # The front comes from a user's photo and the text is silent: the truth is in the picture, + # and naming glasses either way makes the model add or drop them. + return 'the same face exactly as in image 1, with nothing added to it or removed from it' + return 'bare eyes, no glasses' + + +def size(view): + return SIZES[view] + + +def uses_front(view): + """Whether the view is an edit of the card's front picture (false for the front itself and for scenes).""" + return view in _VIEW_TEXT + + +def _who(card): + return (card.get('description') or card['name']).strip().rstrip('.') + + +def view_prompt(card, view): + """The prompt for one view of a card (a dict with name, description, kind).""" + who = _who(card) + if card.get('kind') == 'setting': + return ('%s. A wide view of the empty place, no people, even daylight, natural colours, sharp focus, ' + 'no text.' % who) + if view == 'front': + return ('%s. The person stands alone, the whole body from head to feet, facing the camera, neutral pose ' + 'with the arms relaxed, %s.' % (who, BACKDROP)) + if view not in _VIEW_TEXT: + raise ValueError('unknown view') + eyewear = _eyewear(card) + if view == 'face': + return ('Same person as image 1: same face, hair, clothes and body; plain light grey studio backdrop; ' + '%s. Treat image 1 as the identity and wardrobe reference. Photograph this exact person ' + 'from the requested camera angle. The head, shoulders, torso and feet must all follow ' + 'that angle; do not keep a frontal head on a rotated body. Preserve the facial proportions, ' + 'age, hairline, hairstyle, %s, clothing details and accessories from image 1. Keep even ' + 'soft studio light. Show only one person in one photograph, no text, no collage.' + % (_VIEW_TEXT[view], eyewear)) + if view == 'side': + return (_VIEW_TEXT[view] + eyewear + '.保持同一张脸、发型、衣服和配件。全身站立，双脚入画，' + '浅灰纯色背景，均匀柔光，只画一个人的一张照片。') + return '%s %s.%s' % (_VIEW_TEXT[view], eyewear, SAME) + + +def photo_prompt(): + """Edit a user's photo (image 1) into the clean front picture of a card.""" + return ('Image 1 shows the person. Make it a clean character reference photo: the whole body from head to ' + 'feet. Turn the head and body to face us. Neutral pose with the arms relaxed, %s. Keep exactly ' + 'the same person, face, hair and clothes as in image 1. Preserve all existing accessories; ' + 'add no new accessories.' % BACKDROP) diff --git a/freevideo_engine/image_gen/compose.py b/freevideo_engine/image_gen/compose.py new file mode 100644 index 00000000..f0901ca8 --- /dev/null +++ b/freevideo_engine/image_gen/compose.py @@ -0,0 +1,157 @@ +"""First frames with placed people: a scene, some character pictures and where each one stands. + +With a scene picture the people are cut out of their plain backdrops and pasted at the asked places; +the result is a rough layout, and the image model renders it as one clean photograph (the layout is +passed as a reference, so the places and sizes carry over). With only a scene description the same +places are written into the prompt. Either way the answer is a normal generate request. +""" +from pathlib import Path +import uuid + +from PIL import Image + +MAX_PEOPLE = 3 # the model takes 4 pictures: the people (or none) plus the layout or the scene +TOLERANCE = 9 + + +def cut_out(picture, step=TOLERANCE): + """The person on their plain backdrop as an RGBA picture. + + The backdrop is what the edge reaches by changing less than `step` between neighbouring pixels, + so a vignette or gradient backdrop goes as well; it is found on a small copy and scaled back. + """ + from collections import deque + rgb = picture.convert('RGB') + scale = min(1., 384 / max(rgb.size)) + small = rgb.resize((max(2, round(rgb.width * scale)), max(2, round(rgb.height * scale))), Image.BILINEAR) + width, height = small.size + pixels = list(small.getdata()) + gone = bytearray(width * height) + queue = deque() + for x in range(width): + queue.extend((x, width * (height - 1) + x)) + for y in range(height): + queue.extend((width * y, width * y + width - 1)) + for index in queue: + gone[index] = 1 + while queue: + index = queue.popleft() + r, g, b = pixels[index] + x = index % width + for other in (index - 1 if x else -1, index + 1 if x < width - 1 else -1, index - width, index + width): + if 0 <= other < len(pixels) and not gone[other]: + r2, g2, b2 = pixels[other] + if abs(r - r2) + abs(g - g2) + abs(b - b2) < step: + gone[other] = 1 + queue.append(other) + mask = Image.new('L', small.size) + mask.putdata([0 if flag else 255 for flag in gone]) + mask = mask.resize(rgb.size, Image.BILINEAR).point(lambda v: 255 if v > 127 else 0) + result = rgb.convert('RGBA') + result.putalpha(mask) + box = mask.getbbox() + return result.crop(box) if box else result + + +def collage(scene, people, width, height): + """The scene scaled to cover the canvas with each person pasted: people = [(picture, x, height, feet)] in 0-1.""" + canvas = scene.convert('RGB') + scale = max(width / canvas.width, height / canvas.height) + canvas = canvas.resize((max(width, round(canvas.width * scale)), max(height, round(canvas.height * scale)))) + left, top = (canvas.width - width) // 2, (canvas.height - height) // 2 + canvas = canvas.crop((left, top, left + width, top + height)) + for picture, x, tall, feet in people: + person = cut_out(picture) + size = max(1, round(tall * height)) + person = person.resize((max(1, round(person.width * size / person.height)), size)) + at = (round(x * width - person.width / 2), round(feet * height - person.height)) + canvas.paste(person, at, person) + return canvas + + +def _check(people): + if not isinstance(people, list) or not 1 <= len(people) <= MAX_PEOPLE: + raise ValueError('invalid_request') + rows = [] + for person in people: + if not isinstance(person, dict) or not isinstance(person.get('image'), str): + raise ValueError('invalid_media') + x = person.get('x', .5) + tall = person.get('height', .62) + feet = person.get('feet', .9) + if any(type(v) not in (int, float) or not 0 <= v <= 1 for v in (x, tall, feet)) or tall < .05: + raise ValueError('invalid_request') + rows.append(dict(image=person['image'], x=float(x), height=float(tall), feet=float(feet), + name=str(person.get('name') or '')[:40])) + return rows + + +def _where(x): + return 'on the left' if x < .38 else 'on the right' if x > .62 else 'in the middle' + + +def _positions(people): + """Left-to-right sentence; ordering words place the people better than percentages do.""" + order = sorted(range(len(people)), key=lambda i: people[i]['x']) + if len(people) == 1: + return 'The person shown in image 1 stands %s of the frame, full body, feet on the floor.' % _where(people[0]['x']) + words = ['on the left', 'in the middle', 'on the right'] if len(people) == 3 else ['on the left', 'on the right'] + rows = ['%s stands the person shown in image %d' % (word, i + 1) for word, i in zip(words, order)] + rows[0] = rows[0][0].upper() + rows[0][1:] + return '; '.join(rows) + '. Show everyone head to toe, standing apart with empty floor between them, with natural shadows.' + + +def build(value, input_directory): + """A generate request for a first frame. value: width, height, people=[{image, x, height, feet, name}], + and the scene as scene_image (a file name) or scene (text). Pictures are file names in the input directory + (or freevideo-image/); the layout picture, if any, is written to freevideo-image/.""" + if not isinstance(value, dict): + raise ValueError('invalid_request') + people = _check(value.get('people')) + width, height = value.get('width'), value.get('height') + if any(type(v) is not int or not 256 <= v <= 2048 for v in (width, height)): + raise ValueError('invalid_dimensions') + width, height = width // 32 * 32, height // 32 * 32 + text = value.get('scene') if isinstance(value.get('scene'), str) else '' + directory = Path(input_directory).resolve() + + def load(name): + folder, _, leaf = name.rpartition('/') + if not leaf or folder not in ('', 'freevideo-image') or any(c in name for c in ('\\', ':', '\0')): + raise ValueError('invalid_media') + path = (directory / folder / leaf) if folder else directory / leaf + if path.is_symlink() or not path.is_file() or path.resolve().parent != (directory / folder if folder else directory): + raise ValueError('invalid_media') + try: + with Image.open(path) as image: + return image.convert('RGBA') + except OSError as error: + raise ValueError('invalid_media') from error + + names = [person['image'] for person in people] + scene_name = value.get('scene_image') + intro = ' '.join('Image %d is the character %s.' % (i + 1, p['name'] or str(i + 1)) for i, p in enumerate(people)) + if isinstance(scene_name, str) and scene_name: + layout = collage(load(scene_name), [(load(p['image']), p['x'], p['height'], p['feet']) for p in people], + width, height) + folder = directory / 'freevideo-image' + folder.mkdir(exist_ok=True) + leaf = 'layout-' + uuid.uuid4().hex + '.png' + layout.save(folder / leaf) + names.append('freevideo-image/' + leaf) + n = len(people) + n = len(people) + prompt = ('%s Image %d is a rough layout with the characters pasted on the scene. Render image %d as one clean ' + 'photograph of the same place with the same people. Keep the scene, keep everyone at exactly the same ' + 'position and size, keep each face, hairstyle and clothes as in their own image. Remove every cut-out ' + 'edge and white outline. Relight everyone to the lighting of the room, ground them with soft contact ' + 'shadows on the floor and slight floor reflections so they stand in the space instead of floating.' + % (intro, n + 1, n + 1)) + else: + if not text.strip(): + raise ValueError('invalid_prompt') + prompt = '%s %s Keep each face and outfit exactly as in their own image. %s' % ( + intro, _positions(people), 'The scene: ' + text.strip().rstrip('.') + '.') + if len(prompt) > 2000: + raise ValueError('invalid_prompt') + return dict(prompt=prompt, width=width, height=height, seed=value.get('seed'), images=names) diff --git a/freevideo_engine/image_gen/image_models.json b/freevideo_engine/image_gen/image_models.json new file mode 100644 index 00000000..a394c493 --- /dev/null +++ b/freevideo_engine/image_gen/image_models.json @@ -0,0 +1,19 @@ +{ + "format": 1, + "model_id": "qwen-image-2.1-turbo", + "family": "image", + "repo": "Comfy-Org/Qwen-Image-2.1", + "revision": "main", + "groups": { + "image": [ + {"file": "diffusion_models/qwen_image_2.1_turbo_int8_convrot.safetensors", "component": "dit", "bytes": 7256787829, "sha256": "8144a0182964aa8403ad6dbc011db837304167d235161b9aabaa8532e967314c", "sources": [{"id": "official", "url": "https://huggingface.co/Comfy-Org/Qwen-Image-2.1/resolve/main/diffusion_models/qwen_image_2.1_turbo_int8_convrot.safetensors"}, {"id": "modelscope", "url": "https://www.modelscope.cn/models/Comfy-Org/Qwen-Image-2.1/resolve/master/diffusion_models/qwen_image_2.1_turbo_int8_convrot.safetensors"}]}, + {"file": "text_encoders/qwen3vl_8b_int8_convrot.safetensors", "component": "text_encoder", "bytes": 9350798360, "sha256": "8bfd0f6e12abf2d2d697ecc888e5e90b0d6741d6708f05799f53afa560452e8f", "sources": [{"id": "official", "url": "https://huggingface.co/Comfy-Org/Qwen-Image-2.1/resolve/main/text_encoders/qwen3vl_8b_int8_convrot.safetensors"}, {"id": "modelscope", "url": "https://www.modelscope.cn/models/Comfy-Org/Qwen-Image-2.1/resolve/master/text_encoders/qwen3vl_8b_int8_convrot.safetensors"}]}, + {"file": "vae/qwen_image_2.1_vae_bf16.safetensors", "component": "vae", "bytes": 675509688, "sha256": "bb21f7473051e1ac368515dd3f2e15cd44d7a11748ee8823e1ddca3e4876b7c9", "sources": [{"id": "official", "url": "https://huggingface.co/Comfy-Org/Qwen-Image-2.1/resolve/main/vae/qwen_image_2.1_vae_bf16.safetensors"}, {"id": "modelscope", "url": "https://www.modelscope.cn/models/Comfy-Org/Qwen-Image-2.1/resolve/master/vae/qwen_image_2.1_vae_bf16.safetensors"}]}, + {"file": "text_encoder/config.json", "component": "text_encoder_config", "bytes": 1542, "sha256": "89050cba6b66c4f68f52305b54fc37efe28d95232246b94aa8dc0cfa36f3b734", "sources": [{"id": "official", "url": "https://huggingface.co/Qwen/Qwen-Image-2.1-Turbo/resolve/main/text_encoder/config.json"}, {"id": "modelscope", "url": "https://www.modelscope.cn/models/Qwen/Qwen-Image-2.1-Turbo/resolve/master/text_encoder/config.json"}]}, + {"file": "processor/tokenizer.json", "component": "processor", "bytes": 11422650, "sha256": "be75606093db2094d7cd20f3c2f385c212750648bd6ea4fb2bf507a6a4c55506", "sources": [{"id": "official", "url": "https://huggingface.co/Qwen/Qwen-Image-2.1-Turbo/resolve/main/processor/tokenizer.json"}, {"id": "modelscope", "url": "https://www.modelscope.cn/models/Qwen/Qwen-Image-2.1-Turbo/resolve/master/processor/tokenizer.json"}]}, + {"file": "processor/chat_template.jinja", "component": "processor", "bytes": 5292, "sha256": "3636d0f0bd6bef02654cdffdc447b79cb2cef8ab02cc75267345946291a489e4", "sources": [{"id": "official", "url": "https://huggingface.co/Qwen/Qwen-Image-2.1-Turbo/resolve/main/processor/chat_template.jinja"}, {"id": "modelscope", "url": "https://www.modelscope.cn/models/Qwen/Qwen-Image-2.1-Turbo/resolve/master/processor/chat_template.jinja"}]}, + {"file": "processor/processor_config.json", "component": "processor", "bytes": 1300, "sha256": "14932921ca485d458a04dafd8069fbb0a4505622a48208d19ed247115801385b", "sources": [{"id": "official", "url": "https://huggingface.co/Qwen/Qwen-Image-2.1-Turbo/resolve/main/processor/processor_config.json"}, {"id": "modelscope", "url": "https://www.modelscope.cn/models/Qwen/Qwen-Image-2.1-Turbo/resolve/master/processor/processor_config.json"}]}, + {"file": "processor/tokenizer_config.json", "component": "processor", "bytes": 705, "sha256": "fd42a0b35a8d13cd473e88e310b50cbf1bccaa9eb60ed0fd873af5f8f0c59307", "sources": [{"id": "official", "url": "https://huggingface.co/Qwen/Qwen-Image-2.1-Turbo/resolve/main/processor/tokenizer_config.json"}, {"id": "modelscope", "url": "https://www.modelscope.cn/models/Qwen/Qwen-Image-2.1-Turbo/resolve/master/processor/tokenizer_config.json"}]} + ] + } +} diff --git a/freevideo_engine/image_gen/image_pack.py b/freevideo_engine/image_gen/image_pack.py new file mode 100644 index 00000000..b0618664 --- /dev/null +++ b/freevideo_engine/image_gen/image_pack.py @@ -0,0 +1,651 @@ +"""Optional, hash-pinned image weights under /image_gen/. + +Downloads use the engine's source ranking, proxy policy and resumable transfers. +Status reads verification stamps without hashing gigabytes. Repair hashes every +file and retains mismatches; removal touches only unchanged files we created. +Installer.start('install', background=True) yields to the existing engine lease +or an optional busy callback. The CLI writes atomic, read-only status snapshots. +""" +from collections import deque +from contextlib import ExitStack, contextmanager +import json +from pathlib import Path, PurePosixPath +import re +import signal +import sys +import threading +import time +from urllib.parse import urlsplit + +from .. import disk_space, model_repair, network, provision +from ..locking import runtime_lock +from ..monitoring import save + +LEDGER = 'download-verified.json' +OWNERSHIP = 'owned-files.json' +RESERVE = 256 * 2**20 +FAMILY = 'image-models' +BACKGROUND_RESERVE = 2 * 2**30 +GENERATION_POLL_SECONDS = .25 +GENERATION_RESUME_SECONDS = 10. +SPEED_WINDOW_SECONDS = 10. +STATUS_INTERVAL_SECONDS = .25 +OFFLINE_TRANSFER_REASONS = frozenset(('curl-5', 'curl-6', 'curl-7')) + + +class PackError(Exception): + def __init__(self, reason, **details): + super().__init__(reason) + self.reason = reason + self.details = details + + +class _GenerationPause(Exception): + """Unwind the transfer so curl is stopped before waiting for generation.""" + + +class _Background: + def __init__(self, installer, stop): + self.installer, self.stop = installer, stop + self.resume_at = None + self.paused = False + + def generating(self): + if self.installer.generation_busy and self.installer.generation_busy(): + return True + # Probe the existing lease exclusively, without borrowing an inherited + # descriptor. This also detects shared device work. Release immediately: + # downloads never hold engine.lock or delay a generation acquiring it. + # A holder is the only cross-process signal; non-generation holders also + # cause a conservative pause. In-process callers can supply busy(). + path = (Path(self.installer.engine_root) / 'engine.lock' + if self.installer.engine_root is not None else None) + try: + with runtime_lock(path, inherit=False): + return False + except BlockingIOError: + return True + + def check(self): + if self.stop.is_set(): + raise InterruptedError('cancelled') + if self.generating(): + self.paused = True + self.resume_at = None + raise _GenerationPause() + + def wait(self): + while True: + try: + self.check() + except _GenerationPause: + pass + else: + if self.paused and self.resume_at is None: + self.resume_at = time.monotonic() + GENERATION_RESUME_SECONDS + if not self.paused or time.monotonic() >= self.resume_at: + self.resume_at = None + self.paused = False + self.installer._set(state='running', reason=None) + return + # Reached only after the active transfer has been reaped, or at an + # in-process checkpoint, so paused really means no active download. + self.installer._set(state='paused_for_generation', reason='busy_video') + self.stop.wait(GENERATION_POLL_SECONDS) + + +def catalog(): + spec = json.loads(Path(__file__).with_name('image_models.json').read_text(encoding='utf-8')) + validate_catalog(spec) + return spec + + +def _relative(name): + if (not isinstance(name, str) or not name or '\\' in name or ':' in name + or name.startswith('/') or any(p in ('', '.', '..') for p in name.split('/'))): + raise PackError('invalid_catalog') + return PurePosixPath(name) + + +def validate_catalog(spec): + try: + if (spec['format'] != 1 or spec['family'] != 'image' or not spec['model_id'] + or not isinstance(spec['repo'], str) or not spec['repo'] + or not isinstance(spec['revision'], str) or not spec['revision']): + raise PackError('invalid_catalog') + found = spec['groups']['image'] + if not found or len({r['file'] for r in found}) != len(found): + raise PackError('invalid_catalog') + for row in found: + _relative(row['file']) + if (type(row['bytes']) is not int or row['bytes'] <= 0 + or not re.fullmatch('[0-9a-f]{64}', row['sha256']) + or row['component'] not in ('dit', 'text_encoder', 'vae', 'text_encoder_config', 'processor')): + raise PackError('invalid_catalog') + if {s['id'] for s in row['sources']} != {'official', 'modelscope'} or len(row['sources']) != 2: + raise PackError('invalid_catalog') + for source in row['sources']: + url = urlsplit(source['url']) + if url.scheme not in ('https', 'http') or not url.netloc or url.username or url.password: + raise PackError('invalid_catalog') + except (KeyError, TypeError, ValueError) as error: + raise PackError('invalid_catalog') from error + + +def rows(spec=None): + spec = spec if spec is not None else catalog() + validate_catalog(spec) + return [dict(row, family='image', model_id=spec['model_id'], repo=spec['repo'], revision=spec['revision']) + for row in spec['groups']['image']] + + +def base(root=None): + if root is None: + from ..paths import installed_model_root + root = installed_model_root() + return Path(root).expanduser().absolute() / 'image_gen' + + +def _path(folder, name): + """Refuse links at the pack root or below it, including partials/receipts.""" + relative = _relative(name) + path = folder + if path.is_symlink(): + raise PackError('unsafe_path') + for part in relative.parts: + path /= part + if path.is_symlink(): + raise PackError('unsafe_path') + return path + + +def target(row, root=None): + return _path(base(root), row['file']) + + +def _read(folder, name): + try: + value = json.loads(_path(folder, name).read_text(encoding='utf-8')) + return value if isinstance(value, dict) else {} + except (OSError, ValueError): + return {} + + +def _intact(path, row, ledger): + try: + prior = ledger.get(str(path), {}) + return path.is_file() and path.stat().st_size == row['bytes'] and isinstance(prior, dict) and all( + prior.get(k) == v for k, v in provision.file_identity(path, row).items()) + except OSError: + return False + + +def _owned_intact(path, stamp): + return path.is_file() and isinstance(stamp, dict) and all( + stamp.get(k) == v for k, v in provision.file_identity(path, stamp).items()) + + +def _notice(): + return (Path(__file__).parent / 'licenses' / 'NOTICE.txt').read_bytes() + + +def _notice_ready(path): + return path.is_file() and path.stat().st_size == len(_notice()) and path.read_bytes() == _notice() + + +def status(root=None): + """Missing files and unique retained bytes; no network or full weight reads.""" + folder = base(root) + ledger = _read(folder, LEDGER) + files, missing, downloaded, verified_bytes = [], [], 0, 0 + for row in rows(): + path = target(row, root) + intact = _intact(path, row, ledger) + size = path.stat().st_size if path.is_file() else 0 + part = _path(folder, row['file'] + '.partial') + partial = part.stat().st_size if part.is_file() else 0 + retained = row['bytes'] if intact else min(row['bytes'], max(size, partial)) + downloaded += retained + verified_bytes += row['bytes'] if intact else 0 + if not intact: + missing.append(row['file']) + files.append(dict(file=row['file'], bytes=row['bytes'], ready=intact, downloaded_bytes=retained)) + notice = _path(folder, 'NOTICE.txt') + notice_ready = _notice_ready(notice) + total = sum(r['bytes'] for r in files) + available = not missing and notice_ready + return dict(installable=True, ready=available, total_bytes=total, verified_bytes=verified_bytes, + downloaded_bytes=downloaded, bytes=sum(r['bytes'] for r in files if not r['ready']), + remaining_bytes=sum(r['bytes'] - r['downloaded_bytes'] for r in files if not r['ready']), + missing=missing, notice_ready=notice_ready, files=files, + state='ready' if available else 'partial' if downloaded else 'missing') + + +def ready(root=None): + return status(root)['ready'] + + +@contextmanager +def _lease(root, engine_root=None, destructive=False): + if root is None and engine_root is None: + from ..paths import data_root + engine_root = data_root() + folder = base(root) + _path(folder, 'pack.lock') + _path(folder, 'logs/model-mismatches.json') + with ExitStack() as stack: + if engine_root is not None: + engine = Path(engine_root) + stack.enter_context(runtime_lock(engine / 'setup.lock', inherit=False)) + if destructive: + stack.enter_context(runtime_lock(engine / 'engine.lock', inherit=False)) + stack.enter_context(runtime_lock(engine / 'launcher/host-setup.lock', inherit=False)) + stack.enter_context(runtime_lock(folder / 'pack.lock', inherit=False)) + yield folder + + +def verify(root=None, progress=None, check=None): + """Explicit full integrity check; bad stamps are discarded, files are kept.""" + with _lease(root) as folder: + _verify(folder, root, progress or (lambda v: None), check or (lambda: None), full=True) + return status(root) + + +def _verify(folder, root, progress, check, full=False): + ledger = _read(folder, LEDGER) + todo = [] + total = sum(r['bytes'] for r in rows()) + done = 0 + for row in rows(): + check() + path = target(row, root) + def checked(current, size): + check() + progress(dict(phase='verify', done=done + current, total=total)) + before = model_repair.identity(path) + if not provision.verified(path, row, ledger, mode='full' if full else 'auto', progress=checked): + ledger.pop(str(path), None) + if path.exists(): + model_repair.record(folder, path, row, before) + todo.append(row) + done += row['bytes'] + save(_path(folder, LEDGER), ledger) + return ledger, todo + + +def _download_plan(engine_root=None): + from ..adaln_assets import _download_plan as installed + plan = installed(engine_root) + plan = dict(plan, sources={k: [dict(r) for r in v] for k, v in plan.get('sources', {}).items()}) + # Older installations measured these same hosts for video weights. Keep + # that measured ranking until an image-specific probe is available. + if FAMILY not in plan['sources']: + inherited = plan['sources'].get('vdn-models', []) + plan['sources'][FAMILY] = [dict(r) for r in inherited if r['id'] in ('official', 'modelscope')] + known = {r['id'] for r in plan['sources'][FAMILY]} + plan['sources'][FAMILY] += [dict(id=s) for s in ('official', 'modelscope') if s not in known] + return plan + + +def _own(folder, owned, path, row=None): + owned[path.relative_to(folder).as_posix()] = provision.file_identity(path, row or {}) + save(_path(folder, OWNERSHIP), owned) + + +def _transfer(plan, row, path, progress, check, background=None): + from ..download_settings import current, SwitchRequested + failures = [] + if background is not None: + plan = dict(plan) + prior_event = plan.get('event_callback') + def event(value): + if value.get('action') in ('fallback', 'route-retry'): + failures.append((value.get('reason'), value.get('http_status'))) + if prior_event: + prior_event(value) + plan['event_callback'] = event + while True: + if background is not None: + background.wait() + settings = current(plan) + # An explicit choice supersedes affinity to the old partial's source. + plan['manual_source_switch'] = settings['source'] in ('official', 'modelscope') + def resource_check(): + check() + latest = current(plan) + if any(latest[k] != settings[k] for k in ('source', 'proxy_mode')): + raise SwitchRequested(latest) + plan['resource_check'] = resource_check + try: + network.download(network.model_urls(plan, row), path, row['sha256'], progress, + network=plan, size=row['bytes'], category=FAMILY, + keep_partial=True, stall_seconds=30) + return + except SwitchRequested: + if background is None: + check() # Resume with the newly selected policy. + except _GenerationPause: + # network.download's finally terminates/reaps curl, retaining the + # partial and source affinity. The next attempt uses its Range. + if background is None: + raise + except ConnectionError as error: + if background is None: + raise + raise PackError('offline') from error + except RuntimeError as error: + if background is None: + raise + offline = failures and all(reason in OFFLINE_TRANSFER_REASONS and + status in (None, '', '000') for reason, status in failures) + raise PackError('offline' if offline else 'source_failed') from error + + +def install(root=None, progress=None, check=None, *, engine_root=None, repair=False, _background=None): + """Install only missing weights. Repair also hashes and retains bad originals.""" + progress = progress or (lambda v: None) + check = check or (lambda: None) + transfer_check = _background.check if _background is not None else check + try: + with _lease(root, engine_root, destructive=repair) as folder: + ledger, todo = _verify(folder, root, progress, check, full=repair) + total = sum(r['bytes'] for r in todo) + retained = 0 + for row in todo: + partial = _path(folder, row['file'] + '.partial') + if partial.is_file() and partial.stat().st_size < row['bytes']: + retained += partial.stat().st_size + reserve = BACKGROUND_RESERVE if _background is not None else RESERVE + needed = total - retained + reserve + free = disk_space.free_bytes(folder) + if total and free < needed: + raise PackError('disk_space', needed_bytes=needed, free_bytes=free) + plan = _download_plan(engine_root) + plan.update(quiet=True, resource_check=check) + owned = _read(folder, OWNERSHIP) + notice = _path(folder, 'NOTICE.txt') + if notice.exists() and not _notice_ready(notice): + # Preserve an unexpected text rather than overwriting it. + model_repair.rename_noreplace(notice, notice.with_name('NOTICE.txt.replaced-' + str(time.time_ns()))) + if not notice.exists(): + notice.write_bytes(_notice()) + _own(folder, owned, notice) + done = 0 + pending_prefixes = 0 + if _background is not None: + done = sum(r['bytes'] for r in rows()) - total + total += done + pending_prefixes = sum(min(row['bytes'], _path(folder, row['file'] + '.partial').stat().st_size) + for row in todo if _path(folder, row['file'] + '.partial').is_file()) + progress(dict(phase='download', done=done + pending_prefixes, total=total)) + for row in todo: + check() + path = target(row, root) + if path.exists(): + if not repair: + raise PackError('needs_repair') + approved = model_repair.replacement(path, row, model_repair.read(folder), set()) + if approved is None: + raise PackError('changed') + model_repair.retain(folder, path, approved) + owned.pop(row['file'], None) + save(_path(folder, OWNERSHIP), owned) + partial = _path(folder, row['file'] + '.partial') + if _background is not None and partial.is_file(): + pending_prefixes -= min(row['bytes'], partial.stat().st_size) + affinity = _path(folder, row['file'] + '.partial.source.json') + may_own = {p: not p.exists() or _owned_intact(p, owned.get(p.relative_to(folder).as_posix())) + for p in (partial, affinity)} + def moved(current, size, speed, **kwargs): + progress(dict(phase='download', done=done + current + pending_prefixes, total=total)) + transfer_check() + try: + _transfer(plan, row, path, moved, transfer_check, _background) + finally: + # Completed files and cancellable prefixes have separate identities. + for artifact in (partial, affinity): + if may_own[artifact] and _path(folder, artifact.relative_to(folder).as_posix()).is_file(): + _own(folder, owned, artifact) + ledger[str(path)] = provision.file_identity(path, row) + save(_path(folder, LEDGER), ledger) + _own(folder, owned, path, row) + done += row['bytes'] + progress(dict(phase='download', done=done + pending_prefixes, total=total)) + check() + return status(root) + except (PackError, InterruptedError): + raise + except BlockingIOError as error: + raise PackError('busy') from error + except Exception as error: + raise PackError('download_failed') from error + + +def repair(root=None, progress=None, check=None, *, engine_root=None): + return install(root, progress, check, engine_root=engine_root, repair=True) + + +def removal_plan(root=None): + folder = base(root) + owned = _read(folder, OWNERSHIP) + allowed = {'NOTICE.txt'} + for row in rows(): + allowed.update((row['file'], row['file'] + '.partial', row['file'] + '.partial.source.json')) + for name in allowed: + _path(folder, name) + files, skipped = [], [] + for name, stamp in owned.items(): + if name not in allowed or not isinstance(stamp, dict): + continue + path = _path(folder, name) + if not path.exists(): + continue + if _owned_intact(path, stamp): + files.append(dict(file=name, identity=stamp, bytes=path.stat().st_size)) + else: + skipped.append(name) + return dict(files=files, bytes=sum(f['bytes'] for f in files), skipped=skipped) + + +def remove(root=None, check=None, *, engine_root=None): + """Delete our unchanged downloads/partials/notice; keep imports and repairs.""" + check = check or (lambda: None) + try: + with _lease(root, engine_root, destructive=True) as folder: + plan = removal_plan(root) + owned, ledger = _read(folder, OWNERSHIP), _read(folder, LEDGER) + released, removed = 0, [] + for entry in plan['files']: + check() + path = _path(folder, entry['file']) + stamp = entry['identity'] + if not _owned_intact(path, stamp): + plan['skipped'].append(entry['file']) + continue + path.unlink() + removed.append(entry['file']) + released += entry['bytes'] + owned.pop(entry['file'], None) + ledger.pop(str(path), None) + save(_path(folder, OWNERSHIP), owned) + save(_path(folder, LEDGER), ledger) + return dict(released_bytes=released, removed=removed, skipped=plan['skipped']) + except BlockingIOError as error: + raise PackError('busy') from error + + +class Installer: + """One background action per pack; cancel keeps completed files and prefixes.""" + def __init__(self, root=None, *, engine_root=None, busy=None): + self.root, self.engine_root = root, engine_root + self.generation_busy = busy + self.lock = threading.Lock() + self.stop = threading.Event() + self.thread = None + self.value = dict(state='idle', done=0, total=0, bytes_per_second=0, reason=None) + self.samples = deque() + self.transferred = 0 + self.last_done = None + self.background = False + + @property + def busy(self): + return self.thread is not None and self.thread.is_alive() + + def state(self): + with self.lock: + return dict(self.value, bytes_per_second=self._rate(time.monotonic())) + + def _rate(self, now): + if self.value['state'] != 'running' or not self.samples: + return 0 + cutoff = now - SPEED_WINDOW_SECONDS + while len(self.samples) > 1 and self.samples[1][0] <= cutoff: + self.samples.popleft() + start, amount = self.samples[0] + if start < cutoff: + if len(self.samples) > 1: + end, next_amount = self.samples[1] + amount += (next_amount - amount) * (cutoff - start) / (end - start) + else: + amount = self.transferred + start = cutoff + return max(0, self.transferred - amount) / (now - start) if now > start else 0 + + def _progress(self, value): + value = dict(value) + with self.lock: + if value.get('phase') == 'download': + now, done = time.monotonic(), value['done'] + if self.last_done is not None: + self.transferred += max(0, done - self.last_done) + self.last_done = done + if not self.samples or now > self.samples[-1][0]: + self.samples.append((now, self.transferred)) + self.value.update(bytes_per_second=self._rate(now)) + elif self.background: + # Verification byte counts measure hashing, not downloaded bytes. + value.pop('done', None) + value.pop('total', None) + self._set(**value) + + def _set(self, **value): + with self.lock: + if 'state' in value and value['state'] != self.value['state'] and self.samples: + self.samples.append((time.monotonic(), self.transferred)) + self.value.update(value) + + def start(self, action='install', *, background=False): + if action not in ('install', 'repair', 'remove', 'verify'): + raise ValueError('Unknown image pack action') + if background and action != 'install': + raise ValueError('Background mode is only available for install') + with self.lock: + if self.busy: + return dict(self.value) + self.stop = threading.Event() + self.background = background + self.samples.clear() + self.transferred, self.last_done = 0, None + self.value = dict(state='running', action=action, done=0, total=0, bytes_per_second=0, reason=None) + self.thread = threading.Thread(target=self._run, args=(action, self.stop, background), + name='freevideo-image-pack', daemon=True) + self.thread.start() + return dict(self.value) + + def cancel(self): + self.stop.set() + return self.state() + + def _run(self, action, stop, background): + def check(): + if stop.is_set(): + raise InterruptedError('cancelled') + try: + control = _Background(self, stop) if background else None + if background: + initial = status(self.root) + self._set(done=initial['downloaded_bytes'], total=initial['total_bytes']) + options = dict(root=self.root, check=control.wait if control else check) + if control: + options['_background'] = control + if action != 'verify': + options['engine_root'] = self.engine_root + if action != 'remove': + options['progress'] = self._progress + result = globals()[action](**options) + check() + if background: + self._set(done=result['total_bytes'], total=result['total_bytes']) + self._set(state='done', reason=None, result=result) + except InterruptedError: + self._set(state='cancelled', reason=None) + except PackError as error: + reason = error.reason + if background and reason not in ('disk_space', 'offline', 'source_failed', 'busy_video', 'failed'): + reason = 'source_failed' if reason == 'download_failed' else 'failed' + self._set(state='failed', reason=reason, **error.details) + except Exception: + self._set(state='failed', reason='failed') + + +def main(argv=None): + """Download worker and read-only status reader for installer/launcher callers.""" + import argparse + argv = list(sys.argv[1:] if argv is None else argv) + if argv and argv[0] == '--status': + argv[0] = 'status' + parser = argparse.ArgumentParser(description=__doc__) + commands = parser.add_subparsers(dest='command', required=True) + download = commands.add_parser('download') + download.add_argument('--background', action='store_true') + download.add_argument('--root', type=Path) + download.add_argument('--engine-root', type=Path) + download.add_argument('--status-file', type=Path, required=True) + reader = commands.add_parser('status') + reader.add_argument('--status-file', type=Path, required=True) + args = parser.parse_args(argv) + if args.command == 'status': + try: + sys.stdout.write(args.status_file.read_text(encoding='utf-8')) + return 0 + except OSError: + print('Image pack status is unavailable.', file=sys.stderr) + return 1 + + installer = Installer(args.root, engine_root=args.engine_root) + received = threading.Event() + def cancelled(signum, frame): + received.set() + # A signal can interrupt state() while its lock is held in this thread. + # Set only the cancellation event here; the worker publishes the state. + installer.stop.set() + handlers = {} + try: + for name in ('SIGTERM', 'SIGBREAK', 'SIGINT'): + if hasattr(signal, name): + signum = getattr(signal, name) + handlers[signum] = signal.signal(signum, cancelled) + installer.start('install', background=args.background) + if received.is_set(): + installer.cancel() + while True: + save(args.status_file, dict(installer.state(), updated_at=time.time())) + if not installer.busy: + # The thread may finish between snapshotting and checking busy. + save(args.status_file, dict(installer.state(), updated_at=time.time())) + return 0 if installer.state()['state'] == 'done' and not received.is_set() else 1 + installer.thread.join(STATUS_INTERVAL_SECONDS) + except OSError: + print('Could not update image pack status.', file=sys.stderr) + return 1 + finally: + if installer.busy: + installer.cancel() + installer.thread.join() + for signum, handler in handlers.items(): + signal.signal(signum, handler) + + +if __name__ == '__main__': + raise SystemExit(main()) diff --git a/freevideo_engine/image_gen/image_policy.py b/freevideo_engine/image_gen/image_policy.py new file mode 100644 index 00000000..95c92dcc --- /dev/null +++ b/freevideo_engine/image_gen/image_policy.py @@ -0,0 +1,172 @@ +"""Memory tiers for the image model (Qwen-Image 2.1 Turbo, ConvRot int8 files). + +``policy.choose`` is written for the video model's fixed 50 layers and block +size, so the image model gets its own small planner. It decides, for one +request, how many DiT blocks and text-encoder layers stay on the GPU between +uses, whether the VAE decodes in tiles, and whether the host can hold the +weights pinned. Everything else (slots, prefetch, pinned planes) is the +engine's ``LayerOffloader``. + +The numbers are measurements on a B200 under per-process memory caps, 1024x1024 +text-to-image unless noted. A layer that is not resident is copied to the card +once per step (DiT) or once per request (text encoder), so what the plan +cannot keep resident turns into link traffic; ``link_floor_seconds`` reports +that bound. Nothing here was measured on a consumer card yet. +""" +from dataclasses import dataclass, field + +GiB = 2**30 + + +@dataclass(frozen=True) +class ImageShape: + """Weight sizes of the planned model; measure() replaces these once loaded.""" + dit_blocks: int = 32 + dit_block_bytes: int = round(0.2035 * GiB) + te_layers: int = 36 + te_layer_bytes: int = round(0.18 * GiB) + # DiT + text encoder + embedding table as the engine pins them for full streaming. + host_weights_bytes: int = round(14.14 * GiB) + steps: int = 8 + + @classmethod + def measure(cls, dit_blocks, te_layers, *, total_weight_bytes=None): + """Shape from the loaded modules' real layer lists (anything None keeps the default).""" + def size(layer): + return sum(t.numel() * t.element_size() for t in list(layer.parameters()) + list(layer.buffers())) + fields = {} + if dit_blocks: + fields.update(dit_blocks=len(dit_blocks), dit_block_bytes=max(size(b) for b in dit_blocks)) + if te_layers: + fields.update(te_layers=len(te_layers), te_layer_bytes=max(size(b) for b in te_layers)) + if total_weight_bytes: + fields.update(host_weights_bytes=int(total_weight_bytes)) + return cls(**fields) + + +# Peak allocated bytes of each stage with nothing but the fixed parts and the two +# streaming slots resident (blocks all streamed, VAE tiled), from a 36-request +# B200 sweep over budgets 5-24 GiB, 1024^2 to 2048^2, 0-3 references. The slots +# stay allocated when every block is resident, so there is no credit for that. +# Peaks are linear in megapixels; the streamed-stage figures are within 0.1 GiB. +_ENCODE_BASE = 2.7 * GiB # prompt length, not image size, drives this stage +_DENOISE_BASE, _DENOISE_PER_MP = 1.06 * GiB, 0.554 * GiB +_DECODE_BASE, _DECODE_PER_MP = 2.10 * GiB, 0.15 * GiB +_UNTILED_DECODE_PER_MP = 6.5 * GiB # decoding the whole image at once instead of in tiles +# Reference images, about 1 MP each after resizing. Their tokens are recomputed at +# every step (the pipeline keeps no prefix key/value cache: it would hold about +# 2 GiB per reference on the card, and the extra compute is about 1 s per image on +# a B200). The first reference pays for the vision tower and one image's +# activations; more references add to the encode stage faster than to the others. +_REF_ENCODE_FLOOR = 0.25 * GiB +_REF_ENCODE_PER_MP, _REF_ENCODE_OFFSET = 0.7 * GiB, 0.5 * GiB +_REF_DENOISE_PER_MP = 0.6 * GiB +_REF_DECODE_PER_MP = 0.15 * GiB +# Encoding a reference into latents (VAE on the card, one image at a time, in 256-pixel +# tiles like the reference pipeline): 1.86-2.07 GiB with nothing resident, then the +# resident weights on top; flat in the reference's size. Real-cap runs, budgets 6-12 GiB. +# Untiled it peaked 4 GiB higher, which is why this stage once set the residency. +_REF_VAE_ENCODE_BASE = 2.25 * GiB + +MIN_MARGIN = 1.5 * GiB # allocator fragmentation, stage hand-over, driver overhead +MARGIN_FRACTION = 0.06 +MIN_HOST_WORKING = 4 * GiB # beyond the weights: images, tokenizer, the process itself + + +@dataclass(frozen=True) +class ImagePlan: + dit_resident_layers: int + te_resident_layers: int + vae_tiling: bool + host_mode: str # 'pinned' | 'pageable' | 'insufficient' + feasible: bool + peak_allocated_bytes: int # largest stage + stage_peaks: dict + margin_bytes: int + dit_streamed_bytes_per_step: int + te_streamed_bytes: int + link_floor_seconds: object + reasons: list = field(default_factory=list) + + def to_dict(self): + return dict(self.__dict__) + + +def stage_peaks(shape, dit_resident, te_resident, *, width, height, reference_pixels=0, vae_tiling=True): + """Largest allocation of each stage for a given residency (bytes).""" + mp = width * height / 1e6 + ref_mp = reference_pixels / 1e6 + resident = dit_resident * shape.dit_block_bytes + te_resident * shape.te_layer_bytes + encode = _ENCODE_BASE + resident + denoise = _DENOISE_BASE + _DENOISE_PER_MP * mp + resident + decode = _DECODE_BASE + _DECODE_PER_MP * mp + resident + if not vae_tiling: + decode += _UNTILED_DECODE_PER_MP * mp + if ref_mp: + encode += max(_REF_ENCODE_FLOOR, _REF_ENCODE_PER_MP * ref_mp - _REF_ENCODE_OFFSET) + denoise += _REF_DENOISE_PER_MP * ref_mp + decode += _REF_DECODE_PER_MP * ref_mp + stages = dict(encode=int(encode), denoise=int(denoise), decode=int(decode)) + if ref_mp: + stages['reference_encode'] = int(_REF_VAE_ENCODE_BASE + resident) + return stages + + +def plan_image(*, gpu_budget_bytes, host_available_bytes=None, width=1024, height=1024, + reference_pixels=0, shape=None, link_bytes_per_second=None): + """Choose residency for one request. + + DiT blocks come first: each is read every step (8 times per image), a + text-encoder layer once. ``host_available_bytes`` is what the host can lock + or hold now (``offload.live_pin_capacity``); None means it was not read and + the plan assumes it fits. + """ + shape = shape or ImageShape() + margin = int(max(MIN_MARGIN, MARGIN_FRACTION * gpu_budget_bytes)) + limit = int(gpu_budget_bytes) - margin + reasons = [] + kwargs = dict(width=width, height=height, reference_pixels=reference_pixels) + + def peak(kd, kt, tiling=True): + return max(stage_peaks(shape, kd, kt, vae_tiling=tiling, **kwargs).values()) + + kd = kt = 0 + tiling = True + feasible = peak(0, 0) <= limit + if not feasible: + reasons.append('even full streaming with a tiled VAE needs %.1f GiB beyond the %.1f GiB budget' % ( + (peak(0, 0) + margin) / GiB, gpu_budget_bytes / GiB)) + else: + # Largest DiT residency that fits, then text-encoder layers with what is left. + kd = max(k for k in range(shape.dit_blocks + 1) if peak(k, 0) <= limit) + # Encoder layers are read once per request, blocks eight times: only + # when every block is resident does leftover room go to the encoder. + if kd == shape.dit_blocks: + kt = max(k for k in range(shape.te_layers + 1) if peak(kd, k) <= limit) + # Untiled decoding is faster; take it only when it fits beside the same residency. + if peak(kd, kt, tiling=False) <= limit: + tiling = False + peaks = stage_peaks(shape, kd, kt, vae_tiling=tiling, **kwargs) + + streamed_dit = (shape.dit_blocks - kd) * shape.dit_block_bytes + streamed_te = (shape.te_layers - kt) * shape.te_layer_bytes + floor = None + if link_bytes_per_second: + floor = (streamed_dit * shape.steps + streamed_te) / link_bytes_per_second + + host_mode = 'pinned' + if host_available_bytes is not None: + need = shape.host_weights_bytes + MIN_HOST_WORKING + if host_available_bytes >= need: + host_mode = 'pinned' + elif host_available_bytes >= MIN_HOST_WORKING + shape.host_weights_bytes // 2: + host_mode = 'pageable' + reasons.append('not enough lockable host memory for all weights; the rest is read unpinned') + else: + host_mode = 'insufficient' + reasons.append('host memory cannot hold the weights (%.1f GiB needed)' % (need / GiB)) + return ImagePlan(dit_resident_layers=kd, te_resident_layers=kt, vae_tiling=tiling, + host_mode=host_mode, feasible=feasible and host_mode != 'insufficient', + peak_allocated_bytes=max(peaks.values()), stage_peaks=peaks, margin_bytes=margin, + dit_streamed_bytes_per_step=streamed_dit, te_streamed_bytes=streamed_te, + link_floor_seconds=floor, reasons=reasons) diff --git a/freevideo_engine/image_gen/image_runner.py b/freevideo_engine/image_gen/image_runner.py new file mode 100644 index 00000000..cde4150c --- /dev/null +++ b/freevideo_engine/image_gen/image_runner.py @@ -0,0 +1,192 @@ +"""One process's loaded image model and the per-request placement. + +``QwenImageGenerator`` (qwen_pipeline) knows the model; this module decides +where the parts live. The DiT blocks and the text-encoder layers are frozen +CPU weights streamed through ``StreamedStack`` (the engine's LayerOffloader); +what stays on the card between uses comes from ``plan_image``. The vocabulary +table is looked up on the CPU, the VAE and the vision tower visit the card +only for the stage that needs them. +""" +import time +from pathlib import Path + +import torch + +from . import image_policy +from .image_stream import StreamedStack + +GiB = 2**30 +_DIT_FILE = 'diffusion_models/qwen_image_2.1_turbo_int8_convrot.safetensors' +_TE_FILE = 'text_encoders/qwen3vl_8b_int8_convrot.safetensors' +_VAE_FILE = 'vae/qwen_image_2.1_vae_bf16.safetensors' + + +def _non_block_children(module, skip): + for name, child in module.named_children(): + if name != skip: + yield child + + +class _Placement: + """Stage hooks for QwenImageGenerator: nothing but moves between host and card.""" + + def __init__(self, runner): + self.runner = runner + + def enter(self, stage): + runner = self.runner + if stage == 'text': + visual = getattr(runner.text_model.model, 'visual', None) + if visual is not None: + visual.to(runner.device) + elif stage in ('encode', 'decode'): + runner.vae.to(runner.device) + + def leave(self, stage): + runner = self.runner + if stage == 'text': + visual = getattr(runner.text_model.model, 'visual', None) + if visual is not None: + visual.to('cpu') + self._trim() + elif stage in ('encode', 'decode'): + runner.vae.to('cpu') + self._trim() + elif stage == 'denoise': + self._trim() + + @staticmethod + def _trim(): + # Stage peaks are planned on allocated bytes; give the cached blocks back so the + # next stage's workspace is not fragmented by the previous one's. + if torch.cuda.is_available(): + torch.cuda.empty_cache() + + +class ImageRunner: + """Load the pack once, then ``generate`` repeatedly. + + ``gpu_budget_bytes`` caps what a request may plan against (None reads the + card: free memory plus what this process already holds). Residency and VAE + tiling are chosen per request by ``image_policy.plan_image``. + """ + + def __init__(self, pack_dir, *, device='cuda', gpu_budget_bytes=None, log=None, progress=None): + from .qwen_pipeline import QwenImageGenerator + from .qwen_te_weights import load_text_encoder + from .qwen_weights import load_dit, load_vae + + self.pack_dir = Path(pack_dir) + self.device = torch.device(device) + self.gpu_budget_bytes = gpu_budget_bytes + self.log = log or (lambda *args, **kwargs: None) + self.load_seconds = {} + + def path(directory, name): + nested = self.pack_dir / directory / name + return nested if nested.exists() else self.pack_dir / name + + progress = progress or (lambda stage, fraction: None) + progress('load', 0.) + started = time.monotonic() + self.transformer = load_dit(self.pack_dir / _DIT_FILE, device='cpu') + self.load_seconds['dit'] = time.monotonic() - started + progress('load', 1 / 3) + started = time.monotonic() + self.vae = load_vae(self.pack_dir / _VAE_FILE, device='cpu') + self.load_seconds['vae'] = time.monotonic() - started + progress('load', 2 / 3) + started = time.monotonic() + config_dir = self.pack_dir / 'text_encoder' + processor_dir = self.pack_dir / 'processor' + text = load_text_encoder(self.pack_dir / _TE_FILE, config_path=config_dir / 'config.json', device='cpu') + from transformers import Qwen3VLProcessor + processor = Qwen3VLProcessor.from_pretrained(str(processor_dir), local_files_only=True) + self.load_seconds['text'] = time.monotonic() - started + self.text_model = text + + # Parts that always sit on the card: everything of the DiT except its blocks, and the + # encoder's final norm and rotary table. The 1.2 GB vocabulary stays on the host and is + # looked up there (a token id list is tiny), which keeps it out of the encode peak. + for child in _non_block_children(self.transformer, 'transformer_blocks'): + child.to(self.device) + for parameter in self.transformer.parameters(recurse=False): + parameter.data = parameter.data.to(self.device) + language = text.model.language_model + for name, child in language.named_children(): + if name not in ('layers', 'embed_tokens'): + child.to(self.device) + self._embed_hooks = self._lookup_on_host(language.embed_tokens) + + started = time.monotonic() + with torch.no_grad(): # LayerOffloader refuses to be built with autograd on + self.dit_stack = StreamedStack(self.transformer.transformer_blocks, name='dit', device=self.device) + self.te_stack = StreamedStack(language.layers, name='text_encoder', device=self.device) + self.load_seconds['pin'] = time.monotonic() - started + + self.shape = image_policy.ImageShape.measure( + list(self.transformer.transformer_blocks), list(language.layers), + total_weight_bytes=self.dit_stack.pinned_bytes + self.te_stack.pinned_bytes) + self.generator = QwenImageGenerator( + None, placement=_Placement(self), device=self.device, transformer=self.transformer, vae=self.vae, + text_encoder=text, processor=processor, sample_sigmas=None, vae_tile=True) + from .qwen_pipeline import load_sample_sigmas + self.generator.sample_sigmas = load_sample_sigmas(self.pack_dir / _DIT_FILE) + self.last_plan = None + progress('load', 1.) + + def _lookup_on_host(self, embedding): + device = self.device + + def to_host(module, args): + return (args[0].to('cpu'),) + tuple(args[1:]) + + def to_card(module, args, output): + return output.to(device) + + return [embedding.register_forward_pre_hook(to_host), embedding.register_forward_hook(to_card)] + + def budget(self): + if self.gpu_budget_bytes is not None: + return int(self.gpu_budget_bytes) + free, _ = torch.cuda.mem_get_info(self.device) + return int(free + torch.cuda.memory_reserved(self.device)) + + def plan(self, width, height, reference_images=None): + width, height = width // 32 * 32, height // 32 * 32 + reference_pixels = 0 + resolution = self.generator.reference_resolution + for image in reference_images or []: + scale = (resolution * resolution / (image.width * image.height)) ** 0.5 + reference_pixels += round(image.width * scale) * round(image.height * scale) + try: + from ..offload import live_pin_capacity + host_available = live_pin_capacity() + self.dit_stack.pinned_bytes + self.te_stack.pinned_bytes + except Exception: + host_available = None + return image_policy.plan_image( + gpu_budget_bytes=self.budget(), host_available_bytes=host_available, width=width, height=height, + reference_pixels=reference_pixels, shape=self.shape) + + def generate(self, prompt, *, width, height, seed=None, reference_images=None, progress=None, cancel=None): + plan = self.plan(width, height, reference_images) + self.last_plan = plan + if not plan.feasible: + raise MemoryError('; '.join(plan.reasons) or 'the image model does not fit this machine') + self.dit_stack.set_resident(plan.dit_resident_layers) + self.te_stack.set_resident(plan.te_resident_layers) + self.generator.vae_tile = plan.vae_tiling + return self.generator.generate(prompt, width=width, height=height, seed=seed, + reference_images=reference_images, progress=progress, cancel=cancel) + + def stats(self): + return dict(dit=self.dit_stack.stats(), text_encoder=self.te_stack.stats(), + plan=self.last_plan.to_dict() if self.last_plan else None, load_seconds=self.load_seconds) + + def close(self): + for hook in self._embed_hooks: + hook.remove() + self.generator.close() + self.dit_stack.close() + self.te_stack.close() + self.transformer = self.vae = self.text_model = None diff --git a/freevideo_engine/image_gen/image_stream.py b/freevideo_engine/image_gen/image_stream.py new file mode 100644 index 00000000..d9836a23 --- /dev/null +++ b/freevideo_engine/image_gen/image_stream.py @@ -0,0 +1,57 @@ +"""Layer streaming for the image model's DiT blocks and text-encoder layers. + +The engine's ``LayerOffloader`` already does what the image model needs: +frozen CPU weights packed into pinned planes, two CUDA slots, a copy stream and +one layer of prefetch, with weights rebound only while a layer runs. This +module wraps it for a stack of identical layers and adds the one thing the +video path does not need: choosing, per request, how many of the first layers +stay on the card between uses (``set_resident``). + +A stack must be run front to back every pass (the offloader enforces it): 32 +DiT blocks each denoising step, 36 encoder layers once per prompt. +""" +from ..offload import LayerOffloader, pin_layer_weights + + +class StreamedStack: + """``layers`` hold frozen CPU weights; they are moved to the card on use.""" + + def __init__(self, layers, *, name, device='cuda', pin=True, max_pin_bytes=None, prefetch=True): + self.name = name + self.layers = list(layers) + if not self.layers: + raise ValueError('StreamedStack needs at least one layer') + # Pinning is an optimization: a host that refuses it costs transfer + # speed, not the request (pin_layer_weights keeps the unpinned layers + # on their ordinary storage and the offloader stages them). + self.pinned_bytes = pin_layer_weights(self.layers, max_bytes=max_pin_bytes) if pin else 0 + self.offloader = LayerOffloader(self.layers, device=device, prefetch=prefetch, host_prefetch=False) + self.layer_bytes = max(sum(count * parameter.element_size() + for parameter, _, _, _, count, _ in entries) + for entries in self.offloader.layout) + self.resident = 0 + + @property + def pinned_layers(self): + return sum(self.offloader.pinned_layers) + + def set_resident(self, count): + """Keep the first ``count`` layers' weights on the card after their first use. + + Only at a pass boundary (between requests, or between stages). The + planes fill during the next pass through the normal streaming path, so + the first pass after raising it still pays the copies once. + """ + count = max(0, min(int(count), len(self.layers))) + self.offloader.cache_between_steps(count * self.layer_bytes, max_layers=count) + self.resident = count + return count + + def stats(self): + return dict(name=self.name, layers=len(self.layers), layer_bytes=self.layer_bytes, + pinned_layers=self.pinned_layers, pinned_bytes=self.pinned_bytes, + resident_layers=self.resident, cached_bytes=self.offloader.cached_bytes(), + **self.offloader.stats()) + + def close(self): + self.offloader.close() diff --git a/freevideo_engine/image_gen/qwen_dit.py b/freevideo_engine/image_gen/qwen_dit.py new file mode 100644 index 00000000..51109f93 --- /dev/null +++ b/freevideo_engine/image_gen/qwen_dit.py @@ -0,0 +1,820 @@ +# Copyright 2026 Qwen-Image Team, The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import math +from typing import Any + +import torch +import torch.nn as nn +import torch.nn.functional as F + +from ._model_utils import ModelOutput, RMSNorm, TimestepEmbedding, model_config + +# Ported from the supplied Diffusers Qwen-Image 2.1 implementation. +# Changes: plain nn.Module/config/output helpers, segmented torch SDPA only, +# no PEFT, Diffusers cache/attention dispatch, or flex-attention dependency. +_IMG_TOKENS_PER_SLOT = 4 + + +def attention(query, key, value, attn_mask=None): + """Replaceable attention entry point; tensors use (batch, sequence, heads, dim).""" + return F.scaled_dot_product_attention( + query.transpose(1, 2), key.transpose(1, 2), value.transpose(1, 2), + attn_mask=attn_mask, dropout_p=0.0, + ).transpose(1, 2) + + +Transformer2DModelOutput = ModelOutput + + +class QwenImage21KVLayerCache: + """Per-layer KV cache for text and condition-image prefix tokens. + + Stores K and V projections (post-RoPE) for the prefix extracted during the first denoising step. Tensor format: + ``(batch_size, num_prefix_tokens, num_heads, head_dim)``. + """ + + def __init__(self): + self.k: torch.Tensor | None = None + self.v: torch.Tensor | None = None + + def store(self, k: torch.Tensor, v: torch.Tensor): + self.k = k + self.v = v + + def get(self) -> tuple[torch.Tensor, torch.Tensor]: + if self.k is None: + raise RuntimeError("KV cache has not been populated yet.") + return self.k, self.v + + +class QwenImage21KVCache: + """Container for all transformer blocks' prefix KV caches.""" + + def __init__(self, num_layers: int): + self.layer_caches = [QwenImage21KVLayerCache() for _ in range(num_layers)] + + def get_layer(self, layer_idx: int) -> QwenImage21KVLayerCache: + return self.layer_caches[layer_idx] + + +# Copied from diffusers.models.transformers.transformer_qwenimage.apply_rotary_emb_qwen +def apply_rotary_emb_qwen( + x: torch.Tensor, + freqs_cis: torch.Tensor | tuple[torch.Tensor], + use_real: bool = True, + use_real_unbind_dim: int = -1, +) -> tuple[torch.Tensor, torch.Tensor]: + """ + Apply rotary embeddings to input tensors using the given frequency tensor. This function applies rotary embeddings + to the given query or key 'x' tensors using the provided frequency tensor 'freqs_cis'. The input tensors are + reshaped as complex numbers, and the frequency tensor is reshaped for broadcasting compatibility. The resulting + tensors contain rotary embeddings and are returned as real tensors. + + Args: + x (`torch.Tensor`): + Query or key tensor to apply rotary embeddings. [B, S, H, D] xk (torch.Tensor): Key tensor to apply + freqs_cis (`tuple[torch.Tensor]`): Precomputed frequency tensor for complex exponentials. ([S, D], [S, D],) + + Returns: + tuple[torch.Tensor, torch.Tensor]: tuple of modified query tensor and key tensor with rotary embeddings. + """ + if use_real: + cos, sin = freqs_cis # [S, D] + cos = cos[None, None] + sin = sin[None, None] + cos, sin = cos.to(x.device), sin.to(x.device) + + if use_real_unbind_dim == -1: + # Used for flux, cogvideox, hunyuan-dit + x_real, x_imag = x.reshape(*x.shape[:-1], -1, 2).unbind(-1) # [B, S, H, D//2] + x_rotated = torch.stack([-x_imag, x_real], dim=-1).flatten(3) + elif use_real_unbind_dim == -2: + # Used for Stable Audio, OmniGen, CogView4 and Cosmos + x_real, x_imag = x.reshape(*x.shape[:-1], 2, -1).unbind(-2) # [B, S, H, D//2] + x_rotated = torch.cat([-x_imag, x_real], dim=-1) + else: + raise ValueError(f"`use_real_unbind_dim={use_real_unbind_dim}` but should be -1 or -2.") + + out = (x.float() * cos + x_rotated.float() * sin).to(x.dtype) + + return out + else: + x_rotated = torch.view_as_complex(x.float().reshape(*x.shape[:-1], -1, 2)) + freqs_cis = freqs_cis.unsqueeze(1) + x_out = torch.view_as_real(x_rotated * freqs_cis).flatten(3) + + return x_out.type_as(x) + + +class QwenImage21TemporalTimesteps(nn.Module): + r"""Sinusoidal timestep embedding. `cos` occupies the first half of the channels and `sin` the second.""" + + def __init__(self, timestep_dim: int, max_period: int = 10000, time_factor: float = 1000.0): + super().__init__() + self.timestep_dim = timestep_dim + self.time_factor = time_factor + + half = timestep_dim // 2 + freqs = torch.exp(-math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32, device="cpu") / half) + self.register_buffer("freqs", freqs, persistent=False) + + def forward(self, timestep: torch.Tensor) -> torch.Tensor: + timestep = self.time_factor * timestep.float() + args = timestep[:, None] * self.freqs[None].to(timestep.device) + embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1) + if self.timestep_dim % 2: + embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1) + return embedding.to(timestep.dtype) + + +class QwenImage21TimestepProjEmbeddings(nn.Module): + def __init__(self, embedding_dim: int): + super().__init__() + self.time_proj = QwenImage21TemporalTimesteps(timestep_dim=256) + self.timestep_embedder = TimestepEmbedding( + in_channels=256, time_embed_dim=embedding_dim, sample_proj_bias=False + ) + + def forward(self, timestep: torch.Tensor, hidden_states: torch.Tensor) -> torch.Tensor: + timesteps_proj = self.time_proj(timestep) + return self.timestep_embedder(timesteps_proj.to(dtype=hidden_states.dtype)) + + +class QwenImage21ZeroCenterRMSNorm(nn.Module): + r""" + RMSNorm whose learnable weight is stored zero-centered: the effective scale is `weight + 1`, computed in fp32. + Checkpoints therefore store `scale - 1`. + """ + + def __init__(self, dim: int, eps: float = 1e-5): + super().__init__() + self.weight = nn.Parameter(torch.zeros(dim)) + self.eps = eps + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + input_dtype = hidden_states.dtype + hidden_states = hidden_states.float() + rrms = torch.rsqrt(torch.mean(hidden_states**2, dim=-1, keepdim=True) + self.eps) + return (hidden_states * rrms * (self.weight.float() + 1)).to(input_dtype) + + +class QwenImage21TextProjection(nn.Module): + def __init__(self, context_in_dim: int, hidden_size: int, eps: float = 1e-6): + super().__init__() + self.text_norm = QwenImage21ZeroCenterRMSNorm(context_in_dim, eps=eps) + self.in_layer = nn.Linear(context_in_dim, hidden_size, bias=False) + self.act = nn.GELU(approximate="tanh") + self.out_layer = nn.Linear(hidden_size, hidden_size, bias=False) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + hidden_states = self.text_norm(hidden_states) + hidden_states = self.in_layer(hidden_states) + hidden_states = self.act(hidden_states) + return self.out_layer(hidden_states) + + +class QwenImage21SwiGLUFeedForward(nn.Module): + def __init__(self, hidden_size: int, mlp_hidden_size: int): + super().__init__() + self.proj = nn.Linear(hidden_size, mlp_hidden_size, bias=False) + self.out = nn.Linear(mlp_hidden_size, hidden_size, bias=False) + self.gate_layer = nn.Linear(hidden_size, mlp_hidden_size, bias=False) + self.activation_fn = nn.SiLU() + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + return self.out(self.activation_fn(self.gate_layer(hidden_states)) * self.proj(hidden_states)) + + +class QwenImage21AdaLayerNormContinuous(nn.Module): + r""" + Final adaptive norm. Scale only — this variant emits no shift, so `linear` maps to `embedding_dim` rather than `2 * + embedding_dim`. + """ + + def __init__(self, embedding_dim: int, conditioning_embedding_dim: int, eps: float = 1e-6): + super().__init__() + self.silu = nn.SiLU() + self.linear = nn.Linear(conditioning_embedding_dim, embedding_dim, bias=False) + self.norm = nn.LayerNorm(embedding_dim, eps, elementwise_affine=False, bias=False) + + def forward( + self, + hidden_states: torch.Tensor, + conditioning_embedding: torch.Tensor, + target_token_mask: torch.Tensor | None = None, + ) -> torch.Tensor: + scale = self.linear(self.silu(conditioning_embedding).to(hidden_states.dtype)) + scale = _select_modulation_rows(scale, target_token_mask) + return self.norm(hidden_states) * (1 + scale) + + +def _select_modulation_rows(params: torch.Tensor, target_token_mask: torch.Tensor | None) -> torch.Tensor: + r""" + Broadcast per-sample modulation `params` over the token axis. + + With `causal_condition`, `params` holds `batch_size + 1` rows: rows `[0, batch_size)` come from the real timestep + and the trailing row from `t = 0`. Text and condition-image tokens take the `t = 0` row, target-image tokens take + their own sample's row. + + Args: + params (`torch.Tensor`): `(batch_size, dim)` without `causal_condition`, else `(batch_size + 1, dim)`. + target_token_mask (`torch.Tensor`, *optional*): `(seq_len,)` bool, `True` at target-image positions. `None` + disables the split and every token uses its own sample's row. + """ + if target_token_mask is None: + return params.unsqueeze(1) + real, zero = params[:-1].unsqueeze(1), params[-1:].unsqueeze(0) + return torch.where(target_token_mask.view(1, -1, 1), real, zero) + + +def build_qwenimage21_block_causal_mask(image_ids, encoder_hidden_states_mask, batch_size, device): + """Dense boolean form for inspection/small checks; normal prefill uses segments.""" + positions = torch.arange(image_ids.numel(), device=device) + same_image = (image_ids[:, None] == image_ids[None, :]) & (image_ids[:, None] >= 0) + allowed = (positions[:, None] >= positions[None, :]) | same_image + allowed = allowed[None, None].expand(batch_size, 1, -1, -1) + if encoder_hidden_states_mask is not None: + allowed = allowed & encoder_hidden_states_mask[:, None, None, :].bool() + return allowed + + +def _qwenimage21_prefix_segments(image_ids: torch.Tensor, prefix_len: int) -> list[tuple[int, int, bool]]: + """Split the prefix into `(start, end, is_text)` runs of equal `image_ids`. + + This is the block-causal structure consumed by the segmented SDPA processor. It only depends on + `image_ids` and `prefix_len`, so the model derives it once per forward rather than in every processor call — + `tolist()` is a device sync, and there is one processor call per layer. + """ + prefix_ids = image_ids[:prefix_len].tolist() + segments = [] + start = 0 + for index in range(1, prefix_len + 1): + if index == prefix_len or prefix_ids[index] != prefix_ids[start]: + segments.append((start, index, prefix_ids[start] < 0)) + start = index + return segments + + +def _qwenimage21_prepare_qkv( + attn: "QwenImage21Attention", + hidden_states: torch.Tensor, + rotary_emb: torch.Tensor | None, + layer_cache: QwenImage21KVLayerCache | None, + kv_cache_mode: str | None, + cache_write_slice: slice | None, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, int]: + """Shared QKV projection, norm, RoPE and KV-cache bookkeeping for both processors.""" + query = attn.to_q(hidden_states) + key = attn.to_k(hidden_states) + value = attn.to_v(hidden_states) + + query = query.unflatten(-1, (attn.heads, -1)) + key = key.unflatten(-1, (attn.heads, -1)) + value = value.unflatten(-1, (attn.heads, -1)) + + query = attn.norm_q(query).to(value.dtype) + key = attn.norm_k(key).to(value.dtype) + + if rotary_emb is not None: + query = apply_rotary_emb_qwen(query, rotary_emb, use_real=False) + key = apply_rotary_emb_qwen(key, rotary_emb, use_real=False) + + if layer_cache is not None: + if kv_cache_mode == "extract" and cache_write_slice is not None: + # `clone()`, not `contiguous()`: at batch size 1 the prefix slice already counts as contiguous + # (size-1 dims are ignored), so `contiguous()` returns the same view and the cache would pin the + # whole prefill K/V for every step of the denoising loop. + layer_cache.store( + key[:, cache_write_slice].clone(), + value[:, cache_write_slice].clone(), + ) + elif kv_cache_mode == "cached": + cached_k, cached_v = layer_cache.get() + key = torch.cat([cached_k, key], dim=1) + value = torch.cat([cached_v, value], dim=1) + + seq_len_q = query.shape[1] + return query, key, value, seq_len_q + + +class QwenImage21AttnProcessor: + r""" + Attention processor for Qwen-Image 2.1 that needs neither `flex_attention` nor a compiled model. + + The prefill decomposes the block-causal mask into one attention call per prefix segment plus one for the target + image. The segment boundaries are computed once + per forward by the model and passed in as `segments`. + """ + + def __call__( + self, + attn: "QwenImage21Attention", + hidden_states: torch.Tensor, + attention_mask: Any | None = None, + rotary_emb: torch.Tensor | None = None, + layer_cache: QwenImage21KVLayerCache | None = None, + kv_cache_mode: str | None = None, + cache_write_slice: slice | None = None, + segments: list[tuple[int, int, bool]] | None = None, + key_valid: torch.Tensor | None = None, + ) -> torch.Tensor: + query, key, value, seq_len_q = _qwenimage21_prepare_qkv( + attn, hidden_states, rotary_emb, layer_cache, kv_cache_mode, cache_write_slice + ) + + if segments is None: + # decode: full attention over [cached prefix, target] + hidden_states = attention( + query, + key, + value, + attn_mask=attention_mask, + ) + else: + # prefill: every segment attends to the keys `[0, end)` (everything before it plus its own block); text + # segments additionally get a causal triangle over their own keys; padded text keys are dropped. + prefix_len = segments[-1][1] if segments else 0 + outputs = [] + for start, end, is_text in segments: + seg_mask = None + if is_text: + seg_len = end - start + seg_mask = torch.cat( + [ + torch.ones(seg_len, start, dtype=torch.bool, device=query.device), + torch.tril(torch.ones(seg_len, seg_len, dtype=torch.bool, device=query.device)), + ], + dim=1, + )[None, None] + if key_valid is not None: + seg_key_valid = key_valid[:, None, None, :end] + seg_mask = seg_key_valid if seg_mask is None else (seg_mask & seg_key_valid) + outputs.append( + attention( + query[:, start:end], + key[:, :end], + value[:, :end], + attn_mask=seg_mask, + ) + ) + outputs.append( + attention( + query[:, prefix_len:], + key, + value, + attn_mask=None if key_valid is None else key_valid[:, None, None, :], + ) + ) + hidden_states = torch.cat(outputs, dim=1) + hidden_states = hidden_states[:, :seq_len_q] + hidden_states = hidden_states.flatten(2, 3).type_as(query) + + hidden_states = attn.to_out[0](hidden_states) + return attn.to_out[1](hidden_states) + + +class QwenImage21Attention(torch.nn.Module): + r""" + Attention module for [`QwenImage21TransformerBlock`]. Projection layout matches the legacy + [`~models.attention_processor.Attention`] so Qwen-Image 2.x checkpoints load into it unchanged. + """ + + # Prefill stays segmented; it never builds a full joint attention matrix. + _default_processor_cls = QwenImage21AttnProcessor + _available_processors = [QwenImage21AttnProcessor] + + def __init__(self, dim: int, heads: int, dim_head: int, eps: float = 1e-6, processor: Any | None = None): + super().__init__() + self.heads = heads + self.inner_dim = heads * dim_head + self.to_q = nn.Linear(dim, self.inner_dim, bias=False) + self.to_k = nn.Linear(dim, self.inner_dim, bias=False) + self.to_v = nn.Linear(dim, self.inner_dim, bias=False) + self.to_out = nn.ModuleList([nn.Linear(self.inner_dim, dim, bias=False), nn.Dropout(0.0)]) + self.norm_q = RMSNorm(dim_head, eps=eps) + self.norm_k = RMSNorm(dim_head, eps=eps) + + self.set_processor(processor if processor is not None else self._default_processor_cls()) + + def set_processor(self, processor): + self.processor = processor + + def forward(self, hidden_states: torch.Tensor, **kwargs) -> torch.Tensor: + return self.processor(self, hidden_states, **kwargs) + + +class QwenImage21TransformerBlock(nn.Module): + r""" + Single-stream block. Modulation is not learned per block — the parent model computes one shared `modulation` tensor + and every block slices its own scales and gates out of it. + """ + + def __init__( + self, + dim: int, + num_attention_heads: int, + attention_head_dim: int, + mlp_ratio: int = 3, + eps: float = 1e-6, + ): + super().__init__() + self.img_norm1 = nn.LayerNorm(dim, elementwise_affine=False, eps=eps) + self.attn = QwenImage21Attention(dim=dim, heads=num_attention_heads, dim_head=attention_head_dim, eps=eps) + self.img_norm2 = nn.LayerNorm(dim, elementwise_affine=False, eps=eps) + self.img_mlp = QwenImage21SwiGLUFeedForward(hidden_size=dim, mlp_hidden_size=dim * mlp_ratio) + + def _modulate( + self, + hidden_states: torch.Tensor, + mod_params: torch.Tensor, + target_token_mask: torch.Tensor | None, + ) -> tuple[torch.Tensor, torch.Tensor]: + scale, gate = mod_params.chunk(2, dim=-1) + scale = _select_modulation_rows(scale, target_token_mask) + gate = _select_modulation_rows(gate, target_token_mask) + return hidden_states * (1 + scale), gate + + def forward( + self, + hidden_states: torch.Tensor, + modulation: torch.Tensor, + rotary_emb: torch.Tensor | None = None, + attention_mask: Any | None = None, + target_token_mask: torch.Tensor | None = None, + layer_cache: QwenImage21KVLayerCache | None = None, + kv_cache_mode: str | None = None, + cache_write_slice: slice | None = None, + segments: list[tuple[int, int, bool]] | None = None, + key_valid: torch.Tensor | None = None, + ) -> torch.Tensor: + mod1, mod2 = modulation.chunk(2, dim=-1) + + img_modulated, img_gate1 = self._modulate(self.img_norm1(hidden_states), mod1, target_token_mask) + attn_output = self.attn( + hidden_states=img_modulated, + attention_mask=attention_mask, + rotary_emb=rotary_emb, + layer_cache=layer_cache, + kv_cache_mode=kv_cache_mode, + cache_write_slice=cache_write_slice, + segments=segments, + key_valid=key_valid, + ) + hidden_states = hidden_states + img_gate1.tanh() * attn_output + + img_modulated2, img_gate2 = self._modulate(self.img_norm2(hidden_states), mod2, target_token_mask) + hidden_states = hidden_states + img_gate2.tanh() * self.img_mlp(img_modulated2) + + if hidden_states.dtype == torch.float16: + hidden_states = hidden_states.clip(-65504, 65504) + + return hidden_states + + +class QwenImage21Rope(nn.Module): + r""" + 3-axis (frame, height, width) rotary embedding over the joint text/image sequence. + + Text tokens advance a shared position on all three axes. Every image block freezes the frame axis at the position + reached by the preceding text and lays its tokens out on a height/width grid centred on zero, so a block's spatial + positions do not depend on where it sits in the sequence. + """ + + def __init__(self, theta: int, axes_dim: list[int]): + super().__init__() + self.theta = theta + self.axes_dim = axes_dim + + pos_index = torch.arange(8192, device="cpu") + neg_index = torch.arange(1024, device="cpu").flip(0) * -1 - 1 + self.freqs = [ + torch.cat([self.rope_params(pos_index, dim, theta), self.rope_params(neg_index, dim, theta)], dim=0) + for dim in axes_dim + ] + + def rope_params(self, index: torch.Tensor, dim: int, theta: int = 10000) -> torch.Tensor: + freqs = torch.outer(index, 1.0 / torch.pow(theta, torch.arange(0, dim, 2, device="cpu").to(torch.float32).div(dim))) + return torch.polar(torch.ones_like(freqs), freqs) + + def forward( + self, img_shapes: list[tuple[int, int, int]], image_pad_mask: torch.Tensor, device: torch.device + ) -> torch.Tensor: + self.freqs = [freq.to(device) for freq in self.freqs] + + frame_index, height_index, width_index = [], [], [] + image_height_index, image_width_index = [], [] + cursor, position = 0, 0 + total_len = image_pad_mask.shape[-1] + is_image_token = image_pad_mask.tolist() + + for _, height, width in img_shapes: + block_start = is_image_token.index(True, cursor) + text_len = block_start - cursor + frame_index.extend(range(position, position + text_len)) + position += text_len + + cursor = block_start + height * width + frame_index.extend([position] * (height * width)) + position += max(height, width) + + image_height_index.extend([h for h in range(-(height - height // 2), height // 2) for _ in range(width)]) + image_width_index.extend([w for _ in range(height) for w in range(-(width - width // 2), width // 2)]) + + if cursor < total_len: + frame_index.extend(range(position, position + total_len - cursor)) + + frame_index = torch.tensor(frame_index, dtype=torch.long, device=device) + height_index = frame_index.clone() + width_index = frame_index.clone() + height_index[image_pad_mask] = torch.tensor(image_height_index, dtype=torch.long, device=device) + width_index[image_pad_mask] = torch.tensor(image_width_index, dtype=torch.long, device=device) + + return torch.cat([self.freqs[0][frame_index], self.freqs[1][height_index], self.freqs[2][width_index]], dim=-1) + + +class QwenImageDiT(nn.Module): + r""" + The single-stream Transformer used by Qwen-Image 2.1. + + Text and image latents share one sequence: condition-image tokens are substituted into the text stream at the + positions the vision-language encoder reserved for them, and the target image's tokens are appended. A single + shared `modulation` projection feeds every block, so blocks hold no modulation parameters of their own. + + Two behaviours distinguish 2.1: + + - **Block-causal attention** — attention follows `(q_idx >= kv_idx) or same_image_block`, so the sequence is causal + while each image block stays internally bidirectional. Exact multi-pass SDPA prefill processes each + prefix block separately, then evaluates all target queries together. + - `causal_condition` — text and condition-image tokens are modulated from `t = 0` instead of the sampled timestep, + which also makes their activations timestep-independent and so cacheable across denoising steps. + + Args: + patch_size (`int`, defaults to `1`): + Side length of the latent patch folded into the channel dim. 2.1 consumes latents unpatched. + in_channels (`int`, defaults to `64`): + Latent channels of the input. + out_channels (`int`, *optional*, defaults to `64`): + Latent channels of the output. Falls back to `in_channels`. + num_layers (`int`, defaults to `32`): + Number of single-stream blocks. + attention_head_dim (`int`, defaults to `128`): + Channels per attention head. + num_attention_heads (`int`, defaults to `32`): + Number of attention heads. + context_in_dim (`int`, defaults to `4096`): + Channel dim of `encoder_hidden_states`. + mlp_ratio (`int`, defaults to `3`): + Feed-forward expansion factor. + axes_dims_rope (`tuple[int]`, defaults to `(16, 56, 56)`): + Rotary dims for the frame, height and width axes. + eps (`float`, defaults to `1e-6`): + Epsilon for the norm layers. + causal_condition (`bool`, defaults to `True`): + Modulate text and condition-image tokens from `t = 0`. Required for KV caching. + """ + + @model_config + def __init__( + self, + patch_size: int = 1, + in_channels: int = 64, + out_channels: int | None = 64, + num_layers: int = 32, + attention_head_dim: int = 128, + num_attention_heads: int = 32, + context_in_dim: int = 4096, + mlp_ratio: int = 3, + axes_dims_rope: tuple[int, int, int] = (16, 56, 56), + eps: float = 1e-6, + causal_condition: bool = True, + ): + super().__init__() + self.out_channels = out_channels or in_channels + self.inner_dim = num_attention_heads * attention_head_dim + + self.pos_embed = QwenImage21Rope(theta=10000, axes_dim=list(axes_dims_rope)) + self.time_text_embed = QwenImage21TimestepProjEmbeddings(embedding_dim=self.inner_dim) + self.txt_in = QwenImage21TextProjection(context_in_dim, self.inner_dim, eps=eps) + self.img_in = nn.Linear(in_channels * patch_size * patch_size, self.inner_dim, bias=False) + + # One shared modulation for every block: [mod1.scale, mod1.gate, mod2.scale, mod2.gate]. + self.modulation = nn.Sequential(nn.SiLU(), nn.Linear(self.inner_dim, 4 * self.inner_dim, bias=False)) + + self.transformer_blocks = nn.ModuleList( + [ + QwenImage21TransformerBlock( + dim=self.inner_dim, + num_attention_heads=num_attention_heads, + attention_head_dim=attention_head_dim, + mlp_ratio=mlp_ratio, + eps=eps, + ) + for _ in range(num_layers) + ] + ) + + self.norm_out = QwenImage21AdaLayerNormContinuous(self.inner_dim, self.inner_dim, eps=eps) + self.proj_out = nn.Linear(self.inner_dim, patch_size * patch_size * self.out_channels, bias=False) + + @staticmethod + def build_token_metadata( + image_pad_mask: torch.Tensor, img_shapes: list[tuple[int, int, int]] + ) -> tuple[torch.Tensor, torch.Tensor]: + r""" + Label every token of the joint sequence with the image block it belongs to. + + Block boundaries come from the token counts in `img_shapes`, not from runs of `True` in `image_pad_mask`: two + condition images that happen to sit next to each other with no text between them form one run but must stay + separate blocks, otherwise they would attend to each other bidirectionally. + + Args: + image_pad_mask (`torch.Tensor`): `(seq_len,)` bool, `True` at image-token positions. + img_shapes (`list[tuple[int, int, int]]`): Per-image `(frame, height, width)` in latent tokens, condition + images first and the target image last. + + Returns: + `tuple[torch.Tensor, torch.Tensor]`: `image_ids` `(seq_len,)` with `-1` at text positions and a unique id + per image block, and `target_token_mask` `(seq_len,)` marking the target image's tokens. + """ + image_positions = image_pad_mask.nonzero(as_tuple=True)[0] + block_lengths = [math.prod(shape) for shape in img_shapes] + if sum(block_lengths) != image_positions.numel(): + raise ValueError( + f"img_shapes accounts for {sum(block_lengths)} image tokens but image_pad_mask marks " + f"{image_positions.numel()}." + ) + + image_ids = torch.full_like(image_pad_mask, -1, dtype=torch.long) + block_ids = torch.repeat_interleave( + torch.arange(len(block_lengths), device=image_pad_mask.device), + torch.tensor(block_lengths, device=image_pad_mask.device), + ) + image_ids[image_positions] = block_ids + + target_token_mask = torch.zeros_like(image_pad_mask) + target_token_mask[image_positions[-block_lengths[-1] :]] = True + return image_ids, target_token_mask + + def forward( + self, + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor, + timestep: torch.Tensor, + img_shapes: list[list[tuple[int, int, int]]], + img_mask: torch.Tensor, + encoder_hidden_states_mask: torch.Tensor | None = None, + attention_kwargs: dict[str, Any] | None = None, + kv_cache: QwenImage21KVCache | None = None, + kv_cache_mode: str | None = None, + return_dict: bool = True, + ) -> torch.Tensor | Transformer2DModelOutput: + r""" + Args: + hidden_states (`torch.Tensor` of shape `(batch_size, image_sequence_length, in_channels)`): + Packed latents, condition images first and the target image last. + encoder_hidden_states (`torch.Tensor` of shape `(batch_size, text_sequence_length, context_in_dim)`): + Text embeddings from the vision-language encoder. + timestep (`torch.Tensor`): + Current denoising step, scaled to `[0, 1]`. + img_shapes (`list[list[tuple[int, int, int]]]`): + Per-sample list of `(frame, height, width)` in latent tokens, condition images first and the target + image last. All samples must share a layout. + img_mask (`torch.Tensor` of shape `(batch_size, vlm_sequence_length)`): + `True` at the vision-language encoder's image slots, each standing for a `2x2` group of latent tokens. + encoder_hidden_states_mask (`torch.Tensor`, *optional*): + `(batch_size, text_sequence_length)` bool marking valid text tokens. Padded positions are excluded from + attention. + kv_cache (`QwenImage21KVCache`, *optional*): + Cache container. Pass together with `kv_cache_mode` to enable prefix KV caching. + kv_cache_mode (`str`, *optional*): + `"extract"` to prefill the cache (first denoising step), `"cached"` to decode from it (later steps). + Requires `causal_condition=True`. + attention_kwargs (`dict`, *optional*): + Accepted for call compatibility. Non-unit LoRA scale is unsupported. + return_dict (`bool`, *optional*, defaults to `True`): + Whether or not to return a [`~models.modeling_outputs.Transformer2DModelOutput`] instead of a plain + tuple. + + Returns: + [`~models.modeling_outputs.Transformer2DModelOutput`] or `tuple`: + [`~models.modeling_outputs.Transformer2DModelOutput`] if `return_dict` is True, otherwise a plain + `tuple` whose first element is the denoised latents. + """ + + if attention_kwargs and attention_kwargs.get("scale", 1.0) != 1.0: + raise ValueError("LoRA scaling is not supported by the standalone DiT") + batch_size = hidden_states.shape[0] + if kv_cache is not None and not self.config.causal_condition: + raise ValueError( + "kv_cache requires `causal_condition=True`. The cache is only valid because text and condition-image " + "tokens modulate from t=0, which makes their activations independent of the denoising step." + ) + if kv_cache is not None and kv_cache_mode not in ("extract", "cached"): + raise ValueError( + f"kv_cache_mode must be 'extract' or 'cached' when kv_cache is provided, got {kv_cache_mode!r}." + ) + if kv_cache is None and kv_cache_mode is not None: + raise ValueError(f"kv_cache_mode is {kv_cache_mode!r} but no kv_cache was passed to hold the prefix.") + + hidden_states = self.img_in(hidden_states) + encoder_hidden_states = self.txt_in(encoder_hidden_states) + + # Each vision-language image slot stands for 2x2 latent tokens, so expand those positions four-fold and drop + # the actual latents into them. Samples share a layout, hence the single row. + repeats = torch.where(img_mask, _IMG_TOKENS_PER_SLOT, 1)[0] + image_pad_mask = torch.repeat_interleave(img_mask[0], repeats) + + target_tokens = math.prod(img_shapes[0][-1]) + joint_hidden_states = torch.cat( + [ + encoder_hidden_states, + encoder_hidden_states.new_zeros(batch_size, target_tokens // 4, encoder_hidden_states.shape[2]), + ], + dim=1, + ) + joint_hidden_states = joint_hidden_states.repeat_interleave(repeats, dim=1) + joint_hidden_states[:, image_pad_mask] = hidden_states + + rotary_emb = self.pos_embed(img_shapes[0], image_pad_mask, device=hidden_states.device) + image_ids, target_token_mask = self.build_token_metadata(image_pad_mask, img_shapes[0]) + + timestep = timestep.to(hidden_states.dtype) + if self.config.causal_condition: + # Extra t=0 row; text and condition-image tokens modulate from it. `modulation_mask` selects which row + # each token reads, and is `None` when every token shares the sampled timestep. + timestep = torch.cat([timestep, timestep.new_zeros(1)], dim=0) + modulation_mask = target_token_mask + else: + modulation_mask = None + temb = self.time_text_embed(timestep, hidden_states) + modulation = self.modulation(temb) + + # Right-padded prompt positions must never be attended to, on any path. Text positions of the joint sequence + # line up, in order, with the non-image positions of the vision-language sequence — the two are interleaved, + # so the mask cannot be sliced off as a prefix. + joint_key_valid = None + if encoder_hidden_states_mask is not None: + joint_key_valid = torch.ones( + batch_size, image_pad_mask.shape[0], dtype=torch.bool, device=hidden_states.device + ) + text_positions = (~image_pad_mask).nonzero(as_tuple=True)[0] + vlm_text_positions = ~img_mask[0][: encoder_hidden_states_mask.shape[1]] + joint_key_valid[:, text_positions] = encoder_hidden_states_mask.bool()[:, vlm_text_positions] + + prefix_len = int((~target_token_mask).sum()) + + if kv_cache_mode == "cached": + # decode: only the target image's queries are recomputed. The block-causal mask degenerates to full + # attention for target rows (they see the entire prefix + their own block), so only the padding mask is + # needed. + joint_hidden_states = joint_hidden_states[:, prefix_len:] + rotary_emb = rotary_emb[prefix_len:] + modulation_mask = modulation_mask[prefix_len:] + attention_mask = None if joint_key_valid is None else joint_key_valid[:, None, None, :] + cache_write_slice = None + block_segments, block_key_valid = None, None + else: + # Prefill evaluates the joint sequence with exact per-segment causality. + attention_mask = None + block_segments = _qwenimage21_prefix_segments(image_ids, prefix_len) + cache_write_slice = slice(0, prefix_len) if kv_cache_mode == "extract" else None + block_key_valid = joint_key_valid + + for index_block, block in enumerate(self.transformer_blocks): + layer_cache = kv_cache.get_layer(index_block) if kv_cache is not None else None + joint_hidden_states = block( + hidden_states=joint_hidden_states, + modulation=modulation, + rotary_emb=rotary_emb, + attention_mask=attention_mask, + target_token_mask=modulation_mask, + layer_cache=layer_cache, + kv_cache_mode=kv_cache_mode, + cache_write_slice=cache_write_slice, + segments=block_segments, + key_valid=block_key_valid, + ) + + joint_hidden_states = self.norm_out(joint_hidden_states, temb, modulation_mask) + output = self.proj_out(joint_hidden_states) + + if not return_dict: + return (output,) + + return Transformer2DModelOutput(sample=output) + + +# Compatibility name for callers migrating from the reference implementation. +QwenImage21Transformer2DModel = QwenImageDiT diff --git a/freevideo_engine/image_gen/qwen_pipeline.py b/freevideo_engine/image_gen/qwen_pipeline.py new file mode 100644 index 00000000..e4c90f0f --- /dev/null +++ b/freevideo_engine/image_gen/qwen_pipeline.py @@ -0,0 +1,375 @@ +# Copyright 2026 Qwen-Image Team, The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Qwen-Image 2.1 Turbo orchestration, with no diffusers runtime dependency. + +Placement protocol: ``enter(stage)`` makes that stage's model executable on the +generator's device; ``leave(stage)`` parks/releases it, including on errors or +cancellation. Stages are ``text`` (QwenTextEncoder.model), ``encode`` (reference +VAE encoding, only for edits), ``denoise`` (transformer), and ``decode`` (VAE). +Models are available on the generator before enter is called; a lazy vision +tower is materialized on CPU before enter('text'). The default placement does +nothing, so callers using CUDA must explicitly arrange model placement. +Progress uses only the public text/denoise/decode labels, in that order. + +Instances are sequential and must not be shared between concurrent jobs. +The pack needs the three Comfy weight files plus local Qwen3-VL config and +processor assets (text_encoder/ and processor/, or text_encoders/ for both). +T1 supplies qwen_weights.load_dit and qwen_weights.load_vae; injected components +allow orchestration to be tested independently of those implementations. +""" + +import json +import math +import sys +from collections.abc import Mapping +from contextlib import contextmanager +from pathlib import Path +from typing import Protocol + +import numpy as np +import torch +from PIL import Image + +from .qwen_text import QwenTextEncoder + + +TURBO_SIGMAS = (1.0, 0.978453, 0.95418, 0.926626, 0.89508, 0.845148, 0.704534, 0.414568, 0.0) + + +class ImagePlacement(Protocol): + def enter(self, stage: str) -> None: ... + def leave(self, stage: str) -> None: ... + + +class NoopPlacement: + def enter(self, stage): + pass + + def leave(self, stage): + pass + + +class ImageGenerationCancelled(RuntimeError): + """The caller cancelled the job; no partial image is decoded or returned.""" + + +def normalize_dimensions(width, height): + """Floor to 32 pixels: 16x VAE compression plus the DiT's 2x2 grouping. + + This is a multiple of the contract's 16-pixel grid, with the stronger + divisibility requirement imposed by the reference DiT and pipeline. + """ + if any(isinstance(v, bool) or not isinstance(v, int) or v < 32 for v in (width, height)): + raise ValueError("width and height must be integers of at least 32 pixels") + return width // 32 * 32, height // 32 * 32 + + +def _validated_sigmas(values): + if not isinstance(values, (list, tuple)): + raise ValueError("sample_sigmas must be a list") + sigmas = tuple(float(value) for value in values) + if len(sigmas) == 8: + sigmas += (0.0,) + if (len(sigmas) != 9 or sigmas[0] != 1.0 or sigmas[-1] != 0.0 + or not all(math.isfinite(v) for v in sigmas) + or not all(a > b >= 0 for a, b in zip(sigmas, sigmas[1:]))): + raise ValueError("Turbo requires eight strictly decreasing sigmas from 1, then terminal 0") + return sigmas + + +def load_sample_sigmas(path): + """Read sample_sigmas from checkpoint metadata; absence uses the Turbo grid. + + Recognized metadata keys are sample_sigmas and config (a JSON object with + sample_sigmas). A present but invalid schedule is an error, never a fallback. + """ + from safetensors import safe_open + + with safe_open(str(path), framework="pt", device="cpu") as checkpoint: + metadata = checkpoint.metadata() or {} + if "sample_sigmas" in metadata: + return _validated_sigmas(json.loads(metadata["sample_sigmas"])) + if "config" in metadata: + config = json.loads(metadata["config"]) + if "sample_sigmas" in config: + return _validated_sigmas(config["sample_sigmas"]) + return TURBO_SIGMAS + + +def euler_step(latents, prediction, sigma, next_sigma): + # Match FlowMatchEulerDiscreteScheduler: FP32 sample, native prediction + # multiplication, and final cast to the prediction's dtype. + return (latents.float() + (next_sigma - sigma) * prediction).to(prediction.dtype) + + +def _config(model, name): + config = model.config + return config[name] if isinstance(config, Mapping) else getattr(config, name) + + +def _dtype(model): + for value in model.parameters(): + if value.is_floating_point(): + return value.dtype + return torch.float32 + + +def _encoded_latents(output): + if isinstance(output, torch.Tensor): + return output + if hasattr(output, "latent_dist"): + return output.latent_dist.mode() + if hasattr(output, "latents"): + return output.latents + raise TypeError("VAE.encode must return a Tensor, latent_dist, or latents") + + +def _sample(output): + if isinstance(output, torch.Tensor): + return output + if isinstance(output, tuple): + return output[0] + return output.sample + + +class QwenImageGenerator: + vae_scale_factor = 16 + + def __init__(self, pack_dir: Path | None, *, placement: ImagePlacement | None = None, device="cuda", + transformer=None, vae=None, text_encoder=None, processor=None, sample_sigmas=None, + reference_resolution=1024, vae_tile=False): + self.device = torch.device(device) + self.placement = placement if placement is not None else NoopPlacement() + self.pack_dir = Path(pack_dir) if pack_dir is not None else None + self.transformer, self.vae = transformer, vae + if text_encoder is not None and processor is not None: + text_encoder = QwenTextEncoder(text_encoder, processor) + self.text_encoder = text_encoder + self._closed = False + self.vae_tile = vae_tile + if not isinstance(reference_resolution, int) or reference_resolution < 32: + raise ValueError("reference_resolution must be at least 32") + self.reference_resolution = reference_resolution + if any(model is None for model in (self.transformer, self.vae, self.text_encoder)): + if self.pack_dir is None: + raise ValueError("pack_dir is required when components are not supplied") + self._load_components() + if sample_sigmas is None and self.pack_dir is not None: + dit_path = self._weight_path("diffusion_models", "qwen_image_2.1_turbo_int8_convrot.safetensors") + if dit_path.exists(): + sample_sigmas = load_sample_sigmas(dit_path) + elif (self.pack_dir / "model_index.json").exists(): + sample_sigmas = json.loads((self.pack_dir / "model_index.json").read_text()).get("sample_sigmas") + self.sample_sigmas = _validated_sigmas(sample_sigmas if sample_sigmas is not None else TURBO_SIGMAS) + + def _weight_path(self, directory, filename): + nested = self.pack_dir / directory / filename + return nested if nested.exists() else self.pack_dir / filename + + def _load_components(self): + if self.transformer is None: + from .qwen_weights import load_dit + + self.transformer = load_dit( + self._weight_path("diffusion_models", "qwen_image_2.1_turbo_int8_convrot.safetensors"), device="cpu" + ) + if self.vae is None: + from .qwen_weights import load_vae + + self.vae = load_vae(self._weight_path("vae", "qwen_image_2.1_vae_bf16.safetensors"), device="cpu") + if self.text_encoder is None: + from transformers import Qwen3VLProcessor + from .qwen_te_weights import load_text_encoder + + config_dir = self.pack_dir / "text_encoder" + if not config_dir.exists(): + config_dir = self.pack_dir / "text_encoders" + processor_dir = self.pack_dir / "processor" + if not processor_dir.exists(): + processor_dir = config_dir + model = load_text_encoder( + self._weight_path("text_encoders", "qwen3vl_8b_int8_convrot.safetensors"), + config_path=config_dir, device="cpu", + ) + processor = Qwen3VLProcessor.from_pretrained(str(processor_dir), local_files_only=True) + self.text_encoder = QwenTextEncoder(model, processor) + + @contextmanager + def _stage(self, stage): + self.placement.enter(stage) + try: + yield + finally: + self.placement.leave(stage) + + @staticmethod + def _pack_latents(latents, batch_size, num_channels_latents, height, width): + return latents.reshape(batch_size, num_channels_latents, height * width).transpose(1, 2) + + @staticmethod + def _unpack_latents(latents, height, width, vae_scale_factor=16): + batch_size, _, channels = latents.shape + height = 2 * (int(height) // (vae_scale_factor * 2)) + width = 2 * (int(width) // (vae_scale_factor * 2)) + return latents.transpose(1, 2).reshape(batch_size, channels, 1, height, width) + + def _stats(self, latents): + channels = _config(self.vae, "z_dim") + mean = torch.as_tensor(_config(self.vae, "latents_mean"), device=latents.device, dtype=latents.dtype) + std = torch.as_tensor(_config(self.vae, "latents_std"), device=latents.device, dtype=latents.dtype) + return mean.reshape(1, channels, 1, 1, 1), std.reshape(1, channels, 1, 1, 1) + + def _reference_inputs(self, images): + pil_images, vae_images, sizes = [], [], [] + for image in images or []: + if not isinstance(image, Image.Image): + raise ValueError("reference_images must be a flat list of PIL images") + image = image.convert("RGBA") + ratio = image.width / image.height + width = math.sqrt(self.reference_resolution**2 * ratio) + height = width / ratio + width, height = round(width / 32) * 32, round(height / 32) * 32 + if min(width, height) < 32: + raise ValueError("Reference image aspect ratio produces an empty latent grid") + resized = image.resize((width, height), Image.Resampling.LANCZOS) + pil_images.append(resized) + # Preserve alpha for the four-channel VAE; only the encoder's copy + # is composited over white in QwenTextEncoder. + # Batch axis first, then channels: the layout the reference pipeline feeds the VAE + # (same values as a per-image permute, a different memory order, and with it a + # different convolution path). + array = np.stack([np.asarray(resized)], axis=0).astype(np.float32) / 255.0 + vae_images.append(torch.from_numpy(array).permute(0, 3, 1, 2).unsqueeze(2) * 2 - 1) + sizes.append((width, height)) + return pil_images or None, vae_images, sizes + + @torch.no_grad() + def generate(self, prompt: str, *, width: int, height: int, seed: int | None = None, steps: int = 8, + reference_images=None, progress=None, cancel=None) -> Image.Image: + if self._closed: + raise RuntimeError("Generator is closed") + if not isinstance(prompt, str): + raise ValueError("prompt must be a string") + if isinstance(steps, bool) or steps != 8: + raise ValueError("Qwen-Image Turbo supports the checkpoint's fixed eight-step profile") + if reference_images is not None and not isinstance(reference_images, list): + raise ValueError("reference_images must be a flat list of PIL images") + width, height = normalize_dimensions(width, height) + + def check_cancel(): + if cancel is not None and cancel(): + raise ImageGenerationCancelled("Image generation cancelled") + + def report(stage, fraction): + if progress is not None: + progress(stage, fraction) + check_cancel() + + check_cancel() + images, vae_images, sizes = self._reference_inputs(reference_images) + if hasattr(self.text_encoder, "prepare_images"): + self.text_encoder.prepare_images(images) + check_cancel() + with self._stage("text"): + report("text", 0.0) + if hasattr(self.text_encoder, "encode_prompt_with_image_mask"): + embeds, mask, image_mask = self.text_encoder.encode_prompt_with_image_mask( + prompt, images, device=self.device + ) + else: + if images: + raise TypeError("Image conditioning needs encode_prompt_with_image_mask") + embeds, mask = self.text_encoder.encode_prompt(prompt) + image_mask = torch.zeros(embeds.shape[:2], dtype=torch.bool, device=embeds.device) + embeds = embeds.to(self.device) + mask = mask.to(self.device) if mask is not None else None + image_mask = image_mask.to(device=self.device, dtype=torch.bool) + report("text", 1.0) + + channels = _config(self.transformer, "in_channels") + if channels != _config(self.vae, "z_dim"): + raise ValueError("DiT in_channels must equal VAE z_dim (2.1 latents are unpatched)") + ref_latents = [] + if vae_images: + with self._stage("encode"): + for image in vae_images: + check_cancel() + # Always tiled, whatever the memory plan: the latents (and so the picture) + # must not depend on the card, and the tiled pass needs far less memory. + encoded = _encoded_latents( + self.vae.encode(image.to(device=self.device, dtype=_dtype(self.vae)), tile=True)) + mean, std = self._stats(encoded) + encoded = (encoded - mean) / std + ref_latents.append(self._pack_latents(encoded, 1, channels, *encoded.shape[-2:]).to(embeds.dtype)) + check_cancel() + ref_latents = torch.cat(ref_latents, dim=1) if ref_latents else None + ref_tokens = 0 if ref_latents is None else ref_latents.shape[1] + if image_mask.sum().item() * 4 != ref_tokens: + raise ValueError("Vision image slots and reference VAE latent grid differ") + + generator = torch.Generator(device=self.device) + generator.manual_seed(seed) if seed is not None else generator.seed() + latent_height, latent_width = height // 16, width // 16 + latents = torch.randn((1, 1, channels, latent_height, latent_width), generator=generator, + device=self.device, dtype=embeds.dtype) + latents = self._pack_latents(latents, 1, channels, latent_height, latent_width) + img_shapes = [[*[(1, h // 16, w // 16) for w, h in sizes], (1, latent_height, latent_width)]] + image_mask = torch.cat([image_mask, image_mask.new_ones(1, latents.shape[1] // 4)], dim=1) + sigmas = torch.tensor(self.sample_sigmas, device=self.device, dtype=torch.float32) + timesteps = sigmas[:-1] * 1000 + with self._stage("denoise"): + report("denoise", 0.0) + for i, timestep in enumerate(timesteps): + check_cancel() + model_input = torch.cat([ref_latents, latents], dim=1) if ref_latents is not None else latents + prediction = _sample(self.transformer( + hidden_states=model_input, + timestep=timestep.expand(1).to(latents.dtype) / 1000, + encoder_hidden_states=embeds, + encoder_hidden_states_mask=mask, + img_shapes=img_shapes, + img_mask=image_mask, + attention_kwargs=None, + kv_cache=None, + kv_cache_mode=None, + return_dict=False, + ))[:, -latents.shape[1]:] + check_cancel() + latents = euler_step(latents, prediction, sigmas[i], sigmas[i + 1]) + report("denoise", (i + 1) / steps) + + check_cancel() + with self._stage("decode"): + report("decode", 0.0) + latents = self._unpack_latents(latents, height, width).to(dtype=_dtype(self.vae)) + mean, std = self._stats(latents) + latents = latents * std + mean + kwargs = {"return_dict": False} + if self.vae_tile: + kwargs["tile"] = True + pixels = _sample(self.vae.decode(latents, **kwargs))[:, :, 0] + check_cancel() + array = (pixels / 2 + 0.5).clamp(0, 1).float().cpu().permute(0, 2, 3, 1).numpy()[0] + image = Image.fromarray((array * 255).round().astype(np.uint8)) + report("decode", 1.0) + return image + + def close(self) -> None: + """Release owned model references; all active stages have already left.""" + self.text_encoder = self.transformer = self.vae = None + self._closed = True + # Int8Linear keeps the last activation alive for Q/K/V sharing. + module = sys.modules.get("freevideo_engine.int8_ops") + if module is not None: + module.Int8Linear._last.clear() diff --git a/freevideo_engine/image_gen/qwen_te_weights.py b/freevideo_engine/image_gen/qwen_te_weights.py new file mode 100644 index 00000000..abfb066d --- /dev/null +++ b/freevideo_engine/image_gen/qwen_te_weights.py @@ -0,0 +1,211 @@ +"""T2: load the Comfy Qwen3-VL ConvRot encoder without a BF16 model copy. + +The DiT loader is supplied by T1. Encoder parameters start on meta and each +checkpoint tensor is installed directly. Only the vocabulary is dequantized, +in bounded row chunks; CUDA projections use the existing engine Int8Linear. +No downloads occur here. Supply a local Qwen3-VL config (or ``config=``) and +local processor/tokenizer assets separately. Vision weights stay unopened +until ``model.load_vision()`` is called before the image-conditioned text stage. +""" + +import json +import weakref +from pathlib import Path + +import torch +from safetensors import safe_open + + +_CONVROT = {"format": "int8_tensorwise", "convrot": True, "convrot_groupsize": 256} + + +def _quant_metadata(st, name): + value = st.get_tensor(name + ".comfy_quant") + if value.dtype != torch.uint8 or value.ndim != 1: + raise ValueError(f"Invalid comfy_quant JSON tensor: {name}") + try: + meta = json.loads(bytes(value.tolist()).decode("utf-8")) + except (ValueError, UnicodeDecodeError) as error: + raise ValueError(f"Invalid comfy_quant JSON: {name}") from error + if meta != _CONVROT: + raise ValueError(f"Unsupported ConvRot format for {name}: {meta}") + + +def _row_scales(st, name, rows): + scales = st.get_tensor(name + ".weight_scale") + if scales.dtype != torch.float32 or tuple(scales.shape) not in ((rows,), (rows, 1), (1, rows)): + raise ValueError(f"Expected FP32 per-row scales for {name}") + scales = scales.reshape(-1).contiguous() + if not torch.isfinite(scales).all() or (scales <= 0).any(): + raise ValueError(f"Invalid per-row scales for {name}") + return scales + + +def _dequantize_vocabulary(st, name, shape, dtype, chunk_rows): + weight = st.get_slice(name + ".weight") + if list(weight.get_shape()) != list(shape) or weight.get_dtype() != "I8" or shape[1] % 256: + raise ValueError(f"Expected group-256 int8 vocabulary for {name}") + scales = _row_scales(st, name, shape[0]) + # Exactly int8_ops.hadamard16: the regular symmetric, orthogonal H256, + # rather than the Sylvester Hadamard convention used by some quantizers. + h4 = torch.tensor([[1, 1, 1, -1], [1, 1, -1, 1], [1, -1, 1, 1], [-1, 1, 1, 1]], dtype=torch.float32) + h16 = torch.kron(h4, h4) / 4 + full = torch.empty(shape, dtype=dtype) + for start in range(0, shape[0], chunk_rows): + end = min(start + chunk_rows, shape[0]) + rotated = (weight[start:end].float() * scales[start:end, None]).reshape(-1, shape[1] // 256, 16, 16) + full[start:end] = (h16 @ rotated @ h16).reshape(-1, shape[1]).to(dtype) + return full + + +def _text_name(model, name): + candidates = [name] + if name.startswith("model.") and not name.startswith("model.language_model."): + candidates.append(name.replace("model.", "model.language_model.", 1)) + for candidate in candidates: + try: + model.get_submodule(candidate) + return candidate + except AttributeError: + pass + raise ValueError(f"Unknown Qwen3-VL checkpoint module: {name}") + + +def _assign_tensor(model, name, value, dtype): + prefix, _, leaf = name.rpartition(".") + try: + module = model.get_submodule(prefix) if prefix else model + except AttributeError as error: + raise ValueError(f"Unknown checkpoint tensor: {name}") from error + target = module._parameters.get(leaf) + is_parameter = target is not None + if not is_parameter: + target = module._buffers.get(leaf) + if target is None or target.shape != value.shape: + raise ValueError(f"Unknown tensor or shape mismatch: {name}, {tuple(value.shape)}") + if not value.is_floating_point(): + raise ValueError(f"Unexpected non-floating unquantized tensor: {name}") + value = value.to(dtype=dtype) + if is_parameter: + module._parameters[leaf] = torch.nn.Parameter(value, requires_grad=False) + else: + module._buffers[leaf] = value + + +def _install_text_weights(model, st, *, dtype, chunk_rows): + keys = set(st.keys()) + handled = set() + for key in sorted(k for k in keys if k.endswith(".comfy_quant")): + name = key[:-len(".comfy_quant")] + if name.startswith("model.visual.") or name == "lm_head": + continue + _quant_metadata(st, name) + full_name = _text_name(model, name) + module = model.get_submodule(full_name) + shape = module.weight.shape if isinstance(module, (torch.nn.Linear, torch.nn.Embedding)) else () + if not shape: + raise ValueError(f"Quantized target is not Linear/Embedding: {name}") + if isinstance(module, torch.nn.Embedding): + module.weight = torch.nn.Parameter( + _dequantize_vocabulary(st, name, shape, dtype, chunk_rows), requires_grad=False + ) + else: + from ..int8_ops import Int8Linear + + weight = st.get_tensor(name + ".weight") + if weight.dtype != torch.int8 or weight.shape != shape or shape[1] % 256: + raise ValueError(f"Expected group-256 int8 matrix for {name}") + scales = _row_scales(st, name, shape[0]) + replacement = Int8Linear.__new__(Int8Linear) + torch.nn.Module.__init__(replacement) + replacement.register_buffer("weight_int8", weight.contiguous()) + replacement.register_buffer("weight_scale", scales) + replacement.original = torch.nn.Module() + replacement.original.register_parameter("bias", module.bias) + replacement.in_features, replacement.out_features = module.in_features, module.out_features + replacement.input_dtype = dtype + parent, _, child = full_name.rpartition(".") + setattr(model.get_submodule(parent), child, replacement) + handled.update({key, name + ".weight", name + ".weight_scale"}) + + for key in sorted(keys - handled): + if key.startswith(("model.visual.", "lm_head.")): + continue + if key.endswith((".comfy_quant", ".weight_scale")): + raise ValueError(f"Unpaired quantization tensor: {key}") + name, _, leaf = key.rpartition(".") + full_name = _text_name(model, name) + module = model.get_submodule(full_name) + if leaf == "weight" and hasattr(module, "weight_int8"): + raise ValueError(f"Duplicate quantized weight: {key}") + if leaf == "bias" and hasattr(module, "weight_int8"): + full_name += ".original" + _assign_tensor(model, full_name + "." + leaf, st.get_tensor(key), dtype) + missing = [name for name, value in model.state_dict().items() if value.is_meta] + if missing: + raise ValueError("Missing text encoder tensors: " + ", ".join(missing)) + + +def load_text_encoder(path, *, config=None, config_path=None, device="cpu", dtype=torch.bfloat16, + chunk_rows=8192): + """Return a Qwen3VLForConditionalGeneration with int8 language projections. + + ``config`` accepts a transformers Qwen3VLConfig or its dictionary. Otherwise + ``config_path`` (default: checkpoint sibling config.json) must be local. + ``lm_head`` is Identity because only hidden states are consumed. The lazy + ``load_vision(device='cpu')`` method installs only BF16 visual tensors and + returns that module, allowing the caller's placement policy to own moves. + CPU supports loading/inspection; Int8Linear execution requires CUDA. + """ + from accelerate import init_empty_weights + from transformers import Qwen3VLConfig, Qwen3VLForConditionalGeneration + + path = Path(path) + if not isinstance(chunk_rows, int) or chunk_rows <= 0: + raise ValueError("chunk_rows must be a positive integer") + if dtype not in (torch.float32, torch.float16, torch.bfloat16): + raise ValueError("Unsupported text encoder floating dtype") + if config is None: + config_file = Path(config_path) if config_path else path.with_name("config.json") + if not config_file.exists(): + raise FileNotFoundError(f"Local Qwen3-VL config required: {config_file}") + config = Qwen3VLConfig.from_pretrained(str(config_file), local_files_only=True) + elif isinstance(config, dict): + config = Qwen3VLConfig.from_dict(config) + config._attn_implementation = "sdpa" + config.text_config.use_cache = False + with init_empty_weights(): + model = Qwen3VLForConditionalGeneration(config) + # Retain the constructor, not the meta vision module or unused logits head. + vision_type = type(model.model.visual) + model.model.visual = None + model.lm_head = torch.nn.Identity() + with safe_open(str(path), framework="pt", device="cpu") as st: + _install_text_weights(model, st, dtype=dtype, chunk_rows=chunk_rows) + model.eval().requires_grad_(False).to(device=device) + model_ref = weakref.ref(model) + + def load_vision(device="cpu"): + owner = model_ref() + if owner is None: + raise RuntimeError("Text encoder has been released") + if owner.model.visual is not None: + return owner.model.visual + with init_empty_weights(): + visual = vision_type(config.vision_config) + with safe_open(str(path), framework="pt", device="cpu") as st: + visual_keys = [key for key in st.keys() if key.startswith("model.visual.")] + for key in sorted(visual_keys): + value = st.get_tensor(key) + if value.dtype != torch.bfloat16: + raise ValueError(f"Expected BF16 visual tensor: {key}") + _assign_tensor(visual, key[len("model.visual."):], value, torch.bfloat16) + missing = [name for name, value in visual.state_dict().items() if value.is_meta] + if missing: + raise ValueError("Missing vision tensors: " + ", ".join(missing)) + visual.eval().requires_grad_(False).to(device=device) + owner.model.visual = visual + return visual + + model.load_vision = load_vision + return model diff --git a/freevideo_engine/image_gen/qwen_text.py b/freevideo_engine/image_gen/qwen_text.py new file mode 100644 index 00000000..43674565 --- /dev/null +++ b/freevideo_engine/image_gen/qwen_text.py @@ -0,0 +1,134 @@ +# Copyright 2026 Qwen-Image Team, The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Qwen3-VL prompt conditioning, ported from the Qwen-Image 2.1 pipeline. + +``encode_prompt`` returns the public (embeddings, attention_mask) pair. +``encode_prompt_with_image_mask`` additionally returns the interleaved image +slots required by the 2.1 DiT. No models are moved between devices here. +The wrapper and underlying Qwen3-VL model are used sequentially: the temporary +RMSNorm hook, like the reference implementation, is not thread safe. +""" + +import torch +from PIL import Image + + +class QwenTextEncoder: + def __init__(self, model, processor): + self.model = model + self.processor = processor + self.sys_prompt = "Comprehend and analyze the provided prompt." + self.prompt_template_t2i = ( + f"<|im_start|>system\n{self.sys_prompt}<|im_end|>\n" + "<|im_start|>user\n{}<|im_end|>\n" + "<|im_start|>assistant\n" + ) + self.prompt_template_ti2i = ( + f"<|im_start|>system\n{self.sys_prompt}<|im_end|>\n" + "<|im_start|>user\n<|vision_start|><|image_pad|><|vision_end|>{}<|im_end|>\n" + "<|im_start|>assistant\n" + ) + sys_message = [{"role": "system", "content": [{"type": "text", "text": self.sys_prompt}]}] + sys_tokens = processor.apply_chat_template(sys_message, tokenize=True, return_dict=False) + self._drop_idx = len(sys_tokens[0]) + self._img_token_id = processor.tokenizer.encode("<|image_pad|>")[0] + + def prepare_images(self, images): + """Materialize the lazy BF16 vision tower on CPU before placement.enter.""" + if images and hasattr(self.model, "load_vision"): + self.model.load_vision() + + @torch.no_grad() + def encode_prompt(self, prompt, images=None): + embeds, mask, _ = self.encode_prompt_with_image_mask(prompt, images) + return embeds, mask + + @torch.no_grad() + def encode_prompt_with_image_mask(self, prompt, images=None, *, device=None): + if device is None: + device = next(self.model.parameters()).device + prompt = [prompt] if isinstance(prompt, str) else prompt + if not prompt or not all(isinstance(p, str) for p in prompt): + raise ValueError("prompt must be a string or a nonempty list of strings") + prompt = [" " if not p else p for p in prompt] + if images is not None and not images: + images = None + is_t2i = images is None + if is_t2i: + prompts = [self.prompt_template_t2i.format(t) for t in prompt] + else: + self.prepare_images(images) + prompts, condition_pil_list = [], [] + for t in prompt: + replace = "<|vision_start|><|image_pad|><|vision_end|>" + for i in range(2, len(images) + 1): + replace += f" <|vision_start|><|image_pad|><|vision_end|>" + template = self.prompt_template_ti2i.replace( + "<|vision_start|><|image_pad|><|vision_end|>", replace + ) + prompts.append(template.format(t)) + for _ in prompt: + for img in images: + if not isinstance(img, Image.Image): + img = Image.fromarray(img) + if img.mode == "RGBA": + white = Image.new("RGB", img.size, (255, 255, 255)) + white.paste(img, mask=img.getchannel("A")) + img = white + condition_pil_list.append(img) + + processor_kwargs = dict(text=prompts, padding=True, padding_side="left", return_tensors="pt") + if not is_t2i: + processor_kwargs["images"] = condition_pil_list + model_inputs = self.processor(**processor_kwargs).to(device) + forward_kwargs = dict( + input_ids=model_inputs.input_ids, + attention_mask=model_inputs.attention_mask, + output_hidden_states=True, + ) + if not is_t2i and hasattr(model_inputs, "pixel_values"): + forward_kwargs.update(pixel_values=model_inputs.pixel_values, image_grid_thw=model_inputs.image_grid_thw) + if hasattr(model_inputs, "mm_token_type_ids"): + forward_kwargs["mm_token_type_ids"] = model_inputs.mm_token_type_ids + + # Transformers 5.15 ties hidden_states[-1] to the normalized output. + # Return the norm's input for this call, exactly as the reference does. + text_model = getattr(self.model.model, "language_model", self.model.model) + handle = text_model.norm.register_forward_hook(lambda module, args, output: args[0]) + try: + outputs = self.model(**forward_kwargs) + finally: + handle.remove() + hidden_states = outputs.hidden_states[-1] + bool_mask = model_inputs.attention_mask.bool() + valid_lengths = bool_mask.sum(dim=1) + selected = hidden_states[bool_mask] + split_hidden_states = list(torch.split(selected, valid_lengths.tolist(), dim=0)) + split_hidden_states = [e[self._drop_idx :] for e in split_hidden_states] + image_pad_mask = [ + (sample_ids[sample_mask.bool()] == self._img_token_id)[self._drop_idx :] + for sample_ids, sample_mask in zip(model_inputs.input_ids, model_inputs.attention_mask) + ] + attn_mask_list = [torch.ones(e.size(0), dtype=torch.long, device=e.device) for e in split_hidden_states] + max_seq_len = max(e.size(0) for e in split_hidden_states) + prompt_embeds = torch.stack( + [torch.cat([u, u.new_zeros(max_seq_len - u.size(0), u.size(1))]) for u in split_hidden_states] + ) + encoder_attention_mask = torch.stack( + [torch.cat([u, u.new_zeros(max_seq_len - u.size(0))]) for u in attn_mask_list] + ) + image_pad_mask = torch.stack([torch.cat([u, u.new_zeros(max_seq_len - u.size(0))]) for u in image_pad_mask]) + if encoder_attention_mask.all(): + encoder_attention_mask = None + return prompt_embeds, encoder_attention_mask, image_pad_mask diff --git a/freevideo_engine/image_gen/qwen_vae.py b/freevideo_engine/image_gen/qwen_vae.py new file mode 100644 index 00000000..e00ecc0e --- /dev/null +++ b/freevideo_engine/image_gen/qwen_vae.py @@ -0,0 +1,1644 @@ +# Copyright 2026 The Qwen Team and The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import torch +import torch.nn as nn +import torch.nn.functional as F +import re + +from ._model_utils import AutoencoderKLOutput, ModelOutput, DiagonalGaussianDistribution, model_config + +# Ported from the supplied Diffusers Qwen-Image 2.1 implementation. +# Changes: plain nn.Module/config/output helpers, no Diffusers loader hooks; +# Comfy key mapping and single-tensor loading are implemented below. +DecoderOutput = ModelOutput + + +CACHE_T = 2 + + +# Copied from diffusers.models.autoencoders.autoencoder_kl_wan.AvgDown3D with AvgDown3D->QwenImage21AvgDown3D +class QwenImage21AvgDown3D(nn.Module): + def __init__( + self, + in_channels, + out_channels, + factor_t, + factor_s=1, + ): + super().__init__() + factor = factor_t * factor_s * factor_s + if in_channels * factor % out_channels != 0: + raise ValueError( + f"`in_channels` ({in_channels}) times the downsampling factor ({factor}) must be divisible by " + f"`out_channels` ({out_channels})." + ) + self.in_channels = in_channels + self.out_channels = out_channels + self.factor_t = factor_t + self.factor_s = factor_s + self.factor = factor + self.group_size = in_channels * factor // out_channels + + def forward(self, x: torch.Tensor) -> torch.Tensor: + pad_t = (self.factor_t - x.shape[2] % self.factor_t) % self.factor_t + pad = (0, 0, 0, 0, pad_t, 0) + x = F.pad(x, pad) + B, C, T, H, W = x.shape + x = x.view( + B, + C, + T // self.factor_t, + self.factor_t, + H // self.factor_s, + self.factor_s, + W // self.factor_s, + self.factor_s, + ) + x = x.permute(0, 1, 3, 5, 7, 2, 4, 6).contiguous() + x = x.view( + B, + C * self.factor, + T // self.factor_t, + H // self.factor_s, + W // self.factor_s, + ) + x = x.view( + B, + self.out_channels, + self.group_size, + T // self.factor_t, + H // self.factor_s, + W // self.factor_s, + ) + x = x.mean(dim=2) + return x + + +# Copied from diffusers.models.autoencoders.autoencoder_kl_wan.DupUp3D with DupUp3D->QwenImage21DupUp3D +class QwenImage21DupUp3D(nn.Module): + def __init__( + self, + in_channels: int, + out_channels: int, + factor_t, + factor_s=1, + ): + super().__init__() + self.in_channels = in_channels + self.out_channels = out_channels + + self.factor_t = factor_t + self.factor_s = factor_s + self.factor = self.factor_t * self.factor_s * self.factor_s + + assert out_channels * self.factor % in_channels == 0 + self.repeats = out_channels * self.factor // in_channels + + def forward(self, x: torch.Tensor, first_chunk=False) -> torch.Tensor: + x = x.repeat_interleave(self.repeats, dim=1) + x = x.view( + x.size(0), + self.out_channels, + self.factor_t, + self.factor_s, + self.factor_s, + x.size(2), + x.size(3), + x.size(4), + ) + x = x.permute(0, 1, 5, 2, 6, 3, 7, 4).contiguous() + x = x.view( + x.size(0), + self.out_channels, + x.size(2) * self.factor_t, + x.size(4) * self.factor_s, + x.size(6) * self.factor_s, + ) + if first_chunk: + x = x[:, :, self.factor_t - 1 :, :, :] + return x + + +class QwenImage21CausalConv3d(nn.Conv2d): + r""" + A custom 3D causal convolution layer with feature caching support. + + This layer extends the standard Conv3D layer by ensuring causality in the time dimension and handling feature + caching for efficient inference. + + Args: + in_channels (int): Number of channels in the input image + out_channels (int): Number of channels produced by the convolution + kernel_size (int or tuple): Size of the convolving kernel + stride (int or tuple, optional): Stride of the convolution. Default: 1 + padding (int or tuple, optional): Zero-padding added to all three sides of the input. Default: 0 + """ + + def __init__( + self, + in_channels: int, + out_channels: int, + kernel_size: int | tuple[int | int | int], + stride: int | tuple[int | int | int] = 1, + padding: int | tuple[int | int | int] = 0, + ) -> None: + super().__init__( + in_channels=in_channels, + out_channels=out_channels, + kernel_size=kernel_size, + stride=stride, + padding=padding, + ) + + # Set up causal padding + self._padding = (self.padding[1], self.padding[1], self.padding[0], self.padding[0]) + self.padding = (0, 0) + + def forward(self, x, cache_x=None) -> torch.Tensor: + padding = list(self._padding) + if cache_x is not None: + raise ValueError( + "This convolution is the image specialization of Wan's causal 3D one: it folds the single frame away " + "and has no temporal context to prepend, so it cannot take a feature cache." + ) + x = x.squeeze(2) # Remove the temporal dimension + x = F.pad(x, padding) + x = super().forward(x) + x = x.unsqueeze(2) # Add the temporal dimension back + return x + + +# Copied from diffusers.models.autoencoders.autoencoder_kl_wan.WanRMS_norm with Wan->QwenImage21 +class QwenImage21RMS_norm(nn.Module): + r""" + A custom RMS normalization layer. + + Args: + dim (int): The number of dimensions to normalize over. + channel_first (bool, optional): Whether the input tensor has channels as the first dimension. + Default is True. + images (bool, optional): Whether the input represents image data. Default is True. + bias (bool, optional): Whether to include a learnable bias term. Default is False. + """ + + def __init__(self, dim: int, channel_first: bool = True, images: bool = True, bias: bool = False) -> None: + super().__init__() + broadcastable_dims = (1, 1, 1) if not images else (1, 1) + shape = (dim, *broadcastable_dims) if channel_first else (dim,) + + self.channel_first = channel_first + self.scale = dim**0.5 + self.gamma = nn.Parameter(torch.ones(shape)) + self.bias = nn.Parameter(torch.zeros(shape)) if bias else 0.0 + + def forward(self, x) -> torch.Tensor: + needs_fp32_normalize = x.dtype in (torch.float16, torch.bfloat16) or any( + t in str(x.dtype) for t in ("float4_", "float8_") + ) + normalized = F.normalize(x.float() if needs_fp32_normalize else x, dim=(1 if self.channel_first else -1)).to( + x.dtype + ) + + return normalized * self.scale * self.gamma + self.bias + + +# Copied from diffusers.models.autoencoders.autoencoder_kl_wan.WanUpsample with Wan->QwenImage21 +class QwenImage21Upsample(nn.Upsample): + r""" + Perform upsampling while ensuring the output tensor has the same data type as the input. + + Args: + x (torch.Tensor): Input tensor to be upsampled. + + Returns: + torch.Tensor: Upsampled tensor with the same data type as the input. + """ + + def forward(self, x) -> torch.Tensor: + return super().forward(x.float()).type_as(x) + + +class QwenImage21Resample(nn.Module): + r""" + A custom resampling module for 2D and 3D data. + + Args: + dim (int): The number of input/output channels. + mode (str): The resampling mode. Must be one of: + - 'none': No resampling (identity operation). + - 'upsample2d': 2D upsampling with nearest-exact interpolation and convolution. + - 'upsample3d': 3D upsampling with nearest-exact interpolation, convolution, and causal 3D convolution. + - 'downsample2d': 2D downsampling with zero-padding and convolution. + - 'downsample3d': 3D downsampling with zero-padding, convolution, and causal 3D convolution. + """ + + def __init__(self, dim: int, mode: str, upsample_out_dim: int = None) -> None: + super().__init__() + self.dim = dim + self.mode = mode + + # default to dim //2 + if upsample_out_dim is None: + upsample_out_dim = dim // 2 + + # layers + if mode == "upsample2d": + self.resample = nn.Sequential( + QwenImage21Upsample(scale_factor=(2.0, 2.0), mode="nearest-exact"), + nn.Conv2d(dim, upsample_out_dim, 3, padding=1), + ) + elif mode == "upsample3d": + self.resample = nn.Sequential( + QwenImage21Upsample(scale_factor=(2.0, 2.0), mode="nearest-exact"), + nn.Conv2d(dim, upsample_out_dim, 3, padding=1), + ) + self.time_conv = QwenImage21CausalConv3d(dim, dim * 2, (1, 1), padding=(0, 0)) + + elif mode == "downsample2d": + self.resample = nn.Sequential(nn.ZeroPad2d((0, 1, 0, 1)), nn.Conv2d(dim, dim, 3, stride=(2, 2))) + elif mode == "downsample3d": + self.resample = nn.Sequential(nn.ZeroPad2d((0, 1, 0, 1)), nn.Conv2d(dim, dim, 3, stride=(2, 2))) + self.time_conv = QwenImage21CausalConv3d(dim, dim, (1, 1), stride=(1, 1), padding=(0, 0)) + + else: + self.resample = nn.Identity() + + def forward(self, x, feat_cache=None, feat_idx=None) -> torch.Tensor: + if feat_idx is None: + feat_idx = [0] + b, c, t, h, w = x.size() + if self.mode == "upsample3d": + if feat_cache is not None: + idx = feat_idx[0] + if feat_cache[idx] is None: + feat_cache[idx] = "Rep" + feat_idx[0] += 1 + else: + cache_x = x[:, :, -CACHE_T:, :, :].clone() + if cache_x.shape[2] < 2 and feat_cache[idx] is not None and feat_cache[idx] != "Rep": + # cache last frame of last two chunk + cache_x = torch.cat( + [feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x], dim=2 + ) + if cache_x.shape[2] < 2 and feat_cache[idx] is not None and feat_cache[idx] == "Rep": + cache_x = torch.cat([torch.zeros_like(cache_x).to(cache_x.device), cache_x], dim=2) + if feat_cache[idx] == "Rep": + x = self.time_conv(x) + else: + x = self.time_conv(x, feat_cache[idx]) + feat_cache[idx] = cache_x + feat_idx[0] += 1 + + x = x.reshape(b, 2, c, t, h, w) + x = torch.stack((x[:, 0, :, :, :, :], x[:, 1, :, :, :, :]), 3) + x = x.reshape(b, c, t * 2, h, w) + t = x.shape[2] + x = x.permute(0, 2, 1, 3, 4).reshape(b * t, c, h, w) + x = self.resample(x) + x = x.view(b, t, x.size(1), x.size(2), x.size(3)).permute(0, 2, 1, 3, 4) + + if self.mode == "downsample3d": + if feat_cache is not None: + idx = feat_idx[0] + if feat_cache[idx] is None: + feat_cache[idx] = x.clone() + feat_idx[0] += 1 + else: + cache_x = x[:, :, -1:, :, :].clone() + x = self.time_conv(torch.cat([feat_cache[idx][:, :, -1:, :, :], x], 2)) + feat_cache[idx] = cache_x + feat_idx[0] += 1 + return x + + +class QwenImage21ResidualBlock(nn.Module): + r""" + A custom residual block module. + + Args: + in_dim (int): Number of input channels. + out_dim (int): Number of output channels. + dropout (float, optional): Dropout rate for the dropout layer. Default is 0.0. + """ + + def __init__( + self, + in_dim: int, + out_dim: int, + dropout: float = 0.0, + ) -> None: + super().__init__() + self.in_dim = in_dim + self.out_dim = out_dim + self.nonlinearity = nn.SiLU() + + # layers + self.norm1 = QwenImage21RMS_norm(in_dim, images=False) + self.conv1 = QwenImage21CausalConv3d(in_dim, out_dim, 3, padding=1) + self.norm2 = QwenImage21RMS_norm(out_dim, images=False) + self.dropout = nn.Dropout(dropout) + self.conv2 = QwenImage21CausalConv3d(out_dim, out_dim, 3, padding=1) + self.conv_shortcut = QwenImage21CausalConv3d(in_dim, out_dim, 1) if in_dim != out_dim else nn.Identity() + + def forward(self, x, feat_cache=None, feat_idx=None) -> torch.Tensor: + if feat_idx is None: + feat_idx = [0] + # Apply shortcut connection + h = self.conv_shortcut(x) + + # First normalization and activation + x = self.norm1(x) + x = self.nonlinearity(x) + + if feat_cache is not None: + idx = feat_idx[0] + cache_x = x[:, :, -CACHE_T:, :, :].clone() + if cache_x.shape[2] < 2 and feat_cache[idx] is not None: + cache_x = torch.cat([feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x], dim=2) + + x = self.conv1(x, feat_cache[idx]) + feat_cache[idx] = cache_x + feat_idx[0] += 1 + else: + x = self.conv1(x) + + # Second normalization and activation + x = self.norm2(x) + x = self.nonlinearity(x) + + # Dropout + x = self.dropout(x) + + if feat_cache is not None: + idx = feat_idx[0] + cache_x = x[:, :, -CACHE_T:, :, :].clone() + if cache_x.shape[2] < 2 and feat_cache[idx] is not None: + cache_x = torch.cat([feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x], dim=2) + + x = self.conv2(x, feat_cache[idx]) + feat_cache[idx] = cache_x + feat_idx[0] += 1 + else: + x = self.conv2(x) + + # Add residual connection + return x + h + + +# Copied from diffusers.models.autoencoders.autoencoder_kl_wan.WanAttentionBlock with Wan->QwenImage21 +class QwenImage21AttentionBlock(nn.Module): + r""" + Causal self-attention with a single head. + + Args: + dim (int): The number of channels in the input tensor. + """ + + def __init__(self, dim): + super().__init__() + self.dim = dim + + # layers + self.norm = QwenImage21RMS_norm(dim) + self.to_qkv = nn.Conv2d(dim, dim * 3, 1) + self.proj = nn.Conv2d(dim, dim, 1) + + def forward(self, x) -> torch.Tensor: + identity = x + batch_size, channels, time, height, width = x.size() + + x = x.permute(0, 2, 1, 3, 4).reshape(batch_size * time, channels, height, width) + x = self.norm(x) + + # compute query, key, value + qkv = self.to_qkv(x) + qkv = qkv.reshape(batch_size * time, 1, channels * 3, -1) + qkv = qkv.permute(0, 1, 3, 2).contiguous() + q, k, v = qkv.chunk(3, dim=-1) + + # apply attention + x = F.scaled_dot_product_attention(q, k, v) + + x = x.squeeze(1).permute(0, 2, 1).reshape(batch_size * time, channels, height, width) + + # output projection + x = self.proj(x) + + # Reshape back: [(b*t), c, h, w] -> [b, c, t, h, w] + x = x.view(batch_size, time, channels, height, width) + x = x.permute(0, 2, 1, 3, 4) + + return x + identity + + +class QwenImage21MidBlock(nn.Module): + """ + Middle block for QwenVAE encoder and decoder. + + Args: + dim (int): Number of input/output channels. + dropout (float): Dropout rate. + """ + + def __init__(self, dim: int, dropout: float = 0.0, num_layers: int = 1): + super().__init__() + self.dim = dim + + # Create the components + resnets = [QwenImage21ResidualBlock(dim, dim, dropout)] + attentions = [] + for _ in range(num_layers): + attentions.append(QwenImage21AttentionBlock(dim)) + resnets.append(QwenImage21ResidualBlock(dim, dim, dropout)) + self.attentions = nn.ModuleList(attentions) + self.resnets = nn.ModuleList(resnets) + + self.gradient_checkpointing = False + + def forward(self, x, feat_cache=None, feat_idx=None) -> torch.Tensor: + if feat_idx is None: + feat_idx = [0] + # First residual block + x = self.resnets[0](x, feat_cache=feat_cache, feat_idx=feat_idx) + + # Process through attention and residual blocks + for attn, resnet in zip(self.attentions, self.resnets[1:]): + if attn is not None: + x = attn(x) + + x = resnet(x, feat_cache=feat_cache, feat_idx=feat_idx) + + return x + + +class QwenImage21ResidualDownBlock(nn.Module): + def __init__(self, in_dim, out_dim, dropout, num_res_blocks, temperal_downsample=False, down_flag=False): + super().__init__() + + # Shortcut path with downsample + self.avg_shortcut = QwenImage21AvgDown3D( + in_dim, + out_dim, + factor_t=2 if temperal_downsample else 1, + factor_s=2 if down_flag else 1, + ) + + # Main path with residual blocks and downsample + resnets = [] + for _ in range(num_res_blocks): + resnets.append(QwenImage21ResidualBlock(in_dim, out_dim, dropout)) + in_dim = out_dim + self.resnets = nn.ModuleList(resnets) + + # Add the final downsample block + if down_flag: + mode = "downsample3d" if temperal_downsample else "downsample2d" + self.downsampler = QwenImage21Resample(out_dim, mode=mode) + else: + self.downsampler = None + + def forward(self, x, feat_cache=None, feat_idx=None) -> torch.Tensor: + if feat_idx is None: + feat_idx = [0] + x_copy = x.clone() + for resnet in self.resnets: + x = resnet(x, feat_cache=feat_cache, feat_idx=feat_idx) + if self.downsampler is not None: + x = self.downsampler(x, feat_cache=feat_cache, feat_idx=feat_idx) + + return x + self.avg_shortcut(x_copy) + + +class QwenImage21Encoder3d(nn.Module): + r""" + A 3D encoder module. + + Args: + dim (int): The base number of channels in the first layer. + z_dim (int): The dimensionality of the latent space. + dim_mult (list of int): Multipliers for the number of channels in each block. + num_res_blocks (int): Number of residual blocks in each block. + attn_scales (list of float): Scales at which to apply attention mechanisms. + temperal_downsample (list of bool): Whether to downsample temporally in each block. + dropout (float): Dropout rate for the dropout layers. + """ + + def __init__( + self, + in_channels: int = 3, + dim=128, + z_dim=4, + dim_mult=[1, 2, 4, 4], + num_res_blocks=2, + attn_scales=[], + temperal_downsample=[True, True, False], + dropout=0.0, + is_residual: bool = False, + ): + super().__init__() + self.dim = dim + self.z_dim = z_dim + self.dim_mult = dim_mult + self.num_res_blocks = num_res_blocks + self.attn_scales = attn_scales + self.temperal_downsample = temperal_downsample + self.nonlinearity = nn.SiLU() + + # dimensions + dims = [dim * u for u in [1] + dim_mult] + scale = 1.0 + + # init block + self.conv_in = QwenImage21CausalConv3d(in_channels, dims[0], 3, padding=1) + + # downsample blocks + self.down_blocks = nn.ModuleList([]) + for i, (in_dim, out_dim) in enumerate(zip(dims[:-1], dims[1:])): + # residual (+attention) blocks + if is_residual: + self.down_blocks.append( + QwenImage21ResidualDownBlock( + in_dim, + out_dim, + dropout, + num_res_blocks, + temperal_downsample=temperal_downsample[i] if i != len(dim_mult) - 1 else False, + down_flag=i != len(dim_mult) - 1, + ) + ) + else: + for _ in range(num_res_blocks): + self.down_blocks.append(QwenImage21ResidualBlock(in_dim, out_dim, dropout)) + if scale in attn_scales: + self.down_blocks.append(QwenImage21AttentionBlock(out_dim)) + in_dim = out_dim + + # downsample block + if i != len(dim_mult) - 1: + mode = "downsample3d" if temperal_downsample[i] else "downsample2d" + self.down_blocks.append(QwenImage21Resample(out_dim, mode=mode)) + scale /= 2.0 + + # middle blocks + self.mid_block = QwenImage21MidBlock(out_dim, dropout, num_layers=1) + + # output blocks + self.norm_out = QwenImage21RMS_norm(out_dim, images=False) + self.conv_out = QwenImage21CausalConv3d(out_dim, z_dim, 3, padding=1) + + self.gradient_checkpointing = False + + def forward(self, x, feat_cache=None, feat_idx=None) -> torch.Tensor: + if feat_idx is None: + feat_idx = [0] + if feat_cache is not None: + idx = feat_idx[0] + cache_x = x[:, :, -CACHE_T:, :, :].clone() + if cache_x.shape[2] < 2 and feat_cache[idx] is not None: + # cache last frame of last two chunk + cache_x = torch.cat([feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x], dim=2) + x = self.conv_in(x, feat_cache[idx]) + feat_cache[idx] = cache_x + feat_idx[0] += 1 + else: + x = self.conv_in(x) + + ## downsamples + for layer in self.down_blocks: + if feat_cache is not None: + x = layer(x, feat_cache=feat_cache, feat_idx=feat_idx) + else: + x = layer(x) + + ## middle + x = self.mid_block(x, feat_cache=feat_cache, feat_idx=feat_idx) + + ## head + x = self.norm_out(x) + x = self.nonlinearity(x) + if feat_cache is not None: + idx = feat_idx[0] + cache_x = x[:, :, -CACHE_T:, :, :].clone() + if cache_x.shape[2] < 2 and feat_cache[idx] is not None: + # cache last frame of last two chunk + cache_x = torch.cat([feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x], dim=2) + x = self.conv_out(x, feat_cache[idx]) + feat_cache[idx] = cache_x + feat_idx[0] += 1 + else: + x = self.conv_out(x) + + return x + + +class QwenImage21ResidualUpBlock(nn.Module): + """ + A block that handles upsampling for the QwenVAE decoder. + + Args: + in_dim (int): Input dimension + out_dim (int): Output dimension + num_res_blocks (int): Number of residual blocks + dropout (float): Dropout rate + temperal_upsample (bool): Whether to upsample on temporal dimension + up_flag (bool): Whether to upsample or not + """ + + def __init__( + self, + in_dim: int, + out_dim: int, + num_res_blocks: int, + dropout: float = 0.0, + temperal_upsample: bool = False, + up_flag: bool = False, + ): + super().__init__() + self.in_dim = in_dim + self.out_dim = out_dim + + if up_flag: + self.avg_shortcut = QwenImage21DupUp3D( + in_dim, + out_dim, + factor_t=2 if temperal_upsample else 1, + factor_s=2, + ) + else: + self.avg_shortcut = None + + # create residual blocks + resnets = [] + current_dim = in_dim + for _ in range(num_res_blocks + 1): + resnets.append(QwenImage21ResidualBlock(current_dim, out_dim, dropout)) + current_dim = out_dim + + self.resnets = nn.ModuleList(resnets) + + # Add upsampling layer if needed + if up_flag: + upsample_mode = "upsample3d" if temperal_upsample else "upsample2d" + self.upsampler = QwenImage21Resample(out_dim, mode=upsample_mode, upsample_out_dim=out_dim) + else: + self.upsampler = None + + self.gradient_checkpointing = False + + def forward(self, x, feat_cache=None, feat_idx=None, first_chunk=False) -> torch.Tensor: + if feat_idx is None: + feat_idx = [0] + """ + Forward pass through the upsampling block. + + Args: + x (torch.Tensor): Input tensor + feat_cache (list, optional): Feature cache for causal convolutions + feat_idx (list, optional): Feature index for cache management + + Returns: + torch.Tensor: Output tensor + """ + x_copy = x.clone() + + for resnet in self.resnets: + if feat_cache is not None: + x = resnet(x, feat_cache=feat_cache, feat_idx=feat_idx) + else: + x = resnet(x) + + if self.upsampler is not None: + if feat_cache is not None: + x = self.upsampler(x, feat_cache=feat_cache, feat_idx=feat_idx) + else: + x = self.upsampler(x) + + if self.avg_shortcut is not None: + x = x + self.avg_shortcut(x_copy, first_chunk=first_chunk) + + return x + + +class QwenImage21UpBlock(nn.Module): + """ + A block that handles upsampling for the QwenVAE decoder. + + Args: + in_dim (int): Input dimension + out_dim (int): Output dimension + num_res_blocks (int): Number of residual blocks + dropout (float): Dropout rate + upsample_mode (str, optional): Mode for upsampling ('upsample2d' or 'upsample3d') + """ + + def __init__( + self, + in_dim: int, + out_dim: int, + num_res_blocks: int, + dropout: float = 0.0, + upsample_mode: str | None = None, + ): + super().__init__() + self.in_dim = in_dim + self.out_dim = out_dim + + # Create layers list + resnets = [] + # Add residual blocks and attention if needed + current_dim = in_dim + for _ in range(num_res_blocks + 1): + resnets.append(QwenImage21ResidualBlock(current_dim, out_dim, dropout)) + current_dim = out_dim + + self.resnets = nn.ModuleList(resnets) + + # Add upsampling layer if needed + self.upsamplers = None + if upsample_mode is not None: + self.upsamplers = nn.ModuleList([QwenImage21Resample(out_dim, mode=upsample_mode)]) + + self.gradient_checkpointing = False + + def forward(self, x, feat_cache=None, feat_idx=None, first_chunk=None) -> torch.Tensor: + if feat_idx is None: + feat_idx = [0] + """ + Forward pass through the upsampling block. + + Args: + x (torch.Tensor): Input tensor + feat_cache (list, optional): Feature cache for causal convolutions + feat_idx (list, optional): Feature index for cache management + + Returns: + torch.Tensor: Output tensor + """ + for resnet in self.resnets: + if feat_cache is not None: + x = resnet(x, feat_cache=feat_cache, feat_idx=feat_idx) + else: + x = resnet(x) + + if self.upsamplers is not None: + if feat_cache is not None: + x = self.upsamplers[0](x, feat_cache=feat_cache, feat_idx=feat_idx) + else: + x = self.upsamplers[0](x) + return x + + +class QwenImage21Decoder3d(nn.Module): + r""" + A 3D decoder module. + + Args: + dim (int): The base number of channels in the first layer. + z_dim (int): The dimensionality of the latent space. + dim_mult (list of int): Multipliers for the number of channels in each block. + num_res_blocks (int): Number of residual blocks in each block. + attn_scales (list of float): Scales at which to apply attention mechanisms. + temperal_upsample (list of bool): Whether to upsample temporally in each block. + dropout (float): Dropout rate for the dropout layers. + """ + + def __init__( + self, + dim=128, + z_dim=4, + dim_mult=[1, 2, 4, 4], + num_res_blocks=2, + attn_scales=[], + temperal_upsample=[False, True, True], + dropout=0.0, + out_channels: int = 3, + is_residual: bool = False, + ): + super().__init__() + self.dim = dim + self.z_dim = z_dim + self.dim_mult = dim_mult + self.num_res_blocks = num_res_blocks + self.attn_scales = attn_scales + self.temperal_upsample = temperal_upsample + + self.nonlinearity = nn.SiLU() + + # dimensions + dims = [dim * u for u in [dim_mult[-1]] + dim_mult[::-1]] + + # init block + self.conv_in = QwenImage21CausalConv3d(z_dim, dims[0], 3, padding=1) + + # middle blocks + self.mid_block = QwenImage21MidBlock(dims[0], dropout, num_layers=1) + + # upsample blocks + self.up_blocks = nn.ModuleList([]) + for i, (in_dim, out_dim) in enumerate(zip(dims[:-1], dims[1:])): + # residual (+attention) blocks + if i > 0 and not is_residual: + in_dim = in_dim // 2 + + # determine if we need upsampling + up_flag = i != len(dim_mult) - 1 + # determine upsampling mode, if not upsampling, set to None + upsample_mode = None + if up_flag and temperal_upsample[i]: + upsample_mode = "upsample3d" + elif up_flag: + upsample_mode = "upsample2d" + # Create and add the upsampling block + if is_residual: + up_block = QwenImage21ResidualUpBlock( + in_dim=in_dim, + out_dim=out_dim, + num_res_blocks=num_res_blocks, + dropout=dropout, + temperal_upsample=temperal_upsample[i] if up_flag else False, + up_flag=up_flag, + ) + else: + up_block = QwenImage21UpBlock( + in_dim=in_dim, + out_dim=out_dim, + num_res_blocks=num_res_blocks, + dropout=dropout, + upsample_mode=upsample_mode, + ) + self.up_blocks.append(up_block) + + # output blocks + self.norm_out = QwenImage21RMS_norm(out_dim, images=False) + self.conv_out = QwenImage21CausalConv3d(out_dim, out_channels, 3, padding=1) + + self.gradient_checkpointing = False + + def forward(self, x, feat_cache=None, feat_idx=None, first_chunk=False) -> torch.Tensor: + if feat_idx is None: + feat_idx = [0] + ## conv1 + if feat_cache is not None: + idx = feat_idx[0] + cache_x = x[:, :, -CACHE_T:, :, :].clone() + if cache_x.shape[2] < 2 and feat_cache[idx] is not None: + # cache last frame of last two chunk + cache_x = torch.cat([feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x], dim=2) + x = self.conv_in(x, feat_cache[idx]) + feat_cache[idx] = cache_x + feat_idx[0] += 1 + else: + x = self.conv_in(x) + + ## middle + x = self.mid_block(x, feat_cache=feat_cache, feat_idx=feat_idx) + + ## upsamples + for up_block in self.up_blocks: + x = up_block(x, feat_cache=feat_cache, feat_idx=feat_idx, first_chunk=first_chunk) + + ## head + x = self.norm_out(x) + x = self.nonlinearity(x) + if feat_cache is not None: + idx = feat_idx[0] + cache_x = x[:, :, -CACHE_T:, :, :].clone() + if cache_x.shape[2] < 2 and feat_cache[idx] is not None: + # cache last frame of last two chunk + cache_x = torch.cat([feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x], dim=2) + x = self.conv_out(x, feat_cache[idx]) + feat_cache[idx] = cache_x + feat_idx[0] += 1 + else: + x = self.conv_out(x) + return x + + +# Copied from diffusers.models.autoencoders.autoencoder_kl_wan.patchify with patchify->_patchify +def _patchify(x, patch_size): + if patch_size == 1: + return x + + if x.dim() != 5: + raise ValueError(f"Invalid input shape: {x.shape}") + # x shape: [batch_size, channels, frames, height, width] + batch_size, channels, frames, height, width = x.shape + + # Ensure height and width are divisible by patch_size + if height % patch_size != 0 or width % patch_size != 0: + raise ValueError(f"Height ({height}) and width ({width}) must be divisible by patch_size ({patch_size})") + + # Reshape to [batch_size, channels, frames, height//patch_size, patch_size, width//patch_size, patch_size] + x = x.view(batch_size, channels, frames, height // patch_size, patch_size, width // patch_size, patch_size) + + # Rearrange to [batch_size, channels * patch_size * patch_size, frames, height//patch_size, width//patch_size] + x = x.permute(0, 1, 6, 4, 2, 3, 5).contiguous() + x = x.view(batch_size, channels * patch_size * patch_size, frames, height // patch_size, width // patch_size) + + return x + + +# Copied from diffusers.models.autoencoders.autoencoder_kl_wan.unpatchify with unpatchify->_unpatchify +def _unpatchify(x, patch_size): + if patch_size == 1: + return x + + if x.dim() != 5: + raise ValueError(f"Invalid input shape: {x.shape}") + # x shape: [batch_size, (channels * patch_size * patch_size), frame, height, width] + batch_size, c_patches, frames, height, width = x.shape + channels = c_patches // (patch_size * patch_size) + + # Reshape to [b, c, patch_size, patch_size, f, h, w] + x = x.view(batch_size, channels, patch_size, patch_size, frames, height, width) + + # Rearrange to [b, c, f, h * patch_size, w * patch_size] + x = x.permute(0, 1, 4, 5, 3, 6, 2).contiguous() + x = x.view(batch_size, channels, frames, height * patch_size, width * patch_size) + + return x + + +class QwenImageVAE(nn.Module): + r""" + Qwen-Image 2.1 single-frame VAE with reference-compatible encode/decode and overlap tiling. + + Inputs and latents use (batch, channels, 1, height, width). Latent normalization remains the caller's job; + the checkpoint's latents_mean and latents_std are exposed in config. + """ + + @model_config + def __init__( + self, + base_dim: int = 96, + decoder_base_dim: int | None = 144, + z_dim: int = 64, + dim_mult: list[int] = [1, 2, 4, 8, 8], + num_res_blocks: int = 2, + attn_scales: list[float] = [], + temperal_downsample: list[bool] = [False, True, True, True], + dropout: float = 0.0, + latents_mean: list[float] = [ + 0.5126, + 0.7721, + -0.0631, + 1.3506, + -0.7855, + -2.1025, + -0.3458, + 1.3722, + 1.8873, + -1.7177, + -0.6510, + 0.2732, + 0.7562, + -0.6163, + -1.0277, + 3.8363, + 2.0210, + 0.0472, + 0.9320, + 2.0087, + 2.4954, + -0.1391, + -1.4249, + 1.8464, + -0.5236, + 1.2826, + 3.7046, + -1.3035, + 2.7286, + -1.4518, + -1.9036, + -1.9955, + -0.0342, + -1.0265, + -0.7636, + 3.0555, + 0.0746, + -3.0751, + -0.1076, + 1.7376, + -1.0914, + -1.9435, + -0.2784, + -1.3680, + 0.4809, + -0.4433, + 0.3764, + 0.5729, + -2.0595, + 1.0960, + -1.3260, + -2.0211, + -5.0179, + 0.5275, + 4.0162, + 1.8505, + 0.3026, + 1.9373, + 1.4937, + 0.2632, + 0.5547, + -1.7121, + -0.1562, + 0.0304, + ], + latents_std: list[float] = [ + 3.2001, + 3.2936, + 3.4321, + 3.0091, + 3.1061, + 4.0379, + 4.0705, + 3.7910, + 3.0785, + 3.6500, + 3.9308, + 3.0904, + 2.8778, + 3.7675, + 3.7320, + 5.0756, + 3.2864, + 4.0397, + 3.1317, + 4.0443, + 2.9249, + 3.9454, + 3.0988, + 4.2489, + 3.4896, + 3.8513, + 3.9323, + 3.4719, + 3.7498, + 4.2830, + 3.5694, + 4.2467, + 3.9037, + 3.2947, + 5.0770, + 3.5075, + 3.2700, + 3.4767, + 2.8063, + 5.1125, + 3.5327, + 4.7833, + 3.1286, + 4.1819, + 3.8527, + 3.8312, + 3.5605, + 4.3875, + 3.9624, + 4.0168, + 3.5643, + 4.0550, + 5.5614, + 4.2963, + 4.4080, + 3.4959, + 3.8747, + 3.7608, + 3.5735, + 3.1490, + 3.7662, + 3.6746, + 3.4563, + 3.8161, + ], + is_residual: bool = True, + in_channels: int = 4, + out_channels: int = 4, + patch_size: int | None = None, + scale_factor_temporal: int | None = 8, + scale_factor_spatial: int | None = 16, + ) -> None: + super().__init__() + + self.z_dim = z_dim + self.temperal_downsample = temperal_downsample + self.temperal_upsample = temperal_downsample[::-1] + + if decoder_base_dim is None: + decoder_base_dim = base_dim + + self.encoder = QwenImage21Encoder3d( + in_channels=in_channels, + dim=base_dim, + z_dim=z_dim * 2, + dim_mult=dim_mult, + num_res_blocks=num_res_blocks, + attn_scales=attn_scales, + temperal_downsample=temperal_downsample, + dropout=dropout, + is_residual=is_residual, + ) + self.quant_conv = QwenImage21CausalConv3d(z_dim * 2, z_dim * 2, 1) + self.post_quant_conv = QwenImage21CausalConv3d(z_dim, z_dim, 1) + + self.decoder = QwenImage21Decoder3d( + dim=decoder_base_dim, + z_dim=z_dim, + dim_mult=dim_mult, + num_res_blocks=num_res_blocks, + attn_scales=attn_scales, + temperal_upsample=self.temperal_upsample, + dropout=dropout, + out_channels=out_channels, + is_residual=is_residual, + ) + + self.spatial_compression_ratio = scale_factor_spatial + + # When decoding a batch of video latents at a time, one can save memory by slicing across the batch dimension + # to perform decoding of a single video latent at a time. + self.use_slicing = False + + # When decoding spatially large video latents, the memory requirement is very high. By breaking the video latent + # frames spatially into smaller tiles and performing multiple forward passes for decoding, and then blending the + # intermediate tiles together, the memory requirement can be lowered. + self.use_tiling = False + + # The minimal tile height and width for spatial tiling to be used + self.tile_sample_min_height = 256 + self.tile_sample_min_width = 256 + + # The minimal distance between two spatial tiles + self.tile_sample_stride_height = 192 + self.tile_sample_stride_width = 192 + + # Precompute and cache conv counts for encoder and decoder for clear_cache speedup + self._cached_conv_counts = { + "decoder": sum(isinstance(m, QwenImage21CausalConv3d) for m in self.decoder.modules()) + if self.decoder is not None + else 0, + "encoder": sum(isinstance(m, QwenImage21CausalConv3d) for m in self.encoder.modules()) + if self.encoder is not None + else 0, + } + + # Copied from diffusers.models.autoencoders.autoencoder_kl_wan.AutoencoderKLWan.enable_tiling + def enable_tiling( + self, + tile_sample_min_height: int | None = None, + tile_sample_min_width: int | None = None, + tile_sample_stride_height: float | None = None, + tile_sample_stride_width: float | None = None, + ) -> None: + r""" + Enable tiled VAE decoding. When this option is enabled, the VAE will split the input tensor into tiles to + compute decoding and encoding in several steps. This is useful for saving a large amount of memory and to allow + processing larger images. + + Args: + tile_sample_min_height (`int`, *optional*): + The minimum height required for a sample to be separated into tiles across the height dimension. + tile_sample_min_width (`int`, *optional*): + The minimum width required for a sample to be separated into tiles across the width dimension. + tile_sample_stride_height (`int`, *optional*): + The minimum amount of overlap between two consecutive vertical tiles. This is to ensure that there are + no tiling artifacts produced across the height dimension. + tile_sample_stride_width (`int`, *optional*): + The stride between two consecutive horizontal tiles. This is to ensure that there are no tiling + artifacts produced across the width dimension. + """ + self.use_tiling = True + self.tile_sample_min_height = tile_sample_min_height or self.tile_sample_min_height + self.tile_sample_min_width = tile_sample_min_width or self.tile_sample_min_width + self.tile_sample_stride_height = tile_sample_stride_height or self.tile_sample_stride_height + self.tile_sample_stride_width = tile_sample_stride_width or self.tile_sample_stride_width + + def disable_tiling(self): + self.use_tiling = False + + def enable_slicing(self): + self.use_slicing = True + + def disable_slicing(self): + self.use_slicing = False + + @classmethod + def from_comfy(cls, path, device="cpu", *, config=None): + """Load Comfy BF16 weights without allocating an initialized full VAE.""" + from .qwen_weights import load_vae + return load_vae(path, device=device, config=config) + + # Copied from diffusers.models.autoencoders.autoencoder_kl_wan.AutoencoderKLWan.clear_cache + def clear_cache(self): + # Use cached conv counts for decoder and encoder to avoid re-iterating modules each call + self._conv_num = self._cached_conv_counts["decoder"] + self._conv_idx = [0] + self._feat_map = [None] * self._conv_num + # cache encode + self._enc_conv_num = self._cached_conv_counts["encoder"] + self._enc_conv_idx = [0] + self._enc_feat_map = [None] * self._enc_conv_num + + # Copied from diffusers.models.autoencoders.autoencoder_kl_wan.AutoencoderKLWan._encode with patchify->_patchify + def _encode(self, x: torch.Tensor): + _, _, num_frame, height, width = x.shape + + self.clear_cache() + if self.config.patch_size is not None: + x = _patchify(x, patch_size=self.config.patch_size) + + if self.use_tiling and (width > self.tile_sample_min_width or height > self.tile_sample_min_height): + return self.tiled_encode(x) + + iter_ = 1 + (num_frame - 1) // 4 + for i in range(iter_): + self._enc_conv_idx = [0] + if i == 0: + out = self.encoder(x[:, :, :1, :, :], feat_cache=self._enc_feat_map, feat_idx=self._enc_conv_idx) + else: + out_ = self.encoder( + x[:, :, 1 + 4 * (i - 1) : 1 + 4 * i, :, :], + feat_cache=self._enc_feat_map, + feat_idx=self._enc_conv_idx, + ) + out = torch.cat([out, out_], 2) + + enc = self.quant_conv(out) + self.clear_cache() + return enc + + # Copied from diffusers.models.autoencoders.autoencoder_kl_wan.AutoencoderKLWan.encode + def encode( + self, x: torch.Tensor, return_dict: bool = True, *, tile: bool | None = None + ) -> AutoencoderKLOutput | tuple[DiagonalGaussianDistribution]: + r""" + Encode a batch of images into latents. + + Args: + x (`torch.Tensor`): Input batch of images. + return_dict (`bool`, *optional*, defaults to `True`): + Whether to return a [`~models.autoencoder_kl.AutoencoderKLOutput`] instead of a plain tuple. + + Returns: + The latent representations of the encoded videos. If `return_dict` is True, a + [`~models.autoencoder_kl.AutoencoderKLOutput`] is returned, otherwise a plain `tuple` is returned. + """ + previous_tiling = self.use_tiling + if tile is not None: + self.use_tiling = tile + try: + if self.use_slicing and x.shape[0] > 1: + encoded_slices = [self._encode(x_slice) for x_slice in x.split(1)] + h = torch.cat(encoded_slices) + else: + h = self._encode(x) + finally: + self.use_tiling = previous_tiling + posterior = DiagonalGaussianDistribution(h) + + if not return_dict: + return (posterior,) + return AutoencoderKLOutput(latent_dist=posterior) + + # Copied from diffusers.models.autoencoders.autoencoder_kl_wan.AutoencoderKLWan._decode with unpatchify->_unpatchify + def _decode(self, z: torch.Tensor, return_dict: bool = True): + _, _, num_frame, height, width = z.shape + tile_latent_min_height = self.tile_sample_min_height // self.spatial_compression_ratio + tile_latent_min_width = self.tile_sample_min_width // self.spatial_compression_ratio + + if self.use_tiling and (width > tile_latent_min_width or height > tile_latent_min_height): + return self.tiled_decode(z, return_dict=return_dict) + + self.clear_cache() + x = self.post_quant_conv(z) + for i in range(num_frame): + self._conv_idx = [0] + if i == 0: + out = self.decoder( + x[:, :, i : i + 1, :, :], feat_cache=self._feat_map, feat_idx=self._conv_idx, first_chunk=True + ) + else: + out_ = self.decoder(x[:, :, i : i + 1, :, :], feat_cache=self._feat_map, feat_idx=self._conv_idx) + out = torch.cat([out, out_], 2) + + if self.config.patch_size is not None: + out = _unpatchify(out, patch_size=self.config.patch_size) + + out = torch.clamp(out, min=-1.0, max=1.0) + + self.clear_cache() + if not return_dict: + return (out,) + + return DecoderOutput(sample=out) + + # Copied from diffusers.models.autoencoders.autoencoder_kl_wan.AutoencoderKLWan.decode + def decode(self, z: torch.Tensor, return_dict: bool = True, *, tile: bool | None = None) -> DecoderOutput | torch.Tensor: + r""" + Decode a batch of images. + + Args: + z (`torch.Tensor`): Input batch of latent vectors. + return_dict (`bool`, *optional*, defaults to `True`): + Whether to return a [`~models.vae.DecoderOutput`] instead of a plain tuple. + + Returns: + [`~models.vae.DecoderOutput`] or `tuple`: + If return_dict is True, a [`~models.vae.DecoderOutput`] is returned, otherwise a plain `tuple` is + returned. + """ + previous_tiling = self.use_tiling + if tile is not None: + self.use_tiling = tile + try: + if self.use_slicing and z.shape[0] > 1: + decoded_slices = [self._decode(z_slice).sample for z_slice in z.split(1)] + decoded = torch.cat(decoded_slices) + else: + decoded = self._decode(z).sample + finally: + self.use_tiling = previous_tiling + + if not return_dict: + return (decoded,) + return DecoderOutput(sample=decoded) + + # Copied from diffusers.models.autoencoders.autoencoder_kl_wan.AutoencoderKLWan.blend_v + def blend_v(self, a: torch.Tensor, b: torch.Tensor, blend_extent: int) -> torch.Tensor: + blend_extent = min(a.shape[-2], b.shape[-2], blend_extent) + for y in range(blend_extent): + b[:, :, :, y, :] = a[:, :, :, -blend_extent + y, :] * (1 - y / blend_extent) + b[:, :, :, y, :] * ( + y / blend_extent + ) + return b + + # Copied from diffusers.models.autoencoders.autoencoder_kl_wan.AutoencoderKLWan.blend_h + def blend_h(self, a: torch.Tensor, b: torch.Tensor, blend_extent: int) -> torch.Tensor: + blend_extent = min(a.shape[-1], b.shape[-1], blend_extent) + for x in range(blend_extent): + b[:, :, :, :, x] = a[:, :, :, :, -blend_extent + x] * (1 - x / blend_extent) + b[:, :, :, :, x] * ( + x / blend_extent + ) + return b + + # Copied from diffusers.models.autoencoders.autoencoder_kl_wan.AutoencoderKLWan.tiled_encode + def tiled_encode(self, x: torch.Tensor) -> AutoencoderKLOutput: + r"""Encode a batch of images using a tiled encoder. + + Args: + x (`torch.Tensor`): Input batch of videos. + + Returns: + `torch.Tensor`: + The latent representation of the encoded videos. + """ + + _, _, num_frames, height, width = x.shape + encode_spatial_compression_ratio = self.spatial_compression_ratio + if self.config.patch_size is not None: + assert encode_spatial_compression_ratio % self.config.patch_size == 0 + encode_spatial_compression_ratio = self.spatial_compression_ratio // self.config.patch_size + + latent_height = height // encode_spatial_compression_ratio + latent_width = width // encode_spatial_compression_ratio + + tile_latent_min_height = self.tile_sample_min_height // encode_spatial_compression_ratio + tile_latent_min_width = self.tile_sample_min_width // encode_spatial_compression_ratio + tile_latent_stride_height = self.tile_sample_stride_height // encode_spatial_compression_ratio + tile_latent_stride_width = self.tile_sample_stride_width // encode_spatial_compression_ratio + + blend_height = tile_latent_min_height - tile_latent_stride_height + blend_width = tile_latent_min_width - tile_latent_stride_width + + # Split x into overlapping tiles and encode them separately. + # The tiles have an overlap to avoid seams between tiles. + rows = [] + for i in range(0, height, self.tile_sample_stride_height): + row = [] + for j in range(0, width, self.tile_sample_stride_width): + self.clear_cache() + time = [] + frame_range = 1 + (num_frames - 1) // 4 + for k in range(frame_range): + self._enc_conv_idx = [0] + if k == 0: + tile = x[:, :, :1, i : i + self.tile_sample_min_height, j : j + self.tile_sample_min_width] + else: + tile = x[ + :, + :, + 1 + 4 * (k - 1) : 1 + 4 * k, + i : i + self.tile_sample_min_height, + j : j + self.tile_sample_min_width, + ] + tile = self.encoder(tile, feat_cache=self._enc_feat_map, feat_idx=self._enc_conv_idx) + tile = self.quant_conv(tile) + time.append(tile) + row.append(torch.cat(time, dim=2)) + rows.append(row) + self.clear_cache() + + result_rows = [] + for i, row in enumerate(rows): + result_row = [] + for j, tile in enumerate(row): + # blend the above tile and the left tile + # to the current tile and add the current tile to the result row + if i > 0: + tile = self.blend_v(rows[i - 1][j], tile, blend_height) + if j > 0: + tile = self.blend_h(row[j - 1], tile, blend_width) + result_row.append(tile[:, :, :, :tile_latent_stride_height, :tile_latent_stride_width]) + result_rows.append(torch.cat(result_row, dim=-1)) + + enc = torch.cat(result_rows, dim=3)[:, :, :, :latent_height, :latent_width] + return enc + + # Copied from diffusers.models.autoencoders.autoencoder_kl_wan.AutoencoderKLWan.tiled_decode with unpatchify->_unpatchify + def tiled_decode(self, z: torch.Tensor, return_dict: bool = True) -> DecoderOutput | torch.Tensor: + r""" + Decode a batch of images using a tiled decoder. + + Args: + z (`torch.Tensor`): Input batch of latent vectors. + return_dict (`bool`, *optional*, defaults to `True`): + Whether or not to return a [`~models.vae.DecoderOutput`] instead of a plain tuple. + + Returns: + [`~models.vae.DecoderOutput`] or `tuple`: + If return_dict is True, a [`~models.vae.DecoderOutput`] is returned, otherwise a plain `tuple` is + returned. + """ + _, _, num_frames, height, width = z.shape + sample_height = height * self.spatial_compression_ratio + sample_width = width * self.spatial_compression_ratio + + tile_latent_min_height = self.tile_sample_min_height // self.spatial_compression_ratio + tile_latent_min_width = self.tile_sample_min_width // self.spatial_compression_ratio + tile_latent_stride_height = self.tile_sample_stride_height // self.spatial_compression_ratio + tile_latent_stride_width = self.tile_sample_stride_width // self.spatial_compression_ratio + tile_sample_stride_height = self.tile_sample_stride_height + tile_sample_stride_width = self.tile_sample_stride_width + if self.config.patch_size is not None: + sample_height = sample_height // self.config.patch_size + sample_width = sample_width // self.config.patch_size + tile_sample_stride_height = tile_sample_stride_height // self.config.patch_size + tile_sample_stride_width = tile_sample_stride_width // self.config.patch_size + blend_height = self.tile_sample_min_height // self.config.patch_size - tile_sample_stride_height + blend_width = self.tile_sample_min_width // self.config.patch_size - tile_sample_stride_width + else: + blend_height = self.tile_sample_min_height - tile_sample_stride_height + blend_width = self.tile_sample_min_width - tile_sample_stride_width + + # Split z into overlapping tiles and decode them separately. + # The tiles have an overlap to avoid seams between tiles. + rows = [] + for i in range(0, height, tile_latent_stride_height): + row = [] + for j in range(0, width, tile_latent_stride_width): + self.clear_cache() + time = [] + for k in range(num_frames): + self._conv_idx = [0] + tile = z[:, :, k : k + 1, i : i + tile_latent_min_height, j : j + tile_latent_min_width] + tile = self.post_quant_conv(tile) + decoded = self.decoder( + tile, feat_cache=self._feat_map, feat_idx=self._conv_idx, first_chunk=(k == 0) + ) + time.append(decoded) + row.append(torch.cat(time, dim=2)) + rows.append(row) + self.clear_cache() + + result_rows = [] + for i, row in enumerate(rows): + result_row = [] + for j, tile in enumerate(row): + # blend the above tile and the left tile + # to the current tile and add the current tile to the result row + if i > 0: + tile = self.blend_v(rows[i - 1][j], tile, blend_height) + if j > 0: + tile = self.blend_h(row[j - 1], tile, blend_width) + result_row.append(tile[:, :, :, :tile_sample_stride_height, :tile_sample_stride_width]) + result_rows.append(torch.cat(result_row, dim=-1)) + dec = torch.cat(result_rows, dim=3)[:, :, :, :sample_height, :sample_width] + + if self.config.patch_size is not None: + dec = _unpatchify(dec, patch_size=self.config.patch_size) + + dec = torch.clamp(dec, min=-1.0, max=1.0) + + if not return_dict: + return (dec,) + return DecoderOutput(sample=dec) + + def forward( + self, + sample: torch.Tensor, + sample_posterior: bool = False, + return_dict: bool = True, + generator: torch.Generator | None = None, + ) -> DecoderOutput | torch.Tensor: + """ + Args: + sample (`torch.Tensor`): Input sample. + sample_posterior (`bool`, *optional*, defaults to `False`): + Whether to sample from the posterior instead of taking its mode. + return_dict (`bool`, *optional*, defaults to `True`): + Whether or not to return a [`DecoderOutput`] instead of a plain tuple. + generator (`torch.Generator`, *optional*): + Generator used when `sample_posterior` is `True`. + + Returns: + [`~models.autoencoders.vae.DecoderOutput`] or `tuple`: + [`~models.autoencoders.vae.DecoderOutput`] if `return_dict` is True, otherwise a plain `tuple`. + """ + x = sample + posterior = self.encode(x).latent_dist + + if sample_posterior: + z = posterior.sample(generator=generator) + else: + z = posterior.mode() + dec = self.decode(z, return_dict=return_dict) + return dec + + +AutoencoderKLQwenImage21 = QwenImageVAE + + +# Comfy Wan2.2-style image-specialized VAE -> this model/Diffusers layout. +# The block/resnet indices are retained, not regrouped like the older Wan VAE. +# Source layout: https://github.com/Comfy-Org/ComfyUI/blob/master/comfy/ldm/wan/vae2_2.py +# Each row is independently testable. Order matters: specific leaves precede +# the generic residual/block/container rows. +VAE_KEY_RENAMES = ( + (r"^conv1\.", "quant_conv."), + (r"^conv2\.", "post_quant_conv."), + (r"^(encoder|decoder)\.conv1\.", r"\1.conv_in."), + (r"\.head\.0\.", ".norm_out."), + (r"\.head\.2\.", ".conv_out."), + (r"\.middle\.0\.", ".mid_block.resnets.0."), + (r"\.middle\.1\.", ".mid_block.attentions.0."), + (r"\.middle\.2\.", ".mid_block.resnets.1."), + (r"\.residual\.0\.", ".norm1."), + (r"\.residual\.2\.", ".conv1."), + (r"\.residual\.3\.", ".norm2."), + (r"\.residual\.6\.", ".conv2."), + (r"\.shortcut\.", ".conv_shortcut."), + (r"^encoder\.downsamples\.(\d+)\.downsamples\.(\d+)\.(resample|time_conv)\.", r"encoder.down_blocks.\1.downsampler.\3."), + (r"^encoder\.downsamples\.(\d+)\.downsamples\.(\d+)\.", r"encoder.down_blocks.\1.resnets.\2."), + (r"^decoder\.upsamples\.(\d+)\.upsamples\.(\d+)\.(resample|time_conv)\.", r"decoder.up_blocks.\1.upsampler.\3."), + (r"^decoder\.upsamples\.(\d+)\.upsamples\.(\d+)\.", r"decoder.up_blocks.\1.resnets.\2."), +) + + +def comfy_vae_key_to_diffusers(key): + """Translate a name; loaders additionally reject keys absent from the model.""" + for pattern, replacement in VAE_KEY_RENAMES: + key = re.sub(pattern, replacement, key) + return key + + +def vae_key_map(model): + """Complete Comfy -> model key table for this VAE's configured block counts.""" + if not model.config.is_residual: + raise ValueError("Comfy Qwen-Image 2.1 VAE mapping requires residual blocks") + result = {} + for target in model.state_dict(): + source = target + if source.startswith("quant_conv."): + source = source.replace("quant_conv.", "conv1.", 1) + elif source.startswith("post_quant_conv."): + source = source.replace("post_quant_conv.", "conv2.", 1) + else: + source = re.sub(r"^(encoder|decoder)\.conv_in\.", r"\1.conv1.", source) + source = source.replace(".norm_out.", ".head.0.").replace(".conv_out.", ".head.2.") + source = source.replace(".mid_block.resnets.0.", ".middle.0.") + source = source.replace(".mid_block.attentions.0.", ".middle.1.") + source = source.replace(".mid_block.resnets.1.", ".middle.2.") + source = re.sub(r"^encoder\.down_blocks\.(\d+)\.resnets\.(\d+)\.", r"encoder.downsamples.\1.downsamples.\2.", source) + source = re.sub(r"^decoder\.up_blocks\.(\d+)\.resnets\.(\d+)\.", r"decoder.upsamples.\1.upsamples.\2.", source) + source = re.sub(r"^encoder\.down_blocks\.(\d+)\.downsampler\.", lambda m: f"encoder.downsamples.{m[1]}.downsamples.{model.config.num_res_blocks}.", source) + source = re.sub(r"^decoder\.up_blocks\.(\d+)\.upsampler\.", lambda m: f"decoder.upsamples.{m[1]}.upsamples.{model.config.num_res_blocks + 1}.", source) + # Root conv1 is not a residual conv. Replace internal leaves only. + if ".middle." in source or ".downsamples." in source or ".upsamples." in source: + source = source.replace(".norm1.", ".residual.0.").replace(".conv1.", ".residual.2.") + source = source.replace(".norm2.", ".residual.3.").replace(".conv2.", ".residual.6.") + source = source.replace(".conv_shortcut.", ".shortcut.") + if source in result or comfy_vae_key_to_diffusers(source) != target: + raise ValueError(f"VAE key mapping is not one-to-one: {source} -> {target}") + result[source] = target + return result diff --git a/freevideo_engine/image_gen/qwen_weights.py b/freevideo_engine/image_gen/qwen_weights.py new file mode 100644 index 00000000..36ce06bb --- /dev/null +++ b/freevideo_engine/image_gen/qwen_weights.py @@ -0,0 +1,223 @@ +"""Strict, tensor-at-a-time loaders for the original Comfy Qwen-Image files. + +The DiT skeleton is created on meta. Quantized rows are never dequantized or +requantized; the fused gate_up is sliced gate-first, as verified by the B200 +prototype. Nonquantized parameters are BF16; row scales remain FP32. +""" +import inspect +import json + +import torch +from torch import nn + +from ..streamed_weights import _open_safetensors +from .qwen_dit import QwenImageDiT +from .qwen_vae import QwenImageVAE, comfy_vae_key_to_diffusers, vae_key_map + + +CONVROT_METADATA = {"format": "int8_tensorwise", "convrot": True, "convrot_groupsize": 256} +COMFY_ATTENTION_HINT = {"attention": "comfy_kitchen_int8"} + + +def _config(metadata, model_class, override): + """Optional JSON `config` metadata; absent metadata uses official 2.1 defaults.""" + try: + value = json.loads(metadata["config"]) if "config" in metadata else {} + except (ValueError, TypeError) as exc: + raise ValueError("Checkpoint config metadata must be a JSON object") from exc + if not isinstance(value, dict): + raise ValueError("Checkpoint config metadata must be a JSON object") + value = {key: item for key, item in value.items() if not key.startswith("_")} + if override is not None: + value.update(override) + allowed = set(inspect.signature(model_class.__init__).parameters) - {"self"} + unknown = set(value) - allowed + if unknown: + raise ValueError(f"Unknown {model_class.__name__} config fields: {sorted(unknown)}") + return value + + +def _keys(checkpoint, prefix=""): + result = {} + for disk_key in checkpoint.keys(): + key = disk_key[len(prefix):] if prefix and disk_key.startswith(prefix) else disk_key + if key in result: + raise ValueError(f"Duplicate normalized checkpoint key: {key}") + result[key] = disk_key + return result + + +def _shape_dtype(checkpoint, key): + tensor_slice = checkpoint.get_slice(key) + return tuple(tensor_slice.get_shape()), tensor_slice.get_dtype() + + +def _check_tensor(checkpoint, key, shape, dtypes): + actual_shape, dtype = _shape_dtype(checkpoint, key) + if actual_shape != tuple(shape) or dtype not in dtypes: + raise ValueError(f"Unsupported tensor {key}: shape={actual_shape}, dtype={dtype}; expected {tuple(shape)}, {dtypes}") + + +def _assign(model, name, tensor, device, dtype): + """Install only this tensor; no full state_dict or initialized weight copy.""" + parent, _, child = name.rpartition(".") + module = model.get_submodule(parent) if parent else model + value = tensor.to(device=device, dtype=dtype).contiguous() + if child in module._parameters: + module._parameters[child] = nn.Parameter(value, requires_grad=False) + elif child in module._buffers: + module._buffers[child] = value + else: + raise ValueError(f"Checkpoint target is not a parameter or buffer: {name}") + + +def _quant_targets(model, name): + if name.endswith("img_mlp.gate_up"): + prefix = name[:-len("gate_up")] + targets = [prefix + "gate_layer", prefix + "proj"] + else: + targets = [name] + modules = [] + for target in targets: + try: + module = model.get_submodule(target) + except AttributeError as exc: + raise ValueError(f"Unknown quantized DiT layer: {name}") from exc + if not isinstance(module, nn.Linear) or module.in_features % 256: + raise ValueError(f"ConvRot target must be a Linear with group-256 input: {name}") + modules.append(module) + if len(modules) == 2 and modules[0].weight.shape != modules[1].weight.shape: + raise ValueError(f"Fused gate_up has incompatible gate/proj shapes: {name}") + return targets, modules + + +def load_dit(path, device="cpu", *, config=None): + """Load a Comfy Turbo int8_convrot DiT; CPU supports storage, CUDA execution. + + The optional config override is for small models/tests. Production files + without JSON config metadata use the official 32-block Qwen-Image 2.1 config. + An unknown marker, tensor layout, missing key, or extra key is an error. + """ + with _open_safetensors(path, framework="pt", device="cpu") as checkpoint: + with torch.device("meta"): + model = QwenImageDiT(**_config(checkpoint.metadata() or {}, QwenImageDiT, config)) + keys = _keys(checkpoint, "model.diffusion_model.") + expected = model.state_dict() + quantized = {} + consumed, replaced = set(), set() + names = sorted(key[:-len(".comfy_quant")] for key in keys if key.endswith(".comfy_quant")) + if not names: + raise ValueError("Expected Comfy int8_convrot DiT with .comfy_quant markers") + # Validate the complete key/shape/dtype contract before reading large matrices. + for name in names: + marker_key = keys[name + ".comfy_quant"] + marker_shape, marker_dtype = _shape_dtype(checkpoint, marker_key) + if marker_dtype != "U8" or len(marker_shape) != 1 or not 0 < marker_shape[0] <= 4096: + raise ValueError(f"Invalid Comfy quantization marker tensor: {marker_key}") + try: + marker = json.loads(bytes(checkpoint.get_tensor(marker_key).tolist()).decode("utf-8")) + except (ValueError, UnicodeError) as exc: + raise ValueError(f"Invalid Comfy quantization JSON: {marker_key}") from exc + if marker != CONVROT_METADATA or type(marker.get("convrot")) is not bool or type(marker.get("convrot_groupsize")) is not int: + raise ValueError(f"Unsupported Comfy quantization format for {name}: {marker!r}") + targets, modules = _quant_targets(model, name) + rows = sum(module.out_features for module in modules) + width = modules[0].in_features + for suffix, shape, dtypes in (("weight", (rows, width), {"I8"}), ("weight_scale", (rows, 1), {"F32"})): + key = name + "." + suffix + if key not in keys: + raise ValueError(f"Missing quantized DiT tensor: {key}") + _check_tensor(checkpoint, keys[key], shape, dtypes) + consumed.add(key) + consumed.add(name + ".comfy_quant") + for target in targets: + if target + ".weight" in replaced: + raise ValueError(f"Multiple checkpoint tensors target the same Linear: {target}") + replaced.add(target + ".weight") + quantized[name] = targets + # Comfy writes one attention hint per block ({"attention": "comfy_kitchen_int8"}); this + # engine runs the reference BF16 attention, so the hint is validated and not used. + for key in sorted(k for k in keys if k.endswith(".attn.comfy_attention.config")): + try: + hint = json.loads(bytes(checkpoint.get_tensor(keys[key]).tolist()).decode("utf-8")) + except (ValueError, UnicodeError) as exc: + raise ValueError(f"Invalid Comfy attention hint: {key}") from exc + if hint != COMFY_ATTENTION_HINT: + raise ValueError(f"Unsupported Comfy attention format for {key}: {hint!r}") + consumed.add(key) + ordinary = set(expected) - replaced + missing, extra = ordinary - set(keys), set(keys) - consumed - ordinary + if missing or extra: + raise ValueError(f"DiT checkpoint keys differ: missing={sorted(missing)}, unexpected={sorted(extra)}") + for key in ordinary: + _check_tensor(checkpoint, keys[key], expected[key].shape, {"BF16", "F32"}) + # Reuse the engine's wrapper construction and exact runtime member names. + from ..int8_ops import install_cached_linears + specs = {} + for targets in quantized.values(): + for target in targets: + shape = list(expected[target + ".weight"].shape) + specs[target] = dict(storage="int8", convrot_group=256, weight_shape=shape, + scale_shape=[shape[0], 1], input_dtype="bfloat16") + install_cached_linears(model, specs, dict(kind="convrot", group=256)) + del expected + # Load scales and weights together, releasing each view before the next layer. + for name, targets in quantized.items(): + scales = checkpoint.get_tensor(keys[name + ".weight_scale"]) + if not torch.isfinite(scales).all() or not (scales > 0).all(): + raise ValueError(f"DiT weight scales must be finite and positive: {name}") + weight = checkpoint.get_tensor(keys[name + ".weight"]) + offset = 0 + for target in targets: + rows = model.get_submodule(target).out_features + _assign(model, target + ".weight_int8", weight[offset:offset + rows], device, torch.int8) + _assign(model, target + ".weight_scale", scales[offset:offset + rows], device, torch.float32) + offset += rows + del weight, scales + for key in sorted(ordinary): + # A quantized Linear's original bias (if a future supported config has + # one) lives under original.bias in the engine wrapper. + parent, _, leaf = key.rpartition(".") + target = key if parent not in specs or leaf != "bias" else parent + ".original.bias" + _assign(model, target, checkpoint.get_tensor(keys[key]), device, torch.bfloat16) + model.to(device=device) + return model.eval().requires_grad_(False) + + +def load_vae(path, device="cpu", *, config=None): + """Load the Comfy BF16 VAE, accepting only explicit one-to-one key mappings. + + Single-frame Comfy Conv3d weights may retain a size-one temporal axis. It is + squeezed without numerical conversion; any other shape is rejected. + """ + with _open_safetensors(path, framework="pt", device="cpu") as checkpoint: + with torch.device("meta"): + model = QwenImageVAE(**_config(checkpoint.metadata() or {}, QwenImageVAE, config)) + table = vae_key_map(model) + expected = model.state_dict() + disk_targets = {} + for key in checkpoint.keys(): + target = comfy_vae_key_to_diffusers(key) + if key not in table or target not in expected: + raise ValueError(f"Unknown Comfy VAE key: {key}") + if target in disk_targets: + raise ValueError(f"Duplicate Comfy VAE target: {target}") + shape, dtype = _shape_dtype(checkpoint, key) + wanted = tuple(expected[target].shape) + module = model.get_submodule(target.rpartition(".")[0]) + if target.endswith(".weight") and isinstance(module, nn.Conv2d) and len(shape) == 5 and len(wanted) == 4 and shape[2] == 1: + shape = shape[:2] + shape[3:] + if shape != wanted or dtype != "BF16": + raise ValueError(f"Unsupported Comfy VAE tensor {key}: shape={shape}, dtype={dtype}; expected {wanted}, BF16") + disk_targets[target] = key + missing = set(expected) - set(disk_targets) + if missing: + raise ValueError(f"Missing Comfy VAE keys: {sorted(missing)}") + del expected + for target, key in disk_targets.items(): + tensor = checkpoint.get_tensor(key) + if tensor.ndim == 5 and isinstance(model.get_submodule(target.rpartition(".")[0]), nn.Conv2d): + tensor = tensor.squeeze(2) + _assign(model, target, tensor, device, torch.bfloat16) + del tensor + return model.eval().requires_grad_(False) diff --git a/freevideo_engine/image_gen/service.py b/freevideo_engine/image_gen/service.py new file mode 100644 index 00000000..2e6d55ba --- /dev/null +++ b/freevideo_engine/image_gen/service.py @@ -0,0 +1,501 @@ +"""One image job at a time, with a leased worker retained until idle or video work.""" +import atexit +import json +import math +from pathlib import Path +import queue +import secrets +import subprocess +import threading +import time +import uuid + +from . import image_pack + +BODY_LIMIT = 64 * 1024 +CANCEL_GRACE = 5. +MAX_IMAGES = 4 +ERRORS = {'no_memory', 'cancelled', 'generation_failed', 'invalid_request', 'busy'} + + +class Worker: + """One reader owns each pipe; stderr is drained without retaining user text.""" + def __init__(self, process): + self.process = process + self.rows = queue.Queue() + self.writer_mutex = threading.Lock() + self.reader = threading.Thread(target=self.read, daemon=True) + self.error_reader = threading.Thread(target=self.read_errors, daemon=True) + self.reader.start() + self.error_reader.start() + + def read(self): + try: + for line in self.process.stdout: + if len(line) > 128 * 1024: + continue + try: + row = json.loads(line) + except ValueError: + continue + if isinstance(row, dict): + self.rows.put(row) + finally: + self.rows.put(None) + + def read_errors(self): + for _ in iter(lambda: self.process.stderr.read(1024), ''): + pass + + def send(self, **row): + with self.writer_mutex: + self.process.stdin.write(json.dumps(row, ensure_ascii=True) + '\n') + self.process.stdin.flush() + + def close(self): + self.reader.join(timeout=5) + self.error_reader.join(timeout=5) + for stream in (self.process.stdin, self.process.stdout, self.process.stderr): + if not stream.closed: + stream.close() + + +class Service: + def __init__(self, installation, input_directory, *, video_busy=None, idle_seconds=180.): + self.installation = installation + self.input_directory = input_directory + self.video_busy = video_busy or self._video_busy + self.idle_seconds = idle_seconds + self.mutex = threading.RLock() + self.stop = threading.Event() + self.session_stop = threading.Event() + self.requests = queue.Queue() + self.thread = None + self.process = None + self.worker = None + self.key = None + self.state = {} + self._busy = False + self.closed = False + self.yielding = False + atexit.register(self.close) + + @staticmethod + def _video_busy(): + from ..resident_process import OWNER + return OWNER.active + + def busy(self): + return self._busy + + def info(self): + _, machine = self.installation() + spec = image_pack.catalog() + return dict(ready=image_pack.ready(machine['model_root']), model=spec['repo'], + bytes=sum(row['bytes'] for row in spec['groups']['image']), busy=self.busy(), + max_images=MAX_IMAGES) + + def request(self, value): + if not isinstance(value, dict): + raise ValueError('invalid_request') + prompt = value.get('prompt') + if not isinstance(prompt, str) or not prompt.strip() or len(prompt) > 2000: + raise ValueError('invalid_prompt') + dimensions = [value.get(name) for name in ('width', 'height')] + if any(type(size) is not int or not 256 <= size <= 2048 for size in dimensions): + raise ValueError('invalid_dimensions') + seed = value.get('seed') + if seed is not None and (type(seed) is not int or not -2**63 <= seed < 2**64): + raise ValueError('invalid_seed') + images = value.get('images', []) + if not isinstance(images, list): + raise ValueError('invalid_media') + if len(images) > MAX_IMAGES: + raise ValueError('too_many_images') + directory = Path(self.input_directory()).resolve() + paths = [] + for name in images: + # A file in the input directory, or one of our own results as freevideo-image/. + folder, _, leaf = name.rpartition('/') if isinstance(name, str) else ('', '', '') + if (not isinstance(name, str) or not leaf or leaf in ('.', '..') + or folder not in ('', 'freevideo-image') + or any(char in name for char in ('\\', ':', '\0'))): + raise ValueError('invalid_media') + parent = directory / folder if folder else directory + path = parent / leaf + if (parent.is_symlink() or path.is_symlink() or parent.resolve() != parent + or path.resolve().parent != parent or not path.is_file()): + raise ValueError('invalid_media') + paths.append(str(path)) + output = directory / 'freevideo-image' + if output.is_symlink() or output.resolve().parent != directory: + raise ValueError('invalid_media') + output.mkdir(exist_ok=True) + return dict(prompt=prompt, width=dimensions[0] // 32 * 32, height=dimensions[1] // 32 * 32, + seed=seed, images=paths, path=str(output / (uuid.uuid4().hex + '.png'))) + + def start(self, action, value): + if action not in ('generate', 'compose'): + raise ValueError('invalid_request') + with self.mutex: + if self.closed: + raise ValueError('closed') + if self.yielding or self.video_busy(): + raise ValueError('busy_video') + if self.busy() or self.thread is not None and self.thread.is_alive() and self.session_stop.is_set(): + raise ValueError('busy') + root, machine = self.installation() + if not image_pack.ready(machine['model_root']): + raise ValueError('pack_missing') + if action == 'compose': + from . import compose + value = compose.build(value, self.input_directory()) + value = self.request(value) + self.preflight(machine, value) + self.key = secrets.token_urlsafe(24) + self.stop.clear() + self.state = dict(phase='starting') + self._busy = True + if self.thread is None or not self.thread.is_alive(): + self.session_stop = threading.Event() + self.requests = queue.Queue() + self.thread = threading.Thread(target=self.work, args=(root, machine), + name='freevideo-image-gen', daemon=True) + self.requests.put(value) + self.thread.start() + else: + self.requests.put(value) + return dict(job=self.key) + + def memory_budget(self, machine): + """Bound the driver query, without importing Torch or starting a worker.""" + if machine.get('device_backend', 'cuda') != 'cuda': + return None + from ..background_room import Gpu, family + from ..background_compile import COMPILE + from ..gpu_budget import open_bounded + from ..resident_process import OWNER + + def read(): + gpu = Gpu(machine.get('gpu_uuid')) + try: + memory = gpu.memory() + if memory is None: + return None + used, total = memory + if not total or not 0 <= used <= total: + return None + # work() releases idle video/compile models; the image worker + # reuses its own allocator. Never credit an unrelated process. + owned = set() + for process in (self.process, None if OWNER.active else OWNER.process, COMPILE.process): + if process is not None and process.poll() is None: + owned.update(family(process.pid)) + rows = gpu.processes() if owned else {} + if owned and (rows is None or not owned.intersection(rows) + or any(rows.get(pid) is None for pid in owned if pid in rows)): + # An unreported owned pool must not cause a false refusal. + # The worker still checks live capacity before inference. + return total + credit = sum(size for pid, size in rows.items() if pid in owned) + return min(total, total - used + credit) + finally: + gpu.close() + + try: + return open_bounded(read, .5) + except (OSError, RuntimeError, AttributeError, ValueError): + return None + + def preflight(self, machine, value): + budget = self.memory_budget(machine) + if budget is None: + return + from PIL import Image + from .image_policy import plan_image + reference_pixels = 0 + for path in value['images']: + with Image.open(path) as image: + scale = (1024**2 / (image.width * image.height)) ** .5 + reference_pixels += round(image.width * scale) * round(image.height * scale) + plan = plan_image(gpu_budget_bytes=budget, width=value['width'], height=value['height'], + reference_pixels=reference_pixels) + if not plan.feasible: + raise ValueError('no_memory') + + def status(self, key): + with self.mutex: + if (not isinstance(key, str) or not self.key + or not secrets.compare_digest(key.encode('utf-8'), self.key.encode('utf-8'))): + raise ValueError('unknown_job') + return dict(self.state) + + def cancel(self, key): + with self.mutex: + row = self.status(key) + if not self.busy(): + return row + self.stop.set() + self.state['phase'] = 'cancelling' + return dict(phase='cancelling') + + def stop_for_generation(self, only_idle=False): + """Reap even an in-flight spawn before video acquires the runtime lease. + + ``only_idle`` is for other users of the lease (the prompt model): they get it + back from an idle worker but never cut off an image that is being made. + """ + with self.mutex: + if only_idle and (self._busy or self.thread is None or not self.thread.is_alive()): + return False + self.yielding = True + self.stop.set() + self.session_stop.set() + thread = self.thread + try: + if thread is not None and thread is not threading.current_thread(): + thread.join() + finally: + with self.mutex: + self.yielding = False + return True + + def close(self): + with self.mutex: + self.closed = True + self.stop_for_generation() + + def environment(self, root): + from ..comfy_bridge import source_root + from ..comfy_environment import isolated_environment + from .. import processes, resident_process + env = isolated_environment(root, source_root()) + for name in ('FREEVIDEO_RUNTIME_LOCK_FD', processes.HANDLE_ENV, resident_process.ENV, + 'FREEVIDEO_PROMPT_PRESTART_WORKER'): + env.pop(name, None) + return env + + def spawn(self, root, machine, env, descriptor): + from .. import processes + command = processes.module_command('freevideo_engine.image_gen.worker', image_pack.base(machine['model_root']), + '--device', machine.get('device_backend', 'cuda'), + '--idle-seconds', self.idle_seconds) + command[0] = machine['comfy_python'] + env = dict(env, FREEVIDEO_RUNTIME_LOCK_FD=str(descriptor)) + return processes.popen(command, env=env, stdin=subprocess.PIPE, stdout=subprocess.PIPE, + stderr=subprocess.PIPE, text=True, encoding='utf-8', errors='replace', + pass_fds=(descriptor,), supervise=True, start_new_session=True) + + def work(self, root, machine): + from ..locking import runtime_lock + from .. import processes, resident_process + session = None + failure = 'worker_exit' + try: + from ..background_compile import COMPILE + from ..encoder_prewarm import IDLE + COMPILE.stop() + if self.video_busy() or resident_process.OWNER.active: + raise BlockingIOError('busy_video') + resident_process.OWNER.stop_prewarm() + IDLE.stop() + with runtime_lock(Path(root) / 'engine.lock', inherit=False) as descriptor: + try: + if self.session_stop.is_set(): + return + env = self.environment(root) + owner = resident_process.OWNER + if owner.process is not None and owner.process.poll() is None and owner.endpoint: + env[resident_process.ENV] = str(owner.endpoint) + try: + resident_process.release_idle_cache(env, descriptor) + except resident_process.SessionRestarted: + pass + env.pop(resident_process.ENV, None) + if self.session_stop.is_set(): + return + if self.video_busy(): + raise BlockingIOError('busy_video') + session = Worker(self.spawn(root, machine, env, descriptor)) + with self.mutex: + self.worker = session + self.process = session.process + while not self.session_stop.is_set(): + try: + value = self.requests.get(timeout=.05) + except queue.Empty: + with self.mutex: + if session.process.poll() is not None and self.requests.empty(): + self.session_stop.set() + break + continue + if session.process.poll() == 0: + session.process.wait() + session.close() + session = Worker(self.spawn(root, machine, env, descriptor)) + with self.mutex: + self.worker = session + self.process = session.process + self.infer(session, value) + finally: + # Keep the parent's lease until the whole owned process tree exits. + if session is not None: + try: + if session.process.poll() is None: + session.send(action='close') + session.process.wait(timeout=5) + except (OSError, subprocess.TimeoutExpired): + pass + processes.stop(session.process, grace=5) + session.process.wait() + session.close() + except BlockingIOError: + failure = 'busy_video' + except Exception: + failure = 'worker_exit' + finally: + with self.mutex: + self.process = self.worker = None + self.session_stop.set() + if self.busy(): + self.state = (dict(phase='cancelled') if self.stop.is_set() + else dict(phase='failed', error=failure)) + self._busy = False + + def infer(self, session, value): + session.send(action='generate', **value) + cancelled_at = None + deadline = time.monotonic() + 600 + while True: + if self.stop.is_set() and cancelled_at is None: + session.send(action='cancel') + cancelled_at = time.monotonic() + if cancelled_at is not None and time.monotonic() - cancelled_at > CANCEL_GRACE: + raise TimeoutError('cancelled') + if self.session_stop.is_set(): + return + if time.monotonic() > deadline: + raise TimeoutError('worker_exit') + try: + row = session.rows.get(timeout=.05) + except queue.Empty: + continue + if row is None: + raise RuntimeError('worker_exit') + phase = row.get('phase') + if phase == 'progress': + stage, fraction = row.get('stage'), row.get('fraction') + if (stage in ('load', 'text', 'encode', 'denoise', 'decode') and type(fraction) in (int, float) + and math.isfinite(fraction) and 0 <= fraction <= 1): + with self.mutex: + self.state.update(phase='cancelling' if self.stop.is_set() else 'progress', + stage=stage, fraction=fraction) + elif phase in ('result', 'error'): + if self.stop.is_set() or row.get('code') == 'cancelled': + result = dict(phase='cancelled') + elif phase == 'error': + code = row.get('code') + result = dict(phase='failed', error=code if code in ERRORS else 'generation_failed') + if code == 'no_memory': + result['message_zh'] = '内存或显存不足，请降低图像尺寸或减少参考图。' + else: + path = Path(value['path']) + if row.get('path') != str(path) or path.is_symlink() or not path.is_file(): + raise RuntimeError('worker_exit') + from PIL import Image + with Image.open(path) as image: + if image.format != 'PNG' or image.size != (value['width'], value['height']): + raise RuntimeError('worker_exit') + image.verify() + seconds = row.get('seconds') + if type(seconds) not in (int, float) or not math.isfinite(seconds) or seconds < 0: + raise RuntimeError('worker_exit') + result = dict(phase='complete', file='freevideo-image/' + path.name, + seconds=seconds, plan=row.get('plan')) + with self.mutex: + self.state = result + self._busy = False + return + + +def stop_prestart_for_generation(server): + service = getattr(server, '_freevideo_image_gen', None) + if service is not None: + service.stop_for_generation() + + +def get_service(server=None): + """The image service of this server, for callers in the same process (the Agent). + + They call ``start('generate', {...})``, ``status(job)``, ``cancel(job)`` and + ``info()`` directly with the same arguments, results and error codes as the + HTTP routes; only the page token is not needed. None before ``register()``. + """ + if server is None: + from server import PromptServer + server = PromptServer.instance + return getattr(server, '_freevideo_image_gen', None) + + +def release_idle_worker(server): + """Hand engine.lock back from a worker that is only waiting for its idle exit.""" + service = getattr(server, '_freevideo_image_gen', None) + if service is not None: + service.stop_for_generation(only_idle=True) + + +def register(): + from aiohttp import web + from server import PromptServer + import folder_paths + from ..comfy_bridge import installation + server = PromptServer.instance + if server is None or getattr(server, '_freevideo_image_gen', None) is not None: + return + + def video_busy(): + running, pending = server.prompt_queue.get_current_queue() + return bool(running or pending) or Service._video_busy() + + service = Service(installation, folder_paths.get_input_directory, video_busy=video_busy) + server._freevideo_image_gen = service + token = secrets.token_urlsafe(24) + + @server.routes.get('/freevideo/image-gen') + async def info(request): + try: + return web.json_response(dict(service.info(), token=token)) + except (OSError, ValueError): + return web.json_response(dict(error='setup_required'), status=400) + + @server.routes.post('/freevideo/image-gen/{action}') + async def action(request): + header = request.headers.get('X-FreeVideo-Image', '') + if not secrets.compare_digest(header.encode('utf-8'), token.encode('utf-8')): + return web.json_response(dict(error='reload'), status=403) + try: + if request.content_length and request.content_length > BODY_LIMIT: + raise ValueError('invalid_request') + body = bytearray() + while True: + chunk = await request.content.read(min(8192, BODY_LIMIT + 1 - len(body))) + if not chunk: + break + body.extend(chunk) + if len(body) > BODY_LIMIT: + raise ValueError('invalid_request') + value = json.loads(body) + if not isinstance(value, dict): + raise ValueError('invalid_request') + name = request.match_info['action'] + if name == 'status': + result = service.status(value.get('job')) + elif name == 'cancel': + result = service.cancel(value.get('job')) + else: + result = service.start(name, value) + return web.json_response(result) + except (ValueError, OSError) as error: + codes = {'invalid_request', 'invalid_prompt', 'invalid_dimensions', 'invalid_seed', 'invalid_media', + 'too_many_images', 'busy', 'busy_video', 'no_memory', 'pack_missing', 'unknown_job', 'closed'} + return web.json_response(dict(error=str(error) if str(error) in codes else 'invalid_request'), status=400) diff --git a/freevideo_engine/image_gen/worker.py b/freevideo_engine/image_gen/worker.py new file mode 100644 index 00000000..28394881 --- /dev/null +++ b/freevideo_engine/image_gen/worker.py @@ -0,0 +1,158 @@ +"""Persistent ImageRunner owner; stdin/stdout contain one JSON object per line.""" +import argparse +from contextlib import ExitStack, redirect_stdout +import json +from pathlib import Path +import queue +import sys +import threading +import time + + +class Inbox: + """Read cancellation while the main thread is inside a denoising step.""" + def __init__(self, stream, send): + self.stream = stream + self.send = send + self.rows = queue.Queue() + self.cancel = threading.Event() + self.closed = threading.Event() + self.mutex = threading.Lock() + self.busy = False + self.reader = threading.Thread(target=self.read, daemon=True) + self.reader.start() + + def read(self): + try: + while not self.closed.is_set(): + line = self.stream.readline(64 * 1024 + 1) + if not line: + break + try: + if len(line.encode('utf-8')) > 64 * 1024: + raise ValueError('invalid_request') + value = json.loads(line) + if not isinstance(value, dict): + raise ValueError('invalid_request') + action = value.get('action') + with self.mutex: + if action == 'cancel': + if self.busy: + self.cancel.set() + elif action == 'close': + break + elif action == 'generate': + if self.busy: + self.send(phase='error', code='busy') + else: + self.busy = True + self.cancel.clear() + self.rows.put(value) + else: + raise ValueError('invalid_request') + except (ValueError, UnicodeError): + self.send(phase='error', code='invalid_request') + finally: + self.closed.set() + self.cancel.set() + self.rows.put(None) + + +def _flatten(image): + """RGB copy; transparency is shown on white, as the encoder's vision tower sees it.""" + from PIL import Image + if image.mode in ('RGBA', 'LA') or 'transparency' in image.info: + rgba = image.convert('RGBA') + base = Image.new('RGBA', rgba.size, (255, 255, 255, 255)) + return Image.alpha_composite(base, rgba).convert('RGB') + return image.convert('RGB') + + +def load_runner(pack_dir, device, *, progress=None): + # Importing this module in ComfyUI or in CPU protocol tests never loads torch. + from .image_runner import ImageRunner + return ImageRunner(pack_dir, device=device, progress=progress) + + +def serve(pack_dir, *, device='cuda', idle_seconds=180., stream=None, output=None, factory=None): + stream = sys.stdin if stream is None else stream + output = sys.stdout if output is None else output + factory = factory or load_runner + output_mutex = threading.Lock() + + def send(**row): + with output_mutex: + output.write(json.dumps(row, ensure_ascii=True, allow_nan=False) + '\n') + output.flush() + + inbox = Inbox(stream, send) + runner = None + try: + while not inbox.closed.is_set(): + try: + value = inbox.rows.get(timeout=idle_seconds) + except queue.Empty: + break + if value is None or inbox.closed.is_set(): + break + started = time.monotonic() + result = None + try: + from PIL import Image + with ExitStack() as stack, redirect_stdout(sys.stderr): + if inbox.cancel.is_set(): + raise InterruptedError('cancelled') + progress = lambda stage, fraction: send(phase='progress', stage=stage, fraction=fraction) + if runner is None: + runner = factory(pack_dir, device, progress=progress) + references = [] + for path in value.get('images', []): + image = stack.enter_context(Image.open(path)) + references.append(_flatten(image)) + image = runner.generate( + value['prompt'], width=value['width'], height=value['height'], + seed=value.get('seed'), reference_images=references, + progress=progress, + cancel=inbox.cancel.is_set) + if inbox.cancel.is_set(): + raise InterruptedError('cancelled') + # The parent chooses a fresh path under ComfyUI input. + # Exclusive creation never replaces an existing user file. + with Path(value['path']).open('xb') as target: + image.save(target, format='PNG') + plan = runner.stats().get('plan') + result = dict(phase='result', path=value['path'], seconds=time.monotonic() - started, plan=plan) + except MemoryError: + result = dict(phase='error', code='no_memory', message_zh='内存或显存不足，请降低图像尺寸或减少参考图。') + except Exception as error: + # Exceptions may echo prompts, filenames or model paths. + code = ('cancelled' if inbox.cancel.is_set() else + 'no_memory' if type(error).__name__ == 'OutOfMemoryError' else 'generation_failed') + result = dict(phase='error', code=code) + finally: + with inbox.mutex: + inbox.busy = False + if result is not None: + send(**result) + finally: + inbox.closed.set() + inbox.cancel.set() + if runner is not None: + with redirect_stdout(sys.stderr): + runner.close() + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument('pack_dir', type=Path) + parser.add_argument('--device', default='cuda') + parser.add_argument('--idle-seconds', type=float, default=180.) + args = parser.parse_args() + if not 0 < args.idle_seconds < float('inf'): + parser.error('--idle-seconds must be positive and finite') + serve(args.pack_dir, device=args.device, idle_seconds=args.idle_seconds) + return 0 + + +if __name__ == '__main__': + raise SystemExit(main()) diff --git a/freevideo_engine/launch_view.py b/freevideo_engine/launch_view.py index dc490683..f9c020bf 100644 --- a/freevideo_engine/launch_view.py +++ b/freevideo_engine/launch_view.py @@ -15,6 +15,7 @@ from .long_paths import io_path, plain VIEW_ENV = 'FREEVIDEO_RUNTIME_VIEW' +ONLINE_ENGINE = 'freevideo-engine' # the folder an online installation is made in class View(namedtuple('View', 'real root rules', defaults=(os.path,))): @@ -64,9 +65,18 @@ def view_for(real_root) -> View: def for_engine(engine_root) -> View: - """The view of the offline runtime around this engine folder, if it needs one.""" + """The view of the offline runtime around this engine folder, if it needs one. + + An installation made online is its engine folder: it gets a view only where the + system code page cannot spell its folder, and one that cannot be made is not an error. + """ from . import runtime_alias engine = Path(plain(engine_root)) + if engine.name.casefold() == ONLINE_ENGINE: + try: + return view_for(engine) if runtime_alias.needs_view(engine, spelling_only=True) else View(engine, engine) + except OSError: # Includes a junction that cannot be made: the checks after it name what fails. + return View(engine, engine) root = engine.parent # The first two tests read nothing from disk wherever no junction is needed. # The third asks by the extended path: where a junction is needed, the plain diff --git a/freevideo_engine/launcher/FField.qml b/freevideo_engine/launcher/FField.qml index c619e6ee..0aac12f3 100644 --- a/freevideo_engine/launcher/FField.qml +++ b/freevideo_engine/launcher/FField.qml @@ -3,6 +3,19 @@ import QtQuick.Controls.Basic TextField { id: control + // State updates must preserve the draft until the user submits it. + property var value: undefined + Binding { + target: control; property: "text" + value: control.value === undefined ? "" : control.value + when: control.value !== undefined && !control.activeFocus + restoreMode: Binding.RestoreNone + } + // Return keeps focus; sync after editingFinished commits or rejects input. + onAccepted: Qt.callLater(function() { + if (control.value !== undefined) + control.text = control.value + }) implicitHeight: theme.height + 4 color: theme.text; placeholderTextColor: theme.disabled selectionColor: theme.accentDim; selectedTextColor: theme.text diff --git a/freevideo_engine/launcher/Main.qml b/freevideo_engine/launcher/Main.qml index 22495fb6..f0dbf70e 100644 --- a/freevideo_engine/launcher/Main.qml +++ b/freevideo_engine/launcher/Main.qml @@ -471,7 +471,7 @@ ApplicationWindow { FText { text: s.form.new_comfy ? t("Install location", "安装位置") : t("ComfyUI folder", "ComfyUI 目录"); font.weight: Font.DemiBold } RowLayout { Layout.fillWidth: true; spacing: 8 - FField { objectName: "installationPath"; Layout.fillWidth: true; text: s.form.new_comfy ? s.form.destination : s.form.comfy; placeholderText: t("Choose a folder", "选择文件夹"); enabled: !s.busy; onEditingFinished: backend.edit(s.form.new_comfy ? "destination" : "comfy", text) } + FField { objectName: "installationPath"; Layout.fillWidth: true; value: s.form.new_comfy ? s.form.destination : s.form.comfy; placeholderText: t("Choose a folder", "选择文件夹"); enabled: !s.busy; onEditingFinished: backend.edit(s.form.new_comfy ? "destination" : "comfy", text) } FButton { text: t("Browse…", "浏览…"); implicitHeight: theme.height + 4; onClicked: backend.browse(s.form.new_comfy ? "destination" : "comfy") } } FText { visible: !s.form.new_comfy; text: t("Uses your existing Python when available. A separate environment loads only FreeVideo.", "优先使用已有 Python；独立环境仅加载 FreeVideo。"); color: theme.muted; font.pixelSize: theme.micro; Layout.fillWidth: true } @@ -635,6 +635,24 @@ ApplicationWindow { onToggled: backend.edit("sampling_caches", checked) } } + FCard { + objectName: "imageModelCard" + visible: Qt.platform.os !== "osx" && !!s.image_pack && s.image_pack.total_bytes > 0 && s.image_pack.state !== "ready" + Layout.fillWidth: true; padding: 16 + RowLayout { + Layout.fillWidth: true; spacing: 14 + FCheck { + objectName: "imageModelSwitch"; Layout.fillWidth: true + text: t("Image generation models", "图像生成模型") + checked: !!s.form.image_model && s.image_pack.enough_disk; enabled: !s.busy && s.image_pack.enough_disk + onToggled: backend.edit("image_model", checked) + } + FText { + objectName: "imageModelSize"; color: theme.muted + text: s.image_pack.enough_disk ? bytes(s.image_pack.total_bytes).replace(" ", "\u00a0") : t("Not enough disk space", "磁盘空间不足") + } + } + } FText { visible: s.offline.models > 0; text: "✓ " + t("Model packages: ", "已导入模型包：") + s.offline.models; color: theme.success; font.pixelSize: theme.micro; Layout.fillWidth: true } } @@ -844,6 +862,7 @@ ApplicationWindow { } } FCard { + objectName: "installationCard" Layout.fillWidth: true; padding: 22; spacing: 12 FText { text: t("Installation", "安装信息"); font.weight: Font.DemiBold } GridLayout { @@ -864,6 +883,28 @@ ApplicationWindow { FCheck { visible: s.needs_consent; text: s.consent; checked: accepted; onToggled: accepted = checked; Layout.fillWidth: true } FButton { objectName: "bundleLicensesButton"; visible: s.needs_consent; text: t("View licenses", "查看许可证"); flat: true; implicitHeight: theme.heightSm - 2; leftPadding: 32; font.pixelSize: theme.micro + 1; onClicked: licensesOpen = true } } + FCard { + objectName: "imagePackCard" + visible: !!s.image_pack && s.image_pack.visible + Layout.fillWidth: true; padding: 22; spacing: 10 + RowLayout { + Layout.fillWidth: true; spacing: 14 + ColumnLayout { + Layout.fillWidth: true; spacing: 4 + FText { objectName: "imagePackTitle"; text: s.image_pack.title; font.weight: Font.DemiBold; Layout.fillWidth: true } + FText { objectName: "imagePackStatus"; text: s.image_pack.detail; color: theme.muted; Layout.fillWidth: true } + } + FText { objectName: "imagePackSize"; text: bytes(s.image_pack.total_bytes).replace(" ", "\u00a0"); color: theme.muted } + FButton { objectName: "imagePackInstall"; visible: !s.image_pack.ready && !s.image_pack.busy; text: s.image_pack.downloaded_bytes > 0 ? t("Resume", "继续") : t("Download", "下载"); enabled: s.image_pack.can_act; onClicked: backend.imagePackAction("install") } + FButton { objectName: "imagePackCancel"; visible: s.image_pack.busy; text: t("Pause", "暂停"); flat: true; onClicked: backend.imagePackAction("cancel") } + FButton { objectName: "imagePackRepair"; visible: s.image_pack.repairable && s.image_pack.downloaded_bytes > 0 && !s.image_pack.busy; text: t("Repair", "修复"); flat: true; enabled: s.image_pack.can_act; onClicked: backend.imagePackAction("repair") } + FButton { objectName: "imagePackRemove"; visible: s.image_pack.removable_bytes > 0 && !s.image_pack.busy; text: t("Remove", "删除"); flat: true; enabled: s.image_pack.can_act; onClicked: backend.imagePackAction("remove") } + } + FMeter { + objectName: "imagePackProgress"; visible: s.image_pack.busy; Layout.fillWidth: true; active: visible + fraction: s.image_pack.install.total > 0 ? s.image_pack.install.done / s.image_pack.install.total : -1 + } + } } Item { height: 18 } @@ -1159,16 +1200,16 @@ ApplicationWindow { Layout.fillWidth: true; visible: settingsTab === "advanced"; spacing: 20; enabled: !s.busy FGroup { Layout.fillWidth: true; title: t("ComfyUI address", "ComfyUI 地址") - FField { objectName: "listenAddress"; Layout.fillWidth: true; text: s.form.url; enabled: !s.form.lan_access; onEditingFinished: backend.edit("url", text) } + FField { objectName: "listenAddress"; Layout.fillWidth: true; value: s.form.url; enabled: !s.form.lan_access; onEditingFinished: backend.edit("url", text) } FText { objectName: "listenAddressHint"; text: t("With 0.0.0.0 or this computer's LAN IP, devices on the same network can open FreeVideo without the access link.", "填 0.0.0.0 或本机的局域网 IP 后，同一网络中的设备不需要访问链接就能打开 FreeVideo。"); color: theme.muted; font.pixelSize: theme.micro; Layout.fillWidth: true } FText { visible: Qt.platform.os === "windows"; text: t("If Windows asks about the firewall the first time, choose Allow.", "Windows 首次询问防火墙权限时，请选择“允许”。"); color: theme.muted; font.pixelSize: theme.micro; Layout.fillWidth: true } } FGroup { Layout.fillWidth: true; title: t("Folders · optional", "目录 · 可选"); enabled: !s.portable FText { text: t("Engine folder", "引擎目录"); font.pixelSize: theme.micro; color: theme.muted } - RowLayout { Layout.fillWidth: true; spacing: 8; FField { Layout.fillWidth: true; text: s.form.engine; placeholderText: t("Automatic", "自动选择"); onEditingFinished: backend.edit("engine", text) } FButton { text: t("Browse…", "浏览…"); implicitHeight: theme.height + 4; onClicked: backend.browse("engine") } } + RowLayout { Layout.fillWidth: true; spacing: 8; FField { Layout.fillWidth: true; value: s.form.engine; placeholderText: t("Automatic", "自动选择"); onEditingFinished: backend.edit("engine", text) } FButton { text: t("Browse…", "浏览…"); implicitHeight: theme.height + 4; onClicked: backend.browse("engine") } } FText { text: t("ComfyUI Python", "ComfyUI Python"); font.pixelSize: theme.micro; color: theme.muted; Layout.topMargin: 4 } - FField { Layout.fillWidth: true; text: s.form.python; placeholderText: t("Path to python.exe", "python.exe 的路径"); onEditingFinished: backend.edit("python", text) } + FField { Layout.fillWidth: true; value: s.form.python; placeholderText: t("Path to python.exe", "python.exe 的路径"); onEditingFinished: backend.edit("python", text) } } FGroup { Layout.fillWidth: true; title: t("Maintenance", "维护"); enabled: !s.portable diff --git a/freevideo_engine/launcher_session.py b/freevideo_engine/launcher_session.py index a23e4e8a..9a25faf7 100644 --- a/freevideo_engine/launcher_session.py +++ b/freevideo_engine/launcher_session.py @@ -11,7 +11,7 @@ from . import __version__ from . import disk_space -from .comfy_launcher_runtime import Controller, layout, local_url, loopback_url, managed_frontend, new_layout +from .comfy_launcher_runtime import Controller, layout, local_url, loopback_url, managed_frontend, new_layout, queue_busy from .desktop_runtime import launcher_root, materialize_source from .download_settings import Probe, read as download_preferences, speed_text from .launcher_settings import Store, default_language @@ -85,7 +85,8 @@ def __init__(self, source=None, *, controller=None, store=None, updater=None, sm self.form = dict(comfy='', destination=str(home / 'FreeVideo'), engine='', python='', url='http://127.0.0.1:8188', models='', model_dirs=[], model_method='auto', environment_method='auto', separate=False, repair=False, new_comfy=True, offline_runtime='', offline_models=[], - imported_archives=[], kept_archives=[], sampling_caches=not saved.get('installation')) + imported_archives=[], kept_archives=[], sampling_caches=not saved.get('installation'), + image_model=False) self.form['lan_access'] = False self.form.update({k: v for k, v in saved.items() if k in self.form}) if 'environment_method' not in saved and self.form['offline_runtime']: @@ -111,6 +112,7 @@ def __init__(self, source=None, *, controller=None, store=None, updater=None, sm self.tail = Tail() self.log_source = None self.model_groups = [] + self.image_installer = None self.download_export = self.download_error = self.download_error_key = '' self._download_view = None self.browser_attempted = False @@ -211,6 +213,9 @@ def remember(self, selection, persist=True): self.persist() def edit(self, key, value): + if self.image_installer and self.image_installer.busy and key in ( + 'engine', 'comfy', 'destination', 'new_comfy', 'offline_runtime'): + return if key == 'language': self.language = str(value); self.persist(); return if key == 'token': @@ -220,7 +225,7 @@ def edit(self, key, value): self.token = validate(str(value)); return if key not in self.form or self.controller.busy or self.importer.busy or self.cleaner.busy or self.archives_busy: return - if key in ('separate', 'repair', 'new_comfy', 'sampling_caches', 'lan_access'): + if key in ('separate', 'repair', 'new_comfy', 'sampling_caches', 'lan_access', 'image_model'): value = bool(value) if key == 'url': value = local_url(str(value)) @@ -415,6 +420,11 @@ def activate_offline(self, root): def action(self, name, accepted=False): if self.cleaner.busy or self.archives_busy: return + if self.image_installer and self.image_installer.busy and ( + name == 'setup' or name == 'update-engine' + or name == 'retry' and self.retry_kind() != 'launch' + or name == 'primary' and self.page != 'launcher'): + return # The model download holds the setup lease; launching and opening # FreeVideo continue, anything that would start setup waits for it. # Retrying a failed start only launches again. @@ -554,7 +564,8 @@ def can_repair(self): """An installed engine this launcher can check and reinstall again.""" return (bool(self.selected) and getattr(self.controller, 'fixed_environment', False) is not True and not (self.controller.busy or self.importer.busy or self.cleaner.busy or self.archives_busy or self.closing - or self.engine_updating or self.upgrade.active or self.model_switch)) + or self.engine_updating or self.upgrade.active or self.model_switch + or self.image_installer and self.image_installer.busy)) def repair_unavailable(self): """Why Repair is missing: it appears once a task ends, never in an offline @@ -587,11 +598,113 @@ def retry_kind(self): return '' def cleanup_downloads(self, remove=False): + if self.image_installer and self.image_installer.busy: + return if (self.controller.busy or self.importer.busy or self.archives_busy or self.closing or self.upgrade.active or self.model_switch or self.update_intent or self.engine_updating or self.updater and self.updater.busy): return self.cleaner.start(self.engine_root(), remove=remove) + def _image_location(self): + engine = Path(self.installation_engine() or self.engine_root()) + try: + machine = json.loads(io_path(engine / 'machine.json').read_text(encoding='utf-8')) + except (OSError, ValueError): + machine = {} + if not isinstance(machine, dict): + machine = {} + root = Path(os.environ.get('FREEVIDEO_MODEL_ROOT') or machine.get('model_root') or engine / 'models') + return root.expanduser(), engine + + def image_pack_action(self, action): + from .image_gen import image_pack + if action == 'cancel': + if self.image_installer: + self.image_installer.cancel() + return + if (self.controller.busy or self.importer.busy or self.cleaner.busy or self.archives_busy + or self.closing or self.engine_updating or self.update_intent or self.upgrade.active or self.model_switch + or self.updater and self.updater.busy): + return + root, engine = self._image_location() + if self.image_installer and self.image_installer.busy: + return + # An install is the quiet kind: it stops while a video is generating and carries on afterwards. + self.image_installer = image_pack.Installer(root, engine_root=engine, busy=self.generating) + self.image_installer.start(action, background=action == 'install') + + def generating(self): + """Whether the opened ComfyUI is running or holding a job (read from the download thread).""" + url = (self.controller.selection or self.selected or {}).get('url') or self.form['url'] + return queue_busy(url) + + def image_disk_fits(self, total): + """Whether the chosen installation location has room for the image model.""" + from .image_gen import image_pack + folder = self.form['destination'] if self.form['new_comfy'] else self.form['comfy'] + try: + return disk_space.free_bytes(Path(folder).expanduser()) >= total + image_pack.RESERVE + except (OSError, ValueError): + return True + + def image_pack_view(self): + from .image_gen import image_pack + if not self.installation_engine(): + total = sum(r['bytes'] for r in image_pack.rows()) + return dict(installable=False, ready=False, state='missing', visible=False, busy=False, + downloaded_bytes=0, verified_bytes=0, total_bytes=total, bytes=total, + title=self.t(*NAMES['image']), detail=self.t('Not installed', '未安装'), + install={}, can_act=False, removable_bytes=0, repairable=False) + root, engine = self._image_location() + try: + value = image_pack.status(root) + removable = image_pack.removal_plan(root)['bytes'] + except (OSError, ValueError, image_pack.PackError): + value = dict(installable=False, ready=False, state='invalid', downloaded_bytes=0, verified_bytes=0, + total_bytes=sum(r['bytes'] for r in image_pack.rows()), bytes=0) + removable = 0 + installer = self.image_installer + active = installer.state() if installer and installer.root == root and installer.engine_root == engine else {} + busy = bool(installer and installer.busy) + labels = {'missing': ('Not installed', '未安装'), 'partial': ('Paused', '已暂停'), + 'ready': ('Ready', '已就绪'), 'invalid': ('Needs repair', '需要修复')} + detail = self.t(*labels[value['state']]) + # Repair only answers a pack that is ready or known to be damaged; half a download just continues. + repairable = value['state'] in ('ready', 'invalid') + if active.get('state') == 'running': + labels = {'install': ('Downloading', '正在下载'), 'repair': ('Repairing', '正在修复'), + 'remove': ('Removing', '正在删除'), 'verify': ('Checking', '正在校验')} + detail = self.t(*labels[active['action']]) + elif active.get('state') == 'paused_for_generation': + detail = self.t('Paused while generating', '生成时暂停') + elif active.get('state') == 'cancelled': + detail = self.t('Paused', '已暂停') + elif active.get('state') == 'failed': + labels = {'disk_space': ('Not enough disk space', '磁盘空间不足'), + 'needs_repair': ('Needs repair', '需要修复'), 'busy': ('In use', '正在使用'), + 'unsafe_path': ('Needs repair', '需要修复'), 'busy_video': ('Paused while generating', '生成时暂停'), + 'offline': ('No network', '网络不可用'), 'source_failed': ('Download failed', '下载失败')} + detail = self.t(*labels.get(active.get('reason'), ('Stopped', '未完成'))) + repairable = repairable or active.get('reason') in ('needs_repair', 'unsafe_path') + blocked = (self.controller.busy or self.importer.busy or self.cleaner.busy or self.archives_busy + or self.closing or self.engine_updating or self.update_intent or self.upgrade.active or self.model_switch + or self.updater and self.updater.busy) + return dict(value, visible=bool(self.installation_engine()), title=self.t(*NAMES['image']), + detail=detail, busy=busy, install=active, can_act=not (blocked or busy), removable_bytes=removable, + repairable=repairable) + + def start_image_model_after_install(self): + """The installer's "also install the image model": once ComfyUI opens, the download runs in the background.""" + if (not self.form['image_model'] or sys.platform == 'darwin' or not self.installation_engine() + or self.controller.state.get('status') != 'open' or self.controller.busy): + return + if self.image_installer and self.image_installer.busy: + return + self.image_pack_action('install') + if self.image_installer and self.image_installer.busy: + self.form['image_model'] = False + self.persist() + def upgrade_model(self): """The user's one consent: download int8 beside the current model, switch, then release the old one. @@ -602,6 +715,8 @@ def upgrade_model(self): the previous model. An old model left after the switch is released here too, when the user asks. """ + if self.image_installer and self.image_installer.busy: + return state = self.upgrade.state if (state.get('status') not in ('available', 'failed', 'releasable') or self.upgrade.busy or self.model_switch or self.closing or self.controller.busy or self.importer.busy or self.cleaner.busy or self.archives_busy or self.engine_updating): @@ -1229,6 +1344,10 @@ def close(self): if self.updater: self.updater.cancelled.set() self.upgrade.close() + if self.image_installer: + self.image_installer.cancel() + if self.image_installer.thread: + self.image_installer.thread.join(5) self.controller.close() from .launcher_bridge import remove remove(self.bridge) @@ -1344,6 +1463,7 @@ def tick(self): elif imported['status'] == 'prepared': self.activate_offline(imported['ready_root']) + self.start_image_model_after_install() row = self.controller.state busy = self.controller.busy selection = getattr(self.controller, 'selection', None) @@ -1633,6 +1753,7 @@ def snapshot(self): from .launcher_copy import display, progress_view, source_name zh = self.language.startswith('zh') row = self.controller.state + busy = self.controller.busy or self.importer.busy or self.cleaner.busy task = row.get('task', {}) port_from = int(row['moved_from']) if row.get('status') == 'open' and 'moved_from' in row else 0 port_to = (urlsplit(row.get('selection', {}).get('url') or row['url']).port or 80) if port_from else 0 @@ -1670,10 +1791,12 @@ def snapshot(self): before_models = (row.get('action') == 'install' and row.get('status') == 'running' and not (row.get('task', {}).get('model_groups') or row.get('model_groups'))) for name in FAMILIES: + if name == 'image' and name not in by_id: + continue # Optional weights have their own receipt; video ready cannot make them ready. if name == 'sampling' and name not in by_id and not self.form['sampling_caches']: continue item = dict(by_id.get(name, {})) - state = 'ready' if ready else item.get('state', 'waiting') + state = 'ready' if ready and name != 'image' else item.get('state', 'waiting') if state == 'pending': # Part of the group is ready; the rest (often small config files) waits for a download slot. state = 'queued' @@ -1696,6 +1819,13 @@ def snapshot(self): detail=self.t(*shown), done=done, total=total, found=item.get('existing_bytes', 0), download=item.get('download_bytes', 0), rate='%.1f MiB/s' % (item['bytes_per_second']/2**20) if item.get('bytes_per_second') else '')) + image_pack = self.image_pack_view() + image_pack['enough_disk'] = self.image_disk_fits(image_pack['total_bytes']) + if image_pack['visible'] and 'image' not in by_id: + models.append(dict(id='image', title=image_pack['title'], state=image_pack['state'], + detail=image_pack['detail'], done=image_pack['downloaded_bytes'], + total=image_pack['total_bytes'], found=image_pack['verified_bytes'], + download=image_pack['bytes'], rate='')) # A model download names its group, so it shows the group's download, as in the plan and the rows below. groups = {self.t(*NAMES[group['id']]): group for group in self.model_groups if group.get('id') in NAMES and group.get('download_bytes')} @@ -1740,6 +1870,7 @@ def snapshot(self): for group, entries in probe.get('sources', {}).items(): for entry in entries: names = {'edge-models': ('Video model', '视频模型'), 'vdn-models': ('Decoder', '解码器'), + 'image-models': NAMES['image'], 'models': ('Text encoder', '文本编码器'), 'pypi': ('Python packages', 'Python 依赖'), 'github': ('Tools', '安装工具'), 'git': ('Git', 'Git'), 'cuda': ('CUDA', 'CUDA')} if group.startswith('torch-'): @@ -1795,13 +1926,13 @@ def snapshot(self): notice=downloads['copy']['exported'].format(folder=self.download_export) if self.download_export else '') if inspection_read_error: download_state.update(error=self.t(*DOWNLOAD_COPY['error_read']), needs_check=True) - return plain_data(dict(download_list=download_state, sampling_cache_bytes=total_bytes(), version=__version__, + return plain_data(dict(download_list=download_state, sampling_cache_bytes=total_bytes(), version=__version__, image_pack=image_pack, disk=disk, decimal_sizes=sys.platform == 'darwin', zh=self.language.startswith('zh'), form=dict(self.form), - page=self.page, status=row.get('status', 'idle'), busy=self.controller.busy or self.importer.busy or self.cleaner.busy, + page=self.page, status=row.get('status', 'idle'), busy=busy, cleanup=cleanup, can_cleanup=not (self.controller.busy or self.importer.busy or self.cleaner.busy or self.archives_busy or self.closing or self.update_intent or self.engine_updating or self.updater and self.updater.busy - or self.upgrade.active or self.model_switch), + or self.upgrade.active or self.model_switch or image_pack['busy']), model_upgrade=dict(model_upgrade, switching=bool(self.model_switch)), offline=dict(progress_view(self.importer.state, zh), runtime=bool(self.form['offline_runtime']), runtime_supported=runtime_packages_supported(), @@ -1815,7 +1946,9 @@ def snapshot(self): report=dict(self.report), models=models, overall=overall, progress=progress, detail='' if file_text else clean(progress.get('detail', '')), progress_text=file_text, remaining_seconds=remaining, - elapsed=duration(time.monotonic()-self.started) if self.started else '', summary=' · '.join(estimate), + # The elapsed row is only shown while working. An idle clock would + # keep publishing changed snapshots after a launch has failed. + elapsed=duration(time.monotonic()-self.started) if self.started and busy else '', summary=' · '.join(estimate), failure=failure, source=source, source_name=source_name(source, zh), proxy_mode=preferences['proxy_mode'], probe=probe, speeds=speeds, token_set=bool(self.token), diff --git a/freevideo_engine/launcher_settings.py b/freevideo_engine/launcher_settings.py index 186b40f8..10a19eeb 100644 --- a/freevideo_engine/launcher_settings.py +++ b/freevideo_engine/launcher_settings.py @@ -30,7 +30,7 @@ def default_language(): def sanitize(value): value = plain_data(value) if isinstance(value, dict) else {} result = {k: value[k] for k in TEXT if isinstance(value.get(k), str)} - for key in ('separate', 'new_comfy', 'sampling_caches', 'lan_access'): + for key in ('separate', 'new_comfy', 'sampling_caches', 'lan_access', 'image_model'): if type(value.get(key)) is bool: result[key] = value[key] if value.get('model_method') in ('auto', 'manual', 'reuse'): diff --git a/freevideo_engine/launcher_view.py b/freevideo_engine/launcher_view.py index 9d0a03ec..fc6712c1 100644 --- a/freevideo_engine/launcher_view.py +++ b/freevideo_engine/launcher_view.py @@ -320,7 +320,9 @@ def wheel(event): grid = ttk.Frame(app.pages[2]); grid.pack(fill='x', pady=(0, b.space(3))) grid.columnconfigure((0, 1, 2), weight=1, uniform='models') app.model_cards = {} - for index, name in enumerate(FAMILIES): + # Optional pack controls live in the Qt launcher; keep this legacy setup + # grid limited to the families setup itself provisions. + for index, name in enumerate(name for name in FAMILIES if name != 'image'): card = ModelCard(grid, app, name) card.grid(row=0, column=index, sticky='nsew', padx=(0 if index == 0 else b.space(2), 0 if index == 2 else b.space(2))) diff --git a/freevideo_engine/media_conditioning.py b/freevideo_engine/media_conditioning.py index 341f2157..09e60cb6 100644 --- a/freevideo_engine/media_conditioning.py +++ b/freevideo_engine/media_conditioning.py @@ -19,7 +19,70 @@ def _tensor(value, shape, name): return value.detach().to(device='cpu', dtype=torch.float32).contiguous() +def condition_at(value, step=None): + """None selects the second pass; an integer selects a zero-based first-pass step.""" + selected = value + for stage in value.get('early_reference', {}).get('conditions', []): + if step is None or step >= stage['start_step']: + selected = stage['value'] + return selected + + def describe(value, width, height, frames): + info = _describe(value, width, height, frames) + schedule = value.get('early_reference') + if schedule is None: + if any('early_steps' in row for row in value.get('references', [])): + raise ValueError('early_steps requires separately encoded early reference conditions') + return info + if (not isinstance(schedule, dict) or set(schedule) != {'version', 'early_steps', 'conditions'} + or type(schedule['version']) is not int or schedule['version'] != 1): + raise ValueError('Invalid early reference conditioning version/fields') + from .media_request import validate_early_steps + early = schedule['early_steps'] + if not isinstance(early, list) or not early or len(early) > 32: + raise ValueError('Invalid early_steps metadata') + for row in early: + if not isinstance(row, dict) or set(row) != {'reference_index', 'early_steps'}: + raise ValueError('Invalid early_steps metadata entry') + validate_early_steps(dict(row, kind='video')) + if type(row['reference_index']) is not int or not 0 <= row['reference_index'] < 32: + raise ValueError('Invalid early reference index') + if len({row['reference_index'] for row in early}) != len(early): + raise ValueError('Duplicate early reference index') + if sorted(row['early_steps'] for row in early if row['early_steps']) != sorted( + row['early_steps'] for row in value.get('references', []) if 'early_steps' in row): + raise ValueError('Early reference metadata disagrees with active video references') + stages = schedule['conditions'] + expected = sorted({row['early_steps'] for row in early if row['early_steps'] > 0}) + if not isinstance(stages, list) or any(not isinstance(row, dict) or set(row) != {'start_step', 'value'} for row in stages): + raise ValueError('Invalid early reference conditions') + if [row['start_step'] for row in stages] != expected or any(type(row['start_step']) is not int for row in stages): + raise ValueError('Early reference boundaries disagree with early_steps') + counts = [dict(start_step=0, **info)] + previous = value.get('references', []) + for stage in stages: + current = stage['value'] + if not isinstance(current, dict) or 'early_reference' in current or current.get('keyframe_anchors'): + raise ValueError('Invalid nested early reference condition') + current_info = _describe(current, width, height, frames) + kept = [row for row in previous if row.get('early_steps', stage['start_step'] + 1) > stage['start_step']] + actual = current.get('references', []) + if len(kept) != len(actual) or any(set(a) != set(b) or any( + not torch.equal(a[k], b[k]) if isinstance(a[k], torch.Tensor) else a[k] != b[k] + for k in a) for a, b in zip(kept, actual)): + raise ValueError('Early reference conditions must retain exactly the remaining references') + previous = actual + counts.append(dict(start_step=stage['start_step'], **current_info)) + for name in ('text_tokens', 'reference_video_tokens', 'reference_audio_tokens'): + info[name] = max(row[name] for row in counts) + for index, row in enumerate(counts): + row['condition'] = 'base' if index == len(counts) - 1 else 'full' if index == 0 else 'after_' + str(row['start_step']) + info['early_reference'] = dict(early_steps=early, conditions=counts) + return info + + +def _describe(value, width, height, frames): task = value.get('task', 'fl2va' if value.get('keyframe_anchors') else 't2va') prompt, tags = value['prompt_embeds'], value['text_token_tags'] refs = value.get('references', []) @@ -38,6 +101,9 @@ def describe(value, width, height, frames): video_rows = audio_rows = 0 for ref in refs: kind = ref.get('kind') + if 'early_steps' in ref: + from .media_request import validate_early_steps + validate_early_steps(ref) if kind not in ('image', 'video', 'audio'): raise ValueError('Unknown reference modality') if kind in ('image', 'video'): diff --git a/freevideo_engine/media_encoding.py b/freevideo_engine/media_encoding.py index a81c7798..31da96f2 100644 --- a/freevideo_engine/media_encoding.py +++ b/freevideo_engine/media_encoding.py @@ -133,6 +133,8 @@ def prepare(media, canvas, directory): for index, row in enumerate(media.get('references', [])): kind, path = row['kind'], row['path'] ref = {'kind': kind} + if 'early_steps' in row: + ref['early_steps'] = row['early_steps'] numbers[kind] += 1 longer, length, used = False, None, limit / 24. target = directory / ('reference-%02d.npy' % index) @@ -204,6 +206,17 @@ def prepare(media, canvas, directory): return result, {'images': key_images} if key_images else {'minimax_ref_items': ref_items} if ref_items else {} +def reference_options(normalized, options, step): + """Filter whole references, including a video's optional audio presentation.""" + items, offset = [], 0 + for row in normalized: + count = 1 + int(row['kind'] == 'video' and bool(row.get('audio'))) + if 'early_steps' not in row or step < row['early_steps']: + items.extend(options['minimax_ref_items'][offset:offset + count]) + offset += count + return {'minimax_ref_items': items} if items else {} + + def encode_latents(normalized, value, base, canvas, *, resident=None, device='cuda'): import gc import numpy as np @@ -214,7 +227,7 @@ def encode_latents(normalized, value, base, canvas, *, resident=None, device='cu if torch.device(device).type not in ('cuda', 'mps'): raise ValueError('Media encoding requires a CUDA or MPS device') empty_cache = torch.mps.empty_cache if torch.device(device).type == 'mps' else torch.cuda.empty_cache - refs = [{'kind': row['kind']} for row in normalized] + refs = [{key: row[key] for key in ('kind', 'early_steps') if key in row} for row in normalized] if any(row.get('pixels') for row in normalized): print(json.dumps({'event': 'media_encode_phase', 'phase': 'Encoding image/video references'}), flush=True) # Conditioning never decodes: the 9.0 GiB decoder in the shared shards @@ -265,5 +278,9 @@ def encode_latents(normalized, value, base, canvas, *, resident=None, device='cu condition_latents=[row['latent'] for row in refs]) else: value['references'] = refs + for stage in value.get('early_reference', {}).get('conditions', []): + remaining = [ref for ref in refs if 'early_steps' not in ref or stage['start_step'] < ref['early_steps']] + if remaining: + stage['value']['references'] = remaining from .media_conditioning import describe return describe(value, canvas['width'], canvas['height'], canvas['frames']) diff --git a/freevideo_engine/media_request.py b/freevideo_engine/media_request.py index 70e0abf8..fb6ea662 100644 --- a/freevideo_engine/media_request.py +++ b/freevideo_engine/media_request.py @@ -25,6 +25,8 @@ def read(path): value = json.loads(path.read_text(encoding='utf-8')) if not isinstance(value, dict) or value.get('version') != 1: raise ValueError('Expected a version 1 FreeVideo media request') + if 'early_steps' in value: + raise ValueError('early_steps is only allowed on video references') if set(value) - {'version', 'first', 'last', 'references', 'loras', 'conditioning_info'}: raise ValueError('Unknown media request fields') result = {'version': 1} @@ -49,7 +51,9 @@ def file_info(name): for entry in entries: if not isinstance(entry, dict): raise ValueError('Invalid ' + field + ' entry') - allowed = {'path', 'kind'} if field == 'references' else {'path', 'strength', 'alpha'} + if field != 'references' and 'early_steps' in entry: + raise ValueError('early_steps is only allowed on video references') + allowed = {'path', 'kind', 'early_steps'} if field == 'references' else {'path', 'strength', 'alpha'} if set(entry) - allowed: raise ValueError('Unknown ' + field + ' fields') row = file_info(entry.get('path')) @@ -57,6 +61,9 @@ def file_info(name): if entry.get('kind') not in ('image', 'video', 'audio'): raise ValueError('Reference kind must be image, video or audio') row['kind'] = entry['kind'] + if 'early_steps' in entry: + validate_early_steps(entry) + row['early_steps'] = entry['early_steps'] else: if Path(row['path']).suffix.lower() != '.safetensors': raise ValueError('LoRA requires safetensors; pickle adapters are not loaded') @@ -80,13 +87,31 @@ def file_info(name): def task_for(media): if media.get('conditioning_info'): return media['conditioning_info']['task'] - if media.get('references'): + if any(row.get('early_steps') != 0 for row in media.get('references', [])): return 'ref2va' if media.get('first') and media.get('last'): return 'fl2va' return 'i2va' if media.get('first') else 'l2va' if media.get('last') else 't2va' +def validate_early_steps(row): + if row.get('kind') != 'video': + raise ValueError('early_steps is only allowed on video references') + if type(row['early_steps']) is not int or row['early_steps'] < 0: + raise ValueError('Video reference early_steps must be an integer >= 0') + + +def reference_stages(media): + """Active references at each boundary; zero-step inputs never reach an encoder.""" + refs = media.get('references', []) + for row in refs: + if 'early_steps' in row: + validate_early_steps(row) + boundaries = sorted({0} | {row['early_steps'] for row in refs if 'early_steps' in row}) + return [(step, dict(media, references=[row for row in refs + if 'early_steps' not in row or step < row['early_steps']])) for step in boundaries] + + def cache_key(prompt, media, canvas): # LoRA affects the denoiser, not the encoder. Paths do not identify content. if not any(media.get(k) for k in ('first', 'last', 'references')): diff --git a/freevideo_engine/model_guidance.py b/freevideo_engine/model_guidance.py index 99dc560c..1a9b536f 100644 --- a/freevideo_engine/model_guidance.py +++ b/freevideo_engine/model_guidance.py @@ -17,6 +17,9 @@ from .long_paths import io_path, plain_errors from .model_status import FAMILIES, NAMES, family +# The download list covers the install plan; the image pack has its own launcher controls. +LIST_FAMILIES = tuple(name for name in FAMILIES if name != 'image') + # Iris can replace the third-party download copy here, including export guidance. DOWNLOAD_COPY = { 'title': ('Download list', '下载清单'), @@ -136,7 +139,7 @@ def size(value): copy = download_copy(zh) rows = plan.get('download_files', []) groups = [] - for group in FAMILIES: + for group in LIST_FAMILIES: files = [dict(row, name=row['file'], size=size(row['bytes'])) for row in rows if row['group'] == group] amount = sum(row['bytes'] for row in files) groups.append(dict(id=group, title=copy[group], files=files, count=len(files), bytes=amount, @@ -346,7 +349,7 @@ def links(plan=None): return {group: [(row['file'] + ' · ' + source, row[key]) for row in (plan or {}).get('download_files', []) if row['group'] == group for source, key in (('Hugging Face', 'hf_url'), ('ModelScope', 'modelscope_url')) if row[key]] - for group in FAMILIES} + for group in LIST_FAMILIES} @plain_errors diff --git a/freevideo_engine/model_status.py b/freevideo_engine/model_status.py index 4bd36aaf..fe3d7487 100644 --- a/freevideo_engine/model_status.py +++ b/freevideo_engine/model_status.py @@ -2,13 +2,16 @@ import math import threading -FAMILIES = ('video', 'text', 'decoder', 'sampling') +FAMILIES = ('video', 'text', 'decoder', 'sampling', 'image') NAMES = {'video': ('Video model', '视频模型'), 'text': ('Text encoder', '文本编码器'), 'decoder': ('Video & audio decoder', '视频与音频解码器'), - 'sampling': ('Sampling caches', '采样缓存')} + 'sampling': ('Sampling caches', '采样缓存'), + 'image': ('Image generation models', '图像生成模型')} def family(row): + if row.get('family') == 'image' or row.get('model_id') == 'qwen-image-2.1-turbo': + return 'image' if row.get('sampling_file'): return 'sampling' if not row['repo'].startswith('OpenVDN/'): diff --git a/freevideo_engine/modern_launcher.py b/freevideo_engine/modern_launcher.py index a651d0c3..4fa9b5fd 100644 --- a/freevideo_engine/modern_launcher.py +++ b/freevideo_engine/modern_launcher.py @@ -134,10 +134,12 @@ def refresh(self): # The card says what to do; the details keep the cause. from .failure_details import launcher_failure zh = session.language.startswith('zh') - self.value = dict(self.value, error=session.error, failure=launcher_failure(session.error, zh=zh)) + value = dict(self.value, error=session.error, failure=launcher_failure(session.error, zh=zh)) else: - self.value = dict(self.value, error=session.error) - self.changed.emit() + value = dict(self.value, error=session.error) + if value != self.value: + self.value = value + self.changed.emit() @Slot(str, 'QVariant') def edit(self, key, value): @@ -159,6 +161,10 @@ def cleanupDownloads(self, remove=False): def upgradeModel(self): self.invoke(session.upgrade_model) + @Slot(str) + def imagePackAction(self, action): + self.invoke(session.image_pack_action, action) + @Slot() def cancelModelUpgrade(self): self.invoke(session.cancel_model_upgrade) diff --git a/freevideo_engine/network.py b/freevideo_engine/network.py index 452dbb37..7656b782 100644 --- a/freevideo_engine/network.py +++ b/freevideo_engine/network.py @@ -34,6 +34,7 @@ 'modelscope': 'https://www.modelscope.cn'}, 'edge-models': {'official': 'https://huggingface.co', 'hf-mirror': 'https://hf-mirror.com', 'modelscope': 'https://modelscope.ai'}, + 'image-models': {'official': 'https://huggingface.co', 'modelscope': 'https://www.modelscope.cn'}, 'github': {'official': 'https://github.com', 'ghfast': 'https://ghfast.top/https://github.com'}, 'cuda': {'official': 'https://developer.download.nvidia.com/compute/cuda/redist', 'nvidia-cn': 'https://developer.download.nvidia.cn/compute/cuda/redist'}, @@ -96,11 +97,15 @@ def modelscope_mapping(row, spec=None): def model_family(row, spec=None): + if row.get('family') == 'image': + return 'image-models' mirror = modelscope_mapping(row, spec) return mirror['family'] if mirror else 'models' def model_url(row, name, env=None, *, spec=None): + if row.get('family') == 'image': + return next(s['url'] for s in row['sources'] if s['id'] == name) if name == 'modelscope': mirror = modelscope_mapping(row, spec) if not mirror: @@ -114,9 +119,13 @@ def model_url(row, name, env=None, *, spec=None): def model_urls(network, row, env=None, *, spec=None): family = model_family(row, spec) # Plans made before ModelScope support keep their existing HF ordering. - if family not in network.get('sources', {}): + if family != 'image-models' and family not in network.get('sources', {}): family = 'models' names = ['official'] if network.get('mode') == 'official' or row.get('official_only') else ordered(network, family) + if family == 'image-models' and network.get('mode') != 'official' and not row.get('official_only'): + available = {s['id'] for s in row['sources']} + names = [n for n in names if n in available] + names += [s['id'] for s in row['sources'] if s['id'] not in names] return [(name, model_url(row, name, env, spec=spec)) for name in names] @@ -282,6 +291,9 @@ def probe(family, name, spec, versions, timeout, env, *, measure_speed=False): settings = spec['models'] url = model_url({'repo': settings['encoder_repo'], 'revision': settings['encoder_revision'], 'file': settings['encoder_file']}, name, env, spec=spec) + elif family == 'image-models': + from .image_gen.image_pack import rows + url = model_url(rows()[0], name, env) elif family in ('vdn-models', 'edge-models'): settings = spec['models'] mirror = settings[family.split('-')[0] + '_modelscope'] @@ -320,7 +332,7 @@ def probe(family, name, spec, versions, timeout, env, *, measure_speed=False): component = versions['cuda_toolkits']['13.0']['components']['cuda_cudart'] url = base + '/' + component['relative_path'] headers = () - if (family in ('models', 'vdn-models', 'edge-models') and name == 'official' + if (family in ('models', 'vdn-models', 'edge-models', 'image-models') and name == 'official' and urlsplit(url).scheme == 'https' and urlsplit(url).netloc == 'huggingface.co'): token = env.get('HF_TOKEN') or env.get('HUGGING_FACE_HUB_TOKEN') if token: @@ -376,7 +388,7 @@ def plan(spec, versions, layout='unified', *, mode='auto', timeout=5, offline=Fa raise ValueError('Unknown download connection mode') proxy_configured = proxy.configured(env) git_keys = [] if offline or model_only else proxy.git_proxy_keys(env) - families = ['pypi', 'models', 'vdn-models', 'github', 'git', 'cuda'] + families = ['pypi', 'models', 'vdn-models', 'image-models', 'github', 'git', 'cuda'] if spec['models'].get('edge_modelscope'): families.insert(3, 'edge-models') if versions.get('target_system') in ('Windows', 'Darwin'): @@ -387,7 +399,7 @@ def plan(spec, versions, layout='unified', *, mode='auto', timeout=5, offline=Fa if family not in families: families.append(family) if model_only: - families=[family for family in families if family in ('models','vdn-models','edge-models')] + families=[family for family in families if family in ('models','vdn-models','edge-models','image-models')] selections, jobs = {}, [] for family in families: group = 'torch' if family.startswith('torch-') else 'github' if family == 'git' else family @@ -764,7 +776,7 @@ def download(candidates, path, expected, progress=None, *, network=None, env=Non """Resume across ranked sources; hash before accepting, retain rejected bytes.""" network = network or {} path = Path(path) - model = category in ('models', 'vdn-models', 'edge-models') + model = category in ('models', 'vdn-models', 'edge-models', 'image-models') stall_seconds = stall_seconds if stall_seconds is not None else 120 if model else 20 def matches(candidate, source='local'): @@ -780,7 +792,7 @@ def verification(done, total): check() if model: report(done, total) - return hash_file(candidate, algorithm, git_blob, discard_cache=category in ('models', 'vdn-models'), + return hash_file(candidate, algorithm, git_blob, discard_cache=category in ('models', 'vdn-models', 'image-models'), progress=verification if model or check else None) == expected if path.is_file(): if matches(path): diff --git a/freevideo_engine/policy.py b/freevideo_engine/policy.py index 1c1720d7..b7d21d15 100644 --- a/freevideo_engine/policy.py +++ b/freevideo_engine/policy.py @@ -539,12 +539,18 @@ def feasible_alternatives(hardware, canvas, **options): by this same planner, never estimated separately. """ from .geometry import geometry - keep = {key: canvas[key] for key in ('task', 'steps', 'reference_video_tokens', 'reference_audio_tokens') + keep = {key: canvas[key] for key in ('task', 'steps', 'reference_video_tokens', 'reference_audio_tokens', 'early_reference') if key in canvas} def fits(width, height, frames): try: - choose(hardware, **dict(options, canvas=dict(geometry(width, height, frames=frames), **keep))) + candidate = dict(geometry(width, height, frames=frames), **keep) + if canvas.get('early_reference') is not None and canvas.get('sampling_plan'): + from .two_pass import plan + sampling = canvas['sampling_plan'] + candidate['sampling_plan'] = plan(candidate, sampling['enabled'], canvas.get('task', 't2va'), + base_steps=sampling['base_steps'], refine_steps=sampling['refine_steps'] or 3) + choose(hardware, **dict(options, canvas=candidate)) return True except (ResourceBudgetError, ValueError): return False @@ -739,7 +745,8 @@ def choose(hardware: Hardware, *, vram_gib=None, ram_gib=None, attention='auto', if type(count) is not int or count < 0: raise ValueError('Reference token counts must be nonnegative integers') reference_rows += count - effective_tokens = canvas['video_tokens'] + reference_rows if canvas is not None else None + from .two_pass import placement_tokens + effective_tokens = placement_tokens(canvas) if canvas is not None else None compatibility = hardware.cuda_compatibility() if compatibility['error']: raise ValueError(compatibility['error']) @@ -1483,6 +1490,11 @@ def held(group): decoder_offload = gpu_budget < 20 * GiB if reference_rows: notes.append('Reference conditioning adds %d packed rows; activation placement scales with total rows. This is a starting estimate, not a verified reference-capacity result.' % reference_rows) + if canvas is not None and canvas.get('early_reference') is not None: + from .two_pass import stage_tokens + rows = stage_tokens(canvas) + notes.append('Early reference placement: first pass %d rows, refinement %s rows; fixed engine covers %d rows.' + % (rows['first'], rows['refinement'], rows['engine'])) # A full floating video plus its RGB copy grows with output geometry. # The temporal streaming path uses the original decoder tiles and blend; # completed uint8 frames go to the retained disk artifact immediately. diff --git a/freevideo_engine/prediction.py b/freevideo_engine/prediction.py index 21858f8a..b983e089 100644 --- a/freevideo_engine/prediction.py +++ b/freevideo_engine/prediction.py @@ -84,6 +84,8 @@ def _canvas(value, options, text_tokens=None): canvas['task'] = options['engine']['task'] if value.get('sampling_plan', {}).get('enabled') or value.get('sampling_plan', {}).get('version') == 2: canvas['sampling_plan'] = value['sampling_plan'] + if 'early_reference' in value: + canvas['early_reference'] = value['early_reference'] if type(canvas['steps']) is not int or canvas['steps'] <= 0: raise ValueError('Forecast steps must be a positive integer') from .media_request import TASKS @@ -446,7 +448,7 @@ def reject(reason): reject('missing_observed_geometry') continue if any(key in original and key in enriched and original[key] != enriched[key] - for key in ('width', 'height', 'frames', 'steps', 'task', 'reference_video_tokens', 'reference_audio_tokens')): + for key in ('width', 'height', 'frames', 'steps', 'task', 'reference_video_tokens', 'reference_audio_tokens', 'early_reference')): reject('inconsistent_observed_geometry') continue canvas = _canvas(enriched, options, observation.get('text_tokens')) diff --git a/freevideo_engine/prompt_vlm/service.py b/freevideo_engine/prompt_vlm/service.py index 540be612..481c4f07 100644 --- a/freevideo_engine/prompt_vlm/service.py +++ b/freevideo_engine/prompt_vlm/service.py @@ -350,6 +350,16 @@ def close(self): from .. import processes processes.stop(process, grace=5) + @staticmethod + def release_idle_image_worker(): + # An image worker that is only waiting out its idle time still holds the engine lease. + try: + from server import PromptServer + from ..image_gen.service import release_idle_worker + release_idle_worker(PromptServer.instance) + except Exception: + pass + def work(self, root, machine, action, value): from ..locking import runtime_lock from ..monitoring import save @@ -357,6 +367,7 @@ def work(self, root, machine, action, value): try: from ..background_compile import COMPILE COMPILE.stop() + self.release_idle_image_worker() if action in ('rewrite', 'ask'): from ..resident_process import OWNER if OWNER.active: diff --git a/freevideo_engine/reference_sampler.py b/freevideo_engine/reference_sampler.py index 48496b65..6a011f70 100644 --- a/freevideo_engine/reference_sampler.py +++ b/freevideo_engine/reference_sampler.py @@ -29,7 +29,7 @@ @torch.no_grad() def generate_latents(transformer, prompt_embeds, text_token_tags, num_frames, num_steps, seed, device, video_shift=12.0, audio_shift=3.0, runtime=None, step_seconds=None, - conditions=None): + conditions=None, condition_schedule=None, condition_trace=None, timestep_task=None): """Preserve the pinned sampler's RNG order and generated-row updates. Reference visual rows are held at max(video_t, .999) and reference audio @@ -50,31 +50,36 @@ def generate_latents(transformer, prompt_embeds, text_token_tags, num_frames, nu channels = transformer.config.in_channels # 24 frame_h, frame_w = LATENT_H // patch[1], LATENT_W // patch[2] - position_ids, token_tags, video_indices, audio_indices, text_indices, num_condition_rows, num_audio_condition_rows = ( - MiniMaxH3Ref2VAPrepareLayoutStep.build_ref2va_packed_sequence( - text_token_tags, references, condition_latents, audio_conditions, num_latent_frames, LATENT_H, LATENT_W, num_audio_latents, - patch, AUDIO_CHANNELS, AUDIO_TAG, VIDEO_TAG, + def make_layout(text_token_tags, references, condition_latents, audio_conditions): + position_ids, token_tags, video_indices, audio_indices, text_indices, num_condition_rows, num_audio_condition_rows = ( + MiniMaxH3Ref2VAPrepareLayoutStep.build_ref2va_packed_sequence( + text_token_tags, references, condition_latents, audio_conditions, num_latent_frames, LATENT_H, LATENT_W, num_audio_latents, + patch, AUDIO_CHANNELS, AUDIO_TAG, VIDEO_TAG, + ) ) - ) - if torch.device(device).type == 'mps': - # RoPE consumes FP32 coordinates; MPS cannot stage the layout's FP64 - # storage first. Keep the existing CUDA transfer and arithmetic. - position_ids = position_ids.float() - position_ids, token_tags = position_ids.to(device), token_tags.to(device) - video_indices, audio_indices, text_indices = ( - video_indices.to(device), audio_indices.to(device), text_indices.to(device), - ) - - # A hybrid-converted transformer needs the packed layout before every forward; the - # layout is the same at every denoising step, so set it once per generation. - if next(iter_hybrids(transformer), None) is not None: - # frame_size and text_indices unconditionally: carrying them is free, and only - # their consumers are gated (short_conv / text_state). - set_layout(transformer, layout_from_indices( - video_indices[num_condition_rows:], num_latent_frames, frame_h * frame_w, - seq_len=position_ids.shape[0], frame_size=(frame_h, frame_w), - text_indices=text_indices, - )) + if torch.device(device).type == 'mps': + # RoPE consumes FP32 coordinates; MPS cannot stage the layout's FP64 + # storage first. Keep the existing CUDA transfer and arithmetic. + position_ids = position_ids.float() + position_ids, token_tags = position_ids.to(device), token_tags.to(device) + video_indices, audio_indices, text_indices = ( + video_indices.to(device), audio_indices.to(device), text_indices.to(device), + ) + + if next(iter_hybrids(transformer), None) is not None: + # frame_size and text_indices unconditionally: carrying them is free, and only + # their consumers are gated (short_conv / text_state). + set_layout(transformer, layout_from_indices( + video_indices[num_condition_rows:], num_latent_frames, frame_h * frame_w, + seq_len=position_ids.shape[0], frame_size=(frame_h, frame_w), + text_indices=text_indices, + )) + + return (position_ids, token_tags, video_indices, audio_indices, text_indices, + num_condition_rows, num_audio_condition_rows) + + position_ids, token_tags, video_indices, audio_indices, text_indices, num_condition_rows, num_audio_condition_rows = ( + make_layout(text_token_tags, references, condition_latents, audio_conditions)) scheduler = step_grid.scheduler(shift=video_shift) audio_scheduler = step_grid.scheduler(shift=audio_shift) @@ -115,15 +120,58 @@ def synchronize(): runtime.barrier() synchronize() - seq_len = position_ids.shape[0] - for t, audio_t in zip(scheduler.timesteps, audio_scheduler.timesteps): + visual_refs = [row for row in refs if row['kind'] in ('image', 'video')] + audio_refs = [row for row in refs if row.get('audio_latent') is not None] + stages = {row['start_step']: row for row in condition_schedule or []} + current_start = 0 + for step, (t, audio_t) in enumerate(zip(scheduler.timesteps, audio_scheduler.timesteps)): step_started = time.perf_counter() + if step in stages: + stage = stages[step] + value = stage['value'] + # Drop only fixed rows. Never restart a scheduler or draw new noise. + generated_video = video_rows[num_condition_rows:] + generated_audio = audio_rows[num_audio_condition_rows:] + keep_video = [i for i, row in enumerate(visual_refs) if row.get('early_steps', step + 1) > step] + keep_audio = [i for i, row in enumerate(audio_refs) if row.get('early_steps', step + 1) > step] + references = [SimpleNamespace(kind=row['kind'], has_audio=row.get('audio_latent') is not None) + for row in value.get('references', [])] + prompt_embeds = value['prompt_embeds'].to(device=device, dtype=prompt_embeds.dtype) + previous_audio_condition_rows = num_audio_condition_rows + position_ids, token_tags, video_indices, audio_indices, text_indices, num_condition_rows, num_audio_condition_rows = ( + make_layout(value['text_token_tags'], references, + [condition_latents[i] for i in keep_video], [audio_conditions[i] for i in keep_audio])) + video_rows = torch.cat([condition_rows[i] for i in keep_video] + [generated_video]) + # A silent proxy does not change the fixed audio prefix. Reuse + # its existing packed storage, including the current generated + # rows; compact only when reference audio actually disappears. + if num_audio_condition_rows != previous_audio_condition_rows: + audio_rows = torch.cat([audio_conditions[i] for i in keep_audio] + [generated_audio]) + if 'refined_text' in stage: + transformer._freevideo_refined_text = stage['refined_text'].to(device) + if video_rows.shape[0] != video_indices.numel() or audio_rows.shape[0] != audio_indices.numel(): + raise ValueError('Switched reference rows disagree with packed layout') + current_start = step + seq_len = position_ids.shape[0] + if condition_trace is not None: + condition_trace.append(dict(step=step, condition_start_step=current_start, + condition=('base' if current_start == max(stages, default=0) else + 'full' if current_start == 0 else 'after_' + str(current_start)), + reference_video_rows=num_condition_rows, reference_audio_rows=num_audio_condition_rows, + text_rows=text_indices.numel(), total_rows=seq_len)) row_timesteps = torch.full((seq_len,), float(t), dtype=torch.float32, device=device) if num_condition_rows: row_timesteps[video_indices[:num_condition_rows]] = max(float(t), KEYFRAME_NOISE_AUG) row_timesteps[audio_indices[num_audio_condition_rows:]] = float(audio_t) row_timesteps[audio_indices[:num_audio_condition_rows]] = 1.0 timestep, timestep_indices = torch.unique(row_timesteps, sorted=True, return_inverse=True) + if timestep_task is not None: + from .adaln import modality_timesteps + # The cached table covers the initial superset of modalities. Keep + # unused clock entries so removing the last visual/audio reference + # does not change the table or the remaining rows' clock indices. + timestep = modality_timesteps(t[None], audio_t[None], timestep_task)[0] + timestep_indices = torch.searchsorted(timestep, row_timesteps) noise_pred, audio_noise_pred = transformer( hidden_states=video_rows[None], audio_hidden_states=audio_rows[None], diff --git a/freevideo_engine/refine.py b/freevideo_engine/refine.py index c180a252..ee9373d7 100644 --- a/freevideo_engine/refine.py +++ b/freevideo_engine/refine.py @@ -22,7 +22,7 @@ def validate_tail(base_steps, refine_steps): def generate_latents(transformer, prompt_embeds, text_token_tags, num_frames, num_steps, seed, device, *, initial_latents, refine_steps, step_seconds=None, - conditions=None, video_shift=12., audio_shift=3., refine_schedule=None): + conditions=None, video_shift=12., audio_shift=3., refine_schedule=None, timestep_task=None): from .step_grid import scheduler from diffusers.modular_pipelines.minimax_h3.before_denoise import ( MiniMaxH3PrepareLayoutStep, MiniMaxH3Ref2VAPrepareLayoutStep, patchify_video_latents) @@ -126,6 +126,10 @@ def generate_latents(transformer, prompt_embeds, text_token_tags, num_frames, nu row_times[audio_ids[audio_condition_rows:]] = float(audio_t) row_times[audio_ids[:audio_condition_rows]] = 1. times, time_ids = torch.unique(row_times, sorted=True, return_inverse=True) + if timestep_task is not None: + from .adaln import modality_timesteps + times = modality_timesteps(video_t[None], audio_t[None], timestep_task)[0] + time_ids = torch.searchsorted(times, row_times) prediction, _ = transformer(hidden_states=video_rows[None], audio_hidden_states=audio_rows[None], encoder_hidden_states=prompt_embeds[None], timestep=times, timestep_indices=time_ids, token_tags=tags, position_ids=positions, video_indices=video_ids, audio_indices=audio_ids, diff --git a/freevideo_engine/release_notes.json b/freevideo_engine/release_notes.json index 7922bc6f..a7c70bf0 100644 --- a/freevideo_engine/release_notes.json +++ b/freevideo_engine/release_notes.json @@ -1,41 +1,74 @@ { "schema": 1, - "product_version": "0.4.1", + "product_version": "0.4.2", "en": { - "summary": "Fixes Extend on videos made earlier, and improves progress display and several agent details.", + "summary": "Adds optional image generation models: generate and edit pictures, make several views for character cards, and let the agent use them. Also fixes several install and launcher issues.", "changes": [ - "Extend works again on videos made earlier, including ones whose size was changed, and the new part now uses the previous video's picture size automatically. Clicking Extend on a result or in Creations now opens Edit with the video on the timeline and a new part already placed after it, ready to generate.", - "Autopilot can start a new conversation, and earlier ones stay in a list to return to; the agent also keeps listening after it answers, so it still responds when the conversation resumes later.", - "A percentage and the time left appear as soon as a video starts; progress keeps moving while the model loads and runs smoothly to the end, on Autopilot job cards as well.", - "In Edit, an agent can start a new timeline from just a name, and writes a short note on each clip so it is clear at a glance what the clip is.", - "A video added to the timeline from Creations now brings the description it was generated from, so it can be regenerated right away.", - "A video of 16 to 29 seconds requested through an agent is now made in one generation when the computer allows it, the same as in Create.", - "When a specific camera path, placement or order of actions keeps coming out wrong, the agent can offer to make a rough camera draft first and generate from it.", - "The question about freeing the disk space used by old model files now appears after updating.", - "At smaller sizes, more of the model stays in GPU memory, so computers with less memory read less from the drive while generating.", - "When the network is slow or a step takes a while, the install page now shows what it is doing.", - "Third-party download during setup now lists only the files this computer needs, and the list can be exported; after downloading, choose the folder and the files that pass the check are used directly.", - "When other programs use a large part of the GPU during a generation, or the GPU overheats, a short note appears under the result." + "Adds optional image generation models. \"Generate with AI\" in the media panel makes a picture from one sentence on a blank canvas, and \"Edit this\" in the image editor changes the current picture, with up to two more pictures from your media as references.", + "With the image generation models installed, character cards start from a front view and use it as the reference for three-quarter, side, back and close-up views; scene cards are made the same way.", + "The agent can generate or edit pictures and add them to your media, and can first make a frame with people placed as described, then generate the video from it.", + "Setup can install the image generation models at the same time (about 16 GB), and existing installs can download them in the background from the launcher; the download pauses while a video is generating.", + "Online installs into a folder the system code page cannot spell, such as a Chinese folder on an English system, no longer fail at the kernel check.", + "On Linux, when the system compiler is newer than the installed GPU toolkit supports, setup now uses a compatible compiler instead of stopping while building the GPU attention component.", + "When a Windows security feature blocks files FreeVideo needs, the error card now says so and notes that the setting may need to be changed in Windows.", + "After a launcher error, an address or path being typed in the settings is no longer reset to its previous value.", + "When the agent makes a film, a shot that continues the previous one is at least 10 seconds long, 10 seconds by default; a length this computer cannot make in one run is reported before generation starts.", + "In AI Autopilot, the conversation list updates as soon as a job starts or stops, and switching conversations quickly stays on the one selected." ] }, "zh": { - "summary": "修复了之前生成的视频不能“接着生成”的问题，并改进了进度显示和 Agent 的几处体验。", + "summary": "新增可选的图像生成模型，可以生成和修改图片，为角色卡生成多个角度，Agent 也能使用；并修复了几处安装和启动器的问题。", "changes": [ - "之前生成的视频（包括换过尺寸的）都可以“接着生成”，不再提示不能接续；新的一段会自动用上一段的画面大小。在结果或作品上点“接着生成”，会打开剪辑视图：原视频已在时间线上，后面接好了新的一段，点“生成”继续。", - "“AI 全自动”里可以开新对话，旧对话留在列表里，可以回去接着做；Agent 答完以后会一直等下一句，隔一会儿再说也有回应。", - "点“生成”后马上显示百分比和预计剩余时间；加载模型时进度照常前进，平稳走到结束，“AI 全自动”的任务卡也一样。", - "剪辑视图里，Agent 只凭一个名字就能新建时间线，并会在每一段上写一句说明，一眼就能看出这一段是什么。", - "从“作品”放上时间线的视频会带上当时的生成描述，可以直接重新生成。", - "通过 Agent 生成 16–29 秒的视频时，电脑条件允许就一次生成，和创作面板一致。", - "指定的镜头路线、位置或动作顺序总是生成得不对时，Agent 会提议先做一个粗略的镜头草稿，再照草稿生成。", - "更新后询问是否释放旧模型文件所占硬盘空间的提示，现在会正常出现。", - "画面较小时，更多模型留在显存里，内存较小的电脑生成时读取硬盘更少。", - "安装时网络较慢或某一步较久时，安装页会显示当前正在做什么。", - "安装时选第三方下载，现在只列出这台电脑需要的文件，并可以导出下载清单；下载完成后选中那个文件夹，校验通过的文件直接使用。", - "生成期间其他程序占用了较多显卡，或显卡过热时，结果下方会出现一句简短的提示。" + "新增可选的图像生成模型。素材面板的“AI 生成”可以在空白画布上用一句话生成一张图；图片编辑器里的“改这张”可以修改当前图片，还可以再选最多两张素材图作参考。", + "装有图像生成模型时，角色卡先生成正面像，再以它为参考生成四分之三侧面、侧面、背面和脸部特写；场景卡也用它生成。", + "Agent 可以生成或修改图片并放进素材，也可以先按描述的人物站位生成一张首帧，再用它生成视频。", + "安装时可以选择同时安装图像生成模型（约 16 GB）；已经安装 FreeVideo 的，可以在启动器里下载，下载在后台进行，生成视频时自动暂停。", + "英文系统上在线安装到中文文件夹时，内核检查不再因为路径无法识别而失败。", + "Linux 上系统编译器比 GPU 工具包支持的版本更新时，安装会换用兼容的编译器，不再停在编译 GPU 注意力加速组件这一步。", + "Windows 的安全功能拦下运行所需的文件时，错误卡会说明原因，并提示可能需要在 Windows 中调整这项设置。", + "启动器出错后，设置里正在输入的地址或路径不再被换回原值。", + "Agent 拍片时，接上一镜的镜头每段至少 10 秒，没写长度时默认 10 秒；这台电脑一次生成不了的长度会在开始前说明。", + "“AI 全自动”里任务开始或停止后，对话列表立刻更新；快速切换对话时停在所选的那一段。" ] }, "history": [ + { + "product_version": "0.4.1", + "en": { + "summary": "Fixes Extend on videos made earlier, and improves progress display and several agent details.", + "changes": [ + "Extend works again on videos made earlier, including ones whose size was changed, and the new part now uses the previous video's picture size automatically. Clicking Extend on a result or in Creations now opens Edit with the video on the timeline and a new part already placed after it, ready to generate.", + "Autopilot can start a new conversation, and earlier ones stay in a list to return to; the agent also keeps listening after it answers, so it still responds when the conversation resumes later.", + "A percentage and the time left appear as soon as a video starts; progress keeps moving while the model loads and runs smoothly to the end, on Autopilot job cards as well.", + "In Edit, an agent can start a new timeline from just a name, and writes a short note on each clip so it is clear at a glance what the clip is.", + "A video added to the timeline from Creations now brings the description it was generated from, so it can be regenerated right away.", + "A video of 16 to 29 seconds requested through an agent is now made in one generation when the computer allows it, the same as in Create.", + "When a specific camera path, placement or order of actions keeps coming out wrong, the agent can offer to make a rough camera draft first and generate from it.", + "The question about freeing the disk space used by old model files now appears after updating.", + "At smaller sizes, more of the model stays in GPU memory, so computers with less memory read less from the drive while generating.", + "When the network is slow or a step takes a while, the install page now shows what it is doing.", + "Third-party download during setup now lists only the files this computer needs, and the list can be exported; after downloading, choose the folder and the files that pass the check are used directly.", + "When other programs use a large part of the GPU during a generation, or the GPU overheats, a short note appears under the result." + ] + }, + "zh": { + "summary": "修复了之前生成的视频不能“接着生成”的问题，并改进了进度显示和 Agent 的几处体验。", + "changes": [ + "之前生成的视频（包括换过尺寸的）都可以“接着生成”，不再提示不能接续；新的一段会自动用上一段的画面大小。在结果或作品上点“接着生成”，会打开剪辑视图：原视频已在时间线上，后面接好了新的一段，点“生成”继续。", + "“AI 全自动”里可以开新对话，旧对话留在列表里，可以回去接着做；Agent 答完以后会一直等下一句，隔一会儿再说也有回应。", + "点“生成”后马上显示百分比和预计剩余时间；加载模型时进度照常前进，平稳走到结束，“AI 全自动”的任务卡也一样。", + "剪辑视图里，Agent 只凭一个名字就能新建时间线，并会在每一段上写一句说明，一眼就能看出这一段是什么。", + "从“作品”放上时间线的视频会带上当时的生成描述，可以直接重新生成。", + "通过 Agent 生成 16–29 秒的视频时，电脑条件允许就一次生成，和创作面板一致。", + "指定的镜头路线、位置或动作顺序总是生成得不对时，Agent 会提议先做一个粗略的镜头草稿，再照草稿生成。", + "更新后询问是否释放旧模型文件所占硬盘空间的提示，现在会正常出现。", + "画面较小时，更多模型留在显存里，内存较小的电脑生成时读取硬盘更少。", + "安装时网络较慢或某一步较久时，安装页会显示当前正在做什么。", + "安装时选第三方下载，现在只列出这台电脑需要的文件，并可以导出下载清单；下载完成后选中那个文件夹，校验通过的文件直接使用。", + "生成期间其他程序占用了较多显卡，或显卡过热时，结果下方会出现一句简短的提示。" + ] + } + }, { "product_version": "0.4.0", "en": { @@ -219,44 +252,6 @@ "修复：部分广告拦截列表会按文件名拦下 FreeVideo 的一个脚本，导致创作面板不显示，现在页面可以正常加载，页面检查也会指出被浏览器插件拦截的文件。在节点视图里把视频放大到全屏时，不再出现“没有完整加载”的提示。较新的 ComfyUI 不再在每个视频完成后打印资源登记警告。第二遍需要把部分数据放在内存里时（例如 8 GB Windows 显卡），本地优化可以正常运行。" ] } - }, - { - "product_version": "0.3.20", - "en": { - "summary": "Offline installations now open FreeVideo in ComfyUI, and on Linux, GPUs with memory to spare generate faster.", - "changes": [ - "Offline installations imported from the environment package now open FreeVideo in ComfyUI: the launcher installs the FreeVideo node that the package does not carry. Before, ComfyUI started without FreeVideo and the launch stopped at \"could not load FreeVideo\". A retry now also restarts the ComfyUI that an earlier attempt left running, and the message no longer suggests repairing the environment.", - "Offline installations in a longer folder path now also get their desktop and Start menu shortcuts. Windows cannot create a shortcut to a program whose path is 260 characters or longer, and the folders of an offline installation could make the launcher copy's path that long; the copy now goes to a shorter folder in that case.", - "On Linux, when the GPU has memory to spare, attention now runs 28 or all 56 heads at a time instead of 16; how many model layers stay on the GPU, and what is kept in system memory, does not change. On a 96 GB GPU, sampling a 10-second 1344×768 video took about 11% less time; with a 16–20 GB budget the gain is about 6–9%. Smaller GPUs and Windows keep their current settings. The same seed can give a slightly different video than before, as it already does between GPUs with different memory.", - "On Linux, video decoding now uses its FP16 cache on every GPU (it was on for some only). Decoding a 10-second 1344×768 video on a 96 GB GPU took 15.4 s instead of 19.2 s and 4.2 GB less memory, with identical output." - ] - }, - "zh": { - "summary": "离线安装可以在 ComfyUI 中打开 FreeVideo；在 Linux 上，显存有富余的显卡生成更快。", - "changes": [ - "用环境包导入的离线安装现在可以在 ComfyUI 中打开 FreeVideo：启动器会安装环境包里没有带的 FreeVideo 节点。之前 ComfyUI 启动时没有加载 FreeVideo，启动停在“无法加载 FreeVideo”。重试时也会重启上一次留下的 ComfyUI，提示也不再让你去修复运行环境。", - "离线安装放在较长的文件夹路径下时，现在也能创建桌面和开始菜单快捷方式。Windows 无法为路径达到 260 个字符的程序创建快捷方式，而离线安装的目录层级可能让启动器副本的路径超过这个长度；这种情况下，副本会改放到路径较短的位置。", - "在 Linux 上，显存有富余时，注意力改为每次计算 28 个或全部 56 个头（原来每次 16 个）；留在显卡上的模型层数和放在内存里的内容都不变。在 96 GB 显卡上，生成 10 秒 1344×768 视频的采样时间缩短约 11%；显存预算在 16–20 GB 时约缩短 6–9%。显存较小的显卡和 Windows 保持现有设置。同一个种子生成的视频可能与以前略有不同，就像在显存不同的显卡上本来就会有差异一样。", - "在 Linux 上，所有显卡的视频解码都使用 FP16 缓存（以前只有部分显卡开启）。在 96 GB 显卡上解码 10 秒 1344×768 视频从 19.2 秒降到 15.4 秒，少占 4.2 GB 显存，输出完全一致。" - ] - } - }, - { - "product_version": "0.3.19", - "en": { - "summary": "On Windows, GPUs with 8 or 10 GB of memory sample the first pass a little faster when a video uses reference images, a first or last frame, or reference audio.", - "changes": [ - "On Windows GPUs with 8 or 10 GB of memory, the first pass of a video up to 10 seconds long now runs attention 8 heads at a time when the video uses reference images, a first or last frame, or reference audio, as text-only videos already did, instead of 4. With one reference image, a 10-second 1344×768 video took about 8% less time per first-pass step on an 8 GB budget, staying within the same video memory budget. The same seed can give a slightly different video than before, as it already does between GPUs with different memory.", - "On Windows, the model layers kept on the GPU between first-pass steps no longer shrink from step to step, so fewer layers are transferred again on each step." - ] - }, - "zh": { - "summary": "在 Windows 上，8 GB 和 10 GB 显存的显卡生成带参考图、首尾帧或参考音频的视频时，第一遍采样稍快一些。", - "changes": [ - "在 8 GB 和 10 GB 显存的 Windows 显卡上，10 秒以内的视频如果带参考图、首帧或尾帧、参考音频，第一遍注意力改为每次计算 8 个头（原来是 4 个），与纯文本生成一致。带一张参考图生成 10 秒 1344×768 视频时，在 8 GB 显存预算下第一遍每步缩短约 8%，显存仍在原来的预算之内。同一个种子生成的视频可能与以前略有不同，就像在显存不同的显卡上本来就会有差异一样。", - "在 Windows 上，第一遍步与步之间保留在显卡上的模型层不再逐步减少，每步需要重新传输的层更少。" - ] - } } ] } diff --git a/freevideo_engine/resident_models.py b/freevideo_engine/resident_models.py index 4beeae5e..e45cc755 100644 --- a/freevideo_engine/resident_models.py +++ b/freevideo_engine/resident_models.py @@ -355,6 +355,9 @@ def engine(self, request, factory): found = None hit = found is not None rows = canvas.get('video_tokens', 72576) + canvas.get('reference_video_tokens', 0) + canvas.get('reference_audio_tokens', 0) + if canvas.get('early_reference') is not None: + from .two_pass import placement_tokens + rows = placement_tokens(canvas) from .policy import BLOCK_BYTES, activation_bytes, windows_gpu_output_workspace from .system import windows # Use the same measured token/head envelopes as placement. The old diff --git a/freevideo_engine/runtime.py b/freevideo_engine/runtime.py index cda824ba..172dae1a 100644 --- a/freevideo_engine/runtime.py +++ b/freevideo_engine/runtime.py @@ -44,6 +44,11 @@ def _load_safetensors(path): return {name: handle.get_tensor(name) for name in handle.keys()} +def refined_text_memo_capacity(canvas): + """Stage prompts of an early-reference canvas plus the top-level prompt and one spare.""" + return len(((canvas or {}).get('early_reference') or {}).get('conditions') or []) + 2 + + class Engine: @torch.no_grad() def __init__(self, cache, *, base=None, checkpoint=None, @@ -127,7 +132,8 @@ def __init__(self, cache, *, base=None, checkpoint=None, self.stream_refiner_count = 0 self.cache_refined_text = cache_refined_text self.refiner_released = False - self.previous_refined_text = None + self.refined_text_memo = {} + self.refined_text_memo_capacity = refined_text_memo_capacity(self.canvas) started = time.perf_counter() manifest = json.loads((self.cache / 'manifest.json').read_text(encoding='utf-8')) online_lora = manifest.get('online_lora', {}) @@ -582,15 +588,38 @@ def _sample(self, conditioning, seed, frames=243, width=1344, height=768, *, fin offset=progress_offset, total=progress_total) progress.start() continuation_offset = 0 + early_receipt = None + condition_schedule = None + sampling_kwargs = {} + if self.task == 't2va' and (self.canvas or {}).get('early_reference'): + # A zero-step proxy may leave an ordinary text-only request. Keep + # that sampler and its exact RNG path, but still report the cutoff. + early_receipt = dict(self.canvas['early_reference'], steps=[dict(step=i, + condition='base', condition_start_step=0, reference_video_rows=0, reference_audio_rows=0, + text_rows=self.canvas['early_reference']['conditions'][0]['text_tokens']) for i in range(active_steps)]) if self.task.startswith('ref2va'): from .media_conditioning import describe value = torch.load(conditioning, map_location='cpu', weights_only=True) info = describe(value, width, height, frames) if info['task'] != self.task: raise ValueError('Reference modalities changed the precomputed schedule') - conditions = value['references'] + if 'early_reference' in info: + from .media_conditioning import condition_at + early_receipt = dict(info['early_reference'], steps=[]) + if value['early_reference']['conditions']: + sampling_kwargs['timestep_task'] = self.task + if refining: + value = condition_at(value) + row = info['early_reference']['conditions'][-1] + early_receipt['steps'] = [dict(step=i, condition='base', condition_start_step=row['start_step'], + reference_video_rows=row['reference_video_tokens'], reference_audio_rows=row['reference_audio_tokens'], + text_rows=row['text_tokens']) for i in range(active_steps)] + else: + condition_schedule = [dict(stage) for stage in value['early_reference']['conditions']] + sampling_kwargs.update(condition_schedule=condition_schedule, condition_trace=early_receipt['steps']) + conditions = value.get('references') or None if continuation is not None: - if len(conditions) != 1 or conditions[0]['kind'] != 'image' or conditions[0].get('audio_latent') is not None: + if conditions is None or len(conditions) != 1 or conditions[0]['kind'] != 'image' or conditions[0].get('audio_latent') is not None: raise ValueError('Continuation anchor requires exactly one image reference without audio') # Ref2VA packs independent-size references before generated # video rows in both passes. Their latents do not need the @@ -761,10 +790,6 @@ def append(self, value): self.device_backend.synchronize() started = time.perf_counter() refinement = None - if self.cache_refined_text: - refinement = self._prepare_refined_text(prompt) - if active_steps == 1: - trigger_prefetch() checkpoint_seconds = 0. # A budget cut inside a step (Windows): see step_pressure. The watch is # the innermost context, so a narrowed group is put back before the @@ -777,6 +802,25 @@ def append(self, value): from .step_pressure import LayerWatch watch = LayerWatch(self, source, pass_cache=pass_cache) try: + if self.cache_refined_text: + if condition_schedule: + # An engine serves several requests; size the memo for this + # request's stage prompts plus the top-level prompt. + self.refined_text_memo_capacity = max(self.refined_text_memo_capacity, len(condition_schedule) + 2) + for stage in condition_schedule: + # Prepare the final subset too, even when K >= steps: + # refinement always uses it. Keep host refiner weights + # until every distinct prompt has been prepared. + if stage['start_step'] >= self.steps and stage is not condition_schedule[-1]: + continue + self._prepare_refined_text(stage['value']['prompt_embeds'].to(self.device_backend.device, torch.bfloat16), + release_weights=False) + # Share the memo's CPU storage; another device-to-host + # copy here would double every stage's retained text. + stage['refined_text'] = next(reversed(self.refined_text_memo.values())) + refinement = self._prepare_refined_text(prompt) + if active_steps == 1: + trigger_prefetch() with progress.layers(self.transformer.transformer_blocks), self.device_backend.make_offloader( self.offload_layers, prefetch=self.prefetch, weight_source=self.stream_weight_source) as offloader, sampling( self.transformer, continuation, seed, self.device_backend.device, @@ -784,7 +828,7 @@ def append(self, value): (watch.layers(self.transformer.transformer_blocks) if watch is not None else nullcontext()): latents, audio = generate_latents(self.transformer, prompt, tags, frames, self.steps, seed, self.device_backend.device, step_seconds=step_seconds, - conditions=conditions) + conditions=conditions, **sampling_kwargs) if prefix is not None: latents, audio = prefix.finish(latents, audio) stats = offloader.stats() @@ -819,6 +863,8 @@ def append(self, value): start_index=1 if refine_schedule else self.steps - refine_steps, restart_seed=seed, schedule=refine_schedule or 'original-tail', audio_policy='first_pass_preserved_with_audio_clock_conditioning') + if early_receipt is not None: + sampled['early_reference'] = early_receipt # The final latents must be durable before exiting the weight # offloader: its cleanup can fail after all NFE have completed. if sampling_complete_callback is not None: @@ -827,6 +873,10 @@ def append(self, value): checkpoint_seconds = time.perf_counter() - tick finally: self.transformer._freevideo_refined_text = None + if condition_schedule and self.cache_refined_text and not self.stream_refiner_count: + # Also release a partially prepared batch when refinement fails. + self.transformer.token_refiner.refiner_blocks.to('meta') + self.refiner_released = True self.device_backend.synchronize() elapsed = time.perf_counter() - started - checkpoint_seconds self.device_backend.empty_cache() @@ -838,16 +888,18 @@ def append(self, value): return latents, audio, sampled @torch.no_grad() - def _prepare_refined_text(self, prompt): + def _prepare_refined_text(self, prompt, *, release_weights=True): """The upstream text refiner has no timestep or video input. - Compute it once for this request and discard its 0.77 GB of weights. - Restore only those tensors from the immutable cache for a later prompt. - The cached tensor is cleared even when sampling fails. + Keep a bounded LRU of CPU results for the stage and top-level prompts. + A batch retains host weights until its last prompt, then discards them. + Sampling clears the active device tensor even when preparation fails. """ from .fp8 import input_dtype started = time.perf_counter() refiner = self.transformer.token_refiner.refiner_blocks + if prompt.ndim not in (2, 3) or (prompt.ndim == 3 and prompt.shape[0] != 1): + raise ValueError('Expected one text prompt for refinement') cpu = prompt.detach().to('cpu').contiguous() cache_key = str(cpu.dtype) + str(list(cpu.shape)) + hashlib.sha256(cpu.view(torch.uint8).numpy().tobytes()).hexdigest() del cpu @@ -856,9 +908,23 @@ def _prepare_refined_text(self, prompt): # reuses the first pass's tensor instead of reloading the refiner # weights to compute the same values again. memo = (cache_key, tuple(self.config.get(k) for k in ('attention', 'linear_compute', 'fp8_gemm', 'inference_kernels'))) - if self.previous_refined_text is not None and self.previous_refined_text[0] == memo: - text = self.previous_refined_text[1].to(self.device_backend.device) + def remember(text): + self.refined_text_memo[memo] = text.detach().cpu() + while len(self.refined_text_memo) > self.refined_text_memo_capacity: + del self.refined_text_memo[next(iter(self.refined_text_memo))] + + def release(): + if release_weights: + if not self.stream_refiner_count: + refiner.to('meta') + self.refiner_released = True + + if memo in self.refined_text_memo: + cpu_text = self.refined_text_memo.pop(memo) + self.refined_text_memo[memo] = cpu_text + text = cpu_text.to(self.device_backend.device) self.transformer._freevideo_refined_text = text + release() return dict(seconds=time.perf_counter() - started, shape=list(text.shape), cache_hit=False, previous_pass=True, refiner_weights_released=self.refiner_released) cache, note = None, None @@ -879,13 +945,11 @@ def _prepare_refined_text(self, prompt): expected_shape = [1, prompt.shape[-2], self.transformer.context_embedder.out_features] if (isinstance(text, torch.Tensor) and list(text.shape) == expected_shape and str(text.dtype) == entry['metrics'].get('dtype') and bool(torch.isfinite(text).all())): - if not self.stream_refiner_count: - refiner.to('meta') - self.refiner_released = True + release() self.transformer._freevideo_refined_text = text - self.previous_refined_text = (memo, text.detach().cpu()) + remember(text) return dict(seconds=time.perf_counter() - started, shape=list(text.shape), cache_hit=True, - refiner_weights_released=True, cache_receipt=str(cache.receipt(cache_key))) + refiner_weights_released=self.refiner_released, cache_receipt=str(cache.receipt(cache_key))) del text except (OSError, ValueError, KeyError, RuntimeError) as error: note = str(error) @@ -898,11 +962,10 @@ def _prepare_refined_text(self, prompt): _load_safetensors(self.cache / self.online_lora_root).items()}) refiner.load_state_dict(values, strict=True, assign=True) del values + self.refiner_released = False # load_prompt returns [tokens, width]; generate_latents adds the batch # dimension before the original transformer call. prompt = prompt.unsqueeze(0) if prompt.ndim == 2 else prompt - if prompt.ndim != 3 or prompt.shape[0] != 1: - raise ValueError('Expected one text prompt for refinement') text = self.transformer.context_embedder(prompt.to(input_dtype(self.transformer.context_embedder))) with self.device_backend.make_offloader(refiner, prefetch=False, weight_source=self.stream_weight_source if self.stream_refiner_count else None) as loader: @@ -910,27 +973,26 @@ def _prepare_refined_text(self, prompt): stats = loader.stats() del loader self.device_backend.synchronize() - if not self.stream_refiner_count: - refiner.to('meta') - self.refiner_released = True + release() self.transformer._freevideo_refined_text = text - self.previous_refined_text = (memo, text.detach().cpu()) + remember(text) saved = False if cache is not None: try: import uuid cache.root.mkdir(parents=True, exist_ok=True) temporary = cache.root / ('refined-' + uuid.uuid4().hex + '.partial') - torch.save(self.previous_refined_text[1], temporary) + torch.save(self.refined_text_memo[memo], temporary) saved = cache.put(cache_key, temporary, dict(success=True, shape=list(text.shape), dtype=str(text.dtype)), move=True) note = cache.note or note except (OSError, ValueError) as error: note = str(error) - gc.collect() - self.device_backend.empty_cache() - empty_host_cache(torch) + if release_weights: + gc.collect() + self.device_backend.empty_cache() + empty_host_cache(torch) return {'seconds': time.perf_counter() - started, 'shape': list(text.shape), - 'refiner_weights_released': True, 'offload': stats, 'cache_hit': False, 'cache_saved': saved, 'cache_note': note} + 'refiner_weights_released': self.refiner_released, 'offload': stats, 'cache_hit': False, 'cache_saved': saved, 'cache_note': note} def device_tensors(self): """The registered tensors this engine keeps on the GPU, or None. @@ -1000,7 +1062,7 @@ def close(self): if self.schedule_hook is not None: self.schedule_hook.remove() self.transformer = None - self.previous_refined_text = None + self.refined_text_memo.clear() self.offload_layers = [] if self.stream_weight_source is not None: self.stream_weight_source.release_host_views() diff --git a/freevideo_engine/runtime_alias.py b/freevideo_engine/runtime_alias.py index 68b41422..37c4eb3f 100644 --- a/freevideo_engine/runtime_alias.py +++ b/freevideo_engine/runtime_alias.py @@ -517,7 +517,8 @@ def _longest(root, suffixes=()): return _manifest_longest(str(path), info.st_mtime_ns, info.st_size, suffixes) -def needs_view(root, longest=None) -> bool: +def needs_view(root, longest=None, spelling_only=False) -> bool: + """spelling_only: only a folder the system code page cannot spell, never one that is merely long.""" if not _backend.windows(): return False root = _ordinary(resolved(checked_real_path(root))) @@ -530,6 +531,8 @@ def needs_view(root, longest=None) -> bool: # length; the short folder is plain letters. if not _backend.ansi_path(root): return True + if spelling_only: + return False # An installation that runs from where it is keeps doing so; only one beyond its budget is # given the short folder. if longest is None and system_long_paths(): diff --git a/freevideo_engine/system.py b/freevideo_engine/system.py index 3f9b83b2..95a8d60a 100644 --- a/freevideo_engine/system.py +++ b/freevideo_engine/system.py @@ -126,6 +126,9 @@ def residual_host_headroom(canvas=None): rows = (sum(canvas.get(key, 0) for key in ('video_tokens', 'reference_video_tokens', 'reference_audio_tokens')) if canvas is not None else 72576) + if canvas is not None and canvas.get('early_reference') is not None: + from .two_pass import placement_tokens + rows = placement_tokens(canvas) return int(2**30 * max(1., rows / 72576)) diff --git a/freevideo_engine/tuning.py b/freevideo_engine/tuning.py index f4b0533a..fd6fd976 100644 --- a/freevideo_engine/tuning.py +++ b/freevideo_engine/tuning.py @@ -231,6 +231,9 @@ def cached_conditioning(state, prompt, destination): def compatible_profile(entry, base, canvas, tokens): + # Tuning entries have no per-step reference schedule or capacity evidence. + if canvas.get('early_reference') is not None: + return False if not isinstance(entry, dict) or not isinstance(entry.get('id'), str): return False reference_context = {key: canvas.get(key, 0) for key in ('reference_video_tokens', 'reference_audio_tokens')} diff --git a/freevideo_engine/two_pass.py b/freevideo_engine/two_pass.py index 47c29f79..e69ceba7 100644 --- a/freevideo_engine/two_pass.py +++ b/freevideo_engine/two_pass.py @@ -181,6 +181,20 @@ def steps(sampling_plan=None): def same_strategy(left, right): """Do not mix mixed-resolution timings with ordinary full-resolution steps.""" + def reference_schedule(value): + # Worker receipts add step traces and describe subsets at the smaller + # canvas. Neither changes the reference schedule's identity. + source = value.get('geometry') or value + schedule = source.get('early_reference', value.get('early_reference')) + if schedule is None: + return None + return dict(early_steps=schedule.get('early_steps'), conditions=[ + {key: row.get(key) for key in ('start_step', 'task', 'text_tokens', + 'reference_video_tokens', 'reference_audio_tokens')} + for row in schedule.get('conditions', [])]) + + if reference_schedule(left) != reference_schedule(right): + return False a, b = left.get('sampling_plan'), right.get('sampling_plan') try: if steps(a) != steps(b): @@ -191,6 +205,59 @@ def same_strategy(left, right): return a == b if (a or {}).get('enabled') or (b or {}).get('enabled') else True +def stage_tokens(canvas): + """Video/reference activation rows for an encoded early-reference schedule. + + Counts are the scalar metadata from media_conditioning.describe, never + conditioning tensors. Ordinary requests retain their historical envelope. + The loaded engine must cover both passes; pass-local caching cannot shrink + its resident prefix during the first steps. + """ + fields = ('reference_video_tokens', 'reference_audio_tokens') + + def references(row): + counts = [row.get(name, 0) for name in fields] + if any(type(n) is not int or n < 0 for n in counts): + raise ValueError('Reference token counts must be nonnegative integers') + return sum(counts) + + maximum = references(canvas) + full = canvas['video_tokens'] + maximum + schedule = canvas.get('early_reference') + if schedule is None: + return dict(first=full, refinement=full, engine=full) + conditions = schedule.get('conditions') if isinstance(schedule, dict) else None + if not isinstance(conditions, list) or not conditions or any(not isinstance(row, dict) for row in conditions): + raise ValueError('Early reference planning requires described condition counts') + starts = [row.get('start_step') for row in conditions] + if (any(type(n) is not int or n < 0 for n in starts) or starts[0] != 0 + or starts != sorted(set(starts))): + raise ValueError('Invalid early reference planning boundaries') + counts = [references(row) for row in conditions] + if any(max(row.get(name, 0) for row in conditions) > canvas.get(name, 0) for name in fields): + raise ValueError('Early reference counts exceed the request allowance') + sampling = canvas.get('sampling_plan') or {} + first = canvas['video_tokens'] + refinement = None + if sampling.get('enabled'): + steps(sampling) + expected = plan(canvas, task=canvas.get('task', 't2va'), + base_steps=sampling['base_steps'], refine_steps=sampling['refine_steps']) + if any(sampling.get(key) != expected[key] for key in ('first', 'second')): + raise ValueError('Early reference sampling geometry disagrees with the request') + from .geometry import geometry + first = geometry(**sampling['first'])['video_tokens'] + # Permanent image/video/audio references still participate in refinement. + refinement = canvas['video_tokens'] + counts[-1] + first += maximum + return dict(first=first, refinement=refinement, engine=max(first, refinement or 0)) + + +def placement_tokens(canvas): + """Shared activation envelope for fixed GPU and host placement.""" + return stage_tokens(canvas)['engine'] + + def checkpoint_path(): """Setup's copy; an older lazily downloaded copy under models/ stays in use.""" from .paths import installed_model_root, model_root diff --git a/freevideo_engine/windows_gpu_memory.py b/freevideo_engine/windows_gpu_memory.py index f029acec..8786df34 100644 --- a/freevideo_engine/windows_gpu_memory.py +++ b/freevideo_engine/windows_gpu_memory.py @@ -111,6 +111,9 @@ def nonlocal_readout_bytes(canvas, attention_width, text_rows=1024): canvas = canvas or {} video = canvas.get('video_tokens', 72576) rows = video + text_rows + canvas.get('reference_video_tokens', 0) + canvas.get('reference_audio_tokens', 0) + if canvas.get('early_reference') is not None: + from .two_pass import placement_tokens + rows = placement_tokens(canvas) + text_rows return int(1.25 * (rows + video) * attention_width * 2) diff --git a/freevideo_engine/worker.py b/freevideo_engine/worker.py index 5b25e9ff..e50e4310 100644 --- a/freevideo_engine/worker.py +++ b/freevideo_engine/worker.py @@ -194,7 +194,7 @@ def fit_decoder(): metrics.setdefault('decoder_residency', dict(planned_resident_blocks=planned)) metrics['decoder_residency'].update(resident_blocks=fit, allocator_limit_bytes=int(limit)) resume = request.get('resume_decode') is not None - if not resume and options.get('task', 't2va') != 't2va': + if not resume and (options.get('task', 't2va') != 't2va' or request.get('geometry', {}).get('early_reference')): from .media_conditioning import describe canvas = request.get('geometry', {}) value = torch.load(request['conditioning'], map_location='cpu', weights_only=True) diff --git a/pyproject.toml b/pyproject.toml index 0b0a3b92..78fdd3f2 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -32,4 +32,5 @@ include = ["freevideo_engine*"] [tool.setuptools.package-data] "freevideo_engine.prompt_vlm" = ["model.json"] -freevideo_engine = ["stage_estimate_params.json", "build-version.txt", "build-identity.json", "release_notes.json", "dependencies.json", "model_files.json", "model_shares.json", "prepared_models.json", "macos_reference_catalog.json", "superseded_tables.json", "bootstrap_versions.json", "curl_windows.json", "git_linux.json", "uv_downloads.json", "torch_downloads.json", "test_prompts.json", "windows_bootstrap.ps1", "assets/*.png", "assets/*.ico", "launcher/*.qml", "launcher/licenses/*.txt", "licenses/*.txt", "licenses/runtime/*.txt", "subtitle_pack.json"] +"freevideo_engine.image_gen" = ["image_models.json", "licenses/NOTICE.txt"] +freevideo_engine = ["stage_estimate_params.json", "build-version.txt", "build-identity.json", "release_notes.json", "dependencies.json", "model_files.json", "model_shares.json", "prepared_models.json", "macos_reference_catalog.json", "superseded_tables.json", "bootstrap_versions.json", "curl_windows.json", "git_linux.json", "gcc_linux.json", "uv_downloads.json", "torch_downloads.json", "test_prompts.json", "windows_bootstrap.ps1", "assets/*.png", "assets/*.ico", "launcher/*.qml", "launcher/licenses/*.txt", "licenses/*.txt", "licenses/runtime/*.txt", "subtitle_pack.json"] diff --git a/scripts/test_freevideo_entry.mjs b/scripts/test_freevideo_entry.mjs index d386fb3a..d790bb68 100644 --- a/scripts/test_freevideo_entry.mjs +++ b/scripts/test_freevideo_entry.mjs @@ -32,7 +32,8 @@ test('main browser entry registers both views and preserves node hooks', async t graph: {extra: {freevideo_studio: true}, _nodes: [], getNodeById() {}}, }; const api = Object.assign(events, {apiURL: value => value}); - const document = {head: new Element('head'), createElement: tag => new Element(tag)}; + const document = {head: new Element('head'), createElement: tag => new Element(tag), + createRange: () => ({createContextualFragment: () => new Element('span')})}; const dependencies = { app, api, outputDownloadURL, regenerateResult, upscaleResult, referenceInputs, createErrorPanel: () => ({element: new Element('error'), show() {}, clear() {}}), @@ -95,6 +96,7 @@ test('main browser entry registers both views and preserves node hooks', async t .replace("import { regenerateResult, upscaleResult } from './studio_queue.js';", 'const {regenerateResult,upscaleResult} = globalThis.__freevideoEntryTest;') .replace("import { openImageEditor, icon } from './image_editor.js';", 'const {openImageEditor,icon} = globalThis.__freevideoEntryTest;') .replace("import { installNavigation, refreshNavigation, preferredView } from './view_navigation.js';", 'const {installNavigation,refreshNavigation,preferredView} = globalThis.__freevideoEntryTest;') + .replace("import { imageGenerationSupported } from './image_generation.js';", 'const imageGenerationSupported = () => true;') .replace("import { agentInterfaceEnabled } from './agent_connect.js';", 'const agentInterfaceEnabled = async () => false;'); await import('data:text/javascript;base64,' + Buffer.from(source).toString('base64')); const extension = extensions.find(row => row.name === 'FreeVideo.UnifiedMedia'); diff --git a/web/autopilot_chat.js b/web/autopilot_chat.js index 533ac560..4184f844 100644 --- a/web/autopilot_chat.js +++ b/web/autopilot_chat.js @@ -22,7 +22,7 @@ const QUIET = {vlm: ['Checking', '检查中'], enhance: ['Polishing', '润色中 export const about = (seconds, t) => seconds >= 60 ? t(`About ${Math.ceil(seconds / 60)} min`, `约 ${Math.ceil(seconds / 60)} 分钟`) : t('Under 1 min', '不到 1 分钟'); export function clock(seconds) { const s = Math.max(0, Math.round(seconds)); return `${Math.floor(s / 60)}:${String(s % 60).padStart(2, '0')}`; } -export const jobName = (job, t) => job.label || t('Video', '视频'); +export const jobName = (job, t) => job.label || (job.kind === 'image' ? t('Picture', '图片') : t('Video', '视频')); export const jobFraction = job => Math.max(0, Math.min(1, Number(job.progress?.fraction) || 0)); // What a job's last progress stands for `age` seconds after it was made, by the creation panel's rule. export function jobEstimate(job, age) { @@ -219,7 +219,10 @@ export function agentChat(t, api, {open, edit, shown = () => true, changed = () try { const response = await agentGet(api, '/freevideo/agent/chat/list', {signal: closed.signal}); if (stopped) return; - conversations = mergeConversations(conversations, response.ok ? await response.json() : null, response.status); + // The body is awaited first: a conversation chosen meanwhile must stay the current one. + const data = response.ok ? await response.json() : null; + if (stopped) return; + conversations = mergeConversations(conversations, data, response.status); renderConversations(); if (conversations.supported === false && selectedSession !== undefined) { selectedSession = undefined; selection++; void poll(); } } catch { /* List failures do not interrupt reading or sending. */ } @@ -371,7 +374,9 @@ export function agentChat(t, api, {open, edit, shown = () => true, changed = () let mode = null, apply = null; return {element: item, update(job, {latest, retry}) { // Jobs without a picture (a check, a prompt) are one quiet line while they run, and nothing after. - const next = VIDEOS.includes(job.kind || 'generate') ? jobMode(job) + (jobMode(job) === 'done' ? `:${job.play}:${job.thumbnail}:${job.work}` : '') + // A picture takes seconds: the same quiet line while it is made, then the picture. + const card = VIDEOS.includes(job.kind || 'generate') || (job.kind === 'image' && !WORKING.includes(job.status)); + const next = card ? jobMode(job) + (jobMode(job) === 'done' ? `:${job.play}:${job.thumbnail}:${job.work}` : '') : WORKING.includes(job.status) ? 'quiet' : 'none'; if (next !== mode) { mode = next; item.dataset.job = job.id; item.dataset.mode = next.split(':')[0]; item.hidden = next === 'none'; @@ -496,8 +501,11 @@ export function agentChat(t, api, {open, edit, shown = () => true, changed = () away = job ? {session: row.id, job, ids: jobs.map(job => job.id)} : null; } catch { /* The last reading stays until the next one. */ } } + let reportedJob = null; function report() { const current = (state.current ? state.jobs[state.current] : away?.job) || null; + // A job that starts or ends changes a dot in the list, so the list is read again right away. + if ((current?.id || null) !== reportedJob) { reportedJob = current?.id || null; void readConversations(true); } changed({known, agent: state.agent, current, fraction: current ? fractionOf(current) : 0, left: current ? leftOf(current) : null, active: state.active, any: state.items.size > 0 || pending.length > 0}); } diff --git a/web/freevideo.js b/web/freevideo.js index 7678efee..26bb1f43 100644 --- a/web/freevideo.js +++ b/web/freevideo.js @@ -16,6 +16,7 @@ import { reportFileName } from './report_issue.js'; import { rememberPromptDraft, savePromptDraft } from './prompt_draft.js'; import { regenerateResult, upscaleResult } from './studio_queue.js'; import { openImageEditor, icon } from './image_editor.js'; +import { imageGenerationSupported } from './image_generation.js'; const languageOverride = typeof location !== 'undefined' ? new URLSearchParams(location.search).get('freevideo_lang') : null; @@ -112,6 +113,9 @@ function mediaPanel(node, mount = null) { const picker = el("input"); picker.type = "file"; picker.multiple = true; picker.accept = "image/*,video/*,audio/*"; picker.hidden = true; const add = button(text("Upload files", "上传文件"), text("Upload images, video or audio", "上传图片、视频或音频"), () => picker.click()); + const generateImage = button(text('Generate with AI', 'AI 生成'), text('Generate with AI', 'AI 生成'), () => newImage()); + generateImage.prepend(document.createRange().createContextualFragment(icon('generate'))); + generateImage.hidden = !imageGenerationSupported(); const audioPicker = el('input'); audioPicker.type = 'file'; audioPicker.multiple = true; audioPicker.accept = 'audio/*,.wav,.mp3,.flac,.ogg,.m4a,.aac,.opus'; audioPicker.hidden = true; const addAudio = button(text('Reference audio', '参考音频'), @@ -122,7 +126,7 @@ function mediaPanel(node, mount = null) { const option = el("option", label); option.value = value; role.append(option); } role.value = 'reference'; - toolbar.append(add, addAudio, role, picker, audioPicker); panel.append(toolbar, mode, note, audioHelp, connections, cards); + toolbar.append(add, generateImage, addAudio, role, picker, audioPicker); panel.append(toolbar, mode, note, audioHelp, connections, cards); const read = () => { const rows = JSON.parse(serialized.value || "[]"); if (!Array.isArray(rows) || rows.length > 32) throw new Error(text("Use at most 32 media items", "最多使用 32 个素材")); @@ -189,12 +193,32 @@ function mediaPanel(node, mount = null) { const result = await uploadOne(new File([outcome.blob], outcome.name, {type: outcome.blob.type})); const latest = read(), index = latest.findIndex(row => row.file === file); if (index < 0) return; - latest[index] = {...latest[index], file: result.file, source, edit: outcome.edit}; update(latest); - message(text(`Saved the edited image, ${outcome.width}×${outcome.height}. The original is kept.`, - `已保存编辑后的图片（${outcome.width}×${outcome.height}），原图保留。`)); + const hasOriginal = outcome.edit.layers.some(layer => layer.id === 'base'); + latest[index] = {...latest[index], file: result.file, ...(hasOriginal ? {source} : {}), edit: outcome.edit}; update(latest); + message(hasOriginal ? text(`Saved the edited image, ${outcome.width}×${outcome.height}. The original is kept.`, + `已保存编辑后的图片（${outcome.width}×${outcome.height}），原图保留。`) : ''); } catch (error) { message(error.message, true); } finally { node.isUploading = false; add.disabled = false; addAudio.disabled = false; } } + async function newImage() { + if (node.isUploading || read().length >= 32) return; + const library = read().filter(row => mediaKind(row.file) === 'image') + .map(row => ({file: row.file, label: (row.source || row.file).split('/').pop()})); + let outcome; + try { + outcome = await openImageEditor({target: videoSize(), role: role.value, t: text, + resolve: viewURL, upload: uploadOne, library}); + } catch { message(text('The image could not be opened for editing.', '这张图片无法打开编辑。'), true); return; } + if (!outcome?.blob || read().length >= 32) return; + node.isUploading = true; add.disabled = addAudio.disabled = generateImage.disabled = true; + try { + const result = await uploadOne(new File([outcome.blob], outcome.name, {type: outcome.blob.type})); + const latest = read(); + latest.push({file: result.file, role: role.value, enabled: true, edit: outcome.edit}); + update(latest); + } catch (error) { message(error.message, true); } + finally { node.isUploading = false; add.disabled = addAudio.disabled = generateImage.disabled = false; } + } function restoreImage(file) { const rows = read(), at = rows.findIndex(row => row.file === file); if (at < 0 || !rows[at].source) return; diff --git a/web/health.js b/web/health.js index 111fe4b4..3451f84e 100644 --- a/web/health.js +++ b/web/health.js @@ -15,7 +15,7 @@ const launchVisit = query.get('freevideo') === 'launch'; // Every script this folder ships; a test keeps the list in step with web/*.js. export const MODULES = ['agent_activity.js', 'agent_connect.js', 'agent_session.js', 'autopilot.js', 'autopilot_chat.js', 'autopilot_conversations.js', 'browser_extension.js', 'chain_support.js', 'capability_view.js', 'compatibility.js', 'error_panel.js', 'first_run_guide.js', 'fonts.js', 'freevideo.js', 'generation_progress.js', 'editor.js', 'editor_actions.js', 'editor_adjust.js', 'editor_adjust_picture.js', 'editor_picture_fields.js', 'editor_picture_session.js', 'editor_preview_picture.js', 'editor_entry.js', 'editor_export_dialog.js', 'editor_first_frame.js', 'editor_genbar.js', 'editor_geometry.js', 'editor_library.js', 'editor_menu.js', 'editor_model.js', 'editor_frame.js', 'editor_player.js', 'editor_room.js', 'editor_select.js', 'editor_state.js', 'editor_store.js', 'editor_tracks.js', 'editor_generation.js', 'editor_transport.js', 'editor_transport_model.js', 'editor_audio.js', 'editor_audio_model.js', 'editor_audio_notes.js', 'editor_audio_replace.js', 'editor_audio_rules.js', 'editor_speed.js', 'editor_speed_model.js', - 'editor_loader.js', 'editor_text_coverage.js', 'editor_text.js', 'editor_text_model.js', 'editor_text_layout.js', 'editor_notes.js', 'editor_icons.js', 'dialog_keys.js', 'frame_sizes.js', 'health.js', 'image_commands.js', 'image_editor.js', 'image_editor_engine.js', 'launcher.js', 'length_plan.js', 'library.js', 'motion.js', 'output_download.js', 'port_notice.js', 'preview_scene.js', + 'editor_loader.js', 'editor_text_coverage.js', 'editor_text.js', 'editor_text_model.js', 'editor_text_layout.js', 'editor_notes.js', 'editor_icons.js', 'dialog_keys.js', 'frame_sizes.js', 'health.js', 'image_commands.js', 'image_editor.js', 'image_editor_engine.js', 'image_generation.js', 'launcher.js', 'length_plan.js', 'library.js', 'motion.js', 'output_download.js', 'port_notice.js', 'preview_scene.js', 'progress_connection.js', 'prompt_draft.js', 'prompt_enhance.js', 'prompt_references.js', 'report_issue.js', 'studio_recent.js', 'result_actions.js', 'sampling_effort.js', 'setup.js', 'share.js', 'share_layout.js', 'studio.js', 'studio_chain.js', 'studio_extend.js', 'studio_length.js', 'first_frame_row.js', 'studio_queue.js', 'studio_size.js', 'toolbar_position.js', 'updates.js', 'view_navigation.js', 'wordmark.js']; diff --git a/web/image_commands.js b/web/image_commands.js index 639ccad0..7aceb516 100644 --- a/web/image_commands.js +++ b/web/image_commands.js @@ -39,7 +39,7 @@ function travel({history, future, doc = history[history.length - 1], selection = const nextHistory = direction < 0 ? history.slice(0, -1) : [...history, future[future.length - 1]]; const nextFuture = direction < 0 ? [...future, history[history.length - 1]] : future.slice(0, -1); let next = nextHistory[nextHistory.length - 1]; - if (!next.layers.some(l => l.id === next.active)) next = {...next, active: next.layers[next.layers.length - 1].id}; + if (!next.layers.some(l => l.id === next.active)) next = {...next, active: next.layers.at(-1)?.id || null}; const sx = next.width / doc.width, sy = next.height / doc.height; return {history: nextHistory, future: nextFuture, doc: next, selection: sx === 1 && sy === 1 ? selection : scaleSelection(selection, sx, sy), scale: {sx, sy}}; @@ -49,7 +49,7 @@ export function redoStep(state) { return travel(state, 1); } export function insertLayer(doc, layer) { const index = doc.layers.findIndex(l => l.id === doc.active); - if (index < 0) fail('not_found', 'Active layer not found'); + if (index < 0 && doc.layers.length) fail('not_found', 'Active layer not found'); if (doc.layers.some(l => l.id === layer.id)) fail('unsupported', 'Duplicate layer id'); const layers = [...doc.layers]; layers.splice(index + 1, 0, layer); return {...doc, layers, active: layer.id}; @@ -86,7 +86,8 @@ export function deleteLayer(doc, id) { export function moveLayer(doc, id, delta) { layerOf(doc, id); const index = doc.layers.findIndex(l => l.id === id), to = index + delta; - if (!Number.isInteger(delta) || index < 1 || to < 1 || to >= doc.layers.length) fail('range', 'Layer move out of range'); + const first = doc.layers[0]?.id === 'base' ? 1 : 0; + if (!Number.isInteger(delta) || index < first || to < first || to >= doc.layers.length) fail('range', 'Layer move out of range'); const layers = [...doc.layers]; [layers[index], layers[to]] = [layers[to], layers[index]]; return {...doc, layers}; } diff --git a/web/image_editor.css b/web/image_editor.css index e03b6a27..86d84be3 100644 --- a/web/image_editor.css +++ b/web/image_editor.css @@ -61,6 +61,33 @@ button.fv-ie-size:hover{background:var(--fv-hover);color:var(--fv-ink)} .fv-ie-opt-title{font-weight:var(--fv-semibold);padding-right:4px;white-space:nowrap} .fv-ie-opt-label{color:var(--fv-muted);font-size:var(--fv-micro);white-space:nowrap} .fv-ie-opt-note{color:var(--fv-muted);font-size:var(--fv-micro);line-height:1.4} +/* The prompt expands over the picture, leaving the stage's size unchanged. */ +.fv-ie-options.fv-ie-generation{position:relative;z-index:2;display:grid;grid-template-columns:auto minmax(60px,1fr) auto auto;overflow:visible;align-items:start} +.fv-ie-generation>.fv-ie-seg{flex-wrap:nowrap} +.fv-ie-gen-field{display:flex;gap:8px;min-width:0;height:32px} +.fv-ie-gen-box{position:relative;flex:1;min-width:0;height:32px} +.fv-ie-gen-refs{display:flex;flex:none;align-items:center;gap:6px;height:32px} +.fv-ie-gen-refs[hidden]{display:none} +.fv-ie-gen-ref{position:relative;flex:none;width:28px;height:28px} +.fv-ie-gen-ref>canvas,.fv-ie-gen-ref>img{display:block;width:28px;height:28px;border-radius:6px;object-fit:cover;background:var(--fv-raised)} +.fv-ie-gen-ref>b{position:absolute;top:-4px;left:-4px;min-width:14px;height:14px;padding:0 3px;border-radius:7px;background:var(--fv-ink);color:var(--fv-card);font-size:10px;font-weight:var(--fv-semibold);line-height:14px;text-align:center;font-variant-numeric:tabular-nums} +.fv-ie-gen-ref-remove{position:absolute;top:-5px;right:-5px;width:16px;height:16px;padding:0;border:0;border-radius:8px;background:var(--fv-ink);color:var(--fv-card);font-size:12px;line-height:16px;opacity:0;cursor:pointer} +.fv-ie-gen-ref:hover>.fv-ie-gen-ref-remove:not(:disabled),.fv-ie-gen-ref-remove:focus-visible{opacity:1} +.fv-ie-gen-ref-add summary{width:28px;height:28px;border:1px dashed var(--fv-border);border-radius:6px} +.fv-ie-gen-ref-add summary[aria-disabled=true]{opacity:.4;cursor:default;pointer-events:none} +.fv-ie-gen-ref-add .fv-ie-menu-list{top:calc(100% + 6px);right:auto;bottom:auto;left:0} +.fv-ie-options.fv-ie-generation[data-empty=true]{grid-template-columns:auto auto minmax(60px,1fr) auto auto} +.fv-ie-generation[data-empty=true]>.fv-ie-button{grid-column:5} +.fv-ie-generation>#fv-ie-gen-ratio{min-height:32px} +.fv-image-editor .fv-ie-gen-prompt{position:absolute;inset:0 0 auto;width:100%;height:32px;min-height:32px;max-height:80px;resize:none;padding:3px 9px;border:1px solid var(--fv-border);border-radius:var(--fv-r-sm);background:var(--fv-raised);line-height:24px;overflow-y:auto} +.fv-ie-gen-prompt:disabled{opacity:.5} +.fv-ie-gen-note{align-self:center;min-width:0;overflow-wrap:anywhere;font-variant-numeric:tabular-nums} +.fv-ie-generation[data-running=true]>.fv-ie-gen-note{white-space:nowrap} +.fv-ie-generation>.fv-ie-button{grid-column:4;min-height:32px} +.fv-ie-generation[data-message=true]>.fv-ie-gen-note{max-width:260px} +/* Same accent ramp and easing as the shared generation progress meters. */ +.fv-ie-gen-meter{position:absolute;bottom:0;left:0;right:0;height:2px;overflow:hidden;background:var(--fv-raised)} +.fv-ie-gen-meter i{display:block;height:100%;background:linear-gradient(90deg,var(--fv-accent),var(--fv-accent-ramp));transition:width .5s var(--fv-ease-out)} .fv-ie-sep{width:1px;height:22px;background:var(--fv-border);flex:none} .fv-ie-check{display:inline-flex;align-items:center;gap:6px;color:var(--fv-ink);font-size:var(--fv-micro);cursor:pointer;white-space:nowrap} .fv-ie-check input{accent-color:var(--fv-accent);margin:0} @@ -223,11 +250,30 @@ button.fv-ie-size:hover{background:var(--fv-hover);color:var(--fv-ink)} .fv-ie-file,.fv-ie-size{display:none} } @media (max-width:560px){ + .fv-ie-options.fv-ie-generation{grid-template-columns:minmax(0,1fr) auto;gap:8px;min-height:88px} + .fv-ie-generation>.fv-ie-seg{grid-column:1/-1;justify-self:start} + .fv-ie-gen-field{grid-column:1;grid-row:2;flex-wrap:wrap;height:auto} + .fv-ie-gen-refs{flex:0 0 100%;overflow-x:auto;height:36px;padding:4px 5px 0} + .fv-ie-gen-box{flex:1 0 100%} + .fv-image-editor .fv-ie-gen-prompt{font-size:16px} + .fv-ie-generation>.fv-ie-button{grid-column:2;grid-row:2;align-self:end} + .fv-ie-gen-note{grid-column:1/-1;grid-row:3} + .fv-ie-generation[data-running=true]{grid-template-columns:minmax(0,1fr) auto auto} + .fv-ie-generation[data-running=true]>.fv-ie-seg{grid-column:1/-1} + .fv-ie-generation[data-running=true]>.fv-ie-button{grid-column:3} + .fv-ie-generation[data-running=true]>.fv-ie-gen-note{grid-column:2;grid-row:2} + .fv-ie-generation[data-message=true]>.fv-ie-gen-note{max-width:none} + .fv-ie-options.fv-ie-generation[data-empty=true]{grid-template-columns:minmax(0,1fr) auto} + .fv-ie-generation[data-empty=true]>.fv-ie-seg{grid-column:1;grid-row:1} + .fv-ie-generation[data-empty=true]>#fv-ie-gen-ratio{grid-column:2;grid-row:1} + .fv-ie-generation[data-empty=true]>.fv-ie-button{grid-column:2;grid-row:2} + .fv-ie-generation[data-empty=true]>.fv-ie-gen-note{grid-column:1/-1;grid-row:3} .fv-ie-actions .fv-ie-button:not(.fv-ie-primary):not(.fv-ie-danger){padding:5px 9px} .fv-ie-popover{left:12px;width:auto} .fv-ie-reset{display:none} } @media (prefers-reduced-motion:reduce){.fv-image-editor[open],.fv-image-editor::backdrop{animation:none!important}} +@media (prefers-reduced-motion:reduce){.fv-ie-gen-meter i{transition:none}} /* Media cards: the picture is the entry to the editor. */ .fv-panel .fv-preview-wrap{position:relative;margin:8px 0;border-radius:var(--fv-r-sm);overflow:hidden;background:var(--fv-bg)} diff --git a/web/image_editor.js b/web/image_editor.js index b5d7ab1a..0881ec08 100644 --- a/web/image_editor.js +++ b/web/image_editor.js @@ -3,10 +3,12 @@ // FreeVideo's own look. The original file is never changed: saving renders a // new image and keeps the document, so the next visit starts from the original // with every layer and step in place. Closing frees every canvas it made. -import {BLENDS, emptyAdjust, migrate, orientedSize, orientMatrix, canvasRect, layerMatrix, scaleSelection, +import {BLENDS, emptyAdjust, migrate, blankDocument, orientedSize, orientMatrix, canvasRect, layerMatrix, scaleSelection, rasterKinds, measureText, Renderer} from './image_editor_engine.js'; import * as commands from './image_commands.js'; import {FONT_KEYS, fontStack, fontLoaded, loadFont, loadFonts} from './fonts.js'; +import {api} from '../../scripts/api.js'; +import {createImageGenClient, createImageGeneration, imageGenerationSupported, blankCanvasSize, blankCanvasRatio} from './image_generation.js'; const css = document.createElement('link'); css.rel = 'stylesheet'; css.href = new URL('./image_editor.css', import.meta.url).href; @@ -18,6 +20,7 @@ const clamp = (v, lo, hi) => Math.min(hi, Math.max(lo, v)); const SVG = (body, size = 18) => ``; // Line icons on a 24-unit grid; eyedropper, hand, crop, link and unlink follow Lucide (ISC, see THIRD_PARTY_NOTICES.md). export const ICONS = { + generate: '', move: '', crop: '', marquee: '', @@ -123,15 +126,21 @@ const strip = d => JSON.stringify({...d, active: null, layers: d.layers.map(l => * source: original image file; resolve(file) → URL; upload(file) → {file}; * target: {width, height} of the video, if known; library: [{file, label}] other media images. */ -export async function openImageEditor({source, name, target = null, edit = null, role = 'reference', t, resolve, - upload = null, library = [], tool: startTool = null}) { - const baseImage = await loadImage(resolve(source)); +export async function openImageEditor({source = null, name = source || '', target = null, edit = null, role = 'reference', t, resolve, + upload = null, library = [], tool: startTool = null, imageGenClient = null}) { + const hasOriginal = !!source && (!edit || edit.version !== 2 || edit.layers.some(l => l.id === 'base')); + const baseImage = hasOriginal ? await loadImage(resolve(source)) : null; const renderer = new Renderer(resolve); const fresh = !edit; - const original = () => { const d = migrate(null, source, baseImage); d.layers[0].name = t('Original', '原图'); return d; }; + const blankSizeForRatio = ratio => blankCanvasSize(ratio, null); + const blankSize = edit || (target?.width > 0 && target?.height > 0 ? target : blankSizeForRatio('16:9')); + const original = () => { + if (!hasOriginal) return blankDocument(blankSize, edit?.ratio || blankCanvasRatio(target)); + const d = migrate(null, source, baseImage); d.layers[0].name = t('Original', '原图'); return d; + }; let doc = edit ? migrate(edit, source, baseImage) : original(); doc = {...doc, layers: doc.layers.map(l => l.id === 'base' ? {...l, name: t('Original', '原图')} : l)}; - if (fresh && target && role !== 'reference') { + if (fresh && hasOriginal && target && role !== 'reference') { // A keyframe becomes the first or last frame of the video: start from its shape. doc = commands.setRatio(doc, 'video', target); } @@ -164,7 +173,10 @@ export async function openImageEditor({source, name, target = null, edit = null, move: {auto: true}, }; const extension = /\.jpe?g$/i.test(name) ? 'jpg' : 'png'; - const stem = name.split('/').pop().replace(/\.[^.]+$/, '').replace(/-edit$/, ''); + // A picture started from nothing gets a readable name of its own, not "-edit". + const stamp = () => { const d = new Date(), two = n => String(n).padStart(2, '0'); + return `${d.getFullYear()}${two(d.getMonth() + 1)}${two(d.getDate())}-${two(d.getHours())}${two(d.getMinutes())}${two(d.getSeconds())}`; }; + const stem = name ? name.split('/').pop().replace(/\.[^.]+$/, '').replace(/-edit$/, '') : `ai-${stamp()}`; // ----- frame -------------------------------------------------------------- const dialog = el('dialog', null, 'fv-image-editor'); @@ -211,9 +223,13 @@ export async function openImageEditor({source, name, target = null, edit = null, const filePicker = el('input'); filePicker.type = 'file'; filePicker.accept = 'image/*'; filePicker.hidden = true; dialog.append(header, main, confirmBar, status, filePicker); document.body.append(dialog); + const generation = createImageGeneration({client: imageGenClient || createImageGenClient(api.fetchApi.bind(api)), + renderer, t, getDoc: () => doc, getSelection: () => selection, commit: commitInserted, el, button, segmented, + library, resolve, icon, onRatioChange: setBlankCanvasRatio}); // ----- tools ---------------------------------------------------------------- const TOOLS = [ + ['generate', t('Generate with AI', 'AI 生成'), null], null, ['move', t('Move', '移动'), 'V'], ['crop', t('Crop and canvas', '裁剪与画布'), 'C'], null, ['marquee', t('Marquee', '选框'), 'M'], ['lasso', t('Lasso', '套索'), 'L'], ['wand', t('Magic wand', '魔棒'), 'W'], null, ['brush', t('Brush', '画笔'), 'B'], ['eraser', t('Eraser', '橡皮擦'), 'E'], ['heal', t('Spot healing', '污点修复'), 'J'], @@ -224,7 +240,8 @@ export async function openImageEditor({source, name, target = null, edit = null, for (const entry of TOOLS) { if (!entry) { toolbar.append(el('span', null, 'fv-ie-tool-gap')); continue; } const [id, label, key] = entry; - const b = iconButton(id === 'marquee' ? 'marquee' : id, `${label} (${key})`, () => select(id), 'fv-ie-tool'); + if (id === 'generate' && !imageGenerationSupported()) continue; + const b = iconButton(id === 'marquee' ? 'marquee' : id, key ? `${label} (${key})` : label, () => select(id), 'fv-ie-tool'); b.dataset.tool = id; toolButtons[id] = b; toolbar.append(b); } const colors = el('div', null, 'fv-ie-colors'); @@ -288,7 +305,8 @@ export async function openImageEditor({source, name, target = null, edit = null, antsEdge = null; if (opts.clone.source) opts.clone.source = [opts.clone.source[0] * sx, opts.clone.source[1] * sy]; } - if (maskTarget && !active().mask) maskTarget = false; + if (maskTarget && !active()?.mask) maskTarget = false; + if (!doc.layers.length && tool !== 'hand') tool = 'generate'; refresh(true); } const changed = () => history.length > 1 && strip(doc) !== strip(initial); @@ -445,18 +463,26 @@ export async function openImageEditor({source, name, target = null, edit = null, // ----- tool options --------------------------------------------------------------- function controls() { undo.disabled = history.length < 2; redo.disabled = !future.length; + save.disabled = !doc.layers.length; const c = doc.crop; sizeLabel.textContent = `${Math.round(c.w)} × ${Math.round(c.h)}`; zoomLabel.textContent = Math.round(layout.k * 100) + '%'; - for (const [id, b] of Object.entries(toolButtons)) b.setAttribute('aria-pressed', String(id === tool)); + for (const [id, b] of Object.entries(toolButtons)) { + b.disabled = !doc.layers.length && !['generate', 'hand'].includes(id); + b.setAttribute('aria-pressed', String(id === tool)); + } stage.dataset.tool = tool; } function select(id) { + if (!doc.layers.length && !['generate', 'hand'].includes(id)) return; + if (id === 'generate' && !imageGenerationSupported()) return; commitText(); gesture?.cancel?.(); gesture = null; tool = id; renderOptions(); controls(); draw(); showHint(); keepFocus(); + if (id === 'generate') void generation.selected(); } function showHint() { const layer = active(); + if (!layer) { hint.textContent = ''; hint.hidden = true; return; } let text = ''; if (['eraser', 'heal', 'clone', 'blur'].includes(tool) && !rasterKinds.has(layer.kind)) text = t('Pick an image or paint layer to work on its pixels.', '请选择图片图层或绘画图层再处理像素。'); else if (tool === 'clone' && !opts.clone.source) text = t('Alt-click to set the source, then paint.', '先按住 Alt 点击设置取样点，再涂抹。'); @@ -469,7 +495,12 @@ export async function openImageEditor({source, name, target = null, edit = null, original: t('Original', '原图比例'), free: t('Free', '自由'), })[id] || id; function renderOptions() { + generation.unmount(); clearPanel(optionsBar); + optionsBar.classList.remove('fv-ie-generation'); + delete optionsBar.dataset.running; delete optionsBar.dataset.message; delete optionsBar.dataset.empty; + if (tool === 'generate') { generation.mount(optionsBar); return; } + if (!doc.layers.length) return; const add = (...nodes) => optionsBar.append(...nodes); const label = text => el('span', text, 'fv-ie-opt-label'); const sep = () => el('span', null, 'fv-ie-sep'); @@ -579,6 +610,7 @@ export async function openImageEditor({source, name, target = null, edit = null, function renderProps() { clearPanel(props); const layer = active(); + if (!layer) return; const head = el('div', null, 'fv-ie-section-head'); head.append(el('h3', layer.kind === 'text' ? t('Text', '文字') : layer.kind === 'shape' ? t('Shape', '形状') : t('Adjust', '调整')), el('span', layer.name || '', 'fv-ie-props-name')); @@ -698,7 +730,7 @@ export async function openImageEditor({source, name, target = null, edit = null, } function renderLayers() { clearPanel(layersBox); - const layer = active(), index = doc.layers.findIndex(l => l.id === layer.id); + const layer = active(), index = doc.layers.findIndex(l => l.id === layer?.id); const head = el('div', null, 'fv-ie-section-head'); head.append(el('h3', t('Layers', '图层'))); const tools = el('div', null, 'fv-ie-layer-tools'); @@ -715,11 +747,14 @@ export async function openImageEditor({source, name, target = null, edit = null, addMenu.append(summary, menu); const del = iconButton('trash', t('Delete layer', '删除图层'), () => deleteLayer()); const up = iconButton('up', t('Move up', '上移'), () => moveLayer(1)), down = iconButton('down', t('Move down', '下移'), () => moveLayer(-1)); - const maskButton = iconButton('mask', layer.mask ? t('Delete mask', '删除蒙版') : t('Add mask', '添加蒙版'), () => toggleMask(), 'fv-ie-icon' + (layer.mask ? ' fv-ie-active' : '')); - del.disabled = !!layer.locked; up.disabled = index < 1 || index >= doc.layers.length - 1; down.disabled = index <= 1; - maskButton.disabled = !rasterKinds.has(layer.kind); + const maskButton = iconButton('mask', layer?.mask ? t('Delete mask', '删除蒙版') : t('Add mask', '添加蒙版'), () => toggleMask(), 'fv-ie-icon' + (layer?.mask ? ' fv-ie-active' : '')); + const first = doc.layers[0]?.id === 'base' ? 1 : 0; + del.disabled = !layer || layer.locked || doc.layers.length === 1; + up.disabled = index < first || index >= doc.layers.length - 1; down.disabled = index <= first; + maskButton.disabled = !rasterKinds.has(layer?.kind); + const copy = iconButton('copy', t('Duplicate layer', '复制图层'), () => duplicateLayer()); copy.disabled = !layer; tools.append(iconButton('plus', t('New paint layer', '新建绘画图层'), () => addPaintLayer()), addMenu, maskButton, - iconButton('copy', t('Duplicate layer', '复制图层'), () => duplicateLayer()), up, down, del); + copy, up, down, del); head.append(tools); const list = el('ol', null, 'fv-ie-layer-list'); list.setAttribute('aria-label', t('Layers', '图层')); for (const l of [...doc.layers].reverse()) { @@ -800,20 +835,23 @@ export async function openImageEditor({source, name, target = null, edit = null, }; function duplicateLayer() { const layer = active(); + if (!layer) return; commitInserted(commands.duplicateLayer(doc, layer.id, t(`${layer.name} copy`, `${layer.name} 副本`))); } function deleteLayer() { - const layer = active(); if (layer.locked) return; + const layer = active(); if (!layer || layer.locked || doc.layers.length === 1) return; maskTarget = false; commit(commands.deleteLayer(doc, layer.id)); } function moveLayer(delta) { const index = doc.layers.findIndex(l => l.id === doc.active), to = index + delta; - if (index < 1 || to < 1 || to >= doc.layers.length) return; + const first = doc.layers[0]?.id === 'base' ? 1 : 0; + if (index < first || to < first || to >= doc.layers.length) return; commit(commands.moveLayer(doc, doc.active, delta)); } function toggleMask() { const layer = active(); + if (!layer) return; if (!layer.mask && !rasterKinds.has(layer.kind)) return; const next = commands.toggleMask(doc, layer.id, selection); maskTarget = next.maskTarget; @@ -822,6 +860,11 @@ export async function openImageEditor({source, name, target = null, edit = null, } // ----- canvas: shape, orientation, size ---------------------------------------------- + function setBlankCanvasRatio(ratio) { + if (doc.layers.length || generation.state().running) return; + const next = blankDocument(blankSizeForRatio(ratio), ratio); + followSize(doc, next); deselectQuiet(); commit(next); + } function aspectOf(id, r = rect()) { if (id === 'video') return target ? target.width / target.height : null; if (id === 'original') return r.w / r.h; @@ -969,6 +1012,7 @@ export async function openImageEditor({source, name, target = null, edit = null, } function pixelLayer() { const layer = active(); + if (!layer) return null; if (!rasterKinds.has(layer.kind)) { message(t('Pick an image or paint layer first.', '请先选择图片图层或绘画图层。')); return null; } return layer; } @@ -1055,6 +1099,7 @@ export async function openImageEditor({source, name, target = null, edit = null, return null; } function startGesture(event, p) { + if (!doc.layers.length || tool === 'generate') return null; const o = toO(...p), d = toD(o), layer = active(); const mode = event.shiftKey ? 'add' : event.altKey ? 'subtract' : null; if (tool === 'crop') return cropGesture(event, o); @@ -1409,12 +1454,15 @@ export async function openImageEditor({source, name, target = null, edit = null, let messageTimer = null; function message(text) { status.textContent = text; clearTimeout(messageTimer); if (text) messageTimer = setTimeout(() => { status.textContent = ''; }, 5000); } async function leave() { + void generation.cancel(); if (healing) await healing; commitText(); if (changed()) { confirmBar.hidden = false; confirmBar.querySelector('.fv-ie-danger').focus(); return; } close(null); } async function finish() { + if (!doc.layers.length) return; + void generation.cancel(); if (healing) await healing; commitText(); if (strip(doc) === strip(original())) { close(fresh ? null : {restore: true}); return; } @@ -1422,13 +1470,14 @@ export async function openImageEditor({source, name, target = null, edit = null, save.disabled = cancel.disabled = true; message(t('Saving the edited image…', '正在保存编辑后的图片…')); try { const {blob, width, height} = await renderer.export(doc, extension === 'jpg' ? 'image/jpeg' : 'image/png'); - close({blob, name: `${stem}-edit.${extension}`, edit: doc, width, height}); + close({blob, name: name ? `${stem}-edit.${extension}` : `${stem}.${extension}`, edit: doc, width, height}); } catch { message(t('The image could not be rendered. Try a smaller canvas or reload the page.', '图片生成失败，请缩小画布或刷新页面后重试。')); save.disabled = cancel.disabled = false; } } function close(value) { + void generation.dispose(); clearInterval(antsTimer); clearTimeout(messageTimer); observer.disconnect(); dialog.removeEventListener('keydown', keys); dialog.removeEventListener('keyup', keyup); @@ -1444,7 +1493,8 @@ export async function openImageEditor({source, name, target = null, edit = null, // A read-only view of the editor for automation and tests. dialog.freevideoEditor = { - state: () => ({doc, selection, tool, fg, bg, zoom: layout.k, maskTarget, history: history.length, future: future.length, healing: !!healing}), + state: () => ({doc, selection, tool, fg, bg, zoom: layout.k, maskTarget, history: history.length, future: future.length, healing: !!healing, + generation: {...generation.state(), info: generation.state().info ? {installed: generation.state().info.installed} : null, job: undefined}}), pictureBox: () => { const b = stage.getBoundingClientRect(), r = rect(), [x, y] = fromO(r.x, r.y); return {x: b.left + x, y: b.top + y, width: r.w * layout.k, height: r.h * layout.k}; }, }; @@ -1454,8 +1504,8 @@ export async function openImageEditor({source, name, target = null, edit = null, dialog.showModal(); computeLayout(); renderPanels(); - select(startTool || (fresh && target && role !== 'reference' ? 'crop' : 'move')); - save.focus(); + select(!doc.layers.length ? 'generate' : startTool || (fresh && target && role !== 'reference' ? 'crop' : 'move')); + (save.disabled ? stage : save).focus(); // The other fonts load in the background, so switching to one is immediate. loadFonts(FONT_KEYS.flatMap(k => [[k, 400], [k, 700]])).catch(() => {}); return result; diff --git a/web/image_editor_engine.js b/web/image_editor_engine.js index ab233089..4830d152 100644 --- a/web/image_editor_engine.js +++ b/web/image_editor_engine.js @@ -44,6 +44,11 @@ export function newDocument(file, image) { extend: {l: 0, t: 0, r: 0, b: 0, color: '#000000'}, crop: {x: 0, y: 0, w, h}, ratio: 'original', layers: [baseLayer(file, image)], active: 'base'}; } +export function blankDocument({width, height}, ratio = 'free') { + return {version: 2, width, height, rotate: 0, flipX: false, flipY: false, + extend: {l: 0, t: 0, r: 0, b: 0, color: '#000000'}, + crop: {x: 0, y: 0, w: width, h: height}, ratio, layers: [], active: null}; +} // Edits saved before layers: crop and strokes in fractions of the oriented picture. export function migrate(edit, file, image) { if (!edit) return newDocument(file, image); diff --git a/web/image_generation.js b/web/image_generation.js new file mode 100644 index 00000000..5144aedd --- /dev/null +++ b/web/image_generation.js @@ -0,0 +1,422 @@ +// Image generation's wire adapter, cancellable session and options UI. The editor +// owns its document/history; this module inserts a result through its commands. +import * as commands from './image_commands.js'; +import {orientMatrix, rasterKinds, scaleSelection} from './image_editor_engine.js'; + +export const IMAGE_GEN_PROTOCOL = Object.freeze({ + root: '/freevideo/image-gen', header: 'X-FreeVideo-Image', upload: '/upload/image', + installed: 'ready', phase: 'phase', bytes: 'bytes', + token: 'token', job: 'job', prompt: 'prompt', width: 'width', height: 'height', images: 'images', seed: 'seed', + stage: 'stage', fraction: 'fraction', file: 'file', error: 'error', name: 'name', subfolder: 'subfolder', + states: {starting: 'queued', progress: 'running', complete: 'done', failed: 'error', + cancelling: 'running', cancelled: 'error'}, +}); +export const IMAGE_GEN_POLL_MS = 500; +export const IMAGE_GEN_INFO_POLL_MS = 10000; +export const BLANK_CANVAS_RATIOS = Object.freeze(['16:9', '9:16', '1:1', '4:3', '3:4']); +export function imageGenerationSupported(platform = globalThis.navigator?.userAgentData?.platform || globalThis.navigator?.platform || '') { + return !/Mac|iPhone|iPad/i.test(platform); +} +export function blankCanvasRatio(target) { + return BLANK_CANVAS_RATIOS.find(ratio => { + const [w, h] = ratio.split(':').map(Number); + return target?.width > 0 && target?.height > 0 && target.width * h === target.height * w; + }) || BLANK_CANVAS_RATIOS[0]; +} +export function blankCanvasSize(ratio, sizes = null) { + if (!BLANK_CANVAS_RATIOS.includes(ratio)) ratio = BLANK_CANVAS_RATIOS[0]; + if (sizes?.[ratio]) { + const [width, height] = sizes[ratio]; + return {width, height}; + } + const [w, h] = ratio.split(':').map(Number); + const units = Math.floor(Math.sqrt(1024 ** 2 / (w * h * 32 ** 2))); + return {width: w * units * 32, height: h * units * 32}; +} +const failure = code => Object.assign(new Error(code || 'generation_failed'), {code: code || 'generation_failed'}); + +export function createImageGenClient(fetchImpl = globalThis.fetch.bind(globalThis)) { + let token = ''; + const p = IMAGE_GEN_PROTOCOL; + async function request(path, options) { + const response = await fetchImpl(path, options); + const value = await response.json(); + if (!response.ok || value[p.error] && !value[p.phase] && !value.state) throw failure(value[p.error]); + return value; + } + async function post(action, value, retried = false) { + if (!token) await info(); + try { + return await request(`${p.root}/${action}`, {method: 'POST', headers: { + 'Content-Type': 'application/json', [p.header]: token}, body: JSON.stringify(value)}); + } catch (error) { + // The server restarted behind an open editor: take its new token once. + if (error.code !== 'reload' || retried) throw error; + token = ''; + return post(action, value, true); + } + } + async function info() { + const value = await request(p.root, {cache: 'no-store'}); + token = value[p.token]; + return {...value, installed: value.installed ?? value[p.installed] ?? false}; + } + return { + info, + async generate({prompt, width, height, images = [], seed}) { + const value = await post('generate', {[p.prompt]: prompt, [p.width]: width, [p.height]: height, + [p.images]: images, ...(seed === undefined ? {} : {[p.seed]: seed})}); + return {...value, job: value[p.job]}; + }, + async status(job) { + const value = await post('status', {[p.job]: job}); + const phase = value[p.phase]; + return {...value, state: value.state || p.states[phase] || 'error', + stage: value[p.stage] || (phase === 'starting' ? 'load' : ''), + fraction: value[p.fraction], file: value[p.file], error: value[p.error], + ...(phase === 'cancelled' ? {error: 'cancelled'} : {})}; + }, + cancel: job => post('cancel', {[p.job]: job}), + async upload(file) { + // ComfyUI's existing image upload stores this at input/. + // FreeVideo's media upload uses subdirectories the image worker rejects. + const body = new FormData(); body.append('image', file); + body.append('type', 'input'); body.append('overwrite', 'false'); + const value = await request(p.upload, {method: 'POST', body}); + if (!value[p.name] || value[p.subfolder] || /[/\\:]/.test(value[p.name])) throw failure('invalid_media'); + return value[p.name]; + }, + }; +} + +export function imagePackMessage(info, t) { + const name = t(info?.pack_name_en || 'Image generation models', info?.pack_name_zh || '图像生成模型'); + // Match the launcher's binary size display; its optional pack_size can supply + // the exact platform-specific label (macOS uses decimal sizes). + const bytes = info?.[IMAGE_GEN_PROTOCOL.bytes]; + // A non-breaking space keeps the number and its unit on one line. + const size = (info?.pack_size || (bytes > 0 ? (bytes / 2 ** 30).toFixed(1) + ' GiB' : '17.3 GB')) + .replace(/ (?=[KMGT]i?B$)/, '\u00a0'); + return t(`Download "${name}" in the launcher first, ${size}.`, `先在启动器里下载“${name}”，${size}。`); +} +export function imageGenError(code, info, t) { + if (code === 'pack_missing') return imagePackMessage(info, t); + if (code === 'busy_video') return t('A video is generating. Try again when it finishes.', '视频正在生成，完成后再试。'); + if (code === 'no_memory') return t('Not enough GPU memory. Make the picture smaller in "Image size…" and try again.', + '显存不足。在“图像大小…”里缩小后再试。'); + return t('The image could not be generated. Try again.', '没有生成成功，请再试一次。'); +} + +// The service accepts 256..2048 on each side; a larger canvas is generated at its +// fitted size and the result is laid over the canvas like any other layer. +export function fitGenerationSize(width, height) { + const scale = Math.min(1, 2048 / Math.max(width, height)); + const side = value => Math.min(2048, Math.max(256, Math.round(value * scale))); + return {width: side(width), height: side(height)}; +} + +export function generationContext(doc, selection, mode) { + const layer = mode === 'edit' ? doc.layers.find(l => l.id === doc.active && rasterKinds.has(l.kind)) : null; + return {doc, selection, layer, mode, ...fitGenerationSize(doc.crop.w, doc.crop.h)}; +} +export async function uploadGenerationReference(renderer, client, context, alive = () => true) { + let blob; + if (context.layer) { + await renderer.preload(context.doc); + if (!alive()) throw failure('cancelled'); + const raster = renderer.layerImage(context.doc, context.layer, 1, true); + blob = await new Promise((resolve, reject) => raster.toBlob(b => b ? resolve(b) : reject(failure()), 'image/png')); + } else blob = (await renderer.export(context.doc, 'image/png')).blob; + if (!alive()) throw failure('cancelled'); + const file = new File([blob], `freevideo-image-reference-${crypto.randomUUID()}.png`, {type: 'image/png'}); + return client.upload(file); +} + +// "Edit this" can take more pictures beside the layer being edited: that layer is always image 1 and +// the others follow in the order they were added. The page offers fewer than the service accepts. +export const PAGE_REFERENCE_LIMIT = 3; +export function referenceLimit(info) { + const most = Number(info?.max_images); + return Number.isInteger(most) && most > 1 ? Math.min(PAGE_REFERENCE_LIMIT, most) : 1; +} +export function addReference(list, item, limit) { + if (!item?.file || list.length >= limit - 1 || list.some(r => r.file === item.file)) return list; + return [...list, {file: item.file, label: item.label || ''}]; +} +export const removeReference = (list, file) => list.filter(r => r.file !== file); +// Media lives in folders the image worker does not read, so each extra picture is copied beside the +// first reference once and its copy is used again for later requests. +export async function uploadExtraReferences(client, references, render, copies, alive = () => true) { + const names = []; + for (const item of references) { + if (!copies.has(item.file)) { + const blob = await render(item.file); + if (!alive()) throw failure('cancelled'); + const name = await client.upload(new File([blob], `freevideo-image-reference-${crypto.randomUUID()}.png`, {type: 'image/png'})); + if (!alive()) throw failure('cancelled'); + copies.set(item.file, name); + } + names.push(copies.get(item.file)); + } + return names; +} + +export function insertGenerationResult(doc, result, context, t) { + const sx = doc.width / context.doc.width, sy = doc.height / context.doc.height; + const reference = context.mode === 'edit' && context.layer; + const above = doc.layers.find(l => l.id === reference?.id) || doc.layers.at(-1); + let next = commands.addImageLayer({...doc, active: above?.id || null}, { + ...result, name: context.mode === 'new' ? t('Generated', '生成') : t('Edited', '改图')}); + let geometry; + if (reference) { + geometry = {cx: reference.cx * sx, cy: reference.cy * sy, w: reference.w * sx, h: reference.h * sy, + rotation: reference.rotation, flipX: reference.flipX, flipY: reference.flipY}; + } else { + const c = doc.crop, point = orientMatrix(doc).inverse().transformPoint(new DOMPoint(c.x + c.w / 2, c.y + c.h / 2)); + geometry = {cx: point.x, cy: point.y, w: c.w, h: c.h, rotation: -doc.rotate, flipX: doc.flipX, flipY: doc.flipY}; + } + next = {...next, layers: next.layers.map(l => l.id === next.active ? {...l, ...geometry} : l)}; + if (context.selection) next = commands.maskFromSelection(next, next.active, scaleSelection(context.selection, sx, sy)); + return next; +} + +// The worker counts every stage from 0 to 1. One bar for the whole picture: each stage owns a +// slice in the order it runs (the reference encode only exists with reference images), and the +// bar stays below 100% until the result is placed. +const STAGE_SLICES = { + new: [['load', .15], ['text', .15], ['denoise', .5], ['decode', .2]], + edit: [['load', .15], ['text', .15], ['encode', .05], ['denoise', .45], ['decode', .2]], +}; +export function imageGenProgress(stage, fraction, withImages = false) { + const own = Math.min(1, Math.max(0, Number.isFinite(fraction) ? fraction : 0)); + let start = 0; + for (const [name, weight] of STAGE_SLICES[withImages ? 'edit' : 'new']) { + if (name === stage) return Math.min(.99, start + weight * own); + start += weight; + } + return Math.min(.99, own); +} + +// An epoch rejects late uploads/status/results. A generate response arriving +// after Stop still has to be cancelled, even though its job was not known yet. +export function createImageGenSession({client, prepare, apply, onChange = () => {}, + schedule = setTimeout, unschedule = clearTimeout, + visible = () => globalThis.document?.visibilityState !== 'hidden'}) { + let epoch = 0, timer = null, infoTimer = null, disposed = false, checking = null; + let state = {info: null, checking: false, running: false, job: null, fraction: 0, stage: '', error: ''}; + const emit = patch => { state = {...state, ...patch}; if (!disposed) onChange(state); }; + const cancelJob = job => Promise.resolve().then(() => client.cancel(job)).catch(() => {}); + function scheduleInfo() { + unschedule(infoTimer); infoTimer = null; + if (disposed || state.info?.installed) return; + infoTimer = schedule(() => { + infoTimer = null; + if (visible()) void check(); else scheduleInfo(); + }, IMAGE_GEN_INFO_POLL_MS); + } + async function check() { + if (checking || disposed) return checking; + emit({checking: true}); + checking = (async () => { + try { const info = await client.info(); if (!disposed) emit({info}); } + catch (error) { if (!disposed && !state.running) emit({error: error.code || error.message}); } + finally { checking = null; if (!disposed) { emit({checking: false}); scheduleInfo(); } } + })(); + return checking; + } + async function start(request) { + if (state.running || disposed || !request.prompt.trim() || !state.info?.installed) return; + const current = ++epoch; + emit({running: true, job: null, fraction: 0, stage: '', error: ''}); + const alive = () => current === epoch && !disposed; + let withImages = false; + async function poll(job, context) { + try { + const row = await client.status(job); + if (!alive()) return; + const fraction = imageGenProgress(row.stage, row.fraction, withImages); + emit({fraction: Math.max(state.fraction, fraction), stage: row.stage || ''}); + if (row.state === 'done') { + if (!row.file) throw failure(); + await apply(row, context, alive); + if (alive()) emit({running: false, job: null, fraction: 1, stage: ''}); + } else if (row.state === 'error') throw failure(row.error); + else timer = schedule(() => { timer = null; void poll(job, context); }, IMAGE_GEN_POLL_MS); + } catch (error) { + if (alive()) { + void cancelJob(job); + emit({running: false, job: null, stage: '', error: error.code || error.message}); + } + } + } + try { + const {context, images = []} = await prepare(request, alive); + if (!alive()) return; + withImages = images.length > 0; + const {job} = await client.generate({prompt: request.prompt, width: context.width, height: context.height, images}); + if (!job) throw failure(); + if (!alive()) { await cancelJob(job); return; } + emit({job}); + timer = schedule(() => { timer = null; void poll(job, context); }, IMAGE_GEN_POLL_MS); + } catch (error) { + if (alive()) emit({running: false, job: null, stage: '', error: error.code || error.message}); + } + } + function cancel() { + ++epoch; unschedule(timer); timer = null; + const job = state.job; + emit({running: false, job: null, stage: '', fraction: 0, error: ''}); + return job ? cancelJob(job) : Promise.resolve(); + } + return {check, start, cancel, clearError: () => emit({error: ''}), state: () => state, + dispose() { disposed = true; unschedule(infoTimer); infoTimer = null; return cancel(); }}; +} + +export function createImageGeneration({client, renderer, t, getDoc, getSelection, commit, el, button, segmented, + library = [], resolve = file => file, icon = () => '', onRatioChange}) { + let mode = 'edit', prompt = '', mounted = null, references = []; + const copies = new Map(); + // A reference is sent no larger than the largest picture the service makes. + async function referencePNG(file) { + const image = await renderer.image(file); + const k = Math.min(1, 2048 / Math.max(image.naturalWidth, image.naturalHeight)); + const canvas = el('canvas'); canvas.width = Math.max(1, Math.round(image.naturalWidth * k)); canvas.height = Math.max(1, Math.round(image.naturalHeight * k)); + canvas.getContext('2d').drawImage(image, 0, 0, canvas.width, canvas.height); + return new Promise((done, fail) => canvas.toBlob(b => b ? done(b) : fail(failure()), 'image/png')); + } + const session = createImageGenSession({client, + prepare: async (request, alive) => { + const context = generationContext(getDoc(), getSelection(), request.mode); + const images = request.mode === 'edit' ? [await uploadGenerationReference(renderer, client, context, alive), + ...await uploadExtraReferences(client, request.references || [], referencePNG, copies, alive)] : []; + return {context, images}; + }, + apply: async (result, context, alive) => { + const image = await renderer.image(result.file); + if (!alive()) return; + renderer.images.get(result.file).value = image; + commit(insertGenerationResult(getDoc(), {src: result.file, nw: image.naturalWidth, nh: image.naturalHeight}, context, t)); + }, onChange: update, + }); + const limit = () => referenceLimit(session.state().info); + const chosen = () => mode === 'edit' ? references.slice(0, limit() - 1) : []; + const start = () => session.start({mode, prompt, references: chosen()}); + function drawEdited(canvas) { + try { + const doc = getDoc(), layer = generationContext(doc, null, 'edit').layer; + const scale = Math.min(1, 128 / (layer ? Math.max(layer.nw, layer.nh) : Math.max(doc.width, doc.height))); + const image = layer ? renderer.layerImage(doc, layer, scale) : renderer.composite(doc, scale); + const k = Math.max(canvas.width / image.width, canvas.height / image.height); + canvas.getContext('2d').drawImage(image, (canvas.width - image.width * k) / 2, (canvas.height - image.height * k) / 2, image.width * k, image.height * k); + } catch { /* drawn on the next update once its image has loaded */ } + } + function renderReferences(refs, disabled) { + const most = limit(), list = chosen(); + refs.hidden = mode !== 'edit' || most < 2 || !(library.length || list.length); + refs.replaceChildren(); + if (refs.hidden) return; + const slot = (number, picture, name) => { + const item = el('span', null, 'fv-ie-gen-ref'); item.setAttribute('role', 'listitem'); + item.setAttribute('aria-label', name); + item.append(picture, el('b', String(number))); + return item; + }; + const edited = el('canvas'); edited.width = edited.height = 56; drawEdited(edited); + refs.append(slot(1, edited, t('Image 1', '图 1'))); + list.forEach((reference, index) => { + const picture = el('img'); picture.src = resolve(reference.file); picture.alt = ''; + const item = slot(index + 2, picture, `${t('Image', '图')} ${index + 2}`); + const remove = button('×', t('Remove', '移除'), () => { references = removeReference(references, reference.file); update(); }, 'fv-ie-gen-ref-remove'); + remove.setAttribute('aria-label', t('Remove', '移除')); remove.disabled = disabled; + item.append(remove); refs.append(item); + }); + const rest = library.filter(item => !list.some(r => r.file === item.file)); + const menu = el('details', null, 'fv-ie-menu fv-ie-gen-ref-add'); + const summary = el('summary'); summary.innerHTML = icon('plus', 16); + summary.title = t('Reference image', '参考图'); summary.setAttribute('aria-label', summary.title); + const full = disabled || list.length >= most - 1 || !rest.length; + summary.setAttribute('aria-disabled', String(full)); if (full) summary.tabIndex = -1; + summary.onclick = event => { if (full) event.preventDefault(); }; + const choices = el('div', null, 'fv-ie-menu-list'); + for (const item of rest) { + const choice = button('', null, () => { references = addReference(references, item, most); update(); }, 'fv-ie-menu-item'); + const picture = el('img'); picture.src = resolve(item.file); picture.alt = ''; picture.loading = 'lazy'; + choice.append(picture, el('span', item.label)); choices.append(choice); + } + menu.append(summary, choices); refs.append(menu); + } + function update() { + if (!mounted?.bar.isConnected) return; + const {input, action, note, meter, fill, group, bar, refs, ratio} = mounted, state = session.state(); + const missing = !state.running && state.info && !state.info.installed; + const disabled = state.running || state.checking || !state.info?.installed; + const empty = !getDoc().layers.length; + if (empty) mode = 'new'; + bar.dataset.empty = String(empty); + if (ratio) { + ratio.hidden = !empty; ratio.disabled = state.running; + ratio.value = BLANK_CANVAS_RATIOS.includes(getDoc().ratio) ? getDoc().ratio : blankCanvasRatio(getDoc()); + } + group.querySelectorAll('button').forEach(b => { + b.disabled = disabled || (empty && b.dataset.value === 'edit'); + b.setAttribute('aria-pressed', String(b.dataset.value === mode)); + }); input.disabled = disabled; + renderReferences(refs, disabled); + input.placeholder = mode === 'edit' ? t('What to change', '要改什么') : t('Describe the picture', '描述画面'); + input.setAttribute('aria-label', input.placeholder); + action.textContent = state.running ? t('Stop', '停止') : t('Generate', '生成'); + action.classList.toggle('fv-ie-primary', !state.running); + action.disabled = !state.running && (disabled || !prompt.trim()); action.hidden = !!missing; + const percent = Math.floor(state.fraction * 100) + '%'; + note.textContent = missing ? imagePackMessage(state.info, t) : state.error ? imageGenError(state.error, state.info, t) + : state.running ? (state.stage === 'load' ? t('Loading the model', '加载模型') + ' ' : '') + percent : ''; + note.hidden = !note.textContent; + bar.dataset.running = String(state.running); bar.dataset.message = String(!!missing || !!state.error); + meter.hidden = !state.running; meter.setAttribute('aria-valuenow', String(Math.floor(state.fraction * 100))); + fill.style.width = state.fraction * 100 + '%'; + } + function mount(bar) { + bar.classList.add('fv-ie-generation'); + const group = segmented([['edit', t('Edit this', '改这张')], ['new', t('New image', '生成一张')]], mode, + value => { mode = value; update(); }, t('Generate with AI', 'AI 生成')).group; + let ratio = null; + if (!getDoc().layers.length) { + ratio = el('select', null, 'fv-ie-select'); ratio.id = 'fv-ie-gen-ratio'; + ratio.setAttribute('aria-label', t('Shape', '比例')); + for (const value of BLANK_CANVAS_RATIOS) { + const option = el('option', value); option.value = value; ratio.append(option); + } + ratio.onchange = () => { + if (!getDoc().layers.length && !session.state().running) onRatioChange?.(ratio.value); + update(); + }; + } + const field = el('div', null, 'fv-ie-gen-field'); + const refs = el('div', null, 'fv-ie-gen-refs'); refs.setAttribute('role', 'list'); + refs.setAttribute('aria-label', t('Reference images', '参考图')); + const box = el('div', null, 'fv-ie-gen-box'); + const input = el('textarea', null, 'fv-ie-gen-prompt'); input.rows = 1; input.maxLength = 2000; input.value = prompt; + box.append(input); field.append(refs, box); + const note = el('span', '', 'fv-ie-opt-note fv-ie-gen-note'); note.setAttribute('role', 'status'); + const action = button(t('Generate', '生成'), null, () => { + if (session.state().running) void session.cancel(); else void start(); + }, 'fv-ie-button fv-ie-primary'); + const meter = el('div', null, 'fv-ie-gen-meter'), fill = el('i'); meter.append(fill); + meter.setAttribute('role', 'progressbar'); meter.setAttribute('aria-label', t('Generate with AI', 'AI 生成')); + meter.setAttribute('aria-valuemin', '0'); meter.setAttribute('aria-valuemax', '100'); + bar.append(group); + if (ratio) bar.append(ratio); + bar.append(field, note, action, meter); + mounted = {bar, input, action, note, meter, fill, group, refs, ratio}; + const grow = () => { input.style.height = '0px'; input.style.height = Math.min(80, Math.max(32, input.scrollHeight + 2)) + 'px'; }; + input.oninput = () => { prompt = input.value; session.clearError(); grow(); }; + input.onkeydown = event => { + if (event.key === 'Enter' && !event.shiftKey && !event.isComposing && !event.ctrlKey && !event.metaKey) { + event.preventDefault(); void start(); + } + }; + grow(); update(); + } + return {mount, unmount() { mounted = null; }, selected: () => session.check(), cancel: session.cancel, + dispose: session.dispose, state: session.state}; +}