Support dreamzero model & Add Numerical Alignment Workflow for skill

[helloyongyang]

No description provided.

[gemini-code-assist]

Code Review

This pull request adds native support for the DreamZero model, introducing its model structure, inference pipelines, weights loading, a flow scheduler, and a runner for DROID image-to-video-action (I2VA) tasks. The review feedback identifies critical runtime issues, including the use of the non-existent torch.linalg.solve_ex function in the scheduler, a key-filtering bug in the checkpoint loading logic that incorrectly discards valid weights, and a potential crash in the runner when processing state directories that lack the expected files.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

[gemini-code-assist]

torch.linalg.solve_ex is not a valid PyTorch function and will raise an AttributeError at runtime. Use the standard torch.linalg.solve function instead, which directly returns the solution tensor.

Suggested change

	rhos_p = torch.linalg.solve_ex(matrix_r[:-1, :-1], vector_b[:-1])[0].to(x.dtype)
	rhos_p = torch.linalg.solve(matrix_r[:-1, :-1], vector_b[:-1]).to(x.dtype)

[gemini-code-assist]

torch.linalg.solve_ex is not a valid PyTorch function and will raise an AttributeError at runtime. Use the standard torch.linalg.solve function instead, which directly returns the solution tensor.

Suggested change

	rhos_c = torch.linalg.solve_ex(matrix_r, vector_b)[0].to(x.dtype)
	rhos_c = torch.linalg.solve(matrix_r, vector_b).to(x.dtype)

[gemini-code-assist]

The _keep_dreamzero_key method checks the original key before stripping its prefix. This causes all keys starting with "model.diffusion_model." or "diffusion_model." to be filtered out and not loaded, as they do not match the allowed prefixes in the tuple. We should strip the prefix first before checking if the key should be kept.

Suggested change

	def _keep_dreamzero_key(key):
	if key.startswith("action_head.model."):
	return True
	if key.startswith(("blocks.", "patch_embedding.", "text_embedding.", "time_embedding.", "time_projection.", "img_emb.", "head.", "action_encoder.", "state_encoder.", "action_decoder.")):
	return True
	return False
	def _keep_dreamzero_key(key):
	for prefix in ("action_head.model.", "model.diffusion_model.", "diffusion_model."):
	if key.startswith(prefix):
	key = key[len(prefix) :]
	if key.startswith(("blocks.", "patch_embedding.", "text_embedding.", "time_embedding.", "time_projection.", "img_emb.", "head.", "action_encoder.", "state_encoder.", "action_decoder.")):
	return True
	return False

[gemini-code-assist]

If state_path is a directory but does not contain state.json or state.npy, the loop completes without finding a file, and state_path remains a directory. This will cause np.loadtxt to crash with an IsADirectoryError or PermissionError. We should raise a clear FileNotFoundError if none of the expected files are found in the directory.

Suggested change

	if os.path.isdir(state_path):
	for name in ("state.json", "state.npy"):
	candidate = os.path.join(state_path, name)
	if os.path.exists(candidate):
	state_path = candidate
	break
	if os.path.isdir(state_path):
	for name in ("state.json", "state.npy"):
	candidate = os.path.join(state_path, name)
	if os.path.exists(candidate):
	state_path = candidate
	break
	else:
	raise FileNotFoundError(f"Could not find state.json or state.npy in directory: {state_path}")