Update app.py

app.py

@@ -8,21 +8,20 @@ import requests
 import gradio as gr
 import numpy as np
 from lxml.html import fromstring
-#from transformers import pipeline
+from transformers import pipeline
 from torch import multiprocessing as mp
 #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
-#from huggingface_hub import hf_hub_download
-#from safetensors.torch import load_file
-from diffusers import DiffusionPipeline
-#from diffusers.utils import load_image
-#import jax
-#import jax.numpy as jnp
-
-def pipe_t2i():
+from diffusers.pipelines.flux import FluxPipeline
+from diffusers.utils import export_to_gif, load_image
+from huggingface_hub import hf_hub_download
+from safetensors.torch import load_file
+from diffusers import DiffusionPipeline, AnimateDiffPipeline, MotionAdapter, EulerDiscreteScheduler, StableDiffusionXLPipeline, UNet2DConditionModel
+import jax
+import jax.numpy as jnp
+
+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
     
@@ -72,19 +71,20 @@ def generate_random_string(length):
     return ''.join(random.choice(characters) for _ in range(length))
 
 @spaces.GPU(duration=40)
-def Piper(name,posi):
+def Piper(name,posi,neg):
+    global step
+    
     print("starting piper")
     
-    ret1 = pp1(
+    out = pipe(
         posi,
         height=512,
         width=512,
-        num_inference_steps=4,
-        max_sequence_length=256,
-        guidance_scale=0
+        num_inference_steps=step,
+        guidance_scale=1
     )
-    ret1.images[0].save(name)
 
+    export_to_gif(out.frames[0],name)
     return name
 
 css="""
@@ -122,6 +122,7 @@ footer {
 js="""
 function custom(){
     document.querySelector("div#prompt input").setAttribute("maxlength","38")
+    document.querySelector("div#prompt2 input").setAttribute("maxlength","38")
 }
 """
 
@@ -158,13 +159,34 @@ def run(p1,*result):
 def main():
 
     global result
-    global pp1
+    global pipe
+    global device
+    global step
+    global dtype
 
+    device = "cuda"
+    dtype = torch.float16
     result=[]
-    pp1=pipe_t2i()
+    step = 2
 
-    mp.set_start_method("spawn", force=True)
+    base = "stabilityai/stable-diffusion-xl-base-1.0"
+    repo = "ByteDance/SDXL-Lightning"
+    ckpt = "sdxl_lightning_4step_unet.safetensors" # Use the correct ckpt for your step setting!
+
+    unet = UNet2DConditionModel.from_config(base, subfolder="unet").to("cuda", torch.float16)
+    unet.load_state_dict(load_file(hf_hub_download(repo, ckpt), device="cuda"))
+
+    repo = "ByteDance/AnimateDiff-Lightning"
+    ckpt = f"animatediff_lightning_{step}step_diffusers.safetensors"
+
+    adapter = MotionAdapter().to(device, dtype)
+    adapter.load_state_dict(load_file(hf_hub_download(repo ,ckpt), device=device))
+
+    pipe = AnimateDiffPipeline.from_pretrained(base, motion_adapter=adapter, unet=unet, torch_dtype=dtype, variant="fp16").to(device)
+    pipe.scheduler = EulerDiscreteScheduler.from_config(pipe.scheduler.config, timestep_spacing="trailing", beta_schedule="linear")
     
+    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"""
@@ -182,8 +204,7 @@ def main():
             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))
-                result.append(gr.Image(interactive=False,elem_classes="image-container", label="Result", show_label=False, type='filepath', show_share_button=False))
-
+                
         gr.on(
             triggers=[run_button.click, prompt.submit],
             fn=run,inputs=[prompt,*result],outputs=result