FilesExpand file tree

main

/

convert_echo_to_diffusers.py

Copy path

Blame

More file actions

Blame

More file actions

Latest commit

 

History

History

History

179 lines (154 loc) · 6.79 KB

main

/

convert_echo_to_diffusers.py

Copy path

Top

File metadata and controls

• Code
• Blame

179 lines (154 loc) · 6.79 KB

Raw

Copy raw file

Download raw file

Open symbols panel

Edit and raw actions

1

2

3

4

5

6

7

8

9

10

11

12

13

14

15

16

17

18

19

20

21

22

23

24

25

26

27

28

29

30

31

32

33

34

35

36

37

38

39

40

41

42

43

44

45

46

47

48

49

50

51

52

53

54

55

56

57

58

59

60

61

62

63

64

65

66

67

68

69

70

71

72

73

74

75

76

77

78

79

80

81

82

83

84

85

86

87

88

89

90

91

92

93

94

95

96

97

98

99

100

101

102

103

104

105

106

107

108

109

110

111

112

113

114

115

116

117

118

119

120

121

122

123

124

125

126

127

128

129

130

131

132

133

134

135

136

137

138

139

140

141

142

143

144

145

146

147

148

149

150

151

152

153

154

155

156

157

158

159

160

161

162

163

164

165

166

167

168

169

170

171

172

173

174

175

176

177

178

179

#!/usr/bin/env python3

"""Convert the Echo BF16 release checkpoint into a Diffusers Modular Pipeline repository."""

import argparse

import gc

import json

from pathlib import Path

import torch

from convert_ltx2_to_diffusers import (

convert_ltx2_audio_vae,

convert_ltx2_connectors,

convert_ltx2_transformer,

convert_ltx2_video_vae,

convert_ltx2_vocoder,

get_model_state_dict_from_combined_ckpt,

)

from safetensors.torch import load_file

from diffusers import __version__

def resolve_checkpoint(path: Path) -> Path:

path = path.expanduser().resolve()

if path.is_file():

return path

manifest_path = path / "checkpoint.json"

if not manifest_path.is_file():

raise FileNotFoundError(f"Checkpoint manifest not found: {manifest_path}")

manifest = json.loads(manifest_path.read_text(encoding="utf-8"))

precision = str(manifest.get("precision", "")).lower()

if precision not in {"bf16", "bfloat16"}:

raise ValueError(

"Echo conversion supports the BF16 release only. The FP8 and FP4 releases use custom packed "

f"kernels and cannot be represented by this converter, got precision={precision!r}."

)

model_filename = manifest.get("files", {}).get("model")

if not model_filename:

raise ValueError(f"Checkpoint manifest has no `files.model`: {manifest_path}")

return path / model_filename

def save_component(model, output_path: Path, name: str, max_shard_size: str) -> None:

model.to(torch.bfloat16).save_pretrained(

output_path / name,

safe_serialization=True,

max_shard_size=max_shard_size,

)

def component_entry(library: str, class_name: str, repo: str, subfolder: str | None):

return [

None,

None,

{

"type_hint": [library, class_name],

"pretrained_model_name_or_path": repo,

"subfolder": subfolder,

"variant": None,

"revision": None,

},

]

def write_modular_index(output_path: Path, base_model: str, repo_id: str | None) -> None:

echo_repo = repo_id or str(output_path)

index = {

"_class_name": "EchoModularPipeline",

"_diffusers_version": __version__,

"_blocks_class_name": "EchoBlocks",

"text_encoder": component_entry("transformers", "Gemma3ForConditionalGeneration", base_model, None),

"tokenizer": component_entry("transformers", "GemmaTokenizerFast", base_model, None),

"connectors": component_entry("ltx2", "LTX2TextConnectors", echo_repo, "connectors"),

"vae": component_entry("diffusers", "AutoencoderKLLTX2Video", echo_repo, "vae"),

"audio_vae": component_entry("diffusers", "AutoencoderKLLTX2Audio", echo_repo, "audio_vae"),

"transformer": component_entry("diffusers", "LTX2VideoTransformer3DModel", echo_repo, "transformer"),

"vocoder": component_entry("ltx2", "LTX2VocoderWithBWE", echo_repo, "vocoder"),

}

