Update app.py

app.py

@@ -18,27 +18,10 @@ from diffusers.utils import export_to_gif, load_image
 from diffusers.models.modeling_utils import ModelMixin
 from huggingface_hub import hf_hub_download
 from safetensors.torch import load_file, save_file
-from diffusers import DiffusionPipeline, AnimateDiffPipeline, MotionAdapter, EulerDiscreteScheduler, StableDiffusionXLPipeline, UNet2DConditionModel
+from diffusers import DiffusionPipeline, AnimateDiffPipeline, MotionAdapter, EulerDiscreteScheduler, StableDiffusionXLPipeline, UNet2DConditionModel, AutoencoderKL, UNet3DConditionModel
 import jax
 import jax.numpy as jnp
 
-class Model(nn.Module):
-    def __init__(self):
-        super().__init__()
-        self.register_buffer('buffer', torch.ones(1, 1))
-
-    def forward(self, x):
-        new_tensor = torch.randn(1, 1)
-        self.buffer = torch.cat([self.buffer, new_tensor], dim=0)
-        return self
-
-def dict2model(dict):
-    model = Model()
-    m = model(dict)
-    mix = ModelMixin()
-    mix(m)
-    return mix
-
 def forest_schnell():
     PIPE = DiffusionPipeline.from_pretrained("black-forest-labs/FLUX.1-schnell", torch_dtype=torch.bfloat16, token=os.getenv("hf_token")).to("cuda")
     return PIPE
@@ -223,23 +206,14 @@ def main():
 
     adapter = MotionAdapter.from_pretrained("guoyww/animatediff-motion-adapter-v1-5-2")
 
-    repo="stabilityai/sd-vae-ft-mse-original"
-    ckpt="vae-ft-mse-840000-ema-pruned.safetensors"
-    vae = dict2model(load_file(hf_hub_download(repo, ckpt), device=device))
-
-    repo="ByteDance/SDXL-Lightning"
-    ckpt=f"sdxl_lightning_{step}step_unet.safetensors"
-    unet = dict2model(load_file(hf_hub_download(repo, ckpt), device=device))
-    
-    #repo = "SG161222/Realistic_Vision_V6.0_B1_noVAE"
+    vae = AutoencoderKL.from_single_file("stabilityai/sd-vae-ft-mse-original/vae-ft-mse-840000-ema-pruned.safetensors")
+    unet = UNet3DConditionModel.to(device, dtype).load_state_dict(load_file(hf_hub_download("ByteDance/SDXL-Lightning", f"sdxl_lightning_{step}step_unet.safetensors"), device=device))
     
-    #repo = "emilianJR/epiCRealism"
-    #ckpt = "unet/diffusion_pytorch_model.safetensors"
-    #unet = load_file(hf_hub_download(repo, ckpt), device=device)
-
     repo = "ByteDance/AnimateDiff-Lightning"
     ckpt = f"animatediff_lightning_{step}step_diffusers.safetensors"
     base = "black-forest-labs/FLUX.1-schnell"
+    #base = "SG161222/Realistic_Vision_V6.0_B1_noVAE"
+    #base = "emilianJR/epiCRealism"
     
     pipe = AnimateDiffPipeline.from_pretrained(base, vae=vae, motion_adapter=adapter, feature_extractor=None, image_encoder=None, unet=unet, torch_dtype=dtype, token=os.getenv("hf_token")).to(device)
     pipe.scheduler = EulerDiscreteScheduler.from_config(pipe.scheduler.config, timestep_spacing="trailing", beta_schedule="linear")