Update base/app.py

base/app.py

@@ -1,5 +1,5 @@
 import gradio as gr
-from text_to_video import model_t2v_fun,setup_seed
+from text_to_video import model_t2v_fun, setup_seed
 from omegaconf import OmegaConf
 import torch
 import imageio
@@ -16,10 +16,9 @@ from diffusers.schedulers import DDIMScheduler, DDPMScheduler, PNDMScheduler, Eu
 from diffusers.models import AutoencoderKL
 from transformers import CLIPTokenizer, CLIPTextModel, CLIPTextModelWithProjection
 
-
 config_path = "./base/configs/sample.yaml"
 args = OmegaConf.load("./base/configs/sample.yaml")
-device = "cuda" if torch.cuda.is_available() else "cpu"
+device = "cpu"  # Force CPU usage
 
 css = """
 h1 {
@@ -32,52 +31,51 @@ h1 {
 """
 
 sd_path = args.pretrained_path
-unet = get_models(args, sd_path).to(device, dtype=torch.float16)
+unet = get_models(args, sd_path).to(device, dtype=torch.float32)  # Use float32 for CPU
 state_dict = find_model("./pretrained_models/lavie_base.pt")
 unet.load_state_dict(state_dict)
-vae = AutoencoderKL.from_pretrained(sd_path, subfolder="vae", torch_dtype=torch.float16).to(device)
+vae = AutoencoderKL.from_pretrained(sd_path, subfolder="vae", torch_dtype=torch.float32).to(device)  # Use float32 for CPU
 tokenizer_one = CLIPTokenizer.from_pretrained(sd_path, subfolder="tokenizer")
-text_encoder_one = CLIPTextModel.from_pretrained(sd_path, subfolder="text_encoder", torch_dtype=torch.float16).to(device) # huge
+text_encoder_one = CLIPTextModel.from_pretrained(sd_path, subfolder="text_encoder", torch_dtype=torch.float32).to(device)  # Use float32 for CPU
 unet.eval()
 vae.eval()
 text_encoder_one.eval()
 
-def infer(prompt, seed_inp, ddim_steps,cfg, infer_type):
-    if seed_inp!=-1:
+def infer(prompt, seed_inp, ddim_steps, cfg, infer_type):
+    if seed_inp != -1:
         setup_seed(seed_inp)
     else:
         seed_inp = random.choice(range(10000000))
         setup_seed(seed_inp)
     if infer_type == 'ddim':
         scheduler = DDIMScheduler.from_pretrained(sd_path, 
-											   subfolder="scheduler",
-											   beta_start=args.beta_start, 
-											   beta_end=args.beta_end, 
-											   beta_schedule=args.beta_schedule)
+                                               subfolder="scheduler",
+                                               beta_start=args.beta_start, 
+                                               beta_end=args.beta_end, 
+                                               beta_schedule=args.beta_schedule)
     elif infer_type == 'eulerdiscrete':
         scheduler = EulerDiscreteScheduler.from_pretrained(sd_path,
-        									   subfolder="scheduler",
-											   beta_start=args.beta_start,
-											   beta_end=args.beta_end,
-											   beta_schedule=args.beta_schedule)
+                                               subfolder="scheduler",
+                                               beta_start=args.beta_start,
+                                               beta_end=args.beta_end,
+                                               beta_schedule=args.beta_schedule)
     elif infer_type == 'ddpm':
         scheduler = DDPMScheduler.from_pretrained(sd_path,
-											  subfolder="scheduler",
-											  beta_start=args.beta_start,
-											  beta_end=args.beta_end,
-											  beta_schedule=args.beta_schedule)
+                                              subfolder="scheduler",
+                                              beta_start=args.beta_start,
+                                              beta_end=args.beta_end,
+                                              beta_schedule=args.beta_schedule)
     model = VideoGenPipeline(vae=vae, text_encoder=text_encoder_one, tokenizer=tokenizer_one, scheduler=scheduler, unet=unet)
     model.to(device)
-    if device == "cuda":
-        model.enable_xformers_memory_efficient_attention()
-    videos = model(prompt, video_length=16, height = 320, width= 512, num_inference_steps=ddim_steps, guidance_scale=cfg).video
+    # Disable xformers for CPU
+    # if device == "cuda":
+    #     model.enable_xformers_memory_efficient_attention()
+    videos = model(prompt, video_length=8, height=160, width=256, num_inference_steps=ddim_steps, guidance_scale=cfg).video  # Reduced resolution and length
     if not os.path.exists(args.output_folder):
         os.mkdir(args.output_folder)
-    torchvision.io.write_video(args.output_folder + prompt[0:30].replace(' ', '_') + '-'+str(seed_inp)+'-'+str(ddim_steps)+'-'+str(cfg)+ '-.mp4', videos[0], fps=8)
-
-
-    return args.output_folder + prompt[0:30].replace(' ', '_') + '-'+str(seed_inp)+'-'+str(ddim_steps)+'-'+str(cfg)+ '-.mp4'
+    torchvision.io.write_video(args.output_folder + prompt[0:30].replace(' ', '_') + '-' + str(seed_inp) + '-' + str(ddim_steps) + '-' + str(cfg) + '-.mp4', videos[0], fps=4)  # Reduced FPS
 