(output_path / "modular_model_index.json").write_text(json.dumps(index, indent=2) + "\n", encoding="utf-8")

def convert(

checkpoint: Path,

output_path: Path,

base_model: str,

repo_id: str | None,

max_shard_size: str,

) -> None:

source = resolve_checkpoint(checkpoint)

output_path = output_path.expanduser().resolve()

output_path.mkdir(parents=True, exist_ok=True)

print(f"Loading Echo checkpoint from {source}", flush=True)

combined = load_file(str(source), device="cpu")

component_states = {

"dit": get_model_state_dict_from_combined_ckpt(combined, "model.diffusion_model"),

"vae": get_model_state_dict_from_combined_ckpt(combined, "vae"),

"audio_vae": get_model_state_dict_from_combined_ckpt(combined, "audio_vae"),

"vocoder": get_model_state_dict_from_combined_ckpt(combined, "vocoder"),

}

del combined

missing = [name for name, state in component_states.items() if not state]

if missing:

raise ValueError(f"Echo checkpoint is missing components: {missing}")

tensor_counts = {}

dit_state = component_states.pop("dit")

transformer = convert_ltx2_transformer(dict(dit_state), version="2.3")

tensor_counts["transformer"] = len(transformer.state_dict())

save_component(transformer, output_path, "transformer", max_shard_size)

del transformer

gc.collect()

connectors = convert_ltx2_connectors(dict(dit_state), version="2.3")

tensor_counts["connectors"] = len(connectors.state_dict())

save_component(connectors, output_path, "connectors", max_shard_size)

del connectors, dit_state

gc.collect()

vae = convert_ltx2_video_vae(component_states.pop("vae"), version="2.3", timestep_conditioning=False)

tensor_counts["vae"] = len(vae.state_dict())

save_component(vae, output_path, "vae", max_shard_size)

del vae

gc.collect()

audio_vae = convert_ltx2_audio_vae(component_states.pop("audio_vae"), version="2.3")

tensor_counts["audio_vae"] = len(audio_vae.state_dict())

save_component(audio_vae, output_path, "audio_vae", max_shard_size)

del audio_vae

gc.collect()

vocoder = convert_ltx2_vocoder(component_states.pop("vocoder"), version="2.3")

tensor_counts["vocoder"] = len(vocoder.state_dict())

save_component(vocoder, output_path, "vocoder", max_shard_size)

del vocoder

gc.collect()

write_modular_index(output_path, base_model=base_model, repo_id=repo_id)

report = {

"schema": "echo.diffusers.conversion.v1",

"source": str(source),

"precision": "bfloat16",

"tensor_counts": tensor_counts,

"base_model": base_model,

"model_repo": repo_id,

"vocoder_output_sample_rate": 48000,

}

(output_path / "conversion_report.json").write_text(json.dumps(report, indent=2) + "\n", encoding="utf-8")

print(f"Echo conversion complete: {output_path}", flush=True)

def parse_args() -> argparse.Namespace:

parser = argparse.ArgumentParser(description=__doc__)

parser.add_argument("--checkpoint", type=Path, required=True, help="Echo BF16 checkpoint or release folder.")

parser.add_argument("--output-path", type=Path, required=True, help="Destination Diffusers model directory.")

parser.add_argument(

"--base-model",

default="google/gemma-3-12b-it",

help="Gemma model repo/path containing the text encoder and tokenizer at its root.",

)

parser.add_argument(

"--repo-id",

default=None,

help="Future Hub repo id for the converted Echo components. Defaults to the local output path.",

)

parser.add_argument("--max-shard-size", default="5GB")

return parser.parse_args()

if __name__ == "__main__":

args = parse_args()

convert(

checkpoint=args.checkpoint,

output_path=args.output_path,

base_model=args.base_model,

repo_id=args.repo_id,

max_shard_size=args.max_shard_size,

)