Upload 2 files

app.py

@@ -1,7 +1,103 @@
+import spaces
 import gradio as gr
+import torch
+from diffusers import LTXPipeline
+import uuid
+import time
+import types
+from typing import Optional
 
-def greet(name):
-    return "Hello " + name + "!!"
+pipe = LTXPipeline.from_pretrained("a-r-r-o-w/LTX-Video-0.9.1-diffusers", torch_dtype=torch.bfloat16)
+pipe.to("cuda")
+pipe.transformer = torch.compile(pipe.transformer, mode="reduce-overhead", fullgraph=True)
+pipe.vae = torch.compile(pipe.vae, mode="reduce-overhead", fullgraph=True)
 
-demo = gr.Interface(fn=greet, inputs="text", outputs="text")
-demo.launch()
+
+# pipe.vae.decode = vae_decode
+
+HEIGHT = 480
+WIDTH = 640
+N_FRAME = 161
+N_AVG_FRAME = 2
+
+negative_prompt = "worst quality, inconsistent motion, blurry, jittery, distorted"
+
+prepare_latents_original = pipe.prepare_latents
+
+# unpack will have shape B, C, F, H, W with F, H, W are in latent dim
+
+# def prepare_latents_loop(*args, **kwargs):
+#     packed_latents = prepare_latents_original(*args, **kwargs)
+#     unpacked_latents = pipe._unpack_latents(packed_latents, (N_FRAME-1)//8+1, HEIGHT//32, WIDTH//32, 1, 1)
+#     # now average the first n and last n frames
+#     last_n = unpacked_latents[:, :, -N_AVG_FRAME:, :, :]
+#     # 0,1,2,3,4, roll -1 => 1,2,3,4,0
+#     # last n: [3, 4]
+#     # last_next_n: [4, 0]
+#     # then 3 will be 0.75*3 + 0.25*4, and 4 will be 0.75*4+0.25*0
+#     last_next_n = torch.roll(unpacked_latents, shifts=-1, dims=2)[:, :, -N_AVG_FRAME:, :, :]
+#     avg_n = last_n * 0.75 + last_next_n * 0.25
+#     unpacked_latents[:, :, -N_AVG_FRAME:, :, :] = avg_n
+#     # pack the latents back
+#     packed_latents = pipe._pack_latents(unpacked_latents)
+#     return packed_latents
+
+# pipe.prepare_latents = prepare_latents_loop
+
+# with the shift it will become step=0 0,1,2,3 -> step=1 1,2,3,0 -> step=2 2,3,0,1 -> step=3 3,0,1,2 -> step=4 0,1,2,3
+# so we only shift (N_FRAME-1)//8+1 times
+
+def modify_latents_callback(pipeline, step, timestep, callback_kwargs):
+    print("Rolling latents on step", step)
+    latents = callback_kwargs.get("latents")
+    unpacked_latents = pipeline._unpack_latents(latents, (N_FRAME-1)//8+1, HEIGHT//32, WIDTH//32, 1, 1)
+    modified_latents = torch.roll(unpacked_latents, shifts=1, dims=2)
+    modified_latents = pipeline._pack_latents(modified_latents)
+    return {"latents": modified_latents}
+
+@spaces.GPU(duration=120)
+def generate_gif(prompt, use_fixed_seed):
+    seed = 0 if use_fixed_seed else torch.seed()
+    generator = torch.Generator(device="cuda").manual_seed(seed)
+
+    output = pipe(
+        prompt=prompt,
+        negative_prompt=negative_prompt,
+        width=WIDTH,
+        height=HEIGHT,
+        num_frames=N_FRAME,
+        num_inference_steps=50,
+        decode_timestep=0.03,
+        decode_noise_scale=0.025,
+        generator=generator,
+        callback_on_step_end=modify_latents_callback,
+    ).frames[0]
+
+    gif_path = f"/tmp/{uuid.uuid4().hex}.gif"
+
+    bef = time.time()
+    # imageio.mimsave(gif_path, output, format="GIF", fps=24, loop=0)
+    gif_path = f"/tmp/{uuid.uuid4().hex}.webp"
+    output[0].save(gif_path, format="WebP", save_all=True, append_images=output[1:], duration=1000/24, loop=0)
+    print("GIF creation time:", time.time() - bef)
+    return gif_path
+
+with gr.Blocks() as demo:
+    gr.Markdown("## LTX Video → Looping GIF Generator")
+    with gr.Row():
+        with gr.Column():
+            prompt_input = gr.Textbox(label="Prompt", lines=4)
+            use_fixed_seed = gr.Checkbox(label="Use Fixed Seed", value=True)
+            generate_btn = gr.Button("Generate")
+        with gr.Column():
+            gif_output = gr.Image(label="Looping GIF Result", type="filepath")
+
+    generate_btn.click(
+        fn=generate_gif,
+        inputs=[prompt_input, use_fixed_seed],
+        outputs=gif_output,
+        concurrency_limit=1
+    )
+
+demo.queue(max_size=5)
+demo.launch(share=True)

requirements.txt

@@ -0,0 +1,6 @@
+gradio
+diffusers==0.32.2
+transformers
+sentencepiece
+torch
+accelerate