Merge pull request #379 from AlexandrePoisson/patch-1

Update app.py
This commit is contained in:
Yuxuan.Zhang 2024-10-01 12:24:39 +08:00 committed by GitHub
commit 928f812f71
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194

View File

@ -37,13 +37,15 @@ from huggingface_hub import hf_hub_download, snapshot_download
device = "cuda" if torch.cuda.is_available() else "cpu" device = "cuda" if torch.cuda.is_available() else "cpu"
MODEL = "THUDM/CogVideoX-5b"
hf_hub_download(repo_id="ai-forever/Real-ESRGAN", filename="RealESRGAN_x4.pth", local_dir="model_real_esran") hf_hub_download(repo_id="ai-forever/Real-ESRGAN", filename="RealESRGAN_x4.pth", local_dir="model_real_esran")
snapshot_download(repo_id="AlexWortega/RIFE", local_dir="model_rife") snapshot_download(repo_id="AlexWortega/RIFE", local_dir="model_rife")
pipe = CogVideoXPipeline.from_pretrained("/share/official_pretrains/hf_home/CogVideoX-5b", torch_dtype=torch.bfloat16).to(device) pipe = CogVideoXPipeline.from_pretrained(MODEL, torch_dtype=torch.bfloat16).to(device)
pipe.scheduler = CogVideoXDPMScheduler.from_config(pipe.scheduler.config, timestep_spacing="trailing") pipe.scheduler = CogVideoXDPMScheduler.from_config(pipe.scheduler.config, timestep_spacing="trailing")
pipe_video = CogVideoXVideoToVideoPipeline.from_pretrained( pipe_video = CogVideoXVideoToVideoPipeline.from_pretrained(
"/share/official_pretrains/hf_home/CogVideoX-5b", MODEL,
transformer=pipe.transformer, transformer=pipe.transformer,
vae=pipe.vae, vae=pipe.vae,
scheduler=pipe.scheduler, scheduler=pipe.scheduler,
@ -53,9 +55,9 @@ pipe_video = CogVideoXVideoToVideoPipeline.from_pretrained(
).to(device) ).to(device)
pipe_image = CogVideoXImageToVideoPipeline.from_pretrained( pipe_image = CogVideoXImageToVideoPipeline.from_pretrained(
"/share/official_pretrains/hf_home/CogVideoX-5b-I2V", MODEL,
transformer=CogVideoXTransformer3DModel.from_pretrained( transformer=CogVideoXTransformer3DModel.from_pretrained(
"/share/official_pretrains/hf_home/CogVideoX-5b-I2V", subfolder="transformer", torch_dtype=torch.bfloat16 MODEL, subfolder="transformer", torch_dtype=torch.bfloat16
), ),
vae=pipe.vae, vae=pipe.vae,
scheduler=pipe.scheduler, scheduler=pipe.scheduler,
@ -315,7 +317,7 @@ with gr.Blocks() as demo:
"></a> "></a>
</div> </div>
<div style="text-align: center; font-size: 15px; font-weight: bold; color: red; margin-bottom: 20px;"> <div style="text-align: center; font-size: 15px; font-weight: bold; color: red; margin-bottom: 20px;">
This demo is for academic research and experiential use only. This demo is for academic research and experimental use only.
</div> </div>
""") """)
with gr.Row(): with gr.Row():