Stack the enter and build decorators for the flux example (#845)

* Stack the enter and build decorators for the flux example * Format
modal-labs · Aug 17, 2024 · fbe6466 · fbe6466
1 parent 89a6c60
commit fbe6466
Showing 1 changed file with 2 additions and 4 deletions.
diff --git a/06_gpu_and_ml/stable_diffusion/flux.py b/06_gpu_and_ml/stable_diffusion/flux.py
@@ -47,14 +47,12 @@
     image=flux_image,
 )
 class Model:
-    @modal.build()
-    def build(self):
+    @modal.enter()
+    def enter(self):
         from huggingface_hub import snapshot_download
 
         snapshot_download(f"black-forest-labs/FLUX.1-{VARIANT}")
 
-    @modal.enter()
-    def enter(self):
         self.pipe = FluxPipeline.from_pretrained(
             f"black-forest-labs/FLUX.1-{VARIANT}", torch_dtype=torch.bfloat16
         )