Upload 2 files

handler.py

@@ -24,17 +24,17 @@ class EndpointHandler:
         ).to("cuda")
         self.pipe.enable_vae_slicing()
         self.pipe.enable_vae_tiling()
-        #self.pipe.transformer.fuse_qkv_projections()
-        #self.pipe.vae.fuse_qkv_projections()
-        #self.pipe.transformer.to(memory_format=torch.channels_last)
-        #self.pipe.vae.to(memory_format=torch.channels_last)
+        self.pipe.transformer.fuse_qkv_projections()
+        self.pipe.vae.fuse_qkv_projections()
+        self.pipe.transformer.to(memory_format=torch.channels_last)
+        self.pipe.vae.to(memory_format=torch.channels_last)
         apply_cache_on_pipe(self.pipe, residual_diff_threshold=0.12)
-        #self.pipe.transformer = torch.compile(
-        #    self.pipe.transformer, mode="max-autotune-no-cudagraphs",
-        #)
-        #self.pipe.vae = torch.compile(
-        #    self.pipe.vae, mode="max-autotune-no-cudagraphs",
-        #)
+        self.pipe.transformer = torch.compile(
+            self.pipe.transformer, mode="max-autotune-no-cudagraphs",
+        )
+        self.pipe.vae = torch.compile(
+            self.pipe.vae, mode="max-autotune-no-cudagraphs",
+        )
         self.pipe.transformer = autoquant(self.pipe.transformer, error_on_unseen=False)
         self.pipe.vae = autoquant(self.pipe.vae, error_on_unseen=False)
         

requirements.txt

@@ -3,8 +3,8 @@ torch==2.6.0+cu126
 torchvision
 torchaudio
 huggingface_hub
-torchao>=0.9.0
-diffusers>=0.32.2
+torchao==0.9.0
+diffusers==0.32.2
 peft
 transformers<=4.48.3
 numpy<2
@@ -13,4 +13,4 @@ Pillow
 sentencepiece
 protobuf
 triton
-para-attn
+para-attn==0.3.23