+    return args.output_folder + prompt[0:30].replace(' ', '_') + '-' + str(seed_inp) + '-' + str(ddim_steps) + '-' + str(cfg) + '-.mp4'
 
 title = """
     <div style="text-align: center; max-width: 700px; margin: 0 auto;">
@@ -109,12 +107,11 @@ with gr.Blocks(css='style.css') as demo:
     with gr.Column():
         with gr.Row(elem_id="col-container"):
             with gr.Column():
-                    
                 prompt = gr.Textbox(value="a corgi walking in the park at sunrise, oil painting style", label="Prompt", placeholder="enter prompt", show_label=True, elem_id="prompt-in", min_width=200, lines=2)
-                infer_type = gr.Dropdown(['ddpm','ddim','eulerdiscrete'], label='infer_type',value='ddim')
+                infer_type = gr.Dropdown(['ddpm', 'ddim', 'eulerdiscrete'], label='infer_type', value='ddim')
                 ddim_steps = gr.Slider(label='Steps', minimum=50, maximum=300, value=50, step=1)
-                seed_inp = gr.Slider(value=-1,label="seed (for random generation, use -1)",show_label=True,minimum=-1,maximum=2147483647)
-                cfg = gr.Number(label="guidance_scale",value=7.5)
+                seed_inp = gr.Slider(value=-1, label="seed (for random generation, use -1)", show_label=True, minimum=-1, maximum=2147483647)
+                cfg = gr.Number(label="guidance_scale", value=7.5)
 
             with gr.Column():
                 submit_btn = gr.Button("Generate video")
@@ -124,27 +121,25 @@ with gr.Blocks(css='style.css') as demo:
             outputs = [video_out]
 
         ex = gr.Examples(
-            examples = [['a corgi walking in the park at sunrise, oil painting style',400,50,7,'ddim'],
-                    ['a cute teddy bear reading a book in the park, oil painting style, high quality',700,50,7,'ddim'],
-                    ['an epic tornado attacking above a glowing city at night, the tornado is made of smoke, highly detailed',230,50,7,'ddim'],
-                    ['a jar filled with fire, 4K video, 3D rendered, well-rendered',400,50,7,'ddim'],
-                    ['a teddy bear walking in the park, oil painting style, high quality',400,50,7,'ddim'],
-                    ['a teddy bear walking on the street, 2k, high quality',100,50,7,'ddim'],
-                    ['a panda taking a selfie, 2k, high quality',400,50,7,'ddim'],
-                    ['a polar bear playing drum kit in NYC Times Square, 4k, high resolution',400,50,7,'ddim'],
-                    ['jungle river at sunset, ultra quality',400,50,7,'ddim'],
-                    ['a shark swimming in clear Carribean ocean, 2k, high quality',400,50,7,'ddim'],
-                    ['A steam train moving on a mountainside by Vincent van Gogh',230,50,7,'ddim'],
-                    ['a confused grizzly bear in calculus class',1000,50,7,'ddim']],
-            fn = infer,
-            inputs=[prompt, seed_inp, ddim_steps,cfg,infer_type],
+            examples=[['a corgi walking in the park at sunrise, oil painting style', 400, 50, 7, 'ddim'],
+                      ['a cute teddy bear reading a book in the park, oil painting style, high quality', 700, 50, 7, 'ddim'],
+                      ['an epic tornado attacking above a glowing city at night, the tornado is made of smoke, highly detailed', 230, 50, 7, 'ddim'],
+                      ['a jar filled with fire, 4K video, 3D rendered, well-rendered', 400, 50, 7, 'ddim'],
+                      ['a teddy bear walking in the park, oil painting style, high quality', 400, 50, 7, 'ddim'],
+                      ['a teddy bear walking on the street, 2k, high quality', 100, 50, 7, 'ddim'],
+                      ['a panda taking a selfie, 2k, high quality', 400, 50, 7, 'ddim'],
+                      ['a polar bear playing drum kit in NYC Times Square, 4k, high resolution', 400, 50, 7, 'ddim'],
+                      ['jungle river at sunset, ultra quality', 400, 50, 7, 'ddim'],
+                      ['a shark swimming in clear Carribean ocean, 2k, high quality', 400, 50, 7, 'ddim'],
+                      ['A steam train moving on a mountainside by Vincent van Gogh', 230, 50, 7, 'ddim'],
+                      ['a confused grizzly bear in calculus class', 1000, 50, 7, 'ddim']],
+            fn=infer,
+            inputs=[prompt, seed_inp, ddim_steps, cfg, infer_type],
             outputs=[video_out],
             cache_examples=True,
         )
         ex.dataset.headers = [""]
-         
-    submit_btn.click(infer, inputs, outputs)
-
-demo.queue(max_size=12, api_open=False).launch(show_api=False)
 
+    submit_btn.click(infer, inputs, outputs)
 
+demo.queue(max_size=12, api_open=False).launch(show_api=False)