Fix up everything. BF16 still failing though

nod-ai · Jan 31, 2025 · 7fe1c1d · 7fe1c1d
1 parent f6e4a81
commit 7fe1c1d
Show file tree

Hide file tree

Showing 9 changed files with 11 additions and 918 deletions.
diff --git a/sharktank/sharktank/layers/configs/llm_configs.py b/sharktank/sharktank/layers/configs/llm_configs.py
@@ -290,7 +290,7 @@ class ClipTextConfig:
     output_hidden_states: bool = False
     use_return_dict: bool = True
     dtype: torch.dtype = torch.float32
-
+    
     @staticmethod
     def from_hugging_face_clip_text_model_config(
         config: "transformers.CLIPTextConfig", # type: ignore

diff --git a/sharktank/sharktank/models/clip/export.py b/sharktank/sharktank/models/clip/export.py
@@ -55,7 +55,7 @@ def clip_text_model_to_dataset(model: ClipTextModel) -> Dataset:
 def export_clip_text_model_iree_parameters(model: ClipTextModel, output_path: PathLike, dtype=None):
     dataset = clip_text_model_to_dataset(model)
     if dtype:
-        dataset.root_theta = tdataset.root_theta.transform(
+        dataset.root_theta = dataset.root_theta.transform(
             functools.partial(set_float_dtype, dtype=dtype)
         )
     dataset.save(output_path)

diff --git a/sharktank/sharktank/pipelines/flux/__init__.py b/sharktank/sharktank/pipelines/flux/__init__.py
@@ -1,10 +1,7 @@
 """Flux text-to-image generation pipeline."""
 
 from .flux_pipeline import FluxPipeline
-from .export import export_flux_pipeline_mlir , export_flux_pipeline_iree_parameters
 
 __all__ = [
     "FluxPipeline",
-    "export_flux_pipeline_mlir",
-    "export_flux_pipeline_iree_parameters",
-]
+]
diff --git a/sharktank/sharktank/pipelines/flux/export.py b/sharktank/sharktank/pipelines/flux/export.py
diff --git a/sharktank/sharktank/torch_exports/flux/README.md b/sharktank/sharktank/torch_exports/flux/README.md