Spaces:
Running
on
Zero
Running
on
Zero
File size: 4,317 Bytes
a45ef62 b9a7927 8fecf84 a45ef62 8fecf84 a45ef62 8fecf84 a45ef62 b9a7927 a45ef62 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 |
import gc
import os
import torch
import spaces
import gradio as gr
from diffusers import LattePipeline
from transformers import T5EncoderModel, BitsAndBytesConfig
import imageio
from torchvision.utils import save_image
def flush():
gc.collect()
torch.cuda.empty_cache()
def bytes_to_giga_bytes(bytes):
return bytes / 1024 / 1024 / 1024
def initialize_pipeline():
model_id = "maxin-cn/Latte-1"
text_encoder = T5EncoderModel.from_pretrained(
model_id,
subfolder="text_encoder",
quantization_config=BitsAndBytesConfig(load_in_4bit=True, bnb_4bit_compute_dtype=torch.float16),
device_map="auto",
)
pipe = LattePipeline.from_pretrained(
model_id,
text_encoder=text_encoder,
transformer=None,
device_map="balanced",
)
return pipe, text_encoder
@spaces.GPU(duration=120)
def generate_video(
prompt: str,
negative_prompt: str = "",
video_length: int = 16,
num_inference_steps: int = 50,
progress=gr.Progress()
):
# Set random seed for reproducibility
torch.manual_seed(0)
# Initialize the pipeline
progress(0, desc="Initializing pipeline...")
pipe, text_encoder = initialize_pipeline()
# Generate prompt embeddings
progress(0.2, desc="Encoding prompt...")
with torch.no_grad():
prompt_embeds, negative_prompt_embeds = pipe.encode_prompt(
prompt,
negative_prompt=negative_prompt
)
# Clean up first pipeline
progress(0.3, desc="Cleaning up...")
del text_encoder
del pipe
flush()
# Initialize the second pipeline
progress(0.4, desc="Initializing generation pipeline...")
pipe = LattePipeline.from_pretrained(
"maxin-cn/Latte-1",
text_encoder=None,
torch_dtype=torch.float16,
).to("cuda")
# Generate video
progress(0.5, desc="Generating video...")
videos = pipe(
video_length=video_length,
num_inference_steps=num_inference_steps,
negative_prompt=None,
prompt_embeds=prompt_embeds,
negative_prompt_embeds=negative_prompt_embeds,
output_type="pt",
).frames.cpu()
progress(0.8, desc="Post-processing...")
# Convert to video format
videos = (videos.clamp(0, 1) * 255).to(dtype=torch.uint8)
# Save temporary file
temp_output = "temp_output.mp4"
imageio.mimwrite(
temp_output,
videos[0].permute(0, 2, 3, 1),
fps=8,
quality=5
)
# Clean up
progress(0.9, desc="Cleaning up...")
del pipe
flush()
progress(1.0, desc="Done!")
return temp_output
def create_demo():
with gr.Blocks() as demo:
gr.Markdown("""
# Latte Video Generation
Generate short videos using the Latte-1 model.
""")
with gr.Row():
with gr.Column():
prompt = gr.Textbox(
label="Prompt",
value="a cat wearing sunglasses and working as a lifeguard at pool.",
info="Describe what you want to generate"
)
negative_prompt = gr.Textbox(
label="Negative Prompt",
value="",
info="What you don't want to see in the generation"
)
video_length = gr.Slider(
minimum=8,
maximum=32,
step=8,
value=16,
label="Video Length (frames)"
)
steps = gr.Slider(
minimum=20,
maximum=100,
step=10,
value=50,
label="Number of Inference Steps"
)
generate_btn = gr.Button("Generate Video")
with gr.Column():
output_video = gr.Video(label="Generated Video")
generate_btn.click(
fn=generate_video,
inputs=[prompt, negative_prompt, video_length, steps],
outputs=output_video
)
return demo
if __name__ == "__main__":
demo = create_demo()
demo.queue()
demo.launch(share=False)
|