Update app.py

app.py

@@ -22,7 +22,7 @@ formatter = logging.Formatter('\n >>> [%(levelname)s] %(asctime)s %(name)s: %(me
 handler2.setFormatter(formatter)
 root.addHandler(handler2)
 
-def cmd(cmd, assert_success=False, capture_output=False, env=None, dry_run=False):
+def run(cmd, assert_success=False, capture_output=False, env=None, dry_run=False):
     if dry_run:
         print(f"--> {cmd}")
         result = 1
@@ -36,28 +36,34 @@ def cmd(cmd, assert_success=False, capture_output=False, env=None, dry_run=False
 
     return result
 
-cmd("apt install python3-mpi4py")
-cmd("pip install -r req.txt")
+run("apt install python3-mpi4py")
+run"pip install -r req.txt")
 
-import spaces
-import torch
-import gradio as gr
-import numpy as np
-from lxml.html import fromstring
-#from transformers import pipeline
-from torch import multiprocessing as mp, nn
-#from torch.multiprocessing import Pool
-#from pathos.multiprocessing import ProcessPool as Pool
-#from pathos.threading import ThreadPool as Pool
-#from diffusers.pipelines.flux import FluxPipeline
-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, DDIMScheduler, StableDiffusionXLPipeline, UNet2DConditionModel, AutoencoderKL, UNet3DConditionModel
-#import jax
-#import jax.numpy as jnp
-from pyina.launchers import TorqueMpiPool as Pool
+def deps():
+  try:
+    import spaces
+    import torch
+    import gradio as gr
+    import numpy as np
+    from lxml.html import fromstring
+    #from transformers import pipeline
+    from torch import multiprocessing as mp, nn
+    #from torch.multiprocessing import Pool
+    #from pathos.multiprocessing import ProcessPool as Pool
+    #from pathos.threading import ThreadPool as Pool
+    #from diffusers.pipelines.flux import FluxPipeline
+    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, DDIMScheduler, StableDiffusionXLPipeline, UNet2DConditionModel, AutoencoderKL, UNet3DConditionModel
+    #import jax
+    #import jax.numpy as jnp
+    from pyina.launchers import TorqueMpiPool as Pool
+  except:
+    pass
+
+deps()
 
 last_motion=None
 dtype = torch.float16
@@ -67,7 +73,7 @@ device = "cuda"
 #ckpt = f"animatediff_lightning_{step}step_diffusers.safetensors"
 base = "emilianJR/epiCRealism"
 #base = "SG161222/Realistic_Vision_V6.0_B1_noVAE"
-#vae = AutoencoderKL.from_pretrained("stabilityai/sd-vae-ft-mse").to(device, dtype=dtype)
+vae = AutoencoderKL.from_pretrained("stabilityai/sd-vae-ft-mse").to(device, dtype=dtype)
 #unet = UNet2DConditionModel.from_config("emilianJR/epiCRealism",subfolder="unet").to(device, dtype).load_state_dict(load_file(hf_hub_download("emilianJR/epiCRealism", "unet/diffusion_pytorch_model.safetensors"), device=device), strict=False)
 adapter = MotionAdapter.from_pretrained("guoyww/animatediff-motion-adapter-v1-5-3", torch_dtype=dtype, device=device)
 
@@ -76,7 +82,7 @@ fps=10
 time=1
 width=384
 height=768
-step = 25
+step=40
 accu=10
 
 css="""
@@ -223,7 +229,7 @@ def infer(pm):
         export_to_gif(out.frames[0],name,fps=fps)
         return name
 
-def run(i,m,p1,p2,*result):
+def main(i,m,p1,p2,*result):
     
         p1_en = translate(p1,"english")
         p2_en = translate(p2,"english")
@@ -235,7 +241,7 @@ def run(i,m,p1,p2,*result):
         with Pool(f'{ ln }:ppn=2', queue='productionQ', timelimit='5:00:00', workdir='.') as pool:
             return pool.map(infer,arr)
 
-pipe = AnimateDiffPipeline.from_pretrained(base, motion_adapter=adapter, torch_dtype=dtype).to(device)
+pipe = AnimateDiffPipeline.from_pretrained(base, vae=vae, motion_adapter=adapter, torch_dtype=dtype).to(device)
 pipe.scheduler = DDIMScheduler(
         clip_sample=False,
         beta_start=0.00085,
@@ -299,6 +305,6 @@ with gr.Blocks(theme=gr.themes.Soft(),css=css,js=js) as demo:
                 
     gr.on(
             triggers=[run_button.click, prompt.submit, prompt2.submit],
-            fn=run,inputs=[img,motion,prompt,prompt2,*result],outputs=result
+            fn=main,inputs=[img,motion,prompt,prompt2,*result],outputs=result
     )
     demo.queue().launch()