Functions — github.com/Tongyi-MAI/Z-Image

hub / github.com/Tongyi-MAI/Z-Image / functions

Functions102 in github.com/Tongyi-MAI/Z-Image

⨍ Functions 102◇ Types & classes 24

↓ 40 callersMethodget

(self, key, default=None)

src/zimage/scheduler.py:21

↓ 9 callersMethod__init__

( self, in_channels=3, out_channels=3, block_out_channels=(64,), layer

src/zimage/autoencoder.py:205

↓ 6 callersMethod__init__

(self, dim: int, eps: float = 1e-5)

src/zimage/transformer.py:57

↓ 4 callersFunctionformat_bytes

Format bytes to GB string. Args: size: Size in bytes Returns: Formatted string in GB

src/utils/helpers.py:14

↓ 3 callersFunction_native_attention_wrapper

( query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, attn_mask: Optional[torch.Tenso

src/utils/attention.py:410

↓ 3 callersMethodcreate_coordinate_grid

(size, start=None, device=None)

src/zimage/transformer.py:362

↓ 3 callersFunctionload_config

(config_path: str)

src/utils/loader.py:45

↓ 3 callersMethodset_timesteps

( self, num_inference_steps: Optional[int] = None, device: Union[str, torch.device] =

src/zimage/scheduler.py:62

↓ 2 callersFunction_normalize_attn_mask

Normalize an attention mask to shape [batch_size, seq_len_k] (bool).

src/utils/attention.py:119

↓ 2 callersFunction_prepare_for_flash_attn_varlen

( batch_size: int, seq_len_q: int, seq_len_kv: int, attn_mask: Optional[torch.Tensor] = None,

src/utils/attention.py:184

↓ 2 callersFunction_process_mask

(attn_mask: Optional[torch.Tensor], dtype: torch.dtype)

src/utils/attention.py:102

↓ 2 callersMethod_sigma_to_t

(self, sigma)

src/zimage/scheduler.py:146

↓ 2 callersFunctionapply_rotary_emb

(x_in: torch.Tensor, freqs_cis: torch.Tensor)

src/zimage/transformer.py:78

↓ 2 callersFunctiongenerate

( transformer, vae, text_encoder, tokenizer, scheduler, prompt: Union[str, List[str]],

src/zimage/pipeline.py:67

↓ 2 callersFunctionload_from_local_dir

Load all Z-Image components from local directory. Args: model_dir: Path to model directory device: Device to load models on

src/utils/loader.py:80

↓ 2 callersFunctionload_sharded_safetensors

Load sharded safetensors from a directory.

src/utils/loader.py:50

↓ 2 callersMethodprint_available_backends

(cls)

src/utils/attention.py:67

↓ 2 callersFunctionset_attention_backend

(backend: Union[str, AttentionBackend, None])

src/utils/attention.py:508

↓ 2 callersFunctionverify_file_integrity

Verify file integrity using a manifest. Args: base_dir: Base directory for relative file paths manifest: Dictionary of r

src/utils/helpers.py:94

↓ 1 callersFunction_flash_attention

( query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, attn_mask: Optional[torch.Tenso

src/utils/attention.py:197

↓ 1 callersFunction_flash_attention_3

( query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, attn_mask: Optional[torch.Tenso

src/utils/attention.py:277

↓ 1 callersFunction_flash_varlen_attention

( query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, attn_mask: Optional[torch.Tenso

src/utils/attention.py:223

↓ 1 callersFunction_flash_varlen_attention_3

( query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, attn_mask: Optional[torch.Tenso

src/utils/attention.py:309

↓ 1 callersMethod_init_step_index

(self, timestep)

src/zimage/scheduler.py:112

↓ 1 callersFunction_math_attention

(*args, **kwargs)

src/utils/attention.py:462

↓ 1 callersFunction_mps_flash_attention

MPS Flash Attention for Apple Silicon (M1/M2/M3/M4).

src/utils/attention.py:373

↓ 1 callersFunction_native_attention

(*args, **kwargs)

src/utils/attention.py:467

↓ 1 callersFunction_native_flash_attention

( query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, attn_mask: Optional[torch.Tenso

src/utils/attention.py:440

↓ 1 callersFunction_prepare_for_flash_attn_varlen_with_mask

( batch_size: int, seq_len_q: int, attn_mask: torch.Tensor, device: Optional[torch.device] = N

src/utils/attention.py:164

↓ 1 callersFunction_prepare_for_flash_attn_varlen_without_mask

( batch_size: int, seq_len_q: int, seq_len_kv: int, device: Optional[torch.device] = None, )

src/utils/attention.py:148

↓ 1 callersFunctioncalculate_shift

( image_seq_len, base_seq_len: int = BASE_IMAGE_SEQ_LEN, max_seq_len: int = MAX_IMAGE_SEQ_LEN,

src/zimage/pipeline.py:23

↓ 1 callersFunctioncompute_file_md5

Compute MD5 hash of a file.

src/utils/helpers.py:52

↓ 1 callersFunctioncompute_md5

Compute MD5 hash of a file.

src/tools/generate_manifest.py:15

↓ 1 callersMethoddecode

(self, z: torch.FloatTensor, return_dict: bool = True)

src/zimage/autoencoder.py:360

↓ 1 callersFunctiondispatch_attention

( query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, attn_mask: Optional[torch.Tenso

src/utils/attention.py:471

↓ 1 callersFunctionensure_model_weights

Ensure model weights exist and optionally verify integrity. Args: model_path: Path to model directory repo_id: HuggingFa

src/utils/helpers.py:135

↓ 1 callersMethodget

(self, key, default=None)

src/zimage/autoencoder.py:20

↓ 1 callersFunctionget_essential_files

Get list of essential model files.

src/tools/generate_manifest.py:24

↓ 1 callersMethodindex_for_timestep

(self, timestep, schedule_timesteps=None)

src/zimage/scheduler.py:104

↓ 1 callersFunctionis_flash_attn_3_available

()

src/utils/import_utils.py:10

↓ 1 callersFunctionis_flash_attn_available

()

src/utils/import_utils.py:6

↓ 1 callersFunctionis_torch_version

(operator: str, version: str)

src/utils/import_utils.py:14

↓ 1 callersFunctionload_manifest

Load manifest file. Returns dict mapping file paths to MD5 hashes (or None).

src/utils/helpers.py:61

↓ 1 callersFunctionmain

()

batch_inference.py:58

↓ 1 callersFunctionmain

()

inference.py:14

↓ 1 callersFunctionmain

()

src/tools/generate_manifest.py:51

↓ 1 callersMethodpatchify_and_embed

( self, all_image: List[torch.Tensor], all_cap_feats: List[torch.Tensor], patc

src/zimage/transformer.py:369

↓ 1 callersMethodprecompute_freqs_cis

(dim: List[int], end: List[int], theta: float = ROPE_THETA)

src/zimage/transformer.py:236

↓ 1 callersFunctionread_prompts

Read prompts from a text file (one per line, empty lines skipped).

batch_inference.py:14

↓ 1 callersFunctionretrieve_timesteps

( scheduler, num_inference_steps: Optional[int] = None, device: Optional[Union[str, torch.device]]

src/zimage/pipeline.py:36

↓ 1 callersFunctionselect_device

Choose the best available device without repeating detection logic.

batch_inference.py:38

↓ 1 callersFunctionslugify

Create a filesystem-safe slug from the prompt.

batch_inference.py:30

↓ 1 callersMethodstep

Predict the sample at the previous timestep.

src/zimage/scheduler.py:120

↓ 1 callersMethodtime_shift

(self, mu: float, sigma: float, t: torch.Tensor)

src/zimage/scheduler.py:149

↓ 1 callersMethodtimestep_embedding

(t, dim, max_period=MAX_PERIOD)

src/zimage/transformer.py:35

↓ 1 callersMethodunpatchify

(self, x: List[torch.Tensor], size: List[Tuple], patch_size, f_patch_size)

src/zimage/transformer.py:345

Method__call__

(self, ids: torch.Tensor)

src/zimage/transformer.py:247

Method__getattr__

(self, name)

src/zimage/scheduler.py:24

Method__getattr__

(self, name)

src/zimage/autoencoder.py:23

Method__init__

(self, **kwargs)

src/zimage/scheduler.py:18

Method__init__

( self, num_train_timesteps: int = 1000, shift: float = 1.0, use_dynamic_shift

src/zimage/scheduler.py:31

Method__init__

(self, out_size, mid_size=None, frequency_embedding_size=FREQUENCY_EMBEDDING_SIZE)

src/zimage/transformer.py:23

Method__init__

(self, dim: int, hidden_dim: int)

src/zimage/transformer.py:68

Method__init__

(self, dim: int, n_heads: int, n_kv_heads: int, qk_norm: bool = True, eps: float = 1e-5)

src/zimage/transformer.py:89

Method__init__

( self, layer_id: int, dim: int, n_heads: int, n_kv_heads: int,

src/zimage/transformer.py:144

Method__init__

(self, hidden_size, out_channels)

src/zimage/transformer.py:206

Method__init__

( self, theta: float = ROPE_THETA, axes_dims: List[int] = ROPE_AXES_DIMS, axes

src/zimage/transformer.py:223

Method__init__

( self, all_patch_size=(2,), all_f_patch_size=(1,), in_channels=16, di

src/zimage/transformer.py:267

Method__init__

(self, **kwargs)

src/zimage/autoencoder.py:17

Method__init__

(self, in_channels, out_channels=None, dropout=0.0, temb_channels=512, groups=32, eps=1e-6)

src/zimage/autoencoder.py:32

Method__init__

(self, in_channels, heads=1, dim_head=None, groups=32, eps=1e-6)

src/zimage/autoencoder.py:70

Method__init__

(self, channels, with_conv=True, out_channels=None, padding=1)

src/zimage/autoencoder.py:102

Method__init__

(self, channels, with_conv=True, out_channels=None)

src/zimage/autoencoder.py:117

Method__init__

(self, in_channels, out_channels, num_layers=1, resnet_eps=1e-6, resnet_groups=32, add_downsample=True)

src/zimage/autoencoder.py:132

Method__init__

(self, in_channels, out_channels, num_layers=1, resnet_eps=1e-6, resnet_groups=32, add_upsample=True)

src/zimage/autoencoder.py:161

Method__init__

(self, in_channels, resnet_eps=1e-6, resnet_groups=32, attention_head_dim=None)

src/zimage/autoencoder.py:186

Method__init__

( self, in_channels=3, out_channels=3, block_out_channels=(64,), layer

src/zimage/autoencoder.py:256

Method__init__

( self, in_channels: int = 3, out_channels: int = 3, down_block_types: Tuple[s

src/zimage/autoencoder.py:305

Function_check_device_cuda

(query: torch.Tensor, **kwargs)

src/utils/attention.py:87

Function_check_device_mps

(query: torch.Tensor, **kwargs)

src/utils/attention.py:97

Function_check_qkv_dtype_bf16_or_fp16

(query: torch.Tensor, **kwargs)

src/utils/attention.py:92

Functiondecorator

(func)

src/utils/attention.py:78

Methoddtype

(self)

src/zimage/autoencoder.py:357

Methodforward

(self, t)

src/zimage/transformer.py:47

Methodforward

(self, x: torch.Tensor)

src/zimage/transformer.py:62

Methodforward

(self, x)

src/zimage/transformer.py:74

Methodforward

( self, hidden_states: torch.Tensor, attention_mask: Optional[torch.Tensor] = None,

src/zimage/transformer.py:103

Methodforward

( self, x: torch.Tensor, attn_mask: torch.Tensor, freqs_cis: torch.Tensor,

src/zimage/transformer.py:171

Methodforward

(self, x, c)

src/zimage/transformer.py:215

Methodforward

( self, x: List[torch.Tensor], t, cap_feats: List[torch.Tensor], patch

src/zimage/transformer.py:474

Methodforward

(self, input_tensor, temb=None)

src/zimage/autoencoder.py:51

Methodforward

(self, hidden_states)

src/zimage/autoencoder.py:81

Methodforward

(self, hidden_states)

src/zimage/autoencoder.py:109

Methodforward

(self, hidden_states)

src/zimage/autoencoder.py:124

Methodforward

(self, hidden_states)

src/zimage/autoencoder.py:147

Methodforward

(self, hidden_states)

src/zimage/autoencoder.py:174

Methodforward

(self, hidden_states)

src/zimage/autoencoder.py:196

Methodforward

(self, x)

src/zimage/autoencoder.py:244

Methodforward

(self, x)

src/zimage/autoencoder.py:293

Functionprint_memory_stats

Print CUDA memory statistics. Args: stage: Description of current stage

src/utils/helpers.py:28

next →1–100 of 102, ranked by callers