Update app.py

app.py

@@ -206,23 +206,22 @@ def main():
 
     progress=gr.Progress()
 
-    base = "black-forest-labs/FLUX.1-dev"
-    #base="SG161222/Realistic_Vision_V6.0_B1_noVAE"
-    #vae="stabilityai/sd-vae-ft-mse-original"
-    #repo = "ByteDance/SDXL-Lightning"
-    #ckpt = f"sdxl_lightning_{step}step_unet.safetensors"
-
-    #unet = UNet2DConditionModel.from_config(base, subfolder="unet").to(device)
-    #unet.load_state_dict(load_file(hf_hub_download(repo, ckpt), device=device), strict=False)
-
-    #base = "emilianJR/epiCRealism"
+    repo="stabilityai/sd-vae-ft-mse-original"
+    ckpt="vae-ft-mse-840000-ema-pruned.safetensors"
+    vae = load_file(hf_hub_download(repo, ckpt), device=device)
+    
     repo = "ByteDance/AnimateDiff-Lightning"
     ckpt = f"animatediff_lightning_{step}step_diffusers.safetensors"
-
+    
+    #base = "SG161222/Realistic_Vision_V6.0_B1_noVAE"
+    #base = "emilianJR/epiCRealism"
+    base = "black-forest-labs/FLUX.1-schnell"
+    
     pipe = AnimateDiffPipeline.from_pretrained(base, torch_dtype=dtype, token=os.getenv("hf_token")).to(device)
     pipe.scheduler = EulerDiscreteScheduler.from_config(pipe.scheduler.config, timestep_spacing="trailing", beta_schedule="linear")
     pipe.unet.load_state_dict(load_file(hf_hub_download(repo, ckpt), device=device), strict=False)
-
+    pipe.enable_free_init(method="butterworth", use_fast_sampling=False)
+    
     mp.set_start_method("spawn", force=True)
 
     with gr.Blocks(theme=gr.themes.Soft(),css=css,js=js) as demo: