[PR Welcome] Request cache support for LongCat-Image

[DefTruth]

LongCat-Image is supported in diffusers: huggingface/diffusers#12828

[e1ijah1]

Hi, I'd like to work on this issue.

[DefTruth]

cool! PR is welcome!

[DefTruth]

@e1ijah1

• Step 1: Add block adapter for longcat-image at src/cache_dit/caching/block_adapters/__init__.py, for example:

cache-dit/src/cache_dit/caching/block_adapters/__init__.py

Lines 835 to 850 in 08253c7

	def ovis_image_adapter(pipe, **kwargs) -> BlockAdapter:
	try:
	from diffusers import OvisImageTransformer2DModel
	except ImportError:
	OvisImageTransformer2DModel = None # requires diffusers>=0.36.dev
	_relaxed_assert(pipe.transformer, OvisImageTransformer2DModel)
	return BlockAdapter(
	pipe=pipe,
	transformer=pipe.transformer,
	blocks=[
	pipe.transformer.transformer_blocks,
	pipe.transformer.single_transformer_blocks,
	],
	forward_pattern=[
	ForwardPattern.Pattern_1,

• Step 2: Register an example for longcat-image at examples/registers.py, for example:

cache-dit/examples/registers.py

Lines 715 to 764 in 08253c7

	@ExampleRegister.register("zimage", default="Tongyi-MAI/Z-Image-Turbo")
	def zimage_example(args: argparse.Namespace, **kwargs) -> Example:
	from diffusers import ZImagePipeline
	if args.cache:
	# Only warmup 4 steps (total 9 steps) for distilled models
	args.max_warmup_steps = min(4, args.max_warmup_steps)
	steps_computation_mask = (
	cache_dit.steps_mask(
	# slow, medium, fast, ultra.
	mask_policy=args.mask_policy,
	total_steps=9 if args.num_inference_steps is None else args.num_inference_steps,
	)
	if args.mask_policy is not None
	else (
	cache_dit.steps_mask(
	compute_bins=[5, 1, 1], # = 7 (compute steps)
	cache_bins=[1, 1], # = 2 (dynamic cache steps)
	)
	if args.steps_mask
	else None
	)
	)
	return Example(
	args=args,
	init_config=ExampleInitConfig(
	task_type=ExampleType.T2I, # Text to Image
	model_name_or_path=_path("Tongyi-MAI/Z-Image-Turbo"),
	pipeline_class=ZImagePipeline,
	bnb_4bit_components=["text_encoder"],
	extra_optimize_kwargs={
	"steps_computation_mask": steps_computation_mask,
	},
	),
	input_data=ExampleInputData(
	prompt=(
	"Young Chinese woman in red Hanfu, intricate embroidery. Impeccable makeup, "
	"red floral forehead pattern. Elaborate high bun, golden phoenix headdress, "
	"red flowers, beads. Holds round folding fan with lady, trees, bird. Neon "
	"lightning-bolt lamp (⚡️), bright yellow glow, above extended left palm. "
	"Soft-lit outdoor night background, silhouetted tiered pagoda (西安大雁塔), "
	"blurred colorful distant lights."
	),
	height=1024,
	width=1024,
	guidance_scale=0.0, # Guidance should be 0 for the Turbo models
	num_inference_steps=9,
	),
	)

[DefTruth]

Hi~ would you like also add cp/tp support for longcat-image? The arch of longcat-image is similar with FLUX.1.

• context parallelism

cache-dit/src/cache_dit/parallelism/transformers/native_diffusers/context_parallelism/cp_plan_flux.py

Line 42 in ec65b19

class FluxContextParallelismPlanner(ContextParallelismPlanner):

• tensor parallelism

cache-dit/src/cache_dit/parallelism/transformers/native_pytorch/tensor_parallelism/tp_plan_flux.py

Line 31 in ec65b19

class FluxTensorParallelismPlanner(TensorParallelismPlanner):

[e1ijah1]

Hi~ would you like also add cp/tp support for longcat-image? The arch of longcat-image is similar with FLUX.1.

• context parallelism

cache-dit/src/cache_dit/parallelism/transformers/native_diffusers/context_parallelism/cp_plan_flux.py

Line 42 in ec65b19

class FluxContextParallelismPlanner(ContextParallelismPlanner):

• tensor parallelism

cache-dit/src/cache_dit/parallelism/transformers/native_pytorch/tensor_parallelism/tp_plan_flux.py

Line 31 in ec65b19

class FluxTensorParallelismPlanner(TensorParallelismPlanner):

Sure, I'll try to add cp/tp support for these two models.

[DefTruth]

Hi~ would you like also add cp/tp support for longcat-image? The arch of longcat-image is similar with FLUX.1.

• context parallelism

cache-dit/src/cache_dit/parallelism/transformers/native_diffusers/context_parallelism/cp_plan_flux.py
Line 42 in ec65b19
class FluxContextParallelismPlanner(ContextParallelismPlanner):

• tensor parallelism

cache-dit/src/cache_dit/parallelism/transformers/native_pytorch/tensor_parallelism/tp_plan_flux.py
Line 31 in ec65b19
class FluxTensorParallelismPlanner(TensorParallelismPlanner):

Sure, I'll try to add cp/tp support for these two models.

cool~

[DefTruth]

done~