diff --git a/docs/community-results.md b/docs/community-results.md index f80ae50a..5d0ad35f 100644 --- a/docs/community-results.md +++ b/docs/community-results.md @@ -123,11 +123,12 @@ Thanks to the community members who sent diagnostic reports. FreeVideo already r | 5070 Ti Laptop | 12 GB | 32 GB | 1344×768 | 10 s | Light | 0.4.23 | 373 s¹³ | 0.4.23 | Thermal‑limited¹ | | 5060 | 8 GB | 48 GB | 768×1344 | 14 s | Medium | 0.3.25 | 3474 s | 0.3.25 | Already optimal | | 3060 | 12 GB | 32 GB | 896×1120 | 10 s | Light | 0.3.25 | 1045 s | 0.4.0 | 1025 s (−20 s) | +| 3060 Laptop | 6 GB | 40 GB | 960×544 | 15 s | Light⁵ | 0.4.24 | 713 s | 0.4.24 | Thermal‑limited¹ | - Time: end to end, from submission until the video is saved. - After: the estimated time, derived from the report's stage times and from complete planned runs on our test machines (RTX 3090 and RTX 5060 Ti on Windows, RTX PRO 6000 on Linux) under the reporting machine's VRAM and RAM limits, not from reruns on that machine. - Already optimal: weight placement and transfers on this machine are already the fastest they can be; the time depends on the GPU's compute and the chosen quality level. -- ¹ Thermal-limited: the GPU spent most of the run in thermal slowdown (3070 Laptop: 229 MHz on average at 89 °C; A5000 Laptop: 882 MHz on average), so cooling limited the time. +- ¹ Thermal-limited: the GPU spent most of the run in thermal slowdown (3070 Laptop: 229 MHz on average at 89 °C; A5000 Laptop: 882 MHz on average; 3060 Laptop: 1496 MHz on average at 81 °C), so cooling limited the time. - ² This machine still runs the FP8 model; the estimate is for the int8 model, a one-click upgrade in the launcher since 0.3.0. - ³ Preview: only the first pass, at half resolution. With 16 GB of RAM most of the model is read from the drive at every step (about 80 MB/s on this machine), which takes most of the time. - ⁴ The report had no diagnostic file; the estimate replays the plan for this machine's configuration and measures the same request on our test machines. @@ -138,5 +139,5 @@ Thanks to the community members who sent diagnostic reports. FreeVideo already r - ⁹ The drive holding the model files reads at about 80 MB/s; reading them took about 570 s of the 1676 s. - ¹⁰ GPU shared: other programs were using about 29% of the graphics card during the run, so FreeVideo's steps ran at a little over 70% of their usual speed. With the card free, the same request takes about 7200 s. - ¹¹ With 16 GB of RAM, the text encoder and part of the video model cannot stay in memory and are read from the drive; text encoding and the first pass, which wait on these reads, took about 460 s of the 887 s. -- ¹² Reference-video task: the reference video and audio roughly double the sequence the model processes, so each step takes about twice as long as a plain request of the same size. +- ¹² Task with a reference: every step also processes the reference video, images or audio, so each step takes longer than a plain request of the same size, up to about twice as long with a long reference. - ¹³ Preview (first pass only). diff --git a/docs/zh-CN/community-results.md b/docs/zh-CN/community-results.md index 0d646c77..fbcbf252 100644 --- a/docs/zh-CN/community-results.md +++ b/docs/zh-CN/community-results.md @@ -80,12 +80,12 @@ | 5060 笔⁠记⁠本 | 8 GB | 32 GB | 1344×768 | 3 s | 标⁠准 | 0.4.7 | 541 s | 0.4.7 | 已⁠是⁠最⁠优 | | 5060 Ti | 16 GB | 34 GB | 1344×768 | 5.2 s | 标⁠准 | 0.4.8 | 2521 s | 0.4.8 | 已⁠是⁠最⁠优¹² | | 4060 | 8 GB | 32 GB | 768×1344 | 17 s | 轻⁠量 | 0.4.8 | 1942 s | 0.4.8 | 已⁠是⁠最⁠优 | -| M1 Max | 统一内存 | 64 GB | 960×544 | 2.3 s | 轻⁠量 | 0.4.8 | 1226 s | 0.4.8 | 已⁠是⁠最⁠优 | +| M1 Max | 统⁠一⁠内⁠存 | 64 GB | 960×544 | 2.3 s | 轻⁠量 | 0.4.8 | 1226 s | 0.4.8 | 已⁠是⁠最⁠优 | | 3070 | 8 GB | 32 GB | 1344×768 | 5.2 s | 轻⁠量⁵ | 0.4.8 | 285 s | 0.4.8 | 已⁠是⁠最⁠优 | | 3060 | 12 GB | 34 GB | 1664×960 | 6.6 s | 标⁠准 | 0.4.8 | 4025 s | 0.4.8 | 已⁠是⁠最⁠优¹² | -| 4070 Laptop | 8 GB | 34 GB | 1344×768 | 15 s | 轻⁠量⁵ | 0.4.8 | 1444 s | 0.4.8 | 已⁠是⁠最⁠优¹² | -| 5060 Ti | 16 GB | 56 GB | 960×544 | 10 s | 轻⁠量 | 0.4.6 | 904 s | 0.4.9 | 805 s (−99 s) | -| 3070 Laptop | 8 GB | 64 GB | 1344×768 | 10 s | 轻⁠量 | 0.4.9 | 1267 s | 0.4.9 | 已⁠是⁠最⁠优 | +| 4070 笔⁠记⁠本 | 8 GB | 34 GB | 1344×768 | 15 s | 轻⁠量⁵ | 0.4.8 | 1444 s | 0.4.8 | 已⁠是⁠最⁠优¹² | +| 5060 Ti | 16 GB | 56 GB | 960×544 | 10 s | 轻⁠量 | 0.4.6 | 904 s | 0.4.9 | 805 s⁠(−99 s⁠) | +| 3070 笔⁠记⁠本 | 8 GB | 64 GB | 1344×768 | 10 s | 轻⁠量 | 0.4.9 | 1267 s | 0.4.9 | 已⁠是⁠最⁠优 | | 3060 | 12 GB | 64 GB | 768×1344 | 5 s | 轻⁠量 | 0.4.12 | 639 s | 0.4.12 | 已⁠是⁠最⁠优 | | 5080 | 16 GB | 64 GB | 768×1344 | 3 s | 极⁠致 | 0.3.20 | 216 s | 0.3.20 | 已⁠是⁠最⁠优 | | 4060 Ti | 8 GB | 32 GB | 960×544 | 5 s | 轻⁠量 | 0.4.12 | 333 s | 0.4.12 | 已⁠是⁠最⁠优 | @@ -123,11 +123,12 @@ | 5070 Ti 笔⁠记⁠本 | 12 GB | 32 GB | 1344×768 | 10 s | 轻⁠量 | 0.4.23 | 373 s¹³ | 0.4.23 | 散⁠热⁠受⁠限¹ | | 5060 | 8 GB | 48 GB | 768×1344 | 14 s | 标⁠准 | 0.3.25 | 3474 s | 0.3.25 | 已⁠是⁠最⁠优 | | 3060 | 12 GB | 32 GB | 896×1120 | 10 s | 轻⁠量 | 0.3.25 | 1045 s | 0.4.0 | 1025 s⁠(−20 s⁠) | +| 3060 笔⁠记⁠本 | 6 GB | 40 GB | 960×544 | 15 s | 轻⁠量⁵ | 0.4.24 | 713 s | 0.4.24 | 散⁠热⁠受⁠限¹ | - 用时:从提交到视频保存完成的端到端用时。 - 优化后:预估用时,根据报告中的分段用时,结合测试机(Windows RTX 3090 和 RTX 5060 Ti、Linux RTX PRO 6000)在报告机器的显存和内存上限下运行完整规划的实测推算,不是在报告机器上重跑。 - 已是最优:这台机器的权重放置和数据传输已经是最快的方式,用时取决于显卡算力和所选档位。 -- ¹ 散热受限:显卡在生成过程中大部分时间处于温控降频(3070 笔记本平均 229 MHz、89 °C;A5000 笔记本平均 882 MHz),用时主要受散热限制。 +- ¹ 散热受限:显卡在生成过程中大部分时间处于温控降频(3070 笔记本平均 229 MHz、89 °C;A5000 笔记本平均 882 MHz;3060 笔记本平均 1496 MHz、81 °C),用时主要受散热限制。 - ² 这台机器仍在使用 FP8 模型;预估为在启动器中一键升级到 int8 模型(0.3.0 起提供)后的用时。 - ³ 预览:只运行第一遍(半分辨率)。16 GB 内存放不下大部分模型,每一步都要从硬盘读取(这台机器约 80 MB/s),大部分时间花在读盘上。 - ⁴ 报告未附诊断文件;预估按这台机器的配置重放规划,并在测试机上实测同样的请求。 @@ -138,5 +139,5 @@ - ⁹ 模型文件所在的硬盘读取约 80 MB/s,1676 秒里约 570 秒用于读取模型文件。 - ¹⁰ 显卡被占用:生成过程中其他程序占用了约 29% 的显卡,FreeVideo 每一步只有独占时七成多一点的速度;显卡空出来时,同样的请求约需 7200 秒。 - ¹¹ 16 GB 内存留不下文本编码器和视频模型的一部分,需要从硬盘读取;887 秒里约 460 秒是等待这些读取的文本编码和第一遍。 -- ¹² 参考视频任务:参考视频和参考音频使模型处理的序列约增加一倍,每步耗时约为相同尺寸普通请求的两倍。 +- ¹² 带参考的任务:每一步还要处理参考视频、参考图片或参考音频,所以每步耗时比相同尺寸的普通请求长;参考较长时约为两倍。 - ¹³ 预览(只做第一遍)。 diff --git a/freevideo_engine/access_proxy.py b/freevideo_engine/access_proxy.py index 1e645478..7d36a3f5 100644 --- a/freevideo_engine/access_proxy.py +++ b/freevideo_engine/access_proxy.py @@ -38,23 +38,34 @@ COOKIE = 'freevideo_access' DAYS = 30 LOGIN_PATH = '/m/login' -# Temporary English copy; keep all password-page text here for Iris. -PASSWORD_COPY = dict(title='FreeVideo', password='Password', remember='Keep me signed in', - submit='Sign in', incorrect='Incorrect password', retry='Try again later ({minutes} minutes).', - phone='FreeVideo mobile') +# All password-page text lives here; the lockout line matches web/phone_model.js. +PASSWORD_COPY = dict( + en=dict(title='FreeVideo', password='Password', remember='Keep me signed in', + submit='Sign in', incorrect='Incorrect password', retry='Too many attempts. Try again in {minutes} min.', + phone='FreeVideo mobile'), + zh=dict(title='FreeVideo', password='密码', remember='保持登录', + submit='登录', incorrect='密码不正确', retry='尝试次数过多,请 {minutes} 分钟后再试。', + phone='FreeVideo 手机版')) -def password_page(*, error='', retry_after=0, status=401): - copy = PASSWORD_COPY +def password_language(request): + """'zh' when the browser's most preferred language is Chinese, else 'en'.""" + first = request.headers.get('Accept-Language', '').split(',')[0].split(';')[0].strip().lower() + return 'zh' if first == 'zh' or first.startswith('zh-') else 'en' + + +def password_page(request, *, error='', retry_after=0, status=401): + language = password_language(request) + copy = PASSWORD_COPY[language] message = copy['retry'].format(minutes=math.ceil(retry_after / 60)) if retry_after else error response = web.Response(status=status, content_type='text/html', text=( - '{title}' '

{title}

{message}

' '' '' '
').format( - **{k: html.escape(v) for k, v in copy.items()}, message=html.escape(message), retry_after=retry_after, action=LOGIN_PATH)) + **{k: html.escape(v) for k, v in copy.items()}, language=language, message=html.escape(message), retry_after=retry_after, action=LOGIN_PATH)) response.headers['Cache-Control'] = 'no-store' # Chromium sends Origin: null from a no-referrer page, which the login would refuse. response.headers['Referrer-Policy'] = 'same-origin' @@ -191,17 +202,17 @@ async def login(request): await asyncio.sleep(max(0, LOGIN_DELAY - (time.monotonic() - started))) auth.reload() if not auth.enabled or auth.record != record: - return password_page(error=PASSWORD_COPY['incorrect']) + return password_page(request, error=PASSWORD_COPY[password_language(request)]['incorrect']) now = time.monotonic() remaining = auth.attempts.remaining(ip, now) if remaining: auth.log(ip, 'locked') - return password_page(retry_after=remaining, status=429) + return password_page(request, retry_after=remaining, status=429) if not correct: remaining = auth.attempts.failed(ip, now) if remaining: auth.log(ip, 'locked') - return password_page(error=PASSWORD_COPY['incorrect'], retry_after=remaining, + return password_page(request, error=PASSWORD_COPY[password_language(request)]['incorrect'], retry_after=remaining, status=429 if remaining else 401) value = auth.issue(remember, time.time()) auth.attempts.success(ip) @@ -284,7 +295,7 @@ async def handle(request): header = request.headers.get('Authorization', '') if not (header[:7].lower() == 'bearer ' and matches(header[7:].strip(), token)): remaining = auth.attempts.remaining(request.remote or '', time.monotonic()) - return password_page(retry_after=remaining, status=429 if remaining else 401) + return password_page(request, retry_after=remaining, status=429 if remaining else 401) else: offered = request.rel_url.query.get('token') if offered is not None: diff --git a/freevideo_engine/adaptive.py b/freevideo_engine/adaptive.py index 71edc9a5..e5355d6e 100644 --- a/freevideo_engine/adaptive.py +++ b/freevideo_engine/adaptive.py @@ -67,6 +67,26 @@ def second_pass_resumable(error, failure, metrics, attempts): return not any(row.get('decision', {}).get('second_pass_resume') for row in attempts) +def decode_resumable(error, failure, metrics, attempts): + """A CUDA error after sampling finished and its latents were saved and validated. + + Any CUDA error qualifies, including a lost or unknown device state: the + retry is one fresh process that loads only the retained latents and decodes + them with the same placement, so no earlier context is reused and the + sampling is not repeated. It is not repeated itself, and a worker that may + still own the device never gets it. Out-of-memory keeps its own path. + """ + from .decode_resume import completed_sampling + if failure.get('kind') not in ('cuda_error', 'code_error') or getattr(error, 'worker_stop_confirmed', None) is False: + return False + recorded = str((metrics or {}).get('error', '')).lower() + if not (cuda_runtime_error(error) or 'cuda error' in recorded) or 'out of memory' in recorded: + return False + if not completed_sampling(metrics): + return False + return not any(row.get('decision', {}).get('decode_resume') for row in attempts) + + def classify_failure(error, metrics=None): """Only positively identified resource failures are retriable. @@ -651,8 +671,10 @@ def execute_attempts(history, identity, profile, canvas, launch, *, history.finish(attempt, failure['outcome'], details=details) row.update(state=failure['outcome'], failure=failure) persist() - resume = (automatic and index < max_retries and not failure['retryable'] - and second_pass_resumable(error, failure, metrics, report['resource_attempts'])) + decode_resume = (automatic and index < max_retries and not failure['retryable'] + and decode_resumable(error, failure, metrics, report['resource_attempts'])) + resume = decode_resume or (automatic and index < max_retries and not failure['retryable'] + and second_pass_resumable(error, failure, metrics, report['resource_attempts'])) if not automatic or not (failure['retryable'] or resume) or index == max_retries: raise previous_ram_retries = sum(r.get('failure', {}).get('kind') == 'ram_pressure' @@ -660,7 +682,12 @@ def execute_attempts(history, identity, profile, canvas, launch, *, previous_gpu_retries = sum(r.get('failure', {}).get('kind') == 'gpu_oom' for r in report['resource_attempts'][:-1]) from .decode_resume import completed_sampling - if resume: + if decode_resume: + updated, decision = copy.deepcopy(current), dict( + reason='A CUDA error ended the request after sampling finished; decode only the ' + 'retained latents in a fresh worker with the same placement.', + numerical_class='unchanged', decode_resume=True) + elif resume: updated, decision = copy.deepcopy(current), dict( reason='A CUDA call was rejected in the second pass; repeat only the second pass ' 'from the retained first pass with the same placement.', diff --git a/freevideo_engine/agent/errors.py b/freevideo_engine/agent/errors.py index 982299d8..2ddc5bf8 100644 --- a/freevideo_engine/agent/errors.py +++ b/freevideo_engine/agent/errors.py @@ -114,6 +114,11 @@ 'Do not retry. Write the subtitle lines yourself with the subtitles edit, or tell the user the model is not installed.'), 'listen_failed': (409, True, 'Listening to the sound did not finish.', 'Read the project, then submit the listening job again with a new ref; if it fails again, write the lines yourself.'), + 'music_download_failed': (409, True, 'The music model could not be downloaded.', + 'Tell the user the music model needs about 20 GB of free disk and a network connection, then ' + 'submit the music job again with a new ref; if it fails again, write the score yourself.'), + 'music_failed': (409, True, 'Composing the music did not finish.', + 'Read the project, then submit the music job again with a new ref; if it fails again, write the score yourself.'), 'clip_source_unusable': (409, False, 'A clip\'s source is missing, or cannot be read as far as the clip uses it.', 'Read the project and check each clip\'s source: it must exist and play, and in + frames must stay ' 'within its real length. Correct the project, then export the new revision.'), @@ -221,6 +226,8 @@ def recovery(code, *, kind=None, stage='request', resumable=True): return ERRORS[code][3] if kind == 'image' and code == 'disk_full': return 'Tell the user to free disk space, then submit the picture again with a new ref.' + if kind == 'music' and code == 'out_of_memory': + return 'Close other GPU applications, then submit the music job again with a new ref.' if kind == 'image' and code == 'no_memory': return 'Lower the picture dimensions or use fewer reference pictures, then submit with a new ref.' if kind in RECOVERY and stage == 'terminal': diff --git a/freevideo_engine/agent/guide.py b/freevideo_engine/agent/guide.py index a2d26c87..c6706dac 100644 --- a/freevideo_engine/agent/guide.py +++ b/freevideo_engine/agent/guide.py @@ -58,15 +58,17 @@ 'for each shot: the shots are generated separately, so the music would not be continuous and would jump at ' 'every cut. Write music in a shot only when the user asks for it there.', 'When the user wants music over the film, keep every shot at N/A and lay one score on one sound track after ' - 'the cut is joined. If the user gives a music file, import it. If there is none, compose it yourself: write ' - 'a program that synthesizes one audio file (WAV, using the standard wave module and numpy or whatever is ' - 'already installed, nothing to install) as long as the finished cut, with its sections, changes and ending ' - 'on the cut points. Import it with POST /imports (MCP freevideo_import_media, CLI imports add), wait until it ' - 'is ready, and insert import: at frame 0 of the project. There is no volume or fade operation, so build ' - 'them into the file: keep the music clearly quieter than the dialogue and ambience, lower it further while ' - 'someone speaks (take the line times from your own screenplay and shot lengths, or run a listen job in align ' - 'mode to time them), and fade it in at the start and out at the end. The audio goes onto a free sound track by ' - 'itself; read the project afterwards and check that one clip covers the whole film.', + 'the cut is joined. If the user gives a music file, import it with POST /imports (MCP freevideo_import_media, ' + 'CLI imports add), wait until it is ready, and insert import: at frame 0 of the project. If there is none, ' + 'run a music job: POST /projects/{id}/music with the revision you read, a caption describing the music in ' + 'English (genre, instruments, tempo, mood and how it changes over the film; there is no singing) and quiet, ' + 'the frame ranges where someone speaks (from your own screenplay and shot lengths, or a listen job in align ' + 'mode). The job composes one piece as long as the timeline, levels it under the dialogue and ambience, lowers ' + 'it in the quiet ranges, fades both ends and lays it on sound track a2 from frame 0; remove an earlier score ' + 'there first. The first use downloads the music model (about 20 GB), so tell the user it can take a while. ' + 'Read the project afterwards and check that one clip covers the whole film. Only if the job fails with ' + 'music_download_failed or music_failed twice, write a program that synthesizes one WAV file (the standard ' + 'wave module and numpy, nothing to install) built the same way and import it.', ], 'references': [ "User message attachments are reference images by default. Use each name directly in media.references, media.first or media.last according to the user's words; if the intended role is unclear, ask one question. Text can be empty when a message has attachments.", @@ -141,6 +143,7 @@ 'POST /projects/{id}/operations sends revision and 1-64 operations, using whole frames at 24 fps. The batch saves once or changes nothing.', 'On 409 project_changed, recompute edits from error.project and its current_revision, then send a new request. On clip_busy wait for generation to finish.', 'POST /projects/{id}/generate reserves pending parts for clip_ids or a start/frames range and fills them as the returned job finishes. POST /projects/{id}/export renders the revision you read; wait for its job.', + 'Music: POST /projects/{id}/music with the revision you read and a caption; wait for the job, then read the project: result.clip_ids are the new clips on a2. Run it after the cut is joined; a second run needs the earlier score removed from a2.', 'Subtitles: POST /projects/{id}/listen with mode transcribe (writes down what is said) or align (times the lines you give) and the revision you read; wait for the job and read result.lines, each with start and end frames. Then add them with a subtitles operation of cues {start, frames, text}; replace true swaps the whole track. With no speech model the job fails with no_speech_model: write the lines yourself.', 'The user sees one line under each picture clip: its note, or else the first line of its description. A description is generation text, often long and in English, so after you add, generate or change a clip, write a note (operation note) in the user\'s language saying what is in the picture: the subject first, then the setting, for example "Train through morning mist" (in Chinese "小火车穿过晨雾"). A short clip shows only the first few characters, about four in Chinese, and those must already make sense on their own. Do not write where the clip sits (opening, middle, ending, part 2): the order of the clips shows that. Do not write how far along it is: the clip already shows its progress. Keep it true when you change the clip again.', 'Reuse ref or Idempotency-Key only for the identical request. Source video/audio/thumbnail URLs require the same Agent token.', @@ -366,6 +369,8 @@ ('POST', '/projects/{id}/listen'): '听剪辑工程的声音,得到字幕行及其时间:对给定字幕行定时,或直接转写。' '先读取 revision,等待返回的任务完成,再用 subtitles 编辑把字幕行加入时间线;' '没有安装语音模型时返回 no_speech_model,改为自己写字幕行', + ('POST', '/projects/{id}/music'): '为剪辑时间线作一段与其等长的配乐,放在第二条声音轨道的开头,音量低于视频声音并带淡入淡出。' + '剪辑拼好后先读取 revision,并删去该轨道上原有的配乐;首次使用会下载音乐模型(约 20 GB)', ('POST', '/images'): '像图片编辑页面一样打开图片;提交 source,会话会保留源图片副本和图层文档', ('GET', '/images/{id}'): '读取图片编辑会话的图层文档、选区和 revision;编辑前先读取 revision 和图层标识', ('PUT', '/images/{id}'): '替换整个图片文档;先读取会话并提交其 revision。' diff --git a/freevideo_engine/agent/jobs.py b/freevideo_engine/agent/jobs.py index 9f645f0d..12436623 100644 --- a/freevideo_engine/agent/jobs.py +++ b/freevideo_engine/agent/jobs.py @@ -32,7 +32,7 @@ ACTIVE = ('waiting', 'queued', 'running', 'cancelling') CAST_AGAIN = 'POST /freevideo/agent/v1/cast/%s/generate to make the reference pictures again.' DONE = ('succeeded', 'failed', 'cancelled', 'interrupted') -GPU_KINDS = ('vlm', 'enhance', 'image') +GPU_KINDS = ('vlm', 'enhance', 'image', 'music') IMAGE_ERRORS = errors.IMAGE_ERRORS IMAGE_WAIT_ERRORS = ('busy', 'busy_video') IMAGE_WAIT_NOTICES = { @@ -1429,6 +1429,9 @@ def run(): elif row['kind'] == 'image_edit': from .images import run_render result = run_render(self, row) + elif row['kind'] == 'music': + from .music import run_music + result = run_music(self, row) else: result = self.run_assemble(row) return result @@ -1785,26 +1788,94 @@ def _poll_chain(self, row): def poll_generation(self, row): prompt_id, node = row['comfy']['prompt_id'], row['comfy']['node'] state = self.port.prompt_state(prompt_id) + polled_at = self.clock() if state == 'running' and row['status'] == 'queued': self.touch(row, status='running', started_at=self.clock()) + history = None if state in ('running', 'pending') else self.port.history(prompt_id) + status = (history or {}).get('status') or {} + outputs = ((history or {}).get('outputs') or {}).get(node) or {} + summary = (outputs.get('freevideo_summary') or [None])[0] + diagnostics = row.setdefault('generation_poll_diagnostics', dict(count=0)) + previous = diagnostics.get('last') + previous_state = tuple(previous.get(key) for key in + ('queue_state', 'history_entry', 'history_success', 'freevideo_summary')) if isinstance(previous, dict) else None + registered_before = diagnostics.get('video_registered_at') + diagnostics['count'] += 1 + diagnostics['last'] = dict(time=polled_at, queue_state=state, + history_entry=history is not None, history_success=status.get('status_str') == 'success', + freevideo_summary=bool(summary)) + video = summary.get('video') if isinstance(summary, dict) else None + if not isinstance(video, str): + for entries in outputs.values(): + for entry in entries if isinstance(entries, list) else (): + if not isinstance(entry, dict) or not isinstance(entry.get('filename'), str): + continue + candidate = '/'.join(part for part in (entry.get('subfolder'), entry['filename']) if part) + if candidate.startswith('FreeVideo/'): + video = candidate + break + if video: + break + if isinstance(video, str) and not diagnostics.get('video_registered_at'): + try: + if self.saved_video(video): + diagnostics['video_registered_at'] = polled_at + except (OSError, JobError, TypeError, ValueError): + pass + current_state = tuple(diagnostics['last'][key] for key in + ('queue_state', 'history_entry', 'history_success', 'freevideo_summary')) + changed = (previous_state != current_state or + registered_before != diagnostics.get('video_registered_at')) + now = polled_at + due = now - row.get('_generation_poll_saved_at', float('-inf')) >= 60 + should_publish = changed or due + if should_publish: + self.save(row) + row['_generation_poll_saved_at'] = now if state in ('running', 'pending'): + if should_publish: + self.publish_generation_diagnostics(row, node) return - history = self.port.history(prompt_id) if history is None: if state == 'missing' and self.clock() - row.get('queued_at', row['updated_at']) > 10: self.finish(row, 'cancelled' if row.get('_cancel') else 'failed', error=dict(code='cancelled' if row.get('_cancel') else 'lost', message='The request left the ComfyUI queue without a result')) + if should_publish: + self.publish_generation_diagnostics(row, node) return - status = history.get('status') or {} - outputs = (history.get('outputs') or {}).get(node) or {} - summary = (outputs.get('freevideo_summary') or [None])[0] if status.get('status_str') == 'success' and summary: self.finish(row, 'succeeded', result=self.generation_result(row, summary)) + if should_publish: + self.publish_generation_diagnostics(row, node) return if status.get('completed') is False or status.get('status_str') == 'error' or not summary: error = self.port.generation_error(history, node) self.finish(row, 'cancelled' if error['code'] == 'cancelled' else 'failed', error=error) + if should_publish: + self.publish_generation_diagnostics(row, node) + + def publish_generation_diagnostics(self, row, node): + if not hasattr(self.port, 'run_output'): + return + try: + output = self.port.run_output(node) + if output is None: + return + path = Path(output).parent / 'agent-generation-poll.json' + value = dict(agent_generation_jobs=[dict(id=job['id'], status=job['status'], + video_registered_at=job.get('generation_poll_diagnostics', {}).get('video_registered_at'), + generation_poll_diagnostics=job.get('generation_poll_diagnostics')) + for job in self.jobs.values() if job.get('kind') == 'generate' + and job.get('generation_poll_diagnostics')]) + temporary = path.with_name(path.name + '.' + uuid.uuid4().hex + '.tmp') + try: + temporary.write_text(json.dumps(value, allow_nan=False), encoding='utf-8') + os.replace(temporary, path) + finally: + temporary.unlink(missing_ok=True) + except (OSError, TypeError, ValueError, AttributeError): + return def generation_result(self, row, summary): spec = row['spec'] diff --git a/freevideo_engine/agent/music.py b/freevideo_engine/agent/music.py new file mode 100644 index 00000000..8a377c9d --- /dev/null +++ b/freevideo_engine/agent/music.py @@ -0,0 +1,74 @@ +"""Music jobs: one piece as long as the timeline, composed locally and laid on sound track a2.""" +import json +import secrets +import time + +from .. import comfy_editor +from .jobs import JobError + +FOLDER = 'freevideo-music' +# A video that starts while the music is made stops it; the job waits for the video and composes again. +PREEMPTIONS = 8 + + +def _busy(jobs): + port = jobs.port + return jobs.runner_busy() or port.queue_busy() or port.image_busy() or port.vlm_busy() + + +def run_music(jobs, row, *, poll=.5): + from ..music import service + spec, port = row['spec'], jobs.port + try: + project = jobs.projects.get(spec['editor_project'])['project'] + except FileNotFoundError as error: + raise JobError('not_found', 'The project does not exist') from error + frames = comfy_editor.duration(project) + if not frames: + raise JobError('invalid_request', 'The timeline has no clips to score') + if frames / comfy_editor.FPS > service.MAX_SECONDS: + raise JobError('invalid_request', 'The timeline is longer than %d s, the longest piece the music model makes' % service.MAX_SECONDS) + relative = '%s/%s-%s.mp4' % (FOLDER, row['id'], secrets.token_hex(4)) + out = port.output_path(relative, create=True) + seed = spec.get('seed', secrets.randbelow(2**32)) + quiet = [(span['start'], span['frames']) for span in spec.get('quiet', ())] + + def phase(name, fraction=None): + row['_phase'] = name + row['_detail'] = json.dumps(dict(fraction=round(fraction, 3))) if fraction is not None else None + + try: + numbers = _compose(jobs, row, service, out, frames, seed, quiet, phase, poll) + phase('placing') + saved, clip_ids = jobs.projects.lay_music(spec['editor_project'], 'generated:' + relative) + except BaseException: + out.unlink(missing_ok=True) + raise + return dict(editor_project=spec['editor_project'], revision=saved['revision'], clip_ids=clip_ids, + source='generated:' + relative, frames=frames, seed=seed, + generate_seconds=numbers.get('generate_s'), load_seconds=numbers.get('load_s')) + + +def _compose(jobs, row, service, out, frames, seed, quiet, phase, poll): + for attempt in range(PREEMPTIONS + 1): + try: + return service.compose(out, row['spec']['caption'], frames, seed=seed, quiet=quiet, fps=comfy_editor.FPS, + server=getattr(jobs.port, 'server', None), cancelled=lambda: bool(row.get('_cancel')), + phase=phase) + except service.Preempted: + if attempt == PREEMPTIONS: + raise JobError('music_failed', 'Videos kept starting while the music was being made') from None + phase('waiting') + # Wait out the video (and anything queued with it) before the GPU is taken again. + while True: + if row.get('_cancel'): + raise JobError('cancelled', 'Cancelled') from None + time.sleep(poll) + if not _busy(jobs): + break + except InterruptedError: + raise JobError('cancelled', 'Cancelled') from None + except service.MusicError as error: + if error.code == 'invalid_request': + raise JobError('invalid_request', 'The timeline is too long for one piece of music') from None + raise JobError(error.code, 'The music was not made (%s)' % (error.reason or error.code), kind='music') from None diff --git a/freevideo_engine/agent/projects.py b/freevideo_engine/agent/projects.py index 02a5a173..dd7e2f2f 100644 --- a/freevideo_engine/agent/projects.py +++ b/freevideo_engine/agent/projects.py @@ -395,6 +395,50 @@ def listen(self, identity, value, *, jobs, ref=None, client=None): self._remember(ref, digest, result, client) return result + def music(self, identity, value, *, jobs, ref=None, client=None): + """Queue a music job for the revision the caller read; a stale revision starts nothing.""" + schema.validate(schema.project_schema('music'), value) + with self.listen_lock: # the same key sent twice at once queues one job + digest, replay = self._receipt('projects:music', dict(id=identity, body=value), ref, client) + if replay is not None: + return replay + stored = self._load(identity)['project'] + if stored['revision'] != value['revision']: + raise ProjectError('project_changed', project=stored) + if not comfy_editor.duration(stored): + raise _invalid('/revision', 'a revision whose timeline has clips', value['revision']) + job = jobs.submit(dict(kind='music', editor_project=identity, + **{k: v for k, v in value.items() if k != 'ref'}), client=client)[0] + result = dict(job=jobs.view(job), jobs=[jobs.view(job)]) + self._remember(ref, digest, result, client) + return result + + def lay_music(self, identity, source): + """Put a finished score on a2 from frame 0 of the project as it is now; returns the save and the new clip IDs.""" + found = comfy_editor.sources(self.root, [source])[source] + if found['missing'] or not found['frames']: + raise JobError('music_failed', 'The music file is not readable') + kind, path = source.split(':', 1) + operation = dict(op='insert', at=0, media=dict(source=dict(kind=kind, id=path), frames=found['frames'], hasVideo=False, + hasAudio=True, width=0, height=0, targetTrack='a2')) + before, refused = {}, {} + def change(project): + before['ids'] = set(_clips(project)) + if any(track['id'] == 'a2' and track['clips'] for track in project['tracks']): + # An insert there would push the whole cut later; an earlier score is removed first. + refused.update(code='audio_full') + return dict(ok=True, project=project) + reply = editor_runtime.runtime().batch(project, [operation]) + if not reply['ok']: + refused.update(reply['error']) + return dict(ok=True, project=project) # nothing changes; the reason is raised below + return reply + saved = self._update(identity, change) + if refused: + code = STATE_CODES.get(refused.get('code'), 'music_failed') + raise JobError(code, 'Sound track a2 cannot take the music from frame 0 (%s)' % refused.get('code')) + return saved, sorted(set(_clips(saved)) - before['ids']) + def _update(self, identity, change, *, tries=4): """Read, change and save; a page save in between makes it read again. None when nothing changed.""" for _ in range(tries): diff --git a/freevideo_engine/agent/routes.py b/freevideo_engine/agent/routes.py index 65c92c0a..0616aab2 100644 --- a/freevideo_engine/agent/routes.py +++ b/freevideo_engine/agent/routes.py @@ -182,7 +182,7 @@ def features(continuation, vlm): 'queue.shared': True, 'queue.busy': True, 'chain.auto': bool(continuation), 'enhance': bool(vlm.get('ready')), 'drama.film': True, 'images.edit': rules, 'images.layers': rules, 'images.heal': rules, 'images.render': rules, 'projects.read': True, - **{'projects.' + name: rules for name in ('insert', 'split', 'move', 'trim', 'remove', 'export', 'generate', 'listen')}} + **{'projects.' + name: rules for name in ('insert', 'split', 'move', 'trim', 'remove', 'export', 'generate', 'listen', 'music')}} def whitebox_support(jobs): @@ -231,6 +231,7 @@ def capabilities(port, jobs): 'enhance': 'Rewrite a rough idea into a full H3 prompt with the local VLM', 'export': 'Render a saved revision of an editing project to a video', 'listen': 'Listen to the sound of an editing project for subtitle lines and times', + 'music': 'Compose music for an editing project and lay it on its second sound track', 'image_edit': 'Render a saved revision of an image session to a picture'}, limits=dict(width=[256, 4096, 32], height=[256, 4096, 32], seconds=[1.625, 60 if continuation else BLOCK_SECONDS], seconds_tested=[4, 15], @@ -2128,7 +2129,7 @@ async def handler(request): raise JobError('idempotency_conflict', 'Conflicting keys') args = [] if operation == 'create' else [identity] extra = dict(client=request.headers.get('X-FreeVideo-Client', '')[:80] or None) - if operation in ('generate', 'export', 'listen'): + if operation in ('generate', 'export', 'listen', 'music'): extra['jobs'] = jobs if operation == 'generate': def capability(size): @@ -2138,9 +2139,9 @@ def capability(size): return None extra['capability'] = capability result = await asyncio.to_thread(getattr(projects, operation), *args, value, ref=key or ref, **extra) - if operation in ('generate', 'export', 'listen'): + if operation in ('generate', 'export', 'listen', 'music'): await jobs.tick() - return reply(result, 201 if operation in ('create', 'generate', 'export', 'listen') else 200) + return reply(result, 201 if operation in ('create', 'generate', 'export', 'listen', 'music') else 200) return guarded(handler) for operation, (method, suffix, _) in contract.PROJECT_OPERATIONS.items(): diff --git a/freevideo_engine/agent/schema.py b/freevideo_engine/agent/schema.py index a9660bb7..8f37c4fd 100644 --- a/freevideo_engine/agent/schema.py +++ b/freevideo_engine/agent/schema.py @@ -64,7 +64,7 @@ def enum(values, **kw): SAVED = string(43, pattern=r'^\d{4}-\d{2}-\d{2}/[0-9a-f]{32}$') SLUG = string(64, pattern=r'^[A-Za-z0-9][A-Za-z0-9._-]{0,63}$') REF = string(128, pattern=r'^[A-Za-z0-9][A-Za-z0-9._:-]{0,127}$', - description='Idempotency key; reuse only with the same request content. Batch row refs also name dependencies.') + description='Idempotency key; reuse only with the same request content.') FILE = string(1024, pattern=r'^(?!/)(?!.*(?:\\|:))(?!.*(?:^|/)\.\.(?:/|$)).+$') FRAME = {'oneOf': [enum(('first', 'last')), number(0, 86400)]} # Seconds into the video; the bound keeps it a float. WAIT = number(0, default=0, description='x-clamp caps accepted wait values at 25 seconds without rejecting larger values.', **{'x-clamp': 25}) @@ -159,7 +159,7 @@ def media(image=False, *, local=False, clip=False, video=False, early=False): def common(): - return dict(label=string(120), project=SLUG, metadata={'x-json-bytes': 16 * 1024, 'description': 'x-json-bytes limits serialized finite JSON to 16384 bytes, which JSON Schema cannot measure.'}, + return dict(label=string(120), project=SLUG, metadata={'x-json-bytes': 16 * 1024, 'description': 'x-json-bytes: serialized finite JSON of at most 16384 bytes.'}, after=array(JOB), ref=REF) @@ -311,6 +311,23 @@ def listen(*, local=False): **{'x-job-kind': 'listen'}) +MUSIC_FIELDS = dict( + caption=string(512, description='The music in English: genre, instruments, tempo, mood and how it changes over the film. No singing.'), + seed=number(0, 2**32 - 1, integer=True, description='Same seed and caption, same piece.'), + quiet=array(obj(dict(start=EDITOR_FRAME, frames=number(1, 10**7, integer=True)), ('start', 'frames')), 1, 500, + description='Frames where someone speaks; the music is lower there.')) + + +def music(*, local=False): + return obj(dict(kind={'const': 'music', 'type': 'string'}, + editor_project=string(34, pattern=r'^e-[0-9a-f]{32}$', description='The editing project to score (GET /projects).'), + revision=number(1, 2**53 - 1, integer=True, description='The saved revision you read; the music follows the timeline length when the job runs.'), + **MUSIC_FIELDS, **common()), ('kind', 'editor_project', 'revision', 'caption'), + description='Compose one piece of music as long as the timeline with the local music model and lay it on sound track a2 from ' + 'frame 0, quieter than the video sound and faded in and out. The first use downloads the model (about 20 GB).', + **{'x-job-kind': 'music'}) + + def image_edit(*, local=False): return obj(dict(kind={'const': 'image_edit', 'type': 'string'}, image=dict(IMAGE, description='The image session to render (POST /images).'), @@ -321,11 +338,12 @@ def image_edit(*, local=False): KINDS = dict(generate=generate, assemble=assemble, vlm=vlm, enhance=enhance, image=image, - export=export, listen=listen, image_edit=image_edit) + export=export, listen=listen, music=music, image_edit=image_edit) -def job_schema(*, local=False, refs=False): - value = {'oneOf': [f(local=local) for f in KINDS.values()], 'x-discriminator': 'kind', +def job_schema(*, local=False, refs=False, batch=False): + # Music scores a joined cut, so it is not a storyboard row. + value = {'oneOf': [f(local=local) for kind, f in KINDS.items() if not (batch and kind == 'music')], 'x-discriminator': 'kind', 'description': 'x-discriminator selects field diagnostics by kind; all accepted branches use standard oneOf.'} return batch_references(value) if refs else value @@ -341,7 +359,7 @@ def batch_references(value): 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}) + return obj(dict(jobs=array(job_schema(local=local, refs=True, batch=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(): @@ -442,7 +460,7 @@ def route(summary, body=None, query=None, **kw): if method == 'GET': extra['query'] = shape else: - extra.update(body=shape, idempotent=True, **({'status': 201} if name in ('create', 'generate', 'export', 'listen') else {})) + extra.update(body=shape, idempotent=True, **({'status': 201} if name in ('create', 'generate', 'export', 'listen', 'music') else {})) routes[method, suffix] = route(summary, **extra) for name, (method, suffix, summary) in WHITEBOX_OPERATIONS.items(): extra = dict(path=obj({'id': WHITEBOX}, ('id',))) if '{id}' in suffix else {} @@ -825,11 +843,12 @@ def image_schema(name): 'generate': ('POST', '/projects/{id}/generate', 'Make new picture and sound for selected clips or a frame range, as the page Generate does. First get revision and clip IDs; the parts are reserved at once and fill in as the returned job finishes.'), 'export': ('POST', '/projects/{id}/export', 'Render one saved revision to an MP4 in the saved videos. First get the revision; wait for the returned job, then read its result video.'), 'listen': ('POST', '/projects/{id}/listen', 'Listen to the timeline sound for subtitles: mode align times lines you give, transcribe writes down what is said. First get the revision; wait for the returned job, then add the lines with the subtitles edit. Nothing is saved; with no speech model for the language the call fails with no_speech_model, so write the lines by hand.'), + 'music': ('POST', '/projects/{id}/music', 'Compose music as long as the joined cut and lay it on sound track a2 at frame 0, under the video sound. First get the revision; a2 must be empty. Wait for the returned job; the first use downloads the model (20 GB).'), } PROJECT_TOOLS = dict(list='freevideo_projects', get='freevideo_project_read', create='freevideo_project_create', write='freevideo_project_write', operations='freevideo_editor_operations', generate='freevideo_project_generate', export='freevideo_export', - listen='freevideo_project_listen') + listen='freevideo_project_listen', music='freevideo_project_music') SESSION_OPERATIONS = { 'create': ('POST', '/sessions', 'Create a work conversation for messages and plans shared with the user.'), @@ -943,6 +962,8 @@ def project_schema(name): return obj(dict(revision=revision, ref=REF), ('revision',)) if name == 'listen': return obj(dict(revision=revision, **LISTEN_FIELDS, ref=REF), ('revision', 'mode')) + if name == 'music': + return obj(dict(revision=revision, **MUSIC_FIELDS, ref=REF), ('revision', 'caption')) if name != 'operations': return EMPTY fields = { @@ -1095,8 +1116,9 @@ def bad(allowed, got=value, suffix='', **kw): return [diagnostic(path + '/' + field, ', '.join(variants), selected)] return issues(variants[selected], value, path, refs=refs) kind = value.get('kind', 'generate') - if not isinstance(kind, str) or kind not in KINDS: - return [diagnostic(path + '/kind', ', '.join(KINDS), kind)] + kinds = [s.get('x-job-kind') for s in schema['oneOf']] + if not isinstance(kind, str) or kind not in kinds: + return [diagnostic(path + '/kind', ', '.join(kinds), kind)] chosen = next(s for s in schema['oneOf'] if s.get('x-job-kind') == kind) return issues(chosen, value, path, refs=refs) for rule in schema.get('allOf', []): diff --git a/freevideo_engine/agent/spec.py b/freevideo_engine/agent/spec.py index 8fe24d31..e70c8b95 100644 --- a/freevideo_engine/agent/spec.py +++ b/freevideo_engine/agent/spec.py @@ -290,6 +290,9 @@ def normalize(value, *, port=None): elif kind == 'listen': result.update(editor_project=value['editor_project'], revision=value['revision'], mode=value['mode'], language=value.get('language', 'auto'), **{k: value[k] for k in ('lines', 'range') if k in value}) + elif kind == 'music': + result.update(editor_project=value['editor_project'], revision=value['revision'], caption=value['caption'], + **{k: value[k] for k in ('seed', 'quiet') if k in value}) elif kind == 'image_edit': result.update(image=value['image'], revision=value['revision']) else: diff --git a/freevideo_engine/background_compile.py b/freevideo_engine/background_compile.py index 377c09ad..a2aa7657 100644 --- a/freevideo_engine/background_compile.py +++ b/freevideo_engine/background_compile.py @@ -12,7 +12,8 @@ It runs only before the session's first generation. After that, the resident encoder holds GPU memory that a request would release, so the plan would no -longer be the request's. An interrupted run writes no receipt. When the run +longer be the request's. An interrupted run leaves the cache receipt intact +and records diagnostic status separately. When the run was stopped or found no room (background_room), the session waits twice as long before the next try, up to MAX_DELAY; the next start tries again too. """ @@ -24,6 +25,7 @@ from . import processes from .long_paths import io_path +from .precompile_state import read, record START_DELAY = 30. # A compile that is owed (nothing compiled yet for this code, driver or GPU) @@ -70,6 +72,7 @@ def __init__(self): self.thread = None self.lock = None self.attempts = [] + self.root = None def start(self, command, *, environment, cwd, log, lock, busy, delay=START_DELAY): """Run `command` once the session is quiet; False when this session already did. @@ -80,6 +83,7 @@ def start(self, command, *, environment, cwd, log, lock, busy, delay=START_DELAY if self.thread is not None or self.closed.is_set(): return False self.lock = Path(lock) + self.root = Path(log).parent self.thread = threading.Thread(target=self._work, args=(command, environment, cwd, Path(log), busy, delay), name='FreeVideo background compile', daemon=True) self.thread.start() @@ -96,12 +100,15 @@ def stop(self, *, final=False): with self.guard: process, self.process = self.process, None if process is None: + if final and self.thread is not None and self.thread.is_alive(): + record(self.root, 'not_started', reason='foreground_work', final=True) return None row = self.attempts[-1] if process.poll() is None: # Marked before the signal: the worker thread then retries # instead of reading the exit as a finished run. row['stopped'] = True + record(self.root, 'not_started', reason='foreground_work', final=final) processes.stop(process, grace=STOP_GRACE) row['stop_seconds'] = time.monotonic() - tick # Every process of the tree holds the lease; once it is free the @@ -129,20 +136,30 @@ def _released(self, deadline): def _quiet(self, busy, delay): """True after `delay` seconds without work; False once closed.""" since = time.monotonic() + postponed = False while time.monotonic() - since < delay: if self.closed.wait(1.): return False if busy(): since = time.monotonic() + if not postponed: + record(self.root, 'not_started', reason='session_busy') + postponed = True + else: + postponed = False return True def _work(self, command, environment, cwd, log, busy, delay): # The first wait may be shorter than the retry backoff, which always # starts from START_DELAY: a stopped or crowded-out run is not hammered. pause = START_DELAY + self.root = log.parent for attempt in range(ATTEMPTS): if attempt: pause = delay = min(2 * pause, MAX_DELAY) + # A retry retains the actual pressure/stop reason from its last + # attempt instead of replacing it with the backoff mechanism. + record(self.root, 'not_started', reason=None if attempt else 'idle_delay') if not self._quiet(busy, delay): return with self.guard: @@ -152,12 +169,19 @@ def _work(self, command, environment, cwd, log, busy, delay): 'a generation request stops it.') row = dict(started=time.time()) self.attempts.append(row) - with io_path(log).open('ab') as stream: - from .triton_compat import temporary_environment - environment = temporary_environment(environment.get('FREEVIDEO_HOME') or self.lock.parent, environment) - process = self.process = processes.popen( - command, env=environment, cwd=cwd, stdin=subprocess.DEVNULL, stdout=stream, - stderr=subprocess.STDOUT, start_new_session=True, supervise=True) + try: + with io_path(log).open('ab') as stream: + from .triton_compat import temporary_environment + environment = temporary_environment(environment.get('FREEVIDEO_HOME') or self.lock.parent, + environment) + process = self.process = processes.popen( + command, env=environment, cwd=cwd, stdin=subprocess.DEVNULL, stdout=stream, + stderr=subprocess.STDOUT, start_new_session=True, supervise=True) + except BaseException: + record(self.root, 'unknown') + raise + # Only a started process is running, with its own identity. + record(self.root, 'running', started_at=row['started'], shapes=[], pid=process.pid) while process.poll() is None: if busy(): self.stop() @@ -178,7 +202,14 @@ def _work(self, command, environment, cwd, log, busy, delay): row.update(returncode=process.poll(), seconds=time.time() - row['started']) from .precompile import BUSY_EXIT if not stopped and row['returncode'] != BUSY_EXIT: + receipt, _ = read(self.root / 'precompile.json', 2 * 2**20) + record(self.root, 'complete' if row['returncode'] == 0 and receipt.get('success') is True + else 'unknown') return # Compiled, or failed with a receipt that says why. + if not stopped: + # The child normally wrote the resource reason. Keep it even + # if the child exited before publishing diagnostic status. + record(self.root, 'not_started') COMPILE = BackgroundCompile() diff --git a/freevideo_engine/bootstrap.py b/freevideo_engine/bootstrap.py index 34b94f44..f477ceea 100644 --- a/freevideo_engine/bootstrap.py +++ b/freevideo_engine/bootstrap.py @@ -2020,8 +2020,10 @@ def feedback(event): return micromamba, [path for path, _, _ in rows[1:]] def install_git_windows(self): - # An older Git on PATH cannot fetch the pinned sources; use the portable one. - if usable_git(self.env['PATH'], self.env): + # An older Git on PATH, or one whose HTTPS helper cannot load libcurl, cannot + # fetch the pinned sources; use the portable one. + found = usable_git(self.env['PATH'], self.env) + if found and network.git_https_ready(found, self.env): return spec = self.versions['windows']['git'] archive = self.root / 'downloads' / ('MinGit-' + spec['version'] + '.zip') diff --git a/freevideo_engine/chain_join.py b/freevideo_engine/chain_join.py index d1fd11a1..e5a7e659 100644 --- a/freevideo_engine/chain_join.py +++ b/freevideo_engine/chain_join.py @@ -171,6 +171,8 @@ def join_files(sources, target, cancelled=None, frames=None, leads=None, expecte `leads`: optional `leads(following, previous)` returning a seam lead kept by the chain; the shared continuation record is read only when it has none. `expected`: how many seams should be repaired; only reported back, the caller decides. + The returned `seam_notes` list describes unrepaired seams, numbered from one, + without source paths or exception messages. `original_audio`: mux the original selection and close the picture on its source-sample clock; a rounded completed prefix keeps its existing frames. """ @@ -220,6 +222,7 @@ def audio_lead_in(video, previous=None): temporary = target.with_name('.join-' + uuid.uuid4().hex + '.mp4') wanted, frames, samples = frames, 0, 0 used_tail = False + blocks_joined = 0 try: with av.open(str(sources[0])) as first: if len(first.streams.video) != 1: @@ -264,6 +267,7 @@ def audio_lead_in(video, previous=None): fifo = av.AudioFifo() joined_audio = _JoinedAudio(output, audio, fifo, rate, layout, audio_lead_in, apply_audio_lead_in, cancelled, voice) + joined_audio.seams_total = len(sources) - 1 streams = [video] + ([container.streams.audio[0]] if sound is not None else []) shift, count = None, 0 # A decoder returns the encoder's padding after the last sample; the stream's own @@ -309,6 +313,7 @@ def audio_lead_in(video, previous=None): if not count: raise ValueError('A block has no picture') frames += count + blocks_joined += 1 if audio is not None: samples = joined_audio.finish() for packet in audio.encode(None): @@ -347,6 +352,11 @@ def audio_lead_in(video, previous=None): audio_samples=samples, sample_rate=rate, seams_repaired=joined_audio.seams_repaired if joined_audio is not None else 0, seams_expected=expected, + seam_notes=(joined_audio.seam_notes if joined_audio is not None else + [dict(seam=seam, step='no_audio' if seam < blocks_joined else 'not_joined', + reason='previous block has no sound track' if seam < blocks_joined + else 'next block is outside the requested frame range') + for seam in range(1, len(sources))] if original_audio is None else []), voice=[[round(setting[0], 1), round(setting[1], 1)] if setting is not None else None for setting in voice] if voice is not None and audio is not None else None) except (JoinCancelled, _VoiceFailed): @@ -384,10 +394,14 @@ def __init__(self, output, stream, fifo, rate, layout, read_lead, apply_lead, ca self.cancelled, self.next_head = cancelled, None self.warned = False self.written = self.seams_repaired = 0 + self.seam = self.seams_total = 0 + self.seam_notes = [] def start(self, source, stream, limit, last=None, following=None, total_frames=None): import av context = stream.codec_context + self.seam += 1 + self.source_format = context.sample_rate, context.layout.name, getattr(context.layout, 'nb_channels', None) self.source, self.following = str(source), str(following) if following is not None else None self.limit, self.heard, self.kept = limit, 0, 0 # end() has released the previous tail and flushed its resampler. Count @@ -453,31 +467,60 @@ def _hold(self, data): def _release(self): tail = self.tail if tail is not None and len(tail) == SEAM_TAIL and self.following is not None: + stage = 'reading next block sound' try: - head = _audio_head(self.following, self.cancelled) + # Keep the two-argument head reader usable by existing adapters. The + # cancellation callback also carries this call's optional diagnostic sink. + def checked(): + return self.cancelled is not None and self.cancelled() + checked.seam_note = self._note + head = _audio_head(self.following, checked) if head is not None: + stage = 'reading seam record' lead = self.read_lead(self.following, previous=self.source) if self.read_lead is not None else None if lead is not None: + stage = 'copying seam windows' repaired_tail, repaired_head = tail.copy(), head.copy() + stage = 'applying seam record' if self.apply_lead(repaired_tail, repaired_head, lead): tail = repaired_tail self.next_head = self.following, repaired_head self.seams_repaired += 1 else: + self._note('apply_rejected', 'seam record was rejected by apply') self._warn() + else: + self._note('no_record', 'seam record is unavailable' if self.read_lead is not None + else 'seam record reader is unavailable') + else: + self._note('next_head', 'next block effective head is unavailable') except JoinCancelled: raise - except Exception: + except Exception as error: + name = type(error).__name__ + self._note('exception', '%s raised %s' % (stage, name), stage=stage, exception=name) self._warn() + elif self.following is not None or self.seam <= self.seams_total: + if tail is None or len(tail) != SEAM_TAIL: + count = len(tail) if tail is not None else 0 + self._note('previous_tail', 'previous block effective tail has %d samples, needs %d' % + (count, SEAM_TAIL), samples=count, required=SEAM_TAIL) + else: + self._note('not_joined', 'next block is outside the requested frame range') if tail is not None: self._queue(tail) self.tail, self.released = None, True def _warn(self): if not self.warned: - logging.warning('[FreeVideo] A chain sound seam could not be repaired; the blocks are joined as they are') + logging.warning('[FreeVideo] A chain sound seam could not be repaired; the blocks are joined as they are' + + _seam_reasons(self.seam_notes[-1:])) self.warned = True + def _note(self, step, reason, **values): + if not any(note['seam'] == self.seam for note in self.seam_notes): + self.seam_notes.append(dict(seam=self.seam, step=step, reason=reason, **values)) + def _queue(self, data): import av import numpy as np @@ -495,9 +538,17 @@ def end(self): self._native(frame) if not self.released: self._release() + elif self.following is not None or self.seam <= self.seams_total: + rate, layout, channels = self.source_format + self._note('previous_format', 'previous block sound is %d Hz %s, needs 32000 Hz stereo' % (rate, layout), + sample_rate=rate, layout=layout, channels=channels, required_sample_rate=32000, + required_channels=2) self._encode(self.resampler.resample(None)) def finish(self): + for seam in range(self.seam + 1, self.seams_total + 1): + self.seam_notes.append(dict(seam=seam, step='not_joined', + reason='next block is outside the requested frame range')) return _encode(self.output, self.stream, self.fifo, [], 0, None, self.written, self.rate, self.layout, final=True, shaper=self.shaper)[1] @@ -506,14 +557,24 @@ def _audio_head(source, cancelled): """Decode just the next block's kept native head, without changing its files.""" import av import numpy as np + note = getattr(cancelled, 'seam_note', lambda *args, **kwargs: None) _check(cancelled) with av.open(source) as container: if not container.streams.audio: + note('no_audio', 'next block has no sound track') return None stream = container.streams.audio[0] if stream.codec_context.sample_rate != 32000 or stream.codec_context.layout.name != 'stereo': + context = stream.codec_context + note('next_format', 'next block sound is %d Hz %s, needs 32000 Hz stereo' % + (context.sample_rate, context.layout.name), sample_rate=context.sample_rate, + layout=context.layout.name, channels=context.layout.nb_channels, + required_sample_rate=32000, required_channels=2) return None if stream.duration is not None and round(stream.duration * stream.time_base * 32000) < SEAM_HEAD: + count = round(stream.duration * stream.time_base * 32000) + note('next_head', 'next block effective head has %d samples, needs %d' % (count, SEAM_HEAD), + samples=count, required=SEAM_HEAD) return None resampler = av.AudioResampler(format='fltp', layout='stereo', rate=32000) chunks, count = [], 0 @@ -528,9 +589,20 @@ def _audio_head(source, cancelled): count += keep if count == SEAM_HEAD: return np.ascontiguousarray(np.concatenate(chunks), dtype=np.float32) + note('next_head', 'next block effective head has %d samples, needs %d' % (count, SEAM_HEAD), + samples=count, required=SEAM_HEAD) return None +def _seam_reasons(notes): + # A seam with no record, or one the cut never reached, is usually a seam nobody asked to repair: name the others first. + notes = sorted(notes, key=lambda note: note['step'] in ('no_record', 'not_joined')) + reasons = ['seam %d: %s' % (note['seam'], note['reason']) for note in notes[:3]] + if len(notes) > 3: + reasons.append('and %d more' % (len(notes) - 3)) + return ' (' + '; '.join(reasons) + ')' if reasons else '' + + def _encode(output, stream, fifo, resampled, heard, limit, written, rate, layout, final=False, shaper=None): """Queue one block's resampled sound up to its length and encode whole encoder frames; `final` writes the rest.""" import av @@ -633,7 +705,8 @@ def kept(following, previous): # The original selection replaces every block's sound, so there are no seam sounds to repair. if original_audio is None and expected_seams is not None and joined['seams_repaired'] < expected_seams: # A seam sound left unrepaired is never reported as a finished video. - raise ValueError('Only %d of %d seam sounds could be repaired' % (joined['seams_repaired'], expected_seams)) + raise ValueError('Only %d of %d seam sounds could be repaired' % (joined['seams_repaired'], expected_seams) + + _seam_reasons(joined.get('seam_notes', []))) report = dict(success=True, geometry=dict(width=joined['width'], height=joined['height'], frames=joined['frames'], fps=FPS, seconds=joined['frames'] / FPS), request_seconds=_request_seconds(sources), video={}, encoding={}, diff --git a/freevideo_engine/comfy_chain.py b/freevideo_engine/comfy_chain.py index 0dc9b560..210fcc53 100644 --- a/freevideo_engine/comfy_chain.py +++ b/freevideo_engine/comfy_chain.py @@ -124,7 +124,7 @@ def submit(self, prompt_id, prompt, extra_data): # This request belongs to our own server, even with an inherited proxy. handlers = [request.ProxyHandler({})] if tls: - context = ssl.create_default_context() + context = ssl.SSLContext(ssl.PROTOCOL_TLS_CLIENT) context.check_hostname = False context.verify_mode = ssl.CERT_NONE handlers.append(request.HTTPSHandler(context=context)) diff --git a/freevideo_engine/comfy_launcher_runtime.py b/freevideo_engine/comfy_launcher_runtime.py index d5764a7d..c78a85a9 100644 --- a/freevideo_engine/comfy_launcher_runtime.py +++ b/freevideo_engine/comfy_launcher_runtime.py @@ -494,8 +494,16 @@ def launcher_comfy(selected, info): return False +ENGINE_UPDATE_WAITS = 'ComfyUI is running a job, so the engine update waits' + + +def _refused(error): + return isinstance(error, ConnectionRefusedError) or isinstance(getattr(error, 'reason', None), ConnectionRefusedError) + + def queue_busy(url): - """Preserve the legacy check; a pool must confirm every executor is idle.""" + """Preserve the legacy check; a pool must confirm every executor is idle. A refused + connection means nothing listens, so there is no job; a timeout is not that answer.""" from .gpu_pool import pool_enabled pooled = pool_enabled() try: @@ -506,8 +514,8 @@ def queue_busy(url): state = get_json(loopback_url(url) + '/freevideo/pool', timeout=3) return state.get('busy') is not False or state.get('confirmed') is not True return False - except (OSError, ValueError, AttributeError): - return pooled + except (OSError, ValueError, AttributeError) as error: + return pooled and not _refused(error) @plain_errors @@ -622,6 +630,34 @@ def checkout_entry(data, folder): and (io_path(folder) / 'freevideo_engine' / 'comfy_nodes.py').is_file()) +PRESTARTUP_MARK = '# Installed by FreeVideo launcher.\n' + + +def install_prestartup(target, engine): + """ComfyUI runs prestartup_script.py of a node folder before it imports its server. The launcher's + small loader runs the certificate fix from the source; a file that is not the launcher's stays.""" + path = target / 'prestartup_script.py' + if path.is_file() and not path.read_text(encoding='utf-8', errors='replace').startswith(PRESTARTUP_MARK): + return + code = (PRESTARTUP_MARK + + "import json as _json\nfrom pathlib import Path as _Path\n" + "_source = _Path(_json.loads((_Path(__file__).parent / 'freevideo-launcher.json').read_text(encoding='utf-8'))['source'])\n" + + launch_view.entry_lines(launch_view.for_engine(engine)) + + "_script = _source / 'freevideo_engine' / 'windows_ssl.py'\n" + "if _script.is_file():\n" + " exec(compile(_script.read_text(encoding='utf-8'), str(_script), 'exec'))\n").encode() + if path.is_file() and path.read_bytes() != code: + # An older loader, or one the user added to: keep a copy, as for the entry point. + backup = engine / 'launcher' / 'backups' / (target.name + '-' + hashlib.sha256(path.read_bytes()).hexdigest()[:12] + '-prestartup.py') + backup.parent.mkdir(parents=True, exist_ok=True) + if not backup.exists(): + backup.write_bytes(path.read_bytes()) + if not path.is_file() or path.read_bytes() != code: + temporary = path.with_suffix('.py.launcher-tmp') + temporary.write_bytes(code) + os.replace(temporary, path) + + def deploy(root, source, engine): """Update only our entry point; preserve source edits, libraries and templates.""" root, source, engine = map(resolved, (root, source, engine)) @@ -657,6 +693,7 @@ def deploy(root, source, engine): temporary = entry.with_suffix('.py.launcher-tmp') temporary.write_bytes(code) os.replace(temporary, entry) + install_prestartup(target, engine) template = (source / 'example_workflows' / TEMPLATE).read_bytes() digest = hashlib.sha256(template).hexdigest()[:12] saved = [] @@ -1151,8 +1188,7 @@ def _install(self, accepted): # Whoever started it (another launcher, the freevideo command), a ComfyUI of this installation # that is generating keeps its packages: the switch waits for the job. if queue_busy(loopback_url(self.server_address or self.selection['url'])): - raise RuntimeError('ComfyUI is running a job, so the engine update waits. Try again after it finishes; ' - 'the job is left running.') + raise RuntimeError(ENGINE_UPDATE_WAITS + '. Try again after it finishes; the job is left running.') self._stop_server_for_setup(selected['url']) engine_binding.rebind(selected['engine'], machine, selected['source']) self.stage('nodes', label='Install FreeVideo workflow') @@ -1178,12 +1214,11 @@ def _stop_server_for_setup(self, url): keeps loaded after a video (Windows refuses to replace SageAttention's _fused.pyd), and a model switch must not leave the old model in that worker. Stop our own idle server; connecting after setup starts it - again. A running job is never stopped.""" - if not self.owns_server(): - return + again. A running job is never stopped, in a server another program started either.""" if queue_busy(loopback_url(self.server_address or url)): raise RuntimeError('ComfyUI is running a job. Try again after it finishes; the job is left running.') - self.stop_owned_server() + if self.owns_server(): + self.stop_owned_server() def ensure_shortcut(self): from .desktop_shortcut import create_after_install diff --git a/freevideo_engine/comfy_nodes.py b/freevideo_engine/comfy_nodes.py index 77f9dbb3..2ff56df0 100644 --- a/freevideo_engine/comfy_nodes.py +++ b/freevideo_engine/comfy_nodes.py @@ -233,6 +233,11 @@ def execute(cls, text, width, height, seconds, seed, first=None, last=None, stop_image_for_generation(PromptServer.instance) except Exception as error: logging.warning('[FreeVideo] Could not stop image generation before the video: %s', error) + try: # nor music for a cut, which composes again after the video + from .music.service import stop_for_generation as stop_music_for_generation + stop_music_for_generation(PromptServer.instance) + except Exception as error: + logging.warning('[FreeVideo] Could not stop music before the video: %s', error) prewarm = {} result_reused = [False] result_report = [None] diff --git a/freevideo_engine/comfy_service.py b/freevideo_engine/comfy_service.py index 3f4f54dd..11602648 100644 --- a/freevideo_engine/comfy_service.py +++ b/freevideo_engine/comfy_service.py @@ -267,8 +267,9 @@ def request(url, *, headers=None, timeout=3): class Stay(HTTPRedirectHandler): def redirect_request(self, *args, **kwargs): return None - context = ssl.create_default_context() - context.check_hostname, context.verify_mode = False, ssl.CERT_NONE # Our own certificate, possibly self-signed. + # Our own certificate, possibly self-signed: nothing is verified, so the system certificates are not loaded. + context = ssl.SSLContext(ssl.PROTOCOL_TLS_CLIENT) + context.check_hostname, context.verify_mode = False, ssl.CERT_NONE opener = build_opener(ProxyHandler({}), Stay(), HTTPSHandler(context=context)) try: with opener.open(Request(url, headers=headers or {}), timeout=timeout) as reply: diff --git a/freevideo_engine/diagnostic_resources.py b/freevideo_engine/diagnostic_resources.py index d8198617..ea57cffa 100644 --- a/freevideo_engine/diagnostic_resources.py +++ b/freevideo_engine/diagnostic_resources.py @@ -156,6 +156,50 @@ def encoder_prewarm(value): return result +PRECOMPILE_REASONS = ('idle_delay', 'session_busy', 'foreground_work', + 'ram_pressure', 'gpu_busy', 'gpu_memory_pressure', 'gpu_processes_unavailable', + 'gpu_in_use', 'lease_busy', 'unknown') + + +def precompile_shape(value): + """Allowlisted request shape only, never cache/conditioning/model paths.""" + value = mapping(value) + result = {key: value[key] for key in ('width', 'height', 'frames', 'base_steps', 'refine_steps') + if type(value.get(key)) is int and value[key] > 0} + if not all(key in result for key in ('width', 'height', 'frames')): + return {} + if value.get('task') in ('t2va', 'fl2va', 'ref2va'): + result['task'] = value['task'] + if type(value.get('two_pass')) is bool: + result['two_pass'] = value['two_pass'] + result.update({key: number for key, number in numbers(value, ('upscale_factor',)).items() if number > 0}) + tokens = sequence(value.get('text_tokens'))[:8] + if tokens and all(type(n) is int and n > 0 for n in tokens): + result['text_tokens'] = list(tokens) + return result + + +def precompile(value): + """Fresh bounded structure for local and uploaded precompile evidence.""" + value = mapping(value) + state = value.get('state') + result = dict(state=state if state in ('complete', 'running', 'not_started', 'unknown') else 'unknown', + covered_shapes=[], scope='installation_snapshot') + for row in sequence(value.get('covered_shapes'))[:16]: + shape = precompile_shape(row) + if shape and shape not in result['covered_shapes']: + result['covered_shapes'].append(shape) + result.update(numbers(value, ('updated_at', 'started_at'))) + if value.get('receipt') in ('available', 'missing', 'invalid'): + result['receipt'] = value['receipt'] + last = mapping(value.get('last_postpone')) + at = numbers(last, ('at',)) + if at: + result['last_postpone'] = dict(at, reason=last.get('reason') + if last.get('reason') in PRECOMPILE_REASONS else 'unknown') + return result + + def embedding(value): """Tensor metadata only: no indices, tensor values, prompt or media content.""" value = mapping(value) @@ -323,7 +367,14 @@ def kernel_receipt(value): def numbers(value, keys): value = mapping(value) return {k: value[k] for k in keys if type(value.get(k)) in (int, float) - and math.isfinite(value[k]) and value[k] >= 0} + and finite(value[k]) and value[k] >= 0} + + +def finite(number): + try: + return math.isfinite(number) + except OverflowError: # An int beyond the float range, e.g. 10**400. + return False def weight_placement(value): diff --git a/freevideo_engine/diagnostic_summary.py b/freevideo_engine/diagnostic_summary.py index cac5df20..3433864f 100644 --- a/freevideo_engine/diagnostic_summary.py +++ b/freevideo_engine/diagnostic_summary.py @@ -32,6 +32,7 @@ def summary_report(report): result = dict(kind='generation', status=summary.get('status') if summary.get('status') in ('complete', 'reused', 'retrying', 'running') else 'incomplete', version=__version__, hardware={}, config={}, geometry={}, memory={}, stages=[], steps=[], errors=[]) + result['precompile'] = diagnostic.precompile(summary.get('precompile')) for k in ('gpu_name', 'system', 'architecture', 'driver_version', 'torch_version', 'cuda_version'): if _label(hardware.get(k)): result['hardware'][k] = hardware[k] diff --git a/freevideo_engine/diagnostics.py b/freevideo_engine/diagnostics.py index b33fa147..48df0e79 100644 --- a/freevideo_engine/diagnostics.py +++ b/freevideo_engine/diagnostics.py @@ -191,7 +191,7 @@ def report_files(directory): if name in NAMES or name.endswith(('.log', '.jsonl', '.request.json', '.debug.json', '.encoding.json', '.engine.json', '.memory.json', '.gpu.json', '.sampling-memory.json', '.comfy-request.json', - 'comfy-request.json')): + 'comfy-request.json', 'agent-generation-poll.json')): yield Path(current) / name diff --git a/freevideo_engine/failure_details.py b/freevideo_engine/failure_details.py index 8ddf096f..7d8ec3d7 100644 --- a/freevideo_engine/failure_details.py +++ b/freevideo_engine/failure_details.py @@ -639,6 +639,14 @@ def _launcher_failure(value, *, zh=False, other_disk=False, disk_name='', can_re ('This folder holds an offline installation', '这个位置是离线包安装'), ('Open it from the launcher. For an automatic installation, go back to Setup and choose another folder.', '请在启动页直接打开这份安装。如需自动安装,请回到安装设置,换一个安装位置。')), + # offline_packages.inspect_archive and check_runtime_platform. + (r'The file is not a FreeVideo offline package', 'not-offline-package', + ('This is not a FreeVideo offline package', '这不是 FreeVideo 离线包'), + ('Choose a FreeVideo model package or Environment ZIP.', '请选择 FreeVideo 的模型包或运行环境包(ZIP 文件)。')), + (r'The Windows runtime package cannot be used on macOS', 'runtime-platform', + ('This Environment ZIP is for Windows', '这个运行环境包只能用于 Windows'), + ('Import only the model packages and continue; FreeVideo downloads the Mac environment automatically.', + '请只导入模型包并继续安装,FreeVideo 会自动下载 Mac 的运行环境。')), # Both carry ComfyUI's log, whose unrelated lines must not pick the card. (r'ComfyUI is still starting or could not load FreeVideo', 'comfy-nodes', ('ComfyUI did not load FreeVideo', 'ComfyUI 没有加载 FreeVideo'), diff --git a/freevideo_engine/generate.py b/freevideo_engine/generate.py index 1a84f8ae..3a97b5a7 100644 --- a/freevideo_engine/generate.py +++ b/freevideo_engine/generate.py @@ -30,6 +30,42 @@ from .resource_history import ResourceHistory +# A worker that exits after a CUDA error normally does so at once. One whose +# driver teardown never returns would otherwise hold the request forever, so +# the controller ends a worker that recorded this failure and is still running. +FAILED_WORKER_GRACE_SECONDS = 20. + + +class FailureReceipt: + """Whether the worker's metrics file records a request ended by a CUDA error.""" + def __init__(self, command): + self.path = self.stamp = None + self.found = False + self.since = time.time_ns() # A receipt left by an earlier request is not this worker's. + try: + request = json.loads(Path(command[command.index('--request') + 1]).read_text(encoding='utf-8')) + self.path = Path(request['metrics']) if request.get('metrics') else None + except (ValueError, IndexError, KeyError, OSError, TypeError): + pass + + def recorded(self): + if self.path is None or self.found: + return self.found + try: + status = self.path.stat() + stamp = (status.st_mtime_ns, status.st_size) + if stamp == self.stamp or status.st_mtime_ns < self.since: + return False + metrics = json.loads(self.path.read_text(encoding='utf-8')) + except (OSError, ValueError): + return False # Unreadable or half written: read again at the next poll. + self.stamp = stamp + failure = metrics.get('failure') if isinstance(metrics, dict) else None + self.found = (isinstance(failure, dict) and metrics.get('success') is False + and failure.get('kind') == 'cuda_error') + return self.found + + def child(command, env, descriptor, log, *, ram_budget_bytes=None, gpu=None, on_start=None, minimum_available_bytes=None, ram_budget_is_estimate=False): from .ram import ProcessMemory @@ -55,6 +91,8 @@ def child(command, env, descriptor, log, *, ram_budget_bytes=None, gpu=None, on_ process = processes.popen(command, env=env, stdout=stream, stderr=subprocess.STDOUT, pass_fds=(descriptor,), start_new_session=True) pending_error = None + receipt = FailureReceipt(command) + failed_since = None try: if on_start: on_start(process.pid) @@ -90,6 +128,14 @@ def child(command, env, descriptor, log, *, ram_budget_bytes=None, gpu=None, on_ (', '.join(pressure['enforced_reasons']), pressure['working_bytes']/2**30, ram_budget_bytes/2**30, pressure['available_bytes']/2**30, pressure['emergency_floor_bytes']/2**30)) + detail) + if receipt.recorded(): + failed_since = time.monotonic() if failed_since is None else failed_since + if time.monotonic() - failed_since >= FAILED_WORKER_GRACE_SECONDS: + telemetry['worker_unresponsive_after_failure_seconds'] = time.monotonic() - failed_since + telemetry['worker_exit'] = dict(exit_reading(None), memory=memory_sample()) + print(json.dumps({'event': 'worker_ended_after_failure', + 'seconds': telemetry['worker_unresponsive_after_failure_seconds']}), flush=True) + raise WorkerExit(None, log) try: process.wait(timeout=.5) except subprocess.TimeoutExpired: @@ -751,13 +797,14 @@ def interrupted(signum, frame): and (args.comfy_root or comfy_root()).resolve() == Path(machine['comfy_root']).resolve() and args.model_paths and args.model_paths.resolve() == Path(machine['model_paths']).resolve()) if same_inputs: + # The text cache needs only the encoder and its code: no git, no video models. + from .input_cache import InputCache + input_cache = InputCache(tuning.text_identity(machine, hardware)) + report['input_cache'].update(enabled=True, directory=str(input_cache.root)) identity = tuning.identity(machine, hardware) state, reason = tuning.load_state(identity) if active_loras: state, reason = None, 'LoRA uses its own model history; encoder cache remains reusable' - from .input_cache import InputCache - input_cache = InputCache(identity) - report['input_cache'].update(enabled=True, directory=str(input_cache.root)) report['tuning']['note'] = reason else: report['tuning']['note'] = 'Explicit model/encoder paths differ from the optimized installation' @@ -1207,8 +1254,9 @@ def retry(index, row, selected, decision): shutil.copyfile(destination.with_suffix('.engine.json'), artifacts / 'first-pass.engine.json') from .preview import upscale_reads report['preview']['upscale_reads'] = upscale_reads(request['cache'], report['video'], upscaler_checkpoint) - if not preview and any(row.get('keep_source_audio') for row in media.get('references', [])): - from .guide_audio import apply_to_report as keep_source_audio + from .guide_audio import apply_to_report as keep_source_audio, requested as keeps_source_audio + if keeps_source_audio(media): + # A preview gets the guide clip's sound too; its upscale reads the retained first pass, not this MP4. keep_source_audio(destination, media, report) report['success'] = True except BaseException as error: diff --git a/freevideo_engine/guide_audio.py b/freevideo_engine/guide_audio.py index 4633534d..e8bab46e 100644 --- a/freevideo_engine/guide_audio.py +++ b/freevideo_engine/guide_audio.py @@ -80,6 +80,11 @@ def replace_audio(video, guide, output, *, rate=48000, layout='stereo'): return seconds +def requested(media): + """True when the request asks to keep a guide clip's own audio (a preview included).""" + return any(row.get('guide') and row.get('keep_source_audio') for row in media.get('references', [])) + + def apply(video, media): """Replace ``video`` in place when the request's guide asks to keep its audio. Returns True if replaced.""" guide = next((row for row in media.get('references', []) if row.get('guide') and row.get('keep_source_audio')), None) diff --git a/freevideo_engine/hf_auth.py b/freevideo_engine/hf_auth.py index 0b43569d..42d5c3ba 100644 --- a/freevideo_engine/hf_auth.py +++ b/freevideo_engine/hf_auth.py @@ -8,14 +8,28 @@ 'Kept for this session only. You can continue without it.') HELP_ZH = ('可选:填写 Hugging Face 读取 Token,使用账号下载额度,减少匿名限流。' '可能更快,但不保证提速。仅本次打开有效,不填写也能继续。') +ERROR_EN = 'Enter a token that begins with hf_ and contains only letters and digits, or leave the field empty.' +ERROR_ZH = '令牌应以 hf_ 开头,只含英文字母和数字;也可留空。' -def validate(value): +def clean(value): + """Remove surrounding whitespace, paired quotes and an HF_TOKEN= prefix.""" if not isinstance(value, str): - raise ValueError('Enter a Hugging Face token beginning with hf_, or leave it empty.') + raise ValueError(ERROR_EN) value = value.strip() + if len(value) >= 2 and value[0] in ('"', "'") and value[-1] == value[0]: + value = value[1:-1].strip() + if value[:9].upper() == 'HF_TOKEN=': + value = value[9:].strip() + if len(value) >= 2 and value[0] in ('"', "'") and value[-1] == value[0]: + value = value[1:-1].strip() + return value + + +def validate(value): + value = clean(value) if value and not re.fullmatch(r'hf_[A-Za-z0-9]{1,4093}', value): - raise ValueError('Enter a Hugging Face token beginning with hf_, or leave it empty.') + raise ValueError(ERROR_EN) return value diff --git a/freevideo_engine/image_gen/image_models.json b/freevideo_engine/image_gen/image_models.json index 129d9c1f..797515d9 100644 --- a/freevideo_engine/image_gen/image_models.json +++ b/freevideo_engine/image_gen/image_models.json @@ -39,6 +39,629 @@ } ] } + ], + "music": [ + { + "file": "music/acestep-v15-xl-turbo-bf16/transformer/config.json", + "component": "music", + "bytes": 1303, + "sha256": "6b6753b2833681bb0ad3363be71cafaa2b0a751146bbeb4e8ae485a407fa5775", + "repo": "ACE-Step/acestep-v15-xl-turbo-diffusers", + "revision": "200ba991ae448051e14b0183157e35c2d27c9fb0", + "model_id": "ace-step-1.5-xl-turbo", + "sources": [ + { + "id": "official", + "url": "https://huggingface.co/ACE-Step/acestep-v15-xl-turbo-diffusers/resolve/200ba991ae448051e14b0183157e35c2d27c9fb0/transformer/config.json" + }, + { + "id": "hf-mirror", + "url": "https://hf-mirror.com/ACE-Step/acestep-v15-xl-turbo-diffusers/resolve/200ba991ae448051e14b0183157e35c2d27c9fb0/transformer/config.json" + }, + { + "id": "modelscope", + "url": "https://www.modelscope.cn/models/ACE-Step/acestep-v15-xl-turbo-diffusers/resolve/7f2c053810d7af9d2511487360543bca83ca119a/transformer/config.json" + } + ] + }, + { + "file": "music/acestep-v15-xl-turbo-bf16/transformer/diffusion_pytorch_model-00001-of-00002.safetensors", + "component": "music", + "bytes": 4986971904, + "sha256": "814723a63ba0b11dc4a7e8a5d43b8107c947fcdb5ab1e30c65469252272b35d7", + "repo": "ACE-Step/acestep-v15-xl-turbo-diffusers", + "revision": "200ba991ae448051e14b0183157e35c2d27c9fb0", + "model_id": "ace-step-1.5-xl-turbo", + "sources": [ + { + "id": "official", + "url": "https://huggingface.co/ACE-Step/acestep-v15-xl-turbo-diffusers/resolve/200ba991ae448051e14b0183157e35c2d27c9fb0/transformer/diffusion_pytorch_model-00001-of-00002.safetensors" + }, + { + "id": "hf-mirror", + "url": "https://hf-mirror.com/ACE-Step/acestep-v15-xl-turbo-diffusers/resolve/200ba991ae448051e14b0183157e35c2d27c9fb0/transformer/diffusion_pytorch_model-00001-of-00002.safetensors" + }, + { + "id": "modelscope", + "url": "https://www.modelscope.cn/models/ACE-Step/acestep-v15-xl-turbo-diffusers/resolve/7f2c053810d7af9d2511487360543bca83ca119a/transformer/diffusion_pytorch_model-00001-of-00002.safetensors" + } + ] + }, + { + "file": "music/acestep-v15-xl-turbo-bf16/transformer/diffusion_pytorch_model-00002-of-00002.safetensors", + "component": "music", + "bytes": 3350889672, + "sha256": "38ab6a8443be10b3c4bb0b784a5b78d471358d1b4982c139cb7893c455670b67", + "repo": "ACE-Step/acestep-v15-xl-turbo-diffusers", + "revision": "200ba991ae448051e14b0183157e35c2d27c9fb0", + "model_id": "ace-step-1.5-xl-turbo", + "sources": [ + { + "id": "official", + "url": "https://huggingface.co/ACE-Step/acestep-v15-xl-turbo-diffusers/resolve/200ba991ae448051e14b0183157e35c2d27c9fb0/transformer/diffusion_pytorch_model-00002-of-00002.safetensors" + }, + { + "id": "hf-mirror", + "url": "https://hf-mirror.com/ACE-Step/acestep-v15-xl-turbo-diffusers/resolve/200ba991ae448051e14b0183157e35c2d27c9fb0/transformer/diffusion_pytorch_model-00002-of-00002.safetensors" + }, + { + "id": "modelscope", + "url": "https://www.modelscope.cn/models/ACE-Step/acestep-v15-xl-turbo-diffusers/resolve/7f2c053810d7af9d2511487360543bca83ca119a/transformer/diffusion_pytorch_model-00002-of-00002.safetensors" + } + ] + }, + { + "file": "music/acestep-v15-xl-turbo-bf16/condition_encoder/config.json", + "component": "music", + "bytes": 502, + "sha256": "76e9eeba7b3248500cd0f5fb7ecd66129af0c9ea98ffb762bd8879753595233b", + "repo": "ACE-Step/acestep-v15-xl-turbo-diffusers", + "revision": "200ba991ae448051e14b0183157e35c2d27c9fb0", + "model_id": "ace-step-1.5-xl-turbo", + "sources": [ + { + "id": "official", + "url": "https://huggingface.co/ACE-Step/acestep-v15-xl-turbo-diffusers/resolve/200ba991ae448051e14b0183157e35c2d27c9fb0/condition_encoder/config.json" + }, + { + "id": "hf-mirror", + "url": "https://hf-mirror.com/ACE-Step/acestep-v15-xl-turbo-diffusers/resolve/200ba991ae448051e14b0183157e35c2d27c9fb0/condition_encoder/config.json" + }, + { + "id": "modelscope", + "url": "https://www.modelscope.cn/models/ACE-Step/acestep-v15-xl-turbo-diffusers/resolve/7f2c053810d7af9d2511487360543bca83ca119a/condition_encoder/config.json" + } + ] + }, + { + "file": "music/acestep-v15-xl-turbo-bf16/condition_encoder/diffusion_pytorch_model.safetensors", + "component": "music", + "bytes": 1218676312, + "sha256": "ca111ca95858c8d3c3678c986dda03df20778c0de009f57109c1052866d522b2", + "repo": "ACE-Step/acestep-v15-xl-turbo-diffusers", + "revision": "200ba991ae448051e14b0183157e35c2d27c9fb0", + "model_id": "ace-step-1.5-xl-turbo", + "sources": [ + { + "id": "official", + "url": "https://huggingface.co/ACE-Step/acestep-v15-xl-turbo-diffusers/resolve/200ba991ae448051e14b0183157e35c2d27c9fb0/condition_encoder/diffusion_pytorch_model.safetensors" + }, + { + "id": "hf-mirror", + "url": "https://hf-mirror.com/ACE-Step/acestep-v15-xl-turbo-diffusers/resolve/200ba991ae448051e14b0183157e35c2d27c9fb0/condition_encoder/diffusion_pytorch_model.safetensors" + }, + { + "id": "modelscope", + "url": "https://www.modelscope.cn/models/ACE-Step/acestep-v15-xl-turbo-diffusers/resolve/7f2c053810d7af9d2511487360543bca83ca119a/condition_encoder/diffusion_pytorch_model.safetensors" + } + ] + }, + { + "file": "music/acestep-v15-xl-turbo/config.json", + "component": "music", + "bytes": 2407, + "sha256": "e3665376d46477cbd3e76203cd4e241f11a6a6bb5f88464c49f3c6aae51a918b", + "repo": "ACE-Step/acestep-v15-xl-turbo", + "revision": "d4a0b288b83ebb7e25a8c0b32c573c22e134e8ee", + "model_id": "ace-step-1.5-xl-turbo", + "sources": [ + { + "id": "official", + "url": "https://huggingface.co/ACE-Step/acestep-v15-xl-turbo/resolve/d4a0b288b83ebb7e25a8c0b32c573c22e134e8ee/config.json" + }, + { + "id": "hf-mirror", + "url": "https://hf-mirror.com/ACE-Step/acestep-v15-xl-turbo/resolve/d4a0b288b83ebb7e25a8c0b32c573c22e134e8ee/config.json" + }, + { + "id": "modelscope", + "url": "https://www.modelscope.cn/models/ACE-Step/acestep-v15-xl-turbo/resolve/76db8b369a2f0b3d8b68a3258a6e60ca43753546/config.json" + } + ] + }, + { + "file": "music/acestep-v15-xl-turbo/model-00004-of-00004.safetensors", + "component": "music", + "bytes": 4988483464, + "sha256": "a868b05493c73808408f58b33a903371bb2846b53c8b19bd7476991bc8837259", + "repo": "ACE-Step/acestep-v15-xl-turbo", + "revision": "d4a0b288b83ebb7e25a8c0b32c573c22e134e8ee", + "model_id": "ace-step-1.5-xl-turbo", + "sources": [ + { + "id": "official", + "url": "https://huggingface.co/ACE-Step/acestep-v15-xl-turbo/resolve/d4a0b288b83ebb7e25a8c0b32c573c22e134e8ee/model-00004-of-00004.safetensors" + }, + { + "id": "hf-mirror", + "url": "https://hf-mirror.com/ACE-Step/acestep-v15-xl-turbo/resolve/d4a0b288b83ebb7e25a8c0b32c573c22e134e8ee/model-00004-of-00004.safetensors" + }, + { + "id": "modelscope", + "url": "https://www.modelscope.cn/models/ACE-Step/acestep-v15-xl-turbo/resolve/76db8b369a2f0b3d8b68a3258a6e60ca43753546/model-00004-of-00004.safetensors" + } + ] + }, + { + "file": "music/acestep-v15-xl-turbo/silence_latent.pt", + "component": "music", + "bytes": 3841215, + "sha256": "a778e9dd942f5e8b2c09c55370782d318834432b03dabbcdf70e6ed49ad6358b", + "repo": "ACE-Step/acestep-v15-xl-turbo", + "revision": "d4a0b288b83ebb7e25a8c0b32c573c22e134e8ee", + "model_id": "ace-step-1.5-xl-turbo", + "sources": [ + { + "id": "official", + "url": "https://huggingface.co/ACE-Step/acestep-v15-xl-turbo/resolve/d4a0b288b83ebb7e25a8c0b32c573c22e134e8ee/silence_latent.pt" + }, + { + "id": "hf-mirror", + "url": "https://hf-mirror.com/ACE-Step/acestep-v15-xl-turbo/resolve/d4a0b288b83ebb7e25a8c0b32c573c22e134e8ee/silence_latent.pt" + }, + { + "id": "modelscope", + "url": "https://www.modelscope.cn/models/ACE-Step/acestep-v15-xl-turbo/resolve/76db8b369a2f0b3d8b68a3258a6e60ca43753546/silence_latent.pt" + } + ] + }, + { + "file": "music/acestep-5Hz-lm-1.7B/added_tokens.json", + "component": "music", + "bytes": 2217787, + "sha256": "db08b66a515fb5d6acca0b3492d25bb44e0deda6241fc1113ac0679d40558c48", + "repo": "ACE-Step/Ace-Step1.5", + "revision": "19671f406d603126926c1b7e2adc169acbcade22", + "model_id": "ace-step-1.5-xl-turbo", + "sources": [ + { + "id": "official", + "url": "https://huggingface.co/ACE-Step/Ace-Step1.5/resolve/19671f406d603126926c1b7e2adc169acbcade22/acestep-5Hz-lm-1.7B/added_tokens.json" + }, + { + "id": "hf-mirror", + "url": "https://hf-mirror.com/ACE-Step/Ace-Step1.5/resolve/19671f406d603126926c1b7e2adc169acbcade22/acestep-5Hz-lm-1.7B/added_tokens.json" + }, + { + "id": "modelscope", + "url": "https://www.modelscope.cn/models/ACE-Step/Ace-Step1.5/resolve/dd20ee00556360ae895e308efe9d11daf08fe9d2/acestep-5Hz-lm-1.7B/added_tokens.json" + } + ] + }, + { + "file": "music/acestep-5Hz-lm-1.7B/chat_template.jinja", + "component": "music", + "bytes": 4168, + "sha256": "a55ee1b1660128b7098723e0abcd92caa0788061051c62d51cbe87d9cf1974d8", + "repo": "ACE-Step/Ace-Step1.5", + "revision": "19671f406d603126926c1b7e2adc169acbcade22", + "model_id": "ace-step-1.5-xl-turbo", + "sources": [ + { + "id": "official", + "url": "https://huggingface.co/ACE-Step/Ace-Step1.5/resolve/19671f406d603126926c1b7e2adc169acbcade22/acestep-5Hz-lm-1.7B/chat_template.jinja" + }, + { + "id": "hf-mirror", + "url": "https://hf-mirror.com/ACE-Step/Ace-Step1.5/resolve/19671f406d603126926c1b7e2adc169acbcade22/acestep-5Hz-lm-1.7B/chat_template.jinja" + }, + { + "id": "modelscope", + "url": "https://www.modelscope.cn/models/ACE-Step/Ace-Step1.5/resolve/dd20ee00556360ae895e308efe9d11daf08fe9d2/acestep-5Hz-lm-1.7B/chat_template.jinja" + } + ] + }, + { + "file": "music/acestep-5Hz-lm-1.7B/config.json", + "component": "music", + "bytes": 1385, + "sha256": "73515ca4bfd86a5b56f28bf680311d1911e159c85d7f7ceb2188e310007c9010", + "repo": "ACE-Step/Ace-Step1.5", + "revision": "19671f406d603126926c1b7e2adc169acbcade22", + "model_id": "ace-step-1.5-xl-turbo", + "sources": [ + { + "id": "official", + "url": "https://huggingface.co/ACE-Step/Ace-Step1.5/resolve/19671f406d603126926c1b7e2adc169acbcade22/acestep-5Hz-lm-1.7B/config.json" + }, + { + "id": "hf-mirror", + "url": "https://hf-mirror.com/ACE-Step/Ace-Step1.5/resolve/19671f406d603126926c1b7e2adc169acbcade22/acestep-5Hz-lm-1.7B/config.json" + }, + { + "id": "modelscope", + "url": "https://www.modelscope.cn/models/ACE-Step/Ace-Step1.5/resolve/dd20ee00556360ae895e308efe9d11daf08fe9d2/acestep-5Hz-lm-1.7B/config.json" + } + ] + }, + { + "file": "music/acestep-5Hz-lm-1.7B/merges.txt", + "component": "music", + "bytes": 1671853, + "sha256": "8831e4f1a044471340f7c0a83d7bd71306a5b867e95fd870f74d0c5308a904d5", + "repo": "ACE-Step/Ace-Step1.5", + "revision": "19671f406d603126926c1b7e2adc169acbcade22", + "model_id": "ace-step-1.5-xl-turbo", + "sources": [ + { + "id": "official", + "url": "https://huggingface.co/ACE-Step/Ace-Step1.5/resolve/19671f406d603126926c1b7e2adc169acbcade22/acestep-5Hz-lm-1.7B/merges.txt" + }, + { + "id": "hf-mirror", + "url": "https://hf-mirror.com/ACE-Step/Ace-Step1.5/resolve/19671f406d603126926c1b7e2adc169acbcade22/acestep-5Hz-lm-1.7B/merges.txt" + }, + { + "id": "modelscope", + "url": "https://www.modelscope.cn/models/ACE-Step/Ace-Step1.5/resolve/dd20ee00556360ae895e308efe9d11daf08fe9d2/acestep-5Hz-lm-1.7B/merges.txt" + } + ] + }, + { + "file": "music/acestep-5Hz-lm-1.7B/model.safetensors", + "component": "music", + "bytes": 3708521528, + "sha256": "f161689da73e5ecefa28ff780d51c2d92a00f056d021d7933c779ed5c6cd7db8", + "repo": "ACE-Step/Ace-Step1.5", + "revision": "19671f406d603126926c1b7e2adc169acbcade22", + "model_id": "ace-step-1.5-xl-turbo", + "sources": [ + { + "id": "official", + "url": "https://huggingface.co/ACE-Step/Ace-Step1.5/resolve/19671f406d603126926c1b7e2adc169acbcade22/acestep-5Hz-lm-1.7B/model.safetensors" + }, + { + "id": "hf-mirror", + "url": "https://hf-mirror.com/ACE-Step/Ace-Step1.5/resolve/19671f406d603126926c1b7e2adc169acbcade22/acestep-5Hz-lm-1.7B/model.safetensors" + }, + { + "id": "modelscope", + "url": "https://www.modelscope.cn/models/ACE-Step/Ace-Step1.5/resolve/dd20ee00556360ae895e308efe9d11daf08fe9d2/acestep-5Hz-lm-1.7B/model.safetensors" + } + ] + }, + { + "file": "music/acestep-5Hz-lm-1.7B/special_tokens_map.json", + "component": "music", + "bytes": 1824199, + "sha256": "76e233bfc357b0b03d7b6e6ba8e799244f358552575a7cdadb78d3d19106b298", + "repo": "ACE-Step/Ace-Step1.5", + "revision": "19671f406d603126926c1b7e2adc169acbcade22", + "model_id": "ace-step-1.5-xl-turbo", + "sources": [ + { + "id": "official", + "url": "https://huggingface.co/ACE-Step/Ace-Step1.5/resolve/19671f406d603126926c1b7e2adc169acbcade22/acestep-5Hz-lm-1.7B/special_tokens_map.json" + }, + { + "id": "hf-mirror", + "url": "https://hf-mirror.com/ACE-Step/Ace-Step1.5/resolve/19671f406d603126926c1b7e2adc169acbcade22/acestep-5Hz-lm-1.7B/special_tokens_map.json" + }, + { + "id": "modelscope", + "url": "https://www.modelscope.cn/models/ACE-Step/Ace-Step1.5/resolve/dd20ee00556360ae895e308efe9d11daf08fe9d2/acestep-5Hz-lm-1.7B/special_tokens_map.json" + } + ] + }, + { + "file": "music/acestep-5Hz-lm-1.7B/tokenizer.json", + "component": "music", + "bytes": 24321939, + "sha256": "35af56c3f5cb3ea2cc578aa28a8937770981d504f183ac5c8c38baf4bbd4af4d", + "repo": "ACE-Step/Ace-Step1.5", + "revision": "19671f406d603126926c1b7e2adc169acbcade22", + "model_id": "ace-step-1.5-xl-turbo", + "sources": [ + { + "id": "official", + "url": "https://huggingface.co/ACE-Step/Ace-Step1.5/resolve/19671f406d603126926c1b7e2adc169acbcade22/acestep-5Hz-lm-1.7B/tokenizer.json" + }, + { + "id": "hf-mirror", + "url": "https://hf-mirror.com/ACE-Step/Ace-Step1.5/resolve/19671f406d603126926c1b7e2adc169acbcade22/acestep-5Hz-lm-1.7B/tokenizer.json" + }, + { + "id": "modelscope", + "url": "https://www.modelscope.cn/models/ACE-Step/Ace-Step1.5/resolve/dd20ee00556360ae895e308efe9d11daf08fe9d2/acestep-5Hz-lm-1.7B/tokenizer.json" + } + ] + }, + { + "file": "music/acestep-5Hz-lm-1.7B/tokenizer_config.json", + "component": "music", + "bytes": 14072925, + "sha256": "6cd70cdd89425971794f5235562edcc608b0629a6c4686ae51a8b8c8b8ba5e95", + "repo": "ACE-Step/Ace-Step1.5", + "revision": "19671f406d603126926c1b7e2adc169acbcade22", + "model_id": "ace-step-1.5-xl-turbo", + "sources": [ + { + "id": "official", + "url": "https://huggingface.co/ACE-Step/Ace-Step1.5/resolve/19671f406d603126926c1b7e2adc169acbcade22/acestep-5Hz-lm-1.7B/tokenizer_config.json" + }, + { + "id": "hf-mirror", + "url": "https://hf-mirror.com/ACE-Step/Ace-Step1.5/resolve/19671f406d603126926c1b7e2adc169acbcade22/acestep-5Hz-lm-1.7B/tokenizer_config.json" + }, + { + "id": "modelscope", + "url": "https://www.modelscope.cn/models/ACE-Step/Ace-Step1.5/resolve/dd20ee00556360ae895e308efe9d11daf08fe9d2/acestep-5Hz-lm-1.7B/tokenizer_config.json" + } + ] + }, + { + "file": "music/acestep-5Hz-lm-1.7B/vocab.json", + "component": "music", + "bytes": 2776833, + "sha256": "ca10d7e9fb3ed18575dd1e277a2579c16d108e32f27439684afa0e10b1440910", + "repo": "ACE-Step/Ace-Step1.5", + "revision": "19671f406d603126926c1b7e2adc169acbcade22", + "model_id": "ace-step-1.5-xl-turbo", + "sources": [ + { + "id": "official", + "url": "https://huggingface.co/ACE-Step/Ace-Step1.5/resolve/19671f406d603126926c1b7e2adc169acbcade22/acestep-5Hz-lm-1.7B/vocab.json" + }, + { + "id": "hf-mirror", + "url": "https://hf-mirror.com/ACE-Step/Ace-Step1.5/resolve/19671f406d603126926c1b7e2adc169acbcade22/acestep-5Hz-lm-1.7B/vocab.json" + }, + { + "id": "modelscope", + "url": "https://www.modelscope.cn/models/ACE-Step/Ace-Step1.5/resolve/dd20ee00556360ae895e308efe9d11daf08fe9d2/acestep-5Hz-lm-1.7B/vocab.json" + } + ] + }, + { + "file": "music/Qwen3-Embedding-0.6B/added_tokens.json", + "component": "music", + "bytes": 707, + "sha256": "c0284b582e14987fbd3d5a2cb2bd139084371ed9acbae488829a1c900833c680", + "repo": "ACE-Step/Ace-Step1.5", + "revision": "19671f406d603126926c1b7e2adc169acbcade22", + "model_id": "ace-step-1.5-xl-turbo", + "sources": [ + { + "id": "official", + "url": "https://huggingface.co/ACE-Step/Ace-Step1.5/resolve/19671f406d603126926c1b7e2adc169acbcade22/Qwen3-Embedding-0.6B/added_tokens.json" + }, + { + "id": "hf-mirror", + "url": "https://hf-mirror.com/ACE-Step/Ace-Step1.5/resolve/19671f406d603126926c1b7e2adc169acbcade22/Qwen3-Embedding-0.6B/added_tokens.json" + }, + { + "id": "modelscope", + "url": "https://www.modelscope.cn/models/ACE-Step/Ace-Step1.5/resolve/dd20ee00556360ae895e308efe9d11daf08fe9d2/Qwen3-Embedding-0.6B/added_tokens.json" + } + ] + }, + { + "file": "music/Qwen3-Embedding-0.6B/config.json", + "component": "music", + "bytes": 1359, + "sha256": "bb23c1607cfe059a58d8f0196cf1cebb52082b1056b8e358a579da80a5759420", + "repo": "ACE-Step/Ace-Step1.5", + "revision": "19671f406d603126926c1b7e2adc169acbcade22", + "model_id": "ace-step-1.5-xl-turbo", + "sources": [ + { + "id": "official", + "url": "https://huggingface.co/ACE-Step/Ace-Step1.5/resolve/19671f406d603126926c1b7e2adc169acbcade22/Qwen3-Embedding-0.6B/config.json" + }, + { + "id": "hf-mirror", + "url": "https://hf-mirror.com/ACE-Step/Ace-Step1.5/resolve/19671f406d603126926c1b7e2adc169acbcade22/Qwen3-Embedding-0.6B/config.json" + }, + { + "id": "modelscope", + "url": "https://www.modelscope.cn/models/ACE-Step/Ace-Step1.5/resolve/dd20ee00556360ae895e308efe9d11daf08fe9d2/Qwen3-Embedding-0.6B/config.json" + } + ] + }, + { + "file": "music/Qwen3-Embedding-0.6B/merges.txt", + "component": "music", + "bytes": 1671853, + "sha256": "8831e4f1a044471340f7c0a83d7bd71306a5b867e95fd870f74d0c5308a904d5", + "repo": "ACE-Step/Ace-Step1.5", + "revision": "19671f406d603126926c1b7e2adc169acbcade22", + "model_id": "ace-step-1.5-xl-turbo", + "sources": [ + { + "id": "official", + "url": "https://huggingface.co/ACE-Step/Ace-Step1.5/resolve/19671f406d603126926c1b7e2adc169acbcade22/Qwen3-Embedding-0.6B/merges.txt" + }, + { + "id": "hf-mirror", + "url": "https://hf-mirror.com/ACE-Step/Ace-Step1.5/resolve/19671f406d603126926c1b7e2adc169acbcade22/Qwen3-Embedding-0.6B/merges.txt" + }, + { + "id": "modelscope", + "url": "https://www.modelscope.cn/models/ACE-Step/Ace-Step1.5/resolve/dd20ee00556360ae895e308efe9d11daf08fe9d2/Qwen3-Embedding-0.6B/merges.txt" + } + ] + }, + { + "file": "music/Qwen3-Embedding-0.6B/model.safetensors", + "component": "music", + "bytes": 1191586416, + "sha256": "0437e45c94563b09e13cb7a64478fc406947a93cb34a7e05870fc8dcd48e23fd", + "repo": "ACE-Step/Ace-Step1.5", + "revision": "19671f406d603126926c1b7e2adc169acbcade22", + "model_id": "ace-step-1.5-xl-turbo", + "sources": [ + { + "id": "official", + "url": "https://huggingface.co/ACE-Step/Ace-Step1.5/resolve/19671f406d603126926c1b7e2adc169acbcade22/Qwen3-Embedding-0.6B/model.safetensors" + }, + { + "id": "hf-mirror", + "url": "https://hf-mirror.com/ACE-Step/Ace-Step1.5/resolve/19671f406d603126926c1b7e2adc169acbcade22/Qwen3-Embedding-0.6B/model.safetensors" + }, + { + "id": "modelscope", + "url": "https://www.modelscope.cn/models/ACE-Step/Ace-Step1.5/resolve/dd20ee00556360ae895e308efe9d11daf08fe9d2/Qwen3-Embedding-0.6B/model.safetensors" + } + ] + }, + { + "file": "music/Qwen3-Embedding-0.6B/special_tokens_map.json", + "component": "music", + "bytes": 613, + "sha256": "76862e765266b85aa9459767e33cbaf13970f327a0e88d1c65846c2ddd3a1ecd", + "repo": "ACE-Step/Ace-Step1.5", + "revision": "19671f406d603126926c1b7e2adc169acbcade22", + "model_id": "ace-step-1.5-xl-turbo", + "sources": [ + { + "id": "official", + "url": "https://huggingface.co/ACE-Step/Ace-Step1.5/resolve/19671f406d603126926c1b7e2adc169acbcade22/Qwen3-Embedding-0.6B/special_tokens_map.json" + }, + { + "id": "hf-mirror", + "url": "https://hf-mirror.com/ACE-Step/Ace-Step1.5/resolve/19671f406d603126926c1b7e2adc169acbcade22/Qwen3-Embedding-0.6B/special_tokens_map.json" + }, + { + "id": "modelscope", + "url": "https://www.modelscope.cn/models/ACE-Step/Ace-Step1.5/resolve/dd20ee00556360ae895e308efe9d11daf08fe9d2/Qwen3-Embedding-0.6B/special_tokens_map.json" + } + ] + }, + { + "file": "music/Qwen3-Embedding-0.6B/tokenizer.json", + "component": "music", + "bytes": 11423705, + "sha256": "def76fb086971c7867b829c23a26261e38d9d74e02139253b38aeb9df8b4b50a", + "repo": "ACE-Step/Ace-Step1.5", + "revision": "19671f406d603126926c1b7e2adc169acbcade22", + "model_id": "ace-step-1.5-xl-turbo", + "sources": [ + { + "id": "official", + "url": "https://huggingface.co/ACE-Step/Ace-Step1.5/resolve/19671f406d603126926c1b7e2adc169acbcade22/Qwen3-Embedding-0.6B/tokenizer.json" + }, + { + "id": "hf-mirror", + "url": "https://hf-mirror.com/ACE-Step/Ace-Step1.5/resolve/19671f406d603126926c1b7e2adc169acbcade22/Qwen3-Embedding-0.6B/tokenizer.json" + }, + { + "id": "modelscope", + "url": "https://www.modelscope.cn/models/ACE-Step/Ace-Step1.5/resolve/dd20ee00556360ae895e308efe9d11daf08fe9d2/Qwen3-Embedding-0.6B/tokenizer.json" + } + ] + }, + { + "file": "music/Qwen3-Embedding-0.6B/tokenizer_config.json", + "component": "music", + "bytes": 5404, + "sha256": "443bfa629eb16387a12edbf92a76f6a6f10b2af3b53d87ba1550adfcf45f7fa0", + "repo": "ACE-Step/Ace-Step1.5", + "revision": "19671f406d603126926c1b7e2adc169acbcade22", + "model_id": "ace-step-1.5-xl-turbo", + "sources": [ + { + "id": "official", + "url": "https://huggingface.co/ACE-Step/Ace-Step1.5/resolve/19671f406d603126926c1b7e2adc169acbcade22/Qwen3-Embedding-0.6B/tokenizer_config.json" + }, + { + "id": "hf-mirror", + "url": "https://hf-mirror.com/ACE-Step/Ace-Step1.5/resolve/19671f406d603126926c1b7e2adc169acbcade22/Qwen3-Embedding-0.6B/tokenizer_config.json" + }, + { + "id": "modelscope", + "url": "https://www.modelscope.cn/models/ACE-Step/Ace-Step1.5/resolve/dd20ee00556360ae895e308efe9d11daf08fe9d2/Qwen3-Embedding-0.6B/tokenizer_config.json" + } + ] + }, + { + "file": "music/Qwen3-Embedding-0.6B/vocab.json", + "component": "music", + "bytes": 2776833, + "sha256": "ca10d7e9fb3ed18575dd1e277a2579c16d108e32f27439684afa0e10b1440910", + "repo": "ACE-Step/Ace-Step1.5", + "revision": "19671f406d603126926c1b7e2adc169acbcade22", + "model_id": "ace-step-1.5-xl-turbo", + "sources": [ + { + "id": "official", + "url": "https://huggingface.co/ACE-Step/Ace-Step1.5/resolve/19671f406d603126926c1b7e2adc169acbcade22/Qwen3-Embedding-0.6B/vocab.json" + }, + { + "id": "hf-mirror", + "url": "https://hf-mirror.com/ACE-Step/Ace-Step1.5/resolve/19671f406d603126926c1b7e2adc169acbcade22/Qwen3-Embedding-0.6B/vocab.json" + }, + { + "id": "modelscope", + "url": "https://www.modelscope.cn/models/ACE-Step/Ace-Step1.5/resolve/dd20ee00556360ae895e308efe9d11daf08fe9d2/Qwen3-Embedding-0.6B/vocab.json" + } + ] + }, + { + "file": "music/vae/config.json", + "component": "music", + "bytes": 425, + "sha256": "14e019904df567f26df750317a70e2bd08f9f8f3c40ff4a24c97d1cd3f20ccd2", + "repo": "ACE-Step/Ace-Step1.5", + "revision": "19671f406d603126926c1b7e2adc169acbcade22", + "model_id": "ace-step-1.5-xl-turbo", + "sources": [ + { + "id": "official", + "url": "https://huggingface.co/ACE-Step/Ace-Step1.5/resolve/19671f406d603126926c1b7e2adc169acbcade22/vae/config.json" + }, + { + "id": "hf-mirror", + "url": "https://hf-mirror.com/ACE-Step/Ace-Step1.5/resolve/19671f406d603126926c1b7e2adc169acbcade22/vae/config.json" + }, + { + "id": "modelscope", + "url": "https://www.modelscope.cn/models/ACE-Step/Ace-Step1.5/resolve/dd20ee00556360ae895e308efe9d11daf08fe9d2/vae/config.json" + } + ] + }, + { + "file": "music/vae/diffusion_pytorch_model.safetensors", + "component": "music", + "bytes": 337431388, + "sha256": "da17edb604c40deaf09e9b24974e590d1ca83a374070e5d0884cfa4bed9a99b0", + "repo": "ACE-Step/Ace-Step1.5", + "revision": "19671f406d603126926c1b7e2adc169acbcade22", + "model_id": "ace-step-1.5-xl-turbo", + "sources": [ + { + "id": "official", + "url": "https://huggingface.co/ACE-Step/Ace-Step1.5/resolve/19671f406d603126926c1b7e2adc169acbcade22/vae/diffusion_pytorch_model.safetensors" + }, + { + "id": "hf-mirror", + "url": "https://hf-mirror.com/ACE-Step/Ace-Step1.5/resolve/19671f406d603126926c1b7e2adc169acbcade22/vae/diffusion_pytorch_model.safetensors" + }, + { + "id": "modelscope", + "url": "https://www.modelscope.cn/models/ACE-Step/Ace-Step1.5/resolve/dd20ee00556360ae895e308efe9d11daf08fe9d2/vae/diffusion_pytorch_model.safetensors" + } + ] + } ] } } diff --git a/freevideo_engine/image_gen/image_pack.py b/freevideo_engine/image_gen/image_pack.py index 9e0b7e5f..fe60e97e 100644 --- a/freevideo_engine/image_gen/image_pack.py +++ b/freevideo_engine/image_gen/image_pack.py @@ -139,17 +139,18 @@ def validate_catalog(spec): or not isinstance(spec['revision'], str) or not spec['revision']): raise PackError('invalid_catalog') 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)): + # A group of its own holds only its own component; the image group holds neither. + own = {group: spec['groups'].get(group, []) for group in ('segment', 'music')} + found = image_rows + [row for found in own.values() for row in found] + if (not image_rows or len({r['file'] for r in found}) != len(found) or set(spec['groups']) - {'image', *own} + or any(row['component'] in own for row in image_rows) + or any(row['component'] != group for group, found in own.items() for row in found)): raise PackError('invalid_catalog') for row in found: _relative(row['file']) if (type(row['bytes']) is not int or row['bytes'] <= 0 or not re.fullmatch('[0-9a-f]{64}', row['sha256']) - or row['component'] not in ('dit', 'text_encoder', 'vae', 'text_encoder_config', 'processor', 'segment')): + or row['component'] not in ('dit', 'text_encoder', 'vae', 'text_encoder_config', 'processor', 'segment', 'music')): raise PackError('invalid_catalog') ids = [s['id'] for s in row['sources']] required = {'official', 'modelscope'} @@ -226,6 +227,11 @@ def _owned_intact(path, stamp): 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') +MUSIC_NOTICE_TEXT = (b'ACE-Step 1.5, Copyright (c) 2026 ACEStep. Licensed under the MIT License; see ' + b'freevideo_engine/licenses/ACE-Step-MIT.txt.\nhttps://github.com/ace-step/ACE-Step-1.5\n' + b'Qwen3-Embedding-0.6B, Copyright (c) Alibaba Cloud. Licensed under Apache-2.0.\n') +NOTICES = {'image': ('NOTICE.txt', NOTICE_TEXT), 'segment': ('SAM2-NOTICE.txt', SEGMENT_NOTICE_TEXT), + 'music': ('MUSIC-NOTICE.txt', MUSIC_NOTICE_TEXT)} def _notice(): @@ -235,7 +241,7 @@ def _notice(): 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() == (SEGMENT_NOTICE_TEXT if group == 'segment' else NOTICE_TEXT) + return path.is_file() and path.read_bytes() == NOTICES[group][1] except OSError: return False @@ -243,10 +249,10 @@ def _notice_ready(path, *, group='image'): 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, 'SAM2-NOTICE.txt' if group == 'segment' else 'NOTICE.txt') + notice = _path(folder, NOTICES[group][0]) if notice.exists(): return # present (ours or the user's own text): left as it is - notice.write_bytes(SEGMENT_NOTICE_TEXT if group == 'segment' else NOTICE_TEXT) + notice.write_bytes(NOTICES[group][1]) _own(folder, owned, notice) except Exception as error: logging.warning('[FreeVideo] Could not write the image model notice: %s', redacted_launcher_error(error)) @@ -382,7 +388,7 @@ def status(root=None, *, group='image'): 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, 'SAM2-NOTICE.txt' if group == 'segment' else 'NOTICE.txt'), group=group) + notice_ready = _notice_ready(_path(folder, NOTICES[group][0]), group=group) except PackError: notice_ready = False total = sum(r['bytes'] for r in files) @@ -650,7 +656,7 @@ def repair(root=None, progress=None, check=None, *, engine_root=None, group='ima def removal_plan(root=None, *, group='image'): folder = base(root) owned = _read(folder, OWNERSHIP) - notice = 'SAM2-NOTICE.txt' if group == 'segment' else 'NOTICE.txt' + notice = NOTICES[group][0] allowed = {notice} for row in rows(group=group): allowed.update((row['file'], row['file'] + '.partial', row['file'] + '.partial.source.json')) diff --git a/freevideo_engine/input_cache.py b/freevideo_engine/input_cache.py index 52a7145e..a7c0c4f1 100644 --- a/freevideo_engine/input_cache.py +++ b/freevideo_engine/input_cache.py @@ -6,6 +6,7 @@ import copy import hashlib import json +import os from pathlib import Path import shutil import time @@ -25,7 +26,7 @@ def digest(path): class InputCache: - """One receipt per input: no global index, eviction, or RAM-sized cache.""" + """One receipt per input; a hit refreshes its receipt and `prune` evicts the least recently used ones.""" def __init__(self, identity, *, root=None, kind='conditioning'): encoded = json.dumps(identity, sort_keys=True, separators=(',', ':')).encode() self.identity = hashlib.sha256(encoded).hexdigest() @@ -74,6 +75,10 @@ def get(self, key, destination): return None from .file_ops import publish publish(temporary, destination) + try: + os.utime(self.receipt(key)) + except OSError: + pass result = copy.deepcopy(entry['metrics']) # A receipt describes the original encoding. Keep its timings as # provenance, not work performed by this cache-hit request. @@ -127,4 +132,35 @@ def put(self, key, source, metrics, *, move=False): save(receipt, {'identity': self.identity, 'file': target.name, 'bytes': size, 'sha256': checksum, 'metrics': metrics, 'created_epoch': time.time()}, staging='receipt.') + self.prune() return True + + def prune(self, max_bytes=1 << 30, max_entries=4096): + """Drop the least recently used receipts of this kind, across identities, until under both limits. + + Only files this cache wrote are removed; unreadable receipts are left in place.""" + kind = self.root.parent + rows, owners = [], {} + try: + for receipt in kind.glob('*/*.json'): + try: + entry = json.loads(receipt.read_text(encoding='utf-8')) + if not isinstance(entry['file'], str) or Path(entry['file']).name != entry['file']: + continue + data = receipt.parent / entry['file'] + rows.append((receipt.stat().st_mtime, receipt, data)) + owners[data] = owners.get(data, 0) + 1 + except (OSError, ValueError, KeyError, TypeError): + continue + total = sum({data: data.stat().st_size for _, _, data in rows if data.exists()}.values()) + for _, receipt, data in sorted(rows, key=lambda row: row[0]): + if total <= max_bytes and len(rows) <= max_entries: + break + receipt.unlink() + rows = [row for row in rows if row[1] != receipt] + owners[data] -= 1 + if not owners[data] and data.exists(): + total -= data.stat().st_size + data.unlink() + except OSError as error: + self.note = type(error).__name__ + ': input cache pruning skipped' diff --git a/freevideo_engine/launcher/Main.qml b/freevideo_engine/launcher/Main.qml index 49a6e8cb..60d998e7 100644 --- a/freevideo_engine/launcher/Main.qml +++ b/freevideo_engine/launcher/Main.qml @@ -1208,6 +1208,7 @@ ApplicationWindow { FGroup { Layout.fillWidth: true; title: t("Hugging Face token · optional", "Hugging Face 令牌 · 可选") FField { Layout.fillWidth: true; echoMode: TextInput.Password; placeholderText: s.token_set ? t("Token set for this session", "本次已设置令牌") : "hf_…"; enabled: !s.busy; onEditingFinished: backend.edit("token", text) } + FText { objectName: "tokenError"; visible: s.field_error.key === "token"; text: visible ? s.field_error.message : ""; color: theme.danger; font.pixelSize: theme.micro; Layout.fillWidth: true } RowLayout { Layout.fillWidth: true; spacing: 12 FText { text: t("Use your account's download quota. Kept for this session only.", "使用账号下载额度,仅本次打开有效。"); color: theme.muted; font.pixelSize: theme.micro; Layout.fillWidth: true } @@ -1428,6 +1429,7 @@ ApplicationWindow { FGroup { Layout.fillWidth: true; title: t("ComfyUI address", "ComfyUI 地址") FField { objectName: "listenAddress"; Layout.fillWidth: true; value: s.form.url; enabled: !s.form.lan_access; onEditingFinished: backend.edit("url", text) } + FText { objectName: "addressError"; visible: s.field_error.key === "url"; text: visible ? s.field_error.message : ""; color: theme.danger; font.pixelSize: theme.micro; Layout.fillWidth: true } FText { objectName: "listenAddressHint"; text: s.lan_password_set ? t("With 0.0.0.0 or this computer's LAN IP, devices on the same network can open FreeVideo without the password.", "填 0.0.0.0 或本机的局域网 IP 后,同一网络中的设备不用密码就能打开 FreeVideo。") : t("With 0.0.0.0 or this computer's LAN IP, devices on the same network can open FreeVideo without the access link.", "填 0.0.0.0 或本机的局域网 IP 后,同一网络中的设备不需要访问链接就能打开 FreeVideo。"); color: theme.muted; font.pixelSize: theme.micro; Layout.fillWidth: true } FText { visible: Qt.platform.os === "windows"; text: t("Windows will ask for permission. Check \"Private networks\", then click \"Allow access\".", "Windows 会弹出确认框,请勾选“专用网络”,再点“允许访问”。"); color: theme.muted; font.pixelSize: theme.micro; Layout.fillWidth: true } } diff --git a/freevideo_engine/launcher_session.py b/freevideo_engine/launcher_session.py index f71147fe..865d8b34 100644 --- a/freevideo_engine/launcher_session.py +++ b/freevideo_engine/launcher_session.py @@ -13,7 +13,8 @@ from . import __version__ from . import disk_space -from .comfy_launcher_runtime import Controller, layout, local_url, loopback_url, managed_frontend, new_layout, queue_busy +from .comfy_launcher_runtime import (ENGINE_UPDATE_WAITS, Controller, layout, local_url, loopback_url, managed_frontend, + new_layout, queue_busy) from .desktop_runtime import launcher_root, materialize_source from .download_settings import Probe, read as download_preferences, speed_text from .launcher_settings import Store, default_language @@ -131,6 +132,7 @@ def __init__(self, source=None, *, controller=None, store=None, updater=None, sm self.page = 'comfy' self.token = '' self.error = '' + self.field_error = {} self.notice = '' self.action_notice = '' self.pending_update = None @@ -260,6 +262,10 @@ def remember(self, selection, persist=True): self.persist() def edit(self, key, value): + # A rejected entry stays explained until that field is accepted; the text is in the old + # interface language after a switch. + if self.field_error and key in (self.field_error.get('key'), 'language'): + self.field_error = {} if key in ('engine', 'comfy', 'destination', 'new_comfy', 'offline_runtime') and not self.can_manage(): self.explain_busy() return @@ -271,8 +277,12 @@ def edit(self, key, value): if key == 'token': if self.controller.busy: raise ValueError(self.t('Pause installation before changing the token.', '请先暂停安装,再修改 Token。')) - from .hf_auth import validate - self.token = validate(str(value)); return + from .hf_auth import ERROR_EN, ERROR_ZH, validate + try: + self.token = validate(str(value)) + except ValueError: + raise ValueError(self.t(ERROR_EN, ERROR_ZH)) from None + return if key not in self.form: return if self.controller.busy or self.importer.busy or self.cleaner.busy or self.archives_busy: @@ -281,7 +291,16 @@ def edit(self, key, value): if key in ('separate', 'repair', 'new_comfy', 'sampling_caches', 'lan_access', 'image_model', 'prompt_model', 'prompt_hint_dismissed'): value = bool(value) if key == 'url': - value = local_url(str(value)) + try: + value = local_url(str(value)) + except ValueError as exc: + if str(exc).startswith(('Invalid ComfyUI port', 'Port ')): + message = self.t('The ComfyUI port is not valid.', 'ComfyUI 端口无效。') + else: + message = self.t( + 'Enter localhost or an IP address of this computer, for example http://127.0.0.1:8188.', + '请输入 localhost 或本机的 IP 地址,例如 http://127.0.0.1:8188。') + raise ValueError(message) from None if key == 'environment_method' and value not in ('auto', 'manual'): raise ValueError('Unknown environment installation method') if key == 'environment_method' and value == 'manual' and not runtime_packages_supported(): @@ -1288,10 +1307,11 @@ def update(self, token='', source='launcher'): self.update_checked = time.monotonic() self.updater.run('check') - def _queue_busy(self): - """True or False once known for a server this launcher owns, None before.""" + def _queue_busy(self, any_server=False): + """True or False once known for a server this launcher owns, None before. An engine update + also waits for a server another program started: it moves the same installation.""" owns = getattr(self.controller, 'owns_server', None) - if not callable(owns) or owns() is not True: + if not any_server and (not callable(owns) or owns() is not True): return False now = time.monotonic() if (self.queue_thread is None or not self.queue_thread.is_alive()) and ( @@ -1365,7 +1385,7 @@ def _restart_when_idle(self): def _update_engine_when_idle(self): if self.task_busy(): return - busy = self._queue_busy() + busy = self._queue_busy(any_server=True) self.update_waiting = bool(busy) if busy is not False: return @@ -1772,6 +1792,13 @@ def tick(self): # Nothing new to download or approve: deploy, then restart ComfyUI. self.controller.run('install', True) row, busy = self.controller.state, self.controller.busy + if (self.engine_updating and not busy and row.get('status') == 'failed' + and str(row.get('error') or '').startswith(ENGINE_UPDATE_WAITS)): + # A video started after the check: keep waiting for it instead of showing a failure. + self.controller.state = row = dict(row, status='open', error=None, exception=None) + self.engine_updating = self.engine_autoinstall = False + self.update_intent = 'engine' + self.queue_state = None if self.engine_updating and not busy and not self.engine_autoinstall: self.engine_updating = False groups = row.get('task', {}).get('model_groups') or row.get('model_groups') or row.get('plan', {}).get('model_groups') @@ -2324,7 +2351,8 @@ def snapshot(self): and not self.form['offline_runtime'], zh)), video_model_guide=video_instructions(zh), adaln_dialog=self.adaln_dialog_view(error, update), - selected=bool(self.selected), can_repair=self.can_repair(), repair_offered=repair_offered, error=error, retry_kind=self.retry_kind(), notice=self.notice, + selected=bool(self.selected), can_repair=self.can_repair(), repair_offered=repair_offered, error=error, + field_error=dict(self.field_error or {}), retry_kind=self.retry_kind(), notice=self.notice, # Qt reads a list as an array; LEVELS is a tuple. compatibility=dict(self.compatibility, levels=list(self.compatibility.get('levels') or ())), report=dict(self.report), diff --git a/freevideo_engine/licenses/ACE-Step-MIT.txt b/freevideo_engine/licenses/ACE-Step-MIT.txt new file mode 100644 index 00000000..600451d4 --- /dev/null +++ b/freevideo_engine/licenses/ACE-Step-MIT.txt @@ -0,0 +1,21 @@ +MIT License + +Copyright (c) 2026 ACEStep + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. diff --git a/freevideo_engine/licenses/Vector-Quantize-PyTorch-MIT.txt b/freevideo_engine/licenses/Vector-Quantize-PyTorch-MIT.txt new file mode 100644 index 00000000..cad43213 --- /dev/null +++ b/freevideo_engine/licenses/Vector-Quantize-PyTorch-MIT.txt @@ -0,0 +1,21 @@ +MIT License + +Copyright (c) 2020 Phil Wang + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. diff --git a/freevideo_engine/macos_stages.py b/freevideo_engine/macos_stages.py index 25e77633..b5df0c92 100644 --- a/freevideo_engine/macos_stages.py +++ b/freevideo_engine/macos_stages.py @@ -176,7 +176,7 @@ def interrupted(signum, frame): elif value.get('refine_from') and reuse_preview(value, artifacts, result): stages = ['upscale', 'refinement', 'decode'] joined = False - retry = False + retry = metal_retry = avoid_mlx = False for stage in stages: if joined and stage in ('upscale', 'refinement'): continue @@ -185,8 +185,10 @@ def interrupted(signum, frame): request = artifacts / (stage + '-request.json') receipt = artifacts / (stage + '-result.json') while True: - save(request, dict(value, phase=stage, report=str(receipt), **( - {'memory_retry': True} if retry and stage in ('first-pass', 'refinement') else {}))) + retry_options = {'memory_retry': True} if retry and stage in ('first-pass', 'refinement') else {} + if retry_options and avoid_mlx: + retry_options['avoid_mlx'] = True + save(request, dict(value, phase=stage, report=str(receipt), **retry_options)) tick = time.monotonic() try: # Inherit stdout/stderr: progress and errors stay in the same @@ -196,8 +198,15 @@ def interrupted(signum, frame): pass_fds=(descriptor,), check=True) except subprocess.CalledProcessError: failed = read_receipt(receipt) - if retry or not memory_exhausted(stage, failed): + if not memory_exhausted(stage, failed): raise + metal = METAL_FAILURE in str(failed.get('error_message', '')) + if retry: + if not metal_retry or not metal or avoid_mlx: + raise + # A second Metal failure at the smallest partitions + # gets one fresh process using Torch attention. + avoid_mlx = True # Unified memory taken by other applications can leave a # sampling pass short of its plan, and Metal can fail a command # buffer without naming memory. Retry it once in a fresh @@ -205,6 +214,7 @@ def interrupted(signum, frame): # Keeping a saved first pass avoids changing its latents # with those new partitions. retry = True + metal_retry = metal rows = failed.get('sampling_passes') if (stage == 'first-pass' and failed.get('joined_sampling') is True and failed.get('sampling_phase') in ('upscale', 'refinement') @@ -215,16 +225,17 @@ def interrupted(signum, frame): and len(rows[0]['step_seconds']) == sampling['base_steps']): retain_joined_progress(result, failed) retained_joined = failed - metal = METAL_FAILURE in str(failed.get('error_message', '')) result.setdefault('memory_retries', []).append(dict(phase=stage, kind='metal_command_buffer' if metal else 'unified_memory', + attention='torch' if avoid_mlx else 'mlx', error_type=failed.get('error_type'), error_message=failed.get('error_message'), compute=failed.get('compute'), process_seconds=time.monotonic() - tick)) print(json.dumps(dict(event='resource_retry', retry=dict( failed_phase='refinement' if stage == 'refinement' or retained_joined else 'sampling', - kind='metal_command_buffer' if metal else 'unified_memory', attempt=2, max_attempts=2, + kind='metal_command_buffer' if metal else 'unified_memory', + attempt=3 if avoid_mlx else 2, max_attempts=3 if avoid_mlx else 2, reuse_first_pass=stage == 'refinement' or bool(retained_joined)))), flush=True) - if retained_joined: + if stage == 'first-pass' and retained_joined: break receipt.unlink(missing_ok=True) continue @@ -232,6 +243,9 @@ def interrupted(signum, frame): if receipt.is_file(): result['stages'][stage] = dict(read_receipt(receipt), process_seconds=time.monotonic() - tick) + if stage in ('first-pass', 'refinement'): + # A later pass's failure is not a consecutive failed retry. + metal_retry = False break if stage == 'first-pass' and retained_joined: continue @@ -297,7 +311,7 @@ def interrupted(signum, frame): save(path, result) -def stage_compute(resource, canvas, retry=False): +def stage_compute(resource, canvas, retry=False, avoid_mlx=False): """Partitions for this sampling stage from its live unified-memory allowance. The supervisor's policy fixes the reserve and allocator ceiling; the stage @@ -315,12 +329,15 @@ def stage_compute(resource, canvas, retry=False): shape = geometry(canvas['width'], canvas['height'], frames=canvas['frames']) tokens, frames = shape['video_tokens'] + 1024, shape['latent_frames'] result = plan(budget, tokens, frames=frames, allow_bounded=True, force_bounded=retry) - try: - import mlx.core as mx - mx.array([1.0]) # Fails early where Metal is unusable for MLX. - result['attention_impl'] = 'mlx' - except Exception: + if avoid_mlx: result['attention_impl'] = 'torch' + else: + try: + import mlx.core as mx + mx.array([1.0]) # Fails early where Metal is unusable for MLX. + result['attention_impl'] = 'mlx' + except Exception: + result['attention_impl'] = 'torch' result.update(fast_kernels=resource.get('fast_kernels', True), tokens=tokens, frames=frames) return result @@ -394,7 +411,8 @@ def capture_first(video): initial = torch.load(artifacts / 'upscaled.pt', map_location='cpu', weights_only=True) if refining else None print(json.dumps(dict(event='model_load_phase', phase='Loading native MPS model')), flush=True) retry = value.get('memory_retry') is True - compute = stage_compute(resource, sampling['second' if refining else 'first'], retry=retry) + compute = stage_compute(resource, sampling['second' if refining else 'first'], retry=retry, + avoid_mlx=value.get('avoid_mlx') is True) result['compute'] = compute if retry: result['memory_retry'] = True diff --git a/freevideo_engine/modern_launcher.py b/freevideo_engine/modern_launcher.py index e7815247..5aef4a38 100644 --- a/freevideo_engine/modern_launcher.py +++ b/freevideo_engine/modern_launcher.py @@ -143,7 +143,14 @@ def refresh(self): @Slot(str, 'QVariant') def edit(self, key, value): - self.invoke(session.edit, key, value) + def edit_field(): + try: + session.edit(key, value) + except ValueError as exc: + if key not in ('token', 'url'): + raise + session.field_error = dict(key=key, message=str(exc)) + self.invoke(edit_field) @Slot(str, bool) def action(self, name, accepted=False): diff --git a/freevideo_engine/music/__init__.py b/freevideo_engine/music/__init__.py new file mode 100644 index 00000000..b866c70b --- /dev/null +++ b/freevideo_engine/music/__init__.py @@ -0,0 +1 @@ +"""Instrumental text-to-music with ACE-Step 1.5 XL turbo and 1.7B LM.""" diff --git a/freevideo_engine/music/acestep_model.py b/freevideo_engine/music/acestep_model.py new file mode 100644 index 00000000..dddb3a8e --- /dev/null +++ b/freevideo_engine/music/acestep_model.py @@ -0,0 +1,1521 @@ +# Inference subset of ACE-Step 1.5 XL turbo (MIT; see licenses/ACE-Step-MIT.txt). +# Source: ACE-Step-1.5 checkpoints/acestep-v15-xl-turbo/modeling_acestep_v15_xl_turbo.py +# and acestep/models/common/configuration_acestep_v15.py, snapshot 2026-10-29. +# Upstream Apache/Qwen notices below are retained in addition to the MIT notice. +# Configuration copyright 2024 The Qwen team, Alibaba Group and HuggingFace Inc. +# Copyright 2025 The ACESTEO Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +import copy +import math +from typing import Callable, Optional + +import torch +import torch.nn.functional as F +from torch import nn +from .layers import linear, MusicMLP + +# Transformers imports (sorted by submodule, then alphabetically) +from transformers.cache_utils import Cache, DynamicCache, EncoderDecoderCache +from transformers.modeling_flash_attention_utils import FlashAttentionKwargs +from transformers.modeling_outputs import BaseModelOutput +from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel +from transformers.processing_utils import Unpack +from transformers.utils import can_return_tuple +from transformers.models.qwen3.modeling_qwen3 import ( + Qwen3MLP, + Qwen3RMSNorm, + Qwen3RotaryEmbedding, + apply_rotary_pos_emb, + eager_attention_forward, +) + +from transformers.configuration_utils import PretrainedConfig, layer_type_validation +from transformers.modeling_rope_utils import rope_config_validation + +class AceStepConfig(PretrainedConfig): + model_type = 'acestep' + keys_to_ignore_at_inference = ['past_key_values'] + def __init__( + self, + vocab_size=64003, + fsq_dim=2048, + fsq_input_levels=[8, 8, 8, 5, 5, 5], + fsq_input_num_quantizers=1, + hidden_size=2048, + intermediate_size=6144, + num_hidden_layers=24, + num_attention_heads=16, + num_key_value_heads=8, + head_dim=128, + hidden_act='silu', + max_position_embeddings=32768, + initializer_range=0.02, + rms_norm_eps=1e-06, + use_cache=True, + tie_word_embeddings=True, + rope_theta=1000000, + rope_scaling=None, + attention_bias=False, + use_sliding_window=True, + sliding_window=128, + layer_types=None, + attention_dropout=0.0, + num_lyric_encoder_hidden_layers=8, + audio_acoustic_hidden_dim=64, + pool_window_size=5, + text_hidden_dim=1024, + in_channels=192, + data_proportion=0.5, + timestep_mu=-0.4, + timestep_sigma=1.0, + timbre_hidden_dim=64, + num_timbre_encoder_hidden_layers=4, + timbre_fix_frame=750, + patch_size=2, + num_attention_pooler_hidden_layers=2, + num_audio_decoder_hidden_layers=24, + lyric_alignment_layers_config=None, + model_version='turbo', + **kwargs, + ): + self.max_position_embeddings = max_position_embeddings + self.hidden_size = hidden_size + self.intermediate_size = intermediate_size + self.num_hidden_layers = num_hidden_layers + self.num_attention_heads = num_attention_heads + self.use_sliding_window = use_sliding_window + self.sliding_window = sliding_window if self.use_sliding_window else None + self.text_hidden_dim = text_hidden_dim + self.num_lyric_encoder_hidden_layers = num_lyric_encoder_hidden_layers + self.patch_size = patch_size + self.audio_acoustic_hidden_dim = audio_acoustic_hidden_dim + self.pool_window_size = pool_window_size + self.in_channels = in_channels + self.data_proportion = data_proportion + self.timestep_mu = timestep_mu + self.timestep_sigma = timestep_sigma + self.fsq_dim = fsq_dim + self.fsq_input_levels = fsq_input_levels + self.fsq_input_num_quantizers = fsq_input_num_quantizers + self.timbre_hidden_dim = timbre_hidden_dim + self.num_timbre_encoder_hidden_layers = num_timbre_encoder_hidden_layers + self.timbre_fix_frame = timbre_fix_frame + self.num_attention_pooler_hidden_layers = num_attention_pooler_hidden_layers + self.num_audio_decoder_hidden_layers = num_audio_decoder_hidden_layers + self.lyric_alignment_layers_config = lyric_alignment_layers_config + self.vocab_size = vocab_size + if num_key_value_heads is None: + num_key_value_heads = num_attention_heads + self.num_key_value_heads = num_key_value_heads + self.head_dim = head_dim + self.hidden_act = hidden_act + self.initializer_range = initializer_range + self.rms_norm_eps = rms_norm_eps + self.use_cache = use_cache + self.rope_theta = rope_theta + self.rope_scaling = rope_scaling + self.attention_bias = attention_bias + self.attention_dropout = attention_dropout + self.model_version = model_version + if self.rope_scaling is not None and 'type' in self.rope_scaling: + self.rope_scaling['rope_type'] = self.rope_scaling['type'] + rope_config_validation(self) + self.layer_types = layer_types + if self.layer_types is None: + self.layer_types = ['sliding_attention' if bool((i + 1) % 2) else 'full_attention' for i in range(self.num_hidden_layers)] + layer_type_validation(self.layer_types) + super().__init__(tie_word_embeddings=tie_word_embeddings, **kwargs) + +def create_4d_mask( + seq_len: int, + dtype: torch.dtype, + device: torch.device, + attention_mask: Optional[torch.Tensor] = None, # [Batch, Seq_Len] + sliding_window: Optional[int] = None, + is_sliding_window: bool = False, + is_causal: bool = True, +) -> torch.Tensor: + """ + General 4D Attention Mask generator compatible with CPU/Mac/SDPA and Eager mode. + Supports use cases: + 1. Causal Full: is_causal=True, is_sliding_window=False (standard GPT) + 2. Causal Sliding: is_causal=True, is_sliding_window=True (Mistral/Qwen local window) + 3. Bidirectional Full: is_causal=False, is_sliding_window=False (BERT/Encoder) + 4. Bidirectional Sliding: is_causal=False, is_sliding_window=True (Longformer local) + + Returns: + [Batch, 1, Seq_Len, Seq_Len] additive mask (0.0 for keep, -inf for mask) + """ + # ------------------------------------------------------ + # 1. Construct basic geometry mask [Seq_Len, Seq_Len] + # ------------------------------------------------------ + + # Build index matrices + # i (Query): [0, 1, ..., L-1] + # j (Key): [0, 1, ..., L-1] + indices = torch.arange(seq_len, device=device) + # diff = i - j + diff = indices.unsqueeze(1) - indices.unsqueeze(0) + + # Initialize all True (all positions visible) + valid_mask = torch.ones((seq_len, seq_len), device=device, dtype=torch.bool) + + # (A) Handle causality (Causal) + if is_causal: + # i >= j => diff >= 0 + valid_mask = valid_mask & (diff >= 0) + + # (B) Handle sliding window + if is_sliding_window and sliding_window is not None: + if is_causal: + # Causal sliding: only attend to past window steps + # i - j <= window => diff <= window + # (diff >= 0 already handled above) + valid_mask = valid_mask & (diff <= sliding_window) + else: + # Bidirectional sliding: attend past and future window steps + # |i - j| <= window => abs(diff) <= sliding_window + valid_mask = valid_mask & (torch.abs(diff) <= sliding_window) + + # Expand dimensions to [1, 1, Seq_Len, Seq_Len] for broadcasting + valid_mask = valid_mask.unsqueeze(0).unsqueeze(0) + + # ------------------------------------------------------ + # 2. Apply padding mask (Key Masking) + # ------------------------------------------------------ + if attention_mask is not None: + # attention_mask shape: [Batch, Seq_Len] (1=valid, 0=padding) + # We want to mask out invalid keys (columns) + # Expand shape: [Batch, 1, 1, Seq_Len] + padding_mask_4d = attention_mask.view(attention_mask.shape[0], 1, 1, seq_len).to(torch.bool) + + # Broadcasting: Geometry Mask [1, 1, L, L] & Padding Mask [B, 1, 1, L] + # Result shape: [B, 1, L, L] + valid_mask = valid_mask & padding_mask_4d + + # ------------------------------------------------------ + # 3. Convert to additive mask + # ------------------------------------------------------ + # Get the minimal value for current dtype + min_dtype = torch.finfo(dtype).min + + # Create result tensor filled with -inf by default + mask_tensor = torch.full(valid_mask.shape, min_dtype, dtype=dtype, device=device) + + # Set valid positions to 0.0 + mask_tensor.masked_fill_(valid_mask, 0.0) + + return mask_tensor + + +def pack_sequences(hidden1: torch.Tensor, hidden2: torch.Tensor, mask1: torch.Tensor, mask2: torch.Tensor): + """ + Pack two sequences by concatenating and sorting them based on mask values. + + Args: + hidden1: First hidden states tensor of shape [B, L1, D] + hidden2: Second hidden states tensor of shape [B, L2, D] + mask1: First mask tensor of shape [B, L1] + mask2: Second mask tensor of shape [B, L2] + + Returns: + Tuple of (packed_hidden_states, new_mask) where: + - packed_hidden_states: Packed hidden states with valid tokens (mask=1) first, shape [B, L1+L2, D] + - new_mask: New mask tensor indicating valid positions, shape [B, L1+L2] + """ + # Step 1: Concatenate hidden states and masks along sequence dimension + hidden_cat = torch.cat([hidden1, hidden2], dim=1) # [B, L, D] + mask_cat = torch.cat([mask1, mask2], dim=1) # [B, L] + + B, L, D = hidden_cat.shape + + # Step 2: Sort indices so that mask values of 1 come before 0 + sort_idx = mask_cat.argsort(dim=1, descending=True, stable=True) # [B, L] + + # Step 3: Reorder hidden states using sorted indices + hidden_left = torch.gather(hidden_cat, 1, sort_idx.unsqueeze(-1).expand(B, L, D)) + + # Step 4: Create new mask based on valid sequence lengths + lengths = mask_cat.sum(dim=1) # [B] + new_mask = (torch.arange(L, dtype=torch.long, device=hidden_cat.device).unsqueeze(0) < lengths.unsqueeze(1)) + + return hidden_left, new_mask + + +class TimestepEmbedding(nn.Module): + """ + Timestep embedding module for diffusion models. + + Converts timestep values into high-dimensional embeddings using sinusoidal + positional encoding, followed by MLP layers. Used for conditioning diffusion + models on timestep information. + """ + def __init__( + self, + in_channels: int, + time_embed_dim: int, + scale: float = 1000, + ): + super().__init__() + + self.linear_1 = linear(in_channels, time_embed_dim, bias=True) + self.act1 = nn.SiLU() + self.linear_2 = linear(time_embed_dim, time_embed_dim, bias=True) + self.in_channels = in_channels + + self.act2 = nn.SiLU() + self.time_proj = linear(time_embed_dim, time_embed_dim * 6) + self.scale = scale + + def timestep_embedding(self, t, dim, max_period=10000): + """ + Create sinusoidal timestep embeddings. + + Args: + t: A 1-D tensor of N indices, one per batch element. These may be fractional. + dim: The dimension of the output embeddings. + max_period: Controls the minimum frequency of the embeddings. + + Returns: + An (N, D) tensor of positional embeddings. + """ + t = t * self.scale + half = dim // 2 + freqs = torch.exp( + -math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32) / half + ).to(device=t.device) + args = t[:, None].float() * freqs[None] + embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1) + if dim % 2: + embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1) + return embedding + + def forward(self, t): + t_freq = self.timestep_embedding(t, self.in_channels) + temb = self.linear_1(t_freq.to(t.dtype)) + temb = self.act1(temb) + temb = self.linear_2(temb) + timestep_proj = self.time_proj(self.act2(temb)).unflatten(1, (6, -1)) + return temb, timestep_proj + + +class AceStepAttention(nn.Module): + """ + Multi-headed attention module for AceStep model. + + Implements the attention mechanism from 'Attention Is All You Need' paper, + with support for both self-attention and cross-attention modes. Uses RMSNorm + for query and key normalization, and supports sliding window attention for + efficient long-sequence processing. + """ + + def __init__(self, config: AceStepConfig, layer_idx: int, is_cross_attention: bool = False, is_causal: bool = False): + super().__init__() + self.config = config + self.layer_idx = layer_idx + self.head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads) + self.num_key_value_groups = config.num_attention_heads // config.num_key_value_heads + self.scaling = self.head_dim**-0.5 + self.attention_dropout = config.attention_dropout + if is_cross_attention: + is_causal = False + self.is_causal = is_causal + self.is_cross_attention = is_cross_attention + + self.q_proj = linear(config.hidden_size, config.num_attention_heads * self.head_dim, bias=config.attention_bias) + self.k_proj = linear(config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias) + self.v_proj = linear(config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias) + self.o_proj = linear(config.num_attention_heads * self.head_dim, config.hidden_size, bias=config.attention_bias) + # Apply RMS normalization only on the head dimension (unlike OLMo) + self.q_norm = Qwen3RMSNorm(self.head_dim, eps=config.rms_norm_eps) + self.k_norm = Qwen3RMSNorm(self.head_dim, eps=config.rms_norm_eps) + self.attention_type = config.layer_types[layer_idx] + self.sliding_window = config.sliding_window if config.layer_types[layer_idx] == "sliding_attention" else None + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: Optional[torch.Tensor], + past_key_value: Optional[Cache] = None, + cache_position: Optional[torch.LongTensor] = None, + encoder_hidden_states: Optional[torch.Tensor] = None, + position_embeddings: tuple[torch.Tensor, torch.Tensor] = None, + output_attentions: Optional[bool] = False, + **kwargs: Unpack[FlashAttentionKwargs], + ) -> tuple[torch.Tensor, Optional[torch.Tensor], Optional[tuple[torch.Tensor]]]: + input_shape = hidden_states.shape[:-1] + hidden_shape = (*input_shape, -1, self.head_dim) + + # Project and normalize query states + query_states = self.q_norm(self.q_proj(hidden_states).view(hidden_shape)).transpose(1, 2) + + # Determine if this is cross-attention (requires encoder_hidden_states) + is_cross_attention = self.is_cross_attention and encoder_hidden_states is not None + + # Cross-attention path: attend to encoder hidden states + if is_cross_attention: + encoder_hidden_shape = (*encoder_hidden_states.shape[:-1], -1, self.head_dim) + if past_key_value is not None: + is_updated = past_key_value.is_updated.get(self.layer_idx) + # After the first generated token, we can reuse all key/value states from cache + curr_past_key_value = past_key_value.cross_attention_cache + + # Conditions for calculating key and value states + if not is_updated: + # Compute and cache K/V for the first time + key_states = self.k_norm(self.k_proj(encoder_hidden_states).view(encoder_hidden_shape)).transpose(1, 2) + value_states = self.v_proj(encoder_hidden_states).view(encoder_hidden_shape).transpose(1, 2) + # Update cache: save all key/value states to cache for fast auto-regressive generation + key_states, value_states = curr_past_key_value.update(key_states, value_states, self.layer_idx) + # Set flag that this layer's cross-attention cache is updated + past_key_value.is_updated[self.layer_idx] = True + else: + # Reuse cached key/value states for subsequent tokens + key_states = curr_past_key_value.layers[self.layer_idx].keys + value_states = curr_past_key_value.layers[self.layer_idx].values + else: + # No cache used, compute K/V directly + key_states = self.k_norm(self.k_proj(encoder_hidden_states).view(encoder_hidden_shape)).transpose(1, 2) + value_states = self.v_proj(encoder_hidden_states).view(encoder_hidden_shape).transpose(1, 2) + + # Self-attention path: attend to the same sequence + else: + # Project and normalize key/value states for self-attention + key_states = self.k_norm(self.k_proj(hidden_states).view(hidden_shape)).transpose(1, 2) + value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2) + # Apply rotary position embeddings (RoPE) if provided + if position_embeddings is not None: + cos, sin = position_embeddings + query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin) + + # Update cache for auto-regressive generation + if past_key_value is not None: + # Sin and cos are specific to RoPE models; cache_position needed for the static cache + cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position} + key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs) + + attention_interface: Callable = eager_attention_forward + if is_cross_attention and output_attentions: + attention_interface: Callable = eager_attention_forward + elif self.config._attn_implementation != "eager": + attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation] + + attn_output, attn_weights = attention_interface( + self, + query_states, + key_states, + value_states, + attention_mask, + dropout=0.0, + scaling=self.scaling, + sliding_window=self.sliding_window if not self.is_cross_attention else None, + **kwargs, + ) + + attn_output = attn_output.reshape(*input_shape, -1).contiguous() + attn_output = self.o_proj(attn_output) + return attn_output, attn_weights + + +class AceStepEncoderLayer(nn.Module): + """ + Encoder layer for AceStep model. + + Consists of self-attention and MLP (feed-forward) sub-layers with residual connections. + """ + + def __init__(self, config, layer_idx: int): + super().__init__() + self.hidden_size = config.hidden_size + self.config = config + self.layer_idx = layer_idx + + # Self-attention sub-layer + self.self_attn = AceStepAttention( + config=config, + layer_idx=layer_idx, + is_cross_attention=False, + is_causal=False, + ) + self.input_layernorm = Qwen3RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.post_attention_layernorm = Qwen3RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + + # MLP (feed-forward) sub-layer + self.mlp = MusicMLP(config) + self.attention_type = config.layer_types[layer_idx] + + def forward( + self, + hidden_states: torch.Tensor, + position_embeddings: tuple[torch.Tensor, torch.Tensor], + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + output_attentions: Optional[bool] = False, + **kwargs, + ) -> tuple[ + torch.FloatTensor, + Optional[tuple[torch.FloatTensor, torch.FloatTensor]], + ]: + # Self-attention with residual connection + residual = hidden_states + hidden_states = self.input_layernorm(hidden_states) + hidden_states, self_attn_weights = self.self_attn( + hidden_states=hidden_states, + position_embeddings=position_embeddings, + attention_mask=attention_mask, + position_ids=position_ids, + output_attentions=output_attentions, + # Encoders don't use cache + use_cache=False, + past_key_value=None, + **kwargs, + ) + hidden_states = residual + hidden_states + + # MLP with residual connection + residual = hidden_states + hidden_states = self.post_attention_layernorm(hidden_states) + hidden_states = self.mlp(hidden_states) + hidden_states = residual + hidden_states + + outputs = (hidden_states,) + + if output_attentions: + outputs += (self_attn_weights,) + + return outputs + + +class AceStepDiTLayer(nn.Module): + """ + DiT (Diffusion Transformer) layer for AceStep model. + + Implements a transformer layer with three main components: + 1. Self-attention with adaptive layer norm (AdaLN) + 2. Cross-attention (optional) for conditioning on encoder outputs + 3. Feed-forward MLP with adaptive layer norm + + Uses scale-shift modulation from timestep embeddings for adaptive normalization. + """ + def __init__(self, config: AceStepConfig, layer_idx: int, use_cross_attention: bool = True): + super().__init__() + + # 1. Self-attention sub-layer with adaptive normalization + self.self_attn_norm = Qwen3RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.self_attn = AceStepAttention(config=config, layer_idx=layer_idx) + + # 2. Cross-attention sub-layer (optional, for encoder conditioning) + self.use_cross_attention = use_cross_attention + if self.use_cross_attention: + self.cross_attn_norm = Qwen3RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.cross_attn = AceStepAttention(config=config, layer_idx=layer_idx, is_cross_attention=True) + + # 3. Feed-forward MLP sub-layer with adaptive normalization + self.mlp_norm = Qwen3RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.mlp = MusicMLP(config) + + # Scale-shift table for adaptive layer norm modulation (6 values: 3 for self-attn, 3 for MLP) + self.scale_shift_table = nn.Parameter(torch.randn(1, 6, config.hidden_size) / config.hidden_size**0.5) + self.attention_type = config.layer_types[layer_idx] + + def forward( + self, + hidden_states: torch.Tensor, + position_embeddings: tuple[torch.Tensor, torch.Tensor], + temb: torch.Tensor, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_value: Optional[EncoderDecoderCache] = None, + output_attentions: Optional[bool] = False, + use_cache: Optional[bool] = False, + cache_position: Optional[torch.LongTensor] = None, + encoder_hidden_states: Optional[torch.Tensor] = None, + encoder_attention_mask: Optional[torch.Tensor] = None, + **kwargs, + ) -> torch.Tensor: + + # Extract scale-shift parameters for adaptive layer norm from timestep embeddings + # 6 values: (shift_msa, scale_msa, gate_msa, c_shift_msa, c_scale_msa, c_gate_msa) + shift_msa, scale_msa, gate_msa, c_shift_msa, c_scale_msa, c_gate_msa = ( + self.scale_shift_table + temb + ).chunk(6, dim=1) + + # Step 1: Self-attention with adaptive layer norm (AdaLN) + # Apply adaptive normalization: norm(x) * (1 + scale) + shift + norm_hidden_states = (self.self_attn_norm(hidden_states) * (1 + scale_msa) + shift_msa).type_as(hidden_states) + attn_output, self_attn_weights = self.self_attn( + hidden_states=norm_hidden_states, + position_embeddings=position_embeddings, + attention_mask=attention_mask, + position_ids=position_ids, + output_attentions=output_attentions, + use_cache=False, + past_key_value=None, + **kwargs, + ) + # Apply gated residual connection: x = x + attn_output * gate + hidden_states = (hidden_states + attn_output * gate_msa).type_as(hidden_states) + + # Step 2: Cross-attention (if enabled) for conditioning on encoder outputs + if self.use_cross_attention: + norm_hidden_states = self.cross_attn_norm(hidden_states).type_as(hidden_states) + attn_output, cross_attn_weights = self.cross_attn( + hidden_states=norm_hidden_states, + encoder_hidden_states=encoder_hidden_states, + attention_mask=encoder_attention_mask, + past_key_value=past_key_value, + output_attentions=output_attentions, + use_cache=use_cache, + **kwargs, + ) + # Standard residual connection for cross-attention + hidden_states = hidden_states + attn_output + + # Step 3: Feed-forward (MLP) with adaptive layer norm + # Apply adaptive normalization for MLP: norm(x) * (1 + scale) + shift + norm_hidden_states = (self.mlp_norm(hidden_states) * (1 + c_scale_msa) + c_shift_msa).type_as(hidden_states) + ff_output = self.mlp(norm_hidden_states) + # Apply gated residual connection: x = x + mlp_output * gate + hidden_states = (hidden_states + ff_output * c_gate_msa).type_as(hidden_states) + + outputs = (hidden_states,) + if output_attentions: + outputs += (self_attn_weights, cross_attn_weights) + + return outputs + + +class AceStepPreTrainedModel(PreTrainedModel): + config_class = AceStepConfig + base_model_prefix = "model" + _no_split_modules = ["AceStepEncoderLayer", "AceStepDiTLayer"] + _skip_keys_device_placement = ["past_key_values"] + _supports_flash_attn_3 = True + _supports_flash_attn_2 = True + _supports_sdpa = True + _supports_flex_attn = True + _supports_cache_class = True + _supports_quantized_cache = True + _supports_static_cache = True + _supports_attention_backend = True + + def _init_weights(self, module): + """ + Initialize weights for different module types. + + TODO: Support separate initialization for encoders and decoders. + """ + std = self.config.initializer_range + if isinstance(module, nn.Linear): + module.weight.data.normal_(mean=0.0, std=std) + if module.bias is not None: + module.bias.data.zero_() + elif isinstance(module, nn.Embedding): + module.weight.data.normal_(mean=0.0, std=std) + if module.padding_idx is not None: + module.weight.data[module.padding_idx].zero_() + elif isinstance(module, Qwen3RMSNorm): + module.weight.data.fill_(1.0) + + +class AceStepLyricEncoder(AceStepPreTrainedModel): + """ + Encoder for processing lyric text embeddings. + + Encodes lyric text hidden states using a transformer encoder architecture + with bidirectional attention. Projects text embeddings to model hidden size + and processes them through multiple encoder layers. + """ + def __init__(self, config): + super().__init__(config) + + # Project text embeddings to model hidden size + self.embed_tokens = linear(config.text_hidden_dim, config.hidden_size) + self.norm = Qwen3RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.rotary_emb = Qwen3RotaryEmbedding(config=config) + + # Stack of encoder layers + self.layers = nn.ModuleList( + [AceStepEncoderLayer(config, layer_idx) for layer_idx in range(config.num_lyric_encoder_hidden_layers)] + ) + + # Initialize weights and apply final processing + self.post_init() + + @can_return_tuple + def forward( + self, + input_ids: Optional[torch.LongTensor] = None, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + **flash_attn_kwargs: Unpack[FlashAttentionKwargs], + ) -> BaseModelOutput: + output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions + output_hidden_states = ( + output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + ) + + assert input_ids is None, "Only `input_ids` is supported for the lyric encoder." + assert attention_mask is not None, "Attention mask must be provided for the lyric encoder." + assert inputs_embeds is not None, "Inputs embeddings must be provided for the lyric encoder." + + # Project input embeddings: N x T x text_hidden_dim -> N x T x hidden_size + inputs_embeds = self.embed_tokens(inputs_embeds) + # Cache position: only used for mask construction (not for actual caching) + cache_position = torch.arange(0, inputs_embeds.shape[1], device=inputs_embeds.device) + + # Positional IDs + if position_ids is None: + position_ids = cache_position.unsqueeze(0) + + # Attention masks + seq_len = inputs_embeds.shape[1] + dtype = inputs_embeds.dtype + device = inputs_embeds.device + + # 判断是否使用 Flash Attention 2 + is_flash_attn = (self.config._attn_implementation == "flash_attention_2") + + # 初始化 Mask 变量 + full_attn_mask = None + sliding_attn_mask = None + + if is_flash_attn: + # ------------------------------------------------------- + # 场景 A: Flash Attention 模式 + # ------------------------------------------------------- + # FA 不需要 4D Mask。 + # 如果有 padding mask (attention_mask [B, L]),直接传给它即可。 + # 如果没有 padding mask,传 None。 + # 滑动窗口逻辑由 Layer 内部传给 FA kernel 的 sliding_window 参数控制。 + full_attn_mask = attention_mask + + # 这里的逻辑是:如果配置启用了滑动窗口,FA 模式下我们也只需要传基础的 padding mask + # Layer 会自己决定是否调用带 sliding window 的 kernel + sliding_attn_mask = attention_mask if self.config.use_sliding_window else None + + else: + # ------------------------------------------------------- + # 场景 B: CPU / Mac / SDPA (Eager 模式) + # ------------------------------------------------------- + # 必须手动生成 4D Mask [B, 1, L, L] + + # 1. Full Attention (Bidirectional, Global) + # 对应原来的 create_causal_mask + bidirectional + full_attn_mask = create_4d_mask( + seq_len=seq_len, + dtype=dtype, + device=device, + attention_mask=attention_mask, # [B, L] + sliding_window=None, + is_sliding_window=False, + is_causal=False # <--- 关键:双向注意力 + ) + + # 2. Sliding Attention (Bidirectional, Local) + # 对应原来的 create_sliding_window... + bidirectional + if self.config.use_sliding_window: + sliding_attn_mask = create_4d_mask( + seq_len=seq_len, + dtype=dtype, + device=device, + attention_mask=attention_mask, # [B, L] + sliding_window=self.config.sliding_window, + is_sliding_window=True, # <--- 开启滑动窗口 + is_causal=False # <--- 关键:双向注意力 + ) + + # 构建 Mapping + self_attn_mask_mapping = { + "full_attention": full_attn_mask, + "sliding_attention": sliding_attn_mask, + } + + # Initialize hidden states with input embeddings + hidden_states = inputs_embeds + + # Create position embeddings to be shared across all layers + position_embeddings = self.rotary_emb(hidden_states, position_ids) + + # Pass through transformer layers + all_hidden_states = () if output_hidden_states else None + all_self_attns = () if output_attentions else None + + for layer_module in self.layers[: self.config.num_hidden_layers]: + if output_hidden_states: + all_hidden_states += (hidden_states,) + + layer_outputs = layer_module( + hidden_states, + position_embeddings, + self_attn_mask_mapping[layer_module.attention_type], + position_ids, + output_attentions, + **flash_attn_kwargs, + ) + + hidden_states = layer_outputs[0] + + if output_attentions: + all_self_attns += (layer_outputs[1],) + + hidden_states = self.norm(hidden_states) + + if output_hidden_states: + all_hidden_states += (hidden_states,) + + return BaseModelOutput( + last_hidden_state=hidden_states, + hidden_states=all_hidden_states, + attentions=all_self_attns, + ) + + +class AceStepTimbreEncoder(AceStepPreTrainedModel): + """ + Encoder for extracting timbre embeddings from reference audio. + + Processes packed reference audio acoustic features to extract timbre + representations. Uses a special token (CLS-like) to aggregate information + from the entire reference audio sequence. Outputs are unpacked back to + batch format for use in conditioning. + """ + def __init__(self, config): + super().__init__(config) + + # Project acoustic features to model hidden size + self.embed_tokens = linear(config.timbre_hidden_dim, config.hidden_size) + self.norm = Qwen3RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.rotary_emb = Qwen3RotaryEmbedding(config=config) + # Special token for aggregating timbre information (prepended to sequence) + self.special_token = nn.Parameter(torch.randn(1, 1, config.hidden_size)) + self.layers = nn.ModuleList( + [AceStepEncoderLayer(config, layer_idx) for layer_idx in range(config.num_timbre_encoder_hidden_layers)] + ) + + # Initialize weights and apply final processing + self.post_init() + + def unpack_timbre_embeddings(self, timbre_embs_packed, refer_audio_order_mask): + """ + Unpack packed timbre embeddings into batch format. + + Args: + timbre_embs_packed: Packed timbre embeddings of shape [N, d] + refer_audio_order_mask: Order mask indicating batch assignment for each packed embedding + + Returns: + Tuple of (unpacked_embeddings, mask): + - unpacked_embeddings: Unpacked embeddings of shape [B, max_count, d] + - new_mask: Mask indicating valid positions, shape [B, max_count] + """ + N, d = timbre_embs_packed.shape + device = timbre_embs_packed.device + dtype = timbre_embs_packed.dtype + + # Get batch size + B = int(refer_audio_order_mask.max().item() + 1) + + # Calculate element count and positions for each batch + counts = torch.bincount(refer_audio_order_mask, minlength=B) + max_count = counts.max().item() + + # Calculate positions within batch + sorted_indices = torch.argsort(refer_audio_order_mask * N + torch.arange(N, device=device), stable=True) + sorted_batch_ids = refer_audio_order_mask[sorted_indices] + + positions = torch.arange(N, device=device) + batch_starts = torch.cat([torch.tensor([0], device=device), + torch.cumsum(counts, dim=0)[:-1]]) + positions_in_sorted = positions - batch_starts[sorted_batch_ids] + + inverse_indices = torch.empty_like(sorted_indices) + inverse_indices[sorted_indices] = torch.arange(N, device=device) + positions_in_batch = positions_in_sorted[inverse_indices] + + # Use one-hot encoding and matrix multiplication (gradient-friendly approach) + # Create one-hot encoding + indices_2d = refer_audio_order_mask * max_count + positions_in_batch # (N,) + one_hot = F.one_hot(indices_2d, num_classes=B * max_count).to(dtype) # (N, B*max_count) + + # Rearrange using matrix multiplication + timbre_embs_flat = one_hot.t() @ timbre_embs_packed # (B*max_count, d) + timbre_embs_unpack = timbre_embs_flat.reshape(B, max_count, d) + + # Create mask indicating valid positions + mask_flat = (one_hot.sum(dim=0) > 0).long() # (B*max_count,) + new_mask = mask_flat.reshape(B, max_count) + + return timbre_embs_unpack, new_mask + + @can_return_tuple + def forward( + self, + refer_audio_acoustic_hidden_states_packed: Optional[torch.FloatTensor] = None, + refer_audio_order_mask: Optional[torch.LongTensor] = None, + attention_mask: Optional[torch.Tensor] = None, + **flash_attn_kwargs: Unpack[FlashAttentionKwargs], + ) -> BaseModelOutput: + inputs_embeds = refer_audio_acoustic_hidden_states_packed + # Project embeddings: N x T x timbre_hidden_dim -> N x T x hidden_size + inputs_embeds = self.embed_tokens(inputs_embeds) + # XL prepends the learned reference CLS token before its encoder. + inputs_embeds = torch.cat([self.special_token.expand(inputs_embeds.shape[0], 1, -1), inputs_embeds], dim=1) + if attention_mask is not None: + attention_mask = torch.cat([torch.ones(attention_mask.shape[0], 1, device=attention_mask.device, dtype=attention_mask.dtype), attention_mask], dim=1) + # Prepend special token for timbre aggregation (CLS-like token) + # inputs_embeds = torch.cat([self.special_token.expand(inputs_embeds.shape[0], 1, -1), inputs_embeds], dim=1) + # Cache position: only used for mask construction (not for actual caching) + cache_position = torch.arange(0, inputs_embeds.shape[1], device=inputs_embeds.device) + # Positional IDs + position_ids = cache_position.unsqueeze(0) + + seq_len = inputs_embeds.shape[1] + dtype = inputs_embeds.dtype + device = inputs_embeds.device + + # 判断是否使用 Flash Attention 2 + is_flash_attn = (self.config._attn_implementation == "flash_attention_2") + + # 初始化 Mask 变量 + full_attn_mask = None + sliding_attn_mask = None + + if is_flash_attn: + # ------------------------------------------------------- + # 场景 A: Flash Attention 模式 + # ------------------------------------------------------- + # FA 不需要 4D Mask。 + # 如果有 padding mask (attention_mask [B, L]),直接传给它即可。 + # 如果没有 padding mask,传 None。 + # 滑动窗口逻辑由 Layer 内部传给 FA kernel 的 sliding_window 参数控制。 + full_attn_mask = attention_mask + + # 这里的逻辑是:如果配置启用了滑动窗口,FA 模式下我们也只需要传基础的 padding mask + # Layer 会自己决定是否调用带 sliding window 的 kernel + sliding_attn_mask = attention_mask if self.config.use_sliding_window else None + + else: + # ------------------------------------------------------- + # 场景 B: CPU / Mac / SDPA (Eager 模式) + # ------------------------------------------------------- + # 必须手动生成 4D Mask [B, 1, L, L] + + # 1. Full Attention (Bidirectional, Global) + # 对应原来的 create_causal_mask + bidirectional + full_attn_mask = create_4d_mask( + seq_len=seq_len, + dtype=dtype, + device=device, + attention_mask=attention_mask, # [B, L] + sliding_window=None, + is_sliding_window=False, + is_causal=False # <--- 关键:双向注意力 + ) + + # 2. Sliding Attention (Bidirectional, Local) + # 对应原来的 create_sliding_window... + bidirectional + if self.config.use_sliding_window: + sliding_attn_mask = create_4d_mask( + seq_len=seq_len, + dtype=dtype, + device=device, + attention_mask=attention_mask, # [B, L] + sliding_window=self.config.sliding_window, + is_sliding_window=True, # <--- 开启滑动窗口 + is_causal=False # <--- 关键:双向注意力 + ) + + # 构建 Mapping + self_attn_mask_mapping = { + "full_attention": full_attn_mask, + "sliding_attention": sliding_attn_mask, + } + + # Initialize hidden states + hidden_states = inputs_embeds + + # Create position embeddings to be shared across all layers + position_embeddings = self.rotary_emb(hidden_states, position_ids) + + # Pass through transformer layers + for layer_module in self.layers[: self.config.num_hidden_layers]: + layer_outputs = layer_module( + hidden_states, + position_embeddings, + self_attn_mask_mapping[layer_module.attention_type], + position_ids, + **flash_attn_kwargs, + ) + + hidden_states = layer_outputs[0] + + hidden_states = self.norm(hidden_states) + # Extract special token output (first position) as timbre embedding: N x T x D -> N x D + hidden_states = hidden_states[:, 0, :] + # Unpack packed embeddings back to batch format + timbre_embs_unpack, timbre_embs_mask = self.unpack_timbre_embeddings(hidden_states, refer_audio_order_mask) + return timbre_embs_unpack, timbre_embs_mask + + +class Lambda(nn.Module): + """ + Wrapper module for arbitrary lambda functions. + + Allows using lambda functions in nn.Sequential by wrapping them in a Module. + Useful for simple transformations like transpose operations. + """ + def __init__(self, func): + super().__init__() + self.func = func + + def forward(self, x): + return self.func(x) + + +class AceStepDiTModel(AceStepPreTrainedModel): + """ + DiT (Diffusion Transformer) model for AceStep. + + Main diffusion model that generates audio latents conditioned on text, lyrics, + and timbre. Uses patch-based processing with transformer layers, timestep + conditioning, and cross-attention to encoder outputs. + """ + def __init__(self, config: AceStepConfig): + super().__init__(config) + # Rotary position embeddings for transformer layers + self.rotary_emb = Qwen3RotaryEmbedding(config) + # Stack of DiT transformer layers + self.layers = nn.ModuleList( + [AceStepDiTLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)] + ) + + in_channels = config.in_channels + inner_dim = config.hidden_size + patch_size = config.patch_size + self.patch_size = patch_size + + # Input projection: patch embedding using 1D convolution + # Converts sequence into patches for efficient processing + self.proj_in = nn.Sequential( + Lambda(lambda x: x.transpose(1, 2)), # [B, T, C] -> [B, C, T] + nn.Conv1d( + in_channels=in_channels, + out_channels=inner_dim, + kernel_size=patch_size, + stride=patch_size, + padding=0, + ), + Lambda(lambda x: x.transpose(1, 2)), # [B, C, T//patch_size] -> [B, T//patch_size, C] + ) + + # Timestep embeddings for diffusion conditioning + # Two embeddings: one for timestep t, one for timestep difference (t - r) + self.time_embed = TimestepEmbedding(in_channels=256, time_embed_dim=inner_dim) + self.time_embed_r = TimestepEmbedding(in_channels=256, time_embed_dim=inner_dim) + + # Project encoder hidden states to model dimension + condition_dim = getattr(config, "encoder_hidden_size", None) or config.hidden_size + self.condition_embedder = linear(condition_dim, inner_dim, bias=True) + + # Output normalization and projection + # Adaptive layer norm with scale-shift modulation, then de-patchify + self.norm_out = Qwen3RMSNorm(inner_dim, eps=config.rms_norm_eps) + self.proj_out = nn.Sequential( + Lambda(lambda x: x.transpose(1, 2)), # [B, T//patch_size, inner_dim] -> [B, inner_dim, T//patch_size] + nn.ConvTranspose1d( + in_channels=inner_dim, + out_channels=config.audio_acoustic_hidden_dim, + kernel_size=patch_size, + stride=patch_size, + padding=0, + ), + Lambda(lambda x: x.transpose(1, 2)), # [B, out_channels, T] -> [B, T, out_channels] + ) + # Scale-shift table for adaptive output normalization (2 values: shift, scale) + self.scale_shift_table = nn.Parameter(torch.randn(1, 2, inner_dim) / inner_dim**0.5) + + + def forward( + self, + hidden_states: torch.Tensor, + timestep: torch.Tensor, + timestep_r: torch.Tensor, + attention_mask: torch.Tensor, + encoder_hidden_states: torch.Tensor, + encoder_attention_mask: torch.Tensor, + context_latents: torch.Tensor, + use_cache: Optional[bool] = None, + past_key_values: Optional[EncoderDecoderCache] = None, + cache_position: Optional[torch.LongTensor] = None, + position_ids: Optional[torch.LongTensor] = None, + output_attentions: Optional[bool] = False, + return_hidden_states: int = None, + custom_layers_config: Optional[dict] = None, + enable_early_exit: bool = False, + **flash_attn_kwargs: Unpack[FlashAttentionKwargs], + ): + + use_cache = use_cache if use_cache is not None else self.config.use_cache + + if use_cache and past_key_values is None: + past_key_values = EncoderDecoderCache(DynamicCache(), DynamicCache()) + + # Compute timestep embeddings for diffusion conditioning + # Two embeddings: one for timestep t, one for timestep difference (t - r) + temb_t, timestep_proj_t = self.time_embed(timestep) + temb_r, timestep_proj_r = self.time_embed_r(timestep - timestep_r) + # Combine embeddings + temb = temb_t + temb_r + timestep_proj = timestep_proj_t + timestep_proj_r + + # Concatenate context latents (source latents + chunk masks) with hidden states + hidden_states = torch.cat([context_latents, hidden_states], dim=-1) + # Record original sequence length for later restoration after padding + original_seq_len = hidden_states.shape[1] + # Apply padding if sequence length is not divisible by patch_size + # This ensures proper patch extraction + pad_length = 0 + if hidden_states.shape[1] % self.patch_size != 0: + pad_length = self.patch_size - (hidden_states.shape[1] % self.patch_size) + hidden_states = F.pad(hidden_states, (0, 0, 0, pad_length), mode='constant', value=0) + + # Project input to patches and project encoder states + hidden_states = self.proj_in(hidden_states) + encoder_hidden_states = self.condition_embedder(encoder_hidden_states) + + # Cache positions + if cache_position is None: + past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0 + cache_position = torch.arange( + past_seen_tokens, past_seen_tokens + hidden_states.shape[1], device=hidden_states.device + ) + + # Position IDs + if position_ids is None: + position_ids = cache_position.unsqueeze(0) + + + seq_len = hidden_states.shape[1] + encoder_seq_len = encoder_hidden_states.shape[1] + dtype = hidden_states.dtype + device = hidden_states.device + + # 判断是否使用 Flash Attention 2 + is_flash_attn = (self.config._attn_implementation == "flash_attention_2") + + # 初始化 Mask 变量 + full_attn_mask = None + sliding_attn_mask = None + encoder_attention_mask = None + attention_mask = None + if is_flash_attn: + # ------------------------------------------------------- + # 场景 A: Flash Attention 模式 + # ------------------------------------------------------- + # FA 不需要 4D Mask。 + # 如果有 padding mask (attention_mask [B, L]),直接传给它即可。 + # 如果没有 padding mask,传 None。 + # 滑动窗口逻辑由 Layer 内部传给 FA kernel 的 sliding_window 参数控制。 + full_attn_mask = attention_mask + + # 这里的逻辑是:如果配置启用了滑动窗口,FA 模式下我们也只需要传基础的 padding mask + # Layer 会自己决定是否调用带 sliding window 的 kernel + sliding_attn_mask = attention_mask if self.config.use_sliding_window else None + + else: + # ------------------------------------------------------- + # 场景 B: CPU / Mac / SDPA (Eager 模式) + # ------------------------------------------------------- + # 必须手动生成 4D Mask [B, 1, L, L] + + # 1. Full Attention (Bidirectional, Global) + # 对应原来的 create_causal_mask + bidirectional + full_attn_mask = create_4d_mask( + seq_len=seq_len, + dtype=dtype, + device=device, + attention_mask=attention_mask, # [B, L] + sliding_window=None, + is_sliding_window=False, + is_causal=False # <--- 关键:双向注意力 + ) + max_len = max(seq_len, encoder_seq_len) + + encoder_attention_mask = create_4d_mask( + seq_len=max_len, + dtype=dtype, + device=device, + attention_mask=attention_mask, # [B, L] + sliding_window=None, + is_sliding_window=False, + is_causal=False # <--- 关键:双向注意力 + ) + encoder_attention_mask = encoder_attention_mask[:, :, :seq_len, :encoder_seq_len] + # 2. Sliding Attention (Bidirectional, Local) + # 对应原来的 create_sliding_window... + bidirectional + if self.config.use_sliding_window: + sliding_attn_mask = create_4d_mask( + seq_len=seq_len, + dtype=dtype, + device=device, + attention_mask=attention_mask, # [B, L] + sliding_window=self.config.sliding_window, + is_sliding_window=True, # <--- 开启滑动窗口 + is_causal=False # <--- 关键:双向注意力 + ) + + # 构建 Mapping + self_attn_mask_mapping = { + "full_attention": full_attn_mask, + "sliding_attention": sliding_attn_mask, + "encoder_attention_mask": encoder_attention_mask, + } + + # Create position embeddings to be shared across all decoder layers + position_embeddings = self.rotary_emb(hidden_states, position_ids) + all_cross_attentions = () if output_attentions else None + + # Handle early exit for custom layer configurations + max_needed_layer = float('inf') + if custom_layers_config is not None and enable_early_exit: + max_needed_layer = max(custom_layers_config.keys()) + # Force output_attentions to True when early exit is enabled for attention extraction + output_attentions = True + if all_cross_attentions is None: + all_cross_attentions = () + + # Process through transformer layers + for index_block, layer_module in enumerate(self.layers): + + layer_outputs = layer_module( + hidden_states, + position_embeddings, + timestep_proj, + self_attn_mask_mapping[layer_module.attention_type], + position_ids, + past_key_values, + output_attentions, + use_cache, + cache_position, + encoder_hidden_states, + self_attn_mask_mapping["encoder_attention_mask"], + **flash_attn_kwargs, + ) + hidden_states = layer_outputs[0] + + if output_attentions and self.layers[index_block].use_cross_attention: + # layer_outputs structure: (hidden_states, self_attn_weights, cross_attn_weights) + # Extract the last element which is cross_attn_weights + if len(layer_outputs) >= 3: + all_cross_attentions += (layer_outputs[2],) + + if return_hidden_states: + return hidden_states + + # Extract scale-shift parameters for adaptive output normalization + shift, scale = (self.scale_shift_table + temb.unsqueeze(1)).chunk(2, dim=1) + shift = shift.to(hidden_states.device) + scale = scale.to(hidden_states.device) + + # Apply adaptive layer norm: norm(x) * (1 + scale) + shift + hidden_states = (self.norm_out(hidden_states) * (1 + scale) + shift).type_as(hidden_states) + # Project output: de-patchify back to original sequence format + hidden_states = self.proj_out(hidden_states) + + # Crop back to original sequence length to ensure exact length match (remove padding) + hidden_states = hidden_states[:, :original_seq_len, :] + + outputs = (hidden_states, past_key_values) + + if output_attentions: + outputs += (all_cross_attentions,) + return outputs + + +class AceStepConditionEncoder(AceStepPreTrainedModel): + """ + Condition encoder for AceStep model. + + Encodes multiple conditioning inputs (text, lyrics, timbre) and packs them + into a single sequence for cross-attention in the diffusion model. Handles + projection, encoding, and sequence packing. + """ + def __init__(self, config: AceStepConfig): + super().__init__(config) + self.config = config + # Project text embeddings to model hidden size + self.text_projector = linear(config.text_hidden_dim, config.hidden_size, bias=False) + # Encoder for lyric text + self.lyric_encoder = AceStepLyricEncoder(config) + # Encoder for timbre from reference audio + self.timbre_encoder = AceStepTimbreEncoder(config) + + def forward( + self, + # Text inputs + text_hidden_states: Optional[torch.FloatTensor] = None, + text_attention_mask: Optional[torch.Tensor] = None, + # Lyric inputs + lyric_hidden_states: Optional[torch.LongTensor] = None, + lyric_attention_mask: Optional[torch.Tensor] = None, + # Reference audio for timbre + refer_audio_acoustic_hidden_states_packed: Optional[torch.Tensor] = None, + refer_audio_order_mask: Optional[torch.LongTensor] = None, + ): + # Project and encode text + text_hidden_states = self.text_projector(text_hidden_states) + # Encode lyrics + lyric_encoder_outputs = self.lyric_encoder( + inputs_embeds=lyric_hidden_states, + attention_mask=lyric_attention_mask, + ) + lyric_hidden_states = lyric_encoder_outputs.last_hidden_state + # Encode timbre from reference audio + timbre_embs_unpack, timbre_embs_mask = self.timbre_encoder(refer_audio_acoustic_hidden_states_packed, refer_audio_order_mask) + + # Pack sequences: combine lyrics and timbre, then add text + # This creates a single sequence with all conditioning information + encoder_hidden_states, encoder_attention_mask = pack_sequences(lyric_hidden_states, timbre_embs_unpack, lyric_attention_mask, timbre_embs_mask) + encoder_hidden_states, encoder_attention_mask = pack_sequences(encoder_hidden_states, text_hidden_states, encoder_attention_mask, text_attention_mask) + return encoder_hidden_states, encoder_attention_mask + + + +# Inverse FSQ arithmetic: vector-quantize-pytorch (MIT, Copyright Phil Wang). +# See licenses/Vector-Quantize-PyTorch-MIT.txt. No package import is required. +class AudioCodeQuantizer(nn.Module): + """ResidualFSQ inverse for one quantizer, without encoder-only packages. + + Match vector_quantize_pytorch ResidualFSQ.get_output_from_indices: + build the symmetry-preserving codebook in CPU float32, then cast the + buffer along with weights, index, multiply scales, sum and project. + """ + def __init__(self, config): + super().__init__() + if config.fsq_input_num_quantizers != 1: + raise ValueError("only the published single-quantizer checkpoint is supported") + levels = config.fsq_input_levels + with torch.device("cpu"): + level = torch.tensor(levels, dtype=torch.int32) + basis = torch.cumprod(torch.tensor([1] + levels[:-1]), 0, dtype=torch.int32) + ids = torch.arange(math.prod(levels)).unsqueeze(-1) + codebook = ((ids // basis) % level) * (2. / (level - 1).float()) - 1. + scales = torch.tensor(levels).float() ** 0 + self.register_buffer("codebook", codebook, persistent=False) + self.register_buffer("scales", scales, persistent=False) + self.project_out = linear(len(levels), config.fsq_dim) + + def get_output_from_indices(self, indices): + codes = self.codebook[indices[..., 0]].unsqueeze(0) + codes = codes * self.scales.reshape(1, 1, 1, -1) + return self.project_out(codes.sum(dim=0)) + + +class AudioTokenDetokenizer(AceStepPreTrainedModel): + """ + Audio token detokenizer module. + + Converts quantized audio tokens back to continuous acoustic representations. + Expands each token into multiple patches using special tokens, processes them + through encoder layers, and projects to acoustic hidden dimension. + """ + def __init__(self, config): + super().__init__(config) + self.config = config + self.embed_tokens = linear(config.hidden_size, config.hidden_size) + self.norm = Qwen3RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.rotary_emb = Qwen3RotaryEmbedding(config=config) + self.gradient_checkpointing = False + # Special tokens for expanding each quantized token into patches + self.special_tokens = nn.Parameter(torch.randn(1, config.pool_window_size, config.hidden_size) * 0.02) + self.layers = nn.ModuleList( + [AceStepEncoderLayer(config, layer_idx) for layer_idx in range(config.num_attention_pooler_hidden_layers)] + ) + # Project back to acoustic hidden dimension + self.proj_out = linear(config.hidden_size, config.audio_acoustic_hidden_dim) + + # Initialize weights and apply final processing + self.post_init() + + @can_return_tuple + def forward(self, + x, + attention_mask: Optional[torch.Tensor] = None, + **flash_attn_kwargs: Unpack[FlashAttentionKwargs], + ) -> BaseModelOutput: + B, T, D = x.shape + x = self.embed_tokens(x) + # Expand and add special tokens: N x T x D -> N x T x P x D + # Each token is expanded into pool_window_size patches + x = x.unsqueeze(2) # N x T x 1 x D + x = x.repeat(1, 1, self.config.pool_window_size, 1) # N x T x P x D + # Add learnable special tokens to each patch + special_tokens = self.special_tokens.expand(B, T, -1, -1) + x = x + special_tokens + # Reshape for processing: (batch * time) x patches x hidden + x = x.reshape(B * T, self.config.pool_window_size, -1) + + # Cache position: only used for mask construction + cache_position = torch.arange(0, x.shape[1], device=x.device) + # Positional IDs + position_ids = cache_position.unsqueeze(0) + + # Initialize hidden states + hidden_states = x + + # Create position embeddings to be shared across all layers + position_embeddings = self.rotary_emb(hidden_states, position_ids) + + seq_len = x.shape[1] + dtype = x.dtype + device = x.device + + # 判断是否使用 Flash Attention 2 + is_flash_attn = (self.config._attn_implementation == "flash_attention_2") + + # 初始化 Mask 变量 + full_attn_mask = None + sliding_attn_mask = None + + if is_flash_attn: + # ------------------------------------------------------- + # 场景 A: Flash Attention 模式 + # ------------------------------------------------------- + # FA 不需要 4D Mask。 + # 如果有 padding mask (attention_mask [B, L]),直接传给它即可。 + # 如果没有 padding mask,传 None。 + # 滑动窗口逻辑由 Layer 内部传给 FA kernel 的 sliding_window 参数控制。 + full_attn_mask = attention_mask + + # 这里的逻辑是:如果配置启用了滑动窗口,FA 模式下我们也只需要传基础的 padding mask + # Layer 会自己决定是否调用带 sliding window 的 kernel + sliding_attn_mask = attention_mask if self.config.use_sliding_window else None + + else: + # ------------------------------------------------------- + # 场景 B: CPU / Mac / SDPA (Eager 模式) + # ------------------------------------------------------- + # 必须手动生成 4D Mask [B, 1, L, L] + + # 1. Full Attention (Bidirectional, Global) + # 对应原来的 create_causal_mask + bidirectional + full_attn_mask = create_4d_mask( + seq_len=seq_len, + dtype=dtype, + device=device, + attention_mask=attention_mask, # [B, L] + sliding_window=None, + is_sliding_window=False, + is_causal=False # <--- 关键:双向注意力 + ) + + # 2. Sliding Attention (Bidirectional, Local) + # 对应原来的 create_sliding_window... + bidirectional + if self.config.use_sliding_window: + sliding_attn_mask = create_4d_mask( + seq_len=seq_len, + dtype=dtype, + device=device, + attention_mask=attention_mask, # [B, L] + sliding_window=self.config.sliding_window, + is_sliding_window=True, # <--- 开启滑动窗口 + is_causal=False # <--- 关键:双向注意力 + ) + + # 构建 Mapping + self_attn_mask_mapping = { + "full_attention": full_attn_mask, + "sliding_attention": sliding_attn_mask, + } + + for layer_module in self.layers: + layer_outputs = layer_module( + hidden_states, + position_embeddings, + attention_mask=self_attn_mask_mapping[layer_module.attention_type], + **flash_attn_kwargs, + ) + + hidden_states = layer_outputs[0] + + hidden_states = self.norm(hidden_states) + + hidden_states = self.proj_out(hidden_states) + + hidden_states = hidden_states.reshape(B, T * self.config.pool_window_size, -1) + return hidden_states + + +class AceStepConditionGenerationModel(AceStepPreTrainedModel): + """Instrumental text2music with LM hints and the default eight-step XL sampler.""" + + def __init__(self, config): + super().__init__(config) + self.decoder = AceStepDiTModel(config) + encoder_config = copy.deepcopy(config) + for field in ("hidden_size", "intermediate_size", "num_attention_heads", "num_key_value_heads"): + setattr(encoder_config, field, getattr(config, "encoder_" + field, getattr(config, field))) + self.encoder = AceStepConditionEncoder(encoder_config) + self.tokenizer = nn.Module() # Only the inverse FSQ projection is needed. + self.tokenizer.quantizer = AudioCodeQuantizer(encoder_config) + self.detokenizer = AudioTokenDetokenizer(encoder_config) + self.post_init() + + @torch.no_grad() + def generate_audio(self, text_hidden_states, text_attention_mask, + lyric_hidden_states, lyric_attention_mask, + refer_audio_acoustic_hidden_states_packed, + refer_audio_order_mask, src_latents, chunk_masks, seed, + progress=None): + encoder_hidden_states, encoder_attention_mask = self.encoder( + text_hidden_states=text_hidden_states, + text_attention_mask=text_attention_mask, + lyric_hidden_states=lyric_hidden_states, + lyric_attention_mask=lyric_attention_mask, + refer_audio_acoustic_hidden_states_packed=refer_audio_acoustic_hidden_states_packed, + refer_audio_order_mask=refer_audio_order_mask, + ) + context = torch.cat([src_latents, chunk_masks.to(src_latents.dtype)], dim=-1) + bsz, device, dtype = context.shape[0], context.device, context.dtype + generator = torch.Generator(device=device).manual_seed(seed) + xt = torch.randn(src_latents.shape, generator=generator, device=device, dtype=dtype) + cache = EncoderDecoderCache(DynamicCache(), DynamicCache()) + schedule = torch.tensor([1.0 - i / 8 for i in range(8)], device=device, dtype=dtype) + attention_mask = torch.ones(src_latents.shape[:2], device=device, dtype=torch.long) + # Official defaults: ODE/Euler, shift=1, no CFG/EMA/clamping/retake. + # DCW defaults to enabled/double/0.05/0.02/haar, but its optional + # pytorch_wavelets loader returns None in the supported dependency set. + # Port the resulting no-op, rather than introducing a wavelet dependency. + for i in range(8): + t = schedule[i].item() + ts = t * torch.ones((bsz,), device=device, dtype=dtype) + vt, cache = self.decoder( + hidden_states=xt, timestep=ts, timestep_r=ts, + attention_mask=attention_mask, encoder_hidden_states=encoder_hidden_states, + encoder_attention_mask=encoder_attention_mask, context_latents=context, + use_cache=True, past_key_values=cache, + )[:2] + if i == 7: + xt = xt - vt * ts.unsqueeze(-1).unsqueeze(-1) + else: + dt = t - schedule[i + 1].item() + dt_tensor = dt * torch.ones((bsz,), device=device, dtype=dtype).unsqueeze(-1).unsqueeze(-1) + xt = xt - vt * dt_tensor + if progress: + progress(0.15 + (i + 1) * 0.075, desc=f"Generating music ({i + 1}/8)") + return xt diff --git a/freevideo_engine/music/constrained_logits.py b/freevideo_engine/music/constrained_logits.py new file mode 100644 index 00000000..8d42d760 --- /dev/null +++ b/freevideo_engine/music/constrained_logits.py @@ -0,0 +1,993 @@ +# Adapted from ACE-Step 1.5 constrained_logits_processor.py (MIT). +""" +Constrained Logits Processor for ACE-Step Language Model + +This module implements a finite state machine (FSM) based logits processor that constrains +the language model's output to follow specific formats and value ranges during music generation. + +Key Features: +- Enforces structured metadata generation (BPM, duration, keyscale, etc.) +- Validates numeric ranges (BPM: 30-300, Duration: 10-600s) +- Ensures proper formatting for musical metadata +- Prevents generation of invalid tokens or formats +- Supports constrained audio code generation (0-63999) + +The FSM guides the model through different states to ensure outputs conform to expected +schema requirements without post-processing corrections. + +Usage: + processor = ConstrainedLogitsProcessor(tokenizer, mode="metadata") + outputs = model.generate(inputs, logits_processor=[processor]) +""" +from enum import Enum, auto +from typing import Optional, Dict, Any, Tuple, List, Callable, Set +import logging +logger = logging.getLogger(__name__) +from transformers import AutoTokenizer +from transformers.generation.logits_process import LogitsProcessor +import torch +from .lm_constants import VALID_LANGUAGES, KEYSCALE_NOTES, KEYSCALE_ACCIDENTALS, KEYSCALE_MODES, VALID_KEYSCALES, BPM_MIN, BPM_MAX, DURATION_MIN, DURATION_MAX, VALID_TIME_SIGNATURES +MAX_AUDIO_CODE = 63999 + +class FSMState(Enum): + """Finite State Machine states for metadata generation""" + THINK_TAG = auto() + NEWLINE_AFTER_THINK = auto() + BPM_NAME = auto() + BPM_VALUE = auto() + NEWLINE_AFTER_BPM = auto() + CAPTION_NAME = auto() + CAPTION_VALUE = auto() + DURATION_NAME = auto() + DURATION_VALUE = auto() + NEWLINE_AFTER_DURATION = auto() + GENRES_NAME = auto() + GENRES_VALUE = auto() + NEWLINE_AFTER_GENRES = auto() + KEYSCALE_NAME = auto() + KEYSCALE_VALUE = auto() + NEWLINE_AFTER_KEYSCALE = auto() + LANGUAGE_NAME = auto() + LANGUAGE_VALUE = auto() + TIMESIG_NAME = auto() + TIMESIG_VALUE = auto() + NEWLINE_AFTER_TIMESIG = auto() + THINK_END_TAG = auto() + CODES_GENERATION = auto() + COMPLETED = auto() + +class MetadataConstrainedLogitsProcessor(LogitsProcessor): + """ + FSM-driven LogitsProcessor that constrains generation to produce valid metadata. + + This processor enforces the following format: + + bpm: [30-300] + caption: [text without code blocks, ends with period + newline] + duration: [10-600] + keyscale: [A-G][#/♭]? [major/minor] + language: [en/zh/ja/ko/es/fr/de/uk/ru/...] + timesignature: [2/3/4/6] + + + It uses token masking (setting invalid token logits to -inf) to enforce constraints. + For numeric fields, it uses early-blocking to prevent out-of-range values. + For field transitions (e.g., end of numeric value), it compares P(newline) vs P(digit). + For caption field, it blocks code blocks and newlines, and only transitions when + the previous token was a period and newline has the highest probability. + """ + + def __init__(self, tokenizer: AutoTokenizer, enabled: bool=True, debug: bool=False, genres_vocab_path: Optional[str]=None, skip_genres: bool=True, max_duration: Optional[int]=None): + self.tokenizer = tokenizer + self.enabled = enabled + self.debug = debug + assert skip_genres, "genre generation is outside instrumental inference" + self.skip_genres = True + self.max_duration = max_duration if max_duration is not None else DURATION_MAX + self.skip_caption = False + self.skip_language = False + self.caption: Optional[str] = None + self.user_provided_metadata: Dict[str, Optional[str]] = {'bpm': None, 'caption': None, 'duration': None, 'keyscale': None, 'language': None, 'timesignature': None, 'genres': None} + self.metadata_temperature: Optional[float] = None + self.codes_temperature: Optional[float] = None + self.target_duration: Optional[float] = None + self.target_codes: Optional[int] = None + self.codes_count: int = 0 + self.stop_at_reasoning: bool = False + self.generation_phase: str = 'cot' + self.state = FSMState.THINK_TAG + self.position_in_state = 0 + self.accumulated_value = '' + self.accumulated_token_ids: List[int] = [] + self.caption_after_newline = False + self.caption_token_count = 0 + self.caption_ending = False + self.pending_field_name = '' + self.user_field_token_queue: List[int] = [] + self.current_user_field: Optional[str] = None + self._precompute_tokens() + self._char_to_tokens: Dict[str, set] = {} + self._precompute_char_token_mapping() + self.field_specs = {'bpm': {'min': BPM_MIN, 'max': BPM_MAX}, 'duration': {'min': DURATION_MIN, 'max': self.max_duration}, 'timesignature': {'valid_values': VALID_TIME_SIGNATURES}} + self.valid_bpm_values = [str(v) for v in range(self.field_specs['bpm']['min'], self.field_specs['bpm']['max'] + 1)] + self.valid_duration_values = [str(v) for v in range(self.field_specs['duration']['min'], self.field_specs['duration']['max'] + 1)] + self.valid_timesig_values = [str(v) for v in self.field_specs['timesignature']['valid_values']] + self.keyscale_prefix_tree = self._build_keyscale_prefix_tree() + self.bpm_prefix_tree = self._build_numeric_prefix_tree(self.valid_bpm_values, context_prefix_for_matching='bpm:', context_prefix_for_tokenization='bpm: ') + self.duration_prefix_tree = self._build_numeric_prefix_tree(self.valid_duration_values, context_prefix_for_matching='duration:', context_prefix_for_tokenization='duration: ') + self.timesig_prefix_tree = self._build_numeric_prefix_tree(self.valid_timesig_values, context_prefix_for_matching='timesignature:', context_prefix_for_tokenization='timesignature: ') + self.language_prefix_tree = self._build_language_prefix_tree() + self.fixed_strings = {FSMState.THINK_TAG: '', FSMState.NEWLINE_AFTER_THINK: '\n', FSMState.BPM_NAME: 'bpm:', FSMState.CAPTION_NAME: 'caption:', FSMState.DURATION_NAME: 'duration:', FSMState.GENRES_NAME: 'genres:', FSMState.KEYSCALE_NAME: 'keyscale:', FSMState.LANGUAGE_NAME: 'language:', FSMState.TIMESIG_NAME: 'timesignature:', FSMState.THINK_END_TAG: ''} + self._build_state_transitions() + + def _get_next_field_state(self, current_field: str) -> Optional[FSMState]: + field_order = ['bpm', 'caption', 'duration', 'genres', 'keyscale', 'language', 'timesignature'] + field_to_state = {'bpm': FSMState.BPM_NAME, 'caption': FSMState.CAPTION_NAME, 'duration': FSMState.DURATION_NAME, 'genres': FSMState.GENRES_NAME, 'keyscale': FSMState.KEYSCALE_NAME, 'language': FSMState.LANGUAGE_NAME, 'timesignature': FSMState.TIMESIG_NAME} + try: + current_idx = field_order.index(current_field) + except ValueError: + return FSMState.THINK_END_TAG + for i in range(current_idx + 1, len(field_order)): + field = field_order[i] + if field == 'genres' and self.skip_genres: + continue + if field == 'caption' and self.skip_caption: + continue + if field == 'language' and self.skip_language: + continue + return field_to_state[field] + return FSMState.THINK_END_TAG + + def _build_state_transitions(self): + self.next_state = {FSMState.THINK_TAG: FSMState.NEWLINE_AFTER_THINK, FSMState.NEWLINE_AFTER_THINK: FSMState.BPM_NAME, FSMState.THINK_END_TAG: FSMState.CODES_GENERATION, FSMState.CODES_GENERATION: FSMState.COMPLETED} + self.next_state[FSMState.BPM_NAME] = FSMState.BPM_VALUE + self.next_state[FSMState.BPM_VALUE] = self._get_next_field_state('bpm') + if not self.skip_caption: + self.next_state[FSMState.CAPTION_NAME] = FSMState.CAPTION_VALUE + self.next_state[FSMState.CAPTION_VALUE] = self._get_next_field_state('caption') + self.next_state[FSMState.DURATION_NAME] = FSMState.DURATION_VALUE + self.next_state[FSMState.DURATION_VALUE] = self._get_next_field_state('duration') + if not self.skip_genres: + self.next_state[FSMState.GENRES_NAME] = FSMState.GENRES_VALUE + self.next_state[FSMState.GENRES_VALUE] = self._get_next_field_state('genres') + self.next_state[FSMState.KEYSCALE_NAME] = FSMState.KEYSCALE_VALUE + self.next_state[FSMState.KEYSCALE_VALUE] = self._get_next_field_state('keyscale') + if not self.skip_language: + self.next_state[FSMState.LANGUAGE_NAME] = FSMState.LANGUAGE_VALUE + self.next_state[FSMState.LANGUAGE_VALUE] = self._get_next_field_state('language') + self.next_state[FSMState.TIMESIG_NAME] = FSMState.TIMESIG_VALUE + self.next_state[FSMState.TIMESIG_VALUE] = FSMState.THINK_END_TAG + + def set_skip_genres(self, skip: bool): + self.skip_genres = skip + self._build_state_transitions() + + def set_skip_caption(self, skip: bool): + self.skip_caption = skip + self._build_state_transitions() + + def set_skip_language(self, skip: bool): + self.skip_language = skip + self._build_state_transitions() + + @staticmethod + def postprocess_caption(caption: str) -> str: + if not caption: + return caption + lines = caption.split('\n') + cleaned_lines = [] + for line in lines: + stripped = line.strip() + if stripped: + cleaned_lines.append(stripped) + return ' '.join(cleaned_lines) + + def set_stop_at_reasoning(self, stop: bool): + self.stop_at_reasoning = stop + + def set_generation_phase(self, phase: str): + if phase not in ('cot', 'codes'): + raise ValueError(f"Invalid generation phase: {phase!r}. Must be 'cot' or 'codes'") + self.generation_phase = phase + + def set_user_metadata(self, metadata: Optional[Dict[str, Optional[str]]]=None): + if metadata is None: + metadata = {} + for field in ['bpm', 'caption', 'duration', 'keyscale', 'language', 'timesignature', 'genres']: + if field in metadata: + self.user_provided_metadata[field] = metadata[field] + else: + self.user_provided_metadata[field] = None + self._build_state_transitions() + if self.debug: + provided_fields = [k for (k, v) in self.user_provided_metadata.items() if v is not None] + if provided_fields: + logger.debug(f'User provided metadata fields: {provided_fields}') + else: + logger.debug('No user-provided metadata, all fields will be generated') + + def _precompute_tokens(self): + self.digit_tokens = {} + for d in range(10): + tokens = self.tokenizer.encode(str(d), add_special_tokens=False) + if tokens: + self.digit_tokens[d] = tokens[-1] + newline_tokens = self.tokenizer.encode('\n', add_special_tokens=False) + self.newline_token = newline_tokens[-1] if newline_tokens else None + self.note_tokens = {} + for note in KEYSCALE_NOTES: + tokens = self.tokenizer.encode(note, add_special_tokens=False) + if tokens: + self.note_tokens[note] = tokens[-1] + self.sharp_tokens = [] + for s in ['#', '♯']: + tokens = self.tokenizer.encode(s, add_special_tokens=False) + if tokens: + self.sharp_tokens.append(tokens[-1]) + self.flat_tokens = [] + for f in ['b', '♭']: + tokens = self.tokenizer.encode(f, add_special_tokens=False) + if tokens: + self.flat_tokens.append(tokens[-1]) + space_tokens = self.tokenizer.encode(' ', add_special_tokens=False) + self.space_token = space_tokens[-1] if space_tokens else None + self.major_start_tokens = [] + self.minor_start_tokens = [] + for prefix in ['m', 'M']: + tokens = self.tokenizer.encode(prefix, add_special_tokens=False) + if tokens: + if prefix.lower() == 'm': + self.minor_start_tokens.append(tokens[-1]) + self.major_start_tokens.append(tokens[-1]) + self.vocab_size = len(self.tokenizer) + comma_tokens = self.tokenizer.encode(',', add_special_tokens=False) + self.comma_token = comma_tokens[-1] if comma_tokens else None + self.eos_token_id = self.tokenizer.eos_token_id + period_tokens = self.tokenizer.encode('.', add_special_tokens=False) + self.period_token = period_tokens[-1] if period_tokens else None + backtick_tokens = self.tokenizer.encode('`', add_special_tokens=False) + self.backtick_token = backtick_tokens[-1] if backtick_tokens else None + self.valid_languages = VALID_LANGUAGES + self.audio_code_token_ids: Set[int] = set() + self._precompute_audio_code_tokens() + self.audio_code_mask: Optional[torch.Tensor] = None + self.non_audio_code_mask: Optional[torch.Tensor] = None + self._build_audio_code_mask() + self.valid_keyscales = VALID_KEYSCALES.copy() + + def _precompute_audio_code_tokens(self): + import re + audio_code_pattern = re.compile('^<\\|audio_code_(\\d+)\\|>$') + invalid_tokens_count = 0 + for token_id in range(self.vocab_size): + try: + token_text = self.tokenizer.decode([token_id]) + match = audio_code_pattern.match(token_text) + if match: + code_value = int(match.group(1)) + if 0 <= code_value <= MAX_AUDIO_CODE: + self.audio_code_token_ids.add(token_id) + else: + invalid_tokens_count += 1 + if self.debug: + logger.debug(f'Skipping audio code token {token_id} with invalid code value {code_value} (max: {MAX_AUDIO_CODE})') + except Exception: + continue + if invalid_tokens_count > 0: + logger.debug(f'Found {invalid_tokens_count} audio code tokens with values outside valid range [0, {MAX_AUDIO_CODE}]') + if len(self.audio_code_token_ids) == 0: + logger.warning(f'No valid audio code tokens found in vocabulary (range [0, {MAX_AUDIO_CODE}]). Code generation may fail.') + elif self.debug: + logger.debug(f'Found {len(self.audio_code_token_ids)} valid audio code tokens (range [0, {MAX_AUDIO_CODE}])') + + def _extract_code_from_token(self, token_id: int) -> Optional[int]: + import re + audio_code_pattern = re.compile('^<\\|audio_code_(\\d+)\\|>$') + try: + token_text = self.tokenizer.decode([token_id]) + match = audio_code_pattern.match(token_text) + if match: + return int(match.group(1)) + except Exception: + pass + return None + + def _build_audio_code_mask(self): + if not self.audio_code_token_ids: + self.audio_code_mask = None + self.non_audio_code_mask = None + return + mask = torch.zeros(1, self.vocab_size, dtype=torch.float32) + audio_code_indices = list(self.audio_code_token_ids) + mask[0, audio_code_indices] = float('-inf') + self.audio_code_mask = mask + inverse_mask = torch.full((1, self.vocab_size), float('-inf'), dtype=torch.float32) + inverse_mask[0, audio_code_indices] = 0 + if self.eos_token_id is not None: + inverse_mask[0, self.eos_token_id] = 0 + self.non_audio_code_mask = inverse_mask + if self.debug: + logger.debug(f'Built audio code masks for {len(self.audio_code_token_ids)} tokens') + + def _apply_whitelist_inplace(self, scores: torch.Tensor, allowed_tokens: List[int]) -> None: + if not allowed_tokens: + scores.fill_(float('-inf')) + return + allowed_indices = torch.tensor(allowed_tokens, device=scores.device, dtype=torch.long) + saved_values = scores[0, allowed_indices].clone() + scores.fill_(float('-inf')) + scores[0, allowed_indices] = saved_values + + def _build_keyscale_prefix_tree(self) -> Dict[Tuple[int, ...], Set[int]]: + prefix_to_tokens: Dict[Tuple[int, ...], Set[int]] = {} + context_prefix_for_matching = 'keyscale:' + context_prefix_for_tokenization = 'keyscale: ' + context_token_ids = self.tokenizer.encode(context_prefix_for_matching, add_special_tokens=False) + if self.debug: + context_tokens_str = [self.tokenizer.decode([t]) for t in context_token_ids] + logger.debug(f"Context for matching 'keyscale:' tokenizes to {context_token_ids} -> {context_tokens_str}") + for keyscale in self.valid_keyscales: + full_text = context_prefix_for_tokenization + keyscale + full_token_ids = self.tokenizer.encode(full_text, add_special_tokens=False) + context_end_idx = None + if len(full_token_ids) >= len(context_token_ids): + if full_token_ids[:len(context_token_ids)] == context_token_ids: + context_end_idx = len(context_token_ids) + if context_end_idx is None: + if self.debug: + logger.warning(f"Could not find context prefix in full tokenization of '{full_text}', skipping") + continue + keyscale_token_ids = full_token_ids[context_end_idx:] + if not keyscale_token_ids: + if self.debug: + logger.warning(f"No tokens extracted for keyscale '{keyscale}', skipping") + continue + first_token_id = keyscale_token_ids[0] + first_token_str = self.tokenizer.decode([first_token_id]) + first_char = first_token_str.lstrip()[0].upper() if first_token_str.lstrip() else '' + if first_char not in 'ABCDEFG': + if self.debug: + logger.debug(f"Skipping keyscale '{keyscale}': first token is '{first_token_str}' (id={first_token_id}), not a note") + continue + for i in range(len(keyscale_token_ids) + 1): + token_prefix = tuple(keyscale_token_ids[:i]) + if token_prefix not in prefix_to_tokens: + prefix_to_tokens[token_prefix] = set() + if i < len(keyscale_token_ids): + next_token_id = keyscale_token_ids[i] + prefix_to_tokens[token_prefix].add(next_token_id) + elif self.newline_token: + prefix_to_tokens[token_prefix].add(self.newline_token) + if self.debug: + logger.debug(f'Built keyscale prefix tree with {len(prefix_to_tokens)} token sequence prefixes') + empty_prefix = tuple() + if empty_prefix in prefix_to_tokens: + first_tokens = prefix_to_tokens[empty_prefix] + decoded_first = [(t, repr(self.tokenizer.decode([t]))) for t in sorted(first_tokens)] + logger.debug(f'First tokens allowed (empty prefix): {decoded_first}') + return prefix_to_tokens + + def _build_numeric_prefix_tree(self, valid_values: List[str], context_prefix_for_matching: str='', context_prefix_for_tokenization: str='') -> Dict[Tuple[int, ...], Set[int]]: + prefix_to_tokens: Dict[Tuple[int, ...], Set[int]] = {} + context_token_ids = self.tokenizer.encode(context_prefix_for_matching, add_special_tokens=False) if context_prefix_for_matching else [] + for value_str in valid_values: + full_text = context_prefix_for_tokenization + value_str + token_ids = self.tokenizer.encode(full_text, add_special_tokens=False) + context_end_idx = None + if len(token_ids) >= len(context_token_ids): + if token_ids[:len(context_token_ids)] == context_token_ids: + context_end_idx = len(context_token_ids) + if context_end_idx is None: + if self.debug: + logger.warning(f"Could not find context prefix in full tokenization of '{full_text}', skipping") + continue + value_token_ids = token_ids[context_end_idx:] + for i in range(len(value_token_ids) + 1): + token_prefix = tuple(value_token_ids[:i]) + if token_prefix not in prefix_to_tokens: + prefix_to_tokens[token_prefix] = set() + if i < len(value_token_ids): + next_token_id = value_token_ids[i] + prefix_to_tokens[token_prefix].add(next_token_id) + elif self.newline_token: + prefix_to_tokens[token_prefix].add(self.newline_token) + return prefix_to_tokens + + def _build_language_prefix_tree(self) -> Dict[Tuple[int, ...], Set[int]]: + prefix_to_tokens: Dict[Tuple[int, ...], Set[int]] = {} + context_prefix_for_matching = 'language:' + context_prefix_for_tokenization = 'language: ' + context_token_ids = self.tokenizer.encode(context_prefix_for_matching, add_special_tokens=False) + if self.debug: + context_tokens_str = [self.tokenizer.decode([t]) for t in context_token_ids] + logger.debug(f"Context for matching 'language:' tokenizes to {context_token_ids} -> {context_tokens_str}") + for lang in self.valid_languages: + full_text = context_prefix_for_tokenization + lang + full_token_ids = self.tokenizer.encode(full_text, add_special_tokens=False) + context_end_idx = None + if len(full_token_ids) >= len(context_token_ids): + if full_token_ids[:len(context_token_ids)] == context_token_ids: + context_end_idx = len(context_token_ids) + if context_end_idx is None: + if self.debug: + logger.warning(f"Could not find context prefix in full tokenization of '{full_text}', skipping") + continue + lang_token_ids = full_token_ids[context_end_idx:] + if not lang_token_ids: + if self.debug: + logger.warning(f"No tokens extracted for language '{lang}', skipping") + continue + for i in range(len(lang_token_ids) + 1): + token_prefix = tuple(lang_token_ids[:i]) + if token_prefix not in prefix_to_tokens: + prefix_to_tokens[token_prefix] = set() + if i < len(lang_token_ids): + next_token_id = lang_token_ids[i] + prefix_to_tokens[token_prefix].add(next_token_id) + elif self.newline_token: + prefix_to_tokens[token_prefix].add(self.newline_token) + if self.debug: + logger.debug(f'Built language prefix tree with {len(prefix_to_tokens)} token sequence prefixes') + empty_prefix = tuple() + if empty_prefix in prefix_to_tokens: + first_tokens = prefix_to_tokens[empty_prefix] + decoded_first = [(t, repr(self.tokenizer.decode([t]))) for t in sorted(first_tokens)] + logger.debug(f'First tokens allowed for language (empty prefix): {decoded_first}') + return prefix_to_tokens + + def _precompute_char_token_mapping(self): + self._char_to_tokens: Dict[str, set] = {} + self._token_to_text: Dict[int, str] = {} + for token_id in range(self.vocab_size): + try: + text = self.tokenizer.decode([token_id]) + if not text: + continue + text_lower = text.lower() + if text_lower.strip(): + normalized_text = text_lower.rstrip() + else: + normalized_text = ' ' + self._token_to_text[token_id] = normalized_text + first_char = text[0].lower() + if first_char not in self._char_to_tokens: + self._char_to_tokens[first_char] = set() + self._char_to_tokens[first_char].add(token_id) + stripped_text = text.lstrip() + if stripped_text and stripped_text != text: + first_nonspace_char = stripped_text[0].lower() + if first_nonspace_char not in self._char_to_tokens: + self._char_to_tokens[first_nonspace_char] = set() + self._char_to_tokens[first_nonspace_char].add(token_id) + except Exception: + continue + if self.debug: + logger.debug(f'Precomputed char->token mapping for {len(self._char_to_tokens)} unique characters') + + def reset(self): + self.state = FSMState.THINK_TAG + self.position_in_state = 0 + self.accumulated_value = '' + self.accumulated_token_ids = [] + self.codes_count = 0 + self.user_field_token_queue = [] + self.current_user_field = None + self.caption_after_newline = False + self.caption_token_count = 0 + self.caption_ending = False + self.pending_field_name = '' + + def set_target_duration(self, duration: Optional[float]): + self.target_duration = duration + if duration is not None and duration > 0: + self.target_codes = int(duration * 5) + if self.debug: + logger.debug(f'Set target duration: {duration}s -> {self.target_codes} codes') + else: + self.target_codes = None + if self.debug: + logger.debug('Target duration cleared, no duration constraint') + + def set_max_duration(self, max_duration: int): + if max_duration == self.max_duration: + return + old_max = self.max_duration + self.max_duration = max_duration + self.field_specs['duration']['max'] = max_duration + self.valid_duration_values = [str(v) for v in range(self.field_specs['duration']['min'], self.field_specs['duration']['max'] + 1)] + self.duration_prefix_tree = self._build_numeric_prefix_tree(self.valid_duration_values, context_prefix_for_matching='duration:', context_prefix_for_tokenization='duration: ') + if self.debug: + logger.debug(f'Updated max duration: {old_max}s -> {max_duration}s, rebuilt prefix tree with {len(self.valid_duration_values)} values') + + def _get_allowed_tokens_for_fixed_string(self, fixed_str: str) -> List[int]: + remaining = fixed_str[self.position_in_state:] + if not remaining: + return [] + if self.debug: + logger.debug(f'_get_allowed_tokens_for_fixed_string: fixed_str={repr(fixed_str)}, position_in_state={self.position_in_state}, remaining={repr(remaining)}') + best_token = None + best_prefix_len = 0 + for end in range(len(remaining), 0, -1): + prefix = remaining[:end] + tokens = self.tokenizer.encode(prefix, add_special_tokens=False) + if tokens and len(tokens) == 1: + best_token = tokens[0] + best_prefix_len = end + if self.debug: + logger.debug(f'Found single-token match: prefix={repr(prefix)}, token_id={best_token}, token_text={repr(self.tokenizer.decode([best_token]))}') + break + if best_token is not None: + return [best_token] + allowed_tokens = {} + for end in range(1, min(len(remaining) + 1, 20)): + prefix = remaining[:end] + tokens = self.tokenizer.encode(prefix, add_special_tokens=False) + if tokens: + first_token = tokens[0] + decoded_token = self.tokenizer.decode([first_token]) + normalized_prefix = prefix.lstrip().lower() + normalized_decoded = decoded_token.lstrip().lower() + if normalized_decoded.startswith(normalized_prefix) or normalized_prefix.startswith(normalized_decoded): + if first_token not in allowed_tokens or end > allowed_tokens[first_token]: + allowed_tokens[first_token] = end + sorted_tokens = sorted(allowed_tokens.items(), key=lambda x: x[1], reverse=True) + result = [token for (token, _) in sorted_tokens] if sorted_tokens else [] + if self.debug: + logger.debug(f'Fallback: returning {len(result)} tokens: {[(t, repr(self.tokenizer.decode([t]))) for t in result[:5]]}') + if result: + logger.debug(f'Fixed string: {repr(fixed_str)}, position: {self.position_in_state}, remaining: {repr(remaining)}') + return result + + def _get_allowed_digit_tokens(self, min_val: int, max_val: int) -> List[int]: + if not self.accumulated_value: + allowed_digits = set() + for v in range(min_val, max_val + 1): + allowed_digits.add(int(str(v)[0])) + return [self.digit_tokens[d] for d in allowed_digits if d in self.digit_tokens] + current = int(self.accumulated_value) + allowed = [] + for d in range(10): + new_value = int(self.accumulated_value + str(d)) + if new_value > max_val: + continue + if new_value >= min_val: + allowed.append(d) + elif new_value * 10 <= max_val: + allowed.append(d) + return [self.digit_tokens[d] for d in allowed if d in self.digit_tokens] + + def _get_allowed_numeric_tokens(self, prefix_tree: Dict[Tuple[int, ...], Set[int]]) -> List[int]: + token_prefix = tuple(self.accumulated_token_ids) + if token_prefix in prefix_tree: + return list(prefix_tree[token_prefix]) + return [] + + def _should_end_numeric_field(self, logits: torch.Tensor, min_val: int, max_val: int) -> bool: + if not self.accumulated_value: + return False + current = int(self.accumulated_value) + if current < min_val or current > max_val: + return False + probs = torch.softmax(logits, dim=-1) + newline_prob = probs[0, self.newline_token].item() if self.newline_token else 0 + allowed_digits = self._get_allowed_digit_tokens(min_val, max_val) + if not allowed_digits: + return True + max_digit_prob = max((probs[0, t].item() for t in allowed_digits)) + if self.debug: + logger.debug(f'Numeric field decision: newline_prob={newline_prob:.4f}, max_digit_prob={max_digit_prob:.4f}') + return newline_prob > max_digit_prob + + def _should_end_text_field(self, logits: torch.Tensor) -> bool: + if not self.accumulated_value.strip(): + return False + probs = torch.softmax(logits, dim=-1) + newline_prob = probs[0, self.newline_token].item() if self.newline_token else 0 + masked_probs = probs.clone() + if self.newline_token: + masked_probs[0, self.newline_token] = 0 + max_other_prob = masked_probs[0].max().item() + return newline_prob > max_other_prob + + def _get_allowed_keyscale_tokens(self) -> List[int]: + token_prefix = tuple(self.accumulated_token_ids) + if token_prefix in self.keyscale_prefix_tree: + return list(self.keyscale_prefix_tree[token_prefix]) + return [] + + def _is_keyscale_complete(self) -> bool: + token_prefix = tuple(self.accumulated_token_ids) + if token_prefix in self.keyscale_prefix_tree: + return self.newline_token in self.keyscale_prefix_tree[token_prefix] + return False + + def _get_allowed_language_tokens(self) -> List[int]: + token_prefix = tuple(self.accumulated_token_ids) + if token_prefix in self.language_prefix_tree: + return list(self.language_prefix_tree[token_prefix]) + return [] + + def _get_allowed_timesig_tokens(self) -> List[int]: + token_prefix = tuple(self.accumulated_token_ids) + if token_prefix in self.timesig_prefix_tree: + return list(self.timesig_prefix_tree[token_prefix]) + return [] + + def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor) -> torch.FloatTensor: + if not self.enabled: + return self._apply_temperature_scaling(scores) + if self.state == FSMState.COMPLETED: + return self._apply_temperature_scaling(scores) + if self.generation_phase == 'codes' and self.state == FSMState.THINK_TAG: + if self._input_contains_think_end_tag(input_ids): + self.state = FSMState.CODES_GENERATION + self.codes_count = 0 + if self.debug: + logger.debug('Codes phase: detected in input, skipping to CODES_GENERATION') + if self.state == FSMState.CODES_GENERATION: + if self.non_audio_code_mask is not None: + if self.non_audio_code_mask.device != scores.device or self.non_audio_code_mask.dtype != scores.dtype: + self.non_audio_code_mask = self.non_audio_code_mask.to(device=scores.device, dtype=scores.dtype) + scores = scores + self.non_audio_code_mask + if self.target_codes is not None and self.eos_token_id is not None: + if self.codes_count < self.target_codes: + scores[:, self.eos_token_id] = float('-inf') + if self.debug: + logger.debug(f'Codes generation: {self.codes_count}/{self.target_codes}, blocking EOS') + else: + eos_scores = scores[:, self.eos_token_id].clone() + scores.fill_(float('-inf')) + scores[:, self.eos_token_id] = eos_scores + if self.debug: + logger.debug(f'Codes generation: {self.codes_count}/{self.target_codes}, forcing EOS') + return self._apply_temperature_scaling(scores) + batch_size = scores.shape[0] + for b in range(batch_size): + result = self._process_single_sequence(input_ids[b], scores[b:b + 1]) + scores[b] = result[0] + return self._apply_temperature_scaling(scores) + + def _input_contains_think_end_tag(self, input_ids: torch.LongTensor) -> bool: + think_end_tokens = self.tokenizer.encode('', add_special_tokens=False) + if not think_end_tokens: + return False + for b in range(input_ids.shape[0]): + seq = input_ids[b].tolist() + for i in range(len(seq) - len(think_end_tokens) + 1): + if seq[i:i + len(think_end_tokens)] == think_end_tokens: + return True + return False + + def _apply_temperature_scaling(self, scores: torch.FloatTensor) -> torch.FloatTensor: + if self.state == FSMState.CODES_GENERATION or self.state == FSMState.COMPLETED: + temperature = self.codes_temperature + else: + temperature = self.metadata_temperature + if temperature is None: + return scores + if temperature <= 0: + temperature = 1e-06 + return scores / temperature + + def _get_user_provided_field_tokens(self, field_name: str) -> Optional[List[int]]: + value = self.user_provided_metadata.get(field_name) + if value is None: + return None + field_to_prefix = {'bpm': 'bpm: ', 'caption': 'caption: ', 'duration': 'duration: ', 'keyscale': 'keyscale: ', 'language': 'language: ', 'timesignature': 'timesignature: ', 'genres': 'genres: '} + prefix = field_to_prefix[field_name] + full_text = f'{prefix}{value}\n' + tokens = self.tokenizer.encode(full_text, add_special_tokens=False) + prefix_for_matching = field_name + ':' + prefix_tokens = self.tokenizer.encode(prefix_for_matching, add_special_tokens=False) + if len(tokens) >= len(prefix_tokens) and tokens[:len(prefix_tokens)] == prefix_tokens: + return tokens[len(prefix_tokens):] + else: + if self.debug: + logger.warning(f'Could not match prefix tokens for field {field_name}, using all tokens') + return tokens + + def _process_single_sequence(self, input_ids: torch.LongTensor, scores: torch.FloatTensor) -> torch.FloatTensor: + if self.user_field_token_queue: + next_token = self.user_field_token_queue[0] + self._apply_whitelist_inplace(scores, [next_token]) + return scores + if self.state in self.fixed_strings: + fixed_str = self.fixed_strings[self.state] + allowed = self._get_allowed_tokens_for_fixed_string(fixed_str) + if allowed: + if self.state == FSMState.THINK_END_TAG and self.stop_at_reasoning: + remaining_chars = len(fixed_str) - self.position_in_state + if remaining_chars <= 10: + if self.eos_token_id is not None: + self._apply_whitelist_inplace(scores, [self.eos_token_id]) + if self.debug: + logger.debug(f'stop_at_reasoning=True: forcing EOS near end of tag (remaining: {remaining_chars} chars)') + return scores + self._apply_whitelist_inplace(scores, allowed) + else: + if self.state == FSMState.THINK_END_TAG and self.stop_at_reasoning: + if self.eos_token_id is not None: + self._apply_whitelist_inplace(scores, [self.eos_token_id]) + if self.debug: + logger.debug(f'stop_at_reasoning=True: forcing EOS after completing tag') + return scores + old_state = self.state + self._transition_to_next_state() + if self.state in self.fixed_strings: + if self.debug: + logger.warning(f'State transition from {old_state.name} to {self.state.name} still in fixed_strings, avoiding recursion') + return scores + scores.zero_() + return self._process_single_sequence(input_ids, scores) + elif self.state == FSMState.BPM_VALUE: + if self.user_provided_metadata['bpm'] is not None and (not self.user_field_token_queue) and (not self.accumulated_token_ids): + value = self.user_provided_metadata['bpm'] + value_text = f' {value}\n' + value_tokens = self.tokenizer.encode(value_text, add_special_tokens=False) + if value_tokens: + self.user_field_token_queue = value_tokens + self.current_user_field = 'bpm' + self._apply_whitelist_inplace(scores, [value_tokens[0]]) + return scores + allowed = self._get_allowed_numeric_tokens(self.bpm_prefix_tree) + token_prefix = tuple(self.accumulated_token_ids) + if token_prefix in self.bpm_prefix_tree and self.newline_token in self.bpm_prefix_tree[token_prefix]: + allowed = allowed + [self.newline_token] + self._apply_whitelist_inplace(scores, allowed) + elif self.state == FSMState.CAPTION_VALUE: + if self.user_provided_metadata['caption'] is not None and (not self.user_field_token_queue) and (not self.accumulated_value): + value = self.user_provided_metadata['caption'] + value_text = f' {value}\n' + value_tokens = self.tokenizer.encode(value_text, add_special_tokens=False) + if value_tokens: + self.user_field_token_queue = value_tokens + self.current_user_field = 'caption' + self._apply_whitelist_inplace(scores, [value_tokens[0]]) + return scores + if self.caption_after_newline: + top_token_id = torch.argmax(scores[0]).item() + top_token_text = self.tokenizer.decode([top_token_id]) + if len(top_token_text) > 0 and top_token_text[0] not in ' \t': + self.caption_after_newline = False + self.caption_ending = True + self.pending_field_name = '' + return scores + else: + self.caption_after_newline = False + if self.caption_ending: + return scores + if self.backtick_token is not None: + scores[0, self.backtick_token] = float('-inf') + if self.audio_code_mask is not None: + if self.audio_code_mask.device != scores.device or self.audio_code_mask.dtype != scores.dtype: + self.audio_code_mask = self.audio_code_mask.to(device=scores.device, dtype=scores.dtype) + scores = scores + self.audio_code_mask + if self.caption_token_count >= 512: + if self.newline_token is not None: + self._apply_whitelist_inplace(scores, [self.newline_token]) + return scores + return scores + elif self.state == FSMState.DURATION_VALUE: + if self.user_provided_metadata['duration'] is not None and (not self.user_field_token_queue) and (not self.accumulated_token_ids): + value = self.user_provided_metadata['duration'] + value_text = f' {value}\n' + value_tokens = self.tokenizer.encode(value_text, add_special_tokens=False) + if value_tokens: + self.user_field_token_queue = value_tokens + self.current_user_field = 'duration' + self._apply_whitelist_inplace(scores, [value_tokens[0]]) + return scores + if self.target_duration is not None: + target_str = str(int(self.target_duration)) + current_pos = len(self.accumulated_value) + if current_pos < len(target_str): + next_digit = int(target_str[current_pos]) + if next_digit in self.digit_tokens: + self._apply_whitelist_inplace(scores, [self.digit_tokens[next_digit]]) + elif self.newline_token: + self._apply_whitelist_inplace(scores, [self.newline_token]) + else: + allowed = self._get_allowed_numeric_tokens(self.duration_prefix_tree) + token_prefix = tuple(self.accumulated_token_ids) + if token_prefix in self.duration_prefix_tree and self.newline_token in self.duration_prefix_tree[token_prefix]: + allowed = allowed + [self.newline_token] + self._apply_whitelist_inplace(scores, allowed) + elif self.state == FSMState.KEYSCALE_VALUE: + if self.user_provided_metadata['keyscale'] is not None and (not self.user_field_token_queue) and (not self.accumulated_token_ids): + value = self.user_provided_metadata['keyscale'] + value_text = f' {value}\n' + value_tokens = self.tokenizer.encode(value_text, add_special_tokens=False) + if value_tokens: + self.user_field_token_queue = value_tokens + self.current_user_field = 'keyscale' + self._apply_whitelist_inplace(scores, [value_tokens[0]]) + return scores + token_prefix = tuple(self.accumulated_token_ids) + if token_prefix in self.keyscale_prefix_tree and self.newline_token in self.keyscale_prefix_tree[token_prefix]: + if self.newline_token: + self._apply_whitelist_inplace(scores, [self.newline_token]) + else: + allowed = self._get_allowed_keyscale_tokens() + if allowed: + self._apply_whitelist_inplace(scores, allowed) + elif self.newline_token: + self._apply_whitelist_inplace(scores, [self.newline_token]) + elif self.state == FSMState.LANGUAGE_VALUE: + if self.user_provided_metadata['language'] is not None and (not self.user_field_token_queue) and (not self.accumulated_token_ids): + value = self.user_provided_metadata['language'] + value_text = f' {value}\n' + value_tokens = self.tokenizer.encode(value_text, add_special_tokens=False) + if value_tokens: + self.user_field_token_queue = value_tokens + self.current_user_field = 'language' + self._apply_whitelist_inplace(scores, [value_tokens[0]]) + return scores + if not self.accumulated_token_ids: + empty_prefix = tuple() + if empty_prefix in self.language_prefix_tree: + candidate_tokens = list(self.language_prefix_tree[empty_prefix]) + if candidate_tokens: + candidate_indices = torch.tensor(candidate_tokens, device=scores.device, dtype=torch.long) + candidate_scores = scores[0, candidate_indices] + best_idx = torch.argmax(candidate_scores).item() + top_token_id = candidate_tokens[best_idx] + self._apply_whitelist_inplace(scores, [top_token_id]) + if self.debug: + top_token_text = self.tokenizer.decode([top_token_id]) + logger.debug(f'Language field: selected top-1 token {top_token_id} ({repr(top_token_text)}) from {len(candidate_tokens)} candidates') + elif self.newline_token: + self._apply_whitelist_inplace(scores, [self.newline_token]) + elif self.newline_token: + self._apply_whitelist_inplace(scores, [self.newline_token]) + else: + token_prefix = tuple(self.accumulated_token_ids) + if token_prefix in self.language_prefix_tree and self.newline_token in self.language_prefix_tree[token_prefix]: + if self.newline_token: + self._apply_whitelist_inplace(scores, [self.newline_token]) + else: + allowed = self._get_allowed_language_tokens() + if allowed: + self._apply_whitelist_inplace(scores, allowed) + elif self.newline_token: + self._apply_whitelist_inplace(scores, [self.newline_token]) + elif self.state == FSMState.TIMESIG_VALUE: + if self.user_provided_metadata['timesignature'] is not None and (not self.user_field_token_queue) and (not self.accumulated_token_ids): + value = self.user_provided_metadata['timesignature'] + value_text = f' {value}\n' + value_tokens = self.tokenizer.encode(value_text, add_special_tokens=False) + if value_tokens: + self.user_field_token_queue = value_tokens + self.current_user_field = 'timesignature' + self._apply_whitelist_inplace(scores, [value_tokens[0]]) + return scores + token_prefix = tuple(self.accumulated_token_ids) + if token_prefix in self.timesig_prefix_tree and self.newline_token in self.timesig_prefix_tree[token_prefix]: + if self.newline_token: + self._apply_whitelist_inplace(scores, [self.newline_token]) + else: + allowed = self._get_allowed_timesig_tokens() + self._apply_whitelist_inplace(scores, allowed) + return scores + + def _transition_to_next_state(self): + if self.state in self.next_state: + old_state = self.state + next_state = self.next_state[self.state] + self.state = next_state + self.position_in_state = 0 + self.accumulated_value = '' + self.accumulated_token_ids = [] + self.caption_after_newline = False + self.caption_token_count = 0 + self.caption_ending = False + self.pending_field_name = '' + if self.debug: + logger.debug(f'FSM transition: {old_state.name} -> {self.state.name}') + + def update_state(self, generated_token_id: int): + if not self.enabled: + return + if self.state == FSMState.COMPLETED: + return + if self.state == FSMState.CODES_GENERATION: + self.codes_count += 1 + if self.debug and self.target_codes is not None: + logger.debug(f'Codes count: {self.codes_count}/{self.target_codes}') + return + if self.user_field_token_queue: + expected_token = self.user_field_token_queue[0] + if generated_token_id != expected_token: + if self.debug: + logger.warning(f'Expected token {expected_token} but got {generated_token_id} for user-provided field {self.current_user_field}') + self.user_field_token_queue.pop(0) + if not self.user_field_token_queue: + if self.debug: + logger.debug(f'Completed injection of user-provided field: {self.current_user_field}') + field_name = self.current_user_field + self.current_user_field = None + next_state = self._get_next_field_state(field_name) + if next_state: + old_state = self.state + self.state = next_state + self.position_in_state = 0 + self.accumulated_value = '' + self.accumulated_token_ids = [] + if self.debug: + logger.debug(f'FSM transition (after user field injection): {old_state.name} -> {self.state.name}') + else: + self._transition_to_next_state() + return + token_str = self.tokenizer.decode([generated_token_id]) + if self.debug: + logger.debug(f'Generated token: {repr(token_str)} (id={generated_token_id}), state={self.state.name}') + if self.state in self.fixed_strings: + fixed_str = self.fixed_strings[self.state] + self.position_in_state += len(token_str) + if self.position_in_state >= len(fixed_str): + if self.state == FSMState.THINK_END_TAG and self.stop_at_reasoning: + old_state = self.state + self.state = FSMState.COMPLETED + self.position_in_state = 0 + self.accumulated_value = '' + self.accumulated_token_ids = [] + if self.debug: + logger.debug(f'FSM transition (stop_at_reasoning): {old_state.name} -> {self.state.name}') + else: + self._transition_to_next_state() + elif self.state in [FSMState.BPM_VALUE, FSMState.DURATION_VALUE, FSMState.TIMESIG_VALUE]: + if generated_token_id == self.newline_token: + old_state = self.state + self._transition_to_next_state() + if self.state in self.fixed_strings: + return + else: + self.accumulated_token_ids.append(generated_token_id) + if token_str.strip().isdigit(): + self.accumulated_value += token_str.strip() + elif self.state == FSMState.CAPTION_VALUE: + self.caption_token_count += 1 + self.accumulated_value += token_str + if '\n' in token_str: + self.caption_after_newline = True + else: + self.caption_after_newline = False + if self.caption_ending: + self.pending_field_name += token_str + if ':' in token_str or token_str.strip() == ':': + field_name_full = self.pending_field_name.strip() + field_name = field_name_full.rstrip(':').strip().lower() + if self.debug: + logger.debug(f'Detected field name after caption: {repr(field_name)}') + field_name_to_value_state = {'duration': FSMState.DURATION_VALUE, 'genres': FSMState.GENRES_VALUE, 'keyscale': FSMState.KEYSCALE_VALUE, 'language': FSMState.LANGUAGE_VALUE, 'timesignature': FSMState.TIMESIG_VALUE} + if field_name in field_name_to_value_state: + old_state = self.state + self.state = field_name_to_value_state[field_name] + self.position_in_state = 0 + self.accumulated_value = '' + self.accumulated_token_ids = [] + self.caption_ending = False + self.pending_field_name = '' + if self.debug: + logger.debug(f'FSM transition (caption ending): {old_state.name} -> {self.state.name}') + else: + if self.debug: + logger.warning(f'Unknown field name after caption: {repr(field_name)}, forcing transition') + self.caption_ending = False + self.pending_field_name = '' + self._transition_to_next_state() + elif self.state == FSMState.KEYSCALE_VALUE: + if generated_token_id == self.newline_token: + self._transition_to_next_state() + if self.state in self.fixed_strings: + return + else: + self.accumulated_token_ids.append(generated_token_id) + self.accumulated_value += token_str + elif self.state == FSMState.LANGUAGE_VALUE: + if generated_token_id == self.newline_token: + self._transition_to_next_state() + if self.state in self.fixed_strings: + return + else: + self.accumulated_token_ids.append(generated_token_id) + self.accumulated_value += token_str diff --git a/freevideo_engine/music/layers.py b/freevideo_engine/music/layers.py new file mode 100644 index 00000000..8ed84a8e --- /dev/null +++ b/freevideo_engine/music/layers.py @@ -0,0 +1,90 @@ +"""Construction boundary for linear layers and their weight-only int8 replacement. + +Handwritten DiT/FSQ/MLP projections all use linear(); native Transformers model +creation is kept in build_qwen(). Int8Linear is applied to decoder linears only +(quantizable), the same layers upstream quantizes for small cards. +""" +import torch +from torch import nn +from transformers.models.qwen3.modeling_qwen3 import Qwen3MLP + +LINEAR = nn.Linear +_FP32_EPS = torch.finfo(torch.float32).eps + + +def linear(*args, **kwargs): + return LINEAR(*args, **kwargs) + + +def quantizable(name, module): + """Upstream's int8 filter: decoder linears, never the audio tokenizer or detokenizer.""" + parts = name.split(".") + return isinstance(module, nn.Linear) and parts[0] == "decoder" and not {"tokenizer", "detokenizer"} & set(parts) + + +class Int8Linear(nn.Module): + """Weight-only int8 linear with one symmetric scale per output row. + + Rounding and the matmul follow torchao Int8WeightOnlyConfig (the upstream + int8 option) step for step, so qweight, scale and outputs are bit-identical + to it without the extra package. + """ + + def __init__(self, qweight, scale, bias=None): + super().__init__() + self.out_features, self.in_features = qweight.shape + self.register_buffer("qweight", qweight) + self.register_buffer("scale", scale) + self.bias = None if bias is None else nn.Parameter(bias, requires_grad=False) + + @classmethod + def from_weight(cls, weight, bias=None): + # The scale is computed in the weight dtype, then kept as float32. + max_abs = torch.maximum(-weight.amin(dim=1, keepdim=True).clamp(max=0), + weight.amax(dim=1, keepdim=True).clamp(min=0)) + scale = (max_abs / 127.5).clamp(min=_FP32_EPS).to(torch.float32) + qweight = torch.clamp(torch.round(weight * (1.0 / scale)), -128, 127).to(torch.int8) + return cls(qweight, scale, bias) + + def forward(self, x): + dtype = x.dtype + out = torch.mm(x.reshape(-1, x.shape[-1]), self.qweight.t().to(dtype)) + out = out * self.scale.to(out.dtype).flatten() + out = out.reshape(*x.shape[:-1], self.out_features) + if self.bias is not None: + out += self.bias + return out.to(dtype) + + +class MusicMLP(Qwen3MLP): + def __init__(self, config): + nn.Module.__init__(self) + from transformers.activations import ACT2FN + self.config = config + self.hidden_size = config.hidden_size + self.intermediate_size = config.intermediate_size + self.gate_proj = linear(self.hidden_size, self.intermediate_size, bias=False) + self.up_proj = linear(self.hidden_size, self.intermediate_size, bias=False) + self.down_proj = linear(self.intermediate_size, self.hidden_size, bias=False) + self.act_fn = ACT2FN[config.hidden_act] + + +def build_qwen(config=None, *, path=None, device, dtype): + """Only boundary that constructs native Transformers projections. + + An int8 port can replace native linears here before returning the model; + the sampling code does not need to know about their implementation. + """ + from transformers import Qwen3ForCausalLM, Qwen3Model + if path is not None: + model, info = Qwen3ForCausalLM.from_pretrained( + str(path), local_files_only=True, trust_remote_code=False, + dtype=dtype, attn_implementation="sdpa", output_loading_info=True, + ) + if any(info.get(key) for key in ("missing_keys", "unexpected_keys", "mismatched_keys", "error_msgs")): + raise RuntimeError(f"checkpoint_keys: LM: {info}") + # Upstream loads F32, then .to(device).to(dtype), which also casts + # the nonpersistent RoPE buffer. from_pretrained(dtype=...) alone + # intentionally leaves that buffer F32 in Transformers 5. + return model.to(device=device, dtype=dtype).eval().requires_grad_(False) + return Qwen3Model(config) diff --git a/freevideo_engine/music/lm.py b/freevideo_engine/music/lm.py new file mode 100644 index 00000000..d4096961 --- /dev/null +++ b/freevideo_engine/music/lm.py @@ -0,0 +1,432 @@ +"""Instrumental two-phase ACE-Step pt LM inference (MIT). + +Sampling order, CFG, YAML roundtrip and FSM are ported from llm_inference.py. +The model is the built-in Transformers Qwen3 causal LM, loaded locally. +""" +from contextlib import nullcontext +import logging +from typing import Any, Dict, List, Optional, Tuple +import time +import torch +import yaml +from transformers.generation.logits_process import LogitsProcessorList, RepetitionPenaltyLogitsProcessor +from transformers.generation.streamers import BaseStreamer +from .constrained_logits import MetadataConstrainedLogitsProcessor, FSMState +from .lm_constants import DEFAULT_LM_INSTRUCTION +logger = logging.getLogger(__name__) + +def tqdm(items, **kwargs): + return items + +class InstrumentalLM: + + def __init__(self, model, tokenizer, processor, device): + self.llm = model + self.llm_tokenizer = tokenizer + self.constrained_processor = processor + self.device = torch.device(device) + self.disable_tqdm = True + self.use_legacy_cfg_prompt = False + self.tokens = {} + + def _load_model_context(self): + return nullcontext() + + def _compute_max_new_tokens(self, target_duration, generation_phase, fallback_max=4096): + value = int(target_duration * 5) + 10 if target_duration else min(fallback_max, 3500) + return min(value, 4096 - 64) + + def generate(self, caption, seconds): + """The single-item llm_dit route with all GenerationParams defaults.""" + lyrics = "[Instrumental]" + defaults = dict( + temperature=0.85, negative_prompt="NO USER INPUT", top_k=None, + top_p=0.9, repetition_penalty=1.0, use_constrained_decoding=True, + constrained_decoding_debug=False, caption=caption, lyrics=lyrics, + ) + start = time.perf_counter() + cot = self._run_pt_single( + formatted_prompt=self.build_formatted_prompt(caption, lyrics), + cfg_scale=1.0, target_duration=None, + user_metadata={"duration": int(seconds)}, stop_at_reasoning=True, + skip_genres=True, skip_caption=False, skip_language=False, + generation_phase="cot", cot_text="", **defaults, + ) + metadata, _ = self.parse_lm_output(cot) + cot_s = time.perf_counter() - start + cot_text = self._format_metadata_as_cot(metadata) + start = time.perf_counter() + codes = self._run_pt_single( + formatted_prompt=self.build_formatted_prompt_with_cot(caption, lyrics, cot_text), + cfg_scale=2.0, target_duration=seconds, user_metadata=None, + stop_at_reasoning=False, skip_genres=False, + skip_caption=True, skip_language=True, + generation_phase="codes", cot_text=cot_text, **defaults, + ) + _, codes = self.parse_lm_output(codes) + if len(codes.split("<|audio_code_")) - 1 != int(seconds * 5): + raise RuntimeError("LM did not generate the duration-constrained audio codes") + return metadata, codes, {"cot_s": cot_s, "codes_s": time.perf_counter() - start} + + def _has_meaningful_negative_prompt(self, negative_prompt: str) -> bool: + return negative_prompt and negative_prompt.strip() and (negative_prompt.strip() != 'NO USER INPUT') + + def _build_logits_processor(self, repetition_penalty: float) -> LogitsProcessorList: + logits_processor = LogitsProcessorList() + if repetition_penalty != 1.0: + logits_processor.append(RepetitionPenaltyLogitsProcessor(penalty=repetition_penalty)) + return logits_processor + + def _setup_constrained_processor(self, use_constrained_decoding: bool, constrained_decoding_debug: bool, target_duration: Optional[float], user_metadata: Optional[Dict[str, Optional[str]]], stop_at_reasoning: bool, skip_genres: bool, skip_caption: bool, skip_language: bool, generation_phase: str, is_batch: bool=False, metadata_temperature: Optional[float]=None, codes_temperature: Optional[float]=None) -> Optional[MetadataConstrainedLogitsProcessor]: + use_phase_temperatures = not is_batch and (metadata_temperature is not None or codes_temperature is not None) + if not use_constrained_decoding and (not use_phase_temperatures): + return None + self.constrained_processor.reset() + self.constrained_processor.enabled = use_constrained_decoding + self.constrained_processor.debug = constrained_decoding_debug + if use_phase_temperatures: + self.constrained_processor.metadata_temperature = metadata_temperature + self.constrained_processor.codes_temperature = codes_temperature + else: + self.constrained_processor.metadata_temperature = None + self.constrained_processor.codes_temperature = None + self.constrained_processor.set_target_duration(target_duration) + if is_batch: + self.constrained_processor.set_user_metadata(None) + self.constrained_processor.set_stop_at_reasoning(False) + self.constrained_processor.set_skip_genres(True) + self.constrained_processor.set_skip_caption(True) + self.constrained_processor.set_skip_language(True) + else: + self.constrained_processor.set_user_metadata(user_metadata) + self.constrained_processor.set_stop_at_reasoning(stop_at_reasoning) + self.constrained_processor.set_skip_genres(skip_genres) + self.constrained_processor.set_skip_caption(skip_caption) + self.constrained_processor.set_skip_language(skip_language) + self.constrained_processor.set_generation_phase(generation_phase) + return self.constrained_processor + + def _build_unconditional_prompt(self, caption: str, lyrics: str, cot_text: str, negative_prompt: str, generation_phase: str, is_batch: bool=False) -> str: + if is_batch or generation_phase == 'codes': + prompt = self.build_formatted_prompt_with_cot(caption, lyrics, cot_text, is_negative_prompt=True, negative_prompt=negative_prompt) + else: + prompt = self.build_formatted_prompt(caption, lyrics, is_negative_prompt=True, generation_phase='cot', negative_prompt=negative_prompt) + logger.debug('Built CFG unconditional prompt') + return prompt + + def _apply_top_k_filter(self, logits: torch.Tensor, top_k: Optional[int]) -> torch.Tensor: + if top_k is not None and top_k > 0: + indices_to_remove = logits < torch.topk(logits, top_k)[0][..., -1, None] + logits[indices_to_remove] = float('-inf') + return logits + + def _apply_top_p_filter(self, logits: torch.Tensor, top_p: Optional[float]) -> torch.Tensor: + if top_p is not None and 0.0 < top_p < 1.0: + (sorted_logits, sorted_indices) = torch.sort(logits, descending=True) + cumulative_probs = torch.cumsum(torch.softmax(sorted_logits.float(), dim=-1), dim=-1) + sorted_indices_to_remove = cumulative_probs > top_p + sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone() + sorted_indices_to_remove[..., 0] = 0 + indices_to_remove = sorted_indices_to_remove.scatter(1, sorted_indices, sorted_indices_to_remove) + logits[indices_to_remove] = float('-inf') + return logits + + def _sample_tokens(self, logits: torch.Tensor, temperature: float) -> torch.Tensor: + if temperature > 0: + logits = logits.float() / temperature + probs = torch.softmax(logits, dim=-1) + return torch.multinomial(probs, num_samples=1).squeeze(1) + else: + return torch.argmax(logits, dim=-1) + + def _check_eos_token(self, tokens: torch.Tensor, eos_token_id: int, pad_token_id: Optional[int]) -> bool: + if torch.any(tokens == eos_token_id): + return True + if pad_token_id is not None and pad_token_id != eos_token_id: + if torch.any(tokens == pad_token_id): + return True + return False + + def _update_constrained_processor_state(self, constrained_processor: Optional[MetadataConstrainedLogitsProcessor], tokens: torch.Tensor): + if constrained_processor is not None: + for b in range(tokens.shape[0]): + constrained_processor.update_state(tokens[b].item()) + + def _forward_pass(self, model: Any, generated_ids: torch.Tensor, model_kwargs: Dict[str, Any], past_key_values: Optional[Any], use_cache: bool) -> Any: + if past_key_values is None: + outputs = model(input_ids=generated_ids, **model_kwargs, use_cache=use_cache) + else: + outputs = model(input_ids=generated_ids[:, -1:], past_key_values=past_key_values, **model_kwargs, use_cache=use_cache) + return outputs + + def _run_pt_single(self, formatted_prompt: str, temperature: float, cfg_scale: float, negative_prompt: str, top_k: Optional[int], top_p: Optional[float], repetition_penalty: float, use_constrained_decoding: bool, constrained_decoding_debug: bool, target_duration: Optional[float], user_metadata: Optional[Dict[str, Optional[str]]], stop_at_reasoning: bool, skip_genres: bool, skip_caption: bool, skip_language: bool, generation_phase: str, caption: str, lyrics: str, cot_text: str) -> str: + inputs = self.llm_tokenizer(formatted_prompt, return_tensors='pt', padding=False, truncation=True) + constrained_processor = self._setup_constrained_processor(use_constrained_decoding=use_constrained_decoding, constrained_decoding_debug=constrained_decoding_debug, target_duration=target_duration, user_metadata=user_metadata, stop_at_reasoning=stop_at_reasoning, skip_genres=skip_genres, skip_caption=skip_caption, skip_language=skip_language, generation_phase=generation_phase, is_batch=False) + with self._load_model_context(): + inputs = {k: v.to(self.device) for (k, v) in inputs.items()} + max_new_tokens = self._compute_max_new_tokens(target_duration=target_duration, generation_phase=generation_phase, fallback_max=getattr(self.llm.config, 'max_new_tokens', 4096)) + logits_processor = self._build_logits_processor(repetition_penalty) + if cfg_scale > 1.0: + formatted_unconditional_prompt = self._build_unconditional_prompt(caption=caption, lyrics=lyrics, cot_text=cot_text, negative_prompt=negative_prompt, generation_phase=generation_phase, is_batch=False) + batch_texts = [formatted_prompt, formatted_unconditional_prompt] + original_padding_side = self.llm_tokenizer.padding_side + self.llm_tokenizer.padding_side = 'left' + batch_inputs_tokenized = self.llm_tokenizer(batch_texts, return_tensors='pt', padding=True, truncation=True) + self.llm_tokenizer.padding_side = original_padding_side + batch_inputs_tokenized = {k: v.to(self.device) for (k, v) in batch_inputs_tokenized.items()} + batch_input_ids = batch_inputs_tokenized['input_ids'] + batch_attention_mask = batch_inputs_tokenized.get('attention_mask', None) + outputs = self._generate_with_cfg_custom(batch_input_ids=batch_input_ids, batch_attention_mask=batch_attention_mask, max_new_tokens=max_new_tokens, temperature=temperature, cfg_scale=cfg_scale, top_k=top_k, top_p=top_p, repetition_penalty=repetition_penalty, pad_token_id=self.llm_tokenizer.pad_token_id or self.llm_tokenizer.eos_token_id, streamer=None, constrained_processor=constrained_processor) + outputs = outputs[0:1] + else: + outputs = self._generate_with_constrained_decoding(input_ids=inputs['input_ids'], attention_mask=inputs.get('attention_mask'), max_new_tokens=max_new_tokens, temperature=temperature, top_k=top_k, top_p=top_p, repetition_penalty=repetition_penalty, pad_token_id=self.llm_tokenizer.pad_token_id or self.llm_tokenizer.eos_token_id, streamer=None, constrained_processor=constrained_processor) + if isinstance(outputs, torch.Tensor): + if outputs.dim() == 2: + generated_ids = outputs[0] + else: + generated_ids = outputs + else: + generated_ids = outputs[0] + if cfg_scale > 1.0: + input_length = batch_inputs_tokenized['input_ids'].shape[1] + else: + input_length = inputs['input_ids'].shape[1] + generated_ids = generated_ids[input_length:] + self.tokens[generation_phase] = generated_ids.detach().cpu().clone() + if generated_ids.device.type != 'cpu': + generated_ids = generated_ids.cpu() + output_text = self.llm_tokenizer.decode(generated_ids, skip_special_tokens=False) + return output_text + + def _format_metadata_as_cot(self, metadata: Dict[str, Any]) -> str: + cot_items = {} + for key in ['bpm', 'caption', 'duration', 'keyscale', 'language', 'timesignature']: + if key in metadata and metadata[key] is not None: + value = metadata[key] + if key == 'timesignature' and value.endswith('/4'): + value = value.split('/')[0] + if isinstance(value, str) and value.isdigit(): + value = int(value) + cot_items[key] = value + if len(cot_items) > 0: + cot_yaml = yaml.dump(cot_items, allow_unicode=True, sort_keys=True).strip() + else: + cot_yaml = '' + return f'\n{cot_yaml}\n' + + def build_formatted_prompt(self, caption: str, lyrics: str='', is_negative_prompt: bool=False, generation_phase: str='cot', negative_prompt: str='NO USER INPUT') -> str: + if self.llm_tokenizer is None: + raise ValueError('LLM tokenizer is not initialized. Call initialize() first.') + if is_negative_prompt: + has_negative_prompt = self._has_meaningful_negative_prompt(negative_prompt) + if generation_phase == 'cot': + if has_negative_prompt: + prompt = f'# Caption\n{negative_prompt}\n\n# Lyric\n{lyrics}\n' + else: + prompt = f'# Lyric\n{lyrics}\n' + else: + prompt = caption + else: + prompt = f'# Caption\n{caption}\n\n# Lyric\n{lyrics}\n' + return self.llm_tokenizer.apply_chat_template([{'role': 'system', 'content': f'# Instruction\n{DEFAULT_LM_INSTRUCTION}\n\n'}, {'role': 'user', 'content': prompt}], tokenize=False, add_generation_prompt=True) + + def build_formatted_prompt_with_cot(self, caption: str, lyrics: str, cot_text: str, is_negative_prompt: bool=False, negative_prompt: str='NO USER INPUT') -> str: + if self.llm_tokenizer is None: + raise ValueError('LLM tokenizer is not initialized. Call initialize() first.') + if is_negative_prompt: + has_negative_prompt = self._has_meaningful_negative_prompt(negative_prompt) + cot_for_prompt = '\n\n' + user_prompt = negative_prompt if has_negative_prompt else 'NO USER INPUT' + else: + cot_for_prompt = cot_text + user_prompt = f'# Caption\n{caption}\n\n# Lyric\n{lyrics}\n' + formatted = self.llm_tokenizer.apply_chat_template([{'role': 'system', 'content': f'# Instruction\n{DEFAULT_LM_INSTRUCTION}\n\n'}, {'role': 'user', 'content': user_prompt}], tokenize=False, add_generation_prompt=True) + formatted += cot_for_prompt + '\n\n' + return formatted + + def _generate_with_constrained_decoding(self, input_ids: torch.Tensor, attention_mask: Optional[torch.Tensor], max_new_tokens: int, temperature: float, top_k: Optional[int], top_p: Optional[float], repetition_penalty: float, pad_token_id: int, streamer: Optional[BaseStreamer], constrained_processor: Optional[MetadataConstrainedLogitsProcessor]=None) -> torch.Tensor: + model = self.llm + device = self.device + generated_ids = input_ids.clone() + if attention_mask is not None: + attn_mask = attention_mask.clone() + else: + attn_mask = torch.ones_like(input_ids) + model_kwargs = {'attention_mask': attn_mask} + past_key_values = None + use_cache = hasattr(model, 'generation_config') and getattr(model.generation_config, 'use_cache', True) + eos_token_id = self.llm_tokenizer.eos_token_id + if eos_token_id is None: + eos_token_id = pad_token_id + logits_processor = self._build_logits_processor(repetition_penalty) + with torch.inference_mode(): + for step in tqdm(range(max_new_tokens), desc='LLM Constrained Decoding', unit='token', disable=self.disable_tqdm): + outputs = self._forward_pass(model, generated_ids, model_kwargs, past_key_values, use_cache) + next_token_logits = outputs.logits[:, -1, :] + if constrained_processor is not None: + next_token_logits = constrained_processor(generated_ids, next_token_logits) + for processor in logits_processor: + next_token_logits = processor(generated_ids, next_token_logits) + next_token_logits = self._apply_top_k_filter(next_token_logits, top_k) + next_token_logits = self._apply_top_p_filter(next_token_logits, top_p) + next_tokens = self._sample_tokens(next_token_logits, temperature) + self._update_constrained_processor_state(constrained_processor, next_tokens) + should_stop = self._check_eos_token(next_tokens, eos_token_id, pad_token_id) + next_tokens_unsqueezed = next_tokens.unsqueeze(1) + generated_ids = torch.cat([generated_ids, next_tokens_unsqueezed], dim=1) + attn_mask = torch.cat([attn_mask, torch.ones((input_ids.shape[0], 1), device=device, dtype=attn_mask.dtype)], dim=1) + model_kwargs['attention_mask'] = attn_mask + if use_cache and hasattr(outputs, 'past_key_values'): + past_key_values = outputs.past_key_values + if streamer is not None: + streamer.put(next_tokens_unsqueezed) + if should_stop: + break + if streamer is not None: + streamer.end() + del past_key_values + return generated_ids + + def _generate_with_cfg_custom(self, batch_input_ids: torch.Tensor, batch_attention_mask: Optional[torch.Tensor], max_new_tokens: int, temperature: float, cfg_scale: float, top_k: Optional[int], top_p: Optional[float], repetition_penalty: float, pad_token_id: int, streamer: Optional[BaseStreamer], constrained_processor: Optional[MetadataConstrainedLogitsProcessor]=None) -> torch.Tensor: + from .constrained_logits import FSMState + model = self.llm + device = self.device + batch_size = batch_input_ids.shape[0] // 2 + cond_start_idx = 0 + uncond_start_idx = batch_size + generated_ids = batch_input_ids.clone() + if batch_attention_mask is not None: + attention_mask = batch_attention_mask.clone() + else: + attention_mask = torch.ones_like(batch_input_ids) + model_kwargs = {} + if batch_attention_mask is not None: + model_kwargs['attention_mask'] = attention_mask + past_key_values = None + use_cache = hasattr(model, 'generation_config') and getattr(model.generation_config, 'use_cache', True) + eos_token_id = self.llm_tokenizer.eos_token_id + if eos_token_id is None: + eos_token_id = pad_token_id + logits_processor = self._build_logits_processor(repetition_penalty) + seq_finished = torch.zeros(batch_size, dtype=torch.bool, device=device) + with torch.inference_mode(): + for step in tqdm(range(max_new_tokens), desc='LLM CFG Generation', unit='token', disable=self.disable_tqdm): + outputs = self._forward_pass(model, generated_ids, model_kwargs, past_key_values, use_cache) + next_token_logits = outputs.logits[:, -1, :] + cond_logits = next_token_logits[cond_start_idx:cond_start_idx + batch_size] + uncond_logits = next_token_logits[uncond_start_idx:uncond_start_idx + batch_size] + if constrained_processor is not None and constrained_processor.state == FSMState.CODES_GENERATION and (constrained_processor.non_audio_code_mask is not None): + non_audio_mask = constrained_processor.non_audio_code_mask + if non_audio_mask.device != device or non_audio_mask.dtype != torch.float32: + non_audio_mask = non_audio_mask.to(device=device, dtype=torch.float32) + valid_mask = non_audio_mask[0] == 0 + valid_indices = valid_mask.nonzero(as_tuple=False).squeeze(1) + cond_valid = cond_logits[:, valid_indices].float() + uncond_valid = uncond_logits[:, valid_indices].float() + cfg_valid = uncond_valid + cfg_scale * (cond_valid - uncond_valid) + cfg_logits = torch.full((batch_size, cond_logits.shape[1]), float('-inf'), device=device, dtype=torch.float32) + cfg_logits[:, valid_indices] = cfg_valid + else: + cfg_logits = uncond_logits.float() + cfg_scale * (cond_logits.float() - uncond_logits.float()) + cfg_logits = torch.nan_to_num(cfg_logits, nan=float('-inf')) + if constrained_processor is not None: + current_input_ids = generated_ids[cond_start_idx:cond_start_idx + batch_size] + cfg_logits = constrained_processor(current_input_ids, cfg_logits) + current_input_ids = generated_ids[cond_start_idx:cond_start_idx + batch_size] + for processor in logits_processor: + cfg_logits = processor(current_input_ids, cfg_logits) + cfg_logits = self._apply_top_k_filter(cfg_logits, top_k) + cfg_logits = self._apply_top_p_filter(cfg_logits, top_p) + if seq_finished.any(): + for b in range(batch_size): + if seq_finished[b]: + cfg_logits[b, :] = float('-inf') + cfg_logits[b, eos_token_id] = 0.0 + next_tokens = self._sample_tokens(cfg_logits, temperature) + self._update_constrained_processor_state(constrained_processor, next_tokens) + is_eos = next_tokens == eos_token_id + if pad_token_id is not None and pad_token_id != eos_token_id: + is_eos = is_eos | (next_tokens == pad_token_id) + seq_finished = seq_finished | is_eos + next_tokens_unsqueezed = next_tokens.unsqueeze(1) + generated_ids = torch.cat([generated_ids, next_tokens_unsqueezed.repeat(2, 1)], dim=1) + attention_mask = torch.cat([attention_mask, torch.ones((batch_size * 2, 1), device=device, dtype=attention_mask.dtype)], dim=1) + model_kwargs['attention_mask'] = attention_mask + if use_cache and hasattr(outputs, 'past_key_values'): + past_key_values = outputs.past_key_values + if streamer is not None: + streamer.put(next_tokens_unsqueezed) + if seq_finished.all(): + break + if streamer is not None: + streamer.end() + del past_key_values + return generated_ids + + def parse_lm_output(self, output_text: str) -> Tuple[Dict[str, Any], str]: + debug_output_text = output_text.split('')[0] + logger.debug('Parsing LM metadata') + metadata = {} + audio_codes = '' + import re + code_pattern = '<\\|audio_code_\\d+\\|>' + code_matches = re.findall(code_pattern, output_text) + if code_matches: + audio_codes = ''.join(code_matches) + reasoning_patterns = ['(.*?)', '(.*?)', '(.*?)'] + reasoning_text = None + for pattern in reasoning_patterns: + match = re.search(pattern, output_text, re.DOTALL) + if match: + reasoning_text = match.group(1).strip() + break + if not reasoning_text: + lines_before_codes = output_text.split('<|audio_code_')[0] if '<|audio_code_' in output_text else output_text + reasoning_text = lines_before_codes.strip() + if reasoning_text: + lines = reasoning_text.split('\n') + current_key = None + current_value_lines = [] + + def save_current_field(): + nonlocal current_key, current_value_lines + if current_key and current_value_lines: + value = '\n'.join(current_value_lines) + if current_key == 'bpm': + try: + metadata['bpm'] = int(value.strip()) + except: + metadata['bpm'] = value.strip() + elif current_key == 'caption': + metadata['caption'] = MetadataConstrainedLogitsProcessor.postprocess_caption(value) + elif current_key == 'duration': + try: + metadata['duration'] = int(value.strip()) + except: + metadata['duration'] = value.strip() + elif current_key == 'genres': + metadata['genres'] = value.strip() + elif current_key == 'keyscale': + metadata['keyscale'] = value.strip() + elif current_key == 'language': + metadata['language'] = value.strip() + elif current_key == 'timesignature': + metadata['timesignature'] = value.strip() + current_key = None + current_value_lines = [] + for line in lines: + if line.strip().startswith('<'): + continue + if line and (not line[0].isspace()) and (':' in line): + save_current_field() + parts = line.split(':', 1) + if len(parts) == 2: + current_key = parts[0].strip().lower() + first_value = parts[1] + if first_value.strip(): + current_value_lines.append(first_value) + elif line.startswith(' ') or line.startswith('\t'): + if current_key: + current_value_lines.append(line) + save_current_field() + return (metadata, audio_codes) diff --git a/freevideo_engine/music/lm_constants.py b/freevideo_engine/music/lm_constants.py new file mode 100644 index 00000000..4c93b923 --- /dev/null +++ b/freevideo_engine/music/lm_constants.py @@ -0,0 +1,72 @@ +# ACE-Step 1.5 metadata constants (MIT). +""" +Constants for ACE-Step +Centralized constants used across the codebase +""" + +# ============================================================================== +# Language Constants +# ============================================================================== + +# Supported languages for vocal generation and language detection +# Covers major world languages with good TTS support in the underlying model +# 'unknown' is used when language cannot be determined automatically +VALID_LANGUAGES = [ + 'ar', 'az', 'bg', 'bn', 'ca', 'cs', 'da', 'de', 'el', 'en', + 'es', 'fa', 'fi', 'fr', 'he', 'hi', 'hr', 'ht', 'hu', 'id', + 'is', 'it', 'ja', 'ko', 'la', 'lt', 'ms', 'ne', 'nl', 'no', + 'pa', 'pl', 'pt', 'ro', 'ru', 'sa', 'sk', 'sr', 'sv', 'sw', + 'ta', 'te', 'th', 'tl', 'tr', 'uk', 'ur', 'vi', 'yue', 'zh', + 'unknown' +] + + +# ============================================================================== +# Keyscale Constants +# ============================================================================== + +# Musical note names using standard Western notation +KEYSCALE_NOTES = ['A', 'B', 'C', 'D', 'E', 'F', 'G'] + +# Supported accidentals: natural, ASCII sharp/flat, Unicode sharp/flat +KEYSCALE_ACCIDENTALS = ['', '#', 'b', '♯', '♭'] # empty + ASCII sharp/flat + Unicode sharp/flat + +# Major and minor scale modes +KEYSCALE_MODES = ['major', 'minor'] + +# Generate all valid keyscales: 7 notes × 5 accidentals × 2 modes = 70 combinations +# Examples: "C major", "F# minor", "B♭ major" +VALID_KEYSCALES = set() +for note in KEYSCALE_NOTES: + for acc in KEYSCALE_ACCIDENTALS: + for mode in KEYSCALE_MODES: + VALID_KEYSCALES.add(f"{note}{acc} {mode}") + + +# ============================================================================== +# Metadata Range Constants +# ============================================================================== + +# BPM (Beats Per Minute) range - covers most musical styles +# 30 BPM: Very slow ballads, ambient music +# 300 BPM: Fast electronic dance music, extreme metal +BPM_MIN = 30 +BPM_MAX = 300 + +# Duration range (in seconds) - balances quality vs. computational cost +# 10s: Short loops, musical excerpts +# 600s: Full songs, extended compositions (10 minutes) +DURATION_MIN = 10 +DURATION_MAX = 600 + +# Valid time signatures - common musical meter patterns +# 2: 2/4 time (marches, polka) +# 3: 3/4 time (waltzes, ballads) +# 4: 4/4 time (most pop, rock, hip-hop) +# 6: 6/8 time (compound time, folk dances) +VALID_TIME_SIGNATURES = [2, 3, 4, 6] + + +# ============================================================================== + +DEFAULT_LM_INSTRUCTION = "Generate audio semantic tokens based on the given conditions:" diff --git a/freevideo_engine/music/mix.py b/freevideo_engine/music/mix.py new file mode 100644 index 00000000..71d36788 --- /dev/null +++ b/freevideo_engine/music/mix.py @@ -0,0 +1,84 @@ +"""Lay a generated piece under a cut: the cut's exact length, a level well under the video's own sound, +lower again while someone speaks, and a fade at each end. There is no volume or fade edit on the +timeline, so all of it is written into the file.""" +from fractions import Fraction +import os +from pathlib import Path + +import numpy as np + +RATE = 48000 +# Generated video sound sits near -13 dBFS RMS; the score stays about 14 dB under it. +LEVEL_DB = -27. +PEAK_DB = -6. +QUIET_DB = -8. +RAMP_SECONDS = .3 +FADE_IN_SECONDS, FADE_OUT_SECONDS = 1., 2. + + +def _db(value): + return 10. ** (value / 20.) + + +def envelope(length, quiet=(), fps=24): + """Gain per sample: 1, and QUIET_DB lower over each (start, frames) range, with short ramps.""" + gain = np.ones(length, np.float32) + ramp = max(1, round(RAMP_SECONDS * RATE)) + low = _db(QUIET_DB) + for start, frames in quiet: + first, last = round(start * RATE / fps), round((start + frames) * RATE / fps) + begin, end = max(0, first - ramp), min(length, last + ramp) + if begin >= end: + continue + position = np.arange(begin, end, dtype=np.float64) + # Low inside, straight lines across the ramp just before and just after. + into = np.clip((position - (first - ramp)) / ramp, 0, 1) + out_of = np.clip(((last + ramp) - position) / ramp, 0, 1) + gain[begin:end] = np.minimum(gain[begin:end], 1 - (1 - low) * np.minimum(into, out_of)) + return gain + + +def finish(samples, frames, fps=24, quiet=()): + """float32 [2, n] at RATE → float32 [2, frames / fps seconds], levelled, lowered and faded.""" + length = round(frames * RATE / fps) + samples = np.asarray(samples, np.float32)[:, :length] + if samples.shape[1] < length: + samples = np.pad(samples, ((0, 0), (0, length - samples.shape[1]))) + rms = float(np.sqrt(np.mean(np.square(samples, dtype=np.float64)))) + peak = float(np.max(np.abs(samples))) if length else 0. + if rms <= 0 or peak <= 0: + raise ValueError('silent') + gain = min(_db(LEVEL_DB) / rms, _db(PEAK_DB) / peak) + shape = envelope(length, quiet, fps) + # A short cut keeps most of its length at full level: each fade takes at most a quarter of it. + fade_in = min(round(FADE_IN_SECONDS * RATE), length // 4) + fade_out = min(round(FADE_OUT_SECONDS * RATE), length // 4) + if fade_in: + shape[:fade_in] *= np.linspace(0, 1, fade_in, endpoint=False, dtype=np.float32) + if fade_out: + shape[length - fade_out:] *= np.linspace(1, 0, fade_out, dtype=np.float32) + return samples * np.float32(gain) * shape + + +def write_mp4(path, samples, bit_rate=192_000): + """AAC in an audio-only MP4, the container the timeline takes for generated media; replaced atomically.""" + import av + path = Path(path) + samples = np.asarray(samples, np.float32) + partial = path.with_name('.' + path.name + '.partial') + try: + with av.open(str(partial), 'w', format='mp4', options={'movflags': 'faststart'}) as writer: + stream = writer.add_stream('aac', rate=RATE, layout='stereo') + stream.codec_context.bit_rate = bit_rate + step = 1024 * 32 + for offset in range(0, samples.shape[1], step): + frame = av.AudioFrame.from_ndarray(np.ascontiguousarray(samples[:, offset:offset + step]), + format='fltp', layout='stereo') + frame.sample_rate, frame.pts, frame.time_base = RATE, offset, Fraction(1, RATE) + for packet in stream.encode(frame): + writer.mux(packet) + for packet in stream.encode(): + writer.mux(packet) + os.replace(partial, path) + finally: + partial.unlink(missing_ok=True) diff --git a/freevideo_engine/music/pipeline.py b/freevideo_engine/music/pipeline.py new file mode 100644 index 00000000..77c055d7 --- /dev/null +++ b/freevideo_engine/music/pipeline.py @@ -0,0 +1,601 @@ +"""Local instrumental ACE-Step XL turbo + 1.7B LM; no checkpoint Python. + +Prompt formatting, conditioning and overlap-discard decode are adapted from +ACE-Step 1.5 (MIT, see licenses/ACE-Step-MIT.txt). The supported dependency set follows +upstream's no-pytorch_wavelets DCW fallback. All attention uses PyTorch SDPA. +""" +from contextlib import contextmanager +import gc +import json +import math +import os +from pathlib import Path +import re +import tempfile +import time +import warnings +import wave + +import numpy as np +import torch +from safetensors import safe_open + +SAMPLE_RATE = 48000 +# Exact checkpoint keys omitted from the instrumental LM-hint path. Never use +# prefix exemptions here: a renamed or newly added checkpoint tensor must fail. +UNUSED_DIT_KEYS = frozenset({ + 'null_condition_emb', + 'tokenizer.attention_pooler.embed_tokens.bias', + 'tokenizer.attention_pooler.embed_tokens.weight', + 'tokenizer.attention_pooler.layers.0.input_layernorm.weight', + 'tokenizer.attention_pooler.layers.0.mlp.down_proj.weight', + 'tokenizer.attention_pooler.layers.0.mlp.gate_proj.weight', + 'tokenizer.attention_pooler.layers.0.mlp.up_proj.weight', + 'tokenizer.attention_pooler.layers.0.post_attention_layernorm.weight', + 'tokenizer.attention_pooler.layers.0.self_attn.k_norm.weight', + 'tokenizer.attention_pooler.layers.0.self_attn.k_proj.weight', + 'tokenizer.attention_pooler.layers.0.self_attn.o_proj.weight', + 'tokenizer.attention_pooler.layers.0.self_attn.q_norm.weight', + 'tokenizer.attention_pooler.layers.0.self_attn.q_proj.weight', + 'tokenizer.attention_pooler.layers.0.self_attn.v_proj.weight', + 'tokenizer.attention_pooler.layers.1.input_layernorm.weight', + 'tokenizer.attention_pooler.layers.1.mlp.down_proj.weight', + 'tokenizer.attention_pooler.layers.1.mlp.gate_proj.weight', + 'tokenizer.attention_pooler.layers.1.mlp.up_proj.weight', + 'tokenizer.attention_pooler.layers.1.post_attention_layernorm.weight', + 'tokenizer.attention_pooler.layers.1.self_attn.k_norm.weight', + 'tokenizer.attention_pooler.layers.1.self_attn.k_proj.weight', + 'tokenizer.attention_pooler.layers.1.self_attn.o_proj.weight', + 'tokenizer.attention_pooler.layers.1.self_attn.q_norm.weight', + 'tokenizer.attention_pooler.layers.1.self_attn.q_proj.weight', + 'tokenizer.attention_pooler.layers.1.self_attn.v_proj.weight', + 'tokenizer.attention_pooler.norm.weight', + 'tokenizer.attention_pooler.special_token', + 'tokenizer.audio_acoustic_proj.bias', + 'tokenizer.audio_acoustic_proj.weight', + 'tokenizer.quantizer.project_in.bias', + 'tokenizer.quantizer.project_in.weight', + # The BF16 condition encoder file also carries silence_latent.pt's tensor. + 'silence_latent', +}) +# Release folder: config, silence latent and the one F32 shard holding the +# detokenizer and FSQ. Decoder and encoder come from the official BF16 copy; +# every one of its tensors equals the F32 release cast to BF16. +DIT_NAME = "acestep-v15-xl-turbo" +DIT_BF16_NAME = "acestep-v15-xl-turbo-bf16" +DETOKENIZER_SHARD = "model-00004-of-00004.safetensors" +TRANSFORMER_SHARDS = ("diffusion_pytorch_model-00001-of-00002.safetensors", + "diffusion_pytorch_model-00002-of-00002.safetensors") +CONDITION_ENCODER_FILE = "condition_encoder/diffusion_pytorch_model.safetensors" +# Diffusers names of the BF16 copy -> release checkpoint names. +_BF16_RENAMES = ((r"\.to_q\.", ".q_proj."), (r"\.to_k\.", ".k_proj."), (r"\.to_v\.", ".v_proj."), + (r"\.to_out\.0\.", ".o_proj."), (r"\.norm_q\.", ".q_norm."), (r"\.norm_k\.", ".k_norm."), + (r"^proj_in_conv\.", "proj_in.1."), (r"^proj_out_conv\.", "proj_out.1.")) +LM_NAME = "acestep-5Hz-lm-1.7B" +# resident: everything stays on the GPU. staged: one component on the GPU at a +# time. int8: staged, with weight-only int8 decoder linears. stream: staged, +# with DiT blocks read from the files right before their forward. +MODES = ("resident", "staged", "int8", "stream") + + +class BadRequest(ValueError): + """Invalid music generation parameters.""" + + +def validate_request(caption, seconds, seed): + if not isinstance(caption, str) or not caption.strip(): + raise BadRequest("caption must be a nonempty string") + if (isinstance(seconds, bool) or not isinstance(seconds, (int, float)) + or not math.isfinite(seconds) or not 10 <= seconds <= 600): + raise BadRequest("seconds must be between 10 and 600") + if isinstance(seed, bool) or not isinstance(seed, int) or not 0 <= seed <= 2**64 - 1: + raise BadRequest("seed must be an integer between 0 and 2**64-1") + + +def validate_keys(expected, actual, unused=frozenset()): + """Reject every missing active key and every unexpected non-allowlisted key.""" + expected, actual = set(expected), set(actual) + missing = expected - actual + unexpected = actual - expected - unused + if missing or unexpected: + raise RuntimeError(f"checkpoint_keys: missing={sorted(missing)}, unexpected={sorted(unexpected)}") + + +def _files(root): + """Return {component: {checkpoint key: (file, key in file)}} after checking every file exists.""" + paths = { + "lm": root / LM_NAME / "model.safetensors", + "text": root / "Qwen3-Embedding-0.6B/model.safetensors", + "vae": root / "vae/diffusion_pytorch_model.safetensors", + } + dit_files = [root / DIT_BF16_NAME / "transformer" / name for name in TRANSFORMER_SHARDS] + [ + root / DIT_BF16_NAME / CONDITION_ENCODER_FILE, root / DIT_NAME / DETOKENIZER_SHARD] + required = list(paths.values()) + dit_files + [ + root / DIT_NAME / "config.json", root / DIT_NAME / "silence_latent.pt", + root / LM_NAME / "config.json", root / LM_NAME / "tokenizer.json", + root / LM_NAME / "tokenizer_config.json", + root / "Qwen3-Embedding-0.6B/config.json", + root / "Qwen3-Embedding-0.6B/tokenizer.json", + root / "Qwen3-Embedding-0.6B/tokenizer_config.json", + root / "vae/config.json", + ] + missing = [str(p) for p in required if not p.is_file()] + if missing: + raise FileNotFoundError("Missing model files: " + ", ".join(missing)) + # Files are data only: no model code is imported from the model folder. + weights = {kind: _weight_map(path) for kind, path in paths.items()} + weights["dit"] = _dit_weight_map(root) + return weights + + +def _weight_map(path): + with safe_open(str(path), framework="pt", device="cpu") as source: + return {k: (path, k) for k in source.keys()} + + +def _bf16_key(part, key): + for pattern, replacement in _BF16_RENAMES: + key = re.sub(pattern, replacement, key) + if part == "transformer": + return "decoder." + key + return key if key in ("null_condition_emb", "silence_latent") else "encoder." + key + + +def _dit_weight_map(root): + """BF16 decoder/encoder from the official BF16 copy; detokenizer and FSQ from the F32 shard. + + The F32 shard also holds some decoder/encoder tensors; those duplicate the + BF16 copy and are never read. + """ + weights = {} + sources = [("transformer", root / DIT_BF16_NAME / "transformer" / name) for name in TRANSFORMER_SHARDS] + sources += [("condition_encoder", root / DIT_BF16_NAME / CONDITION_ENCODER_FILE)] + for part, path in sources: + for key in _weight_map(path): + name = _bf16_key(part, key) + if name in weights: + raise RuntimeError(f"checkpoint_keys: duplicated key {name}") + weights[name] = (path, key) + for key, value in _weight_map(root / DIT_NAME / DETOKENIZER_SHARD).items(): + if key.split(".")[0] in ("detokenizer", "tokenizer"): + if key in weights: + raise RuntimeError(f"checkpoint_keys: duplicated key {key}") + weights[key] = value + return weights + + +def _read(weights, name, device, dtype): + path, key = weights[name] + with safe_open(str(path), framework="pt", device="cpu") as source: + return source.get_tensor(key).to(device=device, dtype=dtype) + + +def _read_json(path): + with path.open(encoding="utf-8") as stream: + return json.load(stream) + + +def _load_weights(model, weight_map, unused=frozenset()): + state = model.state_dict() + validate_keys(state, weight_map, unused) + if all(t.device.type == "meta" for t in state.values()): + # Shape-only DiT: validate every shape without materializing all 9 GiB. + for shard in {path for path, _ in weight_map.values()}: + with safe_open(str(shard), framework="pt", device="cpu") as source: + for key, target in state.items(): + path, name = weight_map[key] + if path == shard and tuple(source.get_slice(name).get_shape()) != tuple(target.shape): + raise RuntimeError(f"checkpoint_shape: {key}") + return model.eval().requires_grad_(False) + # Close the mapping after each tensor to bound resident host memory. + with torch.no_grad(): + for key, target in state.items(): + path, name = weight_map[key] + with safe_open(str(path), framework="pt", device="cpu") as source: + value = source.get_tensor(name) + if value.shape != target.shape: + raise RuntimeError(f"checkpoint_shape: {key}: {value.shape} != {target.shape}") + target.copy_(value) + del value + return model.eval().requires_grad_(False) + + +def _materialize(model, weight_map, device, dtype, quantize): + """Fill a shape-only (meta) model from local files, one tensor on the host at a time. + + With quantize, each decoder linear becomes Int8Linear as soon as its BF16 + weight arrives, so the BF16 decoder (7.77 GiB) never exists on the device. + Constructor buffers that are not in the files are real CPU data and move. + """ + from .layers import Int8Linear, quantizable + with torch.no_grad(): + for name, module in list(model.named_modules()): + prefix = name + "." if name else "" + if quantize and quantizable(name, module): + bias = None if module.bias is None else _read(weight_map, prefix + "bias", device, dtype) + layer = Int8Linear.from_weight(_read(weight_map, prefix + "weight", device, dtype), bias) + parent, _, leaf = name.rpartition(".") + setattr(model.get_submodule(parent), leaf, layer) + continue + for leaf, parameter in list(module.named_parameters(recurse=False)): + value = _read(weight_map, prefix + leaf, device, parameter.dtype) + module._parameters[leaf] = torch.nn.Parameter(value, requires_grad=False) + for leaf, buffer in list(module.named_buffers(recurse=False)): + if prefix + leaf in weight_map: + module._buffers[leaf] = _read(weight_map, prefix + leaf, device, buffer.dtype) + elif buffer is not None: + module._buffers[leaf] = buffer.to(device) + if any(t.device.type == "meta" for t in list(model.parameters()) + list(model.buffers())): + raise RuntimeError("checkpoint_keys: tensor left unmaterialized") + return model + + +def _construct(kind, root, weights, device, dtype): + from diffusers import AutoencoderOobleck + from transformers import Qwen3Config + from .layers import build_qwen + from transformers.initialization import no_init_weights + from transformers.models.qwen3.modeling_qwen3 import Qwen3RotaryEmbedding + from .acestep_model import AceStepConfig, AceStepConditionGenerationModel, AudioCodeQuantizer + + if kind == "lm": + return build_qwen(path=root / LM_NAME, device=device, dtype=dtype) + + # Ordinary components allocate directly on the target. Sequential DiT uses + # meta parameters with real CPU constructor buffers; blocks materialize + # directly from local shards immediately before their CUDA forward. + previous_dtype = torch.get_default_dtype() + try: + torch.set_default_dtype(dtype) + with torch.device(device), no_init_weights(): + if kind == "dit": + config = AceStepConfig(**_read_json(root / DIT_NAME / "config.json")) + config._attn_implementation = "sdpa" + model = AceStepConditionGenerationModel(config) + elif kind == "text": + config = Qwen3Config(**_read_json(root / "Qwen3-Embedding-0.6B/config.json")) + config._attn_implementation = "sdpa" + model = build_qwen(config, device=device, dtype=dtype) + else: + model = AutoencoderOobleck.from_config(_read_json(root / "vae/config.json")) + finally: + torch.set_default_dtype(previous_dtype) + # Upstream constructs RoPE on CPU before moving to CUDA. Computing pow on + # CUDA can round differently; reproduce that small constructor buffer on CPU. + for module in model.modules(): + if isinstance(module, Qwen3RotaryEmbedding): + with torch.device("cpu"): + inv_freq, scale = module.compute_default_rope_parameters(module.config) + buffer_device = "cpu" if torch.device(device).type == "meta" else device + module.inv_freq = inv_freq.to(device=buffer_device, dtype=dtype) + module.original_inv_freq = module.inv_freq.clone() + module.attention_scaling = scale + for module in model.modules(): + if isinstance(module, AudioCodeQuantizer): + buffer_device = "cpu" if torch.device(device).type == "meta" else device + for name, value in module.named_buffers(): + module._buffers[name] = value.to(device=buffer_device, dtype=dtype) + return _load_weights(model, weights, UNUSED_DIT_KEYS if kind == "dit" else frozenset()) + + +def _prompt(caption, seconds, metadata): + # handler._resolve_generate_music_task uses this instruction whenever LM + # audio codes are supplied, including ordinary instrumental text2music. + metas = (f"- bpm: {metadata.get('bpm') or 'N/A'}\n" + f"- timesignature: {metadata.get('timesignature') or 'N/A'}\n" + f"- keyscale: {metadata.get('keyscale') or 'N/A'}\n" + f"- duration: {int(seconds)} seconds\n") + return ("# Instruction\nGenerate audio semantic tokens based on the given conditions:\n\n" + f"# Caption\n{caption}\n\n# Metas\n{metas}<|endoftext|>\n") + + +class MusicPipeline: + """One request at a time; progress callbacks receive (fraction, desc=...). + + Sequential mode releases each component completely, then reopens local + safetensors for the next stage. This bounds host memory as well as VRAM; + stage weight reads are included in generate time. load preflights local + file headers; full key/shape validation happens when a stage is opened. + """ + + def __init__(self, root, device, mode): + from transformers import PreTrainedTokenizerFast + if mode not in MODES: + raise ValueError(f"mode must be one of {MODES}") + self.root = Path(root) + self.weights = _files(self.root) + self.device = torch.device(device) + self.dtype = torch.bfloat16 if self.device.type == "cuda" else torch.float32 + self.mode = mode + self.sequential = sequential = mode != "resident" + self.components = {} + self.component_bytes = {} + self.tokenizer = PreTrainedTokenizerFast.from_pretrained( + str(self.root / "Qwen3-Embedding-0.6B"), local_files_only=True, + ) + # weights_only disallows pickle globals; this is a tensor, not code. + self.silence = torch.load(self.root / DIT_NAME / "silence_latent.pt", + map_location="cpu", weights_only=True) + if not isinstance(self.silence, torch.Tensor) or self.silence.ndim != 3: + raise RuntimeError("checkpoint_shape: silence_latent") + # Official loader transposes [1,64,T] -> [1,T,64]. + self.silence = self.silence.transpose(1, 2).to(self.dtype) + if self.silence.shape[0] != 1 or self.silence.shape[2] != 64 or self.silence.shape[1] < 750: + raise RuntimeError("checkpoint_shape: silence_latent") + from transformers import AutoTokenizer + from .constrained_logits import MetadataConstrainedLogitsProcessor + self.lm_tokenizer = AutoTokenizer.from_pretrained( + str(self.root / LM_NAME), local_files_only=True, trust_remote_code=False, use_fast=True, + ) + self.lm_processor = MetadataConstrainedLogitsProcessor(self.lm_tokenizer, max_duration=600) + self.stage_times = {} + self.last_lm_tokens = {} + self.last_metadata = {} + # Sequential load validates file headers/keys without loading 10 GiB of + # DiT twice. Shape validation takes place when each component is opened. + for kind in ("lm", "dit", "vae", "text"): + if sequential: + size = 0 + element = torch.tensor([], dtype=self.dtype).element_size() + for shard in {path for path, _ in self.weights[kind].values()}: + with safe_open(str(shard), framework="pt", device="cpu") as source: + for key, (path, name) in self.weights[kind].items(): + if path == shard: + size += math.prod(source.get_slice(name).get_shape()) * element + self.component_bytes[kind] = size + else: + component = _construct(kind, self.root, self.weights[kind], self.device, self.dtype) + self.component_bytes[kind] = sum(t.numel() * t.element_size() for t in component.parameters()) + self.components[kind] = component + self.decode_chunk_size = self._chunk_size() + warnings.warn("ACE-Step default DCW: enabled/double/0.05/0.02/haar; " + "using the official no-pytorch_wavelets no-op path.", RuntimeWarning) + + def _clear(self): + gc.collect() + if self.device.type == "cuda": + torch.cuda.empty_cache() + + @contextmanager + def _component(self, kind): + if kind not in self.components: + # XL decoder weights alone occupy 7.77 GiB. Stream and int8 modes + # build the DiT shape-only and fill it from the files: stream one + # CUDA block per forward, int8 once with int8 decoder linears. + shape_only = kind == "dit" and self.mode in ("stream", "int8") + target = "meta" if shape_only else self.device + started = time.perf_counter() + self.components[kind] = _construct(kind, self.root, self.weights[kind], target, self.dtype) + if shape_only and self.mode == "stream": + self._stream_dit(self.components[kind]) + elif shape_only: + _materialize(self.components[kind], self.weights[kind], self.device, self.dtype, quantize=True) + self._synchronize() + self.stage_times[kind + "_load_s"] = time.perf_counter() - started + try: + yield self.components[kind] + finally: + if self.sequential: + del self.components[kind] + self._clear() + + def _stream_dit(self, dit): + """Materialize exact BF16 weights from local shards, one block at a time. + + Shape-only parameters carry the full validated model structure. Every + nonpersistent constructor buffer is real CPU data, never uninitialized + to_empty storage. At most one block's mapped weights are resident on + the host; all shards are read-only and closed after each block. + """ + weight_map = self.weights["dit"] + + def layout(block, prefix, exclude_layers=False): + entries = [] + for name, parameter in block.named_parameters(): + if exclude_layers and name.startswith("layers."): + continue + parent_name, _, leaf = name.rpartition(".") + parent = block.get_submodule(parent_name) if parent_name else block + entries.append((parent, leaf, parameter, prefix + "." + name)) + return entries + + def materialize(entries): + started = time.perf_counter() + for shard in dict.fromkeys(weight_map[key][0] for _, _, _, key in entries): + with safe_open(str(shard), framework="pt", device="cpu") as source: + for parent, leaf, _, key in entries: + path, name = weight_map[key] + if path != shard: + continue + value = source.get_tensor(name) + tensor = value.to(device=self.device, dtype=self.dtype) + parent._parameters[leaf] = torch.nn.Parameter(tensor, requires_grad=False) + del value, tensor + self.stage_times["dit_weight_load_s"] = ( + self.stage_times.get("dit_weight_load_s", 0.0) + time.perf_counter() - started) + + blocks = [(dit.encoder, "encoder"), (dit.detokenizer, "detokenizer")] + blocks += [(layer, f"decoder.layers.{i}") for i, layer in enumerate(dit.decoder.layers)] + for block, prefix in blocks: + entries = layout(block, prefix) + buffers = [(child, name, value) for child in block.modules() + for name, value in child._buffers.items() if value is not None] + def before(module, args, entries=entries, buffers=buffers): + materialize(entries) + for child, name, value in buffers: + child._buffers[name] = value.to(self.device) + def after(module, args, outputs, entries=entries, buffers=buffers): + for parent, leaf, shape_only, _ in entries: + parent._parameters[leaf] = shape_only + for child, name, value in buffers: + child._buffers[name] = value + block.register_forward_pre_hook(before) + block.register_forward_hook(after, always_call=True) + materialize(layout(dit.tokenizer.quantizer, "tokenizer.quantizer")) + dit.tokenizer.quantizer.to(self.device) + # Decoder projections/timestep embeddings are small and remain resident. + materialize(layout(dit.decoder, "decoder", exclude_layers=True)) + dit.decoder.rotary_emb.to(self.device) + + @contextmanager + def _stage(self, name): + self._synchronize() + start = time.perf_counter() + try: + yield + finally: + self._synchronize() + self.stage_times[name + "_s"] = time.perf_counter() - start + + def _synchronize(self): + if self.device.type == "cuda": + torch.cuda.synchronize(self.device) + + def _chunk_size(self): + # Same thresholds and 512-frame maximum as official tiled_decode. + # Fix the chosen size per instance: offloading must not change the audio. + if hasattr(self, "decode_chunk_size"): + return self.decode_chunk_size + if self.device.type == "cuda": + free, _ = torch.cuda.mem_get_info(self.device) + free += torch.cuda.memory_reserved(self.device) - torch.cuda.memory_allocated(self.device) + # Choose the official chunk for the equivalent fully resident model. + # Keep it across offload modes and allocator caps: changing temporal + # windows changes VAE numerics. Only one VAE window is live at a time. + free -= sum(size for kind, size in self.component_bytes.items() if kind not in self.components) + gib = free / 2**30 + return 512 if gib >= 24 else 384 if gib >= 16 else 256 if gib >= 12 else 128 + return 256 + + def _decode(self, vae, latents): + chunk_size, overlap = self._chunk_size(), 64 + while chunk_size - 2 * overlap <= 0 and overlap > 4: + overlap //= 2 + if latents.shape[-1] <= chunk_size: + return vae.decode(latents).sample + stride = chunk_size - 2 * overlap + pieces = [] + factor = None + for core_start in range(0, latents.shape[-1], stride): + core_end = min(core_start + stride, latents.shape[-1]) + start, end = max(0, core_start - overlap), min(latents.shape[-1], core_end + overlap) + chunk = vae.decode(latents[:, :, start:end]).sample + if factor is None: + factor = chunk.shape[-1] / (end - start) + trim_start = int(round((core_start - start) * factor)) + trim_end = int(round((end - core_end) * factor)) + pieces.append(chunk[:, :, trim_start:chunk.shape[-1] - trim_end].cpu()) + del chunk + audio = torch.cat(pieces, dim=-1) + # Official >=24 GiB decode keeps waveforms on CUDA for the peak limiter. + return audio.to(self.device) if chunk_size == 512 else audio + + @torch.inference_mode() + def generate(self, caption, seconds, seed, progress=None): + validate_request(caption, seconds, seed) + self.stage_times = {} + if progress: + progress(0.0, desc="Planning instrumental music") + from .lm import InstrumentalLM + # Official single-item pt path uses the global torch RNG; seeds=[...] + # is only honored for batches >1 upstream. Seed that RNG explicitly on + # both sides of the reference comparison. Restore caller RNG afterwards. + devices = [self.device.index if self.device.index is not None else torch.cuda.current_device()] if self.device.type == "cuda" else [] + with self._stage("lm"), self._component("lm") as model: + planner = InstrumentalLM(model, self.lm_tokenizer, self.lm_processor, self.device) + with torch.random.fork_rng(devices=devices): + torch.manual_seed(seed) + metadata, codes, times = planner.generate(caption, seconds) + self.last_lm_tokens = planner.tokens + self.last_metadata = metadata.copy() + self.stage_times.update(times) + del planner, model + if self.sequential: + # FSM masks are the only GPU tensors retained by the planner's + # shared processor; return them to CPU before the LM unload. + for name, value in vars(self.lm_processor).items(): + if isinstance(value, torch.Tensor): + setattr(self.lm_processor, name, value.cpu()) + # Match _update_metadata_from_lm + use_cot_caption/language defaults. + # parse_lm_output stores "language", while the updater reads + # "vocal_language". Keep the official effective unknown language. + dit_caption = metadata.get("caption", caption) + language = metadata.get("vocal_language", "unknown") + text = self.tokenizer(_prompt(dit_caption, seconds, metadata), padding="longest", truncation=True, + max_length=256, return_tensors="pt") + lyrics = self.tokenizer(f"# Languages\n{language}\n\n# Lyric\n[Instrumental]<|endoftext|>", + padding="longest", truncation=True, max_length=2048, return_tensors="pt") + with self._stage("text"), self._component("text") as encoder: + hidden = encoder(input_ids=text.input_ids.to(self.device), lyric_attention_mask=None).last_hidden_state + lyric_hidden = encoder.embed_tokens(lyrics.input_ids.to(self.device)) + del encoder + silence = self.silence.to(self.device) + with self._stage("dit"), self._component("dit") as dit: + code_ids = [int(v) for v in re.findall(r"<\|audio_code_(\d+)\|>", codes)] + indices = torch.tensor(code_ids, device=self.device, dtype=torch.long).unsqueeze(0).unsqueeze(-1) + quantized = dit.tokenizer.quantizer.get_output_from_indices(indices) + hints = dit.detokenizer(quantized.to(self.dtype)) + frames = max(128, hints.shape[1]) + if hints.shape[1] < frames: + hints = torch.cat([hints, silence[:, :frames-hints.shape[1]]], dim=1) + src = hints + pred = dit.generate_audio( + text_hidden_states=hidden, text_attention_mask=text.attention_mask.to(self.device).bool(), + lyric_hidden_states=lyric_hidden, lyric_attention_mask=lyrics.attention_mask.to(self.device).bool(), + refer_audio_acoustic_hidden_states_packed=silence[:, :750].contiguous(), + refer_audio_order_mask=torch.tensor([0], device=self.device, dtype=torch.long), + src_latents=src, chunk_masks=torch.ones_like(src), seed=seed, progress=progress, + ) + del dit, quantized, hints, indices + del hidden, lyric_hidden, silence, src + if not torch.isfinite(pred).all() or pred.abs().sum() == 0: + raise RuntimeError("invalid_latents: nonfinite or zero DiT output") + if progress: + progress(0.8, desc="Decoding music") + with self._stage("vae"), self._component("vae") as vae: + audio = self._decode(vae, pred.transpose(1, 2).contiguous()).float() + del vae + peak = audio.abs().amax(dim=[1, 2], keepdim=True) + if torch.any(peak > 1.0): + audio = audio / peak.clamp(min=1.0) + audio = audio[0].cpu() + # generate_music defaults to peak normalization at -1 dB, on CPU, + # after the handler's peak limiter and before AudioSaver writes files. + peak = torch.max(torch.abs(audio)) + if peak >= 1e-6: + audio = audio.clone() * (10 ** (-1.0 / 20.0) / peak) + if progress: + progress(1.0, desc="Music ready") + return audio + + +def load(root, device="cuda", mode="resident"): + return MusicPipeline(root, device, mode) + + +def write_wav(path, audio, rate=SAMPLE_RATE): + """Atomically write stereo [2, samples] audio as little-endian PCM16.""" + path = Path(path) + if isinstance(rate, bool) or not isinstance(rate, int) or rate <= 0: + raise BadRequest("sample rate must be a positive integer") + if isinstance(audio, torch.Tensor): + audio = audio.detach().cpu().numpy() + audio = np.asarray(audio) + if audio.ndim != 2 or audio.shape[0] != 2 or audio.shape[1] == 0 or not np.isfinite(audio).all(): + raise BadRequest("audio must be finite stereo [2, samples]") + pcm = np.rint(np.clip(audio.T, -1, 1) * 32767).astype(" MAX_SECONDS: + raise MusicError('invalid_request', 'too_long') + if not ready(machine['model_root']): + phase('downloading', 0.) + install(machine['model_root'], root, lambda fraction: phase('downloading', fraction), cancelled) + if cancelled(): + raise InterruptedError('cancelled') + phase('composing') + with ACTIVE.lock: + ACTIVE.preempted = False + ACTIVE.idle.clear() + try: + return _run(out, caption, frames, seconds, seed, quiet, fps, server, cancelled, root, machine) + except BlockingIOError: + raise Preempted() from None # other GPU work holds engine.lock; the caller waits and tries again + finally: + with ACTIVE.lock: + ACTIVE.process = None + ACTIVE.idle.set() + + +def _run(out, caption, frames, seconds, seed, quiet, fps, server, cancelled, root, machine): + from ..image_gen import image_pack + from ..locking import runtime_lock + from .. import processes, resident_process + _yield_gpu(server) + with runtime_lock(Path(root) / 'engine.lock', inherit=False, wait_seconds=10) as descriptor: + env = environment(root) + owner = resident_process.OWNER + if owner.process is not None and owner.process.poll() is None and owner.endpoint: + env[resident_process.ENV] = str(owner.endpoint) + try: + resident_process.release_idle_cache(env, descriptor) + owner.record(action='released', reason='music_model') + except resident_process.SessionRestarted: + pass # That worker already exited, which is the release we need. + env.pop(resident_process.ENV, None) + command = processes.module_command( + 'freevideo_engine.music.worker', '--root', image_pack.base(machine['model_root']) / 'music', + '--out', out, '--caption', caption, '--seconds', max(MIN_SECONDS, seconds), '--seed', seed, + '--frames', frames, '--fps', fps, '--quiet', json.dumps([list(span) for span in quiet]), + '--device', machine.get('device_backend', 'cuda')) + command[0] = machine['comfy_python'] + env['FREEVIDEO_RUNTIME_LOCK_FD'] = str(descriptor) + with tempfile.TemporaryFile() as output, tempfile.TemporaryFile() as errors: + with ACTIVE.lock: + # A video that started while the GPU was being freed gets it first. + if ACTIVE.preempted: + raise Preempted() + process = processes.popen(command, env=env, stdin=subprocess.DEVNULL, stdout=output, stderr=errors, + pass_fds=(descriptor,), supervise=True, start_new_session=True) + ACTIVE.process = process + began, timed_out = time.monotonic(), False + try: + while process.poll() is None: + timed_out = time.monotonic() - began > WORKER_SECONDS + if cancelled() or timed_out: + break + time.sleep(POLL_SECONDS) + finally: + # Keep the lease until the whole worker tree has exited. + if process.poll() is None: + processes.stop(process, grace=5) + process.wait() + with ACTIVE.lock: + ACTIVE.process = None + preempted = ACTIVE.preempted + if preempted: + raise Preempted() + if cancelled(): + raise InterruptedError('cancelled') + value = _result(output) + if not value.get('ok'): + errors.seek(0, 2) + errors.seek(max(0, errors.tell() - 4000)) + logging.warning('[FreeVideo] Music worker failed (exit %s, %s): %s', process.returncode, + value.get('error') or 'no result', errors.read().decode('utf-8', 'replace')) + code = value.get('error') + raise MusicError('out_of_memory' if code == 'out_of_memory' else 'music_failed', + code if isinstance(code, str) else 'timeout' if timed_out else 'worker_exit') + return value diff --git a/freevideo_engine/music/worker.py b/freevideo_engine/music/worker.py new file mode 100644 index 00000000..25ed22e6 --- /dev/null +++ b/freevideo_engine/music/worker.py @@ -0,0 +1,118 @@ +"""Run isolated music inference and emit a final machine-readable JSON line.""" +import argparse +import json +from pathlib import Path +import sys +import time + +from .pipeline import MODES, BadRequest, load, validate_request, write_wav + +# Free VRAM at start -> mode. XL DiT + LM + text + VAE resident need about +# 16 GiB; one BF16 DiT stage about STAGED; the int8 DiT stage about INT8. +# Below that, DiT blocks stream from the files (slow, same samples as staged). +RESIDENT_FROM_GIB = 20. +STAGED_FROM_GIB = 10.5 +INT8_FROM_GIB = 6.5 + + +def choose_mode(free_gib): + if free_gib >= RESIDENT_FROM_GIB: + return "resident" + if free_gib >= STAGED_FROM_GIB: + return "staged" + return "int8" if free_gib >= INT8_FROM_GIB else "stream" + + +class _Parser(argparse.ArgumentParser): + def error(self, message): + raise BadRequest(message) + + +def main(argv=None): + phase = "request" + try: + parser = _Parser(description=__doc__) + parser.add_argument("--root", required=True) + parser.add_argument("--out", required=True) + parser.add_argument("--caption", required=True) + parser.add_argument("--seconds", required=True, type=float) + parser.add_argument("--seed", required=True, type=int) + parser.add_argument("--mode", choices=("auto",) + MODES, default="auto") + parser.add_argument("--device", default="cuda") + # With --frames the piece is cut to the timeline, levelled under the video sound and faded (mix.finish). + parser.add_argument("--frames", type=int) + parser.add_argument("--fps", type=float, default=24.) + parser.add_argument("--quiet", default="[]") + args = parser.parse_args(argv) + validate_request(args.caption, args.seconds, args.seed) + quiet = _quiet(args.quiet) + if args.frames is not None and not (0 < args.frames and 0 < args.fps): + raise BadRequest("frames and fps must be positive") + import torch + cuda = args.device == "cuda" and torch.cuda.is_available() + mode = args.mode + if cuda: + torch.cuda.reset_peak_memory_stats() + if mode == "auto": + mode = choose_mode(torch.cuda.mem_get_info()[0] / 2**30) + elif mode == "auto": + mode = "resident" + start = time.perf_counter() + phase = "load" + pipeline = load(args.root, device=args.device, mode=mode) + _synchronize(torch, cuda) + load_s = time.perf_counter() - start + start = time.perf_counter() + phase = "generate" + audio = pipeline.generate(args.caption, args.seconds, args.seed) + _synchronize(torch, cuda) + generate_s = time.perf_counter() - start + if args.frames is not None: + from .mix import finish + audio = finish(audio.numpy(), args.frames, args.fps, quiet) + path = Path(args.out).resolve() + phase = "write" + if path.suffix.lower() == ".mp4": + from .mix import write_mp4 + write_mp4(path, audio) + else: + write_wav(path, audio) + result = dict(ok=True, path=str(path), seconds=audio.shape[1] / 48000, + sample_rate=48000, load_s=load_s, generate_s=generate_s, mode=mode, + sequential=mode != "resident", + lm_s=getattr(pipeline, "stage_times", {}).get("lm_s", 0.0), + peak_vram_gib=torch.cuda.max_memory_allocated() / 2**30 if cuda else None) + code = 0 + except Exception as exc: + import torch + if isinstance(exc, BadRequest): + error = "bad_request" + elif isinstance(exc, FileNotFoundError) and phase == "load": + error = "model_missing" + elif isinstance(exc, torch.cuda.OutOfMemoryError): + error = "out_of_memory" + elif isinstance(exc, OSError) and phase == "write": + error = "write_failed" + else: + error = "inference_failed" + result = dict(ok=False, error=error, detail=str(exc)[:2000]) + code = 1 + print(json.dumps(result, ensure_ascii=False), flush=True) + return code + + +def _quiet(text): + try: + spans = json.loads(text) + return [(int(start), int(frames)) for start, frames in spans] + except (TypeError, ValueError) as error: + raise BadRequest("quiet must be a JSON list of [start, frames]") from error + + +def _synchronize(torch, cuda): + if cuda: + torch.cuda.synchronize() + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/freevideo_engine/network.py b/freevideo_engine/network.py index 0c59af3b..61254492 100644 --- a/freevideo_engine/network.py +++ b/freevideo_engine/network.py @@ -1423,6 +1423,30 @@ def git_version(executable, env=None): return (int(found[1]), int(found[2])) if found else None +# Windows loader codes for a missing DLL, a missing export and a wrong-architecture image. +GIT_LOADER_FAILURES = {0xC0000135, 0xC0000139, 0xC000007B} + + +def git_https_ready(executable, env=None): + """False when this Git's HTTPS helper cannot load its libcurl; True otherwise. + + `git --version` never loads libcurl, so a damaged Git passes it and fails at + the first fetch. A request to a local port nobody listens on needs no network + and reaches the helper: "connection refused" means libcurl loaded. + """ + env = dict(env if env is not None else os.environ, NO_PROXY='*', no_proxy='*') + try: + result = subprocess_output.run_external( + [str(executable), '-c', 'http.proxy=', 'ls-remote', 'https://127.0.0.1:9/probe.git'], env=env, + stdout=subprocess.PIPE, stderr=subprocess.STDOUT, timeout=30, **hidden_console()) + except (OSError, ValueError, subprocess.SubprocessError): + return True # Not a library failure; the version check already ran the executable. + output = (result.stdout or '').lower() + if 'failed to load library' in output or ('libcurl' in output and 'dll' in output): + return False + return (result.returncode & 0xFFFFFFFF) not in GIT_LOADER_FAILURES + + def old_git_message(version): return ('Git %d.%d is too old to download the pinned sources; FreeVideo needs Git %d.%d or newer. ' 'Click Retry installation to use FreeVideo\'s own Git.' % (*version, *MIN_GIT)) diff --git a/freevideo_engine/offline_packages.py b/freevideo_engine/offline_packages.py index 635bc40e..16590f79 100644 --- a/freevideo_engine/offline_packages.py +++ b/freevideo_engine/offline_packages.py @@ -147,7 +147,7 @@ def inventory(rows): def check_runtime_platform(): import sys if sys.platform == 'darwin': - raise ValueError('这是 Windows 运行环境包,不能用于 Mac。请导入模型包后继续安装,安装器会自动下载 Mac 运行环境。') + raise ValueError('The Windows runtime package cannot be used on macOS.') def inspect_archive(path): @@ -158,7 +158,7 @@ def inspect_archive(path): marker = next((name for name in ('offline-package.json', 'runtime.json') if PREFIX + name in names), None) if marker is None: - raise ValueError('不是 FreeVideo 离线包,请选择模型或运行环境 ZIP。') + raise ValueError('The file is not a FreeVideo offline package.') info = archive.getinfo(PREFIX + marker) if info.file_size > 16 * 2**20: raise ValueError('Oversized offline inventory') diff --git a/freevideo_engine/original_audio.py b/freevideo_engine/original_audio.py index b8bc77b5..7e724f5c 100644 --- a/freevideo_engine/original_audio.py +++ b/freevideo_engine/original_audio.py @@ -145,6 +145,19 @@ def _samples(frame): return data if frame.format.is_planar else data.reshape(-1, len(frame.layout.channels)).T +def _fits_flac(data): + """Whether native samples fit FLAC's signed 24-bit grid without rounding.""" + import numpy as np + if data.dtype.kind == 'f': + if not ((data >= -1).all() and (data < 1).all()): + return False + scaled = data.astype(np.float64) * (1 << 23) + return bool((scaled == np.trunc(scaled)).all()) + if data.dtype.itemsize <= 2: + return True + return bool(((data & 255) == 0).all()) + + def inspect(row, *, max_seconds=MAX_SECONDS): """Read a bounded selection on the source's integer decoded-sample clock. @@ -166,10 +179,13 @@ def inspect(row, *, max_seconds=MAX_SECONDS): raise AudioUseError('audio_original_format') if end is not None and end - start > rate * max_seconds: raise _too_long(max_seconds) - count = 0 + count, flac_compatible = 0, True for frame in container.decode(audio=0): if frame.sample_rate != rate or len(frame.layout.channels) != channels: raise AudioUseError('audio_original_format') + begin, stop = max(0, start - count), min(frame.samples, end - count) if end is not None else frame.samples + if begin < stop and flac_compatible: + flac_compatible = _fits_flac(_samples(frame)[:, begin:stop]) count += frame.samples if end is not None and count >= end: break @@ -190,7 +206,7 @@ def inspect(row, *, max_seconds=MAX_SECONDS): raise ValueError('Original audio changed after inspection') return dict(version=VERSION, path=str(path.resolve()), sha256=sha, bytes=path.stat().st_size, sample_rate=rate, channels=channels, start_sample=start, end_sample=end, - codec=codec, full_source=not selection, + codec=codec, full_source=not selection, flac_compatible=flac_compatible, normalization='PyAV-native-to-fltp-32000-v1', decoder_version=av.__version__, library_versions={k: list(v) for k, v in av.library_versions.items()}) @@ -292,6 +308,7 @@ def _audio_stream(container, codec, rate, channels, itemsize): stream.codec_context.format = fmt if codec == 'flac' and itemsize == 4: fmt = stream.codec_context.format = 's32' + stream.codec_context.options = {'bits_per_raw_sample': '24'} return stream, fmt @@ -506,7 +523,7 @@ def saved(path): else: picture = Path(video) copied = False - if info['full_source'] and info['codec'] in ('aac', 'alac', 'mp3', 'flac'): + if info['full_source'] and info['codec'] in ('aac', 'mp3', 'flac'): try: _mux_streams(picture, temporary, rate, copy_from=info['path']) copied = saved(temporary) @@ -517,7 +534,11 @@ def saved(path): if not copied: chunks = _selected_chunks(info) first = next(chunks) - if first.dtype.kind == 'f': + fits_flac = info.get('flac_compatible') + if fits_flac is None: + # Receipts saved before the browser format check still export losslessly. + fits_flac = all(_fits_flac(data) for data in _selected_chunks(info)) + if first.dtype.kind == 'f' and not fits_flac: codec, kind = ('pcm_f64le', first.dtype) if first.dtype.itemsize == 8 else ('pcm_f32le', np.dtype(' change['at']: + return + value.update(schema_version=1, state=change['state'], updated_at=change['at']) + if change['started_at'] is not None: + value['started_at'] = change['started_at'] + if change['shapes'] is not None: + value['covered_shapes'] = change['shapes'] + if change['reason'] is not None: + value['last_postpone'] = dict(reason=change['reason'] if change['reason'] in diagnostic.PRECOMPILE_REASONS + else 'unknown', at=change['at']) + if change['pid'] is not None: + try: + import psutil + value.update(pid=change['pid'], created=psutil.Process(change['pid']).create_time()) + except Exception: + pass # Exited already: without an identity the state reads as unknown. + save(path, value) + + +def covered_shapes(rows): + return [shape for row in diagnostic.sequence(rows)[:16] + if type(diagnostic.mapping(row).get('returncode')) is int and row['returncode'] == 0 + if (shape := diagnostic.precompile_shape(row.get('shape')))] + + +def collect(root=None): + """Read bounded files only; no GPU queries, cache probes or tensor imports.""" + if root is None: + from .system import install_root + root = install_root() + # The ComfyUI parent can select a separate engine through comfyui.json + # without setting FREEVIDEO_HOME in its own process. Match the bridge's + # selection, including paths relative to the source checkout. + if not os.environ.get('FREEVIDEO_HOME'): + source = Path(__file__).resolve().parents[1] + config, _ = read(source / 'comfyui.json', 64 * 1024) + selected = config.get('installation') + if isinstance(selected, str) and selected.strip(): + root = Path(selected).expanduser() + if not root.is_absolute(): + root = source / root + receipt, receipt_status = read(Path(root) / 'precompile.json', 2 * 2**20) + activity, activity_status = read(Path(root) / FILENAME, 64 * 1024) + result = dict(state='not_started' if receipt_status == activity_status == 'missing' else 'unknown', + covered_shapes=covered_shapes(receipt.get('shapes')), receipt=receipt_status, + scope='installation_snapshot') + rows = receipt.get('shapes') + complete = (receipt.get('success') is True and isinstance(rows, list) + and 0 < len(rows) == len(result['covered_shapes'])) + if complete: + result['state'] = 'complete' + if activity.get('schema_version') == 1: + current = diagnostic.precompile(activity) + # A receipt published after the last activity write wins (for example + # if the process exited between saving its receipt and saving status). + finished = diagnostic.numbers(receipt, ('started', 'seconds')) + completed_at = finished.get('started', 0) + finished.get('seconds', 0) + newer_receipt = complete and completed_at > current.get('updated_at', 0) + shapes = result['covered_shapes'] + if not newer_receipt: + result.update(current) + if result['state'] == 'running' and not active(activity): + result['state'] = 'unknown' + for shape in current['covered_shapes']: + if shape not in shapes: + shapes.append(shape) + result['covered_shapes'] = shapes + if current.get('last_postpone'): + result['last_postpone'] = current['last_postpone'] + if result['state'] == 'complete' and not complete: + result['state'] = 'unknown' + return diagnostic.precompile(result) + + +def active(value): + try: + import psutil + if type(value.get('pid')) is not int or value['pid'] <= 0: + return False + process = psutil.Process(value['pid']) + return (process.create_time() == value.get('created') and process.is_running() + and process.status() != psutil.STATUS_ZOMBIE) + except Exception: + return False diff --git a/freevideo_engine/release_notes.json b/freevideo_engine/release_notes.json index c201f9c7..64913c72 100644 --- a/freevideo_engine/release_notes.json +++ b/freevideo_engine/release_notes.json @@ -1,41 +1,78 @@ { "schema": 1, - "product_version": "0.4.26", + "product_version": "0.4.27", "en": { - "summary": "Recent videos open at every window width, when every download source fails the launcher shows the reason for each source, and replacing a person runs on 16 GB Macs.", + "summary": "The Agent can make music for a project, an update started during a generation waits for it to finish, and a preview can keep the guide clip's sound.", "changes": [ - "Creation panel: \"Recent videos\" can be opened at every window width; in a narrow window the list opens over the preview. The @ character menu shows each character's face close-up instead of the full-length picture. When a part length is refused, the length panel says whether it is too short or too long instead of asking to try again. When the connection to FreeVideo drops briefly, the queue no longer reports browser extension interference; a loop that cannot submit its next job stops with the error instead of retrying.", - "Install: when every source for a download fails, the launcher's error card lists each source with its reason and time, and the copied details and the diagnostic report include the same list. Git errors name the folder of the dependency. Account names, passwords, tokens and local paths are removed from the listed addresses. The video text encoder can also be downloaded from ModelScope, so a first install no longer depends on Hugging Face alone for this file; a download that fails the integrity check falls back to Hugging Face.", - "Prompt enhancement: when the connection to FreeVideo is lost, a connection-lost message is shown instead of the browser's raw error, and a failed progress check is retried before the rewrite stops. After a rewrite, characters mentioned with @ are still shown by name in the prompt box. When the prompt or references change during a rewrite, a note asks for a new rewrite instead of an error with a report button.", - "Prompt enhancement: the rewrite tab shows the scene description followed by a separate \"Sound:\" section instead of the model's field names. With several shots, each shot is its own paragraph labelled with its number and start time. A rewrite that is not edited is submitted exactly as the model wrote it. The default prompt in a prompt box that has not been edited uses the same layout and is also submitted unchanged. A prompt filled in with \"Reuse settings\" or opened from a workflow saved by an earlier version also uses this layout and, when not edited, is submitted unchanged.", - "Editor: when the last edit can no longer be found, a blank edit opens with a message saying so; clicking the edit name at the top opens another edit.", - "Settings: a new AI Agent section chooses whether an agent asks before starting a film and before a long reshoot, and how many minutes make a reshoot long. Both ask by default. When an agent's white-model scene fails, the agent receives the reason instead of a generic connection error.", - "Install: when the NVIDIA driver is installed but not running, for example right after it was installed, the launcher's error card asks for a restart instead of showing the raw detection error. When the NVIDIA driver is missing, the card adds one line naming the graphics card Windows lists; when Windows lists only AMD or Intel graphics, the card says FreeVideo supports NVIDIA graphics cards only.", - "Video generation: when a generation falls back to a freshly started worker, the result is the same as in the worker that was already running. After a driver or FreeVideo update, or on another graphics card in the same computer, the first generation no longer fails because an attention method the card cannot run was chosen.", - "Image generation: when Windows runs out of virtual memory while loading the image model, the failure card says so and suggests closing other apps or increasing virtual memory.", - "Linux: when no file manager can open a folder, the launcher and the image history show a message, and the image history also shows the images folder. The default save location follows the desktop's Documents folder. Curly apostrophes in English text are drawn at normal width instead of full width.", - "Mac: encoding a reference video (replace person, extending a video) needs about half the unified memory per tile, so it runs on 16 GB Macs where it stopped with a unified-memory error, including at the lowest resolution. Results are unchanged.", - "Mac: before installing, the launcher checks whether the disk of the chosen folder supports the file operations the Python environment needs. Without hard links the installation copies files instead; when links or renaming are not supported, a card asks for a folder on the Mac's built-in disk before anything is installed." + "Agent: a new music tool makes a soundtrack for a project and places it on an audio track. The music model is downloaded on first use (about 20 GB); cards with little video memory use a built-in 8-bit form. If two attempts fail, the Agent synthesizes the music itself.", + "Updates: when \"Update now\" is pressed while a video is generating, the page shows a waiting screen instead of a failure card, and the update starts when the job ends. The wait also covers a job on a server that another program started.", + "Audio: a kept original audio track is written as lossless audio that browsers can play. A preview made with a guide clip keeps the clip's sound when that was requested.", + "Long videos: when seams cannot be joined, the notes name the seams that failed first, with the step and the format. The how-to for a long video points to \"Extend\" below the finished video.", + "Phone access: the password page is shown in Chinese when the browser's preferred language is Chinese.", + "Speed: the encoded text of a prompt is kept per encoder file and prompt, so a repeated prompt starts sooner.", + "Precompile: the status keeps the order of its writes, stopping responds quickly, and the diagnostic report includes the status.", + "Decoding and reports: a CUDA error while decoding after sampling has finished no longer fails the video; the saved result is decoded once more in a fresh process. A process that keeps running after a CUDA error is ended, so the failure is reported instead of waiting. Diagnostic reports record the polling of Agent generations.", + "Errors: a refused offline package says why, and a slow ComfyUI page is no longer blamed on the browser.", + "Install: on Windows the portable Git is used when the installed Git cannot load libcurl, and the Windows certificate store is used for HTTPS checks, skipping a certificate that cannot be read. A rejected token or address stays explained under its field until it is accepted.", + "Linux: on Arch, the setup names the pacman command for missing basic tools.", + "Mac: after two Metal command buffer failures during sampling, a third attempt runs with Torch attention." ] }, "zh": { - "summary": "任何窗口宽度下都能打开“最近作品”,依赖下载在所有源都失败时,启动器列出每个源的原因;16 GB 的 Mac 上也能替换人物。", + "summary": "Agent 可以为作品制作配乐;生成中点击更新会等任务结束后再更新;预览也可以保留引导片段的声音。", "changes": [ - "创作面板:任何窗口宽度下都能打开“最近作品”;窗口较窄时,列表盖在预览上打开。输入 @ 时,角色头像改为脸部特写,不再是全身图。每段时长被拒绝时,长度面板会说明是过短还是过长,不再提示重试。与 FreeVideo 的连接短暂中断时,队列不再提示浏览器扩展干扰;循环生成提交下一个任务失败时,循环停止并显示错误,不再重试。", - "安装:依赖下载在所有源都失败时,启动器错误卡片逐个列出每个源的失败原因和耗时,复制的详情和诊断报告中也附有这份列表。Git 报错会写明对应依赖的文件夹。列出的地址会去掉账号、密码、令牌和本地路径。视频文本编码器增加魔搭下载源,首次安装下载这个文件时不再只依赖 Hugging Face;下载的文件校验不通过时改从 Hugging Face 下载。", - "提示词增强:与 FreeVideo 的连接中断时,显示连接断开的提示,而不是浏览器的原始错误;改写进度读取失败时先重试,不再立即中止。改写后,人物仍以“@人物名”显示在提示词框中。改写过程中原文或素材有改动时,显示一条提醒,请按新内容重新改写,不再显示错误和报告按钮。", - "提示词增强:“改写”页签只显示画面描述,下面单独列出“声音:”段,不再显示模型的字段名;有多个镜头时,每个镜头单独成段,标出镜头编号和开始时间。未修改的改写结果按模型原文提交。提示词框里未改动的默认提示词也按同样方式显示,提交时仍用原文。用“复用设置”填入的提示词和旧版本保存的工作流里的提示词同样按这种方式显示,未修改时按原文提交。", - "剪辑视图:上次的剪辑已找不到时,打开空白剪辑并给出说明;点击顶部的剪辑名可打开其他剪辑。", - "设置:新增“AI Agent”区块,可选择 Agent 在开拍整部片前、以及耗时较长的重拍前是否先询问,并设置重拍达到多少分钟算较长。两项默认都询问。Agent 搭建的白模场景失败时,Agent 收到具体原因,不再是笼统的连接错误。", - "安装:NVIDIA 显卡驱动已安装但未运行时(例如刚装好还未重启),启动器错误卡片提示重启电脑,不再显示原始的检测失败信息。未找到 NVIDIA 显卡驱动时,提示卡新增一行,写出系统列出的显卡;系统只列出 AMD 或 Intel 显卡时,卡片说明 FreeVideo 只支持 NVIDIA 显卡。", - "视频生成:生成改由新启动的进程完成时,结果与原先运行中的进程相同。驱动或 FreeVideo 更新后,或换用同一台电脑里的另一张显卡时,首次生成不再因选用这张显卡无法运行的注意力方式而失败。", - "图像生成:Windows 虚拟内存不足、无法载入图像模型时,失败卡片会说明原因,并建议关闭其他程序或调大虚拟内存。", - "Linux:没有可用的文件管理器打开文件夹时,启动器和图片历史给出提示,图片历史同时显示图片文件夹的路径;默认保存位置跟随桌面设置的文档文件夹。英文中的弯撇号按西文宽度显示,不再显示为全角。", - "Mac:编码参考视频(替换人物、接着生成)时,每个瓦片所需的统一内存约减半,在 16 GB 的 Mac 上不再因统一内存不足而中止,最低分辨率下也是如此;结果不变。", - "Mac:安装开始前,启动器先检查所选文件夹所在磁盘是否支持 Python 运行环境所需的文件操作。不支持硬链接时改用复制继续安装;不支持链接或改名时,在安装任何内容之前提示选择 Mac 内置硬盘上的文件夹。" + "Agent:新增配乐工具,为作品制作配乐并放到一条音频轨上。配乐模型首次使用时下载(约 20 GB);显存较小的显卡使用内置的 8 位形式。两次尝试都失败时,由 Agent 自己合成。", + "更新:有视频正在生成时点击“立即更新”,页面显示等待画面,不再出现失败卡片,任务结束后开始更新。由其他程序启动的服务上的任务也会等待。", + "音频:保留的原始音轨以浏览器可播放的无损音频写入。使用引导片段的预览在要求保留时,也保留该片段的声音。", + "长视频:接缝无法衔接时,说明先列出失败的接缝及其步骤和格式。长视频的使用指引改为指向成片下方的“接着生成”。", + "手机访问:浏览器的首选语言为中文时,密码页显示中文。", + "速度:提示词的文本编码结果按编码器文件和提示词保存,重复使用同一提示词时更快开始。", + "预编译:状态按写入顺序保存,停止时响应更快,诊断报告包含该状态。", + "解码与诊断报告:采样完成后解码时出现 CUDA 错误,不再整条失败,已保存的结果会在新进程中重新解码一次;出现 CUDA 错误后仍不退出的进程会被结束,直接报告失败,不再一直等待。诊断报告记录 Agent 生成任务的轮询情况。", + "错误提示:离线包被拒绝时说明原因;ComfyUI 页面加载较慢时不再归咎于浏览器。", + "安装:Windows 上已安装的 Git 无法载入 libcurl 时改用便携版 Git;HTTPS 检查使用 Windows 证书库,遇到无法读取的证书时跳过。令牌或地址被拒绝后,说明一直显示在对应输入框下,直到输入被接受。", + "Linux:在 Arch 上缺少基础工具时,安装脚本给出对应的 pacman 命令。", + "Mac:采样时连续两次 Metal 命令缓冲区失败后,用 Torch 注意力再试第三次。" ] }, "history": [ + { + "product_version": "0.4.26", + "en": { + "summary": "Recent videos open at every window width, when every download source fails the launcher shows the reason for each source, and replacing a person runs on 16 GB Macs.", + "changes": [ + "Creation panel: \"Recent videos\" can be opened at every window width; in a narrow window the list opens over the preview. The @ character menu shows each character's face close-up instead of the full-length picture. When a part length is refused, the length panel says whether it is too short or too long instead of asking to try again. When the connection to FreeVideo drops briefly, the queue no longer reports browser extension interference; a loop that cannot submit its next job stops with the error instead of retrying.", + "Install: when every source for a download fails, the launcher's error card lists each source with its reason and time, and the copied details and the diagnostic report include the same list. Git errors name the folder of the dependency. Account names, passwords, tokens and local paths are removed from the listed addresses. The video text encoder can also be downloaded from ModelScope, so a first install no longer depends on Hugging Face alone for this file; a download that fails the integrity check falls back to Hugging Face.", + "Prompt enhancement: when the connection to FreeVideo is lost, a connection-lost message is shown instead of the browser's raw error, and a failed progress check is retried before the rewrite stops. After a rewrite, characters mentioned with @ are still shown by name in the prompt box. When the prompt or references change during a rewrite, a note asks for a new rewrite instead of an error with a report button.", + "Prompt enhancement: the rewrite tab shows the scene description followed by a separate \"Sound:\" section instead of the model's field names. With several shots, each shot is its own paragraph labelled with its number and start time. A rewrite that is not edited is submitted exactly as the model wrote it. The default prompt in a prompt box that has not been edited uses the same layout and is also submitted unchanged. A prompt filled in with \"Reuse settings\" or opened from a workflow saved by an earlier version also uses this layout and, when not edited, is submitted unchanged.", + "Editor: when the last edit can no longer be found, a blank edit opens with a message saying so; clicking the edit name at the top opens another edit.", + "Settings: a new AI Agent section chooses whether an agent asks before starting a film and before a long reshoot, and how many minutes make a reshoot long. Both ask by default. When an agent's white-model scene fails, the agent receives the reason instead of a generic connection error.", + "Install: when the NVIDIA driver is installed but not running, for example right after it was installed, the launcher's error card asks for a restart instead of showing the raw detection error. When the NVIDIA driver is missing, the card adds one line naming the graphics card Windows lists; when Windows lists only AMD or Intel graphics, the card says FreeVideo supports NVIDIA graphics cards only.", + "Video generation: when a generation falls back to a freshly started worker, the result is the same as in the worker that was already running. After a driver or FreeVideo update, or on another graphics card in the same computer, the first generation no longer fails because an attention method the card cannot run was chosen.", + "Image generation: when Windows runs out of virtual memory while loading the image model, the failure card says so and suggests closing other apps or increasing virtual memory.", + "Linux: when no file manager can open a folder, the launcher and the image history show a message, and the image history also shows the images folder. The default save location follows the desktop's Documents folder. Curly apostrophes in English text are drawn at normal width instead of full width.", + "Mac: encoding a reference video (replace person, extending a video) needs about half the unified memory per tile, so it runs on 16 GB Macs where it stopped with a unified-memory error, including at the lowest resolution. Results are unchanged.", + "Mac: before installing, the launcher checks whether the disk of the chosen folder supports the file operations the Python environment needs. Without hard links the installation copies files instead; when links or renaming are not supported, a card asks for a folder on the Mac's built-in disk before anything is installed." + ] + }, + "zh": { + "summary": "任何窗口宽度下都能打开“最近作品”,依赖下载在所有源都失败时,启动器列出每个源的原因;16 GB 的 Mac 上也能替换人物。", + "changes": [ + "创作面板:任何窗口宽度下都能打开“最近作品”;窗口较窄时,列表盖在预览上打开。输入 @ 时,角色头像改为脸部特写,不再是全身图。每段时长被拒绝时,长度面板会说明是过短还是过长,不再提示重试。与 FreeVideo 的连接短暂中断时,队列不再提示浏览器扩展干扰;循环生成提交下一个任务失败时,循环停止并显示错误,不再重试。", + "安装:依赖下载在所有源都失败时,启动器错误卡片逐个列出每个源的失败原因和耗时,复制的详情和诊断报告中也附有这份列表。Git 报错会写明对应依赖的文件夹。列出的地址会去掉账号、密码、令牌和本地路径。视频文本编码器增加魔搭下载源,首次安装下载这个文件时不再只依赖 Hugging Face;下载的文件校验不通过时改从 Hugging Face 下载。", + "提示词增强:与 FreeVideo 的连接中断时,显示连接断开的提示,而不是浏览器的原始错误;改写进度读取失败时先重试,不再立即中止。改写后,人物仍以“@人物名”显示在提示词框中。改写过程中原文或素材有改动时,显示一条提醒,请按新内容重新改写,不再显示错误和报告按钮。", + "提示词增强:“改写”页签只显示画面描述,下面单独列出“声音:”段,不再显示模型的字段名;有多个镜头时,每个镜头单独成段,标出镜头编号和开始时间。未修改的改写结果按模型原文提交。提示词框里未改动的默认提示词也按同样方式显示,提交时仍用原文。用“复用设置”填入的提示词和旧版本保存的工作流里的提示词同样按这种方式显示,未修改时按原文提交。", + "剪辑视图:上次的剪辑已找不到时,打开空白剪辑并给出说明;点击顶部的剪辑名可打开其他剪辑。", + "设置:新增“AI Agent”区块,可选择 Agent 在开拍整部片前、以及耗时较长的重拍前是否先询问,并设置重拍达到多少分钟算较长。两项默认都询问。Agent 搭建的白模场景失败时,Agent 收到具体原因,不再是笼统的连接错误。", + "安装:NVIDIA 显卡驱动已安装但未运行时(例如刚装好还未重启),启动器错误卡片提示重启电脑,不再显示原始的检测失败信息。未找到 NVIDIA 显卡驱动时,提示卡新增一行,写出系统列出的显卡;系统只列出 AMD 或 Intel 显卡时,卡片说明 FreeVideo 只支持 NVIDIA 显卡。", + "视频生成:生成改由新启动的进程完成时,结果与原先运行中的进程相同。驱动或 FreeVideo 更新后,或换用同一台电脑里的另一张显卡时,首次生成不再因选用这张显卡无法运行的注意力方式而失败。", + "图像生成:Windows 虚拟内存不足、无法载入图像模型时,失败卡片会说明原因,并建议关闭其他程序或调大虚拟内存。", + "Linux:没有可用的文件管理器打开文件夹时,启动器和图片历史给出提示,图片历史同时显示图片文件夹的路径;默认保存位置跟随桌面设置的文档文件夹。英文中的弯撇号按西文宽度显示,不再显示为全角。", + "Mac:编码参考视频(替换人物、接着生成)时,每个瓦片所需的统一内存约减半,在 16 GB 的 Mac 上不再因统一内存不足而中止,最低分辨率下也是如此;结果不变。", + "Mac:安装开始前,启动器先检查所选文件夹所在磁盘是否支持 Python 运行环境所需的文件操作。不支持硬链接时改用复制继续安装;不支持链接或改名时,在安装任何内容之前提示选择 Mac 内置硬盘上的文件夹。" + ] + } + }, { "product_version": "0.4.25", "en": { @@ -193,35 +230,6 @@ "不受支持的显卡的提示卡现在写出检测到的显卡和计算能力,并说明 FreeVideo 需要计算能力 8.0 及以上的显卡(GeForce RTX 30 系列及更新),与显存大小无关。错误文字中句子里的斜杠不再被替换成路径占位符。" ] } - }, - { - "product_version": "0.4.19", - "en": { - "summary": "Several pictures can be deleted from the image history at once; prompt enhancement states the reason when it cannot finish.", - "changes": [ - "Install: setup no longer stops when some graphics card details cannot be read; command-line generation works on Windows when FreeVideo is installed in a very deep folder.", - "Image history: after Select, several pictures can be checked and deleted at once, with a second click to confirm; pictures that are in use or cannot be deleted stay in place with the reason shown.", - "Prompt enhancement: when enhancement cannot finish and the original prompt is used, the reason is shown; a rewrite that lacks the video or audio description is sent back to the model to add it.", - "Editor: Extend on a finished video without a description asks for a description first instead of showing an English validation error; when the description box is cleared, the description carried over for that segment is used.", - "The Linux package now includes curl, so installation and downloads work on systems without it.", - "Download list: the check no longer hashes every model file again each time it runs; files already hashed are skipped while their size, times and inode are unchanged.", - "Keeps the failed phase, step count, error type and memory readings of a GPU error that raised the compatibility level, so a diagnostic report sent after a later run still shows where it stopped.", - "The chat no longer shows the notice for an edited waiting message as an error with a report button; it is a plain status line." - ] - }, - "zh": { - "summary": "生图历史可以一次删除多张;提示词增强未能完成时写明原因。", - "changes": [ - "安装:部分显卡信息读不到时,安装不再中断;Windows 上 FreeVideo 安装在很深的文件夹时,命令行生成可以正常运行。", - "生图历史:点“选择”后可以勾选多张一次删除,删除前需要再次确认;被占用或无法删除的图片会保留并写明原因。", - "提示词增强:增强未能完成、改用原提示词时,写明原因;改写结果缺少视频或音频描述时,会请模型补上。", - "剪辑页:在没有描述的成片后“接着生成”时,会提示先写一句描述,不再显示英文校验错误;描述框被清空时,沿用这一段原有的描述。", - "Linux 安装包自带 curl,系统未安装 curl 时也能正常安装和下载。", - "下载清单:检查不再每次重新校验所有模型文件;已经校验过的文件在大小、时间和 inode 都没变时直接跳过。", - "保留导致兼容档位升高的显卡错误的出错阶段、步数、错误类型和内存读数,之后一次运行发送的诊断报告仍能看出它停在哪里。", - "对话中修改了正在等待 Agent 的消息后,提示不再按错误显示、不再带“报告问题”按钮,改为普通状态提示。" - ] - } } ] } diff --git a/freevideo_engine/resource_history.py b/freevideo_engine/resource_history.py index 339bfbe0..abe0b797 100644 --- a/freevideo_engine/resource_history.py +++ b/freevideo_engine/resource_history.py @@ -409,7 +409,14 @@ def finish(self, attempt_id, outcome, *, details=None, observation=None): if observation is not None and outcome != 'success': raise ValueError('Failed attempts cannot certify resource calibration') if outcome == 'discarded': - if (details or {}).get('kind') in ('cuda_error', 'unknown_worker_exit'): + # A first CUDA error in the decode phase comes after sampling finished: the + # sampling partitions a compatibility level changes did not fail, and the + # retry decodes the saved latents. A decode retry that fails again is + # recorded, since only a higher level changes the decoder. + repeated = bool(json.loads(row['details_json']).get('decision_evidence', {}).get('decode_resume')) + after_sampling = ((details or {}).get('kind') == 'cuda_error' + and (details or {}).get('failed_phase') == 'decode' and not repeated) + if (details or {}).get('kind') in ('cuda_error', 'unknown_worker_exit') and not after_sampling: from .compatibility import Store Store(self.path.parent).incident(self._row(row), details['kind'], evidence=details) db.execute('DELETE FROM attempts WHERE id=?',(attempt_id,)) diff --git a/freevideo_engine/support_report.py b/freevideo_engine/support_report.py index 7a082c9b..5211703a 100644 --- a/freevideo_engine/support_report.py +++ b/freevideo_engine/support_report.py @@ -14,6 +14,26 @@ from .diagnostic_resources import kernel_receipt from .monitoring import save +AGENT_VIDEO_STALL_SECONDS = 180 + + +def agent_video_stalls(rows, *, now=None): + now = time.time() if now is None else now + result = [] + for row in rows if isinstance(rows, list) else (): + if not isinstance(row, dict) or row.get('status') != 'running': + continue + registered = row.get('video_registered_at') + diagnostics = mapping(row.get('generation_poll_diagnostics')) + last = mapping(diagnostics.get('last')) + if (type(registered) not in (int, float) or now - registered < AGENT_VIDEO_STALL_SECONDS + or not isinstance(row.get('id'), str) or not last): + continue + result.append(dict(task_id=row['id'], status='running', video_registered_at=registered, + last_poll={key: last.get(key) for key in ('time', 'queue_state', 'history_entry', + 'history_success', 'freevideo_summary')}, poll_count=diagnostics.get('count'))) + return result + def number(value): return value if type(value) in (int, float) and math.isfinite(value) and value >= 0 else None @@ -380,6 +400,7 @@ def count(value): exception=latest.get('exception', []), gpu=latest or mapping(encoding.get('gpu')), memory=mapping(last_stage.get('memory')))) bridge = mapping(bridge) or read_json('comfy-request', output.parent / 'comfy-request.json') + agent_poll = read_json('agent-generation-poll', output.parent / 'agent-generation-poll.json') redactor = Redactor([(str(output.parent), '')]) try: redactor.replacements.append((str(Path.home()), '')) @@ -458,6 +479,14 @@ def count(value): if not request.get('encoding_resources'): request['encoding_resources'] = dict(ram=read_json('.encoding.memory.json')) summary = summarize(request, engine, encoding) + try: + from .precompile_state import collect as precompile_status + precompile = precompile_status() + except Exception: + # Status collection is optional even when the installation is broken. + from .diagnostic_resources import precompile as sanitize_precompile + precompile = sanitize_precompile(None) + summary['precompile'] = precompile if live and not restarted and bridge.get('status') in ('starting', 'running'): summary['status'] = 'running' if bridge.get('status') == 'reused': @@ -474,6 +503,10 @@ def count(value): 'Cached encoder receipts do not describe this request’s peak.'), bridge={k:bridge[k] for k in ('status', 'error', 'bridge_seconds', 'bridge_wall_seconds', 'sampling_cache_install', 'encoder_prewarm') if k in bridge}, collection_notes=notes, log_tails={}) + stalled = agent_video_stalls(bridge.get('agent_generation_jobs') or + agent_poll.get('agent_generation_jobs')) + if stalled: + payload['agent_generation_stalls'] = stalled try: if preview: previous, previous_engine, previous_encoding = preview @@ -575,6 +608,9 @@ def clean(value, key=''): suffix = ('.encoder-retry-' if encoder_retry else '.retry-') + str(_retry['index']) + '.debug.json' target = output.with_suffix(suffix) cleaned = clean(payload) + # Like token_summary, text_tokens is numeric shape metadata; restore + # only this allowlisted section after the generic secret redactor. + cleaned['summary']['precompile'] = precompile if timeline is not None: # Counter slashes in fixed label keys look like paths to the text # scrubber. Restore only the receipt already validated above. diff --git a/freevideo_engine/system.py b/freevideo_engine/system.py index 4173395b..f1ce651f 100644 --- a/freevideo_engine/system.py +++ b/freevideo_engine/system.py @@ -71,17 +71,17 @@ def linux_tools_command(names, which=shutil.which, root=None): The desktop launcher cannot ask for a password in a terminal, so its plan names the command; ./setup.sh installs the same packages itself. """ - manager = next((m for m in ('apt-get', 'dnf', 'yum', 'zypper') if which(m)), None) + manager = next((m for m in ('apt-get', 'dnf', 'yum', 'zypper', 'pacman') if which(m)), None) if manager is None: return None packages = [] for name in names: - package = {'g++': 'build-essential' if manager == 'apt-get' else 'gcc-c++', + package = {'g++': 'build-essential' if manager == 'apt-get' else 'gcc' if manager == 'pacman' else 'gcc-c++', 'make': 'build-essential' if manager == 'apt-get' else 'make'}.get(name, name) if package not in packages: packages.append(package) install = {'apt-get': 'apt-get install -y', 'dnf': 'dnf install -y', 'yum': 'yum install -y', - 'zypper': 'zypper --non-interactive install'}[manager] + 'zypper': 'zypper --non-interactive install', 'pacman': 'pacman -S --needed'}[manager] privilege = '' if (root if root is not None else os.geteuid() == 0) else 'sudo ' update = privilege + 'apt-get update && ' if manager == 'apt-get' else '' return update + privilege + install + ' ' + ' '.join(packages) diff --git a/freevideo_engine/tuning.py b/freevideo_engine/tuning.py index 8c6f9a60..6489c6bb 100644 --- a/freevideo_engine/tuning.py +++ b/freevideo_engine/tuning.py @@ -116,6 +116,54 @@ def identity(machine, hardware): 'encoder': machine['encoder'], 'model_revision': machine['model_revision']} +TEXT_SCHEMA = 1 +TEXT_FILES = ('encode_worker.py', 'encoder_lowmem.py', 'encoder_precision.py', 'encoder_checkpoint.py', + 'encoder_pinning.py', 'encoder_workspace.py', 'encoder_memory.py', 'adaptive.py', + 'conditioning.py', 'media_request.py') +TEXT_PACKAGES = ('torch', 'transformers', 'tokenizers', 'comfy-kitchen', 'safetensors', 'numpy') + + +def encoder_content(path, catalog): + """Cheap content identity of the encoder file: catalog hash and size, plus its header and sampled blocks. + + A same-size edit that touches none of the sampled bytes is not seen; a full 15 GB hash on every run would + cost more than the encode it avoids.""" + size = path.stat().st_size + expected = next((row for row in catalog if row['file'].endswith('/' + path.name)), None) + value = hashlib.sha256() + with path.open('rb') as stream: + head = stream.read(8) + value.update(head) + value.update(stream.read(int.from_bytes(head, 'little') if int.from_bytes(head, 'little') < 2**26 else 0)) + for index in range(16): + stream.seek(size * index // 16) + value.update(stream.read(65536)) + return {'size': size, 'catalog_sha256': expected and expected['sha256'], + 'catalog_bytes': expected and expected['bytes'], 'sampled_sha256': value.hexdigest()} + + +def text_identity(machine, hardware): + """What the text conditioning of a prompt depends on: encoder content, its code and arithmetic, nothing else. + + No git, no timestamps, no video models, no GPU serial number: an unrelated update or a copied folder without + .git keeps the cache; a different encoder file, library, source or GPU architecture does not.""" + package = Path(__file__).parent + catalog = json.loads((package / 'model_files.json').read_text(encoding='utf-8')) + catalog = catalog['files'] if isinstance(catalog, dict) else catalog + encoder = Path(machine['encoder_model_root']) / 'text_encoders' / machine['encoder'] + specs = json.loads((package / 'dependencies.json').read_text(encoding='utf-8')) + texts = Path(machine['comfy_root']) / 'comfy' / 'text_encoders' + assets = {str(p.relative_to(texts)): digest(p) for p in sorted(texts.rglob('*')) + if p.is_file() and p.suffix in ('.py', '.json', '.txt', '.model')} + versions = package_versions(machine['comfy_python']) + return {'schema': TEXT_SCHEMA, 'kind': 'text-conditioning', 'system': hardware.system, + 'gpu': [hardware.gpu_name, list(hardware.capability), hardware.driver_version], + 'encoder': encoder_content(encoder, catalog), + 'encoder_packages': {k: v for k, v in versions.items() if k in TEXT_PACKAGES}, + 'comfy_commit': specs['encoder']['comfy_commit'], 'comfy_text_encoders': assets, + 'source_sha256': {name: digest(package / name) for name in TEXT_FILES}} + + def report_identity_errors(inventory, expected): errors = [] if expected.get('system') and inventory.get('hardware', {}).get('system') != expected['system']: diff --git a/freevideo_engine/windows_ssl.py b/freevideo_engine/windows_ssl.py new file mode 100644 index 00000000..32a99bd2 --- /dev/null +++ b/freevideo_engine/windows_ssl.py @@ -0,0 +1,52 @@ +"""Run by ComfyUI before it imports its server (and aiohttp), through the prestartup loader the launcher +writes: keep one unreadable Windows root certificate from stopping the start.""" +import ssl +import sys +import warnings + + +def tolerate_unreadable_store_certificates(): + """Windows hands Python every trusted root as one block, and OpenSSL rejects the whole block when a + single certificate in it does not parse. Without this, creating any default TLS context raises + "[ASN1] nested asn1 error" and ComfyUI quits on its first import of aiohttp. + + Only after that failure are the certificates loaded one by one, skipping the ones OpenSSL cannot + read. Roots installed by an administrator stay trusted, and a store without a bad certificate takes + exactly the same path as before.""" + original = getattr(ssl.SSLContext, '_load_windows_store_certs', None) + if sys.platform != 'win32' or original is None or getattr(original, 'freevideo_tolerant', False): + return False + + def load(self, storename, purpose): + try: + return original(self, storename, purpose) + except ssl.SSLError: + pass + # Read again on every context, so a root an administrator removes later is not trusted from memory. + certificates = [] + try: + probe = ssl.SSLContext(ssl.PROTOCOL_TLS_CLIENT) + for certificate, encoding, trust in ssl.enum_certificates(storename): + if encoding != 'x509_asn' or not (trust is True or purpose.oid in trust): + continue + try: + probe.load_verify_locations(cadata=certificate) + except ssl.SSLError: + continue + certificates.append(certificate) + except PermissionError: + warnings.warn('unable to enumerate Windows certificate store') + loaded = bytearray(b''.join(certificates)) + if loaded: + self.load_verify_locations(cadata=loaded) + return loaded + + load.freevideo_tolerant = True + ssl.SSLContext._load_windows_store_certs = load + return True + + +try: + tolerate_unreadable_store_certificates() +except Exception as error: + print('[FreeVideo] Windows certificate store check skipped: %s' % error) diff --git a/scripts/bootstrap_linux.sh b/scripts/bootstrap_linux.sh index f6ea584e..ca98981d 100644 --- a/scripts/bootstrap_linux.sh +++ b/scripts/bootstrap_linux.sh @@ -106,6 +106,7 @@ fv_system_deps() { elif command -v dnf >/dev/null 2>&1; then manager=dnf elif command -v yum >/dev/null 2>&1; then manager=yum elif command -v zypper >/dev/null 2>&1; then manager=zypper + elif command -v pacman >/dev/null 2>&1; then manager=pacman else fv_fail "缺少基础工具:${listed}。无法自动安装,请联系系统管理员安装以上工具,完成后重新开始安装。" fv_fail "Missing basic tools: ${listed}. They cannot be installed automatically. Ask your administrator to install them, then start the installation again." @@ -114,7 +115,7 @@ fv_system_deps() { for name in "${fv_missing[@]}"; do package="$name" case "$name" in - g++|make) if [[ "$manager" == apt-get ]]; then package=build-essential; else package=gcc-c++; fi ;; + g++|make) if [[ "$manager" == apt-get ]]; then package=build-essential; elif [[ "$manager" == pacman ]]; then package=gcc; else package=gcc-c++; fi ;; sha256sum) package=coreutils ;; flock) package=util-linux ;; esac @@ -128,6 +129,7 @@ fv_system_deps() { install=(apt-get -o Acquire::Retries=2 -o Acquire::http::Timeout=30 -o Acquire::https::Timeout=30 install -y --no-install-recommends "${packages[@]}") ;; dnf|yum) install=("$manager" --setopt=timeout=30 --setopt=retries=2 install -y "${packages[@]}") ;; zypper) install=(zypper --non-interactive install --no-recommends "${packages[@]}") ;; + pacman) install=(pacman -S --needed --noconfirm "${packages[@]}") ;; esac if (( EUID != 0 )); then privilege=(sudo); fi if (( fv_readonly )); then @@ -142,6 +144,7 @@ fv_system_deps() { # because a minimal system often has an empty one. local manual manual="${privilege[*]:+sudo }$manager install ${packages[*]}" + [[ "$manager" != pacman ]] || manual="${privilege[*]:+sudo }pacman -S --needed ${packages[*]}" [[ "$manager" != apt-get ]] || manual="${privilege[*]:+sudo }apt-get update && ${privilege[*]:+sudo }apt-get install ${packages[*]}" if (( EUID != 0 )) && ! command -v sudo >/dev/null 2>&1; then fv_fail "缺少基础工具:${listed}。此系统未安装 sudo,请联系系统管理员运行以下命令,完成后重新开始安装:" diff --git a/web/autopilot_chat.js b/web/autopilot_chat.js index e6a5fbd1..373d6d85 100644 --- a/web/autopilot_chat.js +++ b/web/autopilot_chat.js @@ -43,6 +43,8 @@ export const STAGE_TEXT = Object.freeze({ const WORKING = ['running', 'queued', 'waiting', 'cancelling'], VIDEOS = ['generate', 'assemble']; // What a job with no picture of its own is doing, in the reader's language. const QUIET = {vlm: ['Checking', '检查中'], enhance: ['Enhancing prompt', '增强提示词']}; +// Other jobs without a card keep their line while they run, and a finished one still marks the tab. +const BUSY = {music: ['Making music', '制作配乐']}; const WHITEBOX = { generate: ['Generate video', '生成视频'], edit: ['Edit prompt', '改描述'], decline: ['Skip', '跳过'], approved: ['Approved', '已确认'], superseded: ['Replaced', '已被替换'], declined: ['Skipped', '已跳过'], @@ -1018,7 +1020,7 @@ export function agentChat(t, api, {open, edit, shown = () => true, changed = () else if (next === 'quiet') { const label = el('span'); item.replaceChildren(el('span', 'fv-chat-dots'), label); apply = next => { label.textContent = stopping.has(next.id) || next.status === 'cancelling' - ? t('Stopping', '正在停止') : t(...(QUIET[next.kind] || ['Working', '进行中'])); }; + ? t('Stopping', '正在停止') : t(...(QUIET[next.kind] || BUSY[next.kind] || ['Working', '进行中'])); }; } else if (next === 'stopped' && (!job.error || job.kind === 'image')) { item.replaceChildren(el('span', null, t('Stopped', '已停止'))); apply = null; } else if (next.startsWith('failed') && job.kind === 'image') { diff --git a/web/health.js b/web/health.js index 9fc040a4..0e2347db 100644 --- a/web/health.js +++ b/web/health.js @@ -660,7 +660,7 @@ async function check() { send('loading', []); const app = await waitFor(comfy, 60000); if (!app) { - const problems = [{text: t('The ComfyUI page did not finish loading within a minute.', 'ComfyUI 页面在一分钟内没有加载完成。'), advice: 'browser'}]; + const problems = [{text: t('The ComfyUI page did not finish loading within a minute.', 'ComfyUI 页面在一分钟内没有加载完成。'), advice: 'reload'}]; show(problems); send('error', problems); return; } // A restarted ComfyUI has not heard from this page; tell it again. diff --git a/web/howto.js b/web/howto.js index 72e1679c..cafdb7ee 100644 --- a/web/howto.js +++ b/web/howto.js @@ -35,8 +35,8 @@ const HOWTOS = [ zh: '填写更长的时长,会自动分段生成。'}, {id: 'generate', anchor: '[data-fv-guide="generate"]', en: 'Choose "Generate video".', zh: '点“生成视频”。'}, {id: 'extend', anchor: '.fv-action-extend', image: 'extend-button', - en: 'Choose "Extend" on a finished video to continue it. Each continued part is at least 5 seconds; with a part shorter than 10 seconds, one video has at most 4 parts.', - zh: '在生成好的视频上点“接着生成”,继续下一段。接着生成的每一段至少 5 秒;有短于 10 秒的段时,一条视频最多 4 段。'}, + en: 'Choose "Extend" below a finished video to continue it. Each continued part is at least 5 seconds; with a part shorter than 10 seconds, one video has at most 4 parts.', + zh: '在生成好的视频下方点“接着生成”,继续下一段。接着生成的每一段至少 5 秒;有短于 10 秒的段时,一条视频最多 4 段。'}, {id: 'longer', anchor: '[data-view="editor"]', en: 'Parts can also be added one by one in "Edit".', zh: '也可在“剪辑视图”中逐段接续。'}]}, {id: 'editor-extend-refs', view: 'editor', title: ['Add reference images when extending in Edit', '在剪辑视图延长时加参考图'], steps: [ diff --git a/web/phone.js b/web/phone.js index 864a2121..0943f353 100644 --- a/web/phone.js +++ b/web/phone.js @@ -107,7 +107,7 @@ pw.onsubmit = async event => { const wait = Number(response.headers.get('X-FreeVideo-Retry-After')); if (response.ok) { locked = false; password.value = ''; pwError.hidden = true; field.classList.remove('bad'); renderShell(); await poller.poll(); } else if (response.status === 429 && wait > 0) lockFor(wait); - else if (response.status === 401) { pwError.textContent = t('Wrong password. Please try again.', '密码不对,请再输一次。'); pwError.hidden = false; field.classList.add('bad'); password.focus(); password.select(); } + else if (response.status === 401) { pwError.textContent = t('Incorrect password.', '密码不正确。'); pwError.hidden = false; field.classList.add('bad'); password.focus(); password.select(); } else { pwError.textContent = t('Connection lost. Reconnecting…', '连接断了,正在重新连接…'); pwError.hidden = false; } } catch { pwError.textContent = t('Connection lost. Reconnecting…', '连接断了,正在重新连接…'); pwError.hidden = false; } finally { pw.dataset.busy = 'false'; tickPassword(); } diff --git a/web/phone_model.js b/web/phone_model.js index 450acee6..7d4f11a4 100644 --- a/web/phone_model.js +++ b/web/phone_model.js @@ -25,8 +25,8 @@ export function mergePhoneJobs(previous, chat, activity, snapshots = []) { } export function retryWords(seconds, t) { const count = Math.ceil(Math.max(0, seconds)); - return count > 90 ? t(`Too many tries. Please try again in ${Math.ceil(count / 60)} min.`, `试错次数太多,请 ${Math.ceil(count / 60)} 分钟后再试。`) - : t(`Too many tries. Please try again in ${count} s.`, `试错次数太多,请 ${count} 秒后再试。`); + return count > 90 ? t(`Too many attempts. Try again in ${Math.ceil(count / 60)} min.`, `尝试次数过多,请 ${Math.ceil(count / 60)} 分钟后再试。`) + : t(`Too many attempts. Try again in ${count} s.`, `尝试次数过多,请 ${count} 秒后再试。`); } export function queueWords(index, t) { const n = index + 1, mod = n % 100;