Spaces:
Running
Running
File size: 3,870 Bytes
fff8451 |
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 |
import torch
import imageio
import os
import gradio as gr
from diffusers.schedulers import EulerAncestralDiscreteScheduler
from transformers import T5EncoderModel, T5Tokenizer
from allegro.pipelines.pipeline_allegro import AllegroPipeline
from allegro.models.vae.vae_allegro import AllegroAutoencoderKL3D
from allegro.models.transformers.transformer_3d_allegro import AllegroTransformer3DModel
import huggingface_hub
weights_dir = './allegro_weights'
os.makedirs(weights_dir, exist_ok=True)
huggingface_hub.snapshot_download(
repo_id='rhymes-ai/Allegro',
allow_patterns=[
'scheduler/**',
'text_encoder/**',
'tokenizer/**',
'transformer/**',
'vae/**',
],
local_dir=weights_dir,
local_dir_use_symlinks=False,
)
def single_inference(user_prompt, save_path, guidance_scale, num_sampling_steps, seed, enable_cpu_offload):
dtype = torch.bfloat16
# Load models
vae = AllegroAutoencoderKL3D.from_pretrained(
"weights_dir/vae/",
torch_dtype=torch.float32
).cuda()
vae.eval()
text_encoder = T5EncoderModel.from_pretrained("weights_dir/text_encoder/", torch_dtype=dtype)
text_encoder.eval()
tokenizer = T5Tokenizer.from_pretrained("weights_dir/tokenizer/")
scheduler = EulerAncestralDiscreteScheduler()
transformer = AllegroTransformer3DModel.from_pretrained("weights_dir/transformer/", torch_dtype=dtype).cuda()
transformer.eval()
allegro_pipeline = AllegroPipeline(
vae=vae,
text_encoder=text_encoder,
tokenizer=tokenizer,
scheduler=scheduler,
transformer=transformer
).to("cuda:0")
positive_prompt = """
(masterpiece), (best quality), (ultra-detailed), (unwatermarked),
{}
emotional, harmonious, vignette, 4k epic detailed, shot on kodak, 35mm photo,
sharp focus, high budget, cinemascope, moody, epic, gorgeous
"""
negative_prompt = """
nsfw, lowres, bad anatomy, bad hands, text, error, missing fingers, extra digit, fewer digits, cropped, worst quality,
low quality, normal quality, jpeg artifacts, signature, watermark, username, blurry.
"""
# Process user prompt
user_prompt = positive_prompt.format(user_prompt.lower().strip())
if enable_cpu_offload:
allegro_pipeline.enable_sequential_cpu_offload()
out_video = allegro_pipeline(
user_prompt,
negative_prompt=negative_prompt,
num_frames=88,
height=720,
width=1280,
num_inference_steps=num_sampling_steps,
guidance_scale=guidance_scale,
max_sequence_length=512,
generator=torch.Generator(device="cuda:0").manual_seed(seed)
).video[0]
# Save video
os.makedirs(os.path.dirname(save_path), exist_ok=True)
imageio.mimwrite(save_path, out_video, fps=15, quality=8)
return save_path
# Gradio interface function
def run_inference(user_prompt, guidance_scale, num_sampling_steps, seed, enable_cpu_offload):
save_path = "./output_videos/generated_video.mp4"
result_path = single_inference(user_prompt, save_path, guidance_scale, num_sampling_steps, seed, enable_cpu_offload)
return result_path
# Create Gradio interface
iface = gr.Interface(
fn=run_inference,
inputs=[
gr.Textbox(label="User Prompt"),
gr.Slider(minimum=0, maximum=20, step=0.1, label="Guidance Scale", value=7.5),
gr.Slider(minimum=10, maximum=200, step=1, label="Number of Sampling Steps", value=100),
gr.Slider(minimum=0, maximum=10000, step=1, label="Random Seed", value=42),
gr.Checkbox(label="Enable CPU Offload", value=False),
],
outputs=gr.Video(label="Generated Video"),
title="Allegro Video Generation",
description="Generate a video based on a text prompt using the Allegro pipeline."
)
# Launch the interface
iface.launch()
|