add triton kernels

Author: sjjeong94

run_cli.sh

@@ -0,0 +1 @@
+python -m flux --name flux-schnell --height 1024 --width 1024 --loop

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: