add triton kernels

ai-compiler-study · Jan 1, 2025 · 0c68d32 · 0c68d32
1 parent 140cbfb
commit 0c68d32
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/run_cli.sh b/run_cli.sh
@@ -0,0 +1 @@
+python -m flux --name flux-schnell --height 1024 --width 1024 --loop
diff --git a/src/flux/model.py b/src/flux/model.py
@@ -4,15 +4,14 @@
 from torch import Tensor, nn
 
 from flux.modules.layers import (
-    DoubleStreamBlock,
     EmbedND,
     LastLayer,
     MLPEmbedder,
-    SingleStreamBlock,
     timestep_embedding,
 )
 from flux.modules.lora import LinearLora, replace_linear_with_lora
 
+from triton_kernels import SingleStreamBlock, DoubleStreamBlock
 
 @dataclass
 class FluxParams:
Original file line number	Diff line number	Diff line change
		@@ -0,0 +1 @@
		python -m flux --name flux-schnell --height 1024 --width 1024 --loop