Update custom_pipeline.py

custom_pipeline.py

@@ -19,7 +19,10 @@ BATCH_SIZE = 4  # Optimal batch size for A100
 
 @torch.jit.script
 def calculate_timestep_shift(image_seq_len: int) -> float:
-    """Optimized timestep shift calculation using TorchScript"""
+    BASE_SEQ_LEN = 256
+    MAX_SEQ_LEN = 4096
+    BASE_SHIFT = 0.5
+    MAX_SHIFT = 1.2
     m = (MAX_SHIFT - BASE_SHIFT) / (MAX_SEQ_LEN - BASE_SEQ_LEN)
     b = BASE_SHIFT - m * BASE_SEQ_LEN
     return image_seq_len * m + b