Add a shortfin pipeline for flux #876
Conversation
| @@ -50,8 +52,12 @@ def clip_text_model_to_dataset(model: ClipTextModel) -> Dataset: | |||
| return Dataset(properties=model.config.to_properties(), root_theta=model.theta) | |||
|
|
|||
|
|
|||
| def export_clip_text_model_iree_parameters(model: ClipTextModel, output_path: PathLike): | |||
| def export_clip_text_model_iree_parameters(model: ClipTextModel, output_path: PathLike, dtype=None): | |||
There was a problem hiding this comment.
Could you put a type hint for the new dtype argument?
I think they are important even if the assumption is that everyone should know what dtype should be.
They also improve code navigation as parsers could do a better job. They can determine types of resulting expressions.
| if not self.hp.use_post_quant_conv: | ||
| sample = rearrange( | ||
| sample, | ||
| "b (h w) (c ph pw) -> b c (h ph) (w pw)", | ||
| h=math.ceil(1024 / 16), | ||
| w=math.ceil(1024 / 16), | ||
| ph=2, | ||
| pw=2, | ||
| ) |
There was a problem hiding this comment.
Is this removal intentional?
| # t5_dataset.root_theta = t5_dataset.root_theta.transform( | ||
| # functools.partial(set_float_dtype, dtype=dtype) | ||
| # ) |
There was a problem hiding this comment.
This should not be commented out.
|
|
||
| self._rng = torch.Generator(device="cpu") | ||
|
|
||
| def _get_noise( |
There was a problem hiding this comment.
This function is available already. We should have just one and reuse it.
| # clip_dataset.root_theta = clip_dataset.root_theta.transform( | ||
| # functools.partial(set_float_dtype, dtype=dtype) | ||
| # ) |
There was a problem hiding this comment.
This should not be commented out.
| bs, c, h, w = img.shape | ||
|
|
||
| # Prepare image and position IDs | ||
| img = rearrange(img, "b c (h ph) (w pw) -> b (h w) (c ph pw)", ph=2, pw=2) |
There was a problem hiding this comment.
I know we are not going to export the pipeline itself, but are einops functions like rearrange exportable to MLIR?
| atol: Optional[float] = None, | ||
| rtol: Optional[float] = None, | ||
| ): | ||
| """Compare pipeline outputs between different dtypes.""" |
There was a problem hiding this comment.
Could mention that this compares against the hugging face implementation.
Also the tests below could mention that in the name.
| } | ||
|
|
||
| ARTIFACT_VERSION = "12032024" | ||
| SDXL_BUCKET = ( |
There was a problem hiding this comment.
SDXL_BUCKET -> FLUX_BUCKET
| SDXL_BUCKET = ( | ||
| f"https://sharkpublic.blob.core.windows.net/sharkpublic/flux.1/{ARTIFACT_VERSION}/" | ||
| ) | ||
| SDXL_WEIGHTS_BUCKET = ( |
| sfnp.bfloat16: "bf16", | ||
| } | ||
|
|
||
| ARTIFACT_VERSION = "12032024" |
No description provided.