diff --git a/freevideo_engine/agent/cli.py b/freevideo_engine/agent/cli.py index 06259869..793729b4 100644 --- a/freevideo_engine/agent/cli.py +++ b/freevideo_engine/agent/cli.py @@ -76,6 +76,12 @@ 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('imports', help='Bring local video or audio into the editing timeline') + imports = command.add_subparsers(dest='operation', required=True) + command = imports.add_parser('add', help='Upload content; poll imports list until ready before insert') + command.add_argument('path') + schema.cli_arguments(command, schema.obj({'ref': schema.REF})) + imports.add_parser('list', help='List imported sources and preparation progress') 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') @@ -280,6 +286,9 @@ def source(value): result = client.get(path, **value) if method == 'GET' else client.post(path, value) elif args.command == 'upload': result = client.upload(args.path) + elif args.command == 'imports': + result = (client.import_media(args.path, **({'ref': args.ref} if hasattr(args, 'ref') else {})) + if args.operation == 'add' else client.get('/imports')) else: result = client.frame(source(args.source), values.get('at', 'last')) except ApiError as error: diff --git a/freevideo_engine/agent/client.py b/freevideo_engine/agent/client.py index aae09919..7bdfc67a 100644 --- a/freevideo_engine/agent/client.py +++ b/freevideo_engine/agent/client.py @@ -281,6 +281,16 @@ def asset(self, identity, target, *, thumbnail=False): def upload(self, path): """Read on the client computer and upload content; never send a local path.""" + return self._upload(path, '/media') + + def import_media(self, path, *, ref=None): + """Import a client file into the editor. Poll /imports until its state is ready.""" + from . import schema + if ref is not None: + schema.validate(schema.obj({'ref': schema.REF}), dict(ref=ref)) + return self._upload(path, '/imports', ref=ref) + + def _upload(self, path, endpoint, *, ref=None): from .schema import MAX_UPLOAD_BYTES, MAX_JSON_UPLOAD_BYTES source = local_file(path) try: @@ -289,7 +299,10 @@ def upload(self, path): raise ApiError(413, 'too_large', 'File exceeds upload limit', dict(next='Reduce the file to at most %d bytes.' % MAX_UPLOAD_BYTES, details=[dict(limit_bytes=MAX_UPLOAD_BYTES)])) if size * 4 // 3 + 4096 <= MAX_JSON_UPLOAD_BYTES: - return self.upload_bytes(source.name, source.read_bytes()) + if endpoint == '/media': + return self.upload_bytes(source.name, source.read_bytes()) + return self.post(endpoint, dict(name=source.name, data_base64=base64.b64encode(source.read_bytes()).decode(), + **({'ref': ref} if ref is not None else {})), timeout=600) import secrets boundary = 'freevideo-' + secrets.token_hex(16) # Only a safe extension crosses a MIME header; local filenames stay local. @@ -298,14 +311,15 @@ def upload(self, path): raise ApiError(422, 'invalid_media', 'Invalid media extension') prefix = ('--%s\r\nContent-Disposition: form-data; name="file"; filename="media%s"\r\n' 'Content-Type: application/octet-stream\r\n\r\n' % (boundary, suffix)).encode('ascii') - ending = ('\r\n--%s--\r\n' % boundary).encode('ascii') + ending = (('\r\n--%s\r\nContent-Disposition: form-data; name="ref"\r\n\r\n%s' % (boundary, ref) + if ref is not None else '') + '\r\n--%s--\r\n' % boundary).encode('ascii') with source.open('rb') as stream: def chunks(): yield prefix while data := stream.read(1024 * 1024): yield data yield ending - return self.request('POST', '/freevideo/agent/v1/media', encoded=chunks(), + return self.request('POST', '/freevideo/agent/v1' + endpoint, encoded=chunks(), content_type='multipart/form-data; boundary=' + boundary, content_length=len(prefix) + size + len(ending), timeout=600) except OSError: diff --git a/freevideo_engine/agent/guide.py b/freevideo_engine/agent/guide.py index 56a6ccfd..c274f10b 100644 --- a/freevideo_engine/agent/guide.py +++ b/freevideo_engine/agent/guide.py @@ -74,8 +74,8 @@ ], 'quality': [ 'Steps per quality: "light" = 8 first-pass steps + 3 refine steps, two-pass, the first pass at half size ' - '(default); "medium" = the same 8 + 3 two-pass with a larger first pass, for more detail; "high" = 20 steps ' - 'with a step cache and "max" = 20 plain steps, each in one pass. The same on a Mac. sampling.two_pass=false ' + '(default); "medium" = the same 8 + 3 two-pass with a larger first pass, for more detail; "high" = 16 steps ' + 'with a step cache and "max" = 20 steps with a step cache, each in one pass. The same on a Mac. sampling.two_pass=false ' 'without base_steps runs 20 steps.', 'Two-pass generation runs only with 8 first-pass steps. Any other base_steps (for example 20 with ' 'two_pass=true) runs in one pass of that many steps; the job notice says so. A preview is always two-pass ' @@ -128,6 +128,14 @@ ], 'project_routes': [dict(operation=name, method=method, path=path, description=summary) for name, (method, path, summary) in PROJECT_OPERATIONS.items()], + 'bring_your_own_footage': [ + 'Render complete clips with ffmpeg/Canvas on your computer; bake overlays into their pixels. FreeVideo receives media, not your code. Use the project dimensions (even pixels), 24 fps and integer frame positions; seconds = frames / 24.', + 'Download generated shots using result.video.url, library video.url, or sources[key].video/audio from GET /projects/{id}, with the same Authorization: Bearer token (Range supported). These URLs are relative to the server origin. Render your changes locally.', + 'POST /imports accepts multipart file then optional ref, or JSON {name, data_base64, ref?}; HTTP never reads a local path. MCP freevideo_import_media and CLI imports add read a local path and upload content, like media upload. Limits: 2 GiB content, 96 MiB JSON body, and the same media decode limits.', + 'The answer is {import: {id, source: "import:", state, fraction, ...}}. As on the page, poll GET /imports (MCP freevideo_imports, CLI imports list) about every two seconds. Match id in items; preparing is not insertable, ready is usable, failed is terminal with reason. Same ref or Idempotency-Key and filename/content reuses one import for that X-FreeVideo-Client, including after restart; changed content conflicts.', + 'The page receiver and transcode queue create an upright 24 fps H.264 preview by dropping/duplicating timestamped frames, with even dimensions and longest side at most 1920; preview audio is 32 kHz stereo AAC. Original dimensions are retained. Inserting the first video follows its original size by default (up to 8192 pixels); set project size_mode="fixed" to keep a chosen canvas. Later clips fit the canvas (cover by default). Export reads the original when available at the project output size/rate (preview fallback if missing).', + 'Insert with {op:"insert", source:"import:", at:48} in POST /projects/{id}/operations using the current revision. Add subtitles/text and audio through project operations, then POST /projects/{id}/export and wait for the job. There is one video track; effects are baked into complete clips, not separate overlay tracks or blend modes.', + ], 'editing_images': [ 'POST /images opens a picture as the image editor does; GET /images/{id} reads its document (layers), selection and revision.', 'POST /images/{id}/operations sends revision and 1-64 operations run with the page\'s own commands; the batch saves once or changes nothing. Layer IDs of new layers are in results[i].new_layer_ids.', @@ -275,6 +283,8 @@ ('POST', '/jobs/{id}/cancel'): '取消任务', ('POST', '/jobs/{id}/resume'): '接着生成一条视频里还没完成的部分', ('POST', '/media'): '上传图片、视频或音频内容', + ('POST', '/imports'): '导入视频或音频到剪辑时间线;轮询 GET /imports 等待就绪', + ('GET', '/imports'): '列出已导入素材和准备进度', ('GET', '/media/{kind}/{file}/file'): '下载上传或处理后的素材', ('POST', '/media/inspect'): '查看素材信息', ('POST', '/media/frame'): '提取视频的一帧', diff --git a/freevideo_engine/agent/imports.py b/freevideo_engine/agent/imports.py new file mode 100644 index 00000000..1dfa5ddd --- /dev/null +++ b/freevideo_engine/agent/imports.py @@ -0,0 +1,142 @@ +"""Content uploads into the editor's import store and software transcode queue.""" +import asyncio +import base64 +import binascii +import hashlib +from pathlib import Path +import tempfile +from types import SimpleNamespace +from urllib.parse import quote, urlencode + +from .. import comfy_editor as editor +from . import media, schema, spec +from .jobs import JobError + + +def invalid(path, expected, code='invalid_media'): + error = spec.SpecError(schema.diagnostic(path, expected, None)) + error.code = code + raise error + + +def view(value): + fields = ('id', 'name', 'kind', 'state', 'fraction', 'reason', 'frames', 'fps', 'width', 'height', + 'preview_width', 'preview_height', 'has_audio') + row = {key: value.get(key) for key in fields} + row['fps'] = value.get('fps', editor.FPS) + row['has_audio'] = value.get('has_audio', False) + row['source'] = 'import:' + value['id'] + for kind in ('preview', 'thumbnail'): + row[kind] = (schema.PREFIX + '/editor/sources/content?' + urlencode(dict(key=row['source'], file=kind)) + if row['state'] == 'ready' else None) + return row + + +def listing(root): + return dict(items=[view(row) for row in editor.list_imports(root)]) + + +def replay(root, receipt): + if receipt is None: + return None + for row in editor.list_imports(root): + value = editor._read_import(root, row['id']) + saved = value.get('agent_receipt') + if saved and (saved.get('client'), saved.get('ref')) == (receipt['client'], receipt['ref']): + if saved['digest'] != receipt['digest']: + invalid('/ref', 'the same filename and content for a replay', 'idempotency_conflict') + return view(row) + return None + + +class Content: + """Adapt a checked upload to the page's streaming receiver, without a server-side path API.""" + def __init__(self, path): + self.path = path + + async def iter_chunked(self, size): + with self.path.open('rb') as stream: + while chunk := await asyncio.to_thread(stream.read, size): + yield chunk + + +async def upload(request, value, root, manager, lock, decoders): + if manager is None: + raise JobError('not_ready', 'Editor import queue unavailable') + # Temporary content stays beside the destination, and is removed on rejection or disconnect. + with tempfile.TemporaryDirectory(prefix='.agent-import-', dir=root) as temporary: + multipart = request.content_type == 'multipart/form-data' + field = '/file' if multipart else '/data_base64' + ref = None if multipart else value.get('ref') + if multipart: + reader = await request.multipart() + part = await reader.next() + if part is None or getattr(part, 'name', None) != 'file' or not getattr(part, 'filename', None): + invalid('/file', 'one nonempty file part, followed by optional ref') + name = part.filename + else: + name = value['name'] + spec.import_name(name) + target = Path(temporary) / ('media' + Path(name).suffix.lower()) + digest = hashlib.sha256() + size = 0 + with target.open('xb') as stream: + if multipart: + while chunk := await part.read_chunk(size=256 * 1024): + size += len(chunk) + if size > schema.MAX_UPLOAD_BYTES: + raise JobError('too_large', 'Upload exceeds limit', limit_bytes=schema.MAX_UPLOAD_BYTES) + digest.update(chunk) + await asyncio.to_thread(stream.write, chunk) + extra = await reader.next() + if extra is not None: + if getattr(extra, 'name', None) != 'ref' or getattr(extra, 'filename', None): + invalid('/file', 'one file part followed only by optional ref') + data = bytearray() + while chunk := await extra.read_chunk(): + data.extend(chunk) + if len(data) > 1024: + spec.fail('/ref', 'a ref of at most 128 characters') + try: + ref = data.decode('utf-8') + except UnicodeDecodeError: + spec.fail('/ref', 'a UTF-8 ref') + if await reader.next() is not None: + invalid('/file', 'one file part followed only by optional ref') + else: + try: + data = base64.b64decode(value['data_base64'], validate=True) + except (binascii.Error, ValueError): + invalid(field, 'valid base64-encoded video or audio content') + size = len(data) + if size > schema.MAX_UPLOAD_BYTES: + raise JobError('too_large', 'Upload exceeds limit', limit_bytes=schema.MAX_UPLOAD_BYTES) + digest.update(data) + await asyncio.to_thread(stream.write, data) + if not size: + invalid(field, 'nonempty video or audio content') + key = request.headers.get('Idempotency-Key') + for where, candidate in (('/ref', ref), ('/headers/Idempotency-Key', key)): + if candidate is not None: + problems = schema.issues(schema.REF, candidate, where) + if problems: + raise spec.SpecError(problems) + if ref and key and ref != key: + invalid('/ref', 'the same value as Idempotency-Key', 'idempotency_conflict') + ref = key or ref + receipt = (dict(ref=ref, client=request.headers.get('X-FreeVideo-Client', '')[:80] or None, + digest=hashlib.sha256(name.encode() + b'\0' + digest.digest()).hexdigest()) if ref else None) + async with decoders: + try: + await asyncio.to_thread(media.validate_upload, target) + except (OSError, ValueError, SyntaxError): + invalid(field, 'decodable video or audio within the media upload limits') + async with lock: + previous = await asyncio.to_thread(replay, root, receipt) + if previous is not None: + return {'import': previous}, 200 + adapted = SimpleNamespace(content_length=size, headers={'X-FreeVideo-Filename': quote(name, safe='')}, + content=Content(target)) + made = await editor.receive(adapted, root, receipt=receipt) + manager.submit(made['id']) + return {'import': view(made)}, 201 diff --git a/freevideo_engine/agent/mcp.py b/freevideo_engine/agent/mcp.py index 62ce0d27..e94b1f88 100644 --- a/freevideo_engine/agent/mcp.py +++ b/freevideo_engine/agent/mcp.py @@ -31,7 +31,12 @@ "frame": "last"}), look at results with freevideo_look, ask the local VLM with freevideo_ask, and join clips with freevideo_assemble. For durable filming, use freevideo_drama_plan → start → get → sheet → choose/reshoot → finish. Cast first: search freevideo_cast_list, reuse the card the user named or one that fits, create a card only if none does. -If paused is set, wait for the user using get with since and wait=25. Only the user may pause/resume.""" +If paused is set, wait for the user using get with since and wait=25. Only the user may pause/resume. +Bring your own footage: render complete video clips on your computer with ffmpeg/Canvas, baking effects into pixels. +freevideo_import_media reads your local file and uploads content (up to 2 GiB); reuse ref for safe retries. +Poll freevideo_imports about every two seconds until state=ready (failed is terminal), then insert its source=import: +at an integer frame in a 24 fps project, add subtitles/audio, and export. GET /guide bring_your_own_footage explains +authenticated source downloads, dimensions and audio. No second video track or blend mode is needed.""" def tool_list(): @@ -53,6 +58,8 @@ def tool_list(): dict(name='freevideo_cancel', title='Cancel a job', description='Cancel a waiting, queued or running job.', annotations={'readOnlyHint': False, 'destructiveHint': True, 'openWorldHint': False}), dict(name='freevideo_resume', title='Resume generation', description='Resume a failed, interrupted or cancelled generation; only unfinished blocks run. Poll the same job ID.', annotations=write), dict(name='freevideo_add_media', title='Add media', description='Read a local image, video or audio file and upload its content; returns the "file" to use in media.', annotations=write), + dict(name='freevideo_import_media', title='Import footage', description='Read a local video or audio file and upload it into the editor. Returns import.source (import:) and state. Poll freevideo_imports until ready, then insert that source into a project. Same ref and content safely replay.', annotations=write), + dict(name='freevideo_imports', title='Imported footage', description='List editor imports with source, state (preparing/ready/failed), fraction and authenticated preview URLs. Wait for ready before insert; failed is terminal.', annotations=read), dict(name='freevideo_look', title='Look at media', description='Return a picture you can see: one frame of a video (at = first, last or seconds) or a contact sheet of timed frames (mode = sheet), or an image file. Use it to review shots.', annotations=read), dict(name='freevideo_edit_image', title='Edit an image', description="Crop/rotate/flip/pad/adjust an image, add text, shapes or image overlays, or remove burned-in captions (remove_text), without the browser. The result is a new input file plus an editor document a person can reopen in FreeVideo's image editor.", annotations=write), dict(name='freevideo_ask', title='Ask the local VLM', description='Ask FreeVideo\'s local vision-language model about images or sampled frames of a video (e.g. "Does this shot match its description? Answer as JSON {score, issues}").', annotations=write), @@ -359,6 +366,10 @@ def tool(self, identity, name, args, token): return c.resume(args['id'], **({'ref': args['ref']} if 'ref' in args else {})), [] if name == 'freevideo_add_media': return c.upload(args['path']), [] + if name == 'freevideo_import_media': + return c.import_media(args['path'], **({'ref': args['ref']} if 'ref' in args else {})), [] + if name == 'freevideo_imports': + return c.get('/imports'), [] if name == 'freevideo_edit_image': made = c.edit(args['source'], args.get('recipe'), document=args.get('document'), target=args.get('target')) return made, self.picture(made['view']) diff --git a/freevideo_engine/agent/routes.py b/freevideo_engine/agent/routes.py index e0fa1135..74f7eb73 100644 --- a/freevideo_engine/agent/routes.py +++ b/freevideo_engine/agent/routes.py @@ -165,7 +165,8 @@ def features(continuation, vlm): from .editor_runtime import available rules = available() return {'capabilities': True, 'plan': True, 'jobs.cancel': True, 'jobs.resume': True, 'takes': True, - 'library.read': True, 'library.download': True, 'errors.validation': True, 'errors.missing_media': True, + 'library.read': True, 'library.download': True, 'imports.create': True, 'imports.read': True, + 'errors.validation': True, 'errors.missing_media': True, 'generate.text': True, 'generate.size': True, 'generate.first_frame': True, 'generate.seed': True, 'queue.shared': True, 'queue.busy': True, 'chain.auto': bool(continuation), 'enhance': bool(vlm.get('ready')), 'drama.film': True, 'images.edit': rules, @@ -196,7 +197,7 @@ def capabilities(port, jobs): service = getattr(port.server, '_freevideo_prompt_vlm', None) if service is not None: try: - vlm = {k: v for k, v in service.info().items() if k in ('ready', 'model', 'license', 'bytes', 'busy')} + vlm = {k: v for k, v in service.info(wait=False).items() if k in ('ready', 'model', 'license', 'bytes', 'busy')} except (OSError, ValueError): vlm = dict(ready=False) running, pending = port.snapshot() @@ -387,6 +388,7 @@ async def off(request): token_failures = set() thumbnails = asyncio.Semaphore(1) upload_decoders = asyncio.Semaphore(1) + import_lock = asyncio.Lock() def expected(root=None): from ..comfy_bridge import installation_root @@ -492,7 +494,7 @@ async def respond(request): if definition.get('multipart') and request.content_type == 'multipart/form-data': pass else: - limit = contract.MAX_JSON_UPLOAD_BYTES if suffix == '/media' else contract.MAX_BODY_BYTES + limit = contract.MAX_JSON_UPLOAD_BYTES if suffix in ('/media', '/imports') else contract.MAX_BODY_BYTES value = await read_body(request, limit, definition.get('optional_body', False)) try: contract.validate(definition['body'], value) @@ -518,8 +520,10 @@ async def respond(request): value['next'] = value['next'].replace('{id}', request.match_info['id']) if error.code == 'too_large': limit = error.details.get('limit_bytes', contract.MAX_UPLOAD_BYTES) + location = ({'path': '/file' if request.content_type == 'multipart/form-data' else '/data_base64'} + if suffix == '/imports' else {}) value.update(message='The upload must be at most %d bytes; a larger payload was received.' % limit, - details=[dict(limit_bytes=limit)], next='Reduce the payload to at most %d bytes and retry.' % limit) + details=[dict(limit_bytes=limit, **location)], next='Reduce the payload to at most %d bytes and retry.' % limit) if error.code == 'busy': entries = (await asyncio.to_thread(jobs.queue_view))['entries'] value = errors.build('busy', owner=entries[0]['owner'] if entries else 'ui') @@ -1466,6 +1470,26 @@ def relocate(value): made.pop('path', None) return assets.public(dict(made, **info, file=relative, view=assets.view(relative))) + @routes.post(PREFIX + '/imports') + @guarded + async def import_media(request): + from . import imports + value = None if request.content_type == 'multipart/form-data' else await body(request) + try: + result, status = await imports.upload(request, value, port.output_root, + getattr(server, '_freevideo_editor', None), import_lock, upload_decoders) + except specs.SpecError: + raise + except (ValueError, AssertionError): + imports.invalid('/file', 'a well-formed multipart upload with video or audio content') + return reply(result, status) + + @routes.get(PREFIX + '/imports') + @guarded + async def list_imports(request): + from . import imports + return reply(await asyncio.to_thread(imports.listing, port.output_root)) + @routes.post(PREFIX + '/media') @guarded async def upload(request): diff --git a/freevideo_engine/agent/schema.py b/freevideo_engine/agent/schema.py index dfca3b7c..79fa2b4d 100644 --- a/freevideo_engine/agent/schema.py +++ b/freevideo_engine/agent/schema.py @@ -14,7 +14,7 @@ S = 'https://json-schema.org/draft/2020-12/schema' PREFIX = '/freevideo/agent/v1' # (two_pass, base_steps, refine_steps) of each level on this platform; the step -# cache a level adds is tiers.step_cache(name, mac) (only high has one, easycache-av), +# cache a level adds is tiers.step_cache(name, mac) (high and max use easycache-av), # its two-pass upscale factor tiers.upscale_factor(name, mac) (medium samples a larger first pass). QUALITY = {name: tiers.plan(name, mac=tiers.on_mac()) for name in tiers.TIERS} ASPECTS = {'16:9': (16, 9), '9:16': (9, 16), '1:1': (1, 1), '4:3': (4, 3), '3:4': (3, 4), @@ -316,6 +316,20 @@ def batch_schema(*, local=False): return obj(dict(jobs=array(job_schema(local=local, refs=True), 1, 64), ref=REF), ('jobs',), description='x-batch additionally checks earlier-row references and at most 64 jobs after takes expansion.', **{'x-batch': True}) +def import_upload(): + return obj(dict(name=string(255, pattern=r'^[^/\\:\x00-\x1f]+$'), + data_base64=string(MAX_JSON_UPLOAD_BYTES), ref=REF), ('name', 'data_base64')) + + +def import_view(): + return obj(dict(id=string(80), source=EDITOR_SOURCE, name=string(255), kind=enum(('video', 'audio')), + state=enum(('preparing', 'ready', 'failed')), fraction=number(0, 1), reason=nullable(string()), + frames=nullable(number(1, integer=True)), fps=number(1), width=nullable(number(0, integer=True)), + height=nullable(number(0, integer=True)), preview_width=nullable(number(0, integer=True)), + preview_height=nullable(number(0, integer=True)), has_audio=nullable({'type': 'boolean'}), + preview=nullable(string()), thumbnail=nullable(string())), ('id', 'source', 'name', 'kind', 'state', 'fraction')) + + def contracts(*, local=False): """Every registered route, including the UI control plane, has a contract.""" def route(summary, body=None, query=None, **kw): @@ -346,6 +360,9 @@ def route(summary, body=None, query=None, **kw): ('POST', '/media'): route('Upload image, video or audio content', body=obj(dict(name=string(255, pattern=r'^[^/\\:\x00]+$'), data_base64=string(MAX_JSON_UPLOAD_BYTES)), ('name', 'data_base64')), status=201, multipart=True), + ('POST', '/imports'): route('Import video or audio into the editing timeline; poll GET /imports until ready', + body=import_upload(), status=201, multipart=True, idempotent=True), + ('GET', '/imports'): route('List editor imports and preparation progress'), ('GET', '/media/{kind}/{file}/file'): route('Download uploaded or derived media'), ('POST', '/media/inspect'): route('Inspect media', body=obj(source, ('source',))), ('POST', '/media/frame'): route('Extract a video frame', body=obj(dict(source, at=dict(FRAME, default='last')), ('source',)), status=201), @@ -1293,6 +1310,7 @@ def schemas(): 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_view=cast_view(), projects={name: project_schema(name) for name in PROJECT_OPERATIONS}, + imports=dict(upload=import_upload(), item=import_view()), images={name: image_schema(name) for name in IMAGE_OPERATIONS}, whitebox={name: whitebox_schema(name) for name in WHITEBOX_OPERATIONS}, sessions=dict(event=session_event(stored=True), meta=session_view(), @@ -1316,6 +1334,8 @@ def openapi(): content = {'application/json': {'schema': deepcopy(contract['body'])}} if contract.get('multipart'): content['multipart/form-data'] = {'schema': obj({'file': {'type': 'string', 'format': 'binary'}}, ('file',))} + if suffix == '/imports': + content['multipart/form-data']['schema']['properties']['ref'] = REF operation['requestBody'] = dict(required=not contract.get('optional_body', False), content=content) operation['parameters'] = [dict(name=k, **{'in': 'query'}, required=k in contract['query'].get('required', []), schema=v) for k, v in contract['query']['properties'].items()] @@ -1356,6 +1376,9 @@ def openapi(): if suffix == '/sessions' and method == 'GET' else session_view()) elif suffix.endswith('/events'): content_type, response_schema = 'text/event-stream', {'type': 'string'} + elif suffix == '/imports': + response_schema = obj({'import': import_view()}, ('import',)) if method == 'POST' else obj( + {'items': array(import_view())}, ('items',)) elif suffix == '/editor/sources/content': content_type, response_schema = 'application/octet-stream', {'type': 'string', 'format': 'binary'} elif suffix.endswith('/file') or suffix.endswith('/thumbnail'): @@ -1411,6 +1434,7 @@ def openapi(): 'freevideo_job': ('GET', '/jobs/{id}'), 'freevideo_wait': ('GET', '/jobs/{id}'), 'freevideo_jobs': ('GET', '/jobs'), 'freevideo_cancel': ('POST', '/jobs/{id}/cancel'), 'freevideo_resume': ('POST', '/jobs/{id}/resume'), 'freevideo_add_media': ('POST', '/media'), + 'freevideo_import_media': ('POST', '/imports'), 'freevideo_imports': ('GET', '/imports'), 'freevideo_edit_image': ('POST', '/media/edit'), 'freevideo_library': ('GET', '/library'), } TOOL_ENDPOINTS.update({'freevideo_drama_' + name: (method, suffix) @@ -1547,6 +1571,8 @@ def _tool_schema(name): return batch_schema(local=True) if name == 'freevideo_add_media': return obj({'path': string(4096)}, ('path',)) + if name == 'freevideo_import_media': + return obj(dict(path=string(4096), ref=REF), ('path',)) if '{id}' in suffix: value['properties']['id'] = contract.get('path', {}).get('properties', {}).get('id', JOB) value['required'] = list(value.get('required', [])) + ['id'] diff --git a/freevideo_engine/agent/spec.py b/freevideo_engine/agent/spec.py index c3fdeb48..a5d51c08 100644 --- a/freevideo_engine/agent/spec.py +++ b/freevideo_engine/agent/spec.py @@ -27,6 +27,22 @@ def fail(path, message): raise SpecError(contract.diagnostic(path, message, None)) +def import_name(name): + """Timeline imports accept video/audio content, with no HTTP filesystem paths.""" + from ..comfy_assets import media_kind + problems = contract.issues(contract.import_upload()['properties']['name'], name, '/name') + if problems: + raise SpecError(problems) + try: + if media_kind(name) not in ('video', 'audio'): + raise ValueError + except ValueError: + error = SpecError(contract.diagnostic('/name', 'a supported video or audio filename', None)) + error.code = 'invalid_media' + raise error from None + return name + + # Field checks for values no route schema describes: screenplays and filming records. def _bool(value, path): if type(value) is not bool: @@ -147,9 +163,7 @@ def sampling(value): cache = tiers.step_cache(quality, mac) if own else None factor = tiers.upscale_factor(quality, mac) if own and two else 2 if custom: - quality = next((k for k, v in QUALITY.items() if v == (two, base, refine) - and tiers.step_cache(k, mac) == cache - and (not two or tiers.upscale_factor(k, mac) == factor)), 'custom') + quality = tiers.name_for(base, two, refine, cache, mac, factor) or 'custom' return dict(quality=quality, two_pass=two, base_steps=base, refine_steps=refine, **({'step_cache': cache} if cache else {}), **({'upscale_factor': factor} if factor != 2 else {}), **({'sampling_note': note} if note else {})) diff --git a/freevideo_engine/comfy_bridge.py b/freevideo_engine/comfy_bridge.py index d112fb37..96010987 100644 --- a/freevideo_engine/comfy_bridge.py +++ b/freevideo_engine/comfy_bridge.py @@ -1486,7 +1486,7 @@ def asset_progress(message): # Reference audio selects its own tables; match the encoder's choice. asset_task = ((ANCHOR_ADALN_TASK if anchored or pictured else ADALN_TASK) if continue_from is not None else engine_task(extra.get('media', {}), media_base)) - # The cache only names the level of a missing table (high or max at 20 steps). + # Tables follow the step count; the cache is carried for exact level recognition. preparation = prepare_sampling_assets(root, machine, dict(planned, step_cache=step_cache) if step_cache else planned, asset_task, progress=asset_progress, interrupted=interrupted, environ=environment) diff --git a/freevideo_engine/comfy_editor.py b/freevideo_engine/comfy_editor.py index 1d88c671..5150f7a2 100644 --- a/freevideo_engine/comfy_editor.py +++ b/freevideo_engine/comfy_editor.py @@ -866,7 +866,7 @@ def _display_name(value): return name -async def receive(request, root): +async def receive(request, root, *, receipt=None): if request.content_length is not None and request.content_length > MAX_IMPORT_BYTES: raise TooLarge('Import exceeds 8 GiB') name = _display_name(request.headers.get('X-FreeVideo-Filename')) @@ -894,6 +894,9 @@ async def receive(request, root): await asyncio.to_thread(os.replace, temporary, destination) value = dict(id=identity, name=name, kind=kind, state='preparing', fraction=0., reason=None, frames=None, fps=FPS, width=None, height=None, has_audio=False) + if receipt is not None: + # Keep the Agent replay key in the same atomic record as the import. + value['agent_receipt'] = receipt await asyncio.to_thread(_write, _import_path(root, identity), value) return {key: value[key] for key in ('id', 'name', 'kind', 'state', 'fraction')} except BaseException: diff --git a/freevideo_engine/decode.py b/freevideo_engine/decode.py index 2c6d79f4..7ed5a3e5 100644 --- a/freevideo_engine/decode.py +++ b/freevideo_engine/decode.py @@ -227,6 +227,9 @@ def decoded(done, total, elapsed): audio = audio[:, :round(kept_frames * 40 / 24) * 800] if continuation is not None: audio = audio[:, continuation['audio_latents'] * 800:] + audio_level = {'audio_rms': float(audio.square().mean().sqrt()) if audio.numel() else 0., + 'audio_peak': float(audio.abs().max()) if audio.numel() else 0., + 'audio_samples': int(audio.shape[-1])} del audio_vae, audio_mean, audio_std gc.collect() torch.cuda.empty_cache() @@ -269,6 +272,7 @@ def decoded(done, total, elapsed): 'vae_linear_compute_cache': linear_compute_cache, 'vae_compute_cache_preparation': compute_cache_metrics, 'resident_vae_cache_hit': vae_cache_hit, 'resident_audio_vae_cache_hit': audio_cache_hit, 'vae_preload_seconds': preload_seconds, + **audio_level, 'audio_load_decode_seconds': audio_seconds, 'encode_seconds': time.perf_counter() - encode_started, 'decoded_artifact_save_seconds': artifact_seconds, 'streamed_video_output': stream_output, 'streamed_vae_weights': stream_weights, diff --git a/freevideo_engine/diagnostic_resources.py b/freevideo_engine/diagnostic_resources.py index cf4ed692..6738492f 100644 --- a/freevideo_engine/diagnostic_resources.py +++ b/freevideo_engine/diagnostic_resources.py @@ -243,7 +243,7 @@ def decoder_measurements(value): result = numbers(value, ('vae_load_seconds', 'video_decode_seconds', 'audio_load_decode_seconds', 'video_postprocess_seconds', 'encode_seconds', 'decoded_artifact_save_seconds', 'decode_save_seconds', 'vae_resident_blocks', - 'vae_pinned_weight_bytes')) + 'vae_pinned_weight_bytes', 'audio_rms', 'audio_peak', 'audio_samples')) for key in ('vae_linear_compute_cache', 'resident_vae_cache_hit', 'resident_audio_vae_cache_hit', 'streamed_video_output', 'streamed_vae_weights'): if type(value.get(key)) is bool: diff --git a/freevideo_engine/failure_details.py b/freevideo_engine/failure_details.py index adf51af2..3b770a74 100644 --- a/freevideo_engine/failure_details.py +++ b/freevideo_engine/failure_details.py @@ -88,7 +88,7 @@ def sampling_table_failure(value, *, zh=False): return None from .launcher_copy import size_text language = 1 if zh else 0 - # Tables follow the step count; the cache tells high from max at 20 steps. + # Tables follow the step count: 16 steps is high, 20 steps is max. level = table_level(details['base_steps'], details['two_pass'], details['refine_steps'], details.get('step_cache'), on_mac(), details.get('scale', 2)) index = list(TIERS).index(level) if level else None diff --git a/freevideo_engine/image_gen/failure.py b/freevideo_engine/image_gen/failure.py index ff30ff5c..a29582f7 100644 --- a/freevideo_engine/image_gen/failure.py +++ b/freevideo_engine/image_gen/failure.py @@ -15,8 +15,8 @@ ERRORS = frozenset(('no_memory', 'busy_video', 'busy', 'pack_missing', 'model_unreadable', 'device_unavailable', 'worker_exit', 'timeout', 'cancelled', - 'invalid_request', 'generation_failed')) -NO_REPORT = frozenset(('busy', 'invalid_request', 'cancelled')) + 'invalid_request', 'invalid_media', 'invalid_points', 'segment_not_installed', 'generation_failed')) +NO_REPORT = frozenset(('busy', 'invalid_request', 'invalid_media', 'invalid_points', 'segment_not_installed', 'cancelled')) STAGES = ('start', 'load', 'text', 'encode', 'denoise', 'decode') KINDS = ('generate', 'edit', 'cast', 'compose', 'download', 'install') MESSAGE_ZH = { @@ -30,6 +30,9 @@ 'timeout': '这一张用时太长,已停止。请再试一次。', 'generation_failed': '图像生成出错了,请再试一次。', 'invalid_request': '图片设置有误,请调整后重试。', + 'invalid_media': '图片无效,请重新选择输入目录中的图片。点选支持不超过 4096×4096 像素总数的图片。', + 'invalid_points': '请选择 1–16 个图片范围内的点,且至少有一个选中点;1 为选中,0 为排除。', + 'segment_not_installed': '点选模型未安装,请先下载点选模型(SAM 2.1 small)。', } _DEVICE = re.compile( r'cuda.*(?:initializ|device.*(?:lost|unavailable|removed)|invalid device ordinal|' diff --git a/freevideo_engine/image_gen/image_models.json b/freevideo_engine/image_gen/image_models.json index e7840ef2..129d9c1f 100644 --- a/freevideo_engine/image_gen/image_models.json +++ b/freevideo_engine/image_gen/image_models.json @@ -14,6 +14,31 @@ {"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": "hf-mirror", "url": "https://hf-mirror.com/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": "hf-mirror", "url": "https://hf-mirror.com/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": "hf-mirror", "url": "https://hf-mirror.com/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"}]} + ], + "segment": [ + { + "file": "sam2/sam2.1_hiera_small.pt", + "component": "segment", + "bytes": 184416285, + "sha256": "6d1aa6f30de5c92224f8172114de081d104bbd23dd9dc5c58996f0cad5dc4d38", + "repo": "facebook/sam2.1-hiera-small", + "revision": "main", + "model_id": "sam2.1-hiera-small", + "sources": [ + { + "id": "official", + "url": "https://huggingface.co/facebook/sam2.1-hiera-small/resolve/main/sam2.1_hiera_small.pt" + }, + { + "id": "hf-mirror", + "url": "https://hf-mirror.com/facebook/sam2.1-hiera-small/resolve/main/sam2.1_hiera_small.pt" + }, + { + "id": "modelscope", + "url": "https://www.modelscope.cn/models/facebook/sam2.1-hiera-small/resolve/master/sam2.1_hiera_small.pt" + } + ] + } ] } } diff --git a/freevideo_engine/image_gen/image_pack.py b/freevideo_engine/image_gen/image_pack.py index 3f9a4bc4..9e0b7e5f 100644 --- a/freevideo_engine/image_gen/image_pack.py +++ b/freevideo_engine/image_gen/image_pack.py @@ -138,17 +138,22 @@ def validate_catalog(spec): 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): + image_rows = spec['groups']['image'] + segment_rows = spec['groups'].get('segment', []) + found = image_rows + segment_rows + if (not image_rows or len({r['file'] for r in found}) != len(found) + or any(row['component'] == 'segment' for row in image_rows) + or any(row['component'] != 'segment' for row in segment_rows)): 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')): + or row['component'] not in ('dit', 'text_encoder', 'vae', 'text_encoder_config', 'processor', 'segment')): raise PackError('invalid_catalog') ids = [s['id'] for s in row['sources']] - if (len(set(ids)) != len(ids) or not {'official', 'modelscope'} <= set(ids) + required = {'official', 'modelscope'} + if (len(set(ids)) != len(ids) or not required <= set(ids) or not set(ids) <= {'official', 'modelscope', 'hf-mirror'}): raise PackError('invalid_catalog') for source in row['sources']: @@ -159,11 +164,13 @@ def validate_catalog(spec): raise PackError('invalid_catalog') from error -def rows(spec=None): +def rows(spec=None, *, group='image'): 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']] + groups = spec['groups'].values() if group is None else [spec['groups'][group]] + return [dict(row, family='image', model_id=row.get('model_id', spec['model_id']), + repo=row.get('repo', spec['repo']), revision=row.get('revision', spec['revision'])) + for found in groups for row in found] def base(root=None): @@ -216,27 +223,30 @@ def _owned_intact(path, stamp): # for people reading the source and is never needed at run time. NOTICE_TEXT = (b'Qwen is licensed under the Qwen RESEARCH LICENSE AGREEMENT, Copyright (c) 2026 Hangzhou Tongyi ' b'Laboratory Technology Co., Ltd. All Rights Reserved.\n') +SEGMENT_NOTICE_TEXT = (b'SAM 2.1 Hiera small, Copyright (c) Meta Platforms, Inc. and affiliates.\n' + b'Licensed under Apache-2.0; see freevideo_engine/licenses/SAM2-Apache-2.0.txt.\n' + b'https://github.com/facebookresearch/sam2\n') def _notice(): return NOTICE_TEXT -def _notice_ready(path): +def _notice_ready(path, *, group='image'): """Informational only: a missing or different notice never makes the pack unusable.""" try: - return path.is_file() and path.read_bytes() == NOTICE_TEXT + return path.is_file() and path.read_bytes() == (SEGMENT_NOTICE_TEXT if group == 'segment' else NOTICE_TEXT) except OSError: return False -def _write_notice(folder, owned): +def _write_notice(folder, owned, *, group='image'): """Put the notice into the model folder when possible; a failure is logged and the install goes on.""" try: - notice = _path(folder, 'NOTICE.txt') + notice = _path(folder, 'SAM2-NOTICE.txt' if group == 'segment' else 'NOTICE.txt') if notice.exists(): return # present (ours or the user's own text): left as it is - notice.write_bytes(NOTICE_TEXT) + notice.write_bytes(SEGMENT_NOTICE_TEXT if group == 'segment' else NOTICE_TEXT) _own(folder, owned, notice) except Exception as error: logging.warning('[FreeVideo] Could not write the image model notice: %s', redacted_launcher_error(error)) @@ -277,7 +287,7 @@ def _recheck(root, folder, pending): try: stamps = deepcopy(_read(folder, LEDGER)) results = [] - for row in rows(): + for row in rows(group=None): path = target(row, root) name = str(path) if name in pending and _stale(path, row, stamps): @@ -349,13 +359,13 @@ def _start_recheck(root, folder, stale): return True -def status(root=None): +def status(root=None, *, group='image'): """Missing files and unique retained bytes; no network and no weight reads on this thread.""" folder = base(root) ledger = _read(folder, LEDGER) files, missing, downloaded, verified_bytes = [], [], 0, 0 stale = [] - for row in rows(): + for row in rows(group=group): path = target(row, root) intact = _intact(path, row, ledger) old = not intact and _stale(path, row, ledger) @@ -372,7 +382,7 @@ def status(root=None): files.append(dict(file=row['file'], bytes=row['bytes'], ready=intact, downloaded_bytes=retained, stale=old)) checking = bool(stale) and _start_recheck(root, folder, stale) try: - notice_ready = _notice_ready(_path(folder, 'NOTICE.txt')) + notice_ready = _notice_ready(_path(folder, 'SAM2-NOTICE.txt' if group == 'segment' else 'NOTICE.txt'), group=group) except PackError: notice_ready = False total = sum(r['bytes'] for r in files) @@ -391,8 +401,8 @@ def checking(root=None) -> bool: return base(root) in _RECHECKS -def ready(root=None): - return status(root)['ready'] +def ready(root=None, *, group='image'): + return status(root, group=group)['ready'] def downloading(root=None, engine_root=None) -> bool: @@ -444,19 +454,19 @@ def _lease(root, engine_root=None, destructive=False): yield folder -def verify(root=None, progress=None, check=None): +def verify(root=None, progress=None, check=None, *, group='image'): """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) + _verify(folder, root, progress or (lambda v: None), check or (lambda: None), full=True, group=group) + return status(root, group=group) -def _verify(folder, root, progress, check, full=False): +def _verify(folder, root, progress, check, full=False, *, group='image'): ledger = _read(folder, LEDGER) todo = [] - total = sum(r['bytes'] for r in rows()) + total = sum(r['bytes'] for r in rows(group=group)) done = 0 - for row in rows(): + for row in rows(group=group): check() path = target(row, root) hashed = False @@ -559,14 +569,14 @@ def resource_check(): raise PackError(reason) from error -def install(root=None, progress=None, check=None, *, engine_root=None, repair=False, _background=None): +def install(root=None, progress=None, check=None, *, engine_root=None, repair=False, _background=None, group='image'): """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) + ledger, todo = _verify(folder, root, progress, check, full=repair, group=group) total = sum(r['bytes'] for r in todo) retained = 0 for row in todo: @@ -581,11 +591,11 @@ def install(root=None, progress=None, check=None, *, engine_root=None, repair=Fa plan = _download_plan(engine_root) plan.update(quiet=True, resource_check=check) owned = _read(folder, OWNERSHIP) - _write_notice(folder, owned) + _write_notice(folder, owned, group=group) done = 0 pending_prefixes = 0 if _background is not None: - done = sum(r['bytes'] for r in rows()) - total + done = sum(r['bytes'] for r in rows(group=group)) - 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()) @@ -624,7 +634,7 @@ def moved(current, size, speed, **kwargs): done += row['bytes'] progress(dict(phase='download', done=done + pending_prefixes, total=total)) check() - return status(root) + return status(root, group=group) except (PackError, InterruptedError): raise except BlockingIOError as error: @@ -633,18 +643,19 @@ def moved(current, size, speed, **kwargs): 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 repair(root=None, progress=None, check=None, *, engine_root=None, group='image'): + return install(root, progress, check, engine_root=engine_root, repair=True, group=group) -def removal_plan(root=None): +def removal_plan(root=None, *, group='image'): folder = base(root) owned = _read(folder, OWNERSHIP) - allowed = {'NOTICE.txt'} - for row in rows(): + notice = 'SAM2-NOTICE.txt' if group == 'segment' else 'NOTICE.txt' + allowed = {notice} + for row in rows(group=group): allowed.update((row['file'], row['file'] + '.partial', row['file'] + '.partial.source.json')) for name in allowed: - if name == 'NOTICE.txt': + if name == notice: continue _path(folder, name) files, skipped = [], [] @@ -654,7 +665,7 @@ def removal_plan(root=None): try: path = _path(folder, name) except PackError: - if name != 'NOTICE.txt': + if name != notice: raise skipped.append(name) continue @@ -667,12 +678,12 @@ def removal_plan(root=None): return dict(files=files, bytes=sum(f['bytes'] for f in files), skipped=skipped) -def remove(root=None, check=None, *, engine_root=None): +def remove(root=None, check=None, *, engine_root=None, group='image'): """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) + plan = removal_plan(root, group=group) owned, ledger = _read(folder, OWNERSHIP), _read(folder, LEDGER) released, removed = 0, [] for entry in plan['files']: @@ -696,8 +707,9 @@ def remove(root=None, check=None, *, engine_root=None): class Installer: """One background action per pack; cancel keeps completed files and prefixes.""" - def __init__(self, root=None, *, engine_root=None, busy=None): + def __init__(self, root=None, *, engine_root=None, busy=None, group='image'): self.root, self.engine_root = root, engine_root + self.group = group self.generation_busy = busy self.lock = threading.Lock() self.stop = threading.Event() @@ -805,9 +817,11 @@ def check(): try: control = _Background(self, stop) if background else None if background: - initial = status(self.root) + initial = status(self.root, group=self.group) self._set(done=initial['downloaded_bytes'], total=initial['total_bytes']) options = dict(root=self.root, check=control.wait if control else check) + if self.group != 'image': + options['group'] = self.group if control: options['_background'] = control if action != 'verify': @@ -833,11 +847,11 @@ def _failed(self, action, error): code = image_pack_error_code(error) try: spec = catalog() - entries = rows(spec) + entries = rows(spec, group=self.group) except Exception: spec, entries = {}, [] try: - pack = status(self.root) + pack = status(self.root, group=self.group) except Exception: pack = dict(ready=False, files=[], missing=[], downloaded_bytes=0) with self.lock: @@ -853,7 +867,8 @@ def _failed(self, action, error): file.update(checks.get(row['file'], {})) files.append(file) pack = {k: pack[k] for k in ('ready', 'missing', 'downloaded_bytes', 'verified_bytes', 'total_bytes') if k in pack} - pack.update(name=spec.get('model_id', 'image'), repo=spec.get('repo'), revision=spec.get('revision'), + identity = entries[0] if entries and self.group == 'segment' else spec + pack.update(name=identity.get('model_id', 'image'), repo=identity.get('repo'), revision=identity.get('revision'), format=spec.get('format'), files=files, transferred_bytes=downloaded, sources=[dict(file=row['file'], id=s['id'], host=urlsplit(s['url']).hostname) for row in entries for s in row['sources']], events=events) @@ -901,6 +916,7 @@ def main(argv=None): commands = parser.add_subparsers(dest='command', required=True) download = commands.add_parser('download') download.add_argument('--background', action='store_true') + download.add_argument('--group', choices=('image', 'segment'), default='image') download.add_argument('--root', type=Path) download.add_argument('--engine-root', type=Path) download.add_argument('--status-file', type=Path, required=True) @@ -915,7 +931,7 @@ def main(argv=None): print('Image pack status is unavailable.', file=sys.stderr) return 1 - installer = Installer(args.root, engine_root=args.engine_root) + installer = Installer(args.root, engine_root=args.engine_root, group=args.group) received = threading.Event() def cancelled(signum, frame): received.set() diff --git a/freevideo_engine/image_gen/image_runner.py b/freevideo_engine/image_gen/image_runner.py index e96c8002..07546fdf 100644 --- a/freevideo_engine/image_gen/image_runner.py +++ b/freevideo_engine/image_gen/image_runner.py @@ -223,7 +223,8 @@ def plan(self, width, height, reference_images=None): 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): + def generate(self, prompt, *, width, height, seed=None, reference_images=None, edit_mask=None, + progress=None, cancel=None): plan = self.plan(width, height, reference_images) self.last_plan = plan if not plan.feasible: @@ -233,7 +234,8 @@ def generate(self, prompt, *, width, height, seed=None, reference_images=None, p self.generator.vae_tile = plan.vae_tiling self.generator.vae_tile_size = plan.vae_tile_size return self.generator.generate(prompt, width=width, height=height, seed=seed, - reference_images=reference_images, progress=progress, cancel=cancel) + reference_images=reference_images, edit_mask=edit_mask, + progress=progress, cancel=cancel) def at_boundary(self): """False when a failure left a weights streamer inside a layer pass; the next request cannot use it.""" diff --git a/freevideo_engine/image_gen/qwen_pipeline.py b/freevideo_engine/image_gen/qwen_pipeline.py index 8bde8b50..15ebbc98 100644 --- a/freevideo_engine/image_gen/qwen_pipeline.py +++ b/freevideo_engine/image_gen/qwen_pipeline.py @@ -257,7 +257,7 @@ def _reference_inputs(self, images): @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: + reference_images=None, edit_mask=None, progress=None, cancel=None) -> Image.Image: if self._closed: raise RuntimeError("Generator is closed") if not isinstance(prompt, str): @@ -267,6 +267,8 @@ def generate(self, prompt: str, *, width: int, height: int, seed: int | None = N 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) + if edit_mask is not None and not reference_images: + raise ValueError("edit_mask needs a reference image: the first one is the picture being edited") def check_cancel(): if cancel is not None and cancel(): @@ -314,6 +316,21 @@ def report(stage, fraction): encoded = (encoded - mean) / std ref_latents.append(self._pack_latents(encoded, 1, channels, *encoded.shape[-2:]).to(embeds.dtype)) check_cancel() + keep_latents = cell_mask = None + if edit_mask is not None: + from . import region_edit + + original = reference_images[0].convert("RGBA") + with self._stage("encode"): + check_cancel() + resized = original.resize((width, height), Image.Resampling.LANCZOS) + array = np.asarray(resized, dtype=np.float32)[None] / 255.0 + tensor = torch.from_numpy(array).permute(0, 3, 1, 2).unsqueeze(2) * 2 - 1 + encoded = _encoded_latents(self.vae.encode(tensor.to(device=self.device, dtype=_dtype(self.vae)), tile=True)) + mean, std = self._stats(encoded) + keep_latents = self._pack_latents((encoded - mean) / std, 1, channels, *encoded.shape[-2:]).to(embeds.dtype) + cell_mask = torch.from_numpy(region_edit.latent_mask(edit_mask, (width, height))).to(self.device) + cell_mask = cell_mask.reshape(1, -1, 1) 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: @@ -325,6 +342,7 @@ def report(stage, fraction): 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) + start_noise = latents.clone() if keep_latents is not None else None 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) @@ -348,6 +366,11 @@ def report(stage, fraction): ))[:, -latents.shape[1]:] check_cancel() latents = euler_step(latents, prediction, sigmas[i], sigmas[i + 1]) + if keep_latents is not None: + # Outside the selection the latent stays on the original's own path to this sigma. + s = sigmas[i + 1] + kept = (1 - s) * keep_latents.float() + s * start_noise.float() + latents = (cell_mask * latents.float() + (1 - cell_mask) * kept).to(prediction.dtype) report("denoise", (i + 1) / steps) check_cancel() @@ -366,6 +389,9 @@ def report(stage, fraction): 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)) + if edit_mask is not None: + image = region_edit.composite(reference_images[0], image, + region_edit.blend_alpha(edit_mask, reference_images[0].size)) report("decode", 1.0) return image diff --git a/freevideo_engine/image_gen/region_edit.py b/freevideo_engine/image_gen/region_edit.py new file mode 100644 index 00000000..a1cb624c --- /dev/null +++ b/freevideo_engine/image_gen/region_edit.py @@ -0,0 +1,90 @@ +"""Edit only a selected region of a picture. + +The caller supplies the original picture and a mask (white = change). The +sampler keeps the unselected part of the latent on the original's own noise +path at every step, and the finished picture is blended over the original with +a feathered edge, so pixels outside the grown selection are the original's. +""" +import numpy as np +from PIL import Image, ImageDraw, ImageFilter + +LATENT_CELL = 16 +DEFAULT_GROW = 24 +DEFAULT_FEATHER = 16 + + +def _odd(radius): + return max(1, int(radius) * 2 + 1) + + +def grown_mask(mask, grow=DEFAULT_GROW): + """Mask as an L image, selection grown by ``grow`` pixels.""" + if not isinstance(mask, Image.Image): + raise ValueError("mask must be a PIL image") + mask = mask.convert("L").point(lambda v: 255 if v >= 128 else 0) + if grow > 0: + mask = mask.filter(ImageFilter.MaxFilter(_odd(grow))) + return mask + + +def lasso_mask(size, points): + """A drawn curve closed by joining its last point to its first, filled.""" + mask = Image.new("L", size, 0) + points = [(float(x), float(y)) for x, y in points] + if len(points) >= 3: + ImageDraw.Draw(mask).polygon(points, fill=255) + return mask + + +def combine_masks(size, added, removed=()): + """Union of the added selections minus the union of the removed ones (L images, 255 = selected).""" + result = np.zeros((size[1], size[0]), dtype=bool) + for mask in added: + result |= np.asarray(mask.convert("L").resize(size)) >= 128 + for mask in removed: + result &= ~(np.asarray(mask.convert("L").resize(size)) >= 128) + return Image.fromarray(result.astype(np.uint8) * 255, "L") + + +def selection_is_empty(mask): + return not np.asarray(mask.convert("L")).max() >= 128 + + +def latent_mask(mask, size, *, grow=DEFAULT_GROW, feather=DEFAULT_FEATHER): + """1.0 on the latent cells the model may change, else 0.0, shaped (rows, cols). + + Covers the grown selection plus the feather, so the soft edge of the final + blend lies inside regenerated content. + """ + width, height = size + base = grown_mask(mask.resize((width, height), Image.Resampling.NEAREST), grow + feather) + array = np.asarray(base, dtype=np.uint8) > 0 + rows, cols = height // LATENT_CELL, width // LATENT_CELL + cells = array[:rows * LATENT_CELL, :cols * LATENT_CELL].reshape(rows, LATENT_CELL, cols, LATENT_CELL) + return cells.any(axis=(1, 3)).astype(np.float32) + + +def blend_alpha(mask, size, *, grow=DEFAULT_GROW, feather=DEFAULT_FEATHER): + """Feathered 0..1 alpha at ``size``: 1 inside the grown selection, 0 beyond the feather.""" + resized = mask.resize(size, Image.Resampling.NEAREST) + soft = grown_mask(resized, grow).filter(ImageFilter.GaussianBlur(feather / 2)) + alpha = np.asarray(soft, dtype=np.float32) / 255.0 + # The blur is centred on the grown edge: widen so the alpha reaches 1 inside the selection + # and is exactly 0 where the latent mask no longer reaches. + hard = np.asarray(grown_mask(resized, grow + feather), dtype=np.uint8) > 0 + alpha = np.where(hard, alpha, 0.0) + return alpha + + +def composite(original, generated, alpha): + """Original pixels where alpha is 0; generated pixels (resized to the original) where 1.""" + original = original.convert("RGB") + if generated.size != original.size: + generated = generated.resize(original.size, Image.Resampling.LANCZOS) + src = np.asarray(original, dtype=np.float32) + new = np.asarray(generated.convert("RGB"), dtype=np.float32) + out = new * alpha[..., None] + src * (1.0 - alpha[..., None]) + result = np.asarray(original).copy() + changed = alpha > 0 + result[changed] = np.clip(np.rint(out[changed]), 0, 255).astype(np.uint8) + return Image.fromarray(result) diff --git a/freevideo_engine/image_gen/segment.py b/freevideo_engine/image_gen/segment.py new file mode 100644 index 00000000..a72b12bc --- /dev/null +++ b/freevideo_engine/image_gen/segment.py @@ -0,0 +1,116 @@ +"""CPU SAM 2.1 Hiera small; one pixel-hashed image encoding per worker.""" +import hashlib +import math + + +# Bound full-resolution candidate logits as well as image decompression. +MAX_PIXELS = 4096 * 4096 +MAX_POINTS = 16 + + +def validate_points(points, size=None): + if not isinstance(points, list) or not 1 <= len(points) <= MAX_POINTS: + raise ValueError('invalid_points') + for point in points: + if (not isinstance(point, (list, tuple)) or len(point) != 3 + or any(type(v) not in (int, float) or type(v) is float and not math.isfinite(v) + for v in point[:2]) + or type(point[2]) is not int or point[2] not in (0, 1) + or any(not 0 <= v < MAX_PIXELS for v in point[:2]) + or size is not None and (point[0] >= size[0] or point[1] >= size[1])): + raise ValueError('invalid_points') + if not any(p[2] for p in points): + raise ValueError('invalid_points') + + +def load_image(path): + from PIL import Image + try: + with Image.open(path) as image: + if image.width * image.height > MAX_PIXELS: + raise ValueError('invalid_media') + # Match region_edit: coordinates and output dimensions are stored pixels. + return image.convert('RGB') + except (OSError, ValueError, Image.DecompressionBombError) as error: + raise ValueError('invalid_media') from error + + +def build_model(checkpoint): + """Literal upstream sam2.1_hiera_s config, without Hydra or iopath.""" + import torch + from ..sam2.modeling.backbones.hieradet import Hiera + from ..sam2.modeling.backbones.image_encoder import ImageEncoder, FpnNeck + from ..sam2.modeling.position_encoding import PositionEmbeddingSine + from ..sam2.modeling.memory_attention import MemoryAttention, MemoryAttentionLayer + from ..sam2.modeling.memory_encoder import MemoryEncoder, MaskDownSampler, Fuser, CXBlock + from ..sam2.modeling.sam.transformer import RoPEAttention + from ..sam2.modeling.sam2_base import SAM2Base + + model = SAM2Base( + image_encoder=ImageEncoder(scalp=1, + trunk=Hiera(embed_dim=96, num_heads=1, stages=[1, 2, 11, 2], + global_att_blocks=[7, 10, 13], window_pos_embed_bkg_spatial_size=[7, 7]), + neck=FpnNeck(position_encoding=PositionEmbeddingSine(256, normalize=True, temperature=10000, warmup_cache=False), + d_model=256, backbone_channel_list=[768, 384, 192, 96], + fpn_top_down_levels=[2, 3], fpn_interp_model='nearest')), + memory_attention=MemoryAttention(d_model=256, pos_enc_at_input=True, num_layers=4, + layer=MemoryAttentionLayer(activation='relu', dim_feedforward=2048, dropout=.1, + pos_enc_at_attn=False, d_model=256, pos_enc_at_cross_attn_keys=True, + pos_enc_at_cross_attn_queries=False, + self_attention=RoPEAttention(rope_theta=10000., feat_sizes=[64, 64], + embedding_dim=256, num_heads=1, downsample_rate=1, dropout=.1), + cross_attention=RoPEAttention(rope_theta=10000., feat_sizes=[64, 64], rope_k_repeat=True, + embedding_dim=256, num_heads=1, downsample_rate=1, dropout=.1, kv_in_dim=64))), + memory_encoder=MemoryEncoder(out_dim=64, + position_encoding=PositionEmbeddingSine(64, normalize=True, temperature=10000, warmup_cache=False), + mask_downsampler=MaskDownSampler(kernel_size=3, stride=2, padding=1), + fuser=Fuser(layer=CXBlock(dim=256, kernel_size=7, padding=3, + layer_scale_init_value=1e-6, use_dwconv=True), num_layers=2)), + num_maskmem=7, image_size=1024, sigmoid_scale_for_mem_enc=20., sigmoid_bias_for_mem_enc=-10., + use_mask_input_as_output_without_sam=True, directly_add_no_mem_embed=True, + no_obj_embed_spatial=True, use_high_res_features_in_sam=True, multimask_output_in_sam=True, + iou_prediction_use_sigmoid=True, use_obj_ptrs_in_encoder=True, add_tpos_enc_to_obj_ptrs=True, + proj_tpos_enc_in_obj_ptrs=True, use_signed_tpos_enc_to_obj_ptrs=True, + only_obj_ptrs_in_the_past_for_eval=True, pred_obj_scores=True, pred_obj_scores_mlp=True, + fixed_no_obj_ptr=True, multimask_output_for_tracking=True, use_multimask_token_for_obj_ptr=True, + multimask_min_pt_num=0, multimask_max_pt_num=1, use_mlp_for_obj_ptr_proj=True, + compile_image_encoder=False) + state = torch.load(checkpoint, map_location='cpu', weights_only=True) + model.load_state_dict(state['model'], strict=True) + return model.cpu().eval() + + +class Segmenter: + def __init__(self, checkpoint): + import torch + from ..sam2.sam2_image_predictor import SAM2ImagePredictor + torch.set_num_threads(4) + torch.set_num_interop_threads(1) + self.predictor = SAM2ImagePredictor(build_model(checkpoint), max_hole_area=0, max_sprinkle_area=0) + self.key = None + self.encode_count = 0 + + def segment(self, image, points): + import numpy as np + import torch + from PIL import Image + validate_points(points, image.size) + image = image.convert('RGB') + key = image.size, hashlib.sha256(image.tobytes()).digest() + # The reference adapter groups positive points before negative points. + ordered = [p for p in points if p[2]] + [p for p in points if not p[2]] + with torch.inference_mode(): + if key != self.key: + self.clear_image() + self.predictor.set_image(np.asarray(image)) + self.key = key + self.encode_count += 1 + masks, scores, _ = self.predictor.predict( + point_coords=np.asarray([p[:2] for p in ordered], dtype=np.float32), + point_labels=np.asarray([p[2] for p in ordered], dtype=np.int32), + multimask_output=True) + return Image.fromarray(masks[int(np.argmax(scores))].astype(bool).astype(np.uint8) * 255) + + def clear_image(self): + self.key = None + self.predictor.reset_predictor() diff --git a/freevideo_engine/image_gen/segment_worker.py b/freevideo_engine/image_gen/segment_worker.py new file mode 100644 index 00000000..ab1e98f1 --- /dev/null +++ b/freevideo_engine/image_gen/segment_worker.py @@ -0,0 +1,81 @@ +"""CPU-only point selection worker. Idle exit releases the entire Torch allocator.""" +import argparse +from contextlib import redirect_stdout +import json +import os +from pathlib import Path +import queue +import sys + +from .worker import Inbox + + +def serve(checkpoint, *, idle_seconds=45., stream=None, output=None, factory=None): + from .segment import Segmenter, load_image, validate_points + stream = sys.stdin if stream is None else stream + output = sys.stdout if output is None else output + factory = factory or Segmenter + + def send(**row): + output.write(json.dumps(row, allow_nan=False) + '\n') + output.flush() + + inbox = Inbox(stream, send) + segmenter = 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 + send(phase='accepted') + stage = 'start' + try: + with redirect_stdout(sys.stderr): + image = load_image(value['image']) + validate_points(value['points'], image.size) + stage = 'load' + if segmenter is None: + segmenter = factory(checkpoint) + stage = 'encode' + mask = segmenter.segment(image, value['points']) + if inbox.cancel.is_set(): + break + with Path(value['path']).open('xb') as target: + mask.save(target, format='PNG') + result = dict(phase='result', path=value['path']) + except Exception as error: + from .failure import classify_exception + code = str(error) if str(error) in ('invalid_media', 'invalid_points') else classify_exception(error, stage=stage) + result = dict(phase='error', code=code) + with inbox.mutex: + inbox.busy = False + send(**result) + finally: + inbox.closed.set() + if segmenter is not None: + segmenter.clear_image() + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument('checkpoint', type=Path) + parser.add_argument('--idle-seconds', type=float, default=45.) + args = parser.parse_args() + if not 0 < args.idle_seconds < float('inf'): + parser.error('idle-seconds must be positive and finite') + # Set before importing Torch or torchvision; this process never uses CUDA/MPS. + os.environ.update(CUDA_VISIBLE_DEVICES='', HIP_VISIBLE_DEVICES='', + OMP_NUM_THREADS='4', MKL_NUM_THREADS='4', OPENBLAS_NUM_THREADS='4') + stream = None + if os.name == 'nt': + from ..win32 import private_stdin + stream = private_stdin() + serve(args.checkpoint, idle_seconds=args.idle_seconds, stream=stream) + return 0 + + +if __name__ == '__main__': + raise SystemExit(main()) diff --git a/freevideo_engine/image_gen/service.py b/freevideo_engine/image_gen/service.py index 0558c067..852cfaa5 100644 --- a/freevideo_engine/image_gen/service.py +++ b/freevideo_engine/image_gen/service.py @@ -1,5 +1,6 @@ """One active image and one waiting job, with a worker leased until idle or video work.""" import atexit +import asyncio from collections import OrderedDict from copy import deepcopy import json @@ -26,7 +27,7 @@ QUEUE_TIMEOUT = 600. QUEUE_POLL = .1 INVALID_ERRORS = frozenset(('invalid_request', 'invalid_prompt', 'invalid_dimensions', 'invalid_seed', - 'invalid_media', 'too_many_images')) + 'invalid_media', 'invalid_points', 'too_many_images')) class WorkerIdleExit(RuntimeError): @@ -122,6 +123,12 @@ def __init__(self, installation, input_directory, *, video_busy=None, idle_secon self.queue_wake = threading.Event() self.finished = OrderedDict() self.pack_installer = None + self.segment_installer = None + self.segment_worker = None + self.segmenting = False + self.segment_stop = threading.Event() + self.segment_done = threading.Event() + self.segment_done.set() atexit.register(self.close) @staticmethod @@ -130,7 +137,7 @@ def _video_busy(): return OWNER.active def busy(self): - return self._busy + return self._busy or self.segmenting def info(self): engine_root, machine = self.installation() @@ -141,6 +148,9 @@ def info(self): downloading=not ready and image_pack.downloading(machine['model_root'], engine_root), bytes=sum(row['bytes'] for row in spec['groups']['image']), busy=self.busy(), max_images=MAX_IMAGES, install=self.pack_state(ready), + segment=dict(ready=image_pack.ready(machine['model_root'], group='segment'), + max_points=16, device='cpu', + install=self.segment_installer.state() if self.segment_installer else None), bytes_per_second=None if ready else observed_speed(engine_root)) def pack_state(self, ready): @@ -157,21 +167,153 @@ def pack_state(self, ready): message_zh=image_pack_failure(value.get('reason'), zh=True)['detail']) return row - def pack_action(self, action): + def pack_action(self, action, *, group='image'): """Start or pause the image model download from the page; it yields to a running video.""" engine_root, machine = self.installation() root = machine['model_root'] with self.mutex: - installer = self.pack_installer + attribute = 'segment_installer' if group == 'segment' else 'pack_installer' + installer = getattr(self, attribute) if action == 'cancel': if installer is not None: installer.cancel() - elif not image_pack.ready(root) and not image_pack.downloading(root, engine_root): + elif not image_pack.ready(root, group=group) and not image_pack.downloading(root, engine_root): if installer is None or installer.root != root or installer.engine_root != engine_root: - installer = self.pack_installer = image_pack.Installer(root, engine_root=engine_root, busy=self.video_busy) + options = dict(group=group) if group != 'image' else {} + installer = image_pack.Installer(root, engine_root=engine_root, busy=self.video_busy, **options) + setattr(self, attribute, installer) installer.start('install', background=True) return self.info() + def media_path(self, name): + """Shared input-name and symlink policy for generation and point selection.""" + directory = Path(self.input_directory()).resolve() + folder, _, leaf = name.rpartition('/') if isinstance(name, str) else ('', '', '') + if (not isinstance(name, str) or not leaf or name.startswith('/') 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') + return str(path) + + def output_path(self): + directory = Path(self.input_directory()).resolve() + output = directory / 'freevideo-image' + if output.is_symlink() or output.resolve().parent != directory: + raise ValueError('invalid_media') + output.mkdir(exist_ok=True) + return str(output / (uuid.uuid4().hex + '.png')) + + def close_segment_worker(self): + from .. import processes + with self.mutex: + session, self.segment_worker = self.segment_worker, None + if session is not None: + processes.stop(session.process, grace=1) + session.process.wait() + session.close() + + def spawn_segment_worker(self, root, machine): + from .. import processes + row = image_pack.rows(group='segment')[0] + command = processes.module_command('freevideo_engine.image_gen.segment_worker', + image_pack.target(row, machine['model_root']), '--idle-seconds', 45.) + command[0] = machine['comfy_python'] + env = dict(self.environment(root), CUDA_VISIBLE_DEVICES='', HIP_VISIBLE_DEVICES='', + OMP_NUM_THREADS='4', MKL_NUM_THREADS='4', OPENBLAS_NUM_THREADS='4') + return Worker(processes.popen(command, env=env, stdin=subprocess.PIPE, stdout=subprocess.PIPE, + stderr=subprocess.PIPE, text=True, encoding='utf-8', errors='replace', + supervise=True, start_new_session=True)) + + def segment(self, value): + """Immediate admission, no waiting queue. The HTTP caller runs this off its event loop.""" + from .segment import validate_points + from ..locking import runtime_lock + if not isinstance(value, dict): + raise ValueError('invalid_request') + validate_points(value.get('points')) + path = self.media_path(value.get('image')) + 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.queued is not None: + raise ValueError('busy') + try: + root, machine = self.installation() + except (OSError, ValueError): + raise ValueError('segment_not_installed') from None + try: + if not image_pack.ready(machine['model_root'], group='segment'): + raise ValueError('segment_not_installed') + except (OSError, image_pack.PackError): + raise ValueError('segment_not_installed') from None + self.segmenting = True + self.segment_stop.clear() + self.segment_done.clear() + # Reap the idle Qwen worker before taking its runtime lease. + self.session_stop.set() + thread = self.thread + try: + if thread is not None: + thread.join() + output = self.output_path() + with runtime_lock(Path(root) / 'engine.lock', inherit=False): + for attempt in range(2): + with self.mutex: + if self.closed or self.segment_stop.is_set() or self.yielding or self.video_busy(): + raise ValueError('busy_video') + session = self.segment_worker + if session is not None and session.process.poll() is not None: + self.close_segment_worker() + session = None + if session is None: + session = self.segment_worker = self.spawn_segment_worker(root, machine) + accepted = False + deadline = time.monotonic() + INFERENCE_TIMEOUT + try: + session.send(action='generate', image=path, points=value['points'], path=output) + while True: + if self.closed or self.segment_stop.is_set() or self.yielding or self.video_busy(): + raise ValueError('busy_video') + if time.monotonic() >= deadline: + raise ValueError('timeout') + try: + row = session.rows.get(timeout=.1) + except queue.Empty: + continue + if row is None: + raise WorkerIdleExit() + if row.get('phase') == 'accepted': + accepted = True + elif row.get('phase') == 'result': + if row.get('path') != output or not Path(output).is_file(): + raise ValueError('worker_exit') + return dict(mask='freevideo-image/' + Path(output).name) + elif row.get('phase') == 'error': + raise ValueError(row['code'] if row.get('code') in ERRORS else 'generation_failed') + except (BrokenPipeError, WorkerIdleExit): + self.close_segment_worker() + if accepted or attempt: + raise ValueError('worker_exit') from None + except BlockingIOError: + raise ValueError('busy') from None + except Exception: + self.close_segment_worker() + if self.segment_stop.is_set(): + raise ValueError('busy_video') from None + raise + finally: + with self.mutex: + self.segmenting = False + self.segment_done.set() + self.queue_wake.set() + def request(self, value): if not isinstance(value, dict): raise ValueError('invalid_request') @@ -191,27 +333,24 @@ def request(self, value): 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'))): + paths = [self.media_path(name) for name in images] + mask = value.get('mask') + mask_path = None + if mask is not None: + # The region to change: a mask drawn on the picture being edited, the first one. + if not paths: 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) + mask_path = self.media_path(mask) + from PIL import Image + try: + with Image.open(paths[0]) as picture, Image.open(mask_path) as drawn: + if drawn.size != picture.size: + raise ValueError('invalid_media') + except OSError: + raise ValueError('invalid_media') from None return dict(prompt=prompt, width=dimensions[0] // 32 * 32, height=dimensions[1] // 32 * 32, - steps=8, seed=seed, images=paths, path=str(output / (uuid.uuid4().hex + '.png'))) + steps=8, seed=seed, images=paths, **(dict(mask=mask_path) if mask_path else {}), + path=self.output_path()) def start(self, action, value): if action not in ('generate', 'compose'): @@ -221,6 +360,8 @@ def start(self, action, value): with self.mutex: if self.closed: raise ValueError('closed') + if self.segmenting: + raise ValueError('busy') if self.queued is not None: raise ValueError('busy') if self.yielding: @@ -503,9 +644,11 @@ def stop_for_generation(self, only_idle=False): 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()): + if only_idle and (self.busy() or (self.segment_worker is None and + (self.thread is None or not self.thread.is_alive()))): return False self.yielding = True + self.segment_stop.set() self.stop.set() self.session_stop.set() if self.queued is not None: @@ -514,6 +657,8 @@ def stop_for_generation(self, only_idle=False): self.queued['video_since'] = time.monotonic() thread = self.thread try: + self.close_segment_worker() + self.segment_done.wait() if thread is not None and thread is not threading.current_thread(): thread.join() finally: @@ -570,6 +715,7 @@ def work(self, root, machine): failure = 'worker_exit' failure_error = None try: + self.close_segment_worker() from ..background_compile import COMPILE from ..encoder_prewarm import IDLE COMPILE.stop() @@ -667,7 +813,7 @@ def respawn(): with self.mutex: self.process = self.worker = None self.session_stop.set() - if self.busy(): + if self._busy: context = self.context or Context('generate', {}) fitted = self.state.get('fitted') self.state = (dict(phase='cancelled') if self.stop.is_set() @@ -915,6 +1061,10 @@ async def action(request): result = service.cancel(value.get('job')) elif name in ('pack-install', 'pack-cancel'): result = service.pack_action(name.split('-')[1]) + elif name in ('segment-install', 'segment-cancel'): + result = service.pack_action(name.split('-')[1], group='segment') + elif name == 'segment': + result = await asyncio.to_thread(service.segment, value) else: result = service.start(name, value) return web.json_response(result) diff --git a/freevideo_engine/image_gen/worker.py b/freevideo_engine/image_gen/worker.py index 7c5503c8..d97d807f 100644 --- a/freevideo_engine/image_gen/worker.py +++ b/freevideo_engine/image_gen/worker.py @@ -120,9 +120,13 @@ def progress(stage, fraction): for path in value.get('images', []): image = stack.enter_context(Image.open(path)) references.append(_flatten(image)) + region = {} + if value.get('mask'): + region['edit_mask'] = stack.enter_context(Image.open(value['mask'])) + region['edit_mask'].load() image = runner.generate( value['prompt'], width=value['width'], height=value['height'], - seed=value.get('seed'), reference_images=references, + seed=value.get('seed'), reference_images=references, **region, progress=progress, cancel=inbox.cancel.is_set) if inbox.cancel.is_set(): diff --git a/freevideo_engine/launcher/Main.qml b/freevideo_engine/launcher/Main.qml index dee2e399..e8578e28 100644 --- a/freevideo_engine/launcher/Main.qml +++ b/freevideo_engine/launcher/Main.qml @@ -82,6 +82,7 @@ ApplicationWindow { var more = rate > 0 || s.bytes_per_second > 0 ? sizeAndTime(left, rate).split(" · ")[1] : "" return t("Downloading the prompt model · ", "正在下载提示词模型 · ") + gb(m.install.done) + " / " + total + (more ? " · " + more : "") } + if (m.checking) return t("Checking the prompt model…", "正在检查提示词模型…") if (m.failure) return m.failure if (m.install.state === "cancelled") return t("Prompt model download paused · ", "提示词模型下载已暂停 · ") + done + " / " + total return t("You can also download the prompt model (", "可以再下载提示词模型(") + sizeAndTime(m.total_bytes).replace(" · ", t(", ", ",")) @@ -670,7 +671,7 @@ ApplicationWindow { } RowLayout { objectName: "promptModelRow" - visible: !!s.prompt_model && !s.prompt_model.ready + visible: !!s.prompt_model && !s.prompt_model.ready && !s.prompt_model.checking Layout.fillWidth: true; spacing: 14 FCheck { objectName: "promptModelSwitch"; Layout.fillWidth: true @@ -962,7 +963,7 @@ ApplicationWindow { } FCard { objectName: "promptModelCard" - visible: !!s.prompt_model && !s.prompt_model.ready && (s.prompt_model.suggest || s.prompt_model.busy + visible: !!s.prompt_model && !s.prompt_model.ready && (s.prompt_model.suggest || s.prompt_model.busy || s.prompt_model.checking || ["failed", "cancelled"].indexOf(s.prompt_model.install.state) >= 0) Layout.fillWidth: true; padding: 16; spacing: 10 FText { @@ -974,8 +975,8 @@ ApplicationWindow { fraction: s.prompt_model.install.total > 0 ? s.prompt_model.install.done / s.prompt_model.install.total : -1 } Flow { - Layout.fillWidth: true; spacing: 8 - FButton { objectName: "promptModelDownload"; visible: !s.prompt_model.busy + visible: !s.prompt_model.checking; Layout.fillWidth: true; spacing: 8 + FButton { objectName: "promptModelDownload"; visible: !s.prompt_model.busy && !s.prompt_model.checking text: s.prompt_model.suggest ? t("Download", "下载") : t("Resume", "继续"); enabled: !s.busy; onClicked: backend.promptModelAction("install") } FButton { objectName: "promptModelPause"; visible: s.prompt_model.busy; text: t("Pause", "暂停"); onClicked: backend.promptModelAction("cancel") } FButton { objectName: "promptModelDismiss"; visible: s.prompt_model.suggest; text: t("Don't ask again", "不再提示"); flat: true; onClicked: backend.dismissPromptHint() } diff --git a/freevideo_engine/launcher_session.py b/freevideo_engine/launcher_session.py index 3856bf52..15f24273 100644 --- a/freevideo_engine/launcher_session.py +++ b/freevideo_engine/launcher_session.py @@ -815,6 +815,9 @@ def prompt_model_action(self, action): engine = self._image_location()[1] if self.prompt_installer and self.prompt_installer.busy: return + from .prompt_vlm import assets + if assets.checking(engine): + return if not (self.prompt_installer and self.prompt_installer.root == engine): self.prompt_installer = prompt_installer.Installer(engine, busy=self.generating) self.prompt_installer.start() @@ -826,6 +829,7 @@ def prompt_model_view(self): total = sum(row['bytes'] for row in spec['files']) engine = self._image_location()[1] if self.installation_engine() else None ready = bool(engine) and assets.ready(engine) + checking = bool(engine) and not ready and assets.checking(engine) installer = self.prompt_installer active = installer.state() if installer and engine and installer.root == engine else {} busy = bool(active and installer.busy) @@ -841,9 +845,9 @@ def prompt_model_view(self): failure = {'disk_space': self.t('Not enough disk space.', '磁盘空间不足。'), 'busy': self.t('Another download of this model is running.', '这个模型正在被另一处下载。')}.get( active.get('reason'), self.t('The download stopped. Resume to continue.', '下载中断了,点继续可接着下载。')) - return dict(ready=ready, total_bytes=total, enough_disk=enough, busy=busy, install=active, failure=failure, + return dict(ready=ready, checking=checking, total_bytes=total, enough_disk=enough, busy=busy, install=active, failure=failure, downloaded_bytes=active.get('done', 0), title=self.t(*NAMES['text']), - suggest=bool(engine) and not ready and not busy and enough and not self.form['prompt_hint_dismissed'] + suggest=bool(engine) and not ready and not checking and not busy and enough and not self.form['prompt_hint_dismissed'] and active.get('state') not in ('failed', 'cancelled')) def dismiss_prompt_hint(self): diff --git a/freevideo_engine/licenses/SAM2-Apache-2.0.txt b/freevideo_engine/licenses/SAM2-Apache-2.0.txt new file mode 100644 index 00000000..261eeb9e --- /dev/null +++ b/freevideo_engine/licenses/SAM2-Apache-2.0.txt @@ -0,0 +1,201 @@ + Apache License + Version 2.0, January 2004 + http://www.apache.org/licenses/ + + TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION + + 1. Definitions. + + "License" shall mean the terms and conditions for use, reproduction, + and distribution as defined by Sections 1 through 9 of this document. + + "Licensor" shall mean the copyright owner or entity authorized by + the copyright owner that is granting the License. + + "Legal Entity" shall mean the union of the acting entity and all + other entities that control, are controlled by, or are under common + control with that entity. For the purposes of this definition, + "control" means (i) the power, direct or indirect, to cause the + direction or management of such entity, whether by contract or + otherwise, or (ii) ownership of fifty percent (50%) or more of the + outstanding shares, or (iii) beneficial ownership of such entity. + + "You" (or "Your") shall mean an individual or Legal Entity + exercising permissions granted by this License. + + "Source" form shall mean the preferred form for making modifications, + including but not limited to software source code, documentation + source, and configuration files. + + "Object" form shall mean any form resulting from mechanical + transformation or translation of a Source form, including but + not limited to compiled object code, generated documentation, + and conversions to other media types. + + "Work" shall mean the work of authorship, whether in Source or + Object form, made available under the License, as indicated by a + copyright notice that is included in or attached to the work + (an example is provided in the Appendix below). + + "Derivative Works" shall mean any work, whether in Source or Object + form, that is based on (or derived from) the Work and for which the + editorial revisions, annotations, elaborations, or other modifications + represent, as a whole, an original work of authorship. For the purposes + of this License, Derivative Works shall not include works that remain + separable from, or merely link (or bind by name) to the interfaces of, + the Work and Derivative Works thereof. + + "Contribution" shall mean any work of authorship, including + the original version of the Work and any modifications or additions + to that Work or Derivative Works thereof, that is intentionally + submitted to Licensor for inclusion in the Work by the copyright owner + or by an individual or Legal Entity authorized to submit on behalf of + the copyright owner. For the purposes of this definition, "submitted" + means any form of electronic, verbal, or written communication sent + to the Licensor or its representatives, including but not limited to + communication on electronic mailing lists, source code control systems, + and issue tracking systems that are managed by, or on behalf of, the + Licensor for the purpose of discussing and improving the Work, but + excluding communication that is conspicuously marked or otherwise + designated in writing by the copyright owner as "Not a Contribution." + + "Contributor" shall mean Licensor and any individual or Legal Entity + on behalf of whom a Contribution has been received by Licensor and + subsequently incorporated within the Work. + + 2. Grant of Copyright License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + copyright license to reproduce, prepare Derivative Works of, + publicly display, publicly perform, sublicense, and distribute the + Work and such Derivative Works in Source or Object form. + + 3. Grant of Patent License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + (except as stated in this section) patent license to make, have made, + use, offer to sell, sell, import, and otherwise transfer the Work, + where such license applies only to those patent claims licensable + by such Contributor that are necessarily infringed by their + Contribution(s) alone or by combination of their Contribution(s) + with the Work to which such Contribution(s) was submitted. If You + institute patent litigation against any entity (including a + cross-claim or counterclaim in a lawsuit) alleging that the Work + or a Contribution incorporated within the Work constitutes direct + or contributory patent infringement, then any patent licenses + granted to You under this License for that Work shall terminate + as of the date such litigation is filed. + + 4. Redistribution. You may reproduce and distribute copies of the + Work or Derivative Works thereof in any medium, with or without + modifications, and in Source or Object form, provided that You + meet the following conditions: + + (a) You must give any other recipients of the Work or + Derivative Works a copy of this License; and + + (b) You must cause any modified files to carry prominent notices + stating that You changed the files; and + + (c) You must retain, in the Source form of any Derivative Works + that You distribute, all copyright, patent, trademark, and + attribution notices from the Source form of the Work, + excluding those notices that do not pertain to any part of + the Derivative Works; and + + (d) If the Work includes a "NOTICE" text file as part of its + distribution, then any Derivative Works that You distribute must + include a readable copy of the attribution notices contained + within such NOTICE file, excluding those notices that do not + pertain to any part of the Derivative Works, in at least one + of the following places: within a NOTICE text file distributed + as part of the Derivative Works; within the Source form or + documentation, if provided along with the Derivative Works; or, + within a display generated by the Derivative Works, if and + wherever such third-party notices normally appear. The contents + of the NOTICE file are for informational purposes only and + do not modify the License. You may add Your own attribution + notices within Derivative Works that You distribute, alongside + or as an addendum to the NOTICE text from the Work, provided + that such additional attribution notices cannot be construed + as modifying the License. + + You may add Your own copyright statement to Your modifications and + may provide additional or different license terms and conditions + for use, reproduction, or distribution of Your modifications, or + for any such Derivative Works as a whole, provided Your use, + reproduction, and distribution of the Work otherwise complies with + the conditions stated in this License. + + 5. Submission of Contributions. Unless You explicitly state otherwise, + any Contribution intentionally submitted for inclusion in the Work + by You to the Licensor shall be under the terms and conditions of + this License, without any additional terms or conditions. + Notwithstanding the above, nothing herein shall supersede or modify + the terms of any separate license agreement you may have executed + with Licensor regarding such Contributions. + + 6. Trademarks. This License does not grant permission to use the trade + names, trademarks, service marks, or product names of the Licensor, + except as required for reasonable and customary use in describing the + origin of the Work and reproducing the content of the NOTICE file. + + 7. Disclaimer of Warranty. Unless required by applicable law or + agreed to in writing, Licensor provides the Work (and each + Contributor provides its Contributions) on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or + implied, including, without limitation, any warranties or conditions + of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A + PARTICULAR PURPOSE. You are solely responsible for determining the + appropriateness of using or redistributing the Work and assume any + risks associated with Your exercise of permissions under this License. + + 8. Limitation of Liability. In no event and under no legal theory, + whether in tort (including negligence), contract, or otherwise, + unless required by applicable law (such as deliberate and grossly + negligent acts) or agreed to in writing, shall any Contributor be + liable to You for damages, including any direct, indirect, special, + incidental, or consequential damages of any character arising as a + result of this License or out of the use or inability to use the + Work (including but not limited to damages for loss of goodwill, + work stoppage, computer failure or malfunction, or any and all + other commercial damages or losses), even if such Contributor + has been advised of the possibility of such damages. + + 9. Accepting Warranty or Additional Liability. While redistributing + the Work or Derivative Works thereof, You may choose to offer, + and charge a fee for, acceptance of support, warranty, indemnity, + or other liability obligations and/or rights consistent with this + License. However, in accepting such obligations, You may act only + on Your own behalf and on Your sole responsibility, not on behalf + of any other Contributor, and only if You agree to indemnify, + defend, and hold each Contributor harmless for any liability + incurred by, or claims asserted against, such Contributor by reason + of your accepting any such warranty or additional liability. + + END OF TERMS AND CONDITIONS + + APPENDIX: How to apply the Apache License to your work. + + To apply the Apache License to your work, attach the following + boilerplate notice, with the fields enclosed by brackets "[]" + replaced with your own identifying information. (Don't include + the brackets!) The text should be enclosed in the appropriate + comment syntax for the file format. We also recommend that a + file or class name and description of purpose be included on the + same "printed page" as the copyright notice for easier + identification within third-party archives. + + Copyright [yyyy] [name of copyright owner] + + 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. diff --git a/freevideo_engine/licenses/SAM2-THIRD_PARTY.txt b/freevideo_engine/licenses/SAM2-THIRD_PARTY.txt new file mode 100644 index 00000000..a24424de --- /dev/null +++ b/freevideo_engine/licenses/SAM2-THIRD_PARTY.txt @@ -0,0 +1,21 @@ +SAM 2.1 Hiera small inference code and checkpoint +Copyright (c) Meta Platforms, Inc. and affiliates. Apache License, Version 2.0. +Full license: SAM2-Apache-2.0.txt in this directory. +Source: https://github.com/facebookresearch/sam2 +Pinned revision: 2b90b9f5ceec907a1c18123530e92e794ad901a4 +Imported from the pinned region-sam/upstream/sam2 delivery. + +FreeVideo modifications (2026-10-18): inference modules are namespaced as +freevideo_engine.sam2; Hydra initialization and the hub-download constructor are +omitted; Hiera uses Python open and weights_only=True instead of iopath. The +fixed sam2.1_hiera_s configuration is constructed in image_gen/segment.py. +RoPE frequency construction stays on CPU and CUDA position-cache warmup is disabled. +No upstream inference mathematics or image transforms are changed. Automatic +masks, video predictors, training code, CUDA extension and setup scripts are +not included. Postprocessing holes/sprinkles is disabled, as in the reference. + +Checkpoint: facebook/sam2.1-hiera-small/sam2.1_hiera_small.pt +184416285 bytes +SHA256: 6d1aa6f30de5c92224f8172114de081d104bbd23dd9dc5c58996f0cad5dc4d38 +https://huggingface.co/facebook/sam2.1-hiera-small +The official SAM2 README licenses model checkpoints under Apache-2.0 as well. diff --git a/freevideo_engine/prompt_vlm/assets.py b/freevideo_engine/prompt_vlm/assets.py index 408906d0..a2f3177e 100644 --- a/freevideo_engine/prompt_vlm/assets.py +++ b/freevideo_engine/prompt_vlm/assets.py @@ -2,6 +2,7 @@ import json from pathlib import Path import shutil +import threading from .. import disk_space FAMILY = 'image-models' @@ -17,15 +18,107 @@ def directory(root): def ready(root): + """Cheap, safe on the launcher's UI thread: version, sizes and modified times only. + + Files whose modified time moved are checked by content once, in the + background; this returns False until that check has accepted them.""" + return _state(root, wait=False) + + +def settle(root): + """Same answer, but a moved modified time is checked by content in the + calling thread. For request handlers, never for a UI thread or under a lock.""" + return ready(root) or _state(root, wait=True) + + +_RECHECK = threading.Lock() +_BOOKKEEPING = threading.Lock() +_OUTCOME = {} +_RUNNING = set() + + +def _state(root, wait): + spec, folder = catalog(), directory(root) + try: + receipt = json.loads((folder / 'verified.json').read_text(encoding='utf-8')) + if receipt.get('revision') != spec['revision']: + return False + stats = {row['file']: (folder / row['file']).stat() for row in spec['files']} + if any(stats[row['file']].st_size != row['bytes'] for row in spec['files']): + return False + moved = [row for row in spec['files'] + if receipt['files'].get(row['file']) != stats[row['file']].st_mtime_ns] + if not moved: + return True + key = (str(folder), spec['revision'], + tuple((row['file'], stats[row['file']].st_size, stats[row['file']].st_mtime_ns) for row in moved)) + known = _OUTCOME.get(key) + if known is not None: + return known + if wait: + return _recheck(folder, spec, moved, key) + with _BOOKKEEPING: + if key in _RUNNING: + return False + _RUNNING.add(key) + threading.Thread(target=_recheck, args=(folder, spec, moved, key), daemon=True, + name='prompt-model-recheck').start() + return False + except (OSError, ValueError, KeyError, TypeError): + return False + + +def checking(root): + """True while the files' content is being checked after a modified time moved. + + A cheap look at the same file identity ready() uses, for a page that must not + call the model missing, or offer its download, in the meantime.""" spec, folder = catalog(), directory(root) try: receipt = json.loads((folder / 'verified.json').read_text(encoding='utf-8')) - return receipt.get('revision') == spec['revision'] and all( - (folder / row['file']).stat().st_size == row['bytes'] and - receipt['files'].get(row['file']) == (folder / row['file']).stat().st_mtime_ns - for row in spec['files']) + if receipt.get('revision') != spec['revision']: + return False + stats = {row['file']: (folder / row['file']).stat() for row in spec['files']} + if any(stats[row['file']].st_size != row['bytes'] for row in spec['files']): + return False + key = (str(folder), spec['revision'], tuple( + (row['file'], stats[row['file']].st_size, stats[row['file']].st_mtime_ns) + for row in spec['files'] if receipt['files'].get(row['file']) != stats[row['file']].st_mtime_ns)) except (OSError, ValueError, KeyError, TypeError): return False + with _BOOKKEEPING: + return key in _RUNNING and key not in _OUTCOME + + +def _recheck(folder, spec, moved, key): + """Accept files whose modified time changed (copied folder, unpacked archive, + antivirus) when their content still matches the catalog, then refresh the + receipt. The outcome is remembered per file identity, so one set of files is + hashed once whatever the answer.""" + from .. import provision + from ..monitoring import save + ledger_path = folder / 'download-verified.json' + try: + with _RECHECK: + if key in _OUTCOME: + return _OUTCOME[key] + try: + ledger = json.loads(ledger_path.read_text(encoding='utf-8')) + except (OSError, ValueError): + ledger = {} + ok = all(provision.verified(folder / row['file'], row, ledger) for row in moved) + if ok: + save(ledger_path, ledger) + save(folder / 'verified.json', dict(revision=spec['revision'], files={ + row['file']: (folder / row['file']).stat().st_mtime_ns for row in spec['files']})) + _OUTCOME[key] = ok + return ok + except Exception: + _OUTCOME[key] = False + return False + finally: + with _BOOKKEEPING: + _RUNNING.discard(key) def remote(spec, row): diff --git a/freevideo_engine/prompt_vlm/int8.py b/freevideo_engine/prompt_vlm/int8.py index 4da67f83..0ba3205a 100644 --- a/freevideo_engine/prompt_vlm/int8.py +++ b/freevideo_engine/prompt_vlm/int8.py @@ -38,6 +38,16 @@ def __init__(self, source, device): bias = source.bias self.register_buffer('bias', None if bias is None else bias.detach().to(device)) + def offload(self): + """Keep the int8 rows in (pinned, when possible) host memory; a Streamer lends them to the GPU.""" + host = self.weight_int8.cpu() + try: + host = host.pin_memory() + except RuntimeError: + pass # Pageable memory copies slower but works. + self.host = host + self.weight_int8 = None + def forward(self, x): shape = x.shape rows = x.reshape(-1, self.in_features) @@ -121,3 +131,87 @@ def install(layer, device): source = getattr(parent, name) setattr(parent, name, Int8Linear(source, device)) del source + + +def layer_bytes(stored): + """{layer index: GPU bytes of its int8 projections and their scales}.""" + import re + projection = re.compile(r'model\.language_model\.layers\.(\d+)\.(self_attn\.[qkvo]_proj|mlp\.(gate|up|down)_proj)\.weight$') + rows = {} + for name, row in stored.items(): + match = projection.match(name) + if match: + index = int(match.group(1)) + rows[index] = rows.get(index, 0) + row['bytes'] // 2 + 4 * row['shape'][0] + return rows + + +PREFETCH = 2 # Layers copied ahead of the one running. + + +def resident_layers(weights, stored, room): + """How many leading layers stay on the GPU as int8 when the rest are streamed + from host memory one layer at a time; None when even the fixed part and the + staging buffers do not fit `room`. Returns (layers, host_bytes).""" + sizes = layer_bytes(stored) + count = len(sizes) + if not count: + return None + projections = sum(row['bytes'] for name, row in stored.items() if name.startswith('model.language_model.layers.') + and name.endswith('_proj.weight')) + fixed = weights - projections + largest = max(sizes.values()) + free = room - fixed - (PREFETCH + 1) * largest + kept, used = 0, 0 + while kept < count and used + sizes[kept] <= free: + used += sizes[kept] + kept += 1 + if free < 0: + return None + return kept, sum(sizes[i] for i in range(kept, count)) + + +class Streamer: + """Lends host-held int8 rows to the GPU one layer ahead of use. + + The copies run on their own stream into pinned memory; a layer's + forward waits only for its own copy and drops the GPU rows afterwards. + """ + + def __init__(self, torch, layers): + self.torch = torch + self.layers = layers + self.linears = [[getattr(parent, name) for parent, name in projections(layer)] for layer in layers] + self.stream = torch.cuda.Stream() + self.pending = {} + for position, layer in enumerate(layers): + layer.register_forward_pre_hook(lambda module, args, kwargs, p=position: self.before(p), with_kwargs=True) + layer.register_forward_hook(lambda module, args, output, p=position: self.after(p)) + + def fetch(self, position): + if position in self.pending: + return + torch = self.torch + with torch.cuda.stream(self.stream): + tensors = [linear.host.to('cuda', non_blocking=True) for linear in self.linears[position]] + event = torch.cuda.Event() + event.record(self.stream) + self.pending[position] = (event, tensors) + + def before(self, position): + torch = self.torch + self.fetch(position) + event, tensors = self.pending.pop(position) + current = torch.cuda.current_stream() + current.wait_event(event) + for linear, tensor in zip(self.linears[position], tensors): + tensor.record_stream(current) + linear.weight_int8 = tensor + count = len(self.layers) + for ahead in range(1, PREFETCH + 1): + self.fetch((position + ahead) % count) # The next token starts at the first streamed layer. + return None + + def after(self, position): + for linear in self.linears[position]: + linear.weight_int8 = None diff --git a/freevideo_engine/prompt_vlm/load.py b/freevideo_engine/prompt_vlm/load.py index 8d299d0e..9119c961 100644 --- a/freevideo_engine/prompt_vlm/load.py +++ b/freevideo_engine/prompt_vlm/load.py @@ -36,8 +36,11 @@ def stored(folder): return rows -def resident(folder, torch, device, precision): - """Returns the model on `device`, with int8 projections if precision == 'int8'.""" +def resident(folder, torch, device, precision, keep=None): + """Returns the model on `device`, with int8 projections if precision == 'int8'. + + With `keep`, layers from that index on hold their int8 rows in host memory + and are streamed to the GPU as they run.""" from accelerate import init_empty_weights from torch import nn from transformers import AutoConfig, GenerationConfig, Qwen3VLForConditionalGeneration @@ -89,6 +92,9 @@ def each_layer(layer, merge): merge(layer) if precision == 'int8': int8.install(layer, device) + if keep is not None and layer.self_attn.layer_idx >= keep: + for parent, name in int8.projections(layer): + getattr(parent, name).offload() core.prepare(model, each_layer) finally: for reader in readers.values(): @@ -96,5 +102,7 @@ def each_layer(layer, merge): if tied: model.lm_head.weight = model.model.language_model.embed_tokens.weight model = model.to(device) + if keep is not None: + model.streamer = int8.Streamer(torch, list(model.model.language_model.layers)[keep:]) torch.cuda.empty_cache() # Merge and quantize temporaries; KV space comes next. return model diff --git a/freevideo_engine/prompt_vlm/service.py b/freevideo_engine/prompt_vlm/service.py index d82ff367..1a8acbff 100644 --- a/freevideo_engine/prompt_vlm/service.py +++ b/freevideo_engine/prompt_vlm/service.py @@ -217,10 +217,11 @@ def spawn(root, machine, env, descriptor=None): pass_fds=() if descriptor is None else (descriptor,), supervise=True, start_new_session=True) - def info(self): + def info(self, wait=True): + """wait=False answers from the cheap check, for callers on an event loop or UI thread.""" root, machine = self.installation() spec = assets.catalog() - return dict(ready=assets.ready(root), model=spec['repo'], license=spec['license'], + return dict(ready=assets.settle(root) if wait else assets.ready(root), model=spec['repo'], license=spec['license'], bytes=sum(row['bytes'] for row in spec['files']), busy=self.busy(), ram_bytes=self.ram_bytes(machine)) @@ -228,11 +229,12 @@ def ram_bytes(self, machine): """Available RAM the model usually needs on this computer, for the download card. The worker decides from free VRAM when it loads: 8 GB and larger NVIDIA - cards hold the whole model. Asks the driver once, without CUDA; anything + cards keep their layers on the GPU; 6 GB cards stream from RAM the layers + that do not fit, which needs more of it. Asks the driver once, without CUDA; anything unknown states the larger figure. """ if self.ram is None: - from .worker import MAPPED_RAM, RESIDENT_RAM + from .worker import MAPPED_RAM, RESIDENT_RAM, STREAMED_RAM self.ram = MAPPED_RAM if machine.get('device_backend', 'cuda') == 'cuda': from .. import processes @@ -245,6 +247,8 @@ def ram_bytes(self, machine): total = totals.get(machine.get('gpu_uuid')) or (len(totals) == 1 and next(iter(totals.values()))) if total and total >= 7.5 * 1024: self.ram = RESIDENT_RAM + elif total and total >= 5.5 * 1024: + self.ram = STREAMED_RAM except (OSError, RuntimeError, ValueError, subprocess.SubprocessError): pass return self.ram @@ -271,15 +275,15 @@ def start(self, action, value): if value.get('accept_download') is not True: raise ValueError('consent_required') elif action == 'ask': - if not assets.ready(root): + if not assets.settle(root): raise ValueError('model_missing') value = self.ask_images(rules.ask_request(value)) elif action == 'split': - if not assets.ready(root): + if not assets.settle(root): raise ValueError('model_missing') value = rules.split_request(value) else: - if not assets.ready(root): + if not assets.settle(root): raise ValueError('model_missing') value = rules.request(value) from ..comfy_assets import resolve_assets @@ -627,6 +631,7 @@ def stop_prestart_for_generation(server): def register(): + import asyncio from aiohttp import web from server import PromptServer import folder_paths @@ -641,7 +646,7 @@ def register(): @server.routes.get('/freevideo/prompt-vlm') async def info(request): try: - return web.json_response(dict(service.info(), token=token)) + return web.json_response(dict(await asyncio.to_thread(service.info), token=token)) except (OSError, ValueError): return web.json_response(dict(error='setup_required'), status=400) @@ -670,6 +675,10 @@ async def action(request): running, pending = server.prompt_queue.get_current_queue() if running or pending: raise ValueError('busy') + if name in ('ask', 'split', 'rewrite'): + # A moved modified time is checked by content here, off the event loop; + # start() then finds the remembered answer. + await asyncio.to_thread(assets.settle, service.installation()[0]) result = service.start(name, value) return web.json_response(result) except (ValueError, OSError) as error: diff --git a/freevideo_engine/prompt_vlm/worker.py b/freevideo_engine/prompt_vlm/worker.py index 34595365..3a2b8162 100644 --- a/freevideo_engine/prompt_vlm/worker.py +++ b/freevideo_engine/prompt_vlm/worker.py @@ -150,7 +150,7 @@ def failure(error, mode='rewrite'): Tokenizers may quote only a fragment of either draft: replacing the full prompt in an exception is not sufficient to make that exception private. """ - allowed = {'ram_space', 'image_size', 'context_length', 'invalid_output'} + allowed = {'ram_space', 'image_size', 'context_length', 'invalid_output', 'out_of_memory'} message = str(error) fallback = 'ask_failed' if mode == 'ask' else 'rewrite_failed' code = message if message in allowed else 'out_of_memory' if 'out of memory' in message.lower() else fallback @@ -205,6 +205,7 @@ def tokenize(processor, messages, max_pixels=512 * 32 * 32): # CUDA context alone take about 2.5 GB of commit). Other placements map and # merge the whole model in RAM; theirs is checked after the imports. RESIDENT_RAM = 4 * 2**30 +STREAMED_RAM = 6 * 2**30 # download-card figure for 6 GB cards: resident part plus the streamed layers' host copy MAPPED_RAM = 11 * 2**30 @@ -263,6 +264,8 @@ def load_model(folder, torch, stats, checkpoint): shapes = stored(folder) weights = sum(row['bytes'] for row in shapes.values()) resident = None + keep = None + host_bytes = 0 weight_budget = 0 if torch.cuda.is_available(): from .. import gpu_budget @@ -277,7 +280,19 @@ def load_model(folder, torch, stats, checkpoint): kv = kv_bytes(shapes, reserved_tokens(stats)) stats.update(gpu_allocator_budget_bytes=free, gpu_weight_budget_bytes=weight_budget, gpu_kv_reserve_bytes=kv) resident = placement(weights, int8.saved_bytes(shapes), kv, free) - stats['ram_required_bytes'] = RESIDENT_RAM if resident else MAPPED_RAM + if not resident: + # Too big even as int8: keep what fits on the GPU and stream the other + # layers' int8 rows from host memory, a layer ahead of use. + plan = int8.resident_layers(weights, shapes, free - kv - WORKING_VRAM) + if plan: + resident, (keep, host_bytes) = 'int8', plan + stats.update(streamed_layers=len(int8.layer_bytes(shapes)) - keep if keep is not None else 0, + streamed_host_bytes=host_bytes) + if not resident and weight_budget <= 2 * 2**30: + # Not even the streaming tier fits: the model would run on the CPU at + # about 0.2 tokens/s and hit the request deadline, so say so at once. + raise ValueError('out_of_memory') + stats['ram_required_bytes'] = (RESIDENT_RAM + host_bytes) if resident else MAPPED_RAM available = stats['available_memory_at_start_bytes' if resident else 'available_memory_before_load_bytes'] if available < stats['ram_required_bytes']: raise ValueError('ram_space') @@ -285,7 +300,7 @@ def load_model(folder, torch, stats, checkpoint): loading = time.monotonic() if resident: from .load import resident as load_resident - model = load_resident(folder, torch, 'cuda:0', resident) + model = load_resident(folder, torch, 'cuda:0', resident, keep) mapping, device = {'': 0}, 'cuda:0' stats.update(weights=resident, kv_positions=kv_positions(reserved_tokens(stats))) else: diff --git a/freevideo_engine/release_notes.json b/freevideo_engine/release_notes.json index 65ad725f..66aca191 100644 --- a/freevideo_engine/release_notes.json +++ b/freevideo_engine/release_notes.json @@ -1,39 +1,72 @@ { "schema": 1, - "product_version": "0.4.13", + "product_version": "0.4.14", "en": { - "summary": "High is faster than Max again, videos at another size of the same shape can be continued, LoRAs on int8 models apply what they can, Settings adds step-by-step how-tos and an offline tutorial, and the installers are smaller.", + "summary": "Max is faster, a picture can now be changed only inside a region you select, agents can import their own clips into the timeline, and long videos can take more requests into the queue while they generate.", "changes": [ - "High: the step cache is back with added checks, so High runs faster than Max again.", - "Long videos: parts of a long video and continued parts use 8 + 3 steps, so a second part is no longer refused.", - "Continue: a video at another size of the same shape now continues from its own ending. Videos that are too short or have too low a frame rate show a clear message.", - "LoRA: on int8 models, the parts of a LoRA that cannot be applied are skipped and the rest takes effect, instead of the whole LoRA being refused.", - "Help: \"User guide\" in Settings shows three how-tos on the real interface, and the new \"Tutorial\" opens an offline illustrated manual in English and Chinese.", - "Upgrade: a FreeVideo tab left open from an older version switches to the current workflow and keeps the prompt, size, length, seed, assets and LoRAs.", - "Agent: assembled videos are saved to \"Your creations\", and their task cards can be opened.", - "Install: the Windows and Linux installers are smaller, the prompt enhancement model can also download from ModelScope, and the GPU acceleration step shows one steady name and description while it builds." + "Quality levels: Max now uses 20 steps with a guarded step cache and finishes sooner; High uses 16 steps with the step cache. The level names stay the same.", + "Picture editing: when changing a picture, drag a box, click an object, or draw a closed outline to select a region. Shift adds to the selection and Alt removes from it. Only the selected region changes, and everything outside it stays exactly the same. Click selection downloads a 184 MB model from ModelScope the first time it is used.", + "Studio: while a long or multi-part video is generating, the button stays available as \"Add to queue\". Up to five requests can wait, and each starts when the current video finishes or is stopped.", + "Agent: agents can import clips they rendered through HTTP, MCP and the command line, place them on the timeline by frame, and export. Before an agent is connected, the autopilot page says what connecting one is for.", + "Editing view: a reference picture that cannot be read is flagged right away so another one can be chosen.", + "Upgrade: when an old FreeVideo tab is replaced, two-pass sampling, first- and second-pass steps, the step cache and the upscale factor are carried over.", + "Prompt enhancement: after the model folder is copied, an offline package is unpacked, or antivirus software scans the files, FreeVideo no longer asks to download 8.9 GB again when the files are unchanged. It now also works on 6 GB graphics cards, finishing in under a minute instead of timing out.", + "Chinese interface: browser translation no longer changes FreeVideo's own labels, so the size menu keeps showing MP.", + "Feedback: feedback sent for a failed generation now includes the engine's error message.", + "Install: the Windows library checks also handle files in deeply nested folders." ] }, "zh": { - "summary": "精细档重新比极致档快;同比例、不同尺寸的视频可以接着生成;int8 模型上的 LoRA 能用的部分照常生效;设置中新增分步演示和离线教程;安装包变小。", + "summary": "极致档更快;改图时可以只改选中的区域;Agent 可以把自己渲染的片段导入时间线;长视频生成中可以继续加入队列。", "changes": [ - "精细档:重新启用步缓存,并加上检查,精细档重新比极致档快。", - "长视频:长视频的分段和接着生成的分段统一使用 8 + 3 步,第二段不再被拒绝。", - "接着生成:同比例、不同尺寸的视频改为从它自己的结尾接着生成;太短或帧率太低的视频会给出明确提示。", - "LoRA:在 int8 模型上,LoRA 中无法应用的部分会跳过,其余照常生效,不再整个拒绝。", - "帮助:设置中的“使用说明”改为在真实界面上演示三项操作;新增“教程”,可打开中英文离线图文手册。", - "升级:旧版本留在标签页中的 FreeVideo 工作流会换成当前版本,并保留提示词、尺寸、时长、种子、素材和 LoRA。", - "Agent:拼接的成片保存到“我的作品”,任务卡可以打开。", - "安装:Windows 和 Linux 安装包变小;提示词增强模型也可以从魔搭下载;GPU 加速组件构建时,名称和说明保持不变。" + "质量档位:极致档改为 20 步加带检查的步缓存,更快完成;精细档改为 16 步加步缓存。档位名称不变。", + "改图:改一张图时,可以拖框、点选物体或画闭合曲线来选区;按住 Shift 添加,按住 Alt 减去。只改选区里的内容,选区外完全不变。点选第一次使用时从魔搭下载 184 MB 的模型。", + "创作面板:长视频或分段视频生成中,按钮保持可用,显示“加入队列”。最多排 5 条,当前视频完成或停止后自动开始下一条。", + "Agent:Agent 可以通过 HTTP、MCP 和命令行导入自己渲染的片段,按帧放上时间线并导出。未连接时,自动页说明连上之后能做什么。", + "剪辑视图:参考图读不出来时会当场提示,可以换一张。", + "升级:替换旧版本的 FreeVideo 标签页时,二次采样、第一遍和第二遍步数、步缓存和放大倍数都会带过来。", + "提示词增强:复制模型文件夹、解压离线包或被杀毒软件扫描之后,只要文件内容没变,就不再要求重新下载 8.9 GB。6 GB 显卡上现在也能用,一分钟内完成,不再超时。", + "中文界面:浏览器翻译不再改动 FreeVideo 自己的文字,尺寸菜单保持显示 MP。", + "反馈:为失败的生成发送反馈时,会附上引擎给出的报错信息。", + "安装:Windows 上的库文件检查也支持层级很深的文件夹。" ] }, "history": [ + { + "product_version": "0.4.13", + "en": { + "summary": "High is faster than Max again, videos at another size of the same shape can be continued, LoRAs on int8 models apply what they can, Settings adds step-by-step how-tos and an offline tutorial, and the installers are smaller.", + "changes": [ + "High: the step cache is back with added checks, so High runs faster than Max again.", + "Long videos: parts of a long video and continued parts use 8 + 3 steps, so a second part is no longer refused.", + "Continue: a video at another size of the same shape now continues from its own ending. Videos that are too short or have too low a frame rate show a clear message.", + "LoRA: on int8 models, the parts of a LoRA that cannot be applied are skipped and the rest takes effect, instead of the whole LoRA being refused.", + "Help: \"User guide\" in Settings shows three how-tos on the real interface, and the new \"Tutorial\" opens an offline illustrated manual in English and Chinese.", + "Upgrade: a FreeVideo tab left open from an older version switches to the current workflow and keeps the prompt, size, length, seed, assets and LoRAs.", + "Agent: assembled videos are saved to \"Your creations\", and their task cards can be opened.", + "Install: the Windows and Linux installers are smaller, the prompt enhancement model can also download from ModelScope, and the GPU acceleration step shows one steady name and description while it builds." + ] + }, + "zh": { + "summary": "精细档重新比极致档快;同比例、不同尺寸的视频可以接着生成;int8 模型上的 LoRA 能用的部分照常生效;设置中新增分步演示和离线教程;安装包变小。", + "changes": [ + "精细档:重新启用步缓存,并加上检查,精细档重新比极致档快。", + "长视频:长视频的分段和接着生成的分段统一使用 8 + 3 步,第二段不再被拒绝。", + "接着生成:同比例、不同尺寸的视频改为从它自己的结尾接着生成;太短或帧率太低的视频会给出明确提示。", + "LoRA:在 int8 模型上,LoRA 中无法应用的部分会跳过,其余照常生效,不再整个拒绝。", + "帮助:设置中的“使用说明”改为在真实界面上演示三项操作;新增“教程”,可打开中英文离线图文手册。", + "升级:旧版本留在标签页中的 FreeVideo 工作流会换成当前版本,并保留提示词、尺寸、时长、种子、素材和 LoRA。", + "Agent:拼接的成片保存到“我的作品”,任务卡可以打开。", + "安装:Windows 和 Linux 安装包变小;提示词增强模型也可以从魔搭下载;GPU 加速组件构建时,名称和说明保持不变。" + ] + } + }, { "product_version": "0.4.12", "en": { "summary": "Preview white-model scenes for Autopilot, LoRA compatibility checks and drag-to-reorder, clearer failure messages, deleting edits, a faster start for a busy machine, and a Medium level that samples two passes.", "changes": [ - "Autopilot: an optional white-model preview, with ready-made scenes, shows the scene before the video is made; the entry showcase has an official race-car example.", + "Autopilot: an optional white-model preview, with ready-made scenes, shows the scene before the video is made; the entry showcase has an example.", "LoRA: a compatibility check before generating, and LoRAs can be reordered by dragging.", "Failure messages: submission, queue and browser-extension failures share one layout and say what to do next; failed requests can be packaged with Export report or Submit to GitHub.", "Editor: an edit can be deleted, an empty untouched edit is discarded, and reopening an edit restores its cover and shows its finishing stage; continuations are queued on the server.", @@ -50,7 +83,7 @@ "zh": { "summary": "AI 全自动新增白模预览场景,LoRA 兼容性检查和拖动排序,失败提示更清楚,可以删除剪辑,繁忙的机器上启动更快,标准档改为两遍采样。", "changes": [ - "AI 全自动:可选的白模预览和现成场景,在生成视频前先看到场景;入口展示里有官方赛车示例。", + "AI 全自动:可选的白模预览和现成场景,在生成视频前先看到场景;入口展示里有示例。", "LoRA:生成前做兼容性检查,LoRA 可以拖动排序。", "失败提示:提交、排队和浏览器扩展的失败用同一种版式,并说明下一步该做什么;失败的请求可以打包为导出报告或提交到 GitHub。", "编辑器:可以删除剪辑,没改动的空剪辑自动丢弃;重新打开剪辑会恢复封面并显示收尾阶段;续接在服务端排队。", diff --git a/freevideo_engine/request_assets.py b/freevideo_engine/request_assets.py index 82e8ddbd..85136b13 100644 --- a/freevideo_engine/request_assets.py +++ b/freevideo_engine/request_assets.py @@ -59,7 +59,7 @@ def _upscaler(root, sampled): def _prompt_vlm(root): from .prompt_vlm import assets - if assets.ready(root): + if assets.settle(root): return [] spec = assets.catalog() return [dict(id='prompt-vlm', kind='prompt-vlm', files=len(spec['files']), diff --git a/freevideo_engine/runtime_libraries.py b/freevideo_engine/runtime_libraries.py index 19571ea1..2fc9e804 100644 --- a/freevideo_engine/runtime_libraries.py +++ b/freevideo_engine/runtime_libraries.py @@ -57,7 +57,7 @@ def site_packages(environment): def _digest(path, algorithm): digest = hashlib.new(algorithm) - with open(path, 'rb') as handle: + with open(io_path(path), 'rb') as handle: for chunk in iter(lambda: handle.read(1 << 20), b''): digest.update(chunk) return algorithm + '=' + base64.urlsafe_b64encode(digest.digest()).rstrip(b'=').decode('ascii') @@ -83,7 +83,7 @@ def check(site, progress=None): root = os.path.normcase(os.path.normpath(str(site))) result = dict(owners=[], checked=0, changed=[], unrecorded=[], unreadable=[], directories=[]) expected = {} - for distribution in importlib.metadata.distributions(path=[str(site)]): + for distribution in importlib.metadata.distributions(path=[str(io_path(site))]): try: rows = list(csv.reader(io.StringIO(distribution.read_text('RECORD') or ''))) name = distribution.metadata['Name'] @@ -108,10 +108,10 @@ def check(site, progress=None): directories.add(os.path.dirname(path)) result['checked'] += 1 try: - if not os.path.isfile(path): + if not os.path.isfile(io_path(path)): result['changed'].append(dict(entry, state='missing')) continue - size = os.path.getsize(path) + size = os.path.getsize(io_path(path)) matching = [(value, length) for _, value, length in records if length is None or length == size] hashed = {} def matches(value): @@ -133,7 +133,7 @@ def matches(value): recorded = {_canonical(path) for path, _ in expected.values()} for directory in sorted(directories): try: - names = sorted(os.listdir(directory)) + names = sorted(os.listdir(io_path(directory))) except OSError: continue for name in names: @@ -266,7 +266,7 @@ def dependency_closure(python, requirements, additional=()): dependencies = {} site = site_for(python) if site is not None: - for distribution in importlib.metadata.distributions(path=[str(site)]): + for distribution in importlib.metadata.distributions(path=[str(io_path(site))]): try: name = distribution.metadata['Name'] if name: @@ -367,7 +367,7 @@ def installed_bytes(python, name): """Installed size of one distribution in this Python's environment, from its RECORD.""" try: site = site_for(python) - for distribution in importlib.metadata.distributions(path=[str(site)] if site else []): + for distribution in importlib.metadata.distributions(path=[str(io_path(site))] if site else []): if (distribution.metadata['Name'] or '').lower() == name.lower(): rows = csv.reader(io.StringIO(distribution.read_text('RECORD') or '')) return sum(int(row[2]) for row in rows if len(row) > 2 and row[2].isdigit()) diff --git a/freevideo_engine/sam2/__init__.py b/freevideo_engine/sam2/__init__.py new file mode 100644 index 00000000..c5c96c86 --- /dev/null +++ b/freevideo_engine/sam2/__init__.py @@ -0,0 +1,8 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. + +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. + + +# FreeVideo: fixed CPU image model; no Hydra initialization or hub downloads. diff --git a/freevideo_engine/sam2/modeling/__init__.py b/freevideo_engine/sam2/modeling/__init__.py new file mode 100644 index 00000000..5277f461 --- /dev/null +++ b/freevideo_engine/sam2/modeling/__init__.py @@ -0,0 +1,5 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. + +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. diff --git a/freevideo_engine/sam2/modeling/backbones/__init__.py b/freevideo_engine/sam2/modeling/backbones/__init__.py new file mode 100644 index 00000000..5277f461 --- /dev/null +++ b/freevideo_engine/sam2/modeling/backbones/__init__.py @@ -0,0 +1,5 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. + +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. diff --git a/freevideo_engine/sam2/modeling/backbones/hieradet.py b/freevideo_engine/sam2/modeling/backbones/hieradet.py new file mode 100644 index 00000000..a299e52d --- /dev/null +++ b/freevideo_engine/sam2/modeling/backbones/hieradet.py @@ -0,0 +1,316 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. + +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. + +import logging +from functools import partial +from typing import List, Tuple, Union + +import torch +import torch.nn as nn +import torch.nn.functional as F + +from freevideo_engine.sam2.modeling.backbones.utils import ( + PatchEmbed, + window_partition, + window_unpartition, +) + +from freevideo_engine.sam2.modeling.sam2_utils import DropPath, MLP + + +def do_pool(x: torch.Tensor, pool: nn.Module, norm: nn.Module = None) -> torch.Tensor: + if pool is None: + return x + # (B, H, W, C) -> (B, C, H, W) + x = x.permute(0, 3, 1, 2) + x = pool(x) + # (B, C, H', W') -> (B, H', W', C) + x = x.permute(0, 2, 3, 1) + if norm: + x = norm(x) + + return x + + +class MultiScaleAttention(nn.Module): + def __init__( + self, + dim: int, + dim_out: int, + num_heads: int, + q_pool: nn.Module = None, + ): + super().__init__() + + self.dim = dim + self.dim_out = dim_out + self.num_heads = num_heads + self.q_pool = q_pool + self.qkv = nn.Linear(dim, dim_out * 3) + self.proj = nn.Linear(dim_out, dim_out) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + B, H, W, _ = x.shape + # qkv with shape (B, H * W, 3, nHead, C) + qkv = self.qkv(x).reshape(B, H * W, 3, self.num_heads, -1) + # q, k, v with shape (B, H * W, nheads, C) + q, k, v = torch.unbind(qkv, 2) + + # Q pooling (for downsample at stage changes) + if self.q_pool: + q = do_pool(q.reshape(B, H, W, -1), self.q_pool) + H, W = q.shape[1:3] # downsampled shape + q = q.reshape(B, H * W, self.num_heads, -1) + + # Torch's SDPA expects [B, nheads, H*W, C] so we transpose + x = F.scaled_dot_product_attention( + q.transpose(1, 2), + k.transpose(1, 2), + v.transpose(1, 2), + ) + # Transpose back + x = x.transpose(1, 2) + x = x.reshape(B, H, W, -1) + + x = self.proj(x) + + return x + + +class MultiScaleBlock(nn.Module): + def __init__( + self, + dim: int, + dim_out: int, + num_heads: int, + mlp_ratio: float = 4.0, + drop_path: float = 0.0, + norm_layer: Union[nn.Module, str] = "LayerNorm", + q_stride: Tuple[int, int] = None, + act_layer: nn.Module = nn.GELU, + window_size: int = 0, + ): + super().__init__() + + if isinstance(norm_layer, str): + norm_layer = partial(getattr(nn, norm_layer), eps=1e-6) + + self.dim = dim + self.dim_out = dim_out + self.norm1 = norm_layer(dim) + + self.window_size = window_size + + self.pool, self.q_stride = None, q_stride + if self.q_stride: + self.pool = nn.MaxPool2d( + kernel_size=q_stride, stride=q_stride, ceil_mode=False + ) + + self.attn = MultiScaleAttention( + dim, + dim_out, + num_heads=num_heads, + q_pool=self.pool, + ) + self.drop_path = DropPath(drop_path) if drop_path > 0.0 else nn.Identity() + + self.norm2 = norm_layer(dim_out) + self.mlp = MLP( + dim_out, + int(dim_out * mlp_ratio), + dim_out, + num_layers=2, + activation=act_layer, + ) + + if dim != dim_out: + self.proj = nn.Linear(dim, dim_out) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + shortcut = x # B, H, W, C + x = self.norm1(x) + + # Skip connection + if self.dim != self.dim_out: + shortcut = do_pool(self.proj(x), self.pool) + + # Window partition + window_size = self.window_size + if window_size > 0: + H, W = x.shape[1], x.shape[2] + x, pad_hw = window_partition(x, window_size) + + # Window Attention + Q Pooling (if stage change) + x = self.attn(x) + if self.q_stride: + # Shapes have changed due to Q pooling + window_size = self.window_size // self.q_stride[0] + H, W = shortcut.shape[1:3] + + pad_h = (window_size - H % window_size) % window_size + pad_w = (window_size - W % window_size) % window_size + pad_hw = (H + pad_h, W + pad_w) + + # Reverse window partition + if self.window_size > 0: + x = window_unpartition(x, window_size, pad_hw, (H, W)) + + x = shortcut + self.drop_path(x) + # MLP + x = x + self.drop_path(self.mlp(self.norm2(x))) + return x + + +class Hiera(nn.Module): + """ + Reference: https://arxiv.org/abs/2306.00989 + """ + + def __init__( + self, + embed_dim: int = 96, # initial embed dim + num_heads: int = 1, # initial number of heads + drop_path_rate: float = 0.0, # stochastic depth + q_pool: int = 3, # number of q_pool stages + q_stride: Tuple[int, int] = (2, 2), # downsample stride bet. stages + stages: Tuple[int, ...] = (2, 3, 16, 3), # blocks per stage + dim_mul: float = 2.0, # dim_mul factor at stage shift + head_mul: float = 2.0, # head_mul factor at stage shift + window_pos_embed_bkg_spatial_size: Tuple[int, int] = (14, 14), + # window size per stage, when not using global att. + window_spec: Tuple[int, ...] = ( + 8, + 4, + 14, + 7, + ), + # global attn in these blocks + global_att_blocks: Tuple[int, ...] = ( + 12, + 16, + 20, + ), + weights_path=None, + return_interm_layers=True, # return feats from every stage + ): + super().__init__() + + assert len(stages) == len(window_spec) + self.window_spec = window_spec + + depth = sum(stages) + self.q_stride = q_stride + self.stage_ends = [sum(stages[:i]) - 1 for i in range(1, len(stages) + 1)] + assert 0 <= q_pool <= len(self.stage_ends[:-1]) + self.q_pool_blocks = [x + 1 for x in self.stage_ends[:-1]][:q_pool] + self.return_interm_layers = return_interm_layers + + self.patch_embed = PatchEmbed( + embed_dim=embed_dim, + ) + # Which blocks have global att? + self.global_att_blocks = global_att_blocks + + # Windowed positional embedding (https://arxiv.org/abs/2311.05613) + self.window_pos_embed_bkg_spatial_size = window_pos_embed_bkg_spatial_size + self.pos_embed = nn.Parameter( + torch.zeros(1, embed_dim, *self.window_pos_embed_bkg_spatial_size) + ) + self.pos_embed_window = nn.Parameter( + torch.zeros(1, embed_dim, self.window_spec[0], self.window_spec[0]) + ) + + dpr = [ + x.item() for x in torch.linspace(0, drop_path_rate, depth) + ] # stochastic depth decay rule + + cur_stage = 1 + self.blocks = nn.ModuleList() + + for i in range(depth): + dim_out = embed_dim + # lags by a block, so first block of + # next stage uses an initial window size + # of previous stage and final window size of current stage + window_size = self.window_spec[cur_stage - 1] + + if self.global_att_blocks is not None: + window_size = 0 if i in self.global_att_blocks else window_size + + if i - 1 in self.stage_ends: + dim_out = int(embed_dim * dim_mul) + num_heads = int(num_heads * head_mul) + cur_stage += 1 + + block = MultiScaleBlock( + dim=embed_dim, + dim_out=dim_out, + num_heads=num_heads, + drop_path=dpr[i], + q_stride=self.q_stride if i in self.q_pool_blocks else None, + window_size=window_size, + ) + + embed_dim = dim_out + self.blocks.append(block) + + self.channel_list = ( + [self.blocks[i].dim_out for i in self.stage_ends[::-1]] + if return_interm_layers + else [self.blocks[-1].dim_out] + ) + + if weights_path is not None: + with open(weights_path, "rb") as f: + chkpt = torch.load(f, map_location="cpu", weights_only=True) + logging.info("loading Hiera", self.load_state_dict(chkpt, strict=False)) + + def _get_pos_embed(self, hw: Tuple[int, int]) -> torch.Tensor: + h, w = hw + window_embed = self.pos_embed_window + pos_embed = F.interpolate(self.pos_embed, size=(h, w), mode="bicubic") + pos_embed = pos_embed + window_embed.tile( + [x // y for x, y in zip(pos_embed.shape, window_embed.shape)] + ) + pos_embed = pos_embed.permute(0, 2, 3, 1) + return pos_embed + + def forward(self, x: torch.Tensor) -> List[torch.Tensor]: + x = self.patch_embed(x) + # x: (B, H, W, C) + + # Add pos embed + x = x + self._get_pos_embed(x.shape[1:3]) + + outputs = [] + for i, blk in enumerate(self.blocks): + x = blk(x) + if (i == self.stage_ends[-1]) or ( + i in self.stage_ends and self.return_interm_layers + ): + feats = x.permute(0, 3, 1, 2) + outputs.append(feats) + + return outputs + + def get_layer_id(self, layer_name): + # https://github.com/microsoft/unilm/blob/master/beit/optim_factory.py#L33 + num_layers = self.get_num_layers() + + if layer_name.find("rel_pos") != -1: + return num_layers + 1 + elif layer_name.find("pos_embed") != -1: + return 0 + elif layer_name.find("patch_embed") != -1: + return 0 + elif layer_name.find("blocks") != -1: + return int(layer_name.split("blocks")[1].split(".")[1]) + 1 + else: + return num_layers + 1 + + def get_num_layers(self) -> int: + return len(self.blocks) diff --git a/freevideo_engine/sam2/modeling/backbones/image_encoder.py b/freevideo_engine/sam2/modeling/backbones/image_encoder.py new file mode 100644 index 00000000..37e9266b --- /dev/null +++ b/freevideo_engine/sam2/modeling/backbones/image_encoder.py @@ -0,0 +1,134 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. + +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. + +from typing import List, Optional + +import torch +import torch.nn as nn +import torch.nn.functional as F + + +class ImageEncoder(nn.Module): + def __init__( + self, + trunk: nn.Module, + neck: nn.Module, + scalp: int = 0, + ): + super().__init__() + self.trunk = trunk + self.neck = neck + self.scalp = scalp + assert ( + self.trunk.channel_list == self.neck.backbone_channel_list + ), f"Channel dims of trunk and neck do not match. Trunk: {self.trunk.channel_list}, neck: {self.neck.backbone_channel_list}" + + def forward(self, sample: torch.Tensor): + # Forward through backbone + features, pos = self.neck(self.trunk(sample)) + if self.scalp > 0: + # Discard the lowest resolution features + features, pos = features[: -self.scalp], pos[: -self.scalp] + + src = features[-1] + output = { + "vision_features": src, + "vision_pos_enc": pos, + "backbone_fpn": features, + } + return output + + +class FpnNeck(nn.Module): + """ + A modified variant of Feature Pyramid Network (FPN) neck + (we remove output conv and also do bicubic interpolation similar to ViT + pos embed interpolation) + """ + + def __init__( + self, + position_encoding: nn.Module, + d_model: int, + backbone_channel_list: List[int], + kernel_size: int = 1, + stride: int = 1, + padding: int = 0, + fpn_interp_model: str = "bilinear", + fuse_type: str = "sum", + fpn_top_down_levels: Optional[List[int]] = None, + ): + """Initialize the neck + :param trunk: the backbone + :param position_encoding: the positional encoding to use + :param d_model: the dimension of the model + :param neck_norm: the normalization to use + """ + super().__init__() + self.position_encoding = position_encoding + self.convs = nn.ModuleList() + self.backbone_channel_list = backbone_channel_list + self.d_model = d_model + for dim in backbone_channel_list: + current = nn.Sequential() + current.add_module( + "conv", + nn.Conv2d( + in_channels=dim, + out_channels=d_model, + kernel_size=kernel_size, + stride=stride, + padding=padding, + ), + ) + + self.convs.append(current) + self.fpn_interp_model = fpn_interp_model + assert fuse_type in ["sum", "avg"] + self.fuse_type = fuse_type + + # levels to have top-down features in its outputs + # e.g. if fpn_top_down_levels is [2, 3], then only outputs of level 2 and 3 + # have top-down propagation, while outputs of level 0 and level 1 have only + # lateral features from the same backbone level. + if fpn_top_down_levels is None: + # default is to have top-down features on all levels + fpn_top_down_levels = range(len(self.convs)) + self.fpn_top_down_levels = list(fpn_top_down_levels) + + def forward(self, xs: List[torch.Tensor]): + + out = [None] * len(self.convs) + pos = [None] * len(self.convs) + assert len(xs) == len(self.convs) + # fpn forward pass + # see https://github.com/facebookresearch/detectron2/blob/main/detectron2/modeling/backbone/fpn.py + prev_features = None + # forward in top-down order (from low to high resolution) + n = len(self.convs) - 1 + for i in range(n, -1, -1): + x = xs[i] + lateral_features = self.convs[n - i](x) + if i in self.fpn_top_down_levels and prev_features is not None: + top_down_features = F.interpolate( + prev_features.to(dtype=torch.float32), + scale_factor=2.0, + mode=self.fpn_interp_model, + align_corners=( + None if self.fpn_interp_model == "nearest" else False + ), + antialias=False, + ) + prev_features = lateral_features + top_down_features + if self.fuse_type == "avg": + prev_features /= 2 + else: + prev_features = lateral_features + x_out = prev_features + out[i] = x_out + pos[i] = self.position_encoding(x_out).to(x_out.dtype) + + return out, pos diff --git a/freevideo_engine/sam2/modeling/backbones/utils.py b/freevideo_engine/sam2/modeling/backbones/utils.py new file mode 100644 index 00000000..930b1b76 --- /dev/null +++ b/freevideo_engine/sam2/modeling/backbones/utils.py @@ -0,0 +1,93 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. + +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. + +"""Some utilities for backbones, in particular for windowing""" + +from typing import Tuple + +import torch +import torch.nn as nn +import torch.nn.functional as F + + +def window_partition(x, window_size): + """ + Partition into non-overlapping windows with padding if needed. + Args: + x (tensor): input tokens with [B, H, W, C]. + window_size (int): window size. + Returns: + windows: windows after partition with [B * num_windows, window_size, window_size, C]. + (Hp, Wp): padded height and width before partition + """ + B, H, W, C = x.shape + + pad_h = (window_size - H % window_size) % window_size + pad_w = (window_size - W % window_size) % window_size + if pad_h > 0 or pad_w > 0: + x = F.pad(x, (0, 0, 0, pad_w, 0, pad_h)) + Hp, Wp = H + pad_h, W + pad_w + + x = x.view(B, Hp // window_size, window_size, Wp // window_size, window_size, C) + windows = x.permute(0, 1, 3, 2, 4, 5).reshape(-1, window_size, window_size, C) + return windows, (Hp, Wp) + + +def window_unpartition(windows, window_size, pad_hw, hw): + """ + Window unpartition into original sequences and removing padding. + Args: + x (tensor): input tokens with [B * num_windows, window_size, window_size, C]. + window_size (int): window size. + pad_hw (Tuple): padded height and width (Hp, Wp). + hw (Tuple): original height and width (H, W) before padding. + Returns: + x: unpartitioned sequences with [B, H, W, C]. + """ + Hp, Wp = pad_hw + H, W = hw + B = windows.shape[0] // (Hp * Wp // window_size // window_size) + x = windows.reshape( + B, Hp // window_size, Wp // window_size, window_size, window_size, -1 + ) + x = x.permute(0, 1, 3, 2, 4, 5).reshape(B, Hp, Wp, -1) + + if Hp > H or Wp > W: + x = x[:, :H, :W, :] + return x + + +class PatchEmbed(nn.Module): + """ + Image to Patch Embedding. + """ + + def __init__( + self, + kernel_size: Tuple[int, ...] = (7, 7), + stride: Tuple[int, ...] = (4, 4), + padding: Tuple[int, ...] = (3, 3), + in_chans: int = 3, + embed_dim: int = 768, + ): + """ + Args: + kernel_size (Tuple): kernel size of the projection layer. + stride (Tuple): stride of the projection layer. + padding (Tuple): padding size of the projection layer. + in_chans (int): Number of input image channels. + embed_dim (int): embed_dim (int): Patch embedding dimension. + """ + super().__init__() + self.proj = nn.Conv2d( + in_chans, embed_dim, kernel_size=kernel_size, stride=stride, padding=padding + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + x = self.proj(x) + # B C H W -> B H W C + x = x.permute(0, 2, 3, 1) + return x diff --git a/freevideo_engine/sam2/modeling/memory_attention.py b/freevideo_engine/sam2/modeling/memory_attention.py new file mode 100644 index 00000000..f6408fbc --- /dev/null +++ b/freevideo_engine/sam2/modeling/memory_attention.py @@ -0,0 +1,169 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. + +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. + +from typing import Optional + +import torch +from torch import nn, Tensor + +from freevideo_engine.sam2.modeling.sam.transformer import RoPEAttention + +from freevideo_engine.sam2.modeling.sam2_utils import get_activation_fn, get_clones + + +class MemoryAttentionLayer(nn.Module): + + def __init__( + self, + activation: str, + cross_attention: nn.Module, + d_model: int, + dim_feedforward: int, + dropout: float, + pos_enc_at_attn: bool, + pos_enc_at_cross_attn_keys: bool, + pos_enc_at_cross_attn_queries: bool, + self_attention: nn.Module, + ): + super().__init__() + self.d_model = d_model + self.dim_feedforward = dim_feedforward + self.dropout_value = dropout + self.self_attn = self_attention + self.cross_attn_image = cross_attention + + # Implementation of Feedforward model + self.linear1 = nn.Linear(d_model, dim_feedforward) + self.dropout = nn.Dropout(dropout) + self.linear2 = nn.Linear(dim_feedforward, d_model) + + self.norm1 = nn.LayerNorm(d_model) + self.norm2 = nn.LayerNorm(d_model) + self.norm3 = nn.LayerNorm(d_model) + self.dropout1 = nn.Dropout(dropout) + self.dropout2 = nn.Dropout(dropout) + self.dropout3 = nn.Dropout(dropout) + + self.activation_str = activation + self.activation = get_activation_fn(activation) + + # Where to add pos enc + self.pos_enc_at_attn = pos_enc_at_attn + self.pos_enc_at_cross_attn_queries = pos_enc_at_cross_attn_queries + self.pos_enc_at_cross_attn_keys = pos_enc_at_cross_attn_keys + + def _forward_sa(self, tgt, query_pos): + # Self-Attention + tgt2 = self.norm1(tgt) + q = k = tgt2 + query_pos if self.pos_enc_at_attn else tgt2 + tgt2 = self.self_attn(q, k, v=tgt2) + tgt = tgt + self.dropout1(tgt2) + return tgt + + def _forward_ca(self, tgt, memory, query_pos, pos, num_k_exclude_rope=0): + kwds = {} + if num_k_exclude_rope > 0: + assert isinstance(self.cross_attn_image, RoPEAttention) + kwds = {"num_k_exclude_rope": num_k_exclude_rope} + + # Cross-Attention + tgt2 = self.norm2(tgt) + tgt2 = self.cross_attn_image( + q=tgt2 + query_pos if self.pos_enc_at_cross_attn_queries else tgt2, + k=memory + pos if self.pos_enc_at_cross_attn_keys else memory, + v=memory, + **kwds, + ) + tgt = tgt + self.dropout2(tgt2) + return tgt + + def forward( + self, + tgt, + memory, + pos: Optional[Tensor] = None, + query_pos: Optional[Tensor] = None, + num_k_exclude_rope: int = 0, + ) -> torch.Tensor: + + # Self-Attn, Cross-Attn + tgt = self._forward_sa(tgt, query_pos) + tgt = self._forward_ca(tgt, memory, query_pos, pos, num_k_exclude_rope) + # MLP + tgt2 = self.norm3(tgt) + tgt2 = self.linear2(self.dropout(self.activation(self.linear1(tgt2)))) + tgt = tgt + self.dropout3(tgt2) + return tgt + + +class MemoryAttention(nn.Module): + def __init__( + self, + d_model: int, + pos_enc_at_input: bool, + layer: nn.Module, + num_layers: int, + batch_first: bool = True, # Do layers expect batch first input? + ): + super().__init__() + self.d_model = d_model + self.layers = get_clones(layer, num_layers) + self.num_layers = num_layers + self.norm = nn.LayerNorm(d_model) + self.pos_enc_at_input = pos_enc_at_input + self.batch_first = batch_first + + def forward( + self, + curr: torch.Tensor, # self-attention inputs + memory: torch.Tensor, # cross-attention inputs + curr_pos: Optional[Tensor] = None, # pos_enc for self-attention inputs + memory_pos: Optional[Tensor] = None, # pos_enc for cross-attention inputs + num_obj_ptr_tokens: int = 0, # number of object pointer *tokens* + ): + if isinstance(curr, list): + assert isinstance(curr_pos, list) + assert len(curr) == len(curr_pos) == 1 + curr, curr_pos = ( + curr[0], + curr_pos[0], + ) + + assert ( + curr.shape[1] == memory.shape[1] + ), "Batch size must be the same for curr and memory" + + output = curr + if self.pos_enc_at_input and curr_pos is not None: + output = output + 0.1 * curr_pos + + if self.batch_first: + # Convert to batch first + output = output.transpose(0, 1) + curr_pos = curr_pos.transpose(0, 1) + memory = memory.transpose(0, 1) + memory_pos = memory_pos.transpose(0, 1) + + for layer in self.layers: + kwds = {} + if isinstance(layer.cross_attn_image, RoPEAttention): + kwds = {"num_k_exclude_rope": num_obj_ptr_tokens} + + output = layer( + tgt=output, + memory=memory, + pos=memory_pos, + query_pos=curr_pos, + **kwds, + ) + normed_output = self.norm(output) + + if self.batch_first: + # Convert back to seq first + normed_output = normed_output.transpose(0, 1) + curr_pos = curr_pos.transpose(0, 1) + + return normed_output diff --git a/freevideo_engine/sam2/modeling/memory_encoder.py b/freevideo_engine/sam2/modeling/memory_encoder.py new file mode 100644 index 00000000..e211db19 --- /dev/null +++ b/freevideo_engine/sam2/modeling/memory_encoder.py @@ -0,0 +1,181 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. + +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. + +import math +from typing import Tuple + +import torch +import torch.nn as nn +import torch.nn.functional as F + +from freevideo_engine.sam2.modeling.sam2_utils import DropPath, get_clones, LayerNorm2d + + +class MaskDownSampler(nn.Module): + """ + Progressively downsample a mask by total_stride, each time by stride. + Note that LayerNorm is applied per *token*, like in ViT. + + With each downsample (by a factor stride**2), channel capacity increases by the same factor. + In the end, we linearly project to embed_dim channels. + """ + + def __init__( + self, + embed_dim=256, + kernel_size=4, + stride=4, + padding=0, + total_stride=16, + activation=nn.GELU, + ): + super().__init__() + num_layers = int(math.log2(total_stride) // math.log2(stride)) + assert stride**num_layers == total_stride + self.encoder = nn.Sequential() + mask_in_chans, mask_out_chans = 1, 1 + for _ in range(num_layers): + mask_out_chans = mask_in_chans * (stride**2) + self.encoder.append( + nn.Conv2d( + mask_in_chans, + mask_out_chans, + kernel_size=kernel_size, + stride=stride, + padding=padding, + ) + ) + self.encoder.append(LayerNorm2d(mask_out_chans)) + self.encoder.append(activation()) + mask_in_chans = mask_out_chans + + self.encoder.append(nn.Conv2d(mask_out_chans, embed_dim, kernel_size=1)) + + def forward(self, x): + return self.encoder(x) + + +# Lightly adapted from ConvNext (https://github.com/facebookresearch/ConvNeXt) +class CXBlock(nn.Module): + r"""ConvNeXt Block. There are two equivalent implementations: + (1) DwConv -> LayerNorm (channels_first) -> 1x1 Conv -> GELU -> 1x1 Conv; all in (N, C, H, W) + (2) DwConv -> Permute to (N, H, W, C); LayerNorm (channels_last) -> Linear -> GELU -> Linear; Permute back + We use (2) as we find it slightly faster in PyTorch + + Args: + dim (int): Number of input channels. + drop_path (float): Stochastic depth rate. Default: 0.0 + layer_scale_init_value (float): Init value for Layer Scale. Default: 1e-6. + """ + + def __init__( + self, + dim, + kernel_size=7, + padding=3, + drop_path=0.0, + layer_scale_init_value=1e-6, + use_dwconv=True, + ): + super().__init__() + self.dwconv = nn.Conv2d( + dim, + dim, + kernel_size=kernel_size, + padding=padding, + groups=dim if use_dwconv else 1, + ) # depthwise conv + self.norm = LayerNorm2d(dim, eps=1e-6) + self.pwconv1 = nn.Linear( + dim, 4 * dim + ) # pointwise/1x1 convs, implemented with linear layers + self.act = nn.GELU() + self.pwconv2 = nn.Linear(4 * dim, dim) + self.gamma = ( + nn.Parameter(layer_scale_init_value * torch.ones((dim)), requires_grad=True) + if layer_scale_init_value > 0 + else None + ) + self.drop_path = DropPath(drop_path) if drop_path > 0.0 else nn.Identity() + + def forward(self, x): + input = x + x = self.dwconv(x) + x = self.norm(x) + x = x.permute(0, 2, 3, 1) # (N, C, H, W) -> (N, H, W, C) + x = self.pwconv1(x) + x = self.act(x) + x = self.pwconv2(x) + if self.gamma is not None: + x = self.gamma * x + x = x.permute(0, 3, 1, 2) # (N, H, W, C) -> (N, C, H, W) + + x = input + self.drop_path(x) + return x + + +class Fuser(nn.Module): + def __init__(self, layer, num_layers, dim=None, input_projection=False): + super().__init__() + self.proj = nn.Identity() + self.layers = get_clones(layer, num_layers) + + if input_projection: + assert dim is not None + self.proj = nn.Conv2d(dim, dim, kernel_size=1) + + def forward(self, x): + # normally x: (N, C, H, W) + x = self.proj(x) + for layer in self.layers: + x = layer(x) + return x + + +class MemoryEncoder(nn.Module): + def __init__( + self, + out_dim, + mask_downsampler, + fuser, + position_encoding, + in_dim=256, # in_dim of pix_feats + ): + super().__init__() + + self.mask_downsampler = mask_downsampler + + self.pix_feat_proj = nn.Conv2d(in_dim, in_dim, kernel_size=1) + self.fuser = fuser + self.position_encoding = position_encoding + self.out_proj = nn.Identity() + if out_dim != in_dim: + self.out_proj = nn.Conv2d(in_dim, out_dim, kernel_size=1) + + def forward( + self, + pix_feat: torch.Tensor, + masks: torch.Tensor, + skip_mask_sigmoid: bool = False, + ) -> Tuple[torch.Tensor, torch.Tensor]: + ## Process masks + # sigmoid, so that less domain shift from gt masks which are bool + if not skip_mask_sigmoid: + masks = F.sigmoid(masks) + masks = self.mask_downsampler(masks) + + ## Fuse pix_feats and downsampled masks + # in case the visual features are on CPU, cast them to CUDA + pix_feat = pix_feat.to(masks.device) + + x = self.pix_feat_proj(pix_feat) + x = x + masks + x = self.fuser(x) + x = self.out_proj(x) + + pos = self.position_encoding(x).to(x.dtype) + + return {"vision_features": x, "vision_pos_enc": [pos]} diff --git a/freevideo_engine/sam2/modeling/position_encoding.py b/freevideo_engine/sam2/modeling/position_encoding.py new file mode 100644 index 00000000..2241d4cf --- /dev/null +++ b/freevideo_engine/sam2/modeling/position_encoding.py @@ -0,0 +1,239 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. + +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. + +import math +from typing import Any, Optional, Tuple + +import numpy as np + +import torch +from torch import nn + + +class PositionEmbeddingSine(nn.Module): + """ + This is a more standard version of the position embedding, very similar to the one + used by the Attention Is All You Need paper, generalized to work on images. + """ + + def __init__( + self, + num_pos_feats, + temperature: int = 10000, + normalize: bool = True, + scale: Optional[float] = None, + # Following settings only relevant + # for warmping up cache for compilation + warmup_cache: bool = True, + image_size: int = 1024, + strides: Tuple[int] = (4, 8, 16, 32), + ): + super().__init__() + assert num_pos_feats % 2 == 0, "Expecting even model width" + self.num_pos_feats = num_pos_feats // 2 + self.temperature = temperature + self.normalize = normalize + if scale is not None and normalize is False: + raise ValueError("normalize should be True if scale is passed") + if scale is None: + scale = 2 * math.pi + self.scale = scale + + self.cache = {} + if warmup_cache and torch.cuda.is_available(): + # Warmup cache for cuda, to help with compilation + device = torch.device("cuda") + for stride in strides: + cache_key = (image_size // stride, image_size // stride) + self._pe(1, device, *cache_key) + + def _encode_xy(self, x, y): + # The positions are expected to be normalized + assert len(x) == len(y) and x.ndim == y.ndim == 1 + x_embed = x * self.scale + y_embed = y * self.scale + + dim_t = torch.arange(self.num_pos_feats, dtype=torch.float32, device=x.device) + dim_t = self.temperature ** (2 * (dim_t // 2) / self.num_pos_feats) + + pos_x = x_embed[:, None] / dim_t + pos_y = y_embed[:, None] / dim_t + pos_x = torch.stack( + (pos_x[:, 0::2].sin(), pos_x[:, 1::2].cos()), dim=2 + ).flatten(1) + pos_y = torch.stack( + (pos_y[:, 0::2].sin(), pos_y[:, 1::2].cos()), dim=2 + ).flatten(1) + return pos_x, pos_y + + @torch.no_grad() + def encode_boxes(self, x, y, w, h): + pos_x, pos_y = self._encode_xy(x, y) + pos = torch.cat((pos_y, pos_x, h[:, None], w[:, None]), dim=1) + return pos + + encode = encode_boxes # Backwards compatibility + + @torch.no_grad() + def encode_points(self, x, y, labels): + (bx, nx), (by, ny), (bl, nl) = x.shape, y.shape, labels.shape + assert bx == by and nx == ny and bx == bl and nx == nl + pos_x, pos_y = self._encode_xy(x.flatten(), y.flatten()) + pos_x, pos_y = pos_x.reshape(bx, nx, -1), pos_y.reshape(by, ny, -1) + pos = torch.cat((pos_y, pos_x, labels[:, :, None]), dim=2) + return pos + + @torch.no_grad() + def _pe(self, B, device, *cache_key): + H, W = cache_key + if cache_key in self.cache: + return self.cache[cache_key].to(device)[None].repeat(B, 1, 1, 1) + + y_embed = ( + torch.arange(1, H + 1, dtype=torch.float32, device=device) + .view(1, -1, 1) + .repeat(B, 1, W) + ) + x_embed = ( + torch.arange(1, W + 1, dtype=torch.float32, device=device) + .view(1, 1, -1) + .repeat(B, H, 1) + ) + + if self.normalize: + eps = 1e-6 + y_embed = y_embed / (y_embed[:, -1:, :] + eps) * self.scale + x_embed = x_embed / (x_embed[:, :, -1:] + eps) * self.scale + + dim_t = torch.arange(self.num_pos_feats, dtype=torch.float32, device=device) + dim_t = self.temperature ** (2 * (dim_t // 2) / self.num_pos_feats) + + pos_x = x_embed[:, :, :, None] / dim_t + pos_y = y_embed[:, :, :, None] / dim_t + pos_x = torch.stack( + (pos_x[:, :, :, 0::2].sin(), pos_x[:, :, :, 1::2].cos()), dim=4 + ).flatten(3) + pos_y = torch.stack( + (pos_y[:, :, :, 0::2].sin(), pos_y[:, :, :, 1::2].cos()), dim=4 + ).flatten(3) + pos = torch.cat((pos_y, pos_x), dim=3).permute(0, 3, 1, 2) + self.cache[cache_key] = pos[0] + return pos + + @torch.no_grad() + def forward(self, x: torch.Tensor): + B = x.shape[0] + cache_key = (x.shape[-2], x.shape[-1]) + return self._pe(B, x.device, *cache_key) + + +class PositionEmbeddingRandom(nn.Module): + """ + Positional encoding using random spatial frequencies. + """ + + def __init__(self, num_pos_feats: int = 64, scale: Optional[float] = None) -> None: + super().__init__() + if scale is None or scale <= 0.0: + scale = 1.0 + self.register_buffer( + "positional_encoding_gaussian_matrix", + scale * torch.randn((2, num_pos_feats)), + ) + + def _pe_encoding(self, coords: torch.Tensor) -> torch.Tensor: + """Positionally encode points that are normalized to [0,1].""" + # assuming coords are in [0, 1]^2 square and have d_1 x ... x d_n x 2 shape + coords = 2 * coords - 1 + coords = coords @ self.positional_encoding_gaussian_matrix + coords = 2 * np.pi * coords + # outputs d_1 x ... x d_n x C shape + return torch.cat([torch.sin(coords), torch.cos(coords)], dim=-1) + + def forward(self, size: Tuple[int, int]) -> torch.Tensor: + """Generate positional encoding for a grid of the specified size.""" + h, w = size + device: Any = self.positional_encoding_gaussian_matrix.device + grid = torch.ones((h, w), device=device, dtype=torch.float32) + y_embed = grid.cumsum(dim=0) - 0.5 + x_embed = grid.cumsum(dim=1) - 0.5 + y_embed = y_embed / h + x_embed = x_embed / w + + pe = self._pe_encoding(torch.stack([x_embed, y_embed], dim=-1)) + return pe.permute(2, 0, 1) # C x H x W + + def forward_with_coords( + self, coords_input: torch.Tensor, image_size: Tuple[int, int] + ) -> torch.Tensor: + """Positionally encode points that are not normalized to [0,1].""" + coords = coords_input.clone() + coords[:, :, 0] = coords[:, :, 0] / image_size[1] + coords[:, :, 1] = coords[:, :, 1] / image_size[0] + return self._pe_encoding(coords.to(torch.float)) # B x N x C + + +# Rotary Positional Encoding, adapted from: +# 1. https://github.com/meta-llama/codellama/blob/main/llama/model.py +# 2. https://github.com/naver-ai/rope-vit +# 3. https://github.com/lucidrains/rotary-embedding-torch + + +def init_t_xy(end_x: int, end_y: int): + t = torch.arange(end_x * end_y, dtype=torch.float32) + t_x = (t % end_x).float() + t_y = torch.div(t, end_x, rounding_mode="floor").float() + return t_x, t_y + + +def compute_axial_cis(dim: int, end_x: int, end_y: int, theta: float = 10000.0): + freqs_x = 1.0 / (theta ** (torch.arange(0, dim, 4)[: (dim // 4)].float() / dim)) + freqs_y = 1.0 / (theta ** (torch.arange(0, dim, 4)[: (dim // 4)].float() / dim)) + + t_x, t_y = init_t_xy(end_x, end_y) + freqs_x = torch.outer(t_x, freqs_x) + freqs_y = torch.outer(t_y, freqs_y) + freqs_cis_x = torch.polar(torch.ones_like(freqs_x), freqs_x) + freqs_cis_y = torch.polar(torch.ones_like(freqs_y), freqs_y) + return torch.cat([freqs_cis_x, freqs_cis_y], dim=-1) + + +def reshape_for_broadcast(freqs_cis: torch.Tensor, x: torch.Tensor): + ndim = x.ndim + assert 0 <= 1 < ndim + assert freqs_cis.shape == (x.shape[-2], x.shape[-1]) + shape = [d if i >= ndim - 2 else 1 for i, d in enumerate(x.shape)] + return freqs_cis.view(*shape) + + +def apply_rotary_enc( + xq: torch.Tensor, + xk: torch.Tensor, + freqs_cis: torch.Tensor, + repeat_freqs_k: bool = False, +): + xq_ = torch.view_as_complex(xq.float().reshape(*xq.shape[:-1], -1, 2)) + xk_ = ( + torch.view_as_complex(xk.float().reshape(*xk.shape[:-1], -1, 2)) + if xk.shape[-2] != 0 + else None + ) + freqs_cis = reshape_for_broadcast(freqs_cis, xq_) + xq_out = torch.view_as_real(xq_ * freqs_cis).flatten(3) + if xk_ is None: + # no keys to rotate, due to dropout + return xq_out.type_as(xq).to(xq.device), xk + # repeat freqs along seq_len dim to match k seq_len + if repeat_freqs_k: + r = xk_.shape[-2] // xq_.shape[-2] + if freqs_cis.is_cuda: + freqs_cis = freqs_cis.repeat(*([1] * (freqs_cis.ndim - 2)), r, 1) + else: + # torch.repeat on complex numbers may not be supported on non-CUDA devices + # (freqs_cis has 4 dims and we repeat on dim 2) so we use expand + flatten + freqs_cis = freqs_cis.unsqueeze(2).expand(-1, -1, r, -1, -1).flatten(2, 3) + xk_out = torch.view_as_real(xk_ * freqs_cis).flatten(3) + return xq_out.type_as(xq).to(xq.device), xk_out.type_as(xk).to(xk.device) diff --git a/freevideo_engine/sam2/modeling/sam/__init__.py b/freevideo_engine/sam2/modeling/sam/__init__.py new file mode 100644 index 00000000..5277f461 --- /dev/null +++ b/freevideo_engine/sam2/modeling/sam/__init__.py @@ -0,0 +1,5 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. + +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. diff --git a/freevideo_engine/sam2/modeling/sam/mask_decoder.py b/freevideo_engine/sam2/modeling/sam/mask_decoder.py new file mode 100644 index 00000000..50fd358b --- /dev/null +++ b/freevideo_engine/sam2/modeling/sam/mask_decoder.py @@ -0,0 +1,295 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. + +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. + +from typing import List, Optional, Tuple, Type + +import torch +from torch import nn + +from freevideo_engine.sam2.modeling.sam2_utils import LayerNorm2d, MLP + + +class MaskDecoder(nn.Module): + def __init__( + self, + *, + transformer_dim: int, + transformer: nn.Module, + num_multimask_outputs: int = 3, + activation: Type[nn.Module] = nn.GELU, + iou_head_depth: int = 3, + iou_head_hidden_dim: int = 256, + use_high_res_features: bool = False, + iou_prediction_use_sigmoid=False, + dynamic_multimask_via_stability=False, + dynamic_multimask_stability_delta=0.05, + dynamic_multimask_stability_thresh=0.98, + pred_obj_scores: bool = False, + pred_obj_scores_mlp: bool = False, + use_multimask_token_for_obj_ptr: bool = False, + ) -> None: + """ + Predicts masks given an image and prompt embeddings, using a + transformer architecture. + + Arguments: + transformer_dim (int): the channel dimension of the transformer + transformer (nn.Module): the transformer used to predict masks + num_multimask_outputs (int): the number of masks to predict + when disambiguating masks + activation (nn.Module): the type of activation to use when + upscaling masks + iou_head_depth (int): the depth of the MLP used to predict + mask quality + iou_head_hidden_dim (int): the hidden dimension of the MLP + used to predict mask quality + """ + super().__init__() + self.transformer_dim = transformer_dim + self.transformer = transformer + + self.num_multimask_outputs = num_multimask_outputs + + self.iou_token = nn.Embedding(1, transformer_dim) + self.num_mask_tokens = num_multimask_outputs + 1 + self.mask_tokens = nn.Embedding(self.num_mask_tokens, transformer_dim) + + self.pred_obj_scores = pred_obj_scores + if self.pred_obj_scores: + self.obj_score_token = nn.Embedding(1, transformer_dim) + self.use_multimask_token_for_obj_ptr = use_multimask_token_for_obj_ptr + + self.output_upscaling = nn.Sequential( + nn.ConvTranspose2d( + transformer_dim, transformer_dim // 4, kernel_size=2, stride=2 + ), + LayerNorm2d(transformer_dim // 4), + activation(), + nn.ConvTranspose2d( + transformer_dim // 4, transformer_dim // 8, kernel_size=2, stride=2 + ), + activation(), + ) + self.use_high_res_features = use_high_res_features + if use_high_res_features: + self.conv_s0 = nn.Conv2d( + transformer_dim, transformer_dim // 8, kernel_size=1, stride=1 + ) + self.conv_s1 = nn.Conv2d( + transformer_dim, transformer_dim // 4, kernel_size=1, stride=1 + ) + + self.output_hypernetworks_mlps = nn.ModuleList( + [ + MLP(transformer_dim, transformer_dim, transformer_dim // 8, 3) + for i in range(self.num_mask_tokens) + ] + ) + + self.iou_prediction_head = MLP( + transformer_dim, + iou_head_hidden_dim, + self.num_mask_tokens, + iou_head_depth, + sigmoid_output=iou_prediction_use_sigmoid, + ) + if self.pred_obj_scores: + self.pred_obj_score_head = nn.Linear(transformer_dim, 1) + if pred_obj_scores_mlp: + self.pred_obj_score_head = MLP(transformer_dim, transformer_dim, 1, 3) + + # When outputting a single mask, optionally we can dynamically fall back to the best + # multimask output token if the single mask output token gives low stability scores. + self.dynamic_multimask_via_stability = dynamic_multimask_via_stability + self.dynamic_multimask_stability_delta = dynamic_multimask_stability_delta + self.dynamic_multimask_stability_thresh = dynamic_multimask_stability_thresh + + def forward( + self, + image_embeddings: torch.Tensor, + image_pe: torch.Tensor, + sparse_prompt_embeddings: torch.Tensor, + dense_prompt_embeddings: torch.Tensor, + multimask_output: bool, + repeat_image: bool, + high_res_features: Optional[List[torch.Tensor]] = None, + ) -> Tuple[torch.Tensor, torch.Tensor]: + """ + Predict masks given image and prompt embeddings. + + Arguments: + image_embeddings (torch.Tensor): the embeddings from the image encoder + image_pe (torch.Tensor): positional encoding with the shape of image_embeddings + sparse_prompt_embeddings (torch.Tensor): the embeddings of the points and boxes + dense_prompt_embeddings (torch.Tensor): the embeddings of the mask inputs + multimask_output (bool): Whether to return multiple masks or a single + mask. + + Returns: + torch.Tensor: batched predicted masks + torch.Tensor: batched predictions of mask quality + torch.Tensor: batched SAM token for mask output + """ + masks, iou_pred, mask_tokens_out, object_score_logits = self.predict_masks( + image_embeddings=image_embeddings, + image_pe=image_pe, + sparse_prompt_embeddings=sparse_prompt_embeddings, + dense_prompt_embeddings=dense_prompt_embeddings, + repeat_image=repeat_image, + high_res_features=high_res_features, + ) + + # Select the correct mask or masks for output + if multimask_output: + masks = masks[:, 1:, :, :] + iou_pred = iou_pred[:, 1:] + elif self.dynamic_multimask_via_stability and not self.training: + masks, iou_pred = self._dynamic_multimask_via_stability(masks, iou_pred) + else: + masks = masks[:, 0:1, :, :] + iou_pred = iou_pred[:, 0:1] + + if multimask_output and self.use_multimask_token_for_obj_ptr: + sam_tokens_out = mask_tokens_out[:, 1:] # [b, 3, c] shape + else: + # Take the mask output token. Here we *always* use the token for single mask output. + # At test time, even if we track after 1-click (and using multimask_output=True), + # we still take the single mask token here. The rationale is that we always track + # after multiple clicks during training, so the past tokens seen during training + # are always the single mask token (and we'll let it be the object-memory token). + sam_tokens_out = mask_tokens_out[:, 0:1] # [b, 1, c] shape + + # Prepare output + return masks, iou_pred, sam_tokens_out, object_score_logits + + def predict_masks( + self, + image_embeddings: torch.Tensor, + image_pe: torch.Tensor, + sparse_prompt_embeddings: torch.Tensor, + dense_prompt_embeddings: torch.Tensor, + repeat_image: bool, + high_res_features: Optional[List[torch.Tensor]] = None, + ) -> Tuple[torch.Tensor, torch.Tensor]: + """Predicts masks. See 'forward' for more details.""" + # Concatenate output tokens + s = 0 + if self.pred_obj_scores: + output_tokens = torch.cat( + [ + self.obj_score_token.weight, + self.iou_token.weight, + self.mask_tokens.weight, + ], + dim=0, + ) + s = 1 + else: + output_tokens = torch.cat( + [self.iou_token.weight, self.mask_tokens.weight], dim=0 + ) + output_tokens = output_tokens.unsqueeze(0).expand( + sparse_prompt_embeddings.size(0), -1, -1 + ) + tokens = torch.cat((output_tokens, sparse_prompt_embeddings), dim=1) + + # Expand per-image data in batch direction to be per-mask + if repeat_image: + src = torch.repeat_interleave(image_embeddings, tokens.shape[0], dim=0) + else: + assert image_embeddings.shape[0] == tokens.shape[0] + src = image_embeddings + src = src + dense_prompt_embeddings + assert ( + image_pe.size(0) == 1 + ), "image_pe should have size 1 in batch dim (from `get_dense_pe()`)" + pos_src = torch.repeat_interleave(image_pe, tokens.shape[0], dim=0) + b, c, h, w = src.shape + + # Run the transformer + hs, src = self.transformer(src, pos_src, tokens) + iou_token_out = hs[:, s, :] + mask_tokens_out = hs[:, s + 1 : (s + 1 + self.num_mask_tokens), :] + + # Upscale mask embeddings and predict masks using the mask tokens + src = src.transpose(1, 2).view(b, c, h, w) + if not self.use_high_res_features: + upscaled_embedding = self.output_upscaling(src) + else: + dc1, ln1, act1, dc2, act2 = self.output_upscaling + feat_s0, feat_s1 = high_res_features + upscaled_embedding = act1(ln1(dc1(src) + feat_s1)) + upscaled_embedding = act2(dc2(upscaled_embedding) + feat_s0) + + hyper_in_list: List[torch.Tensor] = [] + for i in range(self.num_mask_tokens): + hyper_in_list.append( + self.output_hypernetworks_mlps[i](mask_tokens_out[:, i, :]) + ) + hyper_in = torch.stack(hyper_in_list, dim=1) + b, c, h, w = upscaled_embedding.shape + masks = (hyper_in @ upscaled_embedding.view(b, c, h * w)).view(b, -1, h, w) + + # Generate mask quality predictions + iou_pred = self.iou_prediction_head(iou_token_out) + if self.pred_obj_scores: + assert s == 1 + object_score_logits = self.pred_obj_score_head(hs[:, 0, :]) + else: + # Obj scores logits - default to 10.0, i.e. assuming the object is present, sigmoid(10)=1 + object_score_logits = 10.0 * iou_pred.new_ones(iou_pred.shape[0], 1) + + return masks, iou_pred, mask_tokens_out, object_score_logits + + def _get_stability_scores(self, mask_logits): + """ + Compute stability scores of the mask logits based on the IoU between upper and + lower thresholds. + """ + mask_logits = mask_logits.flatten(-2) + stability_delta = self.dynamic_multimask_stability_delta + area_i = torch.sum(mask_logits > stability_delta, dim=-1).float() + area_u = torch.sum(mask_logits > -stability_delta, dim=-1).float() + stability_scores = torch.where(area_u > 0, area_i / area_u, 1.0) + return stability_scores + + def _dynamic_multimask_via_stability(self, all_mask_logits, all_iou_scores): + """ + When outputting a single mask, if the stability score from the current single-mask + output (based on output token 0) falls below a threshold, we instead select from + multi-mask outputs (based on output token 1~3) the mask with the highest predicted + IoU score. This is intended to ensure a valid mask for both clicking and tracking. + """ + # The best mask from multimask output tokens (1~3) + multimask_logits = all_mask_logits[:, 1:, :, :] + multimask_iou_scores = all_iou_scores[:, 1:] + best_scores_inds = torch.argmax(multimask_iou_scores, dim=-1) + batch_inds = torch.arange( + multimask_iou_scores.size(0), device=all_iou_scores.device + ) + best_multimask_logits = multimask_logits[batch_inds, best_scores_inds] + best_multimask_logits = best_multimask_logits.unsqueeze(1) + best_multimask_iou_scores = multimask_iou_scores[batch_inds, best_scores_inds] + best_multimask_iou_scores = best_multimask_iou_scores.unsqueeze(1) + + # The mask from singlemask output token 0 and its stability score + singlemask_logits = all_mask_logits[:, 0:1, :, :] + singlemask_iou_scores = all_iou_scores[:, 0:1] + stability_scores = self._get_stability_scores(singlemask_logits) + is_stable = stability_scores >= self.dynamic_multimask_stability_thresh + + # Dynamically fall back to best multimask output upon low stability scores. + mask_logits_out = torch.where( + is_stable[..., None, None].expand_as(singlemask_logits), + singlemask_logits, + best_multimask_logits, + ) + iou_scores_out = torch.where( + is_stable.expand_as(singlemask_iou_scores), + singlemask_iou_scores, + best_multimask_iou_scores, + ) + return mask_logits_out, iou_scores_out diff --git a/freevideo_engine/sam2/modeling/sam/prompt_encoder.py b/freevideo_engine/sam2/modeling/sam/prompt_encoder.py new file mode 100644 index 00000000..1b449b75 --- /dev/null +++ b/freevideo_engine/sam2/modeling/sam/prompt_encoder.py @@ -0,0 +1,202 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. + +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. + +from typing import Optional, Tuple, Type + +import torch +from torch import nn + +from freevideo_engine.sam2.modeling.position_encoding import PositionEmbeddingRandom + +from freevideo_engine.sam2.modeling.sam2_utils import LayerNorm2d + + +class PromptEncoder(nn.Module): + def __init__( + self, + embed_dim: int, + image_embedding_size: Tuple[int, int], + input_image_size: Tuple[int, int], + mask_in_chans: int, + activation: Type[nn.Module] = nn.GELU, + ) -> None: + """ + Encodes prompts for input to SAM's mask decoder. + + Arguments: + embed_dim (int): The prompts' embedding dimension + image_embedding_size (tuple(int, int)): The spatial size of the + image embedding, as (H, W). + input_image_size (int): The padded size of the image as input + to the image encoder, as (H, W). + mask_in_chans (int): The number of hidden channels used for + encoding input masks. + activation (nn.Module): The activation to use when encoding + input masks. + """ + super().__init__() + self.embed_dim = embed_dim + self.input_image_size = input_image_size + self.image_embedding_size = image_embedding_size + self.pe_layer = PositionEmbeddingRandom(embed_dim // 2) + + self.num_point_embeddings: int = 4 # pos/neg point + 2 box corners + point_embeddings = [ + nn.Embedding(1, embed_dim) for i in range(self.num_point_embeddings) + ] + self.point_embeddings = nn.ModuleList(point_embeddings) + self.not_a_point_embed = nn.Embedding(1, embed_dim) + + self.mask_input_size = ( + 4 * image_embedding_size[0], + 4 * image_embedding_size[1], + ) + self.mask_downscaling = nn.Sequential( + nn.Conv2d(1, mask_in_chans // 4, kernel_size=2, stride=2), + LayerNorm2d(mask_in_chans // 4), + activation(), + nn.Conv2d(mask_in_chans // 4, mask_in_chans, kernel_size=2, stride=2), + LayerNorm2d(mask_in_chans), + activation(), + nn.Conv2d(mask_in_chans, embed_dim, kernel_size=1), + ) + self.no_mask_embed = nn.Embedding(1, embed_dim) + + def get_dense_pe(self) -> torch.Tensor: + """ + Returns the positional encoding used to encode point prompts, + applied to a dense set of points the shape of the image encoding. + + Returns: + torch.Tensor: Positional encoding with shape + 1x(embed_dim)x(embedding_h)x(embedding_w) + """ + return self.pe_layer(self.image_embedding_size).unsqueeze(0) + + def _embed_points( + self, + points: torch.Tensor, + labels: torch.Tensor, + pad: bool, + ) -> torch.Tensor: + """Embeds point prompts.""" + points = points + 0.5 # Shift to center of pixel + if pad: + padding_point = torch.zeros((points.shape[0], 1, 2), device=points.device) + padding_label = -torch.ones((labels.shape[0], 1), device=labels.device) + points = torch.cat([points, padding_point], dim=1) + labels = torch.cat([labels, padding_label], dim=1) + point_embedding = self.pe_layer.forward_with_coords( + points, self.input_image_size + ) + + point_embedding = torch.where( + (labels == -1).unsqueeze(-1), + torch.zeros_like(point_embedding) + self.not_a_point_embed.weight, + point_embedding, + ) + point_embedding = torch.where( + (labels == 0).unsqueeze(-1), + point_embedding + self.point_embeddings[0].weight, + point_embedding, + ) + point_embedding = torch.where( + (labels == 1).unsqueeze(-1), + point_embedding + self.point_embeddings[1].weight, + point_embedding, + ) + point_embedding = torch.where( + (labels == 2).unsqueeze(-1), + point_embedding + self.point_embeddings[2].weight, + point_embedding, + ) + point_embedding = torch.where( + (labels == 3).unsqueeze(-1), + point_embedding + self.point_embeddings[3].weight, + point_embedding, + ) + return point_embedding + + def _embed_boxes(self, boxes: torch.Tensor) -> torch.Tensor: + """Embeds box prompts.""" + boxes = boxes + 0.5 # Shift to center of pixel + coords = boxes.reshape(-1, 2, 2) + corner_embedding = self.pe_layer.forward_with_coords( + coords, self.input_image_size + ) + corner_embedding[:, 0, :] += self.point_embeddings[2].weight + corner_embedding[:, 1, :] += self.point_embeddings[3].weight + return corner_embedding + + def _embed_masks(self, masks: torch.Tensor) -> torch.Tensor: + """Embeds mask inputs.""" + mask_embedding = self.mask_downscaling(masks) + return mask_embedding + + def _get_batch_size( + self, + points: Optional[Tuple[torch.Tensor, torch.Tensor]], + boxes: Optional[torch.Tensor], + masks: Optional[torch.Tensor], + ) -> int: + """ + Gets the batch size of the output given the batch size of the input prompts. + """ + if points is not None: + return points[0].shape[0] + elif boxes is not None: + return boxes.shape[0] + elif masks is not None: + return masks.shape[0] + else: + return 1 + + def _get_device(self) -> torch.device: + return self.point_embeddings[0].weight.device + + def forward( + self, + points: Optional[Tuple[torch.Tensor, torch.Tensor]], + boxes: Optional[torch.Tensor], + masks: Optional[torch.Tensor], + ) -> Tuple[torch.Tensor, torch.Tensor]: + """ + Embeds different types of prompts, returning both sparse and dense + embeddings. + + Arguments: + points (tuple(torch.Tensor, torch.Tensor) or none): point coordinates + and labels to embed. + boxes (torch.Tensor or none): boxes to embed + masks (torch.Tensor or none): masks to embed + + Returns: + torch.Tensor: sparse embeddings for the points and boxes, with shape + BxNx(embed_dim), where N is determined by the number of input points + and boxes. + torch.Tensor: dense embeddings for the masks, in the shape + Bx(embed_dim)x(embed_H)x(embed_W) + """ + bs = self._get_batch_size(points, boxes, masks) + sparse_embeddings = torch.empty( + (bs, 0, self.embed_dim), device=self._get_device() + ) + if points is not None: + coords, labels = points + point_embeddings = self._embed_points(coords, labels, pad=(boxes is None)) + sparse_embeddings = torch.cat([sparse_embeddings, point_embeddings], dim=1) + if boxes is not None: + box_embeddings = self._embed_boxes(boxes) + sparse_embeddings = torch.cat([sparse_embeddings, box_embeddings], dim=1) + + if masks is not None: + dense_embeddings = self._embed_masks(masks) + else: + dense_embeddings = self.no_mask_embed.weight.reshape(1, -1, 1, 1).expand( + bs, -1, self.image_embedding_size[0], self.image_embedding_size[1] + ) + + return sparse_embeddings, dense_embeddings diff --git a/freevideo_engine/sam2/modeling/sam/transformer.py b/freevideo_engine/sam2/modeling/sam/transformer.py new file mode 100644 index 00000000..3855c1f9 --- /dev/null +++ b/freevideo_engine/sam2/modeling/sam/transformer.py @@ -0,0 +1,310 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. + +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. + +import math +from functools import partial +from typing import Tuple, Type + +import torch +import torch.nn.functional as F +from torch import nn, Tensor + +from freevideo_engine.sam2.modeling.position_encoding import apply_rotary_enc, compute_axial_cis +from freevideo_engine.sam2.modeling.sam2_utils import MLP + + +class TwoWayTransformer(nn.Module): + def __init__( + self, + depth: int, + embedding_dim: int, + num_heads: int, + mlp_dim: int, + activation: Type[nn.Module] = nn.ReLU, + attention_downsample_rate: int = 2, + ) -> None: + """ + A transformer decoder that attends to an input image using + queries whose positional embedding is supplied. + + Args: + depth (int): number of layers in the transformer + embedding_dim (int): the channel dimension for the input embeddings + num_heads (int): the number of heads for multihead attention. Must + divide embedding_dim + mlp_dim (int): the channel dimension internal to the MLP block + activation (nn.Module): the activation to use in the MLP block + """ + super().__init__() + self.depth = depth + self.embedding_dim = embedding_dim + self.num_heads = num_heads + self.mlp_dim = mlp_dim + self.layers = nn.ModuleList() + + for i in range(depth): + self.layers.append( + TwoWayAttentionBlock( + embedding_dim=embedding_dim, + num_heads=num_heads, + mlp_dim=mlp_dim, + activation=activation, + attention_downsample_rate=attention_downsample_rate, + skip_first_layer_pe=(i == 0), + ) + ) + + self.final_attn_token_to_image = Attention( + embedding_dim, num_heads, downsample_rate=attention_downsample_rate + ) + self.norm_final_attn = nn.LayerNorm(embedding_dim) + + def forward( + self, + image_embedding: Tensor, + image_pe: Tensor, + point_embedding: Tensor, + ) -> Tuple[Tensor, Tensor]: + """ + Args: + image_embedding (torch.Tensor): image to attend to. Should be shape + B x embedding_dim x h x w for any h and w. + image_pe (torch.Tensor): the positional encoding to add to the image. Must + have the same shape as image_embedding. + point_embedding (torch.Tensor): the embedding to add to the query points. + Must have shape B x N_points x embedding_dim for any N_points. + + Returns: + torch.Tensor: the processed point_embedding + torch.Tensor: the processed image_embedding + """ + # BxCxHxW -> BxHWxC == B x N_image_tokens x C + bs, c, h, w = image_embedding.shape + image_embedding = image_embedding.flatten(2).permute(0, 2, 1) + image_pe = image_pe.flatten(2).permute(0, 2, 1) + + # Prepare queries + queries = point_embedding + keys = image_embedding + + # Apply transformer blocks and final layernorm + for layer in self.layers: + queries, keys = layer( + queries=queries, + keys=keys, + query_pe=point_embedding, + key_pe=image_pe, + ) + + # Apply the final attention layer from the points to the image + q = queries + point_embedding + k = keys + image_pe + attn_out = self.final_attn_token_to_image(q=q, k=k, v=keys) + queries = queries + attn_out + queries = self.norm_final_attn(queries) + + return queries, keys + + +class TwoWayAttentionBlock(nn.Module): + def __init__( + self, + embedding_dim: int, + num_heads: int, + mlp_dim: int = 2048, + activation: Type[nn.Module] = nn.ReLU, + attention_downsample_rate: int = 2, + skip_first_layer_pe: bool = False, + ) -> None: + """ + A transformer block with four layers: (1) self-attention of sparse + inputs, (2) cross attention of sparse inputs to dense inputs, (3) mlp + block on sparse inputs, and (4) cross attention of dense inputs to sparse + inputs. + + Arguments: + embedding_dim (int): the channel dimension of the embeddings + num_heads (int): the number of heads in the attention layers + mlp_dim (int): the hidden dimension of the mlp block + activation (nn.Module): the activation of the mlp block + skip_first_layer_pe (bool): skip the PE on the first layer + """ + super().__init__() + self.self_attn = Attention(embedding_dim, num_heads) + self.norm1 = nn.LayerNorm(embedding_dim) + + self.cross_attn_token_to_image = Attention( + embedding_dim, num_heads, downsample_rate=attention_downsample_rate + ) + self.norm2 = nn.LayerNorm(embedding_dim) + + self.mlp = MLP( + embedding_dim, mlp_dim, embedding_dim, num_layers=2, activation=activation + ) + self.norm3 = nn.LayerNorm(embedding_dim) + + self.norm4 = nn.LayerNorm(embedding_dim) + self.cross_attn_image_to_token = Attention( + embedding_dim, num_heads, downsample_rate=attention_downsample_rate + ) + + self.skip_first_layer_pe = skip_first_layer_pe + + def forward( + self, queries: Tensor, keys: Tensor, query_pe: Tensor, key_pe: Tensor + ) -> Tuple[Tensor, Tensor]: + # Self attention block + if self.skip_first_layer_pe: + queries = self.self_attn(q=queries, k=queries, v=queries) + else: + q = queries + query_pe + attn_out = self.self_attn(q=q, k=q, v=queries) + queries = queries + attn_out + queries = self.norm1(queries) + + # Cross attention block, tokens attending to image embedding + q = queries + query_pe + k = keys + key_pe + attn_out = self.cross_attn_token_to_image(q=q, k=k, v=keys) + queries = queries + attn_out + queries = self.norm2(queries) + + # MLP block + mlp_out = self.mlp(queries) + queries = queries + mlp_out + queries = self.norm3(queries) + + # Cross attention block, image embedding attending to tokens + q = queries + query_pe + k = keys + key_pe + attn_out = self.cross_attn_image_to_token(q=k, k=q, v=queries) + keys = keys + attn_out + keys = self.norm4(keys) + + return queries, keys + + +class Attention(nn.Module): + """ + An attention layer that allows for downscaling the size of the embedding + after projection to queries, keys, and values. + """ + + def __init__( + self, + embedding_dim: int, + num_heads: int, + downsample_rate: int = 1, + dropout: float = 0.0, + kv_in_dim: int = None, + ) -> None: + super().__init__() + self.embedding_dim = embedding_dim + self.kv_in_dim = kv_in_dim if kv_in_dim is not None else embedding_dim + self.internal_dim = embedding_dim // downsample_rate + self.num_heads = num_heads + assert ( + self.internal_dim % num_heads == 0 + ), "num_heads must divide embedding_dim." + + self.q_proj = nn.Linear(embedding_dim, self.internal_dim) + self.k_proj = nn.Linear(self.kv_in_dim, self.internal_dim) + self.v_proj = nn.Linear(self.kv_in_dim, self.internal_dim) + self.out_proj = nn.Linear(self.internal_dim, embedding_dim) + + self.dropout_p = dropout + + def _separate_heads(self, x: Tensor, num_heads: int) -> Tensor: + b, n, c = x.shape + x = x.reshape(b, n, num_heads, c // num_heads) + return x.transpose(1, 2) # B x N_heads x N_tokens x C_per_head + + def _recombine_heads(self, x: Tensor) -> Tensor: + b, n_heads, n_tokens, c_per_head = x.shape + x = x.transpose(1, 2) + return x.reshape(b, n_tokens, n_heads * c_per_head) # B x N_tokens x C + + def forward(self, q: Tensor, k: Tensor, v: Tensor) -> Tensor: + # Input projections + q = self.q_proj(q) + k = self.k_proj(k) + v = self.v_proj(v) + + # Separate into heads + q = self._separate_heads(q, self.num_heads) + k = self._separate_heads(k, self.num_heads) + v = self._separate_heads(v, self.num_heads) + + dropout_p = self.dropout_p if self.training else 0.0 + # Attention + out = F.scaled_dot_product_attention(q, k, v, dropout_p=dropout_p) + + out = self._recombine_heads(out) + out = self.out_proj(out) + + return out + + +class RoPEAttention(Attention): + """Attention with rotary position encoding.""" + + def __init__( + self, + *args, + rope_theta=10000.0, + # whether to repeat q rope to match k length + # this is needed for cross-attention to memories + rope_k_repeat=False, + feat_sizes=(64, 64), # [w, h] for stride 16 feats at 1024 resolution + **kwargs, + ): + super().__init__(*args, **kwargs) + + self.compute_cis = partial( + compute_axial_cis, dim=self.internal_dim // self.num_heads, theta=rope_theta + ) + freqs_cis = self.compute_cis(end_x=feat_sizes[0], end_y=feat_sizes[1]) + # FreeVideo: this inference snapshot is CPU-only, including construction. + self.freqs_cis = freqs_cis + self.rope_k_repeat = rope_k_repeat + + def forward( + self, q: Tensor, k: Tensor, v: Tensor, num_k_exclude_rope: int = 0 + ) -> Tensor: + # Input projections + q = self.q_proj(q) + k = self.k_proj(k) + v = self.v_proj(v) + + # Separate into heads + q = self._separate_heads(q, self.num_heads) + k = self._separate_heads(k, self.num_heads) + v = self._separate_heads(v, self.num_heads) + + # Apply rotary position encoding + w = h = math.sqrt(q.shape[-2]) + self.freqs_cis = self.freqs_cis.to(q.device) + if self.freqs_cis.shape[0] != q.shape[-2]: + self.freqs_cis = self.compute_cis(end_x=w, end_y=h).to(q.device) + if q.shape[-2] != k.shape[-2]: + assert self.rope_k_repeat + + num_k_rope = k.size(-2) - num_k_exclude_rope + q, k[:, :, :num_k_rope] = apply_rotary_enc( + q, + k[:, :, :num_k_rope], + freqs_cis=self.freqs_cis, + repeat_freqs_k=self.rope_k_repeat, + ) + + dropout_p = self.dropout_p if self.training else 0.0 + # Attention + out = F.scaled_dot_product_attention(q, k, v, dropout_p=dropout_p) + + out = self._recombine_heads(out) + out = self.out_proj(out) + + return out diff --git a/freevideo_engine/sam2/modeling/sam2_base.py b/freevideo_engine/sam2/modeling/sam2_base.py new file mode 100644 index 00000000..8b57dd0a --- /dev/null +++ b/freevideo_engine/sam2/modeling/sam2_base.py @@ -0,0 +1,909 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. + +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. + +import torch +import torch.distributed +import torch.nn.functional as F + +from torch.nn.init import trunc_normal_ + +from freevideo_engine.sam2.modeling.sam.mask_decoder import MaskDecoder +from freevideo_engine.sam2.modeling.sam.prompt_encoder import PromptEncoder +from freevideo_engine.sam2.modeling.sam.transformer import TwoWayTransformer +from freevideo_engine.sam2.modeling.sam2_utils import get_1d_sine_pe, MLP, select_closest_cond_frames + +# a large negative value as a placeholder score for missing objects +NO_OBJ_SCORE = -1024.0 + + +class SAM2Base(torch.nn.Module): + def __init__( + self, + image_encoder, + memory_attention, + memory_encoder, + num_maskmem=7, # default 1 input frame + 6 previous frames + image_size=512, + backbone_stride=16, # stride of the image backbone output + sigmoid_scale_for_mem_enc=1.0, # scale factor for mask sigmoid prob + sigmoid_bias_for_mem_enc=0.0, # bias factor for mask sigmoid prob + # During evaluation, whether to binarize the sigmoid mask logits on interacted frames with clicks + binarize_mask_from_pts_for_mem_enc=False, + use_mask_input_as_output_without_sam=False, # on frames with mask input, whether to directly output the input mask without using a SAM prompt encoder + mask decoder + # The maximum number of conditioning frames to participate in the memory attention (-1 means no limit; if there are more conditioning frames than this limit, + # we only cross-attend to the temporally closest `max_cond_frames_in_attn` conditioning frames in the encoder when tracking each frame). This gives the model + # a temporal locality when handling a large number of annotated frames (since closer frames should be more important) and also avoids GPU OOM. + max_cond_frames_in_attn=-1, + # on the first frame, whether to directly add the no-memory embedding to the image feature + # (instead of using the transformer encoder) + directly_add_no_mem_embed=False, + # whether to use high-resolution feature maps in the SAM mask decoder + use_high_res_features_in_sam=False, + # whether to output multiple (3) masks for the first click on initial conditioning frames + multimask_output_in_sam=False, + # the minimum and maximum number of clicks to use multimask_output_in_sam (only relevant when `multimask_output_in_sam=True`; + # default is 1 for both, meaning that only the first click gives multimask output; also note that a box counts as two points) + multimask_min_pt_num=1, + multimask_max_pt_num=1, + # whether to also use multimask output for tracking (not just for the first click on initial conditioning frames; only relevant when `multimask_output_in_sam=True`) + multimask_output_for_tracking=False, + # Whether to use multimask tokens for obj ptr; Only relevant when both + # use_obj_ptrs_in_encoder=True and multimask_output_for_tracking=True + use_multimask_token_for_obj_ptr: bool = False, + # whether to use sigmoid to restrict ious prediction to [0-1] + iou_prediction_use_sigmoid=False, + # The memory bank's temporal stride during evaluation (i.e. the `r` parameter in XMem and Cutie; XMem and Cutie use r=5). + # For r>1, the (self.num_maskmem - 1) non-conditioning memory frames consist of + # (self.num_maskmem - 2) nearest frames from every r-th frames, plus the last frame. + memory_temporal_stride_for_eval=1, + # whether to apply non-overlapping constraints on the object masks in the memory encoder during evaluation (to avoid/alleviate superposing masks) + non_overlap_masks_for_mem_enc=False, + # whether to cross-attend to object pointers from other frames (based on SAM output tokens) in the encoder + use_obj_ptrs_in_encoder=False, + # the maximum number of object pointers from other frames in encoder cross attention (only relevant when `use_obj_ptrs_in_encoder=True`) + max_obj_ptrs_in_encoder=16, + # whether to add temporal positional encoding to the object pointers in the encoder (only relevant when `use_obj_ptrs_in_encoder=True`) + add_tpos_enc_to_obj_ptrs=True, + # whether to add an extra linear projection layer for the temporal positional encoding in the object pointers to avoid potential interference + # with spatial positional encoding (only relevant when both `use_obj_ptrs_in_encoder=True` and `add_tpos_enc_to_obj_ptrs=True`) + proj_tpos_enc_in_obj_ptrs=False, + # whether to use signed distance (instead of unsigned absolute distance) in the temporal positional encoding in the object pointers + # (only relevant when both `use_obj_ptrs_in_encoder=True` and `add_tpos_enc_to_obj_ptrs=True`) + use_signed_tpos_enc_to_obj_ptrs=False, + # whether to only attend to object pointers in the past (before the current frame) in the encoder during evaluation + # (only relevant when `use_obj_ptrs_in_encoder=True`; this might avoid pointer information too far in the future to distract the initial tracking) + only_obj_ptrs_in_the_past_for_eval=False, + # Whether to predict if there is an object in the frame + pred_obj_scores: bool = False, + # Whether to use an MLP to predict object scores + pred_obj_scores_mlp: bool = False, + # Only relevant if pred_obj_scores=True and use_obj_ptrs_in_encoder=True; + # Whether to have a fixed no obj pointer when there is no object present + # or to use it as an additive embedding with obj_ptr produced by decoder + fixed_no_obj_ptr: bool = False, + # Soft no object, i.e. mix in no_obj_ptr softly, + # hope to make recovery easier if there is a mistake and mitigate accumulation of errors + soft_no_obj_ptr: bool = False, + use_mlp_for_obj_ptr_proj: bool = False, + # add no obj embedding to spatial frames + no_obj_embed_spatial: bool = False, + # extra arguments used to construct the SAM mask decoder; if not None, it should be a dict of kwargs to be passed into `MaskDecoder` class. + sam_mask_decoder_extra_args=None, + compile_image_encoder: bool = False, + ): + super().__init__() + + # Part 1: the image backbone + self.image_encoder = image_encoder + # Use level 0, 1, 2 for high-res setting, or just level 2 for the default setting + self.use_high_res_features_in_sam = use_high_res_features_in_sam + self.num_feature_levels = 3 if use_high_res_features_in_sam else 1 + self.use_obj_ptrs_in_encoder = use_obj_ptrs_in_encoder + self.max_obj_ptrs_in_encoder = max_obj_ptrs_in_encoder + if use_obj_ptrs_in_encoder: + # A conv layer to downsample the mask prompt to stride 4 (the same stride as + # low-res SAM mask logits) and to change its scales from 0~1 to SAM logit scale, + # so that it can be fed into the SAM mask decoder to generate a pointer. + self.mask_downsample = torch.nn.Conv2d(1, 1, kernel_size=4, stride=4) + self.add_tpos_enc_to_obj_ptrs = add_tpos_enc_to_obj_ptrs + if proj_tpos_enc_in_obj_ptrs: + assert add_tpos_enc_to_obj_ptrs # these options need to be used together + self.proj_tpos_enc_in_obj_ptrs = proj_tpos_enc_in_obj_ptrs + self.use_signed_tpos_enc_to_obj_ptrs = use_signed_tpos_enc_to_obj_ptrs + self.only_obj_ptrs_in_the_past_for_eval = only_obj_ptrs_in_the_past_for_eval + + # Part 2: memory attention to condition current frame's visual features + # with memories (and obj ptrs) from past frames + self.memory_attention = memory_attention + self.hidden_dim = image_encoder.neck.d_model + + # Part 3: memory encoder for the previous frame's outputs + self.memory_encoder = memory_encoder + self.mem_dim = self.hidden_dim + if hasattr(self.memory_encoder, "out_proj") and hasattr( + self.memory_encoder.out_proj, "weight" + ): + # if there is compression of memories along channel dim + self.mem_dim = self.memory_encoder.out_proj.weight.shape[0] + self.num_maskmem = num_maskmem # Number of memories accessible + # Temporal encoding of the memories + self.maskmem_tpos_enc = torch.nn.Parameter( + torch.zeros(num_maskmem, 1, 1, self.mem_dim) + ) + trunc_normal_(self.maskmem_tpos_enc, std=0.02) + # a single token to indicate no memory embedding from previous frames + self.no_mem_embed = torch.nn.Parameter(torch.zeros(1, 1, self.hidden_dim)) + self.no_mem_pos_enc = torch.nn.Parameter(torch.zeros(1, 1, self.hidden_dim)) + trunc_normal_(self.no_mem_embed, std=0.02) + trunc_normal_(self.no_mem_pos_enc, std=0.02) + self.directly_add_no_mem_embed = directly_add_no_mem_embed + # Apply sigmoid to the output raw mask logits (to turn them from + # range (-inf, +inf) to range (0, 1)) before feeding them into the memory encoder + self.sigmoid_scale_for_mem_enc = sigmoid_scale_for_mem_enc + self.sigmoid_bias_for_mem_enc = sigmoid_bias_for_mem_enc + self.binarize_mask_from_pts_for_mem_enc = binarize_mask_from_pts_for_mem_enc + self.non_overlap_masks_for_mem_enc = non_overlap_masks_for_mem_enc + self.memory_temporal_stride_for_eval = memory_temporal_stride_for_eval + # On frames with mask input, whether to directly output the input mask without + # using a SAM prompt encoder + mask decoder + self.use_mask_input_as_output_without_sam = use_mask_input_as_output_without_sam + self.multimask_output_in_sam = multimask_output_in_sam + self.multimask_min_pt_num = multimask_min_pt_num + self.multimask_max_pt_num = multimask_max_pt_num + self.multimask_output_for_tracking = multimask_output_for_tracking + self.use_multimask_token_for_obj_ptr = use_multimask_token_for_obj_ptr + self.iou_prediction_use_sigmoid = iou_prediction_use_sigmoid + + # Part 4: SAM-style prompt encoder (for both mask and point inputs) + # and SAM-style mask decoder for the final mask output + self.image_size = image_size + self.backbone_stride = backbone_stride + self.sam_mask_decoder_extra_args = sam_mask_decoder_extra_args + self.pred_obj_scores = pred_obj_scores + self.pred_obj_scores_mlp = pred_obj_scores_mlp + self.fixed_no_obj_ptr = fixed_no_obj_ptr + self.soft_no_obj_ptr = soft_no_obj_ptr + if self.fixed_no_obj_ptr: + assert self.pred_obj_scores + assert self.use_obj_ptrs_in_encoder + if self.pred_obj_scores and self.use_obj_ptrs_in_encoder: + self.no_obj_ptr = torch.nn.Parameter(torch.zeros(1, self.hidden_dim)) + trunc_normal_(self.no_obj_ptr, std=0.02) + self.use_mlp_for_obj_ptr_proj = use_mlp_for_obj_ptr_proj + self.no_obj_embed_spatial = None + if no_obj_embed_spatial: + self.no_obj_embed_spatial = torch.nn.Parameter(torch.zeros(1, self.mem_dim)) + trunc_normal_(self.no_obj_embed_spatial, std=0.02) + + self._build_sam_heads() + self.max_cond_frames_in_attn = max_cond_frames_in_attn + + # Model compilation + if compile_image_encoder: + # Compile the forward function (not the full module) to allow loading checkpoints. + print( + "Image encoder compilation is enabled. First forward pass will be slow." + ) + self.image_encoder.forward = torch.compile( + self.image_encoder.forward, + mode="max-autotune", + fullgraph=True, + dynamic=False, + ) + + @property + def device(self): + return next(self.parameters()).device + + def forward(self, *args, **kwargs): + raise NotImplementedError( + "Please use the corresponding methods in SAM2VideoPredictor for inference or SAM2Train for training/fine-tuning" + "See notebooks/video_predictor_example.ipynb for an inference example." + ) + + def _build_sam_heads(self): + """Build SAM-style prompt encoder and mask decoder.""" + self.sam_prompt_embed_dim = self.hidden_dim + self.sam_image_embedding_size = self.image_size // self.backbone_stride + + # build PromptEncoder and MaskDecoder from SAM + # (their hyperparameters like `mask_in_chans=16` are from SAM code) + self.sam_prompt_encoder = PromptEncoder( + embed_dim=self.sam_prompt_embed_dim, + image_embedding_size=( + self.sam_image_embedding_size, + self.sam_image_embedding_size, + ), + input_image_size=(self.image_size, self.image_size), + mask_in_chans=16, + ) + self.sam_mask_decoder = MaskDecoder( + num_multimask_outputs=3, + transformer=TwoWayTransformer( + depth=2, + embedding_dim=self.sam_prompt_embed_dim, + mlp_dim=2048, + num_heads=8, + ), + transformer_dim=self.sam_prompt_embed_dim, + iou_head_depth=3, + iou_head_hidden_dim=256, + use_high_res_features=self.use_high_res_features_in_sam, + iou_prediction_use_sigmoid=self.iou_prediction_use_sigmoid, + pred_obj_scores=self.pred_obj_scores, + pred_obj_scores_mlp=self.pred_obj_scores_mlp, + use_multimask_token_for_obj_ptr=self.use_multimask_token_for_obj_ptr, + **(self.sam_mask_decoder_extra_args or {}), + ) + if self.use_obj_ptrs_in_encoder: + # a linear projection on SAM output tokens to turn them into object pointers + self.obj_ptr_proj = torch.nn.Linear(self.hidden_dim, self.hidden_dim) + if self.use_mlp_for_obj_ptr_proj: + self.obj_ptr_proj = MLP( + self.hidden_dim, self.hidden_dim, self.hidden_dim, 3 + ) + else: + self.obj_ptr_proj = torch.nn.Identity() + if self.proj_tpos_enc_in_obj_ptrs: + # a linear projection on temporal positional encoding in object pointers to + # avoid potential interference with spatial positional encoding + self.obj_ptr_tpos_proj = torch.nn.Linear(self.hidden_dim, self.mem_dim) + else: + self.obj_ptr_tpos_proj = torch.nn.Identity() + + def _forward_sam_heads( + self, + backbone_features, + point_inputs=None, + mask_inputs=None, + high_res_features=None, + multimask_output=False, + ): + """ + Forward SAM prompt encoders and mask heads. + + Inputs: + - backbone_features: image features of [B, C, H, W] shape + - point_inputs: a dictionary with "point_coords" and "point_labels", where + 1) "point_coords" has [B, P, 2] shape and float32 dtype and contains the + absolute pixel-unit coordinate in (x, y) format of the P input points + 2) "point_labels" has shape [B, P] and int32 dtype, where 1 means + positive clicks, 0 means negative clicks, and -1 means padding + - mask_inputs: a mask of [B, 1, H*16, W*16] shape, float or bool, with the + same spatial size as the image. + - high_res_features: either 1) None or 2) or a list of length 2 containing + two feature maps of [B, C, 4*H, 4*W] and [B, C, 2*H, 2*W] shapes respectively, + which will be used as high-resolution feature maps for SAM decoder. + - multimask_output: if it's True, we output 3 candidate masks and their 3 + corresponding IoU estimates, and if it's False, we output only 1 mask and + its corresponding IoU estimate. + + Outputs: + - low_res_multimasks: [B, M, H*4, W*4] shape (where M = 3 if + `multimask_output=True` and M = 1 if `multimask_output=False`), the SAM + output mask logits (before sigmoid) for the low-resolution masks, with 4x + the resolution (1/4 stride) of the input backbone_features. + - high_res_multimasks: [B, M, H*16, W*16] shape (where M = 3 + if `multimask_output=True` and M = 1 if `multimask_output=False`), + upsampled from the low-resolution masks, with shape size as the image + (stride is 1 pixel). + - ious, [B, M] shape, where (where M = 3 if `multimask_output=True` and M = 1 + if `multimask_output=False`), the estimated IoU of each output mask. + - low_res_masks: [B, 1, H*4, W*4] shape, the best mask in `low_res_multimasks`. + If `multimask_output=True`, it's the mask with the highest IoU estimate. + If `multimask_output=False`, it's the same as `low_res_multimasks`. + - high_res_masks: [B, 1, H*16, W*16] shape, the best mask in `high_res_multimasks`. + If `multimask_output=True`, it's the mask with the highest IoU estimate. + If `multimask_output=False`, it's the same as `high_res_multimasks`. + - obj_ptr: [B, C] shape, the object pointer vector for the output mask, extracted + based on the output token from the SAM mask decoder. + """ + B = backbone_features.size(0) + device = backbone_features.device + assert backbone_features.size(1) == self.sam_prompt_embed_dim + assert backbone_features.size(2) == self.sam_image_embedding_size + assert backbone_features.size(3) == self.sam_image_embedding_size + + # a) Handle point prompts + if point_inputs is not None: + sam_point_coords = point_inputs["point_coords"] + sam_point_labels = point_inputs["point_labels"] + assert sam_point_coords.size(0) == B and sam_point_labels.size(0) == B + else: + # If no points are provide, pad with an empty point (with label -1) + sam_point_coords = torch.zeros(B, 1, 2, device=device) + sam_point_labels = -torch.ones(B, 1, dtype=torch.int32, device=device) + + # b) Handle mask prompts + if mask_inputs is not None: + # If mask_inputs is provided, downsize it into low-res mask input if needed + # and feed it as a dense mask prompt into the SAM mask encoder + assert len(mask_inputs.shape) == 4 and mask_inputs.shape[:2] == (B, 1) + if mask_inputs.shape[-2:] != self.sam_prompt_encoder.mask_input_size: + sam_mask_prompt = F.interpolate( + mask_inputs.float(), + size=self.sam_prompt_encoder.mask_input_size, + align_corners=False, + mode="bilinear", + antialias=True, # use antialias for downsampling + ) + else: + sam_mask_prompt = mask_inputs + else: + # Otherwise, simply feed None (and SAM's prompt encoder will add + # a learned `no_mask_embed` to indicate no mask input in this case). + sam_mask_prompt = None + + sparse_embeddings, dense_embeddings = self.sam_prompt_encoder( + points=(sam_point_coords, sam_point_labels), + boxes=None, + masks=sam_mask_prompt, + ) + ( + low_res_multimasks, + ious, + sam_output_tokens, + object_score_logits, + ) = self.sam_mask_decoder( + image_embeddings=backbone_features, + image_pe=self.sam_prompt_encoder.get_dense_pe(), + sparse_prompt_embeddings=sparse_embeddings, + dense_prompt_embeddings=dense_embeddings, + multimask_output=multimask_output, + repeat_image=False, # the image is already batched + high_res_features=high_res_features, + ) + if self.pred_obj_scores: + is_obj_appearing = object_score_logits > 0 + + # Mask used for spatial memories is always a *hard* choice between obj and no obj, + # consistent with the actual mask prediction + low_res_multimasks = torch.where( + is_obj_appearing[:, None, None], + low_res_multimasks, + NO_OBJ_SCORE, + ) + + # convert masks from possibly bfloat16 (or float16) to float32 + # (older PyTorch versions before 2.1 don't support `interpolate` on bf16) + low_res_multimasks = low_res_multimasks.float() + high_res_multimasks = F.interpolate( + low_res_multimasks, + size=(self.image_size, self.image_size), + mode="bilinear", + align_corners=False, + ) + + sam_output_token = sam_output_tokens[:, 0] + if multimask_output: + # take the best mask prediction (with the highest IoU estimation) + best_iou_inds = torch.argmax(ious, dim=-1) + batch_inds = torch.arange(B, device=device) + low_res_masks = low_res_multimasks[batch_inds, best_iou_inds].unsqueeze(1) + high_res_masks = high_res_multimasks[batch_inds, best_iou_inds].unsqueeze(1) + if sam_output_tokens.size(1) > 1: + sam_output_token = sam_output_tokens[batch_inds, best_iou_inds] + else: + low_res_masks, high_res_masks = low_res_multimasks, high_res_multimasks + + # Extract object pointer from the SAM output token (with occlusion handling) + obj_ptr = self.obj_ptr_proj(sam_output_token) + if self.pred_obj_scores: + # Allow *soft* no obj ptr, unlike for masks + if self.soft_no_obj_ptr: + lambda_is_obj_appearing = object_score_logits.sigmoid() + else: + lambda_is_obj_appearing = is_obj_appearing.float() + + if self.fixed_no_obj_ptr: + obj_ptr = lambda_is_obj_appearing * obj_ptr + obj_ptr = obj_ptr + (1 - lambda_is_obj_appearing) * self.no_obj_ptr + + return ( + low_res_multimasks, + high_res_multimasks, + ious, + low_res_masks, + high_res_masks, + obj_ptr, + object_score_logits, + ) + + def _use_mask_as_output(self, backbone_features, high_res_features, mask_inputs): + """ + Directly turn binary `mask_inputs` into a output mask logits without using SAM. + (same input and output shapes as in _forward_sam_heads above). + """ + # Use -10/+10 as logits for neg/pos pixels (very close to 0/1 in prob after sigmoid). + out_scale, out_bias = 20.0, -10.0 # sigmoid(-10.0)=4.5398e-05 + mask_inputs_float = mask_inputs.float() + high_res_masks = mask_inputs_float * out_scale + out_bias + low_res_masks = F.interpolate( + high_res_masks, + size=(high_res_masks.size(-2) // 4, high_res_masks.size(-1) // 4), + align_corners=False, + mode="bilinear", + antialias=True, # use antialias for downsampling + ) + # a dummy IoU prediction of all 1's under mask input + ious = mask_inputs.new_ones(mask_inputs.size(0), 1).float() + if not self.use_obj_ptrs_in_encoder: + # all zeros as a dummy object pointer (of shape [B, C]) + obj_ptr = torch.zeros( + mask_inputs.size(0), self.hidden_dim, device=mask_inputs.device + ) + else: + # produce an object pointer using the SAM decoder from the mask input + _, _, _, _, _, obj_ptr, _ = self._forward_sam_heads( + backbone_features=backbone_features, + mask_inputs=self.mask_downsample(mask_inputs_float), + high_res_features=high_res_features, + ) + # In this method, we are treating mask_input as output, e.g. using it directly to create spatial mem; + # Below, we follow the same design axiom to use mask_input to decide if obj appears or not instead of relying + # on the object_scores from the SAM decoder. + is_obj_appearing = torch.any(mask_inputs.flatten(1).float() > 0.0, dim=1) + is_obj_appearing = is_obj_appearing[..., None] + lambda_is_obj_appearing = is_obj_appearing.float() + object_score_logits = out_scale * lambda_is_obj_appearing + out_bias + if self.pred_obj_scores: + if self.fixed_no_obj_ptr: + obj_ptr = lambda_is_obj_appearing * obj_ptr + obj_ptr = obj_ptr + (1 - lambda_is_obj_appearing) * self.no_obj_ptr + + return ( + low_res_masks, + high_res_masks, + ious, + low_res_masks, + high_res_masks, + obj_ptr, + object_score_logits, + ) + + def forward_image(self, img_batch: torch.Tensor): + """Get the image feature on the input batch.""" + backbone_out = self.image_encoder(img_batch) + if self.use_high_res_features_in_sam: + # precompute projected level 0 and level 1 features in SAM decoder + # to avoid running it again on every SAM click + backbone_out["backbone_fpn"][0] = self.sam_mask_decoder.conv_s0( + backbone_out["backbone_fpn"][0] + ) + backbone_out["backbone_fpn"][1] = self.sam_mask_decoder.conv_s1( + backbone_out["backbone_fpn"][1] + ) + return backbone_out + + def _prepare_backbone_features(self, backbone_out): + """Prepare and flatten visual features.""" + backbone_out = backbone_out.copy() + assert len(backbone_out["backbone_fpn"]) == len(backbone_out["vision_pos_enc"]) + assert len(backbone_out["backbone_fpn"]) >= self.num_feature_levels + + feature_maps = backbone_out["backbone_fpn"][-self.num_feature_levels :] + vision_pos_embeds = backbone_out["vision_pos_enc"][-self.num_feature_levels :] + + feat_sizes = [(x.shape[-2], x.shape[-1]) for x in vision_pos_embeds] + # flatten NxCxHxW to HWxNxC + vision_feats = [x.flatten(2).permute(2, 0, 1) for x in feature_maps] + vision_pos_embeds = [x.flatten(2).permute(2, 0, 1) for x in vision_pos_embeds] + + return backbone_out, vision_feats, vision_pos_embeds, feat_sizes + + def _prepare_memory_conditioned_features( + self, + frame_idx, + is_init_cond_frame, + current_vision_feats, + current_vision_pos_embeds, + feat_sizes, + output_dict, + num_frames, + track_in_reverse=False, # tracking in reverse time order (for demo usage) + ): + """Fuse the current frame's visual feature map with previous memory.""" + B = current_vision_feats[-1].size(1) # batch size on this frame + C = self.hidden_dim + H, W = feat_sizes[-1] # top-level (lowest-resolution) feature size + device = current_vision_feats[-1].device + # The case of `self.num_maskmem == 0` below is primarily used for reproducing SAM on images. + # In this case, we skip the fusion with any memory. + if self.num_maskmem == 0: # Disable memory and skip fusion + pix_feat = current_vision_feats[-1].permute(1, 2, 0).view(B, C, H, W) + return pix_feat + + num_obj_ptr_tokens = 0 + tpos_sign_mul = -1 if track_in_reverse else 1 + # Step 1: condition the visual features of the current frame on previous memories + if not is_init_cond_frame: + # Retrieve the memories encoded with the maskmem backbone + to_cat_memory, to_cat_memory_pos_embed = [], [] + # Add conditioning frames's output first (all cond frames have t_pos=0 for + # when getting temporal positional embedding below) + assert len(output_dict["cond_frame_outputs"]) > 0 + # Select a maximum number of temporally closest cond frames for cross attention + cond_outputs = output_dict["cond_frame_outputs"] + selected_cond_outputs, unselected_cond_outputs = select_closest_cond_frames( + frame_idx, cond_outputs, self.max_cond_frames_in_attn + ) + t_pos_and_prevs = [(0, out) for out in selected_cond_outputs.values()] + # Add last (self.num_maskmem - 1) frames before current frame for non-conditioning memory + # the earliest one has t_pos=1 and the latest one has t_pos=self.num_maskmem-1 + # We also allow taking the memory frame non-consecutively (with stride>1), in which case + # we take (self.num_maskmem - 2) frames among every stride-th frames plus the last frame. + stride = 1 if self.training else self.memory_temporal_stride_for_eval + for t_pos in range(1, self.num_maskmem): + t_rel = self.num_maskmem - t_pos # how many frames before current frame + if t_rel == 1: + # for t_rel == 1, we take the last frame (regardless of r) + if not track_in_reverse: + # the frame immediately before this frame (i.e. frame_idx - 1) + prev_frame_idx = frame_idx - t_rel + else: + # the frame immediately after this frame (i.e. frame_idx + 1) + prev_frame_idx = frame_idx + t_rel + else: + # for t_rel >= 2, we take the memory frame from every r-th frames + if not track_in_reverse: + # first find the nearest frame among every r-th frames before this frame + # for r=1, this would be (frame_idx - 2) + prev_frame_idx = ((frame_idx - 2) // stride) * stride + # then seek further among every r-th frames + prev_frame_idx = prev_frame_idx - (t_rel - 2) * stride + else: + # first find the nearest frame among every r-th frames after this frame + # for r=1, this would be (frame_idx + 2) + prev_frame_idx = -(-(frame_idx + 2) // stride) * stride + # then seek further among every r-th frames + prev_frame_idx = prev_frame_idx + (t_rel - 2) * stride + out = output_dict["non_cond_frame_outputs"].get(prev_frame_idx, None) + if out is None: + # If an unselected conditioning frame is among the last (self.num_maskmem - 1) + # frames, we still attend to it as if it's a non-conditioning frame. + out = unselected_cond_outputs.get(prev_frame_idx, None) + t_pos_and_prevs.append((t_pos, out)) + + for t_pos, prev in t_pos_and_prevs: + if prev is None: + continue # skip padding frames + # "maskmem_features" might have been offloaded to CPU in demo use cases, + # so we load it back to GPU (it's a no-op if it's already on GPU). + feats = prev["maskmem_features"].to(device, non_blocking=True) + to_cat_memory.append(feats.flatten(2).permute(2, 0, 1)) + # Spatial positional encoding (it might have been offloaded to CPU in eval) + maskmem_enc = prev["maskmem_pos_enc"][-1].to(device) + maskmem_enc = maskmem_enc.flatten(2).permute(2, 0, 1) + # Temporal positional encoding + maskmem_enc = ( + maskmem_enc + self.maskmem_tpos_enc[self.num_maskmem - t_pos - 1] + ) + to_cat_memory_pos_embed.append(maskmem_enc) + + # Construct the list of past object pointers + if self.use_obj_ptrs_in_encoder: + max_obj_ptrs_in_encoder = min(num_frames, self.max_obj_ptrs_in_encoder) + # First add those object pointers from selected conditioning frames + # (optionally, only include object pointers in the past during evaluation) + if not self.training and self.only_obj_ptrs_in_the_past_for_eval: + ptr_cond_outputs = { + t: out + for t, out in selected_cond_outputs.items() + if (t >= frame_idx if track_in_reverse else t <= frame_idx) + } + else: + ptr_cond_outputs = selected_cond_outputs + pos_and_ptrs = [ + # Temporal pos encoding contains how far away each pointer is from current frame + ( + ( + (frame_idx - t) * tpos_sign_mul + if self.use_signed_tpos_enc_to_obj_ptrs + else abs(frame_idx - t) + ), + out["obj_ptr"], + ) + for t, out in ptr_cond_outputs.items() + ] + # Add up to (max_obj_ptrs_in_encoder - 1) non-conditioning frames before current frame + for t_diff in range(1, max_obj_ptrs_in_encoder): + t = frame_idx + t_diff if track_in_reverse else frame_idx - t_diff + if t < 0 or (num_frames is not None and t >= num_frames): + break + out = output_dict["non_cond_frame_outputs"].get( + t, unselected_cond_outputs.get(t, None) + ) + if out is not None: + pos_and_ptrs.append((t_diff, out["obj_ptr"])) + # If we have at least one object pointer, add them to the across attention + if len(pos_and_ptrs) > 0: + pos_list, ptrs_list = zip(*pos_and_ptrs) + # stack object pointers along dim=0 into [ptr_seq_len, B, C] shape + obj_ptrs = torch.stack(ptrs_list, dim=0) + # a temporal positional embedding based on how far each object pointer is from + # the current frame (sine embedding normalized by the max pointer num). + if self.add_tpos_enc_to_obj_ptrs: + t_diff_max = max_obj_ptrs_in_encoder - 1 + tpos_dim = C if self.proj_tpos_enc_in_obj_ptrs else self.mem_dim + obj_pos = torch.tensor(pos_list).to( + device=device, non_blocking=True + ) + obj_pos = get_1d_sine_pe(obj_pos / t_diff_max, dim=tpos_dim) + obj_pos = self.obj_ptr_tpos_proj(obj_pos) + obj_pos = obj_pos.unsqueeze(1).expand(-1, B, self.mem_dim) + else: + obj_pos = obj_ptrs.new_zeros(len(pos_list), B, self.mem_dim) + if self.mem_dim < C: + # split a pointer into (C // self.mem_dim) tokens for self.mem_dim < C + obj_ptrs = obj_ptrs.reshape( + -1, B, C // self.mem_dim, self.mem_dim + ) + obj_ptrs = obj_ptrs.permute(0, 2, 1, 3).flatten(0, 1) + obj_pos = obj_pos.repeat_interleave(C // self.mem_dim, dim=0) + to_cat_memory.append(obj_ptrs) + to_cat_memory_pos_embed.append(obj_pos) + num_obj_ptr_tokens = obj_ptrs.shape[0] + else: + num_obj_ptr_tokens = 0 + else: + # for initial conditioning frames, encode them without using any previous memory + if self.directly_add_no_mem_embed: + # directly add no-mem embedding (instead of using the transformer encoder) + pix_feat_with_mem = current_vision_feats[-1] + self.no_mem_embed + pix_feat_with_mem = pix_feat_with_mem.permute(1, 2, 0).view(B, C, H, W) + return pix_feat_with_mem + + # Use a dummy token on the first frame (to avoid empty memory input to tranformer encoder) + to_cat_memory = [self.no_mem_embed.expand(1, B, self.mem_dim)] + to_cat_memory_pos_embed = [self.no_mem_pos_enc.expand(1, B, self.mem_dim)] + + # Step 2: Concatenate the memories and forward through the transformer encoder + memory = torch.cat(to_cat_memory, dim=0) + memory_pos_embed = torch.cat(to_cat_memory_pos_embed, dim=0) + + pix_feat_with_mem = self.memory_attention( + curr=current_vision_feats, + curr_pos=current_vision_pos_embeds, + memory=memory, + memory_pos=memory_pos_embed, + num_obj_ptr_tokens=num_obj_ptr_tokens, + ) + # reshape the output (HW)BC => BCHW + pix_feat_with_mem = pix_feat_with_mem.permute(1, 2, 0).view(B, C, H, W) + return pix_feat_with_mem + + def _encode_new_memory( + self, + current_vision_feats, + feat_sizes, + pred_masks_high_res, + object_score_logits, + is_mask_from_pts, + ): + """Encode the current image and its prediction into a memory feature.""" + B = current_vision_feats[-1].size(1) # batch size on this frame + C = self.hidden_dim + H, W = feat_sizes[-1] # top-level (lowest-resolution) feature size + # top-level feature, (HW)BC => BCHW + pix_feat = current_vision_feats[-1].permute(1, 2, 0).view(B, C, H, W) + if self.non_overlap_masks_for_mem_enc and not self.training: + # optionally, apply non-overlapping constraints to the masks (it's applied + # in the batch dimension and should only be used during eval, where all + # the objects come from the same video under batch size 1). + pred_masks_high_res = self._apply_non_overlapping_constraints( + pred_masks_high_res + ) + # scale the raw mask logits with a temperature before applying sigmoid + binarize = self.binarize_mask_from_pts_for_mem_enc and is_mask_from_pts + if binarize and not self.training: + mask_for_mem = (pred_masks_high_res > 0).float() + else: + # apply sigmoid on the raw mask logits to turn them into range (0, 1) + mask_for_mem = torch.sigmoid(pred_masks_high_res) + # apply scale and bias terms to the sigmoid probabilities + if self.sigmoid_scale_for_mem_enc != 1.0: + mask_for_mem = mask_for_mem * self.sigmoid_scale_for_mem_enc + if self.sigmoid_bias_for_mem_enc != 0.0: + mask_for_mem = mask_for_mem + self.sigmoid_bias_for_mem_enc + maskmem_out = self.memory_encoder( + pix_feat, mask_for_mem, skip_mask_sigmoid=True # sigmoid already applied + ) + maskmem_features = maskmem_out["vision_features"] + maskmem_pos_enc = maskmem_out["vision_pos_enc"] + # add a no-object embedding to the spatial memory to indicate that the frame + # is predicted to be occluded (i.e. no object is appearing in the frame) + if self.no_obj_embed_spatial is not None: + is_obj_appearing = (object_score_logits > 0).float() + maskmem_features += ( + 1 - is_obj_appearing[..., None, None] + ) * self.no_obj_embed_spatial[..., None, None].expand( + *maskmem_features.shape + ) + + return maskmem_features, maskmem_pos_enc + + def _track_step( + self, + frame_idx, + is_init_cond_frame, + current_vision_feats, + current_vision_pos_embeds, + feat_sizes, + point_inputs, + mask_inputs, + output_dict, + num_frames, + track_in_reverse, + prev_sam_mask_logits, + ): + current_out = {"point_inputs": point_inputs, "mask_inputs": mask_inputs} + # High-resolution feature maps for the SAM head, reshape (HW)BC => BCHW + if len(current_vision_feats) > 1: + high_res_features = [ + x.permute(1, 2, 0).view(x.size(1), x.size(2), *s) + for x, s in zip(current_vision_feats[:-1], feat_sizes[:-1]) + ] + else: + high_res_features = None + if mask_inputs is not None and self.use_mask_input_as_output_without_sam: + # When use_mask_input_as_output_without_sam=True, we directly output the mask input + # (see it as a GT mask) without using a SAM prompt encoder + mask decoder. + pix_feat = current_vision_feats[-1].permute(1, 2, 0) + pix_feat = pix_feat.view(-1, self.hidden_dim, *feat_sizes[-1]) + sam_outputs = self._use_mask_as_output( + pix_feat, high_res_features, mask_inputs + ) + else: + # fused the visual feature with previous memory features in the memory bank + pix_feat = self._prepare_memory_conditioned_features( + frame_idx=frame_idx, + is_init_cond_frame=is_init_cond_frame, + current_vision_feats=current_vision_feats[-1:], + current_vision_pos_embeds=current_vision_pos_embeds[-1:], + feat_sizes=feat_sizes[-1:], + output_dict=output_dict, + num_frames=num_frames, + track_in_reverse=track_in_reverse, + ) + # apply SAM-style segmentation head + # here we might feed previously predicted low-res SAM mask logits into the SAM mask decoder, + # e.g. in demo where such logits come from earlier interaction instead of correction sampling + # (in this case, any `mask_inputs` shouldn't reach here as they are sent to _use_mask_as_output instead) + if prev_sam_mask_logits is not None: + assert point_inputs is not None and mask_inputs is None + mask_inputs = prev_sam_mask_logits + multimask_output = self._use_multimask(is_init_cond_frame, point_inputs) + sam_outputs = self._forward_sam_heads( + backbone_features=pix_feat, + point_inputs=point_inputs, + mask_inputs=mask_inputs, + high_res_features=high_res_features, + multimask_output=multimask_output, + ) + + return current_out, sam_outputs, high_res_features, pix_feat + + def _encode_memory_in_output( + self, + current_vision_feats, + feat_sizes, + point_inputs, + run_mem_encoder, + high_res_masks, + object_score_logits, + current_out, + ): + if run_mem_encoder and self.num_maskmem > 0: + high_res_masks_for_mem_enc = high_res_masks + maskmem_features, maskmem_pos_enc = self._encode_new_memory( + current_vision_feats=current_vision_feats, + feat_sizes=feat_sizes, + pred_masks_high_res=high_res_masks_for_mem_enc, + object_score_logits=object_score_logits, + is_mask_from_pts=(point_inputs is not None), + ) + current_out["maskmem_features"] = maskmem_features + current_out["maskmem_pos_enc"] = maskmem_pos_enc + else: + current_out["maskmem_features"] = None + current_out["maskmem_pos_enc"] = None + + def track_step( + self, + frame_idx, + is_init_cond_frame, + current_vision_feats, + current_vision_pos_embeds, + feat_sizes, + point_inputs, + mask_inputs, + output_dict, + num_frames, + track_in_reverse=False, # tracking in reverse time order (for demo usage) + # Whether to run the memory encoder on the predicted masks. Sometimes we might want + # to skip the memory encoder with `run_mem_encoder=False`. For example, + # in demo we might call `track_step` multiple times for each user click, + # and only encode the memory when the user finalizes their clicks. And in ablation + # settings like SAM training on static images, we don't need the memory encoder. + run_mem_encoder=True, + # The previously predicted SAM mask logits (which can be fed together with new clicks in demo). + prev_sam_mask_logits=None, + ): + current_out, sam_outputs, _, _ = self._track_step( + frame_idx, + is_init_cond_frame, + current_vision_feats, + current_vision_pos_embeds, + feat_sizes, + point_inputs, + mask_inputs, + output_dict, + num_frames, + track_in_reverse, + prev_sam_mask_logits, + ) + + ( + _, + _, + _, + low_res_masks, + high_res_masks, + obj_ptr, + object_score_logits, + ) = sam_outputs + + current_out["pred_masks"] = low_res_masks + current_out["pred_masks_high_res"] = high_res_masks + current_out["obj_ptr"] = obj_ptr + if not self.training: + # Only add this in inference (to avoid unused param in activation checkpointing; + # it's mainly used in the demo to encode spatial memories w/ consolidated masks) + current_out["object_score_logits"] = object_score_logits + + # Finally run the memory encoder on the predicted mask to encode + # it into a new memory feature (that can be used in future frames) + self._encode_memory_in_output( + current_vision_feats, + feat_sizes, + point_inputs, + run_mem_encoder, + high_res_masks, + object_score_logits, + current_out, + ) + + return current_out + + def _use_multimask(self, is_init_cond_frame, point_inputs): + """Whether to use multimask output in the SAM head.""" + num_pts = 0 if point_inputs is None else point_inputs["point_labels"].size(1) + multimask_output = ( + self.multimask_output_in_sam + and (is_init_cond_frame or self.multimask_output_for_tracking) + and (self.multimask_min_pt_num <= num_pts <= self.multimask_max_pt_num) + ) + return multimask_output + + def _apply_non_overlapping_constraints(self, pred_masks): + """ + Apply non-overlapping constraints to the object scores in pred_masks. Here we + keep only the highest scoring object at each spatial location in pred_masks. + """ + batch_size = pred_masks.size(0) + if batch_size == 1: + return pred_masks + + device = pred_masks.device + # "max_obj_inds": object index of the object with the highest score at each location + max_obj_inds = torch.argmax(pred_masks, dim=0, keepdim=True) + # "batch_obj_inds": object index of each object slice (along dim 0) in `pred_masks` + batch_obj_inds = torch.arange(batch_size, device=device)[:, None, None, None] + keep = max_obj_inds == batch_obj_inds + # suppress overlapping regions' scores below -10.0 so that the foreground regions + # don't overlap (here sigmoid(-10.0)=4.5398e-05) + pred_masks = torch.where(keep, pred_masks, torch.clamp(pred_masks, max=-10.0)) + return pred_masks diff --git a/freevideo_engine/sam2/modeling/sam2_utils.py b/freevideo_engine/sam2/modeling/sam2_utils.py new file mode 100644 index 00000000..009e9c1e --- /dev/null +++ b/freevideo_engine/sam2/modeling/sam2_utils.py @@ -0,0 +1,323 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. + +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. + + +import copy +from typing import Tuple + +import numpy as np +import torch +import torch.nn as nn +import torch.nn.functional as F + +from freevideo_engine.sam2.utils.misc import mask_to_box + + +def select_closest_cond_frames(frame_idx, cond_frame_outputs, max_cond_frame_num): + """ + Select up to `max_cond_frame_num` conditioning frames from `cond_frame_outputs` + that are temporally closest to the current frame at `frame_idx`. Here, we take + - a) the closest conditioning frame before `frame_idx` (if any); + - b) the closest conditioning frame after `frame_idx` (if any); + - c) any other temporally closest conditioning frames until reaching a total + of `max_cond_frame_num` conditioning frames. + + Outputs: + - selected_outputs: selected items (keys & values) from `cond_frame_outputs`. + - unselected_outputs: items (keys & values) not selected in `cond_frame_outputs`. + """ + if max_cond_frame_num == -1 or len(cond_frame_outputs) <= max_cond_frame_num: + selected_outputs = cond_frame_outputs + unselected_outputs = {} + else: + assert max_cond_frame_num >= 2, "we should allow using 2+ conditioning frames" + selected_outputs = {} + + # the closest conditioning frame before `frame_idx` (if any) + idx_before = max((t for t in cond_frame_outputs if t < frame_idx), default=None) + if idx_before is not None: + selected_outputs[idx_before] = cond_frame_outputs[idx_before] + + # the closest conditioning frame after `frame_idx` (if any) + idx_after = min((t for t in cond_frame_outputs if t >= frame_idx), default=None) + if idx_after is not None: + selected_outputs[idx_after] = cond_frame_outputs[idx_after] + + # add other temporally closest conditioning frames until reaching a total + # of `max_cond_frame_num` conditioning frames. + num_remain = max_cond_frame_num - len(selected_outputs) + inds_remain = sorted( + (t for t in cond_frame_outputs if t not in selected_outputs), + key=lambda x: abs(x - frame_idx), + )[:num_remain] + selected_outputs.update((t, cond_frame_outputs[t]) for t in inds_remain) + unselected_outputs = { + t: v for t, v in cond_frame_outputs.items() if t not in selected_outputs + } + + return selected_outputs, unselected_outputs + + +def get_1d_sine_pe(pos_inds, dim, temperature=10000): + """ + Get 1D sine positional embedding as in the original Transformer paper. + """ + pe_dim = dim // 2 + dim_t = torch.arange(pe_dim, dtype=torch.float32, device=pos_inds.device) + dim_t = temperature ** (2 * (dim_t // 2) / pe_dim) + + pos_embed = pos_inds.unsqueeze(-1) / dim_t + pos_embed = torch.cat([pos_embed.sin(), pos_embed.cos()], dim=-1) + return pos_embed + + +def get_activation_fn(activation): + """Return an activation function given a string""" + if activation == "relu": + return F.relu + if activation == "gelu": + return F.gelu + if activation == "glu": + return F.glu + raise RuntimeError(f"activation should be relu/gelu, not {activation}.") + + +def get_clones(module, N): + return nn.ModuleList([copy.deepcopy(module) for i in range(N)]) + + +class DropPath(nn.Module): + # adapted from https://github.com/huggingface/pytorch-image-models/blob/main/timm/layers/drop.py + def __init__(self, drop_prob=0.0, scale_by_keep=True): + super(DropPath, self).__init__() + self.drop_prob = drop_prob + self.scale_by_keep = scale_by_keep + + def forward(self, x): + if self.drop_prob == 0.0 or not self.training: + return x + keep_prob = 1 - self.drop_prob + shape = (x.shape[0],) + (1,) * (x.ndim - 1) + random_tensor = x.new_empty(shape).bernoulli_(keep_prob) + if keep_prob > 0.0 and self.scale_by_keep: + random_tensor.div_(keep_prob) + return x * random_tensor + + +# Lightly adapted from +# https://github.com/facebookresearch/MaskFormer/blob/main/mask_former/modeling/transformer/transformer_predictor.py # noqa +class MLP(nn.Module): + def __init__( + self, + input_dim: int, + hidden_dim: int, + output_dim: int, + num_layers: int, + activation: nn.Module = nn.ReLU, + sigmoid_output: bool = False, + ) -> None: + super().__init__() + self.num_layers = num_layers + h = [hidden_dim] * (num_layers - 1) + self.layers = nn.ModuleList( + nn.Linear(n, k) for n, k in zip([input_dim] + h, h + [output_dim]) + ) + self.sigmoid_output = sigmoid_output + self.act = activation() + + def forward(self, x): + for i, layer in enumerate(self.layers): + x = self.act(layer(x)) if i < self.num_layers - 1 else layer(x) + if self.sigmoid_output: + x = F.sigmoid(x) + return x + + +# From https://github.com/facebookresearch/detectron2/blob/main/detectron2/layers/batch_norm.py # noqa +# Itself from https://github.com/facebookresearch/ConvNeXt/blob/d1fa8f6fef0a165b27399986cc2bdacc92777e40/models/convnext.py#L119 # noqa +class LayerNorm2d(nn.Module): + def __init__(self, num_channels: int, eps: float = 1e-6) -> None: + super().__init__() + self.weight = nn.Parameter(torch.ones(num_channels)) + self.bias = nn.Parameter(torch.zeros(num_channels)) + self.eps = eps + + def forward(self, x: torch.Tensor) -> torch.Tensor: + u = x.mean(1, keepdim=True) + s = (x - u).pow(2).mean(1, keepdim=True) + x = (x - u) / torch.sqrt(s + self.eps) + x = self.weight[:, None, None] * x + self.bias[:, None, None] + return x + + +def sample_box_points( + masks: torch.Tensor, + noise: float = 0.1, # SAM default + noise_bound: int = 20, # SAM default + top_left_label: int = 2, + bottom_right_label: int = 3, +) -> Tuple[np.array, np.array]: + """ + Sample a noised version of the top left and bottom right corners of a given `bbox` + + Inputs: + - masks: [B, 1, H,W] boxes, dtype=torch.Tensor + - noise: noise as a fraction of box width and height, dtype=float + - noise_bound: maximum amount of noise (in pure pixesl), dtype=int + + Returns: + - box_coords: [B, num_pt, 2], contains (x, y) coordinates of top left and bottom right box corners, dtype=torch.float + - box_labels: [B, num_pt], label 2 is reserverd for top left and 3 for bottom right corners, dtype=torch.int32 + """ + device = masks.device + box_coords = mask_to_box(masks) + B, _, H, W = masks.shape + box_labels = torch.tensor( + [top_left_label, bottom_right_label], dtype=torch.int, device=device + ).repeat(B) + if noise > 0.0: + if not isinstance(noise_bound, torch.Tensor): + noise_bound = torch.tensor(noise_bound, device=device) + bbox_w = box_coords[..., 2] - box_coords[..., 0] + bbox_h = box_coords[..., 3] - box_coords[..., 1] + max_dx = torch.min(bbox_w * noise, noise_bound) + max_dy = torch.min(bbox_h * noise, noise_bound) + box_noise = 2 * torch.rand(B, 1, 4, device=device) - 1 + box_noise = box_noise * torch.stack((max_dx, max_dy, max_dx, max_dy), dim=-1) + + box_coords = box_coords + box_noise + img_bounds = ( + torch.tensor([W, H, W, H], device=device) - 1 + ) # uncentered pixel coords + box_coords.clamp_(torch.zeros_like(img_bounds), img_bounds) # In place clamping + + box_coords = box_coords.reshape(-1, 2, 2) # always 2 points + box_labels = box_labels.reshape(-1, 2) + return box_coords, box_labels + + +def sample_random_points_from_errors(gt_masks, pred_masks, num_pt=1): + """ + Sample `num_pt` random points (along with their labels) independently from the error regions. + + Inputs: + - gt_masks: [B, 1, H_im, W_im] masks, dtype=torch.bool + - pred_masks: [B, 1, H_im, W_im] masks, dtype=torch.bool or None + - num_pt: int, number of points to sample independently for each of the B error maps + + Outputs: + - points: [B, num_pt, 2], dtype=torch.float, contains (x, y) coordinates of each sampled point + - labels: [B, num_pt], dtype=torch.int32, where 1 means positive clicks and 0 means + negative clicks + """ + if pred_masks is None: # if pred_masks is not provided, treat it as empty + pred_masks = torch.zeros_like(gt_masks) + assert gt_masks.dtype == torch.bool and gt_masks.size(1) == 1 + assert pred_masks.dtype == torch.bool and pred_masks.shape == gt_masks.shape + assert num_pt >= 0 + + B, _, H_im, W_im = gt_masks.shape + device = gt_masks.device + + # false positive region, a new point sampled in this region should have + # negative label to correct the FP error + fp_masks = ~gt_masks & pred_masks + # false negative region, a new point sampled in this region should have + # positive label to correct the FN error + fn_masks = gt_masks & ~pred_masks + # whether the prediction completely match the ground-truth on each mask + all_correct = torch.all((gt_masks == pred_masks).flatten(2), dim=2) + all_correct = all_correct[..., None, None] + + # channel 0 is FP map, while channel 1 is FN map + pts_noise = torch.rand(B, num_pt, H_im, W_im, 2, device=device) + # sample a negative new click from FP region or a positive new click + # from FN region, depend on where the maximum falls, + # and in case the predictions are all correct (no FP or FN), we just + # sample a negative click from the background region + pts_noise[..., 0] *= fp_masks | (all_correct & ~gt_masks) + pts_noise[..., 1] *= fn_masks + pts_idx = pts_noise.flatten(2).argmax(dim=2) + labels = (pts_idx % 2).to(torch.int32) + pts_idx = pts_idx // 2 + pts_x = pts_idx % W_im + pts_y = pts_idx // W_im + points = torch.stack([pts_x, pts_y], dim=2).to(torch.float) + return points, labels + + +def sample_one_point_from_error_center(gt_masks, pred_masks, padding=True): + """ + Sample 1 random point (along with its label) from the center of each error region, + that is, the point with the largest distance to the boundary of each error region. + This is the RITM sampling method from https://github.com/saic-vul/ritm_interactive_segmentation/blob/master/isegm/inference/clicker.py + + Inputs: + - gt_masks: [B, 1, H_im, W_im] masks, dtype=torch.bool + - pred_masks: [B, 1, H_im, W_im] masks, dtype=torch.bool or None + - padding: if True, pad with boundary of 1 px for distance transform + + Outputs: + - points: [B, 1, 2], dtype=torch.float, contains (x, y) coordinates of each sampled point + - labels: [B, 1], dtype=torch.int32, where 1 means positive clicks and 0 means negative clicks + """ + import cv2 + + if pred_masks is None: + pred_masks = torch.zeros_like(gt_masks) + assert gt_masks.dtype == torch.bool and gt_masks.size(1) == 1 + assert pred_masks.dtype == torch.bool and pred_masks.shape == gt_masks.shape + + B, _, _, W_im = gt_masks.shape + device = gt_masks.device + + # false positive region, a new point sampled in this region should have + # negative label to correct the FP error + fp_masks = ~gt_masks & pred_masks + # false negative region, a new point sampled in this region should have + # positive label to correct the FN error + fn_masks = gt_masks & ~pred_masks + + fp_masks = fp_masks.cpu().numpy() + fn_masks = fn_masks.cpu().numpy() + points = torch.zeros(B, 1, 2, dtype=torch.float) + labels = torch.ones(B, 1, dtype=torch.int32) + for b in range(B): + fn_mask = fn_masks[b, 0] + fp_mask = fp_masks[b, 0] + if padding: + fn_mask = np.pad(fn_mask, ((1, 1), (1, 1)), "constant") + fp_mask = np.pad(fp_mask, ((1, 1), (1, 1)), "constant") + # compute the distance of each point in FN/FP region to its boundary + fn_mask_dt = cv2.distanceTransform(fn_mask.astype(np.uint8), cv2.DIST_L2, 0) + fp_mask_dt = cv2.distanceTransform(fp_mask.astype(np.uint8), cv2.DIST_L2, 0) + if padding: + fn_mask_dt = fn_mask_dt[1:-1, 1:-1] + fp_mask_dt = fp_mask_dt[1:-1, 1:-1] + + # take the point in FN/FP region with the largest distance to its boundary + fn_mask_dt_flat = fn_mask_dt.reshape(-1) + fp_mask_dt_flat = fp_mask_dt.reshape(-1) + fn_argmax = np.argmax(fn_mask_dt_flat) + fp_argmax = np.argmax(fp_mask_dt_flat) + is_positive = fn_mask_dt_flat[fn_argmax] > fp_mask_dt_flat[fp_argmax] + pt_idx = fn_argmax if is_positive else fp_argmax + points[b, 0, 0] = pt_idx % W_im # x + points[b, 0, 1] = pt_idx // W_im # y + labels[b, 0] = int(is_positive) + + points = points.to(device) + labels = labels.to(device) + return points, labels + + +def get_next_point(gt_masks, pred_masks, method): + if method == "uniform": + return sample_random_points_from_errors(gt_masks, pred_masks) + elif method == "center": + return sample_one_point_from_error_center(gt_masks, pred_masks) + else: + raise ValueError(f"unknown sampling method {method}") diff --git a/freevideo_engine/sam2/sam2_image_predictor.py b/freevideo_engine/sam2/sam2_image_predictor.py new file mode 100644 index 00000000..74809ec5 --- /dev/null +++ b/freevideo_engine/sam2/sam2_image_predictor.py @@ -0,0 +1,449 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. + +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. + +import logging + +from typing import List, Optional, Tuple, Union + +import numpy as np +import torch +from PIL.Image import Image + +from freevideo_engine.sam2.modeling.sam2_base import SAM2Base + +from freevideo_engine.sam2.utils.transforms import SAM2Transforms + + +class SAM2ImagePredictor: + def __init__( + self, + sam_model: SAM2Base, + mask_threshold=0.0, + max_hole_area=0.0, + max_sprinkle_area=0.0, + **kwargs, + ) -> None: + """ + Uses SAM-2 to calculate the image embedding for an image, and then + allow repeated, efficient mask prediction given prompts. + + Arguments: + sam_model (Sam-2): The model to use for mask prediction. + mask_threshold (float): The threshold to use when converting mask logits + to binary masks. Masks are thresholded at 0 by default. + max_hole_area (int): If max_hole_area > 0, we fill small holes in up to + the maximum area of max_hole_area in low_res_masks. + max_sprinkle_area (int): If max_sprinkle_area > 0, we remove small sprinkles up to + the maximum area of max_sprinkle_area in low_res_masks. + """ + super().__init__() + self.model = sam_model + self._transforms = SAM2Transforms( + resolution=self.model.image_size, + mask_threshold=mask_threshold, + max_hole_area=max_hole_area, + max_sprinkle_area=max_sprinkle_area, + ) + + # Predictor state + self._is_image_set = False + self._features = None + self._orig_hw = None + # Whether the predictor is set for single image or a batch of images + self._is_batch = False + + # Predictor config + self.mask_threshold = mask_threshold + + # Spatial dim for backbone feature maps + self._bb_feat_sizes = [ + (256, 256), + (128, 128), + (64, 64), + ] + + @torch.no_grad() + def set_image( + self, + image: Union[np.ndarray, Image], + ) -> None: + """ + Calculates the image embeddings for the provided image, allowing + masks to be predicted with the 'predict' method. + + Arguments: + image (np.ndarray or PIL Image): The input image to embed in RGB format. The image should be in HWC format if np.ndarray, or WHC format if PIL Image + with pixel values in [0, 255]. + image_format (str): The color format of the image, in ['RGB', 'BGR']. + """ + self.reset_predictor() + # Transform the image to the form expected by the model + if isinstance(image, np.ndarray): + logging.info("For numpy array image, we assume (HxWxC) format") + self._orig_hw = [image.shape[:2]] + elif isinstance(image, Image): + w, h = image.size + self._orig_hw = [(h, w)] + else: + raise NotImplementedError("Image format not supported") + + input_image = self._transforms(image) + input_image = input_image[None, ...].to(self.device) + + assert ( + len(input_image.shape) == 4 and input_image.shape[1] == 3 + ), f"input_image must be of size 1x3xHxW, got {input_image.shape}" + logging.info("Computing image embeddings for the provided image...") + backbone_out = self.model.forward_image(input_image) + _, vision_feats, _, _ = self.model._prepare_backbone_features(backbone_out) + # Add no_mem_embed, which is added to the lowest rest feat. map during training on videos + if self.model.directly_add_no_mem_embed: + vision_feats[-1] = vision_feats[-1] + self.model.no_mem_embed + + feats = [ + feat.permute(1, 2, 0).view(1, -1, *feat_size) + for feat, feat_size in zip(vision_feats[::-1], self._bb_feat_sizes[::-1]) + ][::-1] + self._features = {"image_embed": feats[-1], "high_res_feats": feats[:-1]} + self._is_image_set = True + logging.info("Image embeddings computed.") + + @torch.no_grad() + def set_image_batch( + self, + image_list: List[Union[np.ndarray]], + ) -> None: + """ + Calculates the image embeddings for the provided image batch, allowing + masks to be predicted with the 'predict_batch' method. + + Arguments: + image_list (List[np.ndarray]): The input images to embed in RGB format. The image should be in HWC format if np.ndarray + with pixel values in [0, 255]. + """ + self.reset_predictor() + assert isinstance(image_list, list) + self._orig_hw = [] + for image in image_list: + assert isinstance( + image, np.ndarray + ), "Images are expected to be an np.ndarray in RGB format, and of shape HWC" + self._orig_hw.append(image.shape[:2]) + # Transform the image to the form expected by the model + img_batch = self._transforms.forward_batch(image_list) + img_batch = img_batch.to(self.device) + batch_size = img_batch.shape[0] + assert ( + len(img_batch.shape) == 4 and img_batch.shape[1] == 3 + ), f"img_batch must be of size Bx3xHxW, got {img_batch.shape}" + logging.info("Computing image embeddings for the provided images...") + backbone_out = self.model.forward_image(img_batch) + _, vision_feats, _, _ = self.model._prepare_backbone_features(backbone_out) + # Add no_mem_embed, which is added to the lowest rest feat. map during training on videos + if self.model.directly_add_no_mem_embed: + vision_feats[-1] = vision_feats[-1] + self.model.no_mem_embed + + feats = [ + feat.permute(1, 2, 0).view(batch_size, -1, *feat_size) + for feat, feat_size in zip(vision_feats[::-1], self._bb_feat_sizes[::-1]) + ][::-1] + self._features = {"image_embed": feats[-1], "high_res_feats": feats[:-1]} + self._is_image_set = True + self._is_batch = True + logging.info("Image embeddings computed.") + + def predict_batch( + self, + point_coords_batch: List[np.ndarray] = None, + point_labels_batch: List[np.ndarray] = None, + box_batch: List[np.ndarray] = None, + mask_input_batch: List[np.ndarray] = None, + multimask_output: bool = True, + return_logits: bool = False, + normalize_coords=True, + ) -> Tuple[List[np.ndarray], List[np.ndarray], List[np.ndarray]]: + """This function is very similar to predict(...), however it is used for batched mode, when the model is expected to generate predictions on multiple images. + It returns a tuple of lists of masks, ious, and low_res_masks_logits. + """ + assert self._is_batch, "This function should only be used when in batched mode" + if not self._is_image_set: + raise RuntimeError( + "An image must be set with .set_image_batch(...) before mask prediction." + ) + num_images = len(self._features["image_embed"]) + all_masks = [] + all_ious = [] + all_low_res_masks = [] + for img_idx in range(num_images): + # Transform input prompts + point_coords = ( + point_coords_batch[img_idx] if point_coords_batch is not None else None + ) + point_labels = ( + point_labels_batch[img_idx] if point_labels_batch is not None else None + ) + box = box_batch[img_idx] if box_batch is not None else None + mask_input = ( + mask_input_batch[img_idx] if mask_input_batch is not None else None + ) + mask_input, unnorm_coords, labels, unnorm_box = self._prep_prompts( + point_coords, + point_labels, + box, + mask_input, + normalize_coords, + img_idx=img_idx, + ) + masks, iou_predictions, low_res_masks = self._predict( + unnorm_coords, + labels, + unnorm_box, + mask_input, + multimask_output, + return_logits=return_logits, + img_idx=img_idx, + ) + masks_np = masks.squeeze(0).float().detach().cpu().numpy() + iou_predictions_np = ( + iou_predictions.squeeze(0).float().detach().cpu().numpy() + ) + low_res_masks_np = low_res_masks.squeeze(0).float().detach().cpu().numpy() + all_masks.append(masks_np) + all_ious.append(iou_predictions_np) + all_low_res_masks.append(low_res_masks_np) + + return all_masks, all_ious, all_low_res_masks + + def predict( + self, + point_coords: Optional[np.ndarray] = None, + point_labels: Optional[np.ndarray] = None, + box: Optional[np.ndarray] = None, + mask_input: Optional[np.ndarray] = None, + multimask_output: bool = True, + return_logits: bool = False, + normalize_coords=True, + ) -> Tuple[np.ndarray, np.ndarray, np.ndarray]: + """ + Predict masks for the given input prompts, using the currently set image. + + Arguments: + point_coords (np.ndarray or None): A Nx2 array of point prompts to the + model. Each point is in (X,Y) in pixels. + point_labels (np.ndarray or None): A length N array of labels for the + point prompts. 1 indicates a foreground point and 0 indicates a + background point. + box (np.ndarray or None): A length 4 array given a box prompt to the + model, in XYXY format. + mask_input (np.ndarray): A low resolution mask input to the model, typically + coming from a previous prediction iteration. Has form 1xHxW, where + for SAM, H=W=256. + multimask_output (bool): If true, the model will return three masks. + For ambiguous input prompts (such as a single click), this will often + produce better masks than a single prediction. If only a single + mask is needed, the model's predicted quality score can be used + to select the best mask. For non-ambiguous prompts, such as multiple + input prompts, multimask_output=False can give better results. + return_logits (bool): If true, returns un-thresholded masks logits + instead of a binary mask. + normalize_coords (bool): If true, the point coordinates will be normalized to the range [0,1] and point_coords is expected to be wrt. image dimensions. + + Returns: + (np.ndarray): The output masks in CxHxW format, where C is the + number of masks, and (H, W) is the original image size. + (np.ndarray): An array of length C containing the model's + predictions for the quality of each mask. + (np.ndarray): An array of shape CxHxW, where C is the number + of masks and H=W=256. These low resolution logits can be passed to + a subsequent iteration as mask input. + """ + if not self._is_image_set: + raise RuntimeError( + "An image must be set with .set_image(...) before mask prediction." + ) + + # Transform input prompts + + mask_input, unnorm_coords, labels, unnorm_box = self._prep_prompts( + point_coords, point_labels, box, mask_input, normalize_coords + ) + + masks, iou_predictions, low_res_masks = self._predict( + unnorm_coords, + labels, + unnorm_box, + mask_input, + multimask_output, + return_logits=return_logits, + ) + + masks_np = masks.squeeze(0).float().detach().cpu().numpy() + iou_predictions_np = iou_predictions.squeeze(0).float().detach().cpu().numpy() + low_res_masks_np = low_res_masks.squeeze(0).float().detach().cpu().numpy() + return masks_np, iou_predictions_np, low_res_masks_np + + def _prep_prompts( + self, point_coords, point_labels, box, mask_logits, normalize_coords, img_idx=-1 + ): + + unnorm_coords, labels, unnorm_box, mask_input = None, None, None, None + if point_coords is not None: + assert ( + point_labels is not None + ), "point_labels must be supplied if point_coords is supplied." + point_coords = torch.as_tensor( + point_coords, dtype=torch.float, device=self.device + ) + unnorm_coords = self._transforms.transform_coords( + point_coords, normalize=normalize_coords, orig_hw=self._orig_hw[img_idx] + ) + labels = torch.as_tensor(point_labels, dtype=torch.int, device=self.device) + if len(unnorm_coords.shape) == 2: + unnorm_coords, labels = unnorm_coords[None, ...], labels[None, ...] + if box is not None: + box = torch.as_tensor(box, dtype=torch.float, device=self.device) + unnorm_box = self._transforms.transform_boxes( + box, normalize=normalize_coords, orig_hw=self._orig_hw[img_idx] + ) # Bx2x2 + if mask_logits is not None: + mask_input = torch.as_tensor( + mask_logits, dtype=torch.float, device=self.device + ) + if len(mask_input.shape) == 3: + mask_input = mask_input[None, :, :, :] + return mask_input, unnorm_coords, labels, unnorm_box + + @torch.no_grad() + def _predict( + self, + point_coords: Optional[torch.Tensor], + point_labels: Optional[torch.Tensor], + boxes: Optional[torch.Tensor] = None, + mask_input: Optional[torch.Tensor] = None, + multimask_output: bool = True, + return_logits: bool = False, + img_idx: int = -1, + ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """ + Predict masks for the given input prompts, using the currently set image. + Input prompts are batched torch tensors and are expected to already be + transformed to the input frame using SAM2Transforms. + + Arguments: + point_coords (torch.Tensor or None): A BxNx2 array of point prompts to the + model. Each point is in (X,Y) in pixels. + point_labels (torch.Tensor or None): A BxN array of labels for the + point prompts. 1 indicates a foreground point and 0 indicates a + background point. + boxes (np.ndarray or None): A Bx4 array given a box prompt to the + model, in XYXY format. + mask_input (np.ndarray): A low resolution mask input to the model, typically + coming from a previous prediction iteration. Has form Bx1xHxW, where + for SAM, H=W=256. Masks returned by a previous iteration of the + predict method do not need further transformation. + multimask_output (bool): If true, the model will return three masks. + For ambiguous input prompts (such as a single click), this will often + produce better masks than a single prediction. If only a single + mask is needed, the model's predicted quality score can be used + to select the best mask. For non-ambiguous prompts, such as multiple + input prompts, multimask_output=False can give better results. + return_logits (bool): If true, returns un-thresholded masks logits + instead of a binary mask. + + Returns: + (torch.Tensor): The output masks in BxCxHxW format, where C is the + number of masks, and (H, W) is the original image size. + (torch.Tensor): An array of shape BxC containing the model's + predictions for the quality of each mask. + (torch.Tensor): An array of shape BxCxHxW, where C is the number + of masks and H=W=256. These low res logits can be passed to + a subsequent iteration as mask input. + """ + if not self._is_image_set: + raise RuntimeError( + "An image must be set with .set_image(...) before mask prediction." + ) + + if point_coords is not None: + concat_points = (point_coords, point_labels) + else: + concat_points = None + + # Embed prompts + if boxes is not None: + box_coords = boxes.reshape(-1, 2, 2) + box_labels = torch.tensor([[2, 3]], dtype=torch.int, device=boxes.device) + box_labels = box_labels.repeat(boxes.size(0), 1) + # we merge "boxes" and "points" into a single "concat_points" input (where + # boxes are added at the beginning) to sam_prompt_encoder + if concat_points is not None: + concat_coords = torch.cat([box_coords, concat_points[0]], dim=1) + concat_labels = torch.cat([box_labels, concat_points[1]], dim=1) + concat_points = (concat_coords, concat_labels) + else: + concat_points = (box_coords, box_labels) + + sparse_embeddings, dense_embeddings = self.model.sam_prompt_encoder( + points=concat_points, + boxes=None, + masks=mask_input, + ) + + # Predict masks + batched_mode = ( + concat_points is not None and concat_points[0].shape[0] > 1 + ) # multi object prediction + high_res_features = [ + feat_level[img_idx].unsqueeze(0) + for feat_level in self._features["high_res_feats"] + ] + low_res_masks, iou_predictions, _, _ = self.model.sam_mask_decoder( + image_embeddings=self._features["image_embed"][img_idx].unsqueeze(0), + image_pe=self.model.sam_prompt_encoder.get_dense_pe(), + sparse_prompt_embeddings=sparse_embeddings, + dense_prompt_embeddings=dense_embeddings, + multimask_output=multimask_output, + repeat_image=batched_mode, + high_res_features=high_res_features, + ) + + # Upscale the masks to the original image resolution + masks = self._transforms.postprocess_masks( + low_res_masks, self._orig_hw[img_idx] + ) + low_res_masks = torch.clamp(low_res_masks, -32.0, 32.0) + if not return_logits: + masks = masks > self.mask_threshold + + return masks, iou_predictions, low_res_masks + + def get_image_embedding(self) -> torch.Tensor: + """ + Returns the image embeddings for the currently set image, with + shape 1xCxHxW, where C is the embedding dimension and (H,W) are + the embedding spatial dimension of SAM (typically C=256, H=W=64). + """ + if not self._is_image_set: + raise RuntimeError( + "An image must be set with .set_image(...) to generate an embedding." + ) + assert ( + self._features is not None + ), "Features must exist if an image has been set." + return self._features["image_embed"] + + @property + def device(self) -> torch.device: + return self.model.device + + def reset_predictor(self) -> None: + """ + Resets the image embeddings and other state variables. + """ + self._is_image_set = False + self._features = None + self._orig_hw = None + self._is_batch = False diff --git a/freevideo_engine/sam2/utils/__init__.py b/freevideo_engine/sam2/utils/__init__.py new file mode 100644 index 00000000..5277f461 --- /dev/null +++ b/freevideo_engine/sam2/utils/__init__.py @@ -0,0 +1,5 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. + +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. diff --git a/freevideo_engine/sam2/utils/amg.py b/freevideo_engine/sam2/utils/amg.py new file mode 100644 index 00000000..98684296 --- /dev/null +++ b/freevideo_engine/sam2/utils/amg.py @@ -0,0 +1,348 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. + +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. + +import math +from copy import deepcopy +from itertools import product +from typing import Any, Dict, Generator, ItemsView, List, Tuple + +import numpy as np +import torch + +# Very lightly adapted from https://github.com/facebookresearch/segment-anything/blob/main/segment_anything/utils/amg.py + + +class MaskData: + """ + A structure for storing masks and their related data in batched format. + Implements basic filtering and concatenation. + """ + + def __init__(self, **kwargs) -> None: + for v in kwargs.values(): + assert isinstance( + v, (list, np.ndarray, torch.Tensor) + ), "MaskData only supports list, numpy arrays, and torch tensors." + self._stats = dict(**kwargs) + + def __setitem__(self, key: str, item: Any) -> None: + assert isinstance( + item, (list, np.ndarray, torch.Tensor) + ), "MaskData only supports list, numpy arrays, and torch tensors." + self._stats[key] = item + + def __delitem__(self, key: str) -> None: + del self._stats[key] + + def __getitem__(self, key: str) -> Any: + return self._stats[key] + + def items(self) -> ItemsView[str, Any]: + return self._stats.items() + + def filter(self, keep: torch.Tensor) -> None: + for k, v in self._stats.items(): + if v is None: + self._stats[k] = None + elif isinstance(v, torch.Tensor): + self._stats[k] = v[torch.as_tensor(keep, device=v.device)] + elif isinstance(v, np.ndarray): + self._stats[k] = v[keep.detach().cpu().numpy()] + elif isinstance(v, list) and keep.dtype == torch.bool: + self._stats[k] = [a for i, a in enumerate(v) if keep[i]] + elif isinstance(v, list): + self._stats[k] = [v[i] for i in keep] + else: + raise TypeError(f"MaskData key {k} has an unsupported type {type(v)}.") + + def cat(self, new_stats: "MaskData") -> None: + for k, v in new_stats.items(): + if k not in self._stats or self._stats[k] is None: + self._stats[k] = deepcopy(v) + elif isinstance(v, torch.Tensor): + self._stats[k] = torch.cat([self._stats[k], v], dim=0) + elif isinstance(v, np.ndarray): + self._stats[k] = np.concatenate([self._stats[k], v], axis=0) + elif isinstance(v, list): + self._stats[k] = self._stats[k] + deepcopy(v) + else: + raise TypeError(f"MaskData key {k} has an unsupported type {type(v)}.") + + def to_numpy(self) -> None: + for k, v in self._stats.items(): + if isinstance(v, torch.Tensor): + self._stats[k] = v.float().detach().cpu().numpy() + + +def is_box_near_crop_edge( + boxes: torch.Tensor, crop_box: List[int], orig_box: List[int], atol: float = 20.0 +) -> torch.Tensor: + """Filter masks at the edge of a crop, but not at the edge of the original image.""" + crop_box_torch = torch.as_tensor(crop_box, dtype=torch.float, device=boxes.device) + orig_box_torch = torch.as_tensor(orig_box, dtype=torch.float, device=boxes.device) + boxes = uncrop_boxes_xyxy(boxes, crop_box).float() + near_crop_edge = torch.isclose(boxes, crop_box_torch[None, :], atol=atol, rtol=0) + near_image_edge = torch.isclose(boxes, orig_box_torch[None, :], atol=atol, rtol=0) + near_crop_edge = torch.logical_and(near_crop_edge, ~near_image_edge) + return torch.any(near_crop_edge, dim=1) + + +def box_xyxy_to_xywh(box_xyxy: torch.Tensor) -> torch.Tensor: + box_xywh = deepcopy(box_xyxy) + box_xywh[2] = box_xywh[2] - box_xywh[0] + box_xywh[3] = box_xywh[3] - box_xywh[1] + return box_xywh + + +def batch_iterator(batch_size: int, *args) -> Generator[List[Any], None, None]: + assert len(args) > 0 and all( + len(a) == len(args[0]) for a in args + ), "Batched iteration must have inputs of all the same size." + n_batches = len(args[0]) // batch_size + int(len(args[0]) % batch_size != 0) + for b in range(n_batches): + yield [arg[b * batch_size : (b + 1) * batch_size] for arg in args] + + +def mask_to_rle_pytorch(tensor: torch.Tensor) -> List[Dict[str, Any]]: + """ + Encodes masks to an uncompressed RLE, in the format expected by + pycoco tools. + """ + # Put in fortran order and flatten h,w + b, h, w = tensor.shape + tensor = tensor.permute(0, 2, 1).flatten(1) + + # Compute change indices + diff = tensor[:, 1:] ^ tensor[:, :-1] + change_indices = diff.nonzero() + + # Encode run length + out = [] + for i in range(b): + cur_idxs = change_indices[change_indices[:, 0] == i, 1] + cur_idxs = torch.cat( + [ + torch.tensor([0], dtype=cur_idxs.dtype, device=cur_idxs.device), + cur_idxs + 1, + torch.tensor([h * w], dtype=cur_idxs.dtype, device=cur_idxs.device), + ] + ) + btw_idxs = cur_idxs[1:] - cur_idxs[:-1] + counts = [] if tensor[i, 0] == 0 else [0] + counts.extend(btw_idxs.detach().cpu().tolist()) + out.append({"size": [h, w], "counts": counts}) + return out + + +def rle_to_mask(rle: Dict[str, Any]) -> np.ndarray: + """Compute a binary mask from an uncompressed RLE.""" + h, w = rle["size"] + mask = np.empty(h * w, dtype=bool) + idx = 0 + parity = False + for count in rle["counts"]: + mask[idx : idx + count] = parity + idx += count + parity ^= True + mask = mask.reshape(w, h) + return mask.transpose() # Put in C order + + +def area_from_rle(rle: Dict[str, Any]) -> int: + return sum(rle["counts"][1::2]) + + +def calculate_stability_score( + masks: torch.Tensor, mask_threshold: float, threshold_offset: float +) -> torch.Tensor: + """ + Computes the stability score for a batch of masks. The stability + score is the IoU between the binary masks obtained by thresholding + the predicted mask logits at high and low values. + """ + # One mask is always contained inside the other. + # Save memory by preventing unnecessary cast to torch.int64 + intersections = ( + (masks > (mask_threshold + threshold_offset)) + .sum(-1, dtype=torch.int16) + .sum(-1, dtype=torch.int32) + ) + unions = ( + (masks > (mask_threshold - threshold_offset)) + .sum(-1, dtype=torch.int16) + .sum(-1, dtype=torch.int32) + ) + return intersections / unions + + +def build_point_grid(n_per_side: int) -> np.ndarray: + """Generates a 2D grid of points evenly spaced in [0,1]x[0,1].""" + offset = 1 / (2 * n_per_side) + points_one_side = np.linspace(offset, 1 - offset, n_per_side) + points_x = np.tile(points_one_side[None, :], (n_per_side, 1)) + points_y = np.tile(points_one_side[:, None], (1, n_per_side)) + points = np.stack([points_x, points_y], axis=-1).reshape(-1, 2) + return points + + +def build_all_layer_point_grids( + n_per_side: int, n_layers: int, scale_per_layer: int +) -> List[np.ndarray]: + """Generates point grids for all crop layers.""" + points_by_layer = [] + for i in range(n_layers + 1): + n_points = int(n_per_side / (scale_per_layer**i)) + points_by_layer.append(build_point_grid(n_points)) + return points_by_layer + + +def generate_crop_boxes( + im_size: Tuple[int, ...], n_layers: int, overlap_ratio: float +) -> Tuple[List[List[int]], List[int]]: + """ + Generates a list of crop boxes of different sizes. Each layer + has (2**i)**2 boxes for the ith layer. + """ + crop_boxes, layer_idxs = [], [] + im_h, im_w = im_size + short_side = min(im_h, im_w) + + # Original image + crop_boxes.append([0, 0, im_w, im_h]) + layer_idxs.append(0) + + def crop_len(orig_len, n_crops, overlap): + return int(math.ceil((overlap * (n_crops - 1) + orig_len) / n_crops)) + + for i_layer in range(n_layers): + n_crops_per_side = 2 ** (i_layer + 1) + overlap = int(overlap_ratio * short_side * (2 / n_crops_per_side)) + + crop_w = crop_len(im_w, n_crops_per_side, overlap) + crop_h = crop_len(im_h, n_crops_per_side, overlap) + + crop_box_x0 = [int((crop_w - overlap) * i) for i in range(n_crops_per_side)] + crop_box_y0 = [int((crop_h - overlap) * i) for i in range(n_crops_per_side)] + + # Crops in XYWH format + for x0, y0 in product(crop_box_x0, crop_box_y0): + box = [x0, y0, min(x0 + crop_w, im_w), min(y0 + crop_h, im_h)] + crop_boxes.append(box) + layer_idxs.append(i_layer + 1) + + return crop_boxes, layer_idxs + + +def uncrop_boxes_xyxy(boxes: torch.Tensor, crop_box: List[int]) -> torch.Tensor: + x0, y0, _, _ = crop_box + offset = torch.tensor([[x0, y0, x0, y0]], device=boxes.device) + # Check if boxes has a channel dimension + if len(boxes.shape) == 3: + offset = offset.unsqueeze(1) + return boxes + offset + + +def uncrop_points(points: torch.Tensor, crop_box: List[int]) -> torch.Tensor: + x0, y0, _, _ = crop_box + offset = torch.tensor([[x0, y0]], device=points.device) + # Check if points has a channel dimension + if len(points.shape) == 3: + offset = offset.unsqueeze(1) + return points + offset + + +def uncrop_masks( + masks: torch.Tensor, crop_box: List[int], orig_h: int, orig_w: int +) -> torch.Tensor: + x0, y0, x1, y1 = crop_box + if x0 == 0 and y0 == 0 and x1 == orig_w and y1 == orig_h: + return masks + # Coordinate transform masks + pad_x, pad_y = orig_w - (x1 - x0), orig_h - (y1 - y0) + pad = (x0, pad_x - x0, y0, pad_y - y0) + return torch.nn.functional.pad(masks, pad, value=0) + + +def remove_small_regions( + mask: np.ndarray, area_thresh: float, mode: str +) -> Tuple[np.ndarray, bool]: + """ + Removes small disconnected regions and holes in a mask. Returns the + mask and an indicator of if the mask has been modified. + """ + import cv2 # type: ignore + + assert mode in ["holes", "islands"] + correct_holes = mode == "holes" + working_mask = (correct_holes ^ mask).astype(np.uint8) + n_labels, regions, stats, _ = cv2.connectedComponentsWithStats(working_mask, 8) + sizes = stats[:, -1][1:] # Row 0 is background label + small_regions = [i + 1 for i, s in enumerate(sizes) if s < area_thresh] + if len(small_regions) == 0: + return mask, False + fill_labels = [0] + small_regions + if not correct_holes: + fill_labels = [i for i in range(n_labels) if i not in fill_labels] + # If every region is below threshold, keep largest + if len(fill_labels) == 0: + fill_labels = [int(np.argmax(sizes)) + 1] + mask = np.isin(regions, fill_labels) + return mask, True + + +def coco_encode_rle(uncompressed_rle: Dict[str, Any]) -> Dict[str, Any]: + from pycocotools import mask as mask_utils # type: ignore + + h, w = uncompressed_rle["size"] + rle = mask_utils.frPyObjects(uncompressed_rle, h, w) + rle["counts"] = rle["counts"].decode("utf-8") # Necessary to serialize with json + return rle + + +def batched_mask_to_box(masks: torch.Tensor) -> torch.Tensor: + """ + Calculates boxes in XYXY format around masks. Return [0,0,0,0] for + an empty mask. For input shape C1xC2x...xHxW, the output shape is C1xC2x...x4. + """ + # torch.max below raises an error on empty inputs, just skip in this case + if torch.numel(masks) == 0: + return torch.zeros(*masks.shape[:-2], 4, device=masks.device) + + # Normalize shape to CxHxW + shape = masks.shape + h, w = shape[-2:] + if len(shape) > 2: + masks = masks.flatten(0, -3) + else: + masks = masks.unsqueeze(0) + + # Get top and bottom edges + in_height, _ = torch.max(masks, dim=-1) + in_height_coords = in_height * torch.arange(h, device=in_height.device)[None, :] + bottom_edges, _ = torch.max(in_height_coords, dim=-1) + in_height_coords = in_height_coords + h * (~in_height) + top_edges, _ = torch.min(in_height_coords, dim=-1) + + # Get left and right edges + in_width, _ = torch.max(masks, dim=-2) + in_width_coords = in_width * torch.arange(w, device=in_width.device)[None, :] + right_edges, _ = torch.max(in_width_coords, dim=-1) + in_width_coords = in_width_coords + w * (~in_width) + left_edges, _ = torch.min(in_width_coords, dim=-1) + + # If the mask is empty the right edge will be to the left of the left edge. + # Replace these boxes with [0, 0, 0, 0] + empty_filter = (right_edges < left_edges) | (bottom_edges < top_edges) + out = torch.stack([left_edges, top_edges, right_edges, bottom_edges], dim=-1) + out = out * (~empty_filter).unsqueeze(-1) + + # Return to original shape + if len(shape) > 2: + out = out.reshape(*shape[:-2], 4) + else: + out = out[0] + + return out diff --git a/freevideo_engine/sam2/utils/misc.py b/freevideo_engine/sam2/utils/misc.py new file mode 100644 index 00000000..08e69546 --- /dev/null +++ b/freevideo_engine/sam2/utils/misc.py @@ -0,0 +1,349 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. + +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. + +import os +import warnings +from threading import Thread + +import numpy as np +import torch +from PIL import Image +from tqdm import tqdm + + +def get_sdpa_settings(): + if torch.cuda.is_available(): + old_gpu = torch.cuda.get_device_properties(0).major < 7 + # only use Flash Attention on Ampere (8.0) or newer GPUs + use_flash_attn = torch.cuda.get_device_properties(0).major >= 8 + if not use_flash_attn: + warnings.warn( + "Flash Attention is disabled as it requires a GPU with Ampere (8.0) CUDA capability.", + category=UserWarning, + stacklevel=2, + ) + # keep math kernel for PyTorch versions before 2.2 (Flash Attention v2 is only + # available on PyTorch 2.2+, while Flash Attention v1 cannot handle all cases) + pytorch_version = tuple(int(v) for v in torch.__version__.split(".")[:2]) + if pytorch_version < (2, 2): + warnings.warn( + f"You are using PyTorch {torch.__version__} without Flash Attention v2 support. " + "Consider upgrading to PyTorch 2.2+ for Flash Attention v2 (which could be faster).", + category=UserWarning, + stacklevel=2, + ) + math_kernel_on = pytorch_version < (2, 2) or not use_flash_attn + else: + old_gpu = True + use_flash_attn = False + math_kernel_on = True + + return old_gpu, use_flash_attn, math_kernel_on + + +def get_connected_components(mask): + """ + Get the connected components (8-connectivity) of binary masks of shape (N, 1, H, W). + + Inputs: + - mask: A binary mask tensor of shape (N, 1, H, W), where 1 is foreground and 0 is + background. + + Outputs: + - labels: A tensor of shape (N, 1, H, W) containing the connected component labels + for foreground pixels and 0 for background pixels. + - counts: A tensor of shape (N, 1, H, W) containing the area of the connected + components for foreground pixels and 0 for background pixels. + """ + from freevideo_engine.sam2 import _C + + return _C.get_connected_componnets(mask.to(torch.uint8).contiguous()) + + +def mask_to_box(masks: torch.Tensor): + """ + compute bounding box given an input mask + + Inputs: + - masks: [B, 1, H, W] masks, dtype=torch.Tensor + + Returns: + - box_coords: [B, 1, 4], contains (x, y) coordinates of top left and bottom right box corners, dtype=torch.Tensor + """ + B, _, h, w = masks.shape + device = masks.device + xs = torch.arange(w, device=device, dtype=torch.int32) + ys = torch.arange(h, device=device, dtype=torch.int32) + grid_xs, grid_ys = torch.meshgrid(xs, ys, indexing="xy") + grid_xs = grid_xs[None, None, ...].expand(B, 1, h, w) + grid_ys = grid_ys[None, None, ...].expand(B, 1, h, w) + min_xs, _ = torch.min(torch.where(masks, grid_xs, w).flatten(-2), dim=-1) + max_xs, _ = torch.max(torch.where(masks, grid_xs, -1).flatten(-2), dim=-1) + min_ys, _ = torch.min(torch.where(masks, grid_ys, h).flatten(-2), dim=-1) + max_ys, _ = torch.max(torch.where(masks, grid_ys, -1).flatten(-2), dim=-1) + bbox_coords = torch.stack((min_xs, min_ys, max_xs, max_ys), dim=-1) + + return bbox_coords + + +def _load_img_as_tensor(img_path, image_size): + img_pil = Image.open(img_path) + img_np = np.array(img_pil.convert("RGB").resize((image_size, image_size))) + if img_np.dtype == np.uint8: # np.uint8 is expected for JPEG images + img_np = img_np / 255.0 + else: + raise RuntimeError(f"Unknown image dtype: {img_np.dtype} on {img_path}") + img = torch.from_numpy(img_np).permute(2, 0, 1) + video_width, video_height = img_pil.size # the original video size + return img, video_height, video_width + + +class AsyncVideoFrameLoader: + """ + A list of video frames to be load asynchronously without blocking session start. + """ + + def __init__( + self, + img_paths, + image_size, + offload_video_to_cpu, + img_mean, + img_std, + compute_device, + ): + self.img_paths = img_paths + self.image_size = image_size + self.offload_video_to_cpu = offload_video_to_cpu + self.img_mean = img_mean + self.img_std = img_std + # items in `self.images` will be loaded asynchronously + self.images = [None] * len(img_paths) + # catch and raise any exceptions in the async loading thread + self.exception = None + # video_height and video_width be filled when loading the first image + self.video_height = None + self.video_width = None + self.compute_device = compute_device + + # load the first frame to fill video_height and video_width and also + # to cache it (since it's most likely where the user will click) + self.__getitem__(0) + + # load the rest of frames asynchronously without blocking the session start + def _load_frames(): + try: + for n in tqdm(range(len(self.images)), desc="frame loading (JPEG)"): + self.__getitem__(n) + except Exception as e: + self.exception = e + + self.thread = Thread(target=_load_frames, daemon=True) + self.thread.start() + + def __getitem__(self, index): + if self.exception is not None: + raise RuntimeError("Failure in frame loading thread") from self.exception + + img = self.images[index] + if img is not None: + return img + + img, video_height, video_width = _load_img_as_tensor( + self.img_paths[index], self.image_size + ) + self.video_height = video_height + self.video_width = video_width + # normalize by mean and std + img -= self.img_mean + img /= self.img_std + if not self.offload_video_to_cpu: + img = img.to(self.compute_device, non_blocking=True) + self.images[index] = img + return img + + def __len__(self): + return len(self.images) + + +def load_video_frames( + video_path, + image_size, + offload_video_to_cpu, + img_mean=(0.485, 0.456, 0.406), + img_std=(0.229, 0.224, 0.225), + async_loading_frames=False, + compute_device=torch.device("cuda"), +): + """ + Load the video frames from video_path. The frames are resized to image_size as in + the model and are loaded to GPU if offload_video_to_cpu=False. This is used by the demo. + """ + is_bytes = isinstance(video_path, bytes) + is_str = isinstance(video_path, str) + is_mp4_path = is_str and os.path.splitext(video_path)[-1] in [".mp4", ".MP4"] + if is_bytes or is_mp4_path: + return load_video_frames_from_video_file( + video_path=video_path, + image_size=image_size, + offload_video_to_cpu=offload_video_to_cpu, + img_mean=img_mean, + img_std=img_std, + compute_device=compute_device, + ) + elif is_str and os.path.isdir(video_path): + return load_video_frames_from_jpg_images( + video_path=video_path, + image_size=image_size, + offload_video_to_cpu=offload_video_to_cpu, + img_mean=img_mean, + img_std=img_std, + async_loading_frames=async_loading_frames, + compute_device=compute_device, + ) + else: + raise NotImplementedError( + "Only MP4 video and JPEG folder are supported at this moment" + ) + + +def load_video_frames_from_jpg_images( + video_path, + image_size, + offload_video_to_cpu, + img_mean=(0.485, 0.456, 0.406), + img_std=(0.229, 0.224, 0.225), + async_loading_frames=False, + compute_device=torch.device("cuda"), +): + """ + Load the video frames from a directory of JPEG files (".jpg" format). + + The frames are resized to image_size x image_size and are loaded to GPU if + `offload_video_to_cpu` is `False` and to CPU if `offload_video_to_cpu` is `True`. + + You can load a frame asynchronously by setting `async_loading_frames` to `True`. + """ + if isinstance(video_path, str) and os.path.isdir(video_path): + jpg_folder = video_path + else: + raise NotImplementedError( + "Only JPEG frames are supported at this moment. For video files, you may use " + "ffmpeg (https://ffmpeg.org/) to extract frames into a folder of JPEG files, such as \n" + "```\n" + "ffmpeg -i .mp4 -q:v 2 -start_number 0 /'%05d.jpg'\n" + "```\n" + "where `-q:v` generates high-quality JPEG frames and `-start_number 0` asks " + "ffmpeg to start the JPEG file from 00000.jpg." + ) + + frame_names = [ + p + for p in os.listdir(jpg_folder) + if os.path.splitext(p)[-1] in [".jpg", ".jpeg", ".JPG", ".JPEG"] + ] + frame_names.sort(key=lambda p: int(os.path.splitext(p)[0])) + num_frames = len(frame_names) + if num_frames == 0: + raise RuntimeError(f"no images found in {jpg_folder}") + img_paths = [os.path.join(jpg_folder, frame_name) for frame_name in frame_names] + img_mean = torch.tensor(img_mean, dtype=torch.float32)[:, None, None] + img_std = torch.tensor(img_std, dtype=torch.float32)[:, None, None] + + if async_loading_frames: + lazy_images = AsyncVideoFrameLoader( + img_paths, + image_size, + offload_video_to_cpu, + img_mean, + img_std, + compute_device, + ) + return lazy_images, lazy_images.video_height, lazy_images.video_width + + images = torch.zeros(num_frames, 3, image_size, image_size, dtype=torch.float32) + for n, img_path in enumerate(tqdm(img_paths, desc="frame loading (JPEG)")): + images[n], video_height, video_width = _load_img_as_tensor(img_path, image_size) + if not offload_video_to_cpu: + images = images.to(compute_device) + img_mean = img_mean.to(compute_device) + img_std = img_std.to(compute_device) + # normalize by mean and std + images -= img_mean + images /= img_std + return images, video_height, video_width + + +def load_video_frames_from_video_file( + video_path, + image_size, + offload_video_to_cpu, + img_mean=(0.485, 0.456, 0.406), + img_std=(0.229, 0.224, 0.225), + compute_device=torch.device("cuda"), +): + """Load the video frames from a video file.""" + import decord + + img_mean = torch.tensor(img_mean, dtype=torch.float32)[:, None, None] + img_std = torch.tensor(img_std, dtype=torch.float32)[:, None, None] + # Get the original video height and width + decord.bridge.set_bridge("torch") + video_height, video_width, _ = decord.VideoReader(video_path).next().shape + # Iterate over all frames in the video + images = [] + for frame in decord.VideoReader(video_path, width=image_size, height=image_size): + images.append(frame.permute(2, 0, 1)) + + images = torch.stack(images, dim=0).float() / 255.0 + if not offload_video_to_cpu: + images = images.to(compute_device) + img_mean = img_mean.to(compute_device) + img_std = img_std.to(compute_device) + # normalize by mean and std + images -= img_mean + images /= img_std + return images, video_height, video_width + + +def fill_holes_in_mask_scores(mask, max_area): + """ + A post processor to fill small holes in mask scores with area under `max_area`. + """ + # Holes are those connected components in background with area <= self.max_area + # (background regions are those with mask scores <= 0) + assert max_area > 0, "max_area must be positive" + + input_mask = mask + try: + labels, areas = get_connected_components(mask <= 0) + is_hole = (labels > 0) & (areas <= max_area) + # We fill holes with a small positive mask score (0.1) to change them to foreground. + mask = torch.where(is_hole, 0.1, mask) + except Exception as e: + # Skip the post-processing step on removing small holes if the CUDA kernel fails + warnings.warn( + f"{e}\n\nSkipping the post-processing step due to the error above. You can " + "still use SAM 2 and it's OK to ignore the error above, although some post-processing " + "functionality may be limited (which doesn't affect the results in most cases; see " + "https://github.com/facebookresearch/sam2/blob/main/INSTALL.md).", + category=UserWarning, + stacklevel=2, + ) + mask = input_mask + + return mask + + +def concat_points(old_point_inputs, new_points, new_labels): + """Add new points and labels to previous point inputs (add at the end).""" + if old_point_inputs is None: + points, labels = new_points, new_labels + else: + points = torch.cat([old_point_inputs["point_coords"], new_points], dim=1) + labels = torch.cat([old_point_inputs["point_labels"], new_labels], dim=1) + + return {"point_coords": points, "point_labels": labels} diff --git a/freevideo_engine/sam2/utils/transforms.py b/freevideo_engine/sam2/utils/transforms.py new file mode 100644 index 00000000..04ccdf70 --- /dev/null +++ b/freevideo_engine/sam2/utils/transforms.py @@ -0,0 +1,118 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. + +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. + +import warnings + +import torch +import torch.nn as nn +import torch.nn.functional as F +from torchvision.transforms import Normalize, Resize, ToTensor + + +class SAM2Transforms(nn.Module): + def __init__( + self, resolution, mask_threshold, max_hole_area=0.0, max_sprinkle_area=0.0 + ): + """ + Transforms for SAM2. + """ + super().__init__() + self.resolution = resolution + self.mask_threshold = mask_threshold + self.max_hole_area = max_hole_area + self.max_sprinkle_area = max_sprinkle_area + self.mean = [0.485, 0.456, 0.406] + self.std = [0.229, 0.224, 0.225] + self.to_tensor = ToTensor() + self.transforms = torch.jit.script( + nn.Sequential( + Resize((self.resolution, self.resolution)), + Normalize(self.mean, self.std), + ) + ) + + def __call__(self, x): + x = self.to_tensor(x) + return self.transforms(x) + + def forward_batch(self, img_list): + img_batch = [self.transforms(self.to_tensor(img)) for img in img_list] + img_batch = torch.stack(img_batch, dim=0) + return img_batch + + def transform_coords( + self, coords: torch.Tensor, normalize=False, orig_hw=None + ) -> torch.Tensor: + """ + Expects a torch tensor with length 2 in the last dimension. The coordinates can be in absolute image or normalized coordinates, + If the coords are in absolute image coordinates, normalize should be set to True and original image size is required. + + Returns + Un-normalized coordinates in the range of [0, 1] which is expected by the SAM2 model. + """ + if normalize: + assert orig_hw is not None + h, w = orig_hw + coords = coords.clone() + coords[..., 0] = coords[..., 0] / w + coords[..., 1] = coords[..., 1] / h + + coords = coords * self.resolution # unnormalize coords + return coords + + def transform_boxes( + self, boxes: torch.Tensor, normalize=False, orig_hw=None + ) -> torch.Tensor: + """ + Expects a tensor of shape Bx4. The coordinates can be in absolute image or normalized coordinates, + if the coords are in absolute image coordinates, normalize should be set to True and original image size is required. + """ + boxes = self.transform_coords(boxes.reshape(-1, 2, 2), normalize, orig_hw) + return boxes + + def postprocess_masks(self, masks: torch.Tensor, orig_hw) -> torch.Tensor: + """ + Perform PostProcessing on output masks. + """ + from freevideo_engine.sam2.utils.misc import get_connected_components + + masks = masks.float() + input_masks = masks + mask_flat = masks.flatten(0, 1).unsqueeze(1) # flatten as 1-channel image + try: + if self.max_hole_area > 0: + # Holes are those connected components in background with area <= self.fill_hole_area + # (background regions are those with mask scores <= self.mask_threshold) + labels, areas = get_connected_components( + mask_flat <= self.mask_threshold + ) + is_hole = (labels > 0) & (areas <= self.max_hole_area) + is_hole = is_hole.reshape_as(masks) + # We fill holes with a small positive mask score (10.0) to change them to foreground. + masks = torch.where(is_hole, self.mask_threshold + 10.0, masks) + + if self.max_sprinkle_area > 0: + labels, areas = get_connected_components( + mask_flat > self.mask_threshold + ) + is_hole = (labels > 0) & (areas <= self.max_sprinkle_area) + is_hole = is_hole.reshape_as(masks) + # We fill holes with negative mask score (-10.0) to change them to background. + masks = torch.where(is_hole, self.mask_threshold - 10.0, masks) + except Exception as e: + # Skip the post-processing step if the CUDA kernel fails + warnings.warn( + f"{e}\n\nSkipping the post-processing step due to the error above. You can " + "still use SAM 2 and it's OK to ignore the error above, although some post-processing " + "functionality may be limited (which doesn't affect the results in most cases; see " + "https://github.com/facebookresearch/sam2/blob/main/INSTALL.md).", + category=UserWarning, + stacklevel=2, + ) + masks = input_masks + + masks = F.interpolate(masks, orig_hw, mode="bilinear", align_corners=False) + return masks diff --git a/freevideo_engine/step_cache.py b/freevideo_engine/step_cache.py index e081419d..c54f734c 100644 --- a/freevideo_engine/step_cache.py +++ b/freevideo_engine/step_cache.py @@ -111,6 +111,7 @@ def begin(self, video_indices, audio_indices): def stats(self): full = self.steps - len(self.hits) return dict(kind=self.kind, threshold=self.threshold, steps=self.steps, cached_steps=list(self.hits), + cached_share=round(len(self.hits) / self.steps, 3) if self.steps else 0., full_forwards=full, decisions=self.decisions, source=self.source, disabled_reason=self.disabled_reason, guard=GUARD) diff --git a/freevideo_engine/tiers.py b/freevideo_engine/tiers.py index 8e58e08b..a5c08fda 100644 --- a/freevideo_engine/tiers.py +++ b/freevideo_engine/tiers.py @@ -6,12 +6,11 @@ - light: two-pass, 8 + 3 steps, first pass at half size (2x learned upscale). - medium: two-pass, 8 + 3 steps, a larger first pass (`upscale`, below 2). -- high: one pass of 20 steps with the audio-video EasyCache. -- max: one pass of 20 steps, every step sampled. +- high: one pass of 16 steps with the guarded audio-video EasyCache. +- max: one pass of 20 steps with the same guarded cache. -A level in one pass always samples 20 steps; fewer steps come only from a -user's own setting. Only a two-pass level samples a second time. Single-pass -levels carry REFINE_STEPS only as the neutral value the `refine_steps` field +Each level uses the step count in the table below. Only a two-pass level +samples a second time. Single-pass levels carry REFINE_STEPS only as the neutral value the `refine_steps` field needs; it is never run for them, and their `upscale` is 2 (no factor). `step_cache` is the engine's cache key: None, 'easycache-av' (audio-video EasyCache) or 'fbcache' (first-block cache, no level uses it; a request may @@ -20,7 +19,8 @@ Mac (MPS) runs the same levels (TIERS_MAC). Results made before the audio-video EasyCache was ported there, one pass of 16 steps without a cache, -still count as high (LEGACY_MAC_HIGH). +still count as high (LEGACY_MAC_HIGH). Earlier one-pass 20-step results +without a cache still count as max on every platform (LEGACY_MAX). """ import sys @@ -30,10 +30,11 @@ TIERS = { 'light': dict(steps=8, two_pass=True, step_cache=None, upscale=2), 'medium': dict(steps=8, two_pass=True, step_cache=None, upscale=1.5), - 'high': dict(steps=20, two_pass=False, step_cache='easycache-av', upscale=2), - 'max': dict(steps=20, two_pass=False, step_cache=None, upscale=2), + 'high': dict(steps=16, two_pass=False, step_cache='easycache-av', upscale=2), + 'max': dict(steps=20, two_pass=False, step_cache='easycache-av', upscale=2), } TIERS_MAC = TIERS +LEGACY_MAX = dict(steps=20, two_pass=False, step_cache=None, upscale=2) LEGACY_MAC_HIGH = dict(steps=16, two_pass=False, step_cache=None, upscale=2) # Sampling time of a cached pass as a share of the same pass without a cache, @@ -50,6 +51,8 @@ # is plain 20 steps. # EasyCache-AV, guarded (at most 40% of steps reused, never twice in a row): measured on the RTX PRO 6000, # 864x1184, 5 s, 20 steps, 2 reused steps: 212 s against 238 s for the plain 20 steps (about .88 of the sampling). +# At 16 steps the guarded cache reused 1–4 steps (about 14% on average), +# so the same conservative .88 share also applies at 16 steps. SAMPLING_SHARE = {None: 1., 'easycache-av': .88, 'fbcache': .9} # Sampling time (both passes) of a two-pass request with this upscale factor, @@ -94,6 +97,8 @@ def name_for(steps, two_pass, refine_steps=REFINE_STEPS, step_cache=None, mac=Fa if tier['steps'] == steps and tier['two_pass'] == two_pass and (not two_pass or refine_steps == REFINE_STEPS) and tier['step_cache'] == step_cache and tier['upscale'] == upscale), None) + if found is None and (steps, two_pass, step_cache, upscale) == tuple(LEGACY_MAX.values()): + return 'max' if found is None and mac and (steps, two_pass, step_cache, upscale) == tuple(LEGACY_MAC_HIGH.values()): return 'high' return found @@ -102,10 +107,10 @@ def name_for(steps, two_pass, refine_steps=REFINE_STEPS, step_cache=None, mac=Fa def table_level(steps, two_pass, refine_steps=REFINE_STEPS, step_cache=None, mac=False, upscale=2): """The level whose sampling table a plan needs, for naming a missing download. - Tables follow the step count only. Prefer the exact level (the cache tells - high from max at 20 steps), then the other platform's table (16 steps is - the earlier Mac "high"), then the first level with that step count (light - for 8 + 3, whatever the upscale factor). + Tables follow the step count only: 16 steps names high, 20 names max. + Prefer the exact level or a legacy alias, then the other platform's table, + then the first level with that step count (light for 8 + 3, whatever the + upscale factor). """ found = name_for(steps, two_pass, refine_steps, step_cache, mac, upscale) or name_for( steps, two_pass, refine_steps, step_cache, not mac, upscale) diff --git a/freevideo_engine/torch_download.py b/freevideo_engine/torch_download.py index 5bd5bdd0..724b2e93 100644 --- a/freevideo_engine/torch_download.py +++ b/freevideo_engine/torch_download.py @@ -28,7 +28,7 @@ def installed_versions(python): # an isolated venv with the standard Lib/site-packages layout. directory = Path(python).parent.parent / 'Lib' / 'site-packages' return {dist.metadata['Name'].lower(): dist.version - for dist in importlib.metadata.distributions(path=[str(directory)]) + for dist in importlib.metadata.distributions(path=[str(io_path(directory))]) if dist.metadata['Name']} diff --git a/scripts/test_freevideo_entry.mjs b/scripts/test_freevideo_entry.mjs index 65fe2b15..dc1ee8a7 100644 --- a/scripts/test_freevideo_entry.mjs +++ b/scripts/test_freevideo_entry.mjs @@ -126,7 +126,8 @@ test('main browser entry registers both views and preserves node hooks', async t .replace("import { uploadMedia } from './media_upload.js';", 'const {uploadMedia} = globalThis.__freevideoEntryTest;') .replace("import { createCastLibrary } from './cast_library.js';", 'const {createCastLibrary} = globalThis.__freevideoEntryTest;') .replace("import { agentInterfaceEnabled } from './agent_connect.js';", 'const agentInterfaceEnabled = async () => false;') - .replace("import { installNativeErrorGuard } from './native_errors.js';", 'const installNativeErrorGuard = () => {};'); + .replace("import { installNativeErrorGuard } from './native_errors.js';", 'const installNativeErrorGuard = () => {};') + .replace("import { installTranslateGuard } from './translate_guard.js';", 'const installTranslateGuard = () => {};'); await import('data:text/javascript;base64,' + Buffer.from(source).toString('base64')); const extension = extensions.find(row => row.name === 'FreeVideo.UnifiedMedia'); assert.ok(extension, 'the main extension must register'); diff --git a/web/agent_connect.css b/web/agent_connect.css index d5722ec3..5e78fdc5 100644 --- a/web/agent_connect.css +++ b/web/agent_connect.css @@ -5,6 +5,7 @@ .fv-agent-connect-icon{width:44px;height:44px;border-radius:11px;display:grid;place-items:center;background:var(--fv-selected);color:var(--fv-accent-hover)} .fv-agent-connect-body svg{width:24px;height:24px;fill:none;stroke:currentColor;stroke-linecap:round;stroke-linejoin:round} .fv-agent-connect-body>h2{font-size:19px;line-height:1.3;font-weight:700;color:var(--fv-ink)} +.fv-agent-connect-body>.fv-agent-connect-intro{margin-top:-10px;font-size:13px;line-height:1.5;color:#a5b3c6} .fv-agent-connect-body>.fv-muted{margin-top:-10px;font-size:13.5px;color:var(--fv-muted)} .fv-agent-connect-actions{display:flex;align-items:center;gap:12px;margin-top:10px;flex-wrap:wrap} .fv-studio .fv-agent-connect-actions button{padding:9px 20px;font-size:14px}.fv-agent-connect-actions button svg{width:16px;height:16px} diff --git a/web/agent_connect.js b/web/agent_connect.js index 65bf2cd0..085c7133 100644 --- a/web/agent_connect.js +++ b/web/agent_connect.js @@ -36,6 +36,7 @@ export function agentConnectCard(t, api, shown = () => true) { const body = el('section', 'fv-agent-connect-body'); const mark = el('span', 'fv-agent-connect-icon'); mark.append(icon('plug')); const title = el('h2', null, t('Connect an agent', '连接 Agent')); + const intro = el('p', 'fv-agent-connect-intro', t('Once connected, describe the video you want. It generates and edits it for you.', '连上后,说想做什么视频,它来生成和剪辑。')); const subtitle = el('p', 'fv-muted', t('Claude Code, Codex, FreeToken and others', 'Claude Code、Codex、FreeToken 等')); const actions = el('div', 'fv-agent-connect-actions'); const button = el('button', 'fv-primary'); button.type = 'button'; button.disabled = true; @@ -44,7 +45,7 @@ export function agentConnectCard(t, api, shown = () => true) { actions.append(button, arrow, el('span', 'fv-agent-connect-hint', t('Paste to your agent', '粘贴给 Agent'))); const sentence = el('p', 'fv-agent-sentence'); sentence.hidden = true; const status = el('p', 'fv-agent-connect-status', t('Getting the connection address…', '正在获取连接地址…')); status.setAttribute('role', 'status'); - body.append(mark, title, subtitle, actions, sentence, status); + body.append(mark, title, intro, subtitle, actions, sentence, status); let feedback = false, text = '', failures = 0; let timer, feedbackTimer, controller, generation = 0; diff --git a/web/assets/entry-showcase.webp b/web/assets/entry-showcase.webp index 09593385..4d56b0a1 100644 Binary files a/web/assets/entry-showcase.webp and b/web/assets/entry-showcase.webp differ diff --git a/web/capability_view.js b/web/capability_view.js index 524ce7e4..78139bb1 100644 --- a/web/capability_view.js +++ b/web/capability_view.js @@ -1,5 +1,5 @@ import { referenceItems } from './prompt_references.js'; -import { stepCacheOf, upscaleOf } from './sampling_effort.js'; +import { effortFor, effortsFrom, stepCacheOf, upscaleOf } from './sampling_effort.js'; const PRESETS = [.5, .75, .98, 1.5, 2]; const COLUMN_SPACING = 'equal'; // 'equal' or 'width' @@ -77,11 +77,11 @@ export function capabilityContext(node, state) { return query.toString(); } -// The grid's level with exactly this plan; a step cache tells High from Max at 20 steps, the -// upscale factor Light from Medium (a row without one is 2x). -const gridTier = (grid, steps, twoPass, refine, stepCache, upscale = 2) => Object.values(grid.tiers || {}).find(v => v.steps === steps - && v.two_pass === twoPass && (!v.two_pass || refine === 3) && stepCacheOf(v.step_cache) === stepCacheOf(stepCache) - && (!v.two_pass || upscaleOf(v.upscale) === upscaleOf(upscale))); +// Shared recognition includes legacy Max and the Mac-only legacy High. +const gridTier = (grid, steps, twoPass, refine, stepCache, upscale = 2) => { + const level = effortFor(steps, twoPass, refine, stepCache, effortsFrom(grid), upscale); + return grid.tiers?.[level?.name.toLowerCase()]; +}; export function capabilitySurface(grid, state) { if (grid?.version !== 1 || grid.status !== 'ready' || typeof grid.revision !== 'string' || !grid.machine?.system diff --git a/web/chain_support.js b/web/chain_support.js index 218e024d..81014159 100644 --- a/web/chain_support.js +++ b/web/chain_support.js @@ -1,6 +1,6 @@ // Pure helpers the creative panel's long-video planning shares with the editor store; kept out of // editor_store.js so the panel does not load the editing implementation. -import { stepCacheOf, upscaleOf } from './sampling_effort.js'; +import { effortFor, effortsFrom } from './sampling_effort.js'; const clockDefault = {now: () => Date.now(), setTimeout: (fn, ms) => setTimeout(fn, ms), clearTimeout: id => clearTimeout(id)}; export const chainRoute = (id, action) => '/freevideo/chain' + (id ? '/' + encodeURIComponent(id) : '') + (action ? '/' + action : ''); @@ -38,11 +38,9 @@ function capabilityTier(reply, settings) { const tiers = reply.tiers; if (!tiers || typeof tiers !== 'object') return reply; if (tiers[settings?.tier]) return tiers[settings.tier]; - const steps = Number(settings?.base_steps), twoPass = !!settings?.two_pass, cache = stepCacheOf(settings?.step_cache); - const factor = twoPass ? upscaleOf(settings?.upscale_factor) : 2; - return Object.values(tiers).find(tier => tier.steps === steps && tier.two_pass === twoPass - && (!twoPass || Number(settings?.refine_steps) === 3) && stepCacheOf(tier.step_cache) === cache - && (tier.two_pass ? upscaleOf(tier.upscale) : 2) === factor) || {}; + const level = effortFor(Number(settings?.base_steps), !!settings?.two_pass, Number(settings?.refine_steps), + settings?.step_cache, effortsFrom(reply), settings?.upscale_factor); + return tiers[level?.name.toLowerCase()] || {}; } export function capabilityPoints(reply, {width, height, settings = {}}) { const tier = capabilityTier(reply, settings); diff --git a/web/editor.css b/web/editor.css index 7d2cfb2e..f796eee9 100644 --- a/web/editor.css +++ b/web/editor.css @@ -98,6 +98,8 @@ .fv-ed-genbar{position:relative} .fv-ed-refs{display:flex;align-items:center;gap:6px;flex:none} .fv-ed-refs[hidden],.fv-ed-ref-menu[hidden]{display:none} +.fv-ed-ref-notice{position:absolute;bottom:calc(100% + 8px);max-width:min(560px,calc(100% - 24px));border:1px solid var(--fv-border);background:var(--fv-card);padding:10px 14px;border-radius:8px;z-index:9;color:var(--fv-ink);pointer-events:none} +.fv-ed-ref-notice[hidden]{display:none} .fv-ed-ref-chips{display:flex;align-items:center;gap:6px;max-width:min(240px,32vw);overflow-x:auto;scrollbar-width:thin;scrollbar-color:var(--fv-border) transparent;padding:1px} .fv-ed-ref-chips:empty{display:none} /* A strip that scrolls fades at the hidden end(s); one that fits is untouched. */ diff --git a/web/editor.js b/web/editor.js index 71b52787..44230bfb 100644 --- a/web/editor.js +++ b/web/editor.js @@ -272,11 +272,11 @@ export function openEditor({api, app, from = null, panel = {}, components: given } const content = record => JSON.stringify({...record, id: null, revision: 0}); function notify(text) { if (disposed) return; if (!text) { clearTimeout(noticeTimer); notice.hidden = true; return; } notice.textContent = text; notice.hidden = false; clearTimeout(noticeTimer); noticeTimer = setTimeout(() => { notice.hidden = true; }, 9000); } - // `step` names what was being done, so a failure there says which step went wrong. - function error(value, step = null) { + // `step` names what was being done, so a failure there says which step went wrong; `out` is where the line is said (the bar's own line for a reference picture). + function error(value, step = null, out = notify) { const code = value?.reason || value?.code || value?.message; if (['missing_glyphs', 'MissingGlyphs', 'source_missing'].includes(code)) { - notify(code === 'source_missing' ? t('A source file is missing. Import it again from the library.', '有素材的原文件找不到了,请在素材库重新导入。') + out(code === 'source_missing' ? t('A source file is missing. Import it again from the library.', '有素材的原文件找不到了,请在素材库重新导入。') : t('The chosen font is missing some characters. Choose another font.', '所选字体缺少部分字符,请换一种字体。')); return {ok: false, reason: code}; } @@ -286,7 +286,7 @@ export function openEditor({api, app, from = null, panel = {}, components: given const [en, zh] = (value.said || []).map(text => typeof text === 'string' && text ? text.slice(0, 200) : null); const text = code === 'chain_too_long' ? t(en || 'This video can have up to 5 parts. Make it shorter.', zh || '这条视频最多 5 段,请缩短一些。') : t(en || 'Each continued part must be at least 10 seconds. Make it longer.', zh || '接着生成的每一段至少 10 秒,请加长一些。'); - notify(text); return {ok: false, reason: code, say: text}; + out(text); return {ok: false, reason: code, say: text}; } if (['busy', 'locked', 'text', 'audio_full', 'overlap', 'range', 'no_audio'].includes(code)) { if (['locked', 'audio_full', 'overlap', 'range'].includes(code)) tracks?.flashBlocked?.(value.displayFrame); @@ -294,11 +294,11 @@ export function openEditor({api, app, from = null, panel = {}, components: given text: t('Write a description first.', '先写一句描述。'), overlap: t('That spot is taken. Use an empty spot.', '这里已经有片段了,请放到空白处。'), audio_full: t('Both audio tracks are full here. Move an audio clip first.', '这里两条声音轨都满了,请先挪开一段声音。'), no_audio: t('This clip has no sound. Please choose one with sound.', '这个素材没有声音,请选一段有声音的。')}[code]; - if (said) notify(said); + if (said) out(said); return {ok: false, reason: code}; } const text = step === 'generate' ? editorErrorText(value, t) : step === 'export' ? exportErrorText(value, t) : t('Could not complete. Try again.', '没有完成,请重试'); - notify(text); return {ok: false, reason: code, say: text}; + out(text); return {ok: false, reason: code, say: text}; } const store = createEditorStore(api, {t, notify, standIn: given?.generationStandIn, onGenerationReady: ready => { if (disposed) return; diff --git a/web/editor_genbar.js b/web/editor_genbar.js index 2642e18f..0ee29048 100644 --- a/web/editor_genbar.js +++ b/web/editor_genbar.js @@ -59,8 +59,18 @@ export function createEditorGenbar(host, {state, t, actions, text = ''}) { failureBar.append(failureReason, report); failure.append(failureBar); failure.hidden = true; host.append(failure); let failureId = null, downloading = false; // Reference pictures for a part that starts fresh: thumbnails and a small "+" before the description. - const references = actions.references ? createEditorReferences({t, source: actions.references.source, size: actions.references.size, error: actions.error, notify: actions.references.notify}) : null; - if (references) description.before(references.element); + // What goes wrong with a reference picture is said above the "+", where the eyes are, not in the page-wide line. + const refNotice = el('div', 'fv-ed-ref-notice'); refNotice.hidden = true; refNotice.setAttribute('role', 'status'); + let refNoticeTimer = null; + const sayAtRefs = text => { + clearTimeout(refNoticeTimer); + if (!text || disposed) { refNotice.hidden = true; return; } + refNotice.textContent = text; refNotice.hidden = false; + refNotice.style.left = Math.max(0, (references?.element.getBoundingClientRect().left || 0) - host.getBoundingClientRect().left) + 'px'; + refNoticeTimer = setTimeout(() => { refNotice.hidden = true; }, 9000); + }; + const references = actions.references ? createEditorReferences({t, source: actions.references.source, size: actions.references.size, error: value => actions.error(value, null, sayAtRefs), notify: sayAtRefs}) : null; + if (references) { description.before(references.element); host.append(refNotice); } let status = null, serial = 0, disposed = false, pressing = new Set(), previewSelection = null, estimateCache = null, noteTimer = null, queuedAt = null; // Which parts are generating, or how the job ended: a "Queued" line goes once this moves on. const progressKey = () => ['done', 'failed', 'cancelled', 'interrupted'].includes(status?.state) ? status.state @@ -223,7 +233,7 @@ export function createEditorGenbar(host, {state, t, actions, text = ''}) { }); render(); return {element: host, say, downloadFailure, getText() { expand.commit(); return input.value; }, update(value) { status = value; render(); }, preview(selection) { previewSelection = selection; render(); }, setText(value) { input.value = value || ''; state.setRangeText(input.value); render(); }, - dispose() { disposed = true; ++serial; clearTimeout(noteTimer); unsubscribe(); resize?.disconnect(); expand.dispose(); references?.dispose(); assets?.dispose(); }}; + dispose() { disposed = true; ++serial; clearTimeout(noteTimer); clearTimeout(refNoticeTimer); unsubscribe(); resize?.disconnect(); expand.dispose(); references?.dispose(); assets?.dispose(); }}; } // A missing optional module leaves the editor's existing progress in place. diff --git a/web/editor_references.js b/web/editor_references.js index 9dbb7903..df8044f6 100644 --- a/web/editor_references.js +++ b/web/editor_references.js @@ -69,15 +69,27 @@ export function createEditorReferences({t, source, size = () => null, error = () return holder; }; const write = rows => { try { source.write(rows); } catch (value) { error(value); } }; + // A picture the page cannot draw would only fail minutes later, in the engine; tiff is not drawn by every browser, so it goes on unchecked. + const readable = async file => { + if (!/\.(png|jpe?g|webp|bmp)$/i.test(file.name || '')) return true; + if (typeof Image !== 'function') return true; + const url = URL.createObjectURL(file), image = new Image(); + try { image.src = url; await image.decode(); return image.naturalWidth > 0 && image.naturalHeight > 0; } + catch { return false; } + finally { URL.revokeObjectURL(url); } + }; async function addFiles(files) { if (working) return; working = true; update(); + let skipped = false; try { for (const file of files) { if (full()) break; + if (!await readable(file)) { skipped = true; continue; } const result = await source.upload(file); write([...source.rows(), {file: result.file, role: 'reference', enabled: true}]); } + if (skipped && !disposed) notify(t('Could not read this picture. Please choose another one.', '无法读取这张图片,请换一张。')); } catch (value) { error(value); } finally { working = false; picker.value = ''; update(); } } diff --git a/web/entry_showcase.js b/web/entry_showcase.js index e4202e11..fdfd7896 100644 --- a/web/entry_showcase.js +++ b/web/entry_showcase.js @@ -2,15 +2,15 @@ // Each eligible product version gets one chapter in the existing guide.json seen list. export const SHOWCASE_START_VERSION = '0.4.10'; // The ready-made white-model scene the showcase's button starts. -export const SHOWCASE_SCENE = 'car'; +export const SHOWCASE_SCENE = 'cat'; export const SHOWCASE_ASSET = new URL('./assets/entry-showcase.webp', import.meta.url).href; export const SHOWCASE_TEXT = { title: ['From whitebox animation to finished video', '从白模动画到写实视频'], try: ['Get started', '开始使用'], close: ['Close', '关闭'], example: [ - "A red race car hits the end of a straight at full speed and throws itself into a corner, the rear sliding out in a controlled drift with a little tyre smoke, while the camera orbits half a circle around the car and then follows it out of the bend. Sunny circuit with red and white kerbs, grandstands and green grass, realistic paint and tarmac, strong sense of speed, smooth camera. Roaring engine and the squeal of tyres.", - "一辆红色赛车在直道尽头全速冲入弯道,车尾受控地甩出,带起少许轮胎烟,镜头绕车半圈后跟着它驶出弯道。晴天的赛道,红白相间的路缘石、看台和绿草地,写实的漆面和柏油路,速度感强,镜头平滑。引擎轰鸣,轮胎尖啸。", + "A ginger cat crouches on a living room floor, gathers itself and leaps up a staggered wooden bookshelf, jumping from one shelf to the next between rows of books, and finally settles on the top shelf, while the camera pans sideways and tilts up slightly to follow. Warm indoor daylight, realistic fur and wood, soft shadows, smooth camera motion. Soft room tone and the light thud of paws on wood.", + "一只橘猫在客厅地板上蹲下蓄力,跳上错落的木书架,在成排的书之间一层一层往上跳,最后稳稳落在最高一层,镜头向侧面平移并略微上摇跟随。温暖的室内日光,写实的毛发和木纹,柔和的阴影,镜头平滑。安静的室内环境声,爪子落在木头上的轻响。", ], }; diff --git a/web/failure_issue.js b/web/failure_issue.js index 66d3820f..069d5a49 100644 --- a/web/failure_issue.js +++ b/web/failure_issue.js @@ -20,6 +20,18 @@ export function failureStage(report) { return null; } +const TAIL = 1500, TAIL_HEAD = 500; + +// The engine's failure text, not ComfyUI's traceback after it: details end with +// the bridge's own frames. The error and stage lead that text and the child +// log's error ends it, so a long one gives up its middle. +function cause(report) { + const message = String(report?.message || ''); + const text = message.includes('\n') ? message : String(report?.details || ''); + if (text.length <= TAIL) return text; + return text.slice(0, TAIL_HEAD) + '\n...\n' + text.slice(-(TAIL - TAIL_HEAD - 5)); +} + /** Title, text and summary for the prefilled issue. `lora` marks a LoRA the check judged wrongly. */ export function failureIssue(report, kind, {lora = false, t = en => en} = {}) { const stage = failureStage(report); @@ -32,7 +44,7 @@ export function failureIssue(report, kind, {lora = false, t = en => en} = {}) { const types = (report?.nodes || []).map(node => node.type).filter(Boolean); const summary = JSON.stringify({kind: kind || 'unknown', stage: report?.stage, exception_type: report?.exception_type, known_error: report?.known_error, nodes: types, settings: report?.settings, progress: report?.progress, - tail: String(report?.details || '').slice(-1500)}); + tail: cause(report)}); const title = lora ? t('[Feedback] LoRA compatibility misjudged', '[Feedback] LoRA 兼容判断有误') : t('[Feedback] Generation did not finish', '[Feedback] 生成未完成') + (stage ? t(`: ${stage[0]}`, `:${stage[1]}`) : ''); return {title, text, summary}; diff --git a/web/freevideo.js b/web/freevideo.js index ae5fb800..5f5de117 100644 --- a/web/freevideo.js +++ b/web/freevideo.js @@ -1,5 +1,6 @@ import { createErrorPanel, errorText, createErrorReport, adviceFor } from './error_panel.js'; import { installNativeErrorGuard } from './native_errors.js'; +import { installTranslateGuard } from './translate_guard.js'; import { app } from "../../scripts/app.js"; import { api } from "../../scripts/api.js"; import { openEditorView } from './editor_loader.js'; @@ -846,6 +847,7 @@ api.addEventListener('freevideo_prewarm', event => { if (value) app.graph?.getNodeById(value.node)?.freevideoShowPrewarm?.(value); }); +installTranslateGuard({chinese: cn}); installNativeErrorGuard({api, app, describe: detail => { const {advice} = adviceFor(createErrorReport(detail, {stage: 'execution', graph: app.graph}), text); return [advice.title, advice.summary].filter(Boolean).join('\n'); diff --git a/web/health.js b/web/health.js index 3647b1f8..050c2f65 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', 'caret_rect.js', 'chain_support.js', 'capability_view.js', 'cast_generate.js', 'cast_mention.js', 'cast_overlay.js', 'cast_source.js', 'compatibility.js', 'entry_showcase.js', 'error_panel.js', 'first_run_guide.js', 'fonts.js', 'freevideo.js', 'generation_progress.js', 'editor.js', 'editor_actions.js', 'editor_lifecycle.js', 'editor_adjust.js', 'editor_adjust_picture.js', 'editor_picture_fields.js', 'editor_picture_session.js', 'editor_preview_picture.js', 'editor_entry.js', 'editor_expand.js', 'editor_export_dialog.js', 'editor_first_frame.js', 'editor_media_references.js', 'editor_references.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_failure.js', 'editor_header_fit.js', - 'editor_loader.js', 'editor_text_coverage.js', 'editor_text.js', 'editor_toolbar_more.js', 'editor_text_model.js', 'editor_text_layout.js', 'editor_notes.js', 'editor_icons.js', 'dialog_keys.js', 'done_notify.js', 'failure_issue.js', 'feedback.js', 'frame_sizes.js', 'health.js', 'image_commands.js', 'image_editor.js', 'image_editor_engine.js', 'image_generation.js', 'launch_visit.js', 'launcher.js', 'length_plan.js', 'lora_triggers.js', 'lora_check.js', 'library.js', 'motion.js', 'native_errors.js', 'output_download.js', 'port_notice.js', 'preview_scene.js', + 'editor_loader.js', 'editor_text_coverage.js', 'editor_text.js', 'editor_toolbar_more.js', 'editor_text_model.js', 'editor_text_layout.js', 'editor_notes.js', 'editor_icons.js', 'dialog_keys.js', 'done_notify.js', 'failure_issue.js', 'feedback.js', 'frame_sizes.js', 'health.js', 'image_commands.js', 'image_editor.js', 'image_editor_engine.js', 'image_generation.js', 'launch_visit.js', 'launcher.js', 'length_plan.js', 'lora_triggers.js', 'lora_check.js', 'library.js', 'motion.js', 'native_errors.js', 'translate_guard.js', 'output_download.js', 'port_notice.js', 'preview_scene.js', 'progress_connection.js', 'prompt_draft.js', 'prompt_enhance.js', 'prompt_references.js', 'prompt_terms.js', 'request_assets.js', 'repair_hint.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', 'optional_inputs.js', 'studio_queue.js', 'studio_size.js', 'toolbar_position.js', 'updates.js', 'view_navigation.js', 'wordmark.js', 'workspace_restore.js', 'cast_library.js', 'media_upload.js', 'whitebox_encode.js', 'whitebox_mp4.js', 'whitebox_recording.js', 'whitebox_render.js', 'whitebox_scenes.js', 'whitebox_worker.js']; diff --git a/web/image_editor.css b/web/image_editor.css index 40d9d1a1..013cf4ea 100644 --- a/web/image_editor.css +++ b/web/image_editor.css @@ -135,7 +135,7 @@ button.fv-ie-size:hover{background:var(--fv-hover);color:var(--fv-ink)} .fv-ie-view,.fv-ie-ants,.fv-ie-overlay{position:absolute;inset:0;width:100%;height:100%} .fv-ie-ants{pointer-events:none} .fv-ie-overlay{overflow:visible} -.fv-ie-stage[data-tool=move]{cursor:default}.fv-ie-stage[data-tool=crop],.fv-ie-stage[data-tool=marquee],.fv-ie-stage[data-tool=lasso],.fv-ie-stage[data-tool=wand],.fv-ie-stage[data-tool=gradient],.fv-ie-stage[data-tool=shape]{cursor:crosshair} +.fv-ie-stage[data-tool=move]{cursor:default}.fv-ie-stage[data-tool=crop],.fv-ie-stage[data-tool=marquee],.fv-ie-stage[data-tool=lasso],.fv-ie-stage[data-tool=wand],.fv-ie-stage[data-tool=point],.fv-ie-stage[data-pick=true],.fv-ie-stage[data-tool=gradient],.fv-ie-stage[data-tool=shape]{cursor:crosshair} .fv-ie-stage[data-tool=brush],.fv-ie-stage[data-tool=eraser],.fv-ie-stage[data-tool=heal],.fv-ie-stage[data-tool=clone],.fv-ie-stage[data-tool=blur]{cursor:none} .fv-ie-stage[data-tool=text]{cursor:text}.fv-ie-stage[data-tool=eyedropper]{cursor:copy} .fv-ie-stage[data-tool=hand],.fv-ie-stage[data-pan=true]{cursor:grab}.fv-ie-stage[data-pan=true]:active{cursor:grabbing} @@ -239,7 +239,7 @@ button.fv-ie-size:hover{background:var(--fv-hover);color:var(--fv-ink)} .fv-ie-confirm span{flex:1 1 220px} .fv-ie-status{padding:0 16px;color:var(--fv-muted);font-size:var(--fv-micro)} .fv-ie-status:empty{display:none} -.fv-ie-status:not(:empty){padding:7px 16px;border-top:1px solid var(--fv-border)} +.fv-ie-status:not(:empty){position:absolute;left:0;right:0;bottom:0;padding:7px 16px;border-top:1px solid var(--fv-border);background:var(--fv-panel,#0f1620);z-index:3;pointer-events:none} @media (max-width:900px){ .fv-image-editor{width:100vw;height:100dvh;border:0;border-radius:0} diff --git a/web/image_editor.js b/web/image_editor.js index 5a00f216..c15efc3d 100644 --- a/web/image_editor.js +++ b/web/image_editor.js @@ -4,11 +4,11 @@ // 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, blankDocument, orientedSize, orientMatrix, canvasRect, layerMatrix, scaleSelection, - rasterKinds, measureText, Renderer} from './image_editor_engine.js'; + rasterKinds, measureText, Renderer, traceOutline, simplifyLoop} 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'; +import {createImageGenClient, createImageGeneration, imageGenerationSupported, blankCanvasSize, blankCanvasRatio, segmentAt} from './image_generation.js'; import {productVersion} from './updates.js'; const css = document.createElement('link'); @@ -29,6 +29,7 @@ export const ICONS = { lasso: '', polygon: '', wand: '', + point: '', brush: '', eraser: '', heal: '', @@ -166,7 +167,7 @@ export async function openImageEditor({source = null, name = source || '', targe let tool = 'move'; const opts = { marquee: {shape: 'rect', mode: 'replace', feather: 0}, lasso: {type: 'free', mode: 'replace', feather: 0}, - wand: {tolerance: 24, contiguous: true, mode: 'replace'}, + wand: {tolerance: 24, contiguous: true, mode: 'replace'}, point: {mode: 'replace'}, brush: {size: .02, hardness: .85, opacity: 1}, eraser: {size: .035, hardness: .7, opacity: 1}, heal: {size: .025, hardness: .6, text: false}, clone: {size: .035, hardness: .6, opacity: 1, source: null, offset: null, picking: false}, blur: {size: .05, hardness: .4, strength: .6}, gradient: {shape: 'linear', end: 'bg'}, @@ -225,7 +226,8 @@ export async function openImageEditor({source = null, name = source || '', targe 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)), + const genClient = imageGenClient || createImageGenClient(api.fetchApi.bind(api)); + const generation = createImageGeneration({client: genClient, renderer, t, getDoc: () => doc, getSelection: () => selection, commit: commitInserted, el, button, segmented, library, resolve, icon, onRatioChange: setBlankCanvasRatio, productVersion, fetchApi: (path, options) => api.fetchApi(path, options), resumeKey: `${role}|${source || ''}|${name}`}); @@ -234,7 +236,7 @@ export async function openImageEditor({source = null, name = source || '', targe 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, + ['marquee', t('Marquee', '选框'), 'M'], ['lasso', t('Lasso', '套索'), 'L'], ['wand', t('Magic wand', '魔棒'), 'W'], ['point', t('Click select', '点选'), 'P'], null, ['brush', t('Brush', '画笔'), 'B'], ['eraser', t('Eraser', '橡皮擦'), 'E'], ['heal', t('Spot healing', '污点修复'), 'J'], ['clone', t('Clone stamp', '仿制图章'), 'S'], ['blur', t('Blur', '模糊'), 'R'], ['gradient', t('Gradient', '渐变'), 'G'], null, ['text', t('Text', '文字'), 'T'], ['shape', t('Shape', '形状'), 'U'], ['eyedropper', t('Eyedropper', '吸管'), 'I'], ['hand', t('Hand', '抓手'), 'H'], @@ -243,9 +245,27 @@ export async function openImageEditor({source = null, name = source || '', targe for (const entry of TOOLS) { if (!entry) { toolbar.append(el('span', null, 'fv-ie-tool-gap')); continue; } const [id, label, key] = entry; - if (id === 'generate' && !imageGenerationSupported()) continue; + if ((id === 'generate' || id === 'point') && !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); + if (id === 'point') b.hidden = true; + } + // Click select needs the engine's selection pack; the tool appears once the engine says it offers it. + let pointInfo = null, pointBusy = false, pointTimer = null; + async function refreshPoint() { + try { pointInfo = (await genClient.info()).segment || null; } catch { pointInfo = null; } + if (toolButtons.point) toolButtons.point.hidden = !pointInfo; + if (tool === 'point') { if (!pointInfo) select('move'); else renderOptions(); } + clearTimeout(pointTimer); pointTimer = null; + if (pointInfo?.install?.state === 'running') pointTimer = setTimeout(refreshPoint, 1500); + } + if (toolButtons.point) void refreshPoint(); + // A click before the model is here answers with the one thing to do, not with a line of text. + function flashPointDownload() { + const b = optionsBar.querySelector('.fv-ie-point-download'); + if (!b) return; + b.classList.add('fv-ie-active'); + setTimeout(() => b.classList.remove('fv-ie-active'), 700); } const colors = el('div', null, 'fv-ie-colors'); const fgInput = el('input'); fgInput.type = 'color'; fgInput.className = 'fv-ie-fg'; fgInput.value = fg; @@ -447,15 +467,22 @@ export async function openImageEditor({source = null, name = source || '', targe } // Marching ants: the selection's edge, striped, moving. let antsEdge = null; + const exactEdge = () => layout.k * layout.dpr < 1 && !selection?.parts.some(p => p.shape === 'wand' || p.shape === 'mask'); + const edgeScale = () => exactEdge() ? layout.k * layout.dpr : previewScale(); function selectionEdge() { - const s = previewScale(), key = JSON.stringify([selection, s]) + (selection?.parts.some(p => p.shape === 'wand') ? composited?.doc === doc : ''); - if (antsEdge?.key === key && antsEdge.doc === doc) return antsEdge.canvas; - const mask = renderer.selectionMask(doc, selection, {layer: canvasLayer(), scale: s}); - const shifted = (dx, dy) => { const c = document.createElement('canvas'); c.width = mask.width; c.height = mask.height; const x = c.getContext('2d'); - x.drawImage(mask, 0, 0); x.globalCompositeOperation = 'xor'; x.drawImage(mask, dx, dy); return c; }; - const edge = shifted(1, 0); edge.getContext('2d').drawImage(shifted(0, 1), 0, 0); - antsEdge = {key, doc, canvas: edge}; - return edge; + // Zoomed out, the edge is made at the screen's own pixel size so every side stays one whole line. + const s = edgeScale(), key = JSON.stringify([selection, s]) + (selection?.parts.some(p => p.shape === 'wand' || p.shape === 'mask') ? composited?.doc === doc : ''); + if (antsEdge?.key === key && antsEdge.doc === doc) return antsEdge; + const soft = renderer.selectionMask(doc, selection, {layer: canvasLayer(), scale: s}); + // The outline follows the half-way line of the selection; feathering shows only in the result. + const mask = document.createElement('canvas'); mask.width = soft.width; mask.height = soft.height; + const m = mask.getContext('2d', {willReadFrequently: true}); m.drawImage(soft, 0, 0); + const bits = m.getImageData(0, 0, mask.width, mask.height); + for (let i = 3; i < bits.data.length; i += 4) bits.data[i] = bits.data[i] >= 128 ? 255 : 0; + m.putImageData(bits, 0, 0); + const loops = traceOutline(bits.data, mask.width, mask.height).map(loop => simplifyLoop(loop, 0.75)); + antsEdge = {key, doc, loops, w: mask.width, h: mask.height}; + return antsEdge; } function drawAnts() { const ctx = ants.getContext('2d'); @@ -463,15 +490,17 @@ export async function openImageEditor({source = null, name = source || '', targe if (!selection) return; const edge = selectionEdge(), r = rect(), [x0, y0] = fromO(r.x, r.y), dpr = layout.dpr; ctx.setTransform(dpr, 0, 0, dpr, 0, 0); - ctx.imageSmoothingEnabled = false; - ctx.drawImage(edge, x0, y0, r.w * layout.k, r.h * layout.k); - ctx.globalCompositeOperation = 'source-in'; - const stripe = document.createElement('canvas'); stripe.width = stripe.height = 8; const st = stripe.getContext('2d'); - st.fillStyle = '#fff'; st.fillRect(0, 0, 8, 8); st.fillStyle = '#0b0f15'; st.fillRect(0, 0, 4, 8); - const pattern = ctx.createPattern(stripe, 'repeat'); - pattern.setTransform?.(new DOMMatrix().translateSelf(antsPhase, 0).rotateSelf(45)); - ctx.fillStyle = pattern; ctx.fillRect(0, 0, layout.w, layout.h); - ctx.globalCompositeOperation = 'source-over'; + // One closed polyline per border, every side drawn with the same dash phase. + const sx = (exactEdge() ? edge.w / dpr : r.w * layout.k) / edge.w, sy = (exactEdge() ? edge.h / dpr : r.h * layout.k) / edge.h; + const path = new Path2D(); + for (const loop of edge.loops) { + loop.forEach(([x, y], i) => (i ? path.lineTo(x0 + x * sx, y0 + y * sy) : path.moveTo(x0 + x * sx, y0 + y * sy))); + path.closePath(); + } + ctx.lineWidth = 1.5; ctx.lineJoin = 'round'; ctx.lineCap = 'butt'; + ctx.setLineDash([]); ctx.strokeStyle = '#fff'; ctx.stroke(path); + ctx.setLineDash([4, 4]); ctx.lineDashOffset = -antsPhase; ctx.strokeStyle = '#0b0f15'; ctx.stroke(path); + ctx.setLineDash([]); } function canvasLayer() { const r = rect(), centre = toD([r.x + r.w / 2, r.y + r.h / 2]); @@ -490,6 +519,7 @@ export async function openImageEditor({source = null, name = source || '', targe b.setAttribute('aria-pressed', String(id === tool)); } stage.dataset.tool = tool; + stage.dataset.pick = String(tool === 'generate' && generation.drawsSelection()); } function select(id) { if (!doc.layers.length && !['generate', 'hand'].includes(id)) { message(emptyCanvasReason()); return; } @@ -562,6 +592,20 @@ export async function openImageEditor({source = null, name = source || '', targe add(sep(), segmented([['replace', t('New', '新选区')], ['add', t('Add · Shift', '添加 · Shift')], ['subtract', t('Subtract · Alt', '减去 · Alt')]], o.mode, v => { o.mode = v; }, t('Mode', '方式')).group); if (tool !== 'wand') add(slider(`fv-ie-${tool}-feather`, t('Feather', '羽化'), 0, 80, 1, o.feather, v => { o.feather = v; }, null, ' px').row); add(sep(), ...selectionActions()); + } else if (tool === 'point') { + const o = opts.point, install = pointInfo?.install; + if (pointInfo && !pointInfo.ready) { + const running = install?.state === 'running'; + const share = install?.total > 0 ? Math.min(100, Math.round(100 * (install.done || 0) / install.total)) : 0; + const b = running ? button(t(`Downloading ${share}%`, `正在下载 ${share}%`), null, () => {}, 'fv-ie-button fv-ie-point-download') + : button(t('Download click-select model', '下载点选模型'), install?.state === 'failed' ? t(install.message_en, install.message_zh) : null, + async () => { try { pointInfo = {...pointInfo, install: (await genClient.segmentInstall()).install}; } catch { /* the next refresh shows it */ } void refreshPoint(); renderOptions(); }, 'fv-ie-button fv-ie-point-download'); + if (running) b.disabled = true; + add(b); + } else { + add(segmented([['replace', t('New', '新选区')], ['add', t('Add · Shift', '添加 · Shift')], ['subtract', t('Subtract · Alt', '减去 · Alt')]], o.mode, v => { o.mode = v; }, t('Mode', '方式')).group); + } + add(sep(), ...selectionActions()); } else if (tool === 'brush' || tool === 'eraser') { const o = opts[tool]; add(sizeSlider(o), pct(o, 'hardness', t('Hardness', '硬度')), pct(o, 'opacity', t('Opacity', '不透明度'))); @@ -1171,12 +1215,15 @@ export async function openImageEditor({source = null, name = source || '', targe return null; } function startGesture(event, p) { - if (!doc.layers.length || tool === 'generate') return null; + // In "Edit this" dragging on the picture selects the part to change, as the marquee does. + const choosing = tool === 'generate' && generation.drawsSelection(); + if (!doc.layers.length || tool === 'generate' && !choosing) return null; + const use = choosing ? 'marquee' : tool; 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); - if (tool === 'move') return moveGesture(event, p, d); - if (tool === 'marquee') { + if (use === 'crop') return cropGesture(event, o); + if (use === 'move') return moveGesture(event, p, d); + if (use === 'marquee') { const om = opts.marquee; return dragBox(o, (a, b) => { const [x1, y1] = toD(a), [x2, y2] = toD(b); @@ -1185,15 +1232,32 @@ export async function openImageEditor({source = null, name = source || '', targe addPart(part, mode || om.mode); }, om.shape === 'ellipse'); } - if (tool === 'lasso') return lassoGesture(o, mode || opts.lasso.mode); - if (tool === 'wand') { + if (use === 'lasso') return lassoGesture(o, mode || opts.lasso.mode); + if (use === 'wand') { const src = rasterKinds.has(layer.kind) ? layer : doc.layers[0]; const [u, v] = toLocal(src, d); if (u < 0 || v < 0 || u > 1 || v > 1) { message(t('Click on an image layer.', '请点在图片图层上。')); return null; } addPart({shape: 'wand', layer: src.id, u, v, tolerance: opts.wand.tolerance, contiguous: opts.wand.contiguous, upto: src.ops.length}, mode || opts.wand.mode); return null; } - if (tool === 'eyedropper') { + if (use === 'point') { + const src = rasterKinds.has(layer.kind) ? layer : doc.layers[0]; + const [u, v] = toLocal(src, d); + if (u < 0 || v < 0 || u > 1 || v > 1) { message(t('Click on an image layer.', '请点在图片图层上。')); return null; } + if (!pointInfo?.ready) { flashPointDownload(); return null; } + if (pointBusy) return null; + const how = mode || opts.point.mode; + pointBusy = true; message(t('Finding the object…', '正在找出物体……')); + segmentAt(renderer, genClient, doc, src, [u, v]).then(file => { + if (doc.layers.indexOf(src) < 0) return; + addPart({shape: 'mask', layer: src.id, file, upto: src.ops.length}, how); message(''); + }).catch(error => { + message(error?.code === 'busy_video' ? t('A video is generating. Try again when it finishes.', '视频正在生成,完成后再试。') + : t('Could not find an object here. Try again.', '没能在这里找出物体,请再试一次。')); + }).finally(() => { pointBusy = false; }); + return null; + } + if (use === 'eyedropper') { const {canvas: pic, scale} = picture(), r = rect(); const x = Math.floor((o[0] - r.x) * scale), y = Math.floor((o[1] - r.y) * scale); if (x < 0 || y < 0 || x >= pic.width || y >= pic.height) { message(t('Click on the canvas to pick a color.', '请点在画布上取色。')); return null; } @@ -1204,35 +1268,35 @@ export async function openImageEditor({source = null, name = source || '', targe message(t(`Picked ${hex}.`, `已吸取 ${hex}。`)); return null; } - if (tool === 'text') { + if (use === 'text') { const hit = hitLayer(p); if (hit?.kind === 'text') { setActive(hit.id); renderPanels(); textArea?.focus(); return null; } addText(d); return null; } - if (tool === 'shape') return shapeGesture(o); + if (use === 'shape') return shapeGesture(o); // Pixel tools. let target = layer, base = doc; - if ((tool === 'brush' || tool === 'gradient') && (!rasterKinds.has(layer.kind) || (layer.locked && !maskTarget))) { + if ((use === 'brush' || use === 'gradient') && (!rasterKinds.has(layer.kind) || (layer.locked && !maskTarget))) { // Paint goes on a layer of its own above the picture, so it can be erased or hidden alone. - const reuse = tool === 'brush' && [...doc.layers].reverse().find(l => l.kind === 'paint' && l.visible); + const reuse = use === 'brush' && [...doc.layers].reverse().find(l => l.kind === 'paint' && l.visible); if (reuse) { target = reuse; base = commands.setActive(doc, reuse.id); } else { const n = doc.layers.filter(l => l.kind === 'paint').length + 1; - base = commands.addPaintLayer(doc, tool === 'gradient' ? t('Gradient', '渐变') : t(`Paint ${n}`, `绘画 ${n}`)); + base = commands.addPaintLayer(doc, use === 'gradient' ? t('Gradient', '渐变') : t(`Paint ${n}`, `绘画 ${n}`)); target = base.layers.find(l => l.id === base.active); } } else if (!rasterKinds.has(layer.kind)) { message(t('Pick an image or paint layer to work on its pixels.', '请选择图片图层或绘画图层再处理像素。')); return null; } - if (tool === 'gradient') { + if (use === 'gradient') { return dragLine(o, (a, b) => { const from = toLocal(target, toD(a)), to = toLocal(target, toD(b)); pushOp(target, {type: 'gradient', shape: opts.gradient.shape, from, to, c1: fg, c2: opts.gradient.end === 'bg' ? bg : 'transparent', selection}, base); }); } - if (tool === 'clone' && (event.altKey || opts.clone.picking)) { + if (use === 'clone' && (event.altKey || opts.clone.picking)) { opts.clone.source = d; opts.clone.picking = false; opts.clone.offset = null; renderOptions(); showHint(); drawOverlay(); return null; } - if (tool === 'clone' && !opts.clone.source) { message(t('Alt-click to set the source first.', '请先按住 Alt 点击设置取样点。')); return null; } + if (use === 'clone' && !opts.clone.source) { message(t('Alt-click to set the source first.', '请先按住 Alt 点击设置取样点。')); return null; } return strokeGesture(target, d, base); } function strokeGesture(layer, d, base) { @@ -1477,7 +1541,7 @@ export async function openImageEditor({source = null, name = source || '', targe function actualSize() { viewState = {zoom: 1 / layout.fitK, panX: 0, panY: 0}; computeLayout(); draw(); controls(); } // ----- keys ---------------------------------------------------------------------------------- - const KEYS = {v: 'move', c: 'crop', m: 'marquee', l: 'lasso', w: 'wand', b: 'brush', e: 'eraser', j: 'heal', s: 'clone', r: 'blur', g: 'gradient', t: 'text', u: 'shape', i: 'eyedropper', h: 'hand'}; + const KEYS = {v: 'move', c: 'crop', m: 'marquee', l: 'lasso', w: 'wand', p: 'point', b: 'brush', e: 'eraser', j: 'heal', s: 'clone', r: 'blur', g: 'gradient', t: 'text', u: 'shape', i: 'eyedropper', h: 'hand'}; function keys(event) { const typing = /^(INPUT|TEXTAREA|SELECT)$/.test(event.target.tagName) && !['range', 'checkbox', 'color', 'radio'].includes(event.target.type); const mod = event.ctrlKey || event.metaKey, key = event.key.toLowerCase(); diff --git a/web/image_editor_engine.js b/web/image_editor_engine.js index 896451c8..15ea93cd 100644 --- a/web/image_editor_engine.js +++ b/web/image_editor_engine.js @@ -972,6 +972,23 @@ export class Renderer { return this.images.get(file); } loaded(file) { return this.images.get(file)?.value; } + // A black and white picture as transparent pixels (white kept), made once it has loaded. + maskShape(file) { + const image = this.loaded(file); + if (!image) return null; + this.masks ||= new Map(); + let shape = this.masks.get(file); + if (!shape) { + shape = canvas(image.naturalWidth, image.naturalHeight); + const ctx = context(shape); + ctx.drawImage(image, 0, 0); + const data = ctx.getImageData(0, 0, shape.width, shape.height); + for (let i = 0; i < data.data.length; i += 4) { data.data[i + 3] = data.data[i] >= 128 ? 255 : 0; data.data[i] = data.data[i + 1] = data.data[i + 2] = 255; } + ctx.putImageData(data, 0, 0); + this.masks.set(file, shape); + } + return shape; + } async preload(doc) { for (const l of doc.layers) if (l.src) { const img = await this.image(l.src); this.images.get(l.src).value = img; } await loadFonts(textFonts(doc)).catch(error => console.warn('FreeVideo image editor fonts unavailable; using fallback fonts.', error)); @@ -1103,6 +1120,14 @@ export class Renderer { const toSrc = new DOMMatrix().scaleSelf(scale, scale).multiplySelf(layerMatrix(layer).inverse()).multiplySelf(layerMatrix(src)).scaleSelf(1 / scale, 1 / scale); ctx.setTransform(toSrc); ctx.drawImage(m, 0, 0); } + } else if (part.shape === 'mask') { + // A picture the service made from a click: white where the object is, the size of its layer's pixels. + const src = doc.layers.find(l => l.id === part.layer), shape = this.maskShape(part.file); + if (src && shape) { + const toSrc = new DOMMatrix().scaleSelf(scale, scale).multiplySelf(layerMatrix(layer).inverse()).multiplySelf(layerMatrix(src)) + .scaleSelf(src.nw / shape.width, src.nh / shape.height); + ctx.setTransform(toSrc); ctx.imageSmoothingEnabled = false; ctx.drawImage(shape, 0, 0); + } } else { ctx.setTransform(toLocal); ctx.fillStyle = '#fff'; @@ -1220,3 +1245,57 @@ export class Renderer { this.layers.clear(); this.images.clear(); } } + +// The outline of a black and white picture as closed polylines of its pixel corners, one loop per border +// (inside kept clockwise), with runs along one direction joined. Used to draw the selection's ants as one line. +export function traceOutline(alpha, w, h) { + const on = (x, y) => x >= 0 && y >= 0 && x < w && y < h && alpha[(y * w + x) * 4 + 3] >= 128; + const stride = w + 1, next = new Map(); + const add = (x1, y1, x2, y2) => { const k = y1 * stride + x1; (next.get(k) || next.set(k, []).get(k)).push(y2 * stride + x2); }; + for (let y = 0; y < h; y++) for (let x = 0; x < w; x++) { + if (!on(x, y)) continue; + if (!on(x, y - 1)) add(x, y, x + 1, y); + if (!on(x + 1, y)) add(x + 1, y, x + 1, y + 1); + if (!on(x, y + 1)) add(x + 1, y + 1, x, y + 1); + if (!on(x - 1, y)) add(x, y + 1, x, y); + } + const loops = []; + for (const [start, list] of next) { + while (list.length) { + const loop = [start]; + let at = list.pop(); + while (at !== start) { + loop.push(at); + const out = next.get(at); + if (!out?.length) break; + at = out.pop(); + } + const pts = loop.map(v => [v % stride, Math.floor(v / stride)]), kept = []; + for (let i = 0; i < pts.length; i++) { + const a = pts[(i + pts.length - 1) % pts.length], b = pts[i], c = pts[(i + 1) % pts.length]; + if ((b[0] - a[0]) * (c[1] - b[1]) !== (b[1] - a[1]) * (c[0] - b[0])) kept.push(b); + } + if (kept.length > 1) loops.push(kept); + } + } + return loops; +} + +// Douglas-Peucker on a closed loop: staircase borders become straight runs, within eps pixels of the 0.5 border. +export function simplifyLoop(loop, eps) { + if (loop.length < 4) return loop; + let far = 0, best = -1; + for (let i = 1; i < loop.length; i++) { const d = Math.hypot(loop[i][0] - loop[0][0], loop[i][1] - loop[0][1]); if (d > best) { best = d; far = i; } } + const keep = new Array(loop.length).fill(false); keep[0] = keep[far] = true; + const run = (a, b) => { + const pt = i => loop[i % loop.length], [ax, ay] = pt(a), [bx, by] = pt(b), len = Math.hypot(bx - ax, by - ay) || 1; + let at = -1, worst = eps; + for (let i = a + 1; i < b; i++) { + const [x, y] = pt(i), d = Math.abs((bx - ax) * (ay - y) - (ax - x) * (by - ay)) / len; + if (d > worst) { worst = d; at = i; } + } + if (at >= 0) { keep[at % loop.length] = true; run(a, at); run(at, b); } + }; + run(0, far); run(far, loop.length); + return loop.filter((_, i) => keep[i]); +} diff --git a/web/image_generation.js b/web/image_generation.js index 96bb73aa..811f5344 100644 --- a/web/image_generation.js +++ b/web/image_generation.js @@ -8,7 +8,7 @@ import {about} from './request_assets.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', downloading: 'downloading', checking: 'checking', - token: 'token', job: 'job', prompt: 'prompt', width: 'width', height: 'height', images: 'images', seed: 'seed', + token: 'token', job: 'job', prompt: 'prompt', width: 'width', height: 'height', images: 'images', mask: 'mask', seed: 'seed', stage: 'stage', fraction: 'fraction', file: 'file', error: 'error', name: 'name', subfolder: 'subfolder', states: {queued: 'waiting', starting: 'queued', progress: 'running', complete: 'done', failed: 'error', cancelling: 'running', cancelled: 'error'}, @@ -69,9 +69,9 @@ export function createImageGenClient(fetchImpl = globalThis.fetch.bind(globalThi } return { info, - async generate({prompt, width, height, images = [], seed}) { + async generate({prompt, width, height, images = [], mask, seed}) { const value = await post('generate', {[p.prompt]: prompt, [p.width]: width, [p.height]: height, - [p.images]: images, ...(seed === undefined ? {} : {[p.seed]: seed})}); + [p.images]: images, ...(mask ? {[p.mask]: mask} : {}), ...(seed === undefined ? {} : {[p.seed]: seed})}); return {...value, job: value[p.job]}; }, async status(job) { @@ -87,6 +87,10 @@ export function createImageGenClient(fetchImpl = globalThis.fetch.bind(globalThi // The image models are fetched by this page's own engine; the reply is the new info. async packInstall() { const value = await post('pack-install', {}); return {...value, installed: value.installed ?? value[p.installed] ?? false}; }, async packCancel() { const value = await post('pack-cancel', {}); return {...value, installed: value.installed ?? value[p.installed] ?? false}; }, + // Click select: the engine finds the object under a point; the pack it needs is fetched on demand like the image models. + async segment({image, points}) { return post('segment', {image, points}); }, + segmentInstall: () => post('segment-install', {}), + segmentCancel: () => post('segment-cancel', {}), // The prompts and sizes of a new character's five pictures (the Agent's cards use the same ones). castViews: async description => (await post('cast-views', {description})).views, async upload(file) { @@ -264,6 +268,46 @@ export async function uploadGenerationReference(renderer, client, context, alive return client.upload(file); } +// The selection as a black and white picture the size of the layer being edited, uploaded beside it. +// Nothing is sent when nothing is selected, so the whole picture is edited as before. +export async function uploadGenerationMask(renderer, client, context, alive = () => true) { + if (!context.layer || !context.selection) return null; + const shape = renderer.selectionMask(context.doc, context.selection, {layer: context.layer, scale: 1}); + const mask = globalThis.document.createElement('canvas'); + mask.width = shape.width; mask.height = shape.height; + const draw = mask.getContext('2d', {willReadFrequently: true}); + draw.fillStyle = '#000'; draw.fillRect(0, 0, mask.width, mask.height); + draw.drawImage(shape, 0, 0); + const pixels = draw.getImageData(0, 0, mask.width, mask.height).data; + let any = false; + for (let i = 0; i < pixels.length; i += 4) if (pixels[i] >= 128) { any = true; break; } + if (!any) return null; + const blob = await new Promise((resolve, reject) => mask.toBlob(b => b ? resolve(b) : reject(failure()), 'image/png')); + if (!alive()) throw failure('cancelled'); + return client.upload(new File([blob], `freevideo-image-mask-${crypto.randomUUID()}.png`, {type: 'image/png'})); +} + +// The object under a click, as the name of a black and white picture the size of the layer's pixels. +// The layer is uploaded once for as long as it looks the same. +const segmentUploads = new WeakMap(); +export async function segmentAt(renderer, client, doc, layer, point) { + const look = JSON.stringify([layer.ops.map(o => o.id), layer.adjust, layer.src, layer.mask?.ops.map(o => o.id)]); + let known = segmentUploads.get(layer); + if (known?.look !== look) { + await renderer.preload(doc); + const raster = renderer.layerImage(doc, layer, 1, true); + const blob = await new Promise((resolve, reject) => raster.toBlob(b => b ? resolve(b) : reject(failure()), 'image/png')); + known = {look, name: await client.upload(new File([blob], `freevideo-image-reference-${crypto.randomUUID()}.png`, {type: 'image/png'})), width: raster.width, height: raster.height}; + segmentUploads.set(layer, known); + } + const x = Math.min(known.width - 1, Math.max(0, Math.floor(point[0] * known.width))), y = Math.min(known.height - 1, Math.max(0, Math.floor(point[1] * known.height))); + const value = await client.segment({image: known.name, points: [[x, y, 1]]}); + if (typeof value?.mask !== 'string') throw failure(); + const image = await renderer.image(value.mask); + renderer.images.get(value.mask).value = image; + return value.mask; +} + // "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; @@ -308,7 +352,7 @@ export function insertGenerationResult(doc, result, context, t) { 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)); + if (context.selection && !context.region) next = commands.maskFromSelection(next, next.active, scaleSelection(context.selection, sx, sy)); return next; } @@ -399,10 +443,10 @@ export function createImageGenSession({client, prepare, apply, onChange = () => emit({running: true, job: null, fraction: 0, stage: '', error: '', report_id: null, waiting: '', fitted: null}); const alive = () => current === epoch && !disposed; try { - const {context, images = []} = await prepare(request, alive); + const {context, images = [], mask} = await prepare(request, alive); if (!alive()) return; const withImages = images.length > 0; - const row = await client.generate({prompt: request.prompt, width: context.width, height: context.height, images}); + const row = await client.generate({prompt: request.prompt, width: context.width, height: context.height, images, mask}); const {job} = row; if (!job) throw failure(row.error, row); if (row.fitted) emit({fitted: row.fitted}); @@ -492,7 +536,9 @@ export function createImageGeneration({client, renderer, t, getDoc, getSelection 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}; + // Only the selected part of the edited picture (the first) is changed. + const mask = images.length ? await uploadGenerationMask(renderer, client, context, alive) : null; + return {context: mask ? {...context, region: true} : context, images, mask}; }, apply: async (result, context, alive) => { const image = await renderer.image(result.file); @@ -587,7 +633,7 @@ export function createImageGeneration({client, renderer, t, getDoc, getSelection 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.placeholder = mode === 'edit' ? getSelection() ? t('What to change in the selection', '要改选区里的什么') : t('What to change', '要改什么') : t('Describe the picture', '描述画面'); input.setAttribute('aria-label', input.placeholder); const pack = missing ? packView(state.info, t) : null; action.textContent = pack ? pack.button : state.running ? t('Stop', '停止') : t('Generate', '生成'); @@ -674,7 +720,7 @@ export function createImageGeneration({client, renderer, t, getDoc, getSelection }; grow(); update(); } - return {mount, unmount() { clearAttention(); promptRequired = false; mounted = null; }, + return {mount, drawsSelection: () => mode === 'edit' && getDoc().layers.length > 0, unmount() { clearAttention(); promptRequired = false; mounted = null; }, selected: () => session.check(), cancel: session.cancel, // The blank canvas a reloaded generation was started on, to be set again by the editor. resumedCanvas: resumed?.context?.mode === 'new' diff --git a/web/launcher.js b/web/launcher.js index a0dcfb7f..5fa58417 100644 --- a/web/launcher.js +++ b/web/launcher.js @@ -43,6 +43,21 @@ function currentStudioWorkflow(saved, bundled) { if (finite(height, 256, 4096) && Number.isInteger(height)) node.widgets_values[2] = height; if (finite(seconds, 1.625, 60)) node.widgets_values[3] = seconds; if (finite(seed, 0, 2 ** 53 - 1) && Number.isInteger(seed)) node.widgets_values[4] = seed; + // A tab saved by 0.4.13 or later keeps these at fixed positions (6, 7, 8, 17, 18); an unstamped one moved them. + const [major, minor, patch] = String(old.properties?.freevideo_version).split('.').map(Number); + const stamped = [major, minor, patch].every(Number.isInteger) && (major > 0 || minor > 4 || (minor === 4 && patch >= 13)); + if (stamped && Array.isArray(old.widgets_values)) { + const values = node.widgets_values, kept = old.widgets_values; + // The bundled workflow lists only the first widgets: fill up to the last one with the node's defaults. + const defaults = [8, 3, false, false, '', '', 0, false, false, false, 'none', '2']; + while (values.length < 19) values.push(defaults[values.length - 7]); + if (typeof kept[6] === 'boolean') values[6] = kept[6]; + // The quality level is these steps: High and Max are 20 first-pass steps. + if (Number.isInteger(kept[7]) && finite(kept[7], 1, 32)) values[7] = kept[7]; + if (Number.isInteger(kept[8]) && finite(kept[8], 1, 31)) values[8] = kept[8]; + if (['none', 'easycache-av', 'fbcache'].includes(kept[17])) values[17] = kept[17]; + if (['2', '1.5', '1.2'].includes(kept[18])) values[18] = kept[18]; + } const media = graph.nodes.find(row => row.type === 'FreeVideoMedia'); const assets = (saved.nodes.find(row => row.type === 'FreeVideoMedia')?.widgets_values || [])[0]; try { diff --git a/web/sampling_effort.js b/web/sampling_effort.js index 14af2343..4cb0fabf 100644 --- a/web/sampling_effort.js +++ b/web/sampling_effort.js @@ -1,19 +1,24 @@ // One editor control. The queued request remains an immutable graph snapshot. // Each level is a complete plan (freevideo_engine/tiers.py): Light and Medium are two-pass 8 + 3, // told apart by the learned upscale factor (Medium starts from a larger first pass); High and Max -// sample 20 steps once, High with a step cache. `upscale` is 2 for every single-pass level. +// sample 16 and 20 steps once, both with the guarded step cache. `upscale` is 2 for every single-pass level. const frozen = rows => Object.freeze(rows.map(row => Object.freeze({...row}))); export const SAMPLING_EFFORTS = frozen([ {name: 'Light', zh: '轻量', steps: 8, twoPass: true, stepCache: null, upscale: 2, color: '#3dd6b5'}, {name: 'Medium', zh: '标准', steps: 8, twoPass: true, stepCache: null, upscale: 1.5, color: '#4ea8ff'}, - {name: 'High', zh: '精细', steps: 20, twoPass: false, stepCache: 'easycache-av', upscale: 2, color: '#9d7bff'}, - {name: 'Max', zh: '极致', steps: 20, twoPass: false, stepCache: null, upscale: 2, color: '#ffa94d'}, + {name: 'High', zh: '精细', steps: 16, twoPass: false, stepCache: 'easycache-av', upscale: 2, color: '#9d7bff'}, + {name: 'Max', zh: '极致', steps: 20, twoPass: false, stepCache: 'easycache-av', upscale: 2, color: '#ffa94d'}, ]); -// Mac runs the same levels. -export const SAMPLING_EFFORTS_MAC = SAMPLING_EFFORTS; +// Mac runs the same plans; retain platform identity for its legacy High alias. +const macLevels = new WeakSet(); +export const SAMPLING_EFFORTS_MAC = frozen(SAMPLING_EFFORTS); +macLevels.add(SAMPLING_EFFORTS_MAC); +// Earlier Max sampled 20 plain steps, on every platform. +export const legacyMaxEffort = (steps, twoPass, stepCache = null, levels = SAMPLING_EFFORTS) => + steps === 20 && !twoPass && stepCacheOf(stepCache) === null ? levels.find(tier => tier.name === 'Max') : undefined; // Mac results made before the cache came: High was 16 steps without one. -export const legacyMacEffort = (steps, twoPass, stepCache = null) => - steps === 16 && !twoPass && stepCacheOf(stepCache) === null ? SAMPLING_EFFORTS_MAC.find(tier => tier.name === 'High') : undefined; +export const legacyMacEffort = (steps, twoPass, stepCache = null, levels = SAMPLING_EFFORTS_MAC) => + steps === 16 && !twoPass && stepCacheOf(stepCache) === null ? levels.find(tier => tier.name === 'High') : undefined; export const effortName = (t, tier) => t(tier.name, tier.zh); export const STEP_CACHES = Object.freeze([null, 'easycache-av', 'fbcache']); @@ -34,13 +39,15 @@ export function effortsFrom(grid) { const base = grid?.machine?.system === 'Darwin' ? SAMPLING_EFFORTS_MAC : SAMPLING_EFFORTS; const tiers = grid?.tiers; if (!tiers || typeof tiers !== 'object') return base; - return frozen(base.map(tier => { + const levels = frozen(base.map(tier => { const row = tiers[tier.name.toLowerCase()]; if (!row || !Number.isInteger(row.steps) || typeof row.two_pass !== 'boolean') return tier; return {...tier, steps: row.steps, twoPass: row.two_pass, stepCache: 'step_cache' in row ? stepCacheOf(row.step_cache) : tier.stepCache, upscale: !row.two_pass ? 2 : 'upscale' in row ? upscaleOf(row.upscale) : tier.upscale}; })); + if (base === SAMPLING_EFFORTS_MAC) macLevels.add(levels); + return levels; } // The level a plan matches, or undefined for custom steps; the upscale factor tells the @@ -48,7 +55,9 @@ export function effortsFrom(grid) { export function effortFor(steps, twoPass, refine, stepCache = null, levels = SAMPLING_EFFORTS, upscale = 2) { const cache = stepCacheOf(stepCache), factor = twoPass ? upscaleOf(upscale) : 2; return levels.find(tier => tier.steps === steps && tier.twoPass === twoPass && (!twoPass || refine === 3) - && tier.stepCache === cache && tier.upscale === factor); + && tier.stepCache === cache && tier.upscale === factor) + || legacyMaxEffort(steps, twoPass, cache, levels) + || (macLevels.has(levels) ? legacyMacEffort(steps, twoPass, cache, levels) : undefined); } // The level a result's sampling_plan names. Results made with the earlier 8 + 2 default count as @@ -60,11 +69,15 @@ export function resultEffort(plan) { } // The steps and factor switching two-pass leaves. Off from an 8-step first pass samples the -// single-pass levels' 20 steps; on from a single-pass level gives Light. Typed steps stay. +// Max plan including its cache; on from a single-pass level gives Light without a cache. Typed steps stay. export function passSwitch(on, steps, refine, stepCache = null, levels = SAMPLING_EFFORTS) { - if (!on) return {steps: steps === 8 ? levels.find(tier => !tier.twoPass)?.steps ?? 20 : steps, refine, upscale: 2}; + const max = levels.find(tier => tier.name === 'Max'); + if (!on) return {steps: steps === 8 ? max?.steps ?? 20 : steps, refine, upscale: 2, + stepCache: steps === 8 ? max?.stepCache ?? null : null}; const light = levels.find(tier => tier.twoPass && tier.upscale === 2); - return effortFor(steps, false, refine, stepCache, levels) && light ? {steps: light.steps, refine: 3, upscale: 2} : {steps, refine, upscale: 2}; + return effortFor(steps, false, refine, stepCache, levels) && light + ? {steps: light.steps, refine: 3, upscale: 2, stepCache: null} + : {steps, refine, upscale: 2, stepCache: null}; } // A continuation samples 8 + 3 in two passes: it keeps the factor of a two-pass level, else 2 (widget text). @@ -72,8 +85,8 @@ export const continuationUpscale = (steps, twoPass, refine, stepCache, upscale, String(effortFor(Number(steps), twoPass !== false, Number(refine), stepCache, levels, upscale)?.twoPass ? upscaleOf(upscale) : 2); // The level whose sampling table a plan needs, for naming a missing download (tiers.table_level): -// the exact level (the cache tells High from Max at 20 steps, the factor Light from Medium; 16 steps -// is the earlier Mac High), else the first level with those steps (Light for 8 + 3, whatever the factor). +// the exact level or legacy alias, else the first with those steps (16 is High, 20 is Max, +// Light for 8 + 3 whatever the factor). export function tableLevel(steps, twoPass, refine, stepCache = null, upscale = 2) { return effortFor(steps, twoPass, refine, stepCache, SAMPLING_EFFORTS, upscale) || legacyMacEffort(steps, twoPass, stepCache) || SAMPLING_EFFORTS.find(tier => tier.steps === steps && tier.twoPass === twoPass && (!twoPass || refine === 3)); @@ -182,7 +195,8 @@ export function createSamplingEffort(t, {onChange, onPreview = () => {}, levels // The machine's own table (effortsFrom); the names stay, the plans follow the machine. function setLevels(next) { const same = (a, b) => a.steps === b.steps && a.twoPass === b.twoPass && a.stepCache === b.stepCache && a.upscale === b.upscale; - if (!next || next.length !== efforts.length || next.every((tier, i) => same(tier, efforts[i]))) return; + if (!next || next.length !== efforts.length || (macLevels.has(next) === macLevels.has(efforts) + && next.every((tier, i) => same(tier, efforts[i])))) return; efforts = next; update({baseSteps: currentSteps, refineSteps: currentRefine, twoPass, stepCache, upscale, disabled, reason}); } diff --git a/web/studio.css b/web/studio.css index ecbafad0..360665d3 100644 --- a/web/studio.css +++ b/web/studio.css @@ -222,3 +222,6 @@ .fv-lora-problem{grid-column:var(--lora-lead,2)/span 3;grid-row:4;display:flex;flex-direction:column;align-items:flex-start;gap:6px;font-size:var(--fv-micro)}.fv-lora-problem p{margin:0;color:var(--fv-danger)}.fv-lora-problem.fv-lora-partial{color:var(--fv-muted)}.fv-lora-problem[hidden]{display:none}.fv-lora-why summary{cursor:pointer;color:var(--fv-muted)}.fv-lora-why pre{margin:6px 0 0;white-space:pre-wrap;word-break:break-all;color:var(--fv-muted);font:inherit;font-family:ui-monospace,monospace}.fv-lora-row[data-unusable=true] select{color:var(--fv-muted)} .fv-lora-row[data-unusable=true]{cursor:default}.fv-lora-row[data-unusable=true] :is(input[type=number],input[type=checkbox]){opacity:.45}.fv-lora-row[data-unusable=true] select{opacity:1}.fv-lora-row[data-unusable=true] :is(.fv-lora-trigger,.fv-lora-range){display:none} +/* The queued count follows the part number; when the line is too narrow it wraps out of sight. */ +.fv-studio .fv-chain-position:has(.fv-chain-queued){display:flex;flex-wrap:wrap;column-gap:0;line-height:1.4;max-height:1.4em;overflow:hidden} +.fv-studio .fv-chain-position .fv-chain-queued{white-space:pre} diff --git a/web/studio.js b/web/studio.js index f419e9dd..3e732d49 100644 --- a/web/studio.js +++ b/web/studio.js @@ -818,9 +818,9 @@ export function openStudio(node, {extendFrom = null} = {}) { if (extension.source) return; const wantsTwoPass = mode !== 'off', wantsPreview = mode === 'preview'; if (twoPass.checked !== wantsTwoPass) { - // A level in one pass samples 20 steps; two-pass from one is Light. Typed steps stay. + // Switching off uses Max including its cache; switching on uses uncached Light. Typed steps stay. const next = passSwitch(wantsTwoPass, Number(baseSteps.value), Number(refineSteps.value), stepCacheNow(), effort.levels()); - twoPass.checked = wantsTwoPass; autoSinglePass = false; set(node, 'two_pass', wantsTwoPass); clearStepCache(); setUpscale(next.upscale); + twoPass.checked = wantsTwoPass; autoSinglePass = false; set(node, 'two_pass', wantsTwoPass); set(node, 'step_cache', next.stepCache || 'none'); setUpscale(next.upscale); if (next.steps !== Number(baseSteps.value) && !linked(node, 'base_steps') && widget(node, 'base_steps')) { baseSteps.value = String(next.steps); set(node, 'base_steps', next.steps); if (next.refine !== Number(refineSteps.value) && set(node, 'refine_steps', next.refine)) refineSteps.value = String(next.refine); @@ -1099,6 +1099,11 @@ export function openStudio(node, {extendFrom = null} = {}) { stage.append(stageMedia); previewSpace.append(stage); const progress = createGenerationProgress(t, undefined, {compact: true, api}); stage.append(progress.element); const chainLine = el('div', '', 'fv-muted fv-chain-position'); chainLine.hidden = true; progress.element.prepend(chainLine); + // The queued count sits after the part number; a line too narrow for both keeps the part number only. + const setChainLine = text => { + const main = el('span', text), more = chainController?.waiting; + chainLine.replaceChildren(main, ...(more ? [el('span', t(` · ${more} more queued`, ` · 另有 ${more} 条排队`), 'fv-chain-queued')] : [])); + }; progress.updateReport({report_id: node.freevideoReportId}); cleanup.push(() => progress.dispose()); cleanup.push(createPreviewScene(stage, progress.element, stageMedia)); @@ -1302,13 +1307,13 @@ export function openStudio(node, {extendFrom = null} = {}) { }, })); const queueController = studioQueues.get(node); - let chainWasActive = false; + let chainWasActive = false, stopNote = ''; chainController = createStudioChain({api, node, onProgress: (message, chain) => { const position = chainPosition(chain, message); const joining = readChainResult(chain).joining, waiting = chain.state === 'queued' && message.phase === 'prepare'; - chainLine.textContent = joining ? t('Joining the parts.', '正在接合各段。') + setChainLine(joining ? t('Joining the parts.', '正在接合各段。') : waiting ? t(`Waiting in queue · Part ${position.index + 1} of ${position.count}`, `正在排队 · 第 ${position.index + 1} 段,共 ${position.count} 段`) - : t(`Part ${position.index + 1} of ${position.count}`, `第 ${position.index + 1} 段,共 ${position.count} 段`); + : t(`Part ${position.index + 1} of ${position.count}`, `第 ${position.index + 1} 段,共 ${position.count} 段`)); progress.element.querySelector('.fv-generation-label').hidden = joining || waiting; progress.updateReport(message); showProgress(message); }, @@ -1324,12 +1329,14 @@ export function openStudio(node, {extendFrom = null} = {}) { partList.update(chain, active); busy = active || Boolean(queueState?.running || queueState?.pending.length); queued = chain?.state === 'queued'; chainLine.hidden = !active; - chainLine.textContent = outcome?.joining ? t('Joining the parts.', '正在接合各段。') - : t(`Part ${position.index + 1} of ${position.count}`, `第 ${position.index + 1} 段,共 ${position.count} 段`); + setChainLine(outcome?.joining ? t('Joining the parts.', '正在接合各段。') + : t(`Part ${position.index + 1} of ${position.count}`, `第 ${position.index + 1} 段,共 ${position.count} 段`)); if (active && node.properties.freevideo_chain?.width) previewSize(node.properties.freevideo_chain.width, node.properties.freevideo_chain.height); if (!active) { cancelling = false; hideProgress(); } if (!active) progress.element.querySelector('.fv-generation-label').hidden = false; if (chain?.state === 'cancelled') status.textContent = t('Cancelled', '已取消'); + // After a stop with videos queued the next one starts at once; its line says so until something else is done. + if (stopNote) { status.dataset.error = 'false'; status.textContent = stopNote; } resumeChain.hidden = !['cancelled', 'failed', 'interrupted'].includes(chain?.state); queueActions.hidden = !active && !busy && resumeChain.hidden; updateRunButton(); @@ -1439,15 +1446,18 @@ export function openStudio(node, {extendFrom = null} = {}) { runMode = multiple || extension.source ? 'single' : ['single', 'batch', 'loop'].includes(node.properties?.freevideo_run_mode) ? node.properties.freevideo_run_mode : runMode; const chain = chainController?.state, position = chainPosition(chain); - const chainReason = !chainController?.active ? '' : readChainResult(chain)?.joining + // Only another long video can join a running chain; a single-part request would be submitted + // straight to the busy GPU, so its button waits like a full queue does. + const chainWait = Boolean(chainController?.active && (!chainController.canQueue || !multiple)); + const chainReason = !chainWait ? '' : readChainResult(chain)?.joining ? t('Joining parts…', '正在拼接…') : chain?.state === 'queued' ? t('Queued…', '排队中…') : t(`Generating part ${position.index + 1}…`, `正在生成第 ${position.index + 1} 段…`); - generate.disabled = submitting || chainController?.active || lengthControl.blocked || (runMode === 'loop' && queueState?.looping); + generate.disabled = submitting || chainWait || lengthControl.blocked || (runMode === 'loop' && queueState?.looping); generate.textContent = submitting ? t('Adding…', '正在添加…') : chainReason || (runMode === 'batch' ? t(`Add ${count.value || '…'} videos`, `添加 ${count.value || '…'} 个视频`) : runMode === 'loop' ? (queueState?.looping ? t('Looping', '循环中') : t('Start loop', '开始循环')) - : busy ? t('Add to queue', '添加到队列') + : busy ? t('Add to queue', '加入队列') : previewFirst.checked && !previewFirst.disabled ? t('Generate preview', '生成预览') : t('Generate video', '生成视频')); generate.title = (lengthControl.blocked ? lengthControl.message : submitting ? t('Adding…', '正在添加…') : ''); for (const b of modeButtons) b.setAttribute('aria-pressed', String(b.dataset.mode === runMode)); @@ -1543,7 +1553,12 @@ export function openStudio(node, {extendFrom = null} = {}) { if (cancelling) return; try { cancelling = true; updateRunButton(); - if (chainController.active) { await chainController.cancel(); cancelling = false; updateRunButton(); return; } + if (chainController.active) { + stopNote = chainController.waiting > 0 ? t('Stopped. The next queued video is starting.', '已停止,队列中的下一条开始生成。') : ''; + await chainController.cancel(); cancelling = false; updateRunButton(); + if (stopNote) { status.dataset.error = 'false'; status.textContent = stopNote; } + return; + } await syncQueue(); const id = queueController.state().running?.id; if (!id) { cancelling = false; updateRunButton(); return; } @@ -1579,7 +1594,7 @@ export function openStudio(node, {extendFrom = null} = {}) { let held = null, accepted = false; const onAccepted = () => { accepted = true; }; try { - lastQueueError = null; failure.clear(); node.freevideoFailure = ''; node.freevideoFailureReport = null; + stopNote = ''; lastQueueError = null; failure.clear(); node.freevideoFailure = ''; node.freevideoFailureReport = null; capturing = true; updateRunButton(); status.dataset.error = 'false'; runOptions.open = false; lengthControl.freeze(); if ((!extension.source || prompt.value.trim()) && await enhancer.beforeGenerate() === false) return; @@ -1604,8 +1619,10 @@ export function openStudio(node, {extendFrom = null} = {}) { capturing = false; for (const section of controls.querySelectorAll('.fv-section')) section.inert = false; if (lengthPlan?.count > 1) { + const waitingBefore = chainController.waiting; const chain = await chainController.start(snapshot, lengthOptions, {onAccepted}); if (disposed) return; + if (chain && chainController.waiting > waitingBefore) { status.dataset.error = 'false'; status.textContent = addedNote; return; } if (!chain) { if (source) throw new Error(extendText(t, 'unavailable')); lengthControl.fallback(); if (!seconds.reportValidity()) return; await queueController.start(snapshot, {count: batchCount, repeat: chosenMode === 'loop', seedIndex, onAccepted}); } } else await queueController.start(snapshot, {count: batchCount, repeat: chosenMode === 'loop', seedIndex, onAccepted}); @@ -1685,6 +1702,10 @@ export function openStudio(node, {extendFrom = null} = {}) { b.append(el('span', null, 'fv-shape-icon'), el('span', label)); shapes.append(b); } mp.onchange = () => { try { pixels = Number(mp.value); applySize(); if (pixels !== 1) legacyOption?.remove(); } catch (error) { status.dataset.error = 'true'; status.textContent = error.message; } }; + // The note describes the request just sent; once the draft changes it could read as the new one. + const addedNote = t('Added to the queue. It starts when the current video finishes.', '已加入队列,当前视频完成后自动开始。'); + const clearAddedNote = () => { if (status.textContent === addedNote) status.textContent = ''; }; + seconds.addEventListener('input', clearAddedNote); prompt.addEventListener('input', clearAddedNote); seconds.onchange = () => { if (lengthControl.commit({report: false})) updateCanvas(); }; const mediaChanged = async e => { if (String(e.detail) !== String(mediaNode?.id)) return; refreshEffortEstimate(true); if (selected !== 'input') return; try { await inputRatio(); } catch (error) { status.textContent = error.message; } }; const lorasChanged = e => { if (String(e.detail) === String(loraNode?.id)) refreshEffortEstimate(true); }; diff --git a/web/studio_chain.js b/web/studio_chain.js index 87cbeaac..88b74c31 100644 --- a/web/studio_chain.js +++ b/web/studio_chain.js @@ -3,6 +3,7 @@ import {buildChainRequest} from './length_plan.js?v=20261011-length'; export const CHAIN_PROPERTY = 'freevideo_chain'; const activeStates = new Set(['queued', 'running', 'cancelling']); +const MAX_WAITING = 5; const finite = value => typeof value === 'number' && Number.isFinite(value) && value >= 0; const post = body => ({method: 'POST', headers: {'Content-Type': 'application/json'}, body: JSON.stringify(body)}); @@ -55,14 +56,17 @@ export function singleRunSnapshot(snapshot, node, plan) { export function createStudioChain({api, node, onState = () => {}, onProgress = () => {}, onResult = () => {}, onError = () => {}, pollMs = 2000, now = () => Date.now()}) { - let chain = null, poller = null, disposed = false, epoch = 0, submitting = false; + let chain = null, poller = null, disposed = false, epoch = 0, submitting = false, promoting = false; + // Chains the server took while another was running; the next one is shown when the running one stops. + const waiting = []; let previous = {}, lastMessage = {}, delivered = null, errorShown = null, corrected = false, etaAt = 0, lastIndex = -1; const timers = new Set(); const clock = {now, setTimeout(fn, ms) { const id = setTimeout(() => { timers.delete(id); fn(); }, ms); timers.add(id); return id; }, clearTimeout(id) { clearTimeout(id); timers.delete(id); }}; function clearPolls() { for (const id of timers) clearTimeout(id); timers.clear(); } function refresh() { clearPolls(); return poller?.poll(); } - const active = () => submitting || activeStates.has(chain?.state) || chain?.output_state === 'joining'; + const running = () => activeStates.has(chain?.state) || chain?.output_state === 'joining'; + const active = () => submitting || running(); async function json(path, options, allowMissing = false) { const response = await api.fetchApi(path, options); if (response.status === 404 && allowMissing) return null; @@ -102,6 +106,7 @@ export function createStudioChain({api, node, onState = () => {}, onProgress = ( chain = value; onState(chain); const result = readChainResult(chain); if (active()) progress({}, fresh); + else if (waiting.length) void promote(); if (result.error && errorShown !== result.error) { errorShown = result.error; onError(Object.assign(new Error(result.error), {joining: chain.output_state === 'failed', ...chainPosition(chain)})); } @@ -116,6 +121,24 @@ export function createStudioChain({api, node, onState = () => {}, onProgress = ( } return active(); } + async function promote() { + if (promoting || disposed) return; + promoting = true; + let again = false; + try { + while (waiting.length && !disposed && !running()) { + const next = waiting[0]; + let value; + try { value = (await json(chainRoute(next.id), {cache: 'no-store'}, true))?.chain; } + catch { again = true; break; } + waiting.shift(); + if (!value?.id || value.state === 'cancelled') continue; + adopt(value, next.metadata); break; + } + } finally { promoting = false; } + if (again && !disposed) clock.setTimeout(() => void promote(), 3000); + else if (!disposed && chain) onState(chain); + } function watch(fresh = false) { poller?.dispose(); clearPolls(); const id = chain.id, version = epoch; @@ -135,9 +158,10 @@ export function createStudioChain({api, node, onState = () => {}, onProgress = ( const known = node.properties?.[CHAIN_PROPERTY]?.id; const reply = await json('/freevideo/chains', undefined, true); if (disposed || version !== epoch || chain) return; - const found = reply?.chains?.find(row => String(row.node) === String(node.id) && (known ? row.id === known - : activeStates.has(row.state) && row.client_id && row.client_id === api.clientId)); + const mine = row => String(row.node) === String(node.id) && activeStates.has(row.state) && row.client_id && row.client_id === api.clientId; + const found = reply?.chains?.find(row => String(row.node) === String(node.id) && (known ? row.id === known : mine(row))); if (!found) return; + for (const row of reply.chains) if (row.id !== found.id && mine(row)) waiting.push({id: row.id, metadata: null}); const value = (await json(chainRoute(found.id), undefined, true))?.chain; if (disposed || version !== epoch || chain || !value) return; adopt(value); @@ -177,10 +201,12 @@ export function createStudioChain({api, node, onState = () => {}, onProgress = ( } recover(); return { - get active() { return active(); }, get state() { return chain; }, + get active() { return active(); }, get state() { return chain; }, get waiting() { return waiting.length; }, + // A chain that is running does not stop the next one from being added behind it. + get canQueue() { return !disposed && !submitting && running() && chain?.state !== 'cancelling' && waiting.length < MAX_WAITING; }, async start(snapshot, options, {onAccepted = null} = {}) { - if (disposed || active()) throw new Error('chain_busy'); - submitting = true; epoch++; + if (disposed || submitting || (running() && waiting.length >= MAX_WAITING)) throw new Error('chain_busy'); + submitting = true; try { const body = buildChainRequest(options, {node: node.id, prompt: snapshot.output, extra_data: {extra_pnginfo: {workflow: snapshot.workflow || {}}}, ...(api.clientId ? {client_id: api.clientId} : {})}); @@ -189,7 +215,10 @@ export function createStudioChain({api, node, onState = () => {}, onProgress = ( if (!value.id) throw new Error('chain_missing'); onAccepted?.(); const inputs = snapshot.output[String(node.id)].inputs; - submitting = false; adopt(value, {width: inputs.width, height: inputs.height, seconds: options.seconds}); + const metadata = {width: inputs.width, height: inputs.height, seconds: options.seconds}; + submitting = false; + if (running() && chain.id !== value.id) { waiting.push({id: value.id, metadata}); onState(chain); return value; } + epoch++; adopt(value, metadata); return value; } finally { submitting = false; } }, diff --git a/web/translate_guard.js b/web/translate_guard.js new file mode 100644 index 00000000..7db62f0e --- /dev/null +++ b/web/translate_guard.js @@ -0,0 +1,17 @@ +// Browser auto-translate reads ComfyUI's lang="en" page and rewrites FreeVideo's own text: "MP" +// became 兆帕. FreeVideo writes both languages itself, so its surfaces opt out of translation and +// a Chinese interface says so on the document. +const OWN_CLASS = /(^|\s)fv-/; + +export function markOwn(node) { + if (node?.nodeType !== 1 || node.hasAttribute('translate')) return; + if (OWN_CLASS.test(node.getAttribute('class') || '') || node.id?.startsWith('fv-')) node.setAttribute('translate', 'no'); +} + +export function installTranslateGuard({chinese, doc = document} = {}) { + if (chinese) doc.documentElement.lang = 'zh-CN'; + for (const node of doc.body.children) markOwn(node); + new MutationObserver(records => { + for (const record of records) for (const node of record.addedNodes) markOwn(node); + }).observe(doc.body, {childList: true}); +}