Update app.py

app.py

@@ -13,6 +13,10 @@ from image_datasets.canny_dataset import canny_processor, c_crop
 from src.flux.sampling import denoise_controlnet, get_noise, get_schedule, prepare, unpack
 from src.flux.util import load_ae, load_clip, load_t5, load_flow_model, load_controlnet, load_safetensors
 
+# Define the flux_time_shift function
+def flux_time_shift(shift, base, timestep):
+    return base * (timestep ** shift)
+
 # Define the ModelSamplingFlux class
 class ModelSamplingFlux(torch.nn.Module):
     def __init__(self, model_config=None):