Update app.py

app.py

@@ -22,6 +22,56 @@ from diffusers import DiffusionPipeline, AnimateDiffPipeline, MotionAdapter, Eul
 import jax
 import jax.numpy as jnp
 
+last_motion=None
+fps=20
+time=16
+device = "cuda"
+dtype = torch.float16
+result=[]
+step = 2
+progress=gr.Progress()
+
+css="""
+input, input::placeholder {
+    text-align: center !important;
+}
+*, *::placeholder {
+    font-family: Suez One !important;
+}
+h1,h2,h3,h4,h5,h6 {
+    width: 100%;
+    text-align: center;
+}
+footer {
+    display: none !important;
+}
+#col-container {
+    margin: 0 auto;
+    max-width: 15cm;
+}
+.image-container {
+    aspect-ratio: 576 / 1024 !important;
+}
+.dropdown-arrow {
+    display: none !important;
+}
+*:has(>.btn) {
+    display: flex;
+    justify-content: space-evenly;
+    align-items: center;
+}
+.btn {
+    display: flex;
+}
+"""
+
+js="""
+function custom(){
+    document.querySelector("div#prompt input").setAttribute("maxlength","38")
+    document.querySelector("div#prompt2 input").setAttribute("maxlength","38")
+}
+"""
+
 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
@@ -85,8 +135,6 @@ def Piper(name,positive_prompt,negative,motion):
     global base
     global device
 
-    progress((0, step))
-
     if last_motion != motion:
         pipe.unload_lora_weights()
         if motion != "":
@@ -107,47 +155,6 @@ def Piper(name,positive_prompt,negative,motion):
     export_to_gif(out.frames[0],name,fps=fps)
     return name
 
-css="""
-input, input::placeholder {
-    text-align: center !important;
-}
-*, *::placeholder {
-    font-family: Suez One !important;
-}
-h1,h2,h3,h4,h5,h6 {
-    width: 100%;
-    text-align: center;
-}
-footer {
-    display: none !important;
-}
-#col-container {
-    margin: 0 auto;
-    max-width: 15cm;
-}
-.image-container {
-    aspect-ratio: 576 / 1024 !important;
-}
-.dropdown-arrow {
-    display: none !important;
-}
-*:has(>.btn) {
-    display: flex;
-    justify-content: space-evenly;
-    align-items: center;
-}
-.btn {
-    display: flex;
-}
-"""
-
-js="""
-function custom(){
-    document.querySelector("div#prompt input").setAttribute("maxlength","38")
-    document.querySelector("div#prompt2 input").setAttribute("maxlength","38")
-}
-"""
-
 def infer(pm):
         print("infer: started")
     
@@ -180,65 +187,42 @@ def run(m,p1,p2,*result):
 
         return out
 
-def main():
-
-    global result
-    global pipe
-    global device
-    global step
-    global dtype
-    global progress
-    global fps
-    global time
-    global last_motion
-    global base
-    
-    last_motion=None
-    fps=20
-    time=16
-    device = "cuda"
-    dtype = torch.float16
-    result=[]
-    step = 2
-
-    progress=gr.Progress()
-
-    adapter = MotionAdapter.from_pretrained("guoyww/animatediff-motion-adapter-v1-5-3")
-    vae = AutoencoderKL.from_single_file("https://huggingface.co/stabilityai/sd-vae-ft-mse-original/vae-ft-mse-840000-ema-pruned.safetensors")
-    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")
+vae = AutoencoderKL.from_single_file("https://huggingface.co/stabilityai/sd-vae-ft-mse-original/vae-ft-mse-840000-ema-pruned.safetensors")
+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)
     
-    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"
+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"
     
-    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")
-    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=True)
+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")
+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=True)
     
-    mp.set_start_method("spawn", force=True)
-
-    with gr.Blocks(theme=gr.themes.Soft(),css=css,js=js) as demo:
-        with gr.Column(elem_id="col-container"):
-            gr.Markdown(f"""
-                # MULTI-LANGUAGE IMAGE GENERATOR
-            """)
-            with gr.Row():
-                prompt = gr.Textbox(
-                    elem_id="prompt",
-                    placeholder="INCLUDE",
-                    container=False,
-                    max_lines=1
-                )
-            with gr.Row():
-                prompt2 = gr.Textbox(
-                    elem_id="prompt2",
-                    placeholder="EXCLUDE",
-                    container=False,
-                    max_lines=1
-                )
-            with gr.Row():
+mp.set_start_method("spawn", force=True)
+
+with gr.Blocks(theme=gr.themes.Soft(),css=css,js=js) as demo:
+    with gr.Column(elem_id="col-container"):
+        gr.Markdown(f"""
+            # MULTI-LANGUAGE IMAGE GENERATOR
+        """)
+        with gr.Row():
+            prompt = gr.Textbox(
+                elem_id="prompt",
+                placeholder="INCLUDE",
+                container=False,
+                max_lines=1
+            )
+        with gr.Row():
+            prompt2 = gr.Textbox(
+                elem_id="prompt2",
+                placeholder="EXCLUDE",
+                container=False,
+                max_lines=1
+            )
+        with gr.Row():
                 motion = gr.Dropdown(
                     label='Motion',
                     show_label=False,
@@ -256,17 +240,14 @@ def main():
                     value="",
                     interactive=True
                 )
-            with gr.Row():
+        with gr.Row():
                 run_button = gr.Button("START",elem_classes="btn",scale=0)
-            with gr.Row():
+        with gr.Row():
                 result.append(gr.Image(interactive=False,elem_classes="image-container", label="Result", show_label=False, type='filepath', show_share_button=False))
                 result.append(gr.Image(interactive=False,elem_classes="image-container", label="Result", show_label=False, type='filepath', show_share_button=False))
                 
-        gr.on(
+    gr.on(
             triggers=[run_button.click, prompt.submit, prompt2.submit],
             fn=run,inputs=[motion,prompt,prompt2,*result],outputs=result
-        )
-        demo.queue().launch()
-
-if __name__ == "__main__":
-    main()
+    )
+    demo.queue().launch()