Use session source latents for generated repaint#1170
Closed
Conversation
Contributor
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
📝 WalkthroughWalkthroughAdds a session-backed "most natural" repaint flow and makes Changes
Sequence Diagram(s)sequenceDiagram
participant UI as UI/User
participant API as API Layer
participant Params as GenerationParams
participant RetakeSession as RetakeSession
participant GenHandler as generate_music
participant Service as service_generate_execute
participant Diffusion as Diffusion/DiT
participant RetakeLat as RetakeLatents
participant Saver as RetakeSessionSave
UI->>API: POST (repaint_mode="most natural", source_session_dir)
API->>Params: build GenerationParams (resolve_repaint_mode)
Params->>RetakeSession: load_retake_source_track(session_dir, index)
RetakeSession-->>Params: audio_codes, latents, metadata
Params->>GenHandler: call generate_music with retake inputs
GenHandler->>RetakeLat: build_retake_mask(target_length, regions)
RetakeLat-->>GenHandler: repaint mask
GenHandler->>RetakeLat: align_retake_source_latents(source_latents, target_length)
RetakeLat-->>GenHandler: aligned source latents
GenHandler->>Service: service_generate(..., retake_source_latents, source_latent_mix_ratio)
Service->>RetakeLat: build_retake_step_skip_timesteps(infer_steps, mix_ratio)
RetakeLat-->>Service: retake timestep schedule
Service->>Diffusion: generate_audio(with retake-initialization)
Diffusion-->>Service: pred_latents
Service->>RetakeLat: splice_retake_latents(pred_latents, source_latents, mask, crossfade=25)
RetakeLat-->>Service: final latents
Service->>Diffusion: decode final latents (no waveform splice)
Diffusion-->>Service: audio
Service->>Saver: save_generation_session_artifacts(result, session_output_dir) [optional]
Saver-->>Service: saved paths
Service-->>UI: return audio + extra_outputs (retake metadata, session_output_dir)
Estimated code review effort🎯 4 (Complex) | ⏱️ ~75 minutes Possibly related PRs
Suggested reviewers
Poem
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 10
🧹 Nitpick comments (2)
acestep/core/generation/handler/retake_latents_test.py (1)
18-71: ⚡ Quick winAdd regression tests for empty-region and zero-frame-latent inputs.
Given the new helper responsibilities, it would be valuable to add tests asserting:
build_retake_mask(..., repainting_regions=[])raisesValueError.align_retake_source_latents()raises a clearValueErrorfor zero-frame source latents.Based on learnings: Applies to @(test.py|test.py) : Include at least: one success-path test, one regression/edge-case test for the bug being fixed, one non-target behavior check when relevant.
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@acestep/core/generation/handler/retake_latents_test.py` around lines 18 - 71, Add two regression unit tests: one that calls build_retake_mask with repainting_regions=[] and asserts it raises ValueError (to cover the empty-region edge case), and one that calls align_retake_source_latents with a zero-frame source tensor (e.g., shape (0, channels)) and asserts it raises ValueError (to cover zero-frame latent input). Place them alongside the existing tests in retake_latents_test.py, using the same test naming style (e.g., test_build_retake_mask_raises_on_empty_regions and test_align_retake_source_latents_raises_on_zero_frames) and ensure they exercise the functions build_retake_mask and align_retake_source_latents respectively.acestep/inference.py (1)
452-458: ⚡ Quick winAvoid mutating the caller's
GenerationParamsin place.This branch rewrites
task_type,thinking, CoT flags, andaudio_codeson the sharedparamsobject. If the same instance is reused for another call or inspected after generation, it no longer reflects the original request. Prefer localeffective_*values or a copied dataclass.As per coding guidelines, "Keep data flow explicit (data in, data out); side effects must be obvious and deliberate."
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@acestep/inference.py` around lines 452 - 458, This mutates the shared GenerationParams object (params) by overwriting task_type, thinking, CoT flags and audio_codes; instead create a non-mutating local version and use that for downstream work. Either construct a copy of params (e.g., via dataclasses.replace or a provided copy/clone method) and set task_type="text2music", thinking=False, use_cot_metas=False, use_cot_caption=False, use_cot_language=False, and audio_codes=retake_inputs["audio_codes"] on that copy, or compute local effective_* variables (effective_task_type, effective_thinking, effective_use_cot_metas, etc.) and use them (and audio_code_string_to_use) where the code currently reads params; do not assign back into params.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@acestep/core/generation/handler/generate_music.py`:
- Around line 460-473: The branch that calls splice_retake_latents is using raw
retake_source_latents from load_retake_source_track without normalizing shape or
duration to pred_latents, so before calling splice_retake_latents ensure
retake_source_latents is aligned: reshape/expand to match pred_latents.batch
(repeat or tile if single saved track but batch>1), trim or pad (with zeros or
suitable silence) to match pred_latents.shape[1] (frames), and resample or
adjust frames if sample-rate or length mismatches; perform this normalization
right after retake_source_latents is loaded (before
build_retake_mask/splice_retake_latents) and update any device/dtype to match
pred_latents to avoid misaligned splices or crashes in splice_retake_latents.
In `@acestep/core/generation/handler/retake_latents.py`:
- Around line 88-92: Before attempting to pad in the branch where
latents.shape[1] < target_length, validate that the sequence length is not zero
and raise a clear domain ValueError if it is; specifically, in the
function/block handling "latents" (the variable used in the shown code) add a
guard like "if latents.shape[1] == 0: raise ValueError(...)" with a descriptive
message indicating zero-frame source latents cannot be padded, so the subsequent
pad = latents[:, -1:, :].expand(...) and torch.cat(...) are never executed on an
empty time dimension.
- Line 39: Replace the current truthy fallback so an explicit empty list is
preserved: set regions = repainting_regions if repainting_regions is not None
else [{"start": repainting_start, "end": repainting_end}]. This ensures the
variable repainting_regions (and the assignment to regions) treats [] as a valid
value and only falls back to the single-region using
repainting_start/repainting_end when repainting_regions is None.
In `@acestep/core/generation/handler/retake_session_save.py`:
- Around line 31-58: The code can leave a partially written session dir (root)
when validation fails (e.g., missing pred_latents or per-track validation), so
change save_session_artifacts to write into a temporary directory (e.g.,
root_tmp under session_dir or root with a ".tmp" suffix) and only rename/move it
to the final directory (root) after all validations, _copy_session_audio,
np.save, and _write_json calls succeed; alternatively, ensure a
try/except/finally around the processing that on any exception removes the
partially written root (cleanup) before re-raising. Refer to the existing
symbols root/session_dir, pred_latents, tracks, _copy_session_audio, np.save,
and _write_json to locate the code to wrap or to change to a two-phase
write+rename approach.
In `@acestep/core/generation/handler/retake_session.py`:
- Around line 153-172: Add docstrings to the new helper functions to satisfy the
repo guideline requiring docstrings for new/modified Python functions: add a
brief one-line summary and a short description of parameters and return value
for _read_json(path: Path) -> dict[str, Any], _write_json(path: Path, value:
dict[str, Any]) -> None, _first_value(*values: Any) -> Any, and
_first_text(*values: Any) -> str; ensure each docstring states purpose (e.g.,
read/write JSON, return first non-empty value, convert to text), documents
parameters and return types, and uses triple-quoted string format placed
immediately after each function signature.
In `@acestep/ui/gradio/api/api_routes.py`:
- Around line 505-507: The current expression
float(get_param("source_latent_mix_ratio", default=0.3) or 0.3) treats 0/0.0 as
falsy and replaces it with 0.3; update the handling for source_latent_mix_ratio
so that an explicit 0 or 0.0 is preserved by checking for None/empty rather than
falsiness: call get_param("source_latent_mix_ratio", default=None) (or capture
its return), then set source_latent_mix_ratio = float(value) when value is not
None/empty, otherwise fall back to 0.3; adjust the code around the
source_latent_mix_ratio assignment (referencing get_param and the variable name
source_latent_mix_ratio) accordingly.
- Around line 508-513: The repainting_regions parameter is read as a raw JSON
string from get_param and passed into GenerationParams, causing type errors;
parse/decode it into a Python list[dict] before constructing GenerationParams.
Update the code path in api_routes.py where repainting_regions is retrieved (the
call to get_param("repainting_regions", ...)) to detect a non-None string and
run json.loads (or equivalent safe parsing) into a list, handling parsing errors
(e.g., return a 400 or default to None) so GenerationParams receives a proper
list[dict]; keep the existing save_session_artifacts/to_bool handling unchanged.
In `@acestep/ui/gradio/events/results/batch_queue.py`:
- Around line 180-182: restore_batch_parameters() currently ignores the new
source-session fields persisted into generation_params ("source_session_dir",
"source_track_index", "source_latent_mix_ratio"), so restoring a batch drops
source-session context; update restore_batch_parameters (and any code that
builds its return value) to read these keys from generation_params and include
them in the restored parameter set (same names) so restored batches preserve the
source_session_dir, source_track_index, and source_latent_mix_ratio values.
In `@acestep/ui/gradio/events/results/generation_progress.py`:
- Around line 188-193: The current expression int(source_track_index) if
source_track_index else 1 silently maps falsy values (e.g. 0) to 1; change this
to preserve explicit inputs and validate them: if source_track_index is None
then use the intended default (1), otherwise convert via int(source_track_index)
and raise a ValueError if the result is < 1. Update the handling where
source_track_index is processed (the expression using source_track_index in
generation_progress.py) to perform this explicit None-check, int() cast and >=1
validation and include a clear error message referencing the invalid
source_track_index value.
---
Nitpick comments:
In `@acestep/core/generation/handler/retake_latents_test.py`:
- Around line 18-71: Add two regression unit tests: one that calls
build_retake_mask with repainting_regions=[] and asserts it raises ValueError
(to cover the empty-region edge case), and one that calls
align_retake_source_latents with a zero-frame source tensor (e.g., shape (0,
channels)) and asserts it raises ValueError (to cover zero-frame latent input).
Place them alongside the existing tests in retake_latents_test.py, using the
same test naming style (e.g., test_build_retake_mask_raises_on_empty_regions and
test_align_retake_source_latents_raises_on_zero_frames) and ensure they exercise
the functions build_retake_mask and align_retake_source_latents respectively.
In `@acestep/inference.py`:
- Around line 452-458: This mutates the shared GenerationParams object (params)
by overwriting task_type, thinking, CoT flags and audio_codes; instead create a
non-mutating local version and use that for downstream work. Either construct a
copy of params (e.g., via dataclasses.replace or a provided copy/clone method)
and set task_type="text2music", thinking=False, use_cot_metas=False,
use_cot_caption=False, use_cot_language=False, and
audio_codes=retake_inputs["audio_codes"] on that copy, or compute local
effective_* variables (effective_task_type, effective_thinking,
effective_use_cot_metas, etc.) and use them (and audio_code_string_to_use) where
the code currently reads params; do not assign back into params.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: 40b6f295-d02f-4233-80fc-d749f77c794d
📒 Files selected for processing (31)
acestep/api/http/release_task_models.pyacestep/api/http/release_task_models_test.pyacestep/api/http/release_task_request_builder.pyacestep/api/http/release_task_request_builder_test.pyacestep/api/job_generation_setup.pyacestep/core/generation/handler/generate_music.pyacestep/core/generation/handler/generate_music_execute.pyacestep/core/generation/handler/generate_music_payload.pyacestep/core/generation/handler/generate_music_test.pyacestep/core/generation/handler/resolve_repaint_config_test.pyacestep/core/generation/handler/retake_latents.pyacestep/core/generation/handler/retake_latents_test.pyacestep/core/generation/handler/retake_session.pyacestep/core/generation/handler/retake_session_save.pyacestep/core/generation/handler/retake_session_test.pyacestep/core/generation/handler/service_generate.pyacestep/core/generation/handler/service_generate_execute.pyacestep/core/generation/handler/service_generate_execute_test.pyacestep/inference.pyacestep/ui/gradio/api/api_routes.pyacestep/ui/gradio/events/results/batch_management_background.pyacestep/ui/gradio/events/results/batch_management_helpers.pyacestep/ui/gradio/events/results/batch_management_wrapper.pyacestep/ui/gradio/events/results/batch_queue.pyacestep/ui/gradio/events/results/batch_queue_test.pyacestep/ui/gradio/events/results/generation_progress.pyacestep/ui/gradio/events/wiring/generation_batch_navigation_wiring.pyacestep/ui/gradio/events/wiring/generation_mode_wiring.pyacestep/ui/gradio/events/wiring/generation_mode_wiring_test.pyacestep/ui/gradio/events/wiring/generation_run_wiring.pyacestep/ui/gradio/interfaces/generation_tab_secondary_controls.py
Comment on lines
+31
to
+58
| root = Path(session_dir).expanduser().resolve() | ||
| root.mkdir(parents=True, exist_ok=True) | ||
| extra = result.extra_outputs or {} | ||
| lm_metadata = extra.get("lm_metadata") | ||
| if lm_metadata is not None: | ||
| _write_json(root / "lm_metadata.json", lm_metadata) | ||
|
|
||
| pred_latents = extra.get("pred_latents") | ||
| if pred_latents is None: | ||
| raise ValueError("save_session_artifacts requires pred_latents in generation extra_outputs") | ||
|
|
||
| tracks = [] | ||
| for index, audio in enumerate(result.audios or [], start=1): | ||
| params = dict(audio.get("params") or {}) | ||
| audio_codes = str(params.get("audio_codes") or params.get("audio_code_string") or "").strip() | ||
| if not audio_codes: | ||
| raise ValueError("save_session_artifacts requires per-track audio_codes") | ||
| params_path = root / f"{index:02d}_params.json" | ||
| _copy_session_audio(root, params, audio, index) | ||
| if index - 1 >= pred_latents.shape[0]: | ||
| raise ValueError("save_session_artifacts requires one latent tensor per track") | ||
| latent = pred_latents[index - 1].detach().cpu().float().numpy() | ||
| np.save(root / f"{index:02d}_latents.npy", latent) | ||
| params["session_latents_file"] = f"{index:02d}_latents.npy" | ||
| _write_json(params_path, params) | ||
| tracks.append({"index": index, "params_file": params_path.name}) | ||
|
|
||
| _write_json(root / "session.json", {"source": source, "tracks": tracks}) |
Contributor
There was a problem hiding this comment.
Avoid leaving half-written session directories on validation failure.
This helper creates session_dir, writes lm_metadata.json, and persists tracks as it validates them. If pred_latents is missing or track n fails later validation, the run still leaves behind a partial session that can be mistaken for a reusable source session on the next attempt.
Suggested fix
- root = Path(session_dir).expanduser().resolve()
- root.mkdir(parents=True, exist_ok=True)
- extra = result.extra_outputs or {}
- lm_metadata = extra.get("lm_metadata")
- if lm_metadata is not None:
- _write_json(root / "lm_metadata.json", lm_metadata)
-
- pred_latents = extra.get("pred_latents")
+ root = Path(session_dir).expanduser().resolve()
+ extra = result.extra_outputs or {}
+ lm_metadata = extra.get("lm_metadata")
+ pred_latents = extra.get("pred_latents")
if pred_latents is None:
raise ValueError("save_session_artifacts requires pred_latents in generation extra_outputs")
- tracks = []
+ pending_tracks = []
for index, audio in enumerate(result.audios or [], start=1):
params = dict(audio.get("params") or {})
audio_codes = str(params.get("audio_codes") or params.get("audio_code_string") or "").strip()
if not audio_codes:
raise ValueError("save_session_artifacts requires per-track audio_codes")
- params_path = root / f"{index:02d}_params.json"
- _copy_session_audio(root, params, audio, index)
if index - 1 >= pred_latents.shape[0]:
raise ValueError("save_session_artifacts requires one latent tensor per track")
latent = pred_latents[index - 1].detach().cpu().float().numpy()
- np.save(root / f"{index:02d}_latents.npy", latent)
- params["session_latents_file"] = f"{index:02d}_latents.npy"
- _write_json(params_path, params)
- tracks.append({"index": index, "params_file": params_path.name})
+ pending_tracks.append((index, audio, params, latent))
+ root.mkdir(parents=True, exist_ok=True)
+ if lm_metadata is not None:
+ _write_json(root / "lm_metadata.json", lm_metadata)
+
+ tracks = []
+ for index, audio, params, latent in pending_tracks:
+ params_path = root / f"{index:02d}_params.json"
+ _copy_session_audio(root, params, audio, index)
+ np.save(root / f"{index:02d}_latents.npy", latent)
+ params["session_latents_file"] = f"{index:02d}_latents.npy"
+ _write_json(params_path, params)
+ tracks.append({"index": index, "params_file": params_path.name})
+
_write_json(root / "session.json", {"source": source, "tracks": tracks})🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@acestep/core/generation/handler/retake_session_save.py` around lines 31 - 58,
The code can leave a partially written session dir (root) when validation fails
(e.g., missing pred_latents or per-track validation), so change
save_session_artifacts to write into a temporary directory (e.g., root_tmp under
session_dir or root with a ".tmp" suffix) and only rename/move it to the final
directory (root) after all validations, _copy_session_audio, np.save, and
_write_json calls succeed; alternatively, ensure a try/except/finally around the
processing that on any exception removes the partially written root (cleanup)
before re-raising. Refer to the existing symbols root/session_dir, pred_latents,
tracks, _copy_session_audio, np.save, and _write_json to locate the code to wrap
or to change to a two-phase write+rename approach.
Comment on lines
+505
to
+507
| source_latent_mix_ratio=float( | ||
| get_param("source_latent_mix_ratio", default=0.3) or 0.3, | ||
| ), |
Contributor
There was a problem hiding this comment.
Preserve an explicit 0.0 source mix ratio.
float(get_param(...) or 0.3) treats 0 / 0.0 as “missing”, so callers cannot disable source-latent mixing through the API. That silently changes generation behavior for a valid boundary value.
Suggested fix
+ raw_source_latent_mix_ratio = get_param("source_latent_mix_ratio", default=0.3)
+
params = GenerationParams(
@@
- source_latent_mix_ratio=float(
- get_param("source_latent_mix_ratio", default=0.3) or 0.3,
- ),
+ source_latent_mix_ratio=(
+ 0.3
+ if raw_source_latent_mix_ratio is None
+ else float(raw_source_latent_mix_ratio)
+ ),🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@acestep/ui/gradio/api/api_routes.py` around lines 505 - 507, The current
expression float(get_param("source_latent_mix_ratio", default=0.3) or 0.3)
treats 0/0.0 as falsy and replaces it with 0.3; update the handling for
source_latent_mix_ratio so that an explicit 0 or 0.0 is preserved by checking
for None/empty rather than falsiness: call get_param("source_latent_mix_ratio",
default=None) (or capture its return), then set source_latent_mix_ratio =
float(value) when value is not None/empty, otherwise fall back to 0.3; adjust
the code around the source_latent_mix_ratio assignment (referencing get_param
and the variable name source_latent_mix_ratio) accordingly.
Comment on lines
+508
to
+513
| repainting_regions=get_param("repainting_regions", default=None), | ||
| save_session_artifacts=to_bool( | ||
| get_param("save_session_artifacts", default=False), | ||
| False, | ||
| ), | ||
| session_output_dir=get_param("session_output_dir", default=None), |
Contributor
There was a problem hiding this comment.
Decode repainting_regions before constructing GenerationParams.
For form requests, get_param() returns strings, so this field reaches GenerationParams as raw JSON text instead of list[dict]. The new multi-region repaint path will then fail late or ignore the requested regions.
Suggested fix
+ raw_repainting_regions = get_param("repainting_regions", default=None)
+ if isinstance(raw_repainting_regions, str) and raw_repainting_regions.strip():
+ try:
+ raw_repainting_regions = json.loads(raw_repainting_regions)
+ except json.JSONDecodeError as exc:
+ raise HTTPException(
+ status_code=400,
+ detail="repainting_regions must be valid JSON",
+ ) from exc
+
params = GenerationParams(
@@
- repainting_regions=get_param("repainting_regions", default=None),
+ repainting_regions=raw_repainting_regions,📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| repainting_regions=get_param("repainting_regions", default=None), | |
| save_session_artifacts=to_bool( | |
| get_param("save_session_artifacts", default=False), | |
| False, | |
| ), | |
| session_output_dir=get_param("session_output_dir", default=None), | |
| raw_repainting_regions = get_param("repainting_regions", default=None) | |
| if isinstance(raw_repainting_regions, str) and raw_repainting_regions.strip(): | |
| try: | |
| raw_repainting_regions = json.loads(raw_repainting_regions) | |
| except json.JSONDecodeError as exc: | |
| raise HTTPException( | |
| status_code=400, | |
| detail="repainting_regions must be valid JSON", | |
| ) from exc | |
| params = GenerationParams( | |
| repainting_regions=raw_repainting_regions, | |
| save_session_artifacts=to_bool( | |
| get_param("save_session_artifacts", default=False), | |
| False, | |
| ), | |
| session_output_dir=get_param("session_output_dir", default=None), |
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@acestep/ui/gradio/api/api_routes.py` around lines 508 - 513, The
repainting_regions parameter is read as a raw JSON string from get_param and
passed into GenerationParams, causing type errors; parse/decode it into a Python
list[dict] before constructing GenerationParams. Update the code path in
api_routes.py where repainting_regions is retrieved (the call to
get_param("repainting_regions", ...)) to detect a non-None string and run
json.loads (or equivalent safe parsing) into a list, handling parsing errors
(e.g., return a 400 or default to None) so GenerationParams receives a proper
list[dict]; keep the existing save_session_artifacts/to_bool handling unchanged.
Comment on lines
+180
to
+182
| "source_session_dir": source_session_dir, | ||
| "source_track_index": source_track_index, | ||
| "source_latent_mix_ratio": source_latent_mix_ratio, |
Contributor
There was a problem hiding this comment.
Restore path still drops the new source-session fields.
These values are now persisted into generation_params, but restore_batch_parameters() below never reads or returns them. Restoring a batch will therefore clear the source-session context and make "most natural" edits non-reproducible from history.
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@acestep/ui/gradio/events/results/batch_queue.py` around lines 180 - 182,
restore_batch_parameters() currently ignores the new source-session fields
persisted into generation_params ("source_session_dir", "source_track_index",
"source_latent_mix_ratio"), so restoring a batch drops source-session context;
update restore_batch_parameters (and any code that builds its return value) to
read these keys from generation_params and include them in the restored
parameter set (same names) so restored batches preserve the source_session_dir,
source_track_index, and source_latent_mix_ratio values.
Comment on lines
+188
to
193
| source_session_dir=source_session_dir or None, | ||
| source_track_index=int(source_track_index) if source_track_index else 1, | ||
| source_latent_mix_ratio=( | ||
| float(source_latent_mix_ratio) if source_latent_mix_ratio is not None else 0.3 | ||
| ), | ||
| retake_variance=float(retake_variance) if retake_variance is not None else 0.0, |
Contributor
There was a problem hiding this comment.
Avoid silently coercing invalid source_track_index to track 1.
Line 189 maps falsy values (notably 0) to 1, which can target the wrong source track instead of surfacing invalid input. Please preserve explicit values and validate >= 1 so callers get a clear error path.
Proposed fix
+ parsed_source_track_index = (
+ int(source_track_index) if source_track_index is not None else 1
+ )
+ if parsed_source_track_index < 1:
+ raise ValueError("source_track_index must be >= 1")
+
gen_params = GenerationParams(
@@
- source_track_index=int(source_track_index) if source_track_index else 1,
+ source_track_index=parsed_source_track_index,As per coding guidelines: "Handle errors explicitly" and "Keep data flow explicit (data in, data out)."
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@acestep/ui/gradio/events/results/generation_progress.py` around lines 188 -
193, The current expression int(source_track_index) if source_track_index else 1
silently maps falsy values (e.g. 0) to 1; change this to preserve explicit
inputs and validate them: if source_track_index is None then use the intended
default (1), otherwise convert via int(source_track_index) and raise a
ValueError if the result is < 1. Update the handling where source_track_index is
processed (the expression using source_track_index in generation_progress.py) to
perform this explicit None-check, int() cast and >=1 validation and include a
clear error message referencing the invalid source_track_index value.
ChuxiJ
force-pushed
the
codex/most-natural-repaint
branch
from
9efe0b4 to
c67fb21
Compare
ChuxiJ
force-pushed
the
codex/most-natural-repaint
branch
from
c67fb21 to
e79dcfa
Compare
5 tasks
Contributor
There was a problem hiding this comment.
Actionable comments posted: 5
♻️ Duplicate comments (1)
acestep/ui/gradio/events/results/batch_queue.py (1)
180-182:⚠️ Potential issue | 🟠 Major | ⚡ Quick winRestore path drops newly captured
source_*repaint fields.
capture_current_params()now persists source-session context, butrestore_batch_parameters()never restores it. This clears session-backed repaint state when navigating/restoring historical batches.Suggested fix
def restore_batch_parameters(current_batch_index, batch_queue): @@ if current_batch_index not in batch_queue: gr.Warning(t("messages.no_batch_data")) - return [gr.update()] * 33 + return [gr.update()] * 36 @@ retake_variance = params.get("retake_variance", 0.0) retake_seed = params.get("retake_seed", "") + source_session_dir = params.get("source_session_dir", "") + source_track_index = params.get("source_track_index", 1) + source_latent_mix_ratio = params.get("source_latent_mix_ratio", 0.3) @@ latent_shift, latent_rescale, no_fsq, retake_variance, retake_seed, + source_session_dir, source_track_index, source_latent_mix_ratio, )Also applies to: 188-259
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@acestep/ui/gradio/events/results/batch_queue.py` around lines 180 - 182, restore_batch_parameters() currently omits restoring the session-backed repaint fields added by capture_current_params() (the source_session_dir, source_track_index, source_latent_mix_ratio and any other source_* fields), causing captured source context to be lost when a batch is restored; update restore_batch_parameters() to read and reassign these source_* keys from the saved batch dict back into the live params/state (same identifiers used in capture_current_params()), and ensure any loops/branches in restore_batch_parameters() that restore other param groups (lines ~188-259) also include these source_* fields so session-backed repaint context is preserved across navigation/restores.
🧹 Nitpick comments (1)
acestep/core/generation/handler/retake_latents.py (1)
69-81: ⚡ Quick winDocument raised exceptions in public function docstrings.
These functions raise
ValueErrorbut their docstrings omit aRaisessection. Add explicit exception docs for API clarity and PR-readiness compliance.📝 Suggested docstring update
def align_retake_source_latents( @@ Returns: Tensor shaped ``[batch_size, target_length, C]``. + + Raises: + ValueError: If ``source_latents`` is not shaped ``[T, C]`` or ``[B, T, C]``. """ def splice_retake_latents( @@ Returns: Spliced latents where edit frames come from ``pred_latents`` and non-edit frames come from source latents. + + Raises: + ValueError: If ``repaint_mask`` is not 2D or does not match ``pred_latents`` ``[B, T]``. """ def build_retake_step_skip_timesteps( @@ Returns: Truncated timestep list, or ``None`` when no skip is needed. + + Raises: + ValueError: If ``mix_ratio >= 1.0`` or ``infer_steps < 1``. """As per coding guidelines: "Docstrings are mandatory for all new/modified Python modules, classes, and functions... include ... raised exceptions when relevant."
Also applies to: 107-118, 144-153
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@acestep/core/generation/handler/retake_latents.py` around lines 69 - 81, Update the public function docstring that begins "Align saved source latents to a generated latent batch shape" (in acestep/core/generation/handler/retake_latents.py) to include a Raises section documenting the ValueError(s) the function can throw (e.g., when source_latents has an unexpected shape or when target_length/batch_size are incompatible with the input leading to alignment failure); make the wording explicit and concise (e.g., "Raises: ValueError: If source_latents has unsupported shape or cannot be aligned to the requested target_length and batch_size"). Repeat the same addition for the two other similar docstrings in the file (the blocks noted around the other docstring locations) so all public functions document their ValueError cases.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@acestep/api/http/release_task_models_test.py`:
- Around line 36-45: The test uses a hardcoded "/tmp/source-session" which
triggers Ruff S108; replace the literal with a temporary directory created at
runtime (e.g., use tempfile.mkdtemp() or a tempfile.TemporaryDirectory()
context) and pass that value into GenerateMusicRequest as source_session_dir,
then assert equality against that temp variable; ensure you import tempfile and
clean up the temp directory (or rely on TemporaryDirectory context) to avoid
leaving artifacts.
In `@acestep/core/generation/handler/generate_music.py`:
- Around line 298-305: The handler currently only rewrites task_type "cover" to
"text2music" for session-retake flows; update the logic so that whenever
is_session_retake_mode(repaint_mode) returns true (uses_session_retake) any
non-"text2music" task_type returned by self._resolve_generate_music_task(...) is
forced to "text2music" (e.g., replace the if uses_session_retake and task_type
== "cover": clause with a check like if uses_session_retake and task_type !=
"text2music": task_type = "text2music"), referencing uses_session_retake,
is_session_retake_mode, task_type and _resolve_generate_music_task to locate the
code to change.
In `@acestep/core/generation/handler/retake_session.py`:
- Around line 73-79: Replace the silent clamping of track_index and instead
validate it: convert track_index to int, then if int(track_index) <= 0 raise a
ValueError (e.g., "invalid track_index: must be >= 1") rather than using max(1,
...); keep using that validated index to build params_path/latents_path (refer
to the track_index variable and the params_path/latents_path construction) so
callers get an explicit error for 0 or negative inputs.
In `@acestep/inference.py`:
- Around line 995-1000: When params.save_session_artifacts is true and
params.session_output_dir is falsy, create a unique directory name (e.g., append
a timestamp or UUID) instead of always using "session_artifacts" so each
generation gets its own folder; update the branch in inference where session_dir
is computed (referencing params.save_session_artifacts and
params.session_output_dir) to build a unique session_dir (using save_dir or
os.getcwd() as base) before calling
save_generation_session_artifacts(result=final_result, session_dir=session_dir)
and setting final_result.extra_outputs["session_output_dir"].
In `@acestep/ui/gradio/events/results/batch_queue_test.py`:
- Around line 257-267: The test test_source_session_params_included_in_capture
uses a hardcoded "/tmp/session" which triggers Ruff S108; change the marker to a
neutral, non-/tmp string (e.g., "session_marker" or "test-session") in the call
to _build_args and assert that capture_current_params returns that value,
leaving the other asserts for source_track_index and source_latent_mix_ratio
unchanged so only the directory string is updated; locate this in the test
method test_source_session_params_included_in_capture and the call to
capture_current_params/_build_args.
---
Duplicate comments:
In `@acestep/ui/gradio/events/results/batch_queue.py`:
- Around line 180-182: restore_batch_parameters() currently omits restoring the
session-backed repaint fields added by capture_current_params() (the
source_session_dir, source_track_index, source_latent_mix_ratio and any other
source_* fields), causing captured source context to be lost when a batch is
restored; update restore_batch_parameters() to read and reassign these source_*
keys from the saved batch dict back into the live params/state (same identifiers
used in capture_current_params()), and ensure any loops/branches in
restore_batch_parameters() that restore other param groups (lines ~188-259) also
include these source_* fields so session-backed repaint context is preserved
across navigation/restores.
---
Nitpick comments:
In `@acestep/core/generation/handler/retake_latents.py`:
- Around line 69-81: Update the public function docstring that begins "Align
saved source latents to a generated latent batch shape" (in
acestep/core/generation/handler/retake_latents.py) to include a Raises section
documenting the ValueError(s) the function can throw (e.g., when source_latents
has an unexpected shape or when target_length/batch_size are incompatible with
the input leading to alignment failure); make the wording explicit and concise
(e.g., "Raises: ValueError: If source_latents has unsupported shape or cannot be
aligned to the requested target_length and batch_size"). Repeat the same
addition for the two other similar docstrings in the file (the blocks noted
around the other docstring locations) so all public functions document their
ValueError cases.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: 7d249e7b-bb1e-4512-a7bd-7e0899e2a37a
📒 Files selected for processing (35)
acestep/api/http/release_task_models.pyacestep/api/http/release_task_models_test.pyacestep/api/http/release_task_request_builder.pyacestep/api/http/release_task_request_builder_test.pyacestep/api/job_generation_setup.pyacestep/core/generation/handler/generate_music.pyacestep/core/generation/handler/generate_music_execute.pyacestep/core/generation/handler/generate_music_payload.pyacestep/core/generation/handler/generate_music_test.pyacestep/core/generation/handler/resolve_repaint_config_test.pyacestep/core/generation/handler/retake_latents.pyacestep/core/generation/handler/retake_latents_test.pyacestep/core/generation/handler/retake_session.pyacestep/core/generation/handler/retake_session_save.pyacestep/core/generation/handler/retake_session_test.pyacestep/core/generation/handler/service_generate.pyacestep/core/generation/handler/service_generate_execute.pyacestep/core/generation/handler/service_generate_execute_test.pyacestep/inference.pyacestep/ui/gradio/api/api_routes.pyacestep/ui/gradio/events/results/audio_transfer.pyacestep/ui/gradio/events/results/audio_transfer_test.pyacestep/ui/gradio/events/results/batch_management_background.pyacestep/ui/gradio/events/results/batch_management_helpers.pyacestep/ui/gradio/events/results/batch_management_wrapper.pyacestep/ui/gradio/events/results/batch_queue.pyacestep/ui/gradio/events/results/batch_queue_test.pyacestep/ui/gradio/events/results/generation_progress.pyacestep/ui/gradio/events/results/generation_progress_test.pyacestep/ui/gradio/events/wiring/generation_batch_navigation_wiring.pyacestep/ui/gradio/events/wiring/generation_mode_wiring.pyacestep/ui/gradio/events/wiring/generation_mode_wiring_test.pyacestep/ui/gradio/events/wiring/generation_run_wiring.pyacestep/ui/gradio/events/wiring/results_aux_wiring.pyacestep/ui/gradio/interfaces/generation_tab_secondary_controls.py
🚧 Files skipped from review as they are similar to previous changes (10)
- acestep/ui/gradio/events/wiring/generation_run_wiring.py
- acestep/core/generation/handler/service_generate.py
- acestep/core/generation/handler/resolve_repaint_config_test.py
- acestep/core/generation/handler/generate_music_execute.py
- acestep/ui/gradio/events/wiring/generation_mode_wiring_test.py
- acestep/ui/gradio/api/api_routes.py
- acestep/ui/gradio/interfaces/generation_tab_secondary_controls.py
- acestep/api/http/release_task_request_builder.py
- acestep/core/generation/handler/retake_latents_test.py
- acestep/core/generation/handler/service_generate_execute.py
Comment on lines
+36
to
+45
| req = GenerateMusicRequest( | ||
| audio_code_string="<|audio_code_1|>", | ||
| cover_noise_strength=0.75, | ||
| repaint_mode="most natural", | ||
| source_session_dir="/tmp/source-session", | ||
| ) | ||
| self.assertEqual("<|audio_code_1|>", req.audio_code_string) | ||
| self.assertAlmostEqual(0.75, req.cover_noise_strength) | ||
| self.assertEqual("most natural", req.repaint_mode) | ||
| self.assertEqual("/tmp/source-session", req.source_session_dir) |
Contributor
There was a problem hiding this comment.
Avoid /tmp literals in tests to prevent Ruff S108 failures.
"/tmp/source-session" is only a marker value here, but it triggers S108 and can fail lint.
Suggested change
- source_session_dir="/tmp/source-session",
+ source_session_dir="source-session-dir",
...
- self.assertEqual("/tmp/source-session", req.source_session_dir)
+ self.assertEqual("source-session-dir", req.source_session_dir)📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| req = GenerateMusicRequest( | |
| audio_code_string="<|audio_code_1|>", | |
| cover_noise_strength=0.75, | |
| repaint_mode="most natural", | |
| source_session_dir="/tmp/source-session", | |
| ) | |
| self.assertEqual("<|audio_code_1|>", req.audio_code_string) | |
| self.assertAlmostEqual(0.75, req.cover_noise_strength) | |
| self.assertEqual("most natural", req.repaint_mode) | |
| self.assertEqual("/tmp/source-session", req.source_session_dir) | |
| req = GenerateMusicRequest( | |
| audio_code_string="<|audio_code_1|>", | |
| cover_noise_strength=0.75, | |
| repaint_mode="most natural", | |
| source_session_dir="source-session-dir", | |
| ) | |
| self.assertEqual("<|audio_code_1|>", req.audio_code_string) | |
| self.assertAlmostEqual(0.75, req.cover_noise_strength) | |
| self.assertEqual("most natural", req.repaint_mode) | |
| self.assertEqual("source-session-dir", req.source_session_dir) |
🧰 Tools
🪛 Ruff (0.15.12)
[error] 40-40: Probable insecure usage of temporary file or directory: "/tmp/source-session"
(S108)
[error] 45-45: Probable insecure usage of temporary file or directory: "/tmp/source-session"
(S108)
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@acestep/api/http/release_task_models_test.py` around lines 36 - 45, The test
uses a hardcoded "/tmp/source-session" which triggers Ruff S108; replace the
literal with a temporary directory created at runtime (e.g., use
tempfile.mkdtemp() or a tempfile.TemporaryDirectory() context) and pass that
value into GenerateMusicRequest as source_session_dir, then assert equality
against that temp variable; ensure you import tempfile and clean up the temp
directory (or rely on TemporaryDirectory context) to avoid leaving artifacts.
Comment on lines
+298
to
+305
| uses_session_retake = is_session_retake_mode(repaint_mode) | ||
| task_type, instruction = self._resolve_generate_music_task( | ||
| task_type=task_type, | ||
| audio_code_string=audio_code_string, | ||
| instruction=instruction, | ||
| ) | ||
| if uses_session_retake and task_type == "cover": | ||
| task_type = "text2music" |
Contributor
There was a problem hiding this comment.
Force all session-retake calls onto the text2music path.
Right now only "cover" is rewritten. If a caller reaches this handler with repaint_mode="most natural" and task_type="repaint", it still goes down repaint conditioning instead of the required audio-code/text2music flow.
Proposed fix
- if uses_session_retake and task_type == "cover":
+ if uses_session_retake and task_type != "text2music":
task_type = "text2music"🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@acestep/core/generation/handler/generate_music.py` around lines 298 - 305,
The handler currently only rewrites task_type "cover" to "text2music" for
session-retake flows; update the logic so that whenever
is_session_retake_mode(repaint_mode) returns true (uses_session_retake) any
non-"text2music" task_type returned by self._resolve_generate_music_task(...) is
forced to "text2music" (e.g., replace the if uses_session_retake and task_type
== "cover": clause with a check like if uses_session_retake and task_type !=
"text2music": task_type = "text2music"), referencing uses_session_retake,
is_session_retake_mode, task_type and _resolve_generate_music_task to locate the
code to change.
Comment on lines
+995
to
+1000
| if params.save_session_artifacts: | ||
| session_dir = params.session_output_dir | ||
| if not session_dir: | ||
| session_dir = os.path.join(save_dir or os.getcwd(), "session_artifacts") | ||
| save_generation_session_artifacts(result=final_result, session_dir=session_dir) | ||
| final_result.extra_outputs["session_output_dir"] = session_dir |
Contributor
There was a problem hiding this comment.
Default session_output_dir should be unique per generation.
When save_session_artifacts=True and the caller leaves session_output_dir empty, every run writes into the same session_artifacts folder. That can overwrite or mix prior tracks and make later "most natural" edits read stale artifacts. The Gradio path already creates a unique directory; this API path should too.
Proposed fix
if params.save_session_artifacts:
session_dir = params.session_output_dir
if not session_dir:
- session_dir = os.path.join(save_dir or os.getcwd(), "session_artifacts")
+ session_dir = tempfile.mkdtemp(
+ prefix="session_artifacts_",
+ dir=save_dir or os.getcwd(),
+ )
save_generation_session_artifacts(result=final_result, session_dir=session_dir)
final_result.extra_outputs["session_output_dir"] = session_dir🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@acestep/inference.py` around lines 995 - 1000, When
params.save_session_artifacts is true and params.session_output_dir is falsy,
create a unique directory name (e.g., append a timestamp or UUID) instead of
always using "session_artifacts" so each generation gets its own folder; update
the branch in inference where session_dir is computed (referencing
params.save_session_artifacts and params.session_output_dir) to build a unique
session_dir (using save_dir or os.getcwd() as base) before calling
save_generation_session_artifacts(result=final_result, session_dir=session_dir)
and setting final_result.extra_outputs["session_output_dir"].
Comment on lines
+257
to
+267
| def test_source_session_params_included_in_capture(self): | ||
| """Session-backed repaint inputs must be captured for AutoGen batches.""" | ||
| args = self._build_args( | ||
| source_session_dir="/tmp/session", | ||
| source_track_index=2, | ||
| source_latent_mix_ratio=0.25, | ||
| ) | ||
| result = capture_current_params(*args) | ||
| self.assertEqual("/tmp/session", result["source_session_dir"]) | ||
| self.assertEqual(2, result["source_track_index"]) | ||
| self.assertAlmostEqual(0.25, result["source_latent_mix_ratio"]) |
Contributor
There was a problem hiding this comment.
Use a non-/tmp marker value here to avoid Ruff S108 noise/failures.
This test only checks value round-trip, so a neutral string avoids security-lint violations.
Suggested change
- source_session_dir="/tmp/session",
+ source_session_dir="session-dir",
...
- self.assertEqual("/tmp/session", result["source_session_dir"])
+ self.assertEqual("session-dir", result["source_session_dir"])📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| def test_source_session_params_included_in_capture(self): | |
| """Session-backed repaint inputs must be captured for AutoGen batches.""" | |
| args = self._build_args( | |
| source_session_dir="/tmp/session", | |
| source_track_index=2, | |
| source_latent_mix_ratio=0.25, | |
| ) | |
| result = capture_current_params(*args) | |
| self.assertEqual("/tmp/session", result["source_session_dir"]) | |
| self.assertEqual(2, result["source_track_index"]) | |
| self.assertAlmostEqual(0.25, result["source_latent_mix_ratio"]) | |
| def test_source_session_params_included_in_capture(self): | |
| """Session-backed repaint inputs must be captured for AutoGen batches.""" | |
| args = self._build_args( | |
| source_session_dir="session-dir", | |
| source_track_index=2, | |
| source_latent_mix_ratio=0.25, | |
| ) | |
| result = capture_current_params(*args) | |
| self.assertEqual("session-dir", result["source_session_dir"]) | |
| self.assertEqual(2, result["source_track_index"]) | |
| self.assertAlmostEqual(0.25, result["source_latent_mix_ratio"]) |
🧰 Tools
🪛 Ruff (0.15.12)
[error] 260-260: Probable insecure usage of temporary file or directory: "/tmp/session"
(S108)
[error] 265-265: Probable insecure usage of temporary file or directory: "/tmp/session"
(S108)
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@acestep/ui/gradio/events/results/batch_queue_test.py` around lines 257 - 267,
The test test_source_session_params_included_in_capture uses a hardcoded
"/tmp/session" which triggers Ruff S108; change the marker to a neutral,
non-/tmp string (e.g., "session_marker" or "test-session") in the call to
_build_args and assert that capture_current_params returns that value, leaving
the other asserts for source_track_index and source_latent_mix_ratio unchanged
so only the directory string is updated; locate this in the test method
test_source_session_params_included_in_capture and the call to
capture_current_params/_build_args.
Contributor
There was a problem hiding this comment.
Actionable comments posted: 3
♻️ Duplicate comments (8)
acestep/ui/gradio/api/api_routes.py (2)
508-508:⚠️ Potential issue | 🟠 Major | ⚡ Quick winParse
repainting_regionsbefore buildingGenerationParams.Line 508 forwards raw form values; when sent as JSON text it remains a string instead of
list[dict], which can break multi-region repaint processing later.Suggested patch
+ raw_repainting_regions = get_param("repainting_regions", default=None) + if isinstance(raw_repainting_regions, str) and raw_repainting_regions.strip(): + try: + raw_repainting_regions = json.loads(raw_repainting_regions) + except json.JSONDecodeError as exc: + raise HTTPException( + status_code=400, + detail="repainting_regions must be valid JSON", + ) from exc @@ - repainting_regions=get_param("repainting_regions", default=None), + repainting_regions=raw_repainting_regions,🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@acestep/ui/gradio/api/api_routes.py` at line 508, The code forwards raw form value for repainting_regions into GenerationParams (via the call that uses get_param("repainting_regions", default=None)), which leaves JSON text as a string instead of a list[dict]; update the handler to parse repainting_regions before constructing GenerationParams: read the value from get_param, if it's a non-empty str attempt to json.loads it (import json) and validate it is a list (or list of dicts) and fall back to None on parse/validation failure, then pass the parsed value to GenerationParams instead of the raw string; reference the repainting_regions parameter, the get_param call, and the GenerationParams construction so the change is applied where the value is forwarded.
505-507:⚠️ Potential issue | 🟠 Major | ⚡ Quick winPreserve explicit
0.0forsource_latent_mix_ratio.Line 505–507 currently treats
0as missing and silently resets it to0.3, which changes valid caller intent.Suggested patch
+ raw_source_latent_mix_ratio = get_param("source_latent_mix_ratio", default=None) @@ - source_latent_mix_ratio=float( - get_param("source_latent_mix_ratio", default=0.3) or 0.3, - ), + source_latent_mix_ratio=( + 0.3 + if raw_source_latent_mix_ratio in (None, "") + else float(raw_source_latent_mix_ratio) + ),🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@acestep/ui/gradio/api/api_routes.py` around lines 505 - 507, The code is treating 0.0 as missing because of the "or 0.3" fallback; change the logic around get_param("source_latent_mix_ratio", default=0.3) so you only apply the 0.3 default when the returned value is actually None/empty, not when it is falsy (0.0). Locate the call where source_latent_mix_ratio is set and replace the "float(get_param(... ) or 0.3)" pattern with logic that captures the raw result (e.g., val = get_param(..., default=0.3)) and then use float(val) unless val is None/empty, in which case use 0.3, so explicit 0.0 is preserved.acestep/ui/gradio/events/results/batch_queue.py (1)
180-182:⚠️ Potential issue | 🟠 Major | ⚡ Quick winSource-session fields are captured but still not restored.
On Line 180–182, these values are persisted, but
restore_batch_parameters()does not read/return them, so restoring a batch drops"most natural"session context.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@acestep/ui/gradio/events/results/batch_queue.py` around lines 180 - 182, The code persists "source_session_dir", "source_track_index", and "source_latent_mix_ratio" but restore_batch_parameters() doesn't restore them; update restore_batch_parameters (and any callers that expect its return) to read these three keys from the stored batch dict and include them in its returned parameters structure (e.g., add source_session_dir, source_track_index, source_latent_mix_ratio to the returned dict or dataclass), and ensure downstream code that uses restore_batch_parameters (e.g., batch queue consumer/processor functions) consumes these values so the original session context is fully restored.acestep/api/http/release_task_models_test.py (1)
40-45:⚠️ Potential issue | 🟠 Major | ⚡ Quick winReplace
"/tmp/source-session"literal to avoid Ruff S108.Line 40 and Line 45 still use a temp-dir literal that triggers security lint in tests.
Suggested patch
- source_session_dir="/tmp/source-session", + source_session_dir="source-session-dir", @@ - self.assertEqual("/tmp/source-session", req.source_session_dir) + self.assertEqual("source-session-dir", req.source_session_dir)🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@acestep/api/http/release_task_models_test.py` around lines 40 - 45, Replace the hardcoded "/tmp/source-session" literal with a temp directory created at test runtime (e.g., use tempfile.mkdtemp() or the pytest tmp_path/tmp_path_factory fixture) and assign it to a variable like tmp_dir, then use tmp_dir in the object construction (where source_session_dir is set) and in the assertion (req.source_session_dir) so both occurrences reference the generated temp directory instead of a literal.acestep/core/generation/handler/retake_session.py (2)
73-75:⚠️ Potential issue | 🟠 Major | ⚡ Quick winReject non-positive
source_track_indexinstead of silently loading track 1.
max(1, int(track_index))turns0or negative inputs into track 1, which can retake the wrong source audio without surfacing an error to the caller.As per coding guidelines, "Handle errors explicitly."Suggested fix
- index = max(1, int(track_index)) + index = int(track_index) + if index < 1: + raise ValueError(f"source_track_index must be >= 1, got {track_index!r}")🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@acestep/core/generation/handler/retake_session.py` around lines 73 - 75, The code currently forces non-positive track_index values to 1 via index = max(1, int(track_index)), which silently loads the wrong source; instead validate and reject non-positive inputs: parse/convert track_index to int, if the value <= 0 raise a ValueError (or a domain-specific exception) with a clear message identifying the invalid source_track_index, and only then compute params_path and latents_path using that validated index; update the logic around track_index / index in retake_session.py accordingly so callers receive an explicit error rather than defaulting to track 1.
153-172: 🛠️ Refactor suggestion | 🟠 Major | ⚡ Quick winAdd docstrings to the new private helpers.
_read_json,_write_json,_first_value, and_first_textare new functions and still miss the repo's required Python docstrings.As per coding guidelines, "Docstrings are mandatory for all new or modified Python modules, classes, and functions."
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@acestep/core/generation/handler/retake_session.py` around lines 153 - 172, Add PEP257-style docstrings to each new helper: _read_json, _write_json, _first_value, and _first_text. For _read_json add a one-line summary plus parameters and return describing that it reads JSON from a Path and returns a dict; for _write_json document that it writes a dict to a Path and describe parameters and that it uses ensure_ascii=False, indent=2, default=str; for _first_value document that it returns the first non-empty/non-None/"N/A" value from the varargs and for _first_text document that it returns a string (empty string when no value). Keep each docstring short, use triple-quoted string format immediately below each def, and follow existing repo docstring style (one-line summary + short param/return notes).acestep/inference.py (1)
995-1000:⚠️ Potential issue | 🟠 Major | ⚡ Quick winMake the default session artifact directory unique per run.
When
save_session_artifacts=Trueandsession_output_diris omitted, every generation writes into the samesession_artifactsfolder. That can overwrite the source session a later"most natural"edit is supposed to reuse.Suggested fix
if params.save_session_artifacts: session_dir = params.session_output_dir if not session_dir: - session_dir = os.path.join(save_dir or os.getcwd(), "session_artifacts") + session_dir = tempfile.mkdtemp( + prefix="session_artifacts_", + dir=save_dir or os.getcwd(), + ) save_generation_session_artifacts(result=final_result, session_dir=session_dir) final_result.extra_outputs["session_output_dir"] = session_dir🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@acestep/inference.py` around lines 995 - 1000, The current logic creates a non-unique default session_dir "session_artifacts" when params.save_session_artifacts is true and params.session_output_dir is empty; change it to generate a unique directory per run (e.g., append a timestamp or UUID) when computing session_dir (use save_dir or os.getcwd() as base), then call save_generation_session_artifacts(result=final_result, session_dir=session_dir) and set final_result.extra_outputs["session_output_dir"]=session_dir so each invocation (referenced symbols: params.save_session_artifacts, params.session_output_dir, save_dir, save_generation_session_artifacts, final_result.extra_outputs) writes to a distinct folder.acestep/core/generation/handler/generate_music.py (1)
298-305:⚠️ Potential issue | 🟠 Major | ⚡ Quick winResolve
"auto"here and coerce every session-retake call ontotext2music.This mixin only recognizes explicit
"most natural"and only rewrites"cover". Direct callers that passrepaint_mode="auto"with asource_session_dir, or"most natural"withtask_type="repaint", still miss the session-retake/audio-code path and fall back to normal repaint conditioning.Suggested fix
-from acestep.core.generation.handler.retake_session import is_session_retake_mode +from acestep.core.generation.handler.retake_session import ( + is_session_retake_mode, + resolve_repaint_mode, +)- uses_session_retake = is_session_retake_mode(repaint_mode) + effective_repaint_mode = resolve_repaint_mode(repaint_mode, source_session_dir) + uses_session_retake = is_session_retake_mode(effective_repaint_mode) task_type, instruction = self._resolve_generate_music_task( task_type=task_type, audio_code_string=audio_code_string, instruction=instruction, ) - if uses_session_retake and task_type == "cover": + if uses_session_retake and task_type != "text2music": task_type = "text2music"# Then use `effective_repaint_mode` anywhere this method currently consumes # `repaint_mode` for retake-related behavior.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@acestep/core/generation/handler/generate_music.py` around lines 298 - 305, Compute an effective_repaint_mode (e.g. effective_repaint_mode = repaint_mode) and resolve "auto" into the session-retake mode when appropriate (for example if repaint_mode == "auto" and source_session_dir is present then set effective_repaint_mode = "most natural"); use is_session_retake_mode(effective_repaint_mode) to set uses_session_retake and pass effective_repaint_mode everywhere this function currently uses repaint_mode for retake behavior (notably before calling _resolve_generate_music_task and when deciding task type); then coerce any session-retake call into text2music by changing task_type to "text2music" when uses_session_retake is true or when original task_type == "repaint" and is_session_retake_mode(effective_repaint_mode) is true.
🧹 Nitpick comments (1)
acestep/ui/gradio/events/results/generation_progress.py (1)
456-459: ⚡ Quick winReplace
print()withlogger.warning.Per coding guidelines, use
loguru.loggerinstead ofprint()for error logging:except Exception as e: - print(f"[Auto Score] Failed to prepare tensor data for sample {sample_idx}: {e}") + logger.warning("[Auto Score] Failed to prepare tensor data for sample {}: {}", sample_idx, e) return None🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@acestep/ui/gradio/events/results/generation_progress.py` around lines 456 - 459, Replace the print() call in the exception handler with loguru.logger: change the print(f"[Auto Score] Failed to prepare tensor data for sample {sample_idx}: {e}") to logger.warning(...) (or logger.exception(...) if you want stacktrace) and include sample_idx and the exception details in the message; also ensure loguru.logger is imported/available (referenced as logger) in the module where load_sample_feature_data and the except block live so the handler logs via logger.warning and still returns None.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@acestep/api/http/release_task_request_builder_test.py`:
- Around line 142-147: The test uses hard-coded "/tmp/..." paths for keys like
"source_session_dir" and "session_output_dir" (and the other occurrences at
lines 160-164); replace these literals with a temporary directory from the test
fixture (e.g., use pytest's tmp_path or tempfile.TemporaryDirectory) and set the
dict values to str(tmp_path / "source-session") and str(tmp_path /
"out-session") (or equivalent) so the test uses ephemeral, CI-safe temp dirs;
update both the block containing "source_session_dir"/"session_output_dir" and
the similar block at 160-164 to use the same tmp_path-based values.
In `@acestep/core/generation/handler/generate_music.py`:
- Around line 468-472: The final retake splice currently hardcodes
crossfade_frames=25 in the call to splice_retake_latents, which ignores the
API/UI setting; change that call to pass the repaint_latent_crossfade_frames
variable (or the appropriate config value in scope) instead of 25, ensuring you
preserve any existing default behavior if the variable is None/undefined (e.g.,
fallback to current default). Update the call sites around splice_retake_latents
(the pred_latents retake splice) so the latent crossfade respects
repaint_latent_crossfade_frames.
In `@acestep/inference.py`:
- Around line 452-458: Do not mutate the caller's GenerationParams instance
(params); instead create a local copy (e.g., clone or construct a new
GenerationParams from params) and apply the session/retake overrides (task_type,
thinking, use_cot_metas/use_cot_caption/use_cot_language, audio_codes,
audio_code_string_to_use) to that copy; ensure subsequent serialization uses the
copied object's to_dict() so the original params remains unchanged for later
calls (refer to the GenerationParams type, the params variable, retake_inputs,
audio_codes, audio_code_string_to_use, and any code that calls
params.to_dict()).
---
Duplicate comments:
In `@acestep/api/http/release_task_models_test.py`:
- Around line 40-45: Replace the hardcoded "/tmp/source-session" literal with a
temp directory created at test runtime (e.g., use tempfile.mkdtemp() or the
pytest tmp_path/tmp_path_factory fixture) and assign it to a variable like
tmp_dir, then use tmp_dir in the object construction (where source_session_dir
is set) and in the assertion (req.source_session_dir) so both occurrences
reference the generated temp directory instead of a literal.
In `@acestep/core/generation/handler/generate_music.py`:
- Around line 298-305: Compute an effective_repaint_mode (e.g.
effective_repaint_mode = repaint_mode) and resolve "auto" into the
session-retake mode when appropriate (for example if repaint_mode == "auto" and
source_session_dir is present then set effective_repaint_mode = "most natural");
use is_session_retake_mode(effective_repaint_mode) to set uses_session_retake
and pass effective_repaint_mode everywhere this function currently uses
repaint_mode for retake behavior (notably before calling
_resolve_generate_music_task and when deciding task type); then coerce any
session-retake call into text2music by changing task_type to "text2music" when
uses_session_retake is true or when original task_type == "repaint" and
is_session_retake_mode(effective_repaint_mode) is true.
In `@acestep/core/generation/handler/retake_session.py`:
- Around line 73-75: The code currently forces non-positive track_index values
to 1 via index = max(1, int(track_index)), which silently loads the wrong
source; instead validate and reject non-positive inputs: parse/convert
track_index to int, if the value <= 0 raise a ValueError (or a domain-specific
exception) with a clear message identifying the invalid source_track_index, and
only then compute params_path and latents_path using that validated index;
update the logic around track_index / index in retake_session.py accordingly so
callers receive an explicit error rather than defaulting to track 1.
- Around line 153-172: Add PEP257-style docstrings to each new helper:
_read_json, _write_json, _first_value, and _first_text. For _read_json add a
one-line summary plus parameters and return describing that it reads JSON from a
Path and returns a dict; for _write_json document that it writes a dict to a
Path and describe parameters and that it uses ensure_ascii=False, indent=2,
default=str; for _first_value document that it returns the first
non-empty/non-None/"N/A" value from the varargs and for _first_text document
that it returns a string (empty string when no value). Keep each docstring
short, use triple-quoted string format immediately below each def, and follow
existing repo docstring style (one-line summary + short param/return notes).
In `@acestep/inference.py`:
- Around line 995-1000: The current logic creates a non-unique default
session_dir "session_artifacts" when params.save_session_artifacts is true and
params.session_output_dir is empty; change it to generate a unique directory per
run (e.g., append a timestamp or UUID) when computing session_dir (use save_dir
or os.getcwd() as base), then call
save_generation_session_artifacts(result=final_result, session_dir=session_dir)
and set final_result.extra_outputs["session_output_dir"]=session_dir so each
invocation (referenced symbols: params.save_session_artifacts,
params.session_output_dir, save_dir, save_generation_session_artifacts,
final_result.extra_outputs) writes to a distinct folder.
In `@acestep/ui/gradio/api/api_routes.py`:
- Line 508: The code forwards raw form value for repainting_regions into
GenerationParams (via the call that uses get_param("repainting_regions",
default=None)), which leaves JSON text as a string instead of a list[dict];
update the handler to parse repainting_regions before constructing
GenerationParams: read the value from get_param, if it's a non-empty str attempt
to json.loads it (import json) and validate it is a list (or list of dicts) and
fall back to None on parse/validation failure, then pass the parsed value to
GenerationParams instead of the raw string; reference the repainting_regions
parameter, the get_param call, and the GenerationParams construction so the
change is applied where the value is forwarded.
- Around line 505-507: The code is treating 0.0 as missing because of the "or
0.3" fallback; change the logic around get_param("source_latent_mix_ratio",
default=0.3) so you only apply the 0.3 default when the returned value is
actually None/empty, not when it is falsy (0.0). Locate the call where
source_latent_mix_ratio is set and replace the "float(get_param(... ) or 0.3)"
pattern with logic that captures the raw result (e.g., val = get_param(...,
default=0.3)) and then use float(val) unless val is None/empty, in which case
use 0.3, so explicit 0.0 is preserved.
In `@acestep/ui/gradio/events/results/batch_queue.py`:
- Around line 180-182: The code persists "source_session_dir",
"source_track_index", and "source_latent_mix_ratio" but
restore_batch_parameters() doesn't restore them; update restore_batch_parameters
(and any callers that expect its return) to read these three keys from the
stored batch dict and include them in its returned parameters structure (e.g.,
add source_session_dir, source_track_index, source_latent_mix_ratio to the
returned dict or dataclass), and ensure downstream code that uses
restore_batch_parameters (e.g., batch queue consumer/processor functions)
consumes these values so the original session context is fully restored.
---
Nitpick comments:
In `@acestep/ui/gradio/events/results/generation_progress.py`:
- Around line 456-459: Replace the print() call in the exception handler with
loguru.logger: change the print(f"[Auto Score] Failed to prepare tensor data for
sample {sample_idx}: {e}") to logger.warning(...) (or logger.exception(...) if
you want stacktrace) and include sample_idx and the exception details in the
message; also ensure loguru.logger is imported/available (referenced as logger)
in the module where load_sample_feature_data and the except block live so the
handler logs via logger.warning and still returns None.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: ec2b87ae-d2dc-43cd-b402-c516ece80214
📒 Files selected for processing (41)
acestep/api/http/release_task_models.pyacestep/api/http/release_task_models_test.pyacestep/api/http/release_task_request_builder.pyacestep/api/http/release_task_request_builder_test.pyacestep/api/job_generation_setup.pyacestep/core/generation/handler/generate_music.pyacestep/core/generation/handler/generate_music_execute.pyacestep/core/generation/handler/generate_music_payload.pyacestep/core/generation/handler/generate_music_test.pyacestep/core/generation/handler/resolve_repaint_config_test.pyacestep/core/generation/handler/retake_latents.pyacestep/core/generation/handler/retake_latents_test.pyacestep/core/generation/handler/retake_session.pyacestep/core/generation/handler/retake_session_save.pyacestep/core/generation/handler/retake_session_test.pyacestep/core/generation/handler/service_generate.pyacestep/core/generation/handler/service_generate_execute.pyacestep/core/generation/handler/service_generate_execute_test.pyacestep/inference.pyacestep/ui/gradio/api/api_routes.pyacestep/ui/gradio/events/results/audio_transfer.pyacestep/ui/gradio/events/results/audio_transfer_test.pyacestep/ui/gradio/events/results/batch_management_background.pyacestep/ui/gradio/events/results/batch_management_helpers.pyacestep/ui/gradio/events/results/batch_management_wrapper.pyacestep/ui/gradio/events/results/batch_queue.pyacestep/ui/gradio/events/results/batch_queue_test.pyacestep/ui/gradio/events/results/feature_cache.pyacestep/ui/gradio/events/results/feature_cache_test.pyacestep/ui/gradio/events/results/generation_progress.pyacestep/ui/gradio/events/results/generation_progress_test.pyacestep/ui/gradio/events/results/lrc_utils.pyacestep/ui/gradio/events/results/lrc_utils_test.pyacestep/ui/gradio/events/results/scoring.pyacestep/ui/gradio/events/results/scoring_test.pyacestep/ui/gradio/events/wiring/generation_batch_navigation_wiring.pyacestep/ui/gradio/events/wiring/generation_mode_wiring.pyacestep/ui/gradio/events/wiring/generation_mode_wiring_test.pyacestep/ui/gradio/events/wiring/generation_run_wiring.pyacestep/ui/gradio/events/wiring/results_aux_wiring.pyacestep/ui/gradio/interfaces/generation_tab_secondary_controls.py
✅ Files skipped from review due to trivial changes (2)
- acestep/ui/gradio/events/wiring/generation_batch_navigation_wiring.py
- acestep/core/generation/handler/retake_latents.py
🚧 Files skipped from review as they are similar to previous changes (8)
- acestep/ui/gradio/events/wiring/generation_run_wiring.py
- acestep/core/generation/handler/resolve_repaint_config_test.py
- acestep/ui/gradio/events/results/generation_progress_test.py
- acestep/ui/gradio/events/results/batch_management_wrapper.py
- acestep/ui/gradio/events/results/batch_management_helpers.py
- acestep/ui/gradio/events/wiring/generation_mode_wiring_test.py
- acestep/core/generation/handler/retake_latents_test.py
- acestep/ui/gradio/events/wiring/generation_mode_wiring.py
Comment on lines
+142
to
+147
| "source_session_dir": "/tmp/source-session", | ||
| "source_track_index": 2, | ||
| "source_latent_mix_ratio": 0.25, | ||
| "save_session_artifacts": True, | ||
| "session_output_dir": "/tmp/out-session", | ||
| } |
Contributor
There was a problem hiding this comment.
Avoid "/tmp/..." literals in tests (Ruff S108).
Line 142/146 and Line 160/164 use temp-dir literals that trigger security lint errors and can fail CI.
Suggested patch
- "source_session_dir": "/tmp/source-session",
+ "source_session_dir": "source-session-dir",
@@
- "session_output_dir": "/tmp/out-session",
+ "session_output_dir": "session-output-dir",
@@
- self.assertEqual("/tmp/source-session", request.source_session_dir)
+ self.assertEqual("source-session-dir", request.source_session_dir)
@@
- self.assertEqual("/tmp/out-session", request.session_output_dir)
+ self.assertEqual("session-output-dir", request.session_output_dir)Also applies to: 160-164
🧰 Tools
🪛 Ruff (0.15.12)
[error] 142-142: Probable insecure usage of temporary file or directory: "/tmp/source-session"
(S108)
[error] 146-146: Probable insecure usage of temporary file or directory: "/tmp/out-session"
(S108)
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@acestep/api/http/release_task_request_builder_test.py` around lines 142 -
147, The test uses hard-coded "/tmp/..." paths for keys like
"source_session_dir" and "session_output_dir" (and the other occurrences at
lines 160-164); replace these literals with a temporary directory from the test
fixture (e.g., use pytest's tmp_path or tempfile.TemporaryDirectory) and set the
dict values to str(tmp_path / "source-session") and str(tmp_path /
"out-session") (or equivalent) so the test uses ephemeral, CI-safe temp dirs;
update both the block containing "source_session_dir"/"session_output_dir" and
the similar block at 160-164 to use the same tmp_path-based values.
Contributor
Author
|
Follow-up fix after Gradio testing: Send To Repaint could lose the hidden source session and fall back to ordinary balanced repaint. Added per-audio session markers in the generated JSON sidecar and a JSON fallback resolver for repaint transfers, with tests. Remote Gradio on jieyue:7869 has been restarted from the updated branch. |
Contributor
There was a problem hiding this comment.
Actionable comments posted: 2
🧹 Nitpick comments (1)
acestep/ui/gradio/events/results/generation_progress.py (1)
459-463: ⚡ Quick winAvoid swallowing scoring-prep failures with
print()andexcept Exception.This path turns unexpected feature-cache regressions into a quiet
None, which makes failures harder to diagnose in tests and production. Please switch tologgerand narrow the exception handling to the concrete errors raised byload_sample_feature_data(...).As per coding guidelines: "Error handling: Avoid bare
except:clauses; catch specific exceptions." and "Logging: Usefrom loguru import logger... Avoidprint()in committed code except CLI output."🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@acestep/ui/gradio/events/results/generation_progress.py` around lines 459 - 463, Replace the bare except/print with structured logging and narrow exception handling: import logger from loguru, catch only the concrete exceptions that load_sample_feature_data can raise (e.g., KeyError, IndexError, FileNotFoundError or the specific custom exception(s) it documents) instead of Exception, and call logger.exception/failure with a descriptive message including sample_idx and the exception; for truly unexpected errors re-raise them rather than silently returning None, and only return None for the specific, documented recoverable error cases from load_sample_feature_data.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@acestep/ui/gradio/events/results/generation_progress.py`:
- Around line 489-565: This file exceeds the 200-LOC guideline and the
session/cache persistence helpers should be split into a small helper module:
create a new module (e.g., gradio_persistence) and move the functions
_persist_gradio_source_session, _persist_gradio_feature_cache,
_should_persist_gradio_source_session, _build_gradio_source_session_dir, and
_build_gradio_feature_cache_dir into it; preserve their signatures and move any
required imports/refs (time_module, uuid, os, DEFAULT_RESULTS_DIR,
save_generation_session_artifacts, persist_feature_cache, logger) into the new
module, update generation_progress to import those functions from the new
module, run the module linting/tests to ensure nothing breaks, and add a short
module-level docstring noting the split to satisfy the follow-up documentation
requirement.
In `@acestep/ui/gradio/events/results/source_session_transfer.py`:
- Around line 54-64: The code checks existence with
os.path.expanduser(json_path) but later opens and returns the unexpanded
json_path and session_output_dir, causing mismatches for ~ paths; fix by
expanding json_path and the params' "session_output_dir" before use—call
os.path.expanduser on json_path before open(...) and on
params.get("session_output_dir") before passing to _existing_session_dir (and
anywhere else the raw value is returned or opened, including similar logic
around lines 74-80) so all file operations and returned session_dir consistently
use the expanded path.
---
Nitpick comments:
In `@acestep/ui/gradio/events/results/generation_progress.py`:
- Around line 459-463: Replace the bare except/print with structured logging and
narrow exception handling: import logger from loguru, catch only the concrete
exceptions that load_sample_feature_data can raise (e.g., KeyError, IndexError,
FileNotFoundError or the specific custom exception(s) it documents) instead of
Exception, and call logger.exception/failure with a descriptive message
including sample_idx and the exception; for truly unexpected errors re-raise
them rather than silently returning None, and only return None for the specific,
documented recoverable error cases from load_sample_feature_data.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: dea554b7-d065-4bcb-ac4a-965f9a65f526
📒 Files selected for processing (5)
acestep/ui/gradio/events/results/audio_transfer.pyacestep/ui/gradio/events/results/audio_transfer_test.pyacestep/ui/gradio/events/results/generation_progress.pyacestep/ui/gradio/events/results/generation_progress_test.pyacestep/ui/gradio/events/results/source_session_transfer.py
Comment on lines
+54
to
+64
| json_path = os.path.splitext(audio_path)[0] + ".json" | ||
| if not os.path.exists(os.path.expanduser(json_path)): | ||
| return "", 1 | ||
| try: | ||
| with open(json_path, encoding="utf-8") as file_obj: | ||
| params = json.load(file_obj) | ||
| except (OSError, json.JSONDecodeError): | ||
| return "", 1 | ||
| if not isinstance(params, dict): | ||
| return "", 1 | ||
| session_dir = _existing_session_dir(params.get("session_output_dir")) |
Contributor
There was a problem hiding this comment.
Expand ~ consistently before reading or returning session paths.
os.path.exists(os.path.expanduser(json_path)) accepts ~/..., but the subsequent open(json_path) and returned session_dir still keep the unexpanded value. For home-relative sidecars, this fallback can silently drop the source session even after the path was already validated.
Proposed fix
- json_path = os.path.splitext(audio_path)[0] + ".json"
- if not os.path.exists(os.path.expanduser(json_path)):
+ json_path = os.path.expanduser(os.path.splitext(audio_path)[0] + ".json")
+ if not os.path.exists(json_path):
return "", 1
try:
with open(json_path, encoding="utf-8") as file_obj:
params = json.load(file_obj)
@@
- expanded = os.path.expanduser(session_dir)
- return session_dir if os.path.isdir(expanded) else ""
+ expanded = os.path.abspath(os.path.expanduser(session_dir))
+ return expanded if os.path.isdir(expanded) else ""Also applies to: 74-80
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@acestep/ui/gradio/events/results/source_session_transfer.py` around lines 54
- 64, The code checks existence with os.path.expanduser(json_path) but later
opens and returns the unexpanded json_path and session_output_dir, causing
mismatches for ~ paths; fix by expanding json_path and the params'
"session_output_dir" before use—call os.path.expanduser on json_path before
open(...) and on params.get("session_output_dir") before passing to
_existing_session_dir (and anywhere else the raw value is returned or opened,
including similar logic around lines 74-80) so all file operations and returned
session_dir consistently use the expanded path.
Contributor
Author
|
Follow-up on source eligibility: session persistence is now capability-based instead of task-name-based. Gradio saves a hidden most-natural source session only when the generated result has both per-track audio_codes and final pred_latents. This keeps uploaded/no-code sources on conservative/balanced/aggressive repaint, while allowing any ACE-generated result that actually carries codes + latents to be reused as a most-natural repaint source. Synced and restarted jieyue:7869. |
Contributor
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (2)
acestep/core/generation/handler/retake_session_test.py (1)
111-139: ⚡ Quick winAdd one negative test for latent/audio track count mismatch
Capability gating includes a branch that rejects when
pred_latents.shape[0] < len(audios), but this path is currently untested. Adding it would harden regression coverage for the most-natural source eligibility check.Suggested test addition
+ def test_can_save_session_artifacts_rejects_latent_track_count_mismatch(self): + """Retake source must have at least one latent tensor per output track.""" + result = SimpleNamespace( + success=True, + audios=[ + {"params": {"audio_codes": "<|audio_code_1|>"}}, + {"params": {"audio_codes": "<|audio_code_2|>"}}, + ], + extra_outputs={"pred_latents": torch.ones(1, 6, 3)}, + ) + + self.assertFalse(can_save_generation_session_artifacts(result))🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@acestep/core/generation/handler/retake_session_test.py` around lines 111 - 139, Add a negative unit test exercising the branch in can_save_generation_session_artifacts that rejects when pred_latents.shape[0] < len(audios): create a SimpleNamespace result with success=True, audios containing two entries (each with params including audio_codes), and extra_outputs with pred_latents of shape (1, ..., ...) (i.e., only one latent row), then assert can_save_generation_session_artifacts(result) is False to cover the track-count mismatch path.acestep/ui/gradio/events/results/generation_progress.py (1)
455-459: ⚡ Quick winNarrow this exception handler and switch to
logger.Catching
Exceptionhere will also swallow programmer errors in the new disk-backed feature path, andprint()bypasses the repo's logging pipeline. Catch only the expected load failures and log them throughlogger.warning(...)/logger.exception(...).Proposed fix
def _extract_sample_tensor(extra_outputs, sample_idx): """Slice per-sample tensor data from *extra_outputs* for scoring. Returns ``None`` when data is missing or incomplete. """ try: return load_sample_feature_data(extra_outputs, sample_idx) - except Exception as e: - print(f"[Auto Score] Failed to prepare tensor data for sample {sample_idx}: {e}") + except (KeyError, OSError, RuntimeError, ValueError) as exc: + logger.warning( + "[auto_score] Failed to prepare tensor data for sample {}: {}", + sample_idx, + exc, + ) return NoneAs per coding guidelines: "Error handling: Avoid bare
except:clauses; catch specific exceptions." and "Logging: Usefrom loguru import loggerandlogger.info(),logger.error(), etc. Avoidprint()in committed code except CLI output."🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@acestep/ui/gradio/events/results/generation_progress.py` around lines 455 - 459, Replace the broad except/print around the load_sample_feature_data(extra_outputs, sample_idx) call with targeted handling: import logger from loguru and catch only the expected load-related exceptions (e.g., FileNotFoundError, KeyError, OSError/IOError or a custom LoadError if one exists), log a descriptive warning via logger.warning(...) including sample_idx and the exception, then return None; do not catch Exception broadly—let other exceptions bubble up (or use logger.exception and re-raise) so programmer errors in the disk-backed feature path are not swallowed.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@acestep/core/generation/handler/retake_session_test.py`:
- Around line 30-32: The test uses hardcoded "/tmp/session" which triggers Ruff
S108; update the assertions in retake_session_test to use a neutral path literal
or a temp fixture instead: replace the "/tmp/session" occurrences passed to
resolve_repaint_mode with a non-OS-temp string like "session/path" or use a
pytest tmp_path/tmpdir fixture to supply an isolated path (ensure calls to
resolve_repaint_mode("auto", ...), resolve_repaint_mode("most natural", ...),
and resolve_repaint_mode("aggressive", ...) are updated accordingly).
---
Nitpick comments:
In `@acestep/core/generation/handler/retake_session_test.py`:
- Around line 111-139: Add a negative unit test exercising the branch in
can_save_generation_session_artifacts that rejects when pred_latents.shape[0] <
len(audios): create a SimpleNamespace result with success=True, audios
containing two entries (each with params including audio_codes), and
extra_outputs with pred_latents of shape (1, ..., ...) (i.e., only one latent
row), then assert can_save_generation_session_artifacts(result) is False to
cover the track-count mismatch path.
In `@acestep/ui/gradio/events/results/generation_progress.py`:
- Around line 455-459: Replace the broad except/print around the
load_sample_feature_data(extra_outputs, sample_idx) call with targeted handling:
import logger from loguru and catch only the expected load-related exceptions
(e.g., FileNotFoundError, KeyError, OSError/IOError or a custom LoadError if one
exists), log a descriptive warning via logger.warning(...) including sample_idx
and the exception, then return None; do not catch Exception broadly—let other
exceptions bubble up (or use logger.exception and re-raise) so programmer errors
in the disk-backed feature path are not swallowed.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: d984ea33-bdaf-4012-9eba-08d7347ff96a
📒 Files selected for processing (4)
acestep/core/generation/handler/retake_session_save.pyacestep/core/generation/handler/retake_session_test.pyacestep/ui/gradio/events/results/generation_progress.pyacestep/ui/gradio/events/results/generation_progress_test.py
🚧 Files skipped from review as they are similar to previous changes (1)
- acestep/core/generation/handler/retake_session_save.py
ChuxiJ
changed the title
ChuxiJ
changed the title
ChuxiJ
changed the title
Contributor
Author
|
Closing this draft direction. The repaint changes accumulated conflicting assumptions during investigation; we are restarting from a clean branch off latest main with a narrower goal: reuse well-defined generated-session artifacts for follow-up edits without changing the existing repaint algorithm or adding article-style retake modes. |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Improve generated-source repaint while keeping the existing
mainrepaint algorithm intact.The goal is narrow: when a user repaints audio that was already generated by ACE-Step, reuse the generated result's hidden session artifacts instead of treating that result as an arbitrary uploaded WAV and VAE-encoding the WAV again. Uploaded/external audio continues to use the existing VAE-encode repaint path.
This is not the article-style audio-code
retakeflow, and it does not add a publicmost naturalrepaint mode. Existing Retake variation controls (retake_seed/retake_variance) remain unchanged.Closes #1168.
Refs #1097.
What Changed
Send To Repaintcarries hiddensource_session_dir/source_track_index; users do not need to enter these fields.task_type="repaint"with a valid generated-source session, load the saved per-track source latents and pass them into the normal repaint conditioning entry point.mainwould otherwise get fromVAE(src_audio). After that, existing repaint behavior is unchanged:src_latentsare silenced normallyclean_src_latentscomes from the same target/source latent tensor normallyauto,conservative,balanced, andaggressive;autoresolves to balanced.Non-goals
retake/most naturalrepaint mode.retake_seed/retake_variancesemantics.Validation
Local focused checks:
Result: 67 passed, 1 skipped.
Result: passed.