Update app.py

app.py

@@ -1,31 +1,7 @@
 import gradio as gr
 import os
-import numpy as np
-import argparse
-import imageio
-import torch
 
-from einops import rearrange
-from diffusers import DDIMScheduler, AutoencoderKL
-from transformers import CLIPTextModel, CLIPTokenizer
-# from annotator.canny import CannyDetector
-# from annotator.openpose import OpenposeDetector
-# from annotator.midas import MidasDetector
-# import sys
-# sys.path.insert(0, ".")
-from huggingface_hub import hf_hub_download, snapshot_download
-import controlnet_aux
-from controlnet_aux import OpenposeDetector, CannyDetector, MidasDetector
-from controlnet_aux.open_pose.body import Body
-
-from models.pipeline_controlvideo import ControlVideoPipeline
-from models.util import save_videos_grid, read_video, get_annotation
-from models.unet import UNet3DConditionModel
-from models.controlnet import ControlNetModel3D
-from models.RIFE.IFNet_HDv3 import IFNet
-
-hf_token = os.environ.get('HF_TOKEN')
-device = "cuda"
+from huggingface_hub import snapshot_download
 
 model_ids = [
     'runwayml/stable-diffusion-v1-5',
@@ -37,122 +13,22 @@ for model_id in model_ids:
     model_name = model_id.split('/')[-1]
     snapshot_download(model_id, local_dir=f'checkpoints/{model_name}')
 
-sd_path = "checkpoints/stable-diffusion-v1-5"
-
-inter_path = "checkpoints/flownet.pkl"
-controlnet_dict = {
-    "pose": "checkpoints/sd-controlnet-openpose",
-    "depth": "checkpoints/sd-controlnet-depth",
-    "canny": "checkpoints/sd-controlnet-canny",
-}
+import subprocess
 
+def run_inference(prompt, video_path, condition, video_length):
+    command = "python inference.py --prompt prompt --condition condition --video_path video_path --output_path 'outputs/' --video_length video_length --smoother_steps 19 20"
+    output = subprocess.check_output(command, shell=True, text=True)
+    output = output.strip()  # Remove any leading/trailing whitespace
 
-controlnet_parser_dict = {
-    "pose": OpenposeDetector,
-    "depth": MidasDetector,
-    "canny": CannyDetector,
-}
+    # Process the output as needed
+    print("Command output:", output)
 
-POS_PROMPT = " ,best quality, extremely detailed, HD, ultra-realistic, 8K, HQ, masterpiece, trending on artstation, art, smooth"
-NEG_PROMPT = "longbody, lowres, bad anatomy, bad hands, missing fingers, extra digit, fewer difits, cropped, worst quality, low quality, deformed body, bloated, ugly, unrealistic"
+    return "done"
 
 
-
-def get_args():
-    parser = argparse.ArgumentParser()
-    parser.add_argument("--prompt", type=str, required=True, help="Text description of target video")
-    parser.add_argument("--video_path", type=str, required=True, help="Path to a source video")
-    parser.add_argument("--output_path", type=str, default="./outputs", help="Directory of output")
-    parser.add_argument("--condition", type=str, default="depth", help="Condition of structure sequence")
-    parser.add_argument("--video_length", type=int, default=15, help="Length of synthesized video")
-    parser.add_argument("--height", type=int, default=512, help="Height of synthesized video, and should be a multiple of 32")
-    parser.add_argument("--width", type=int, default=512, help="Width of synthesized video, and should be a multiple of 32")
-    parser.add_argument("--smoother_steps", nargs='+', default=[19, 20], type=int, help="Timesteps at which using interleaved-frame smoother")
-    parser.add_argument("--is_long_video", action='store_true', help="Whether to use hierarchical sampler to produce long video")
-    parser.add_argument("--seed", type=int, default=42, help="Random seed of generator")
     
-    args = parser.parse_args()
-    return args
-
-def infer(prompt, video_path, condition, video_length, is_long_video):
-    #args = get_args()
-    #os.makedirs(args.output_path, exist_ok=True)
-    
-    # Height and width should be a multiple of 32
-    output_path = ""
-    height = 512
-    width = 512
-    height = (height // 32) * 32    
-    width = (width // 32) * 32    
-    smoother_steps = [19, 20]
-    is_long_video = False
-    seed = 42
-   
-    if condition == "pose":
-        pretrained_model_or_path = "lllyasviel/ControlNet"
-        body_model_path = hf_hub_download(pretrained_model_or_path, "annotator/ckpts/body_pose_model.pth", cache_dir="checkpoints")
-        body_estimation = Body(body_model_path)
-        annotator = controlnet_parser_dict[condition](body_estimation)
-    else:
-        annotator = controlnet_parser_dict[condition]()
-
-    tokenizer = CLIPTokenizer.from_pretrained(sd_path, subfolder="tokenizer")
-    text_encoder = CLIPTextModel.from_pretrained(sd_path, subfolder="text_encoder").to(dtype=torch.float16)
-    vae = AutoencoderKL.from_pretrained(sd_path, subfolder="vae").to(dtype=torch.float16)
-    unet = UNet3DConditionModel.from_pretrained_2d(sd_path, subfolder="unet").to(dtype=torch.float16)
-    controlnet = ControlNetModel3D.from_pretrained_2d(controlnet_dict[condition]).to(dtype=torch.float16)
-    interpolater = IFNet(ckpt_path=inter_path).to(dtype=torch.float16)
-    scheduler=DDIMScheduler.from_pretrained(sd_path, subfolder="scheduler")
-
-    pipe = ControlVideoPipeline(
-            vae=vae, text_encoder=text_encoder, tokenizer=tokenizer, unet=unet,
-            controlnet=controlnet, interpolater=interpolater, scheduler=scheduler,
-        )
-    pipe.enable_vae_slicing()
-    pipe.enable_xformers_memory_efficient_attention()
-    pipe.to(device)
-
-    generator = torch.Generator(device="cuda")
-    generator.manual_seed(seed)
-
-    # Step 1. Read a video
-    video = read_video(video_path=video_path, video_length=video_length, width=width, height=height)
-
-    # Save source video
-    original_pixels = rearrange(video, "(b f) c h w -> b c f h w", b=1)
-    save_videos_grid(original_pixels, os.path.join(output_path, "source_video.mp4"), rescale=True)
-
-
-    # Step 2. Parse a video to conditional frames
-    pil_annotation = get_annotation(video, annotator)
-    if condition == "depth" and controlnet_aux.__version__ == '0.0.1':
-        pil_annotation = [pil_annot[0] for pil_annot in pil_annotation]
-
-    # Save condition video
-    video_cond = [np.array(p).astype(np.uint8) for p in pil_annotation]
-    imageio.mimsave(os.path.join(output_path, f"{condition}_condition.mp4"), video_cond, fps=8)
-
-    # Reduce memory (optional)
-    del annotator; torch.cuda.empty_cache()
-
-    # Step 3. inference
-
-    if is_long_video:
-        window_size = int(np.sqrt(video_length))
-        sample = pipe.generate_long_video(prompt + POS_PROMPT, video_length=video_length, frames=pil_annotation, 
-                    num_inference_steps=50, smooth_steps=args.smoother_steps, window_size=window_size,
-                    generator=generator, guidance_scale=12.5, negative_prompt=NEG_PROMPT,
-                    width=width, height=height
-                ).videos
-    else:
-        sample = pipe(prompt + POS_PROMPT, video_length=video_length, frames=pil_annotation, 
-                    num_inference_steps=50, smooth_steps=args.smoother_steps,
-                    generator=generator, guidance_scale=12.5, negative_prompt=NEG_PROMPT,
-                    width=width, height=height
-                ).videos
-    save_videos_grid(sample, f"{output_path}/{prompt}.mp4")
 
-    return f"{output_path}/{prompt}.mp4"
+    #return f"{output_path}/{prompt}.mp4"
 
 with gr.Blocks() as demo:
     with gr.Column():
@@ -160,16 +36,16 @@ with gr.Blocks() as demo:
         video_path = gr.Video(source="upload", type="filepath")
         condition = gr.Textbox(label="Condition", value="depth")
         video_length = gr.Slider(label="video length", minimum=1, maximum=15, step=1, value=2)
-        seed = gr.Number(label="seed", value=42)
+        #seed = gr.Number(label="seed", value=42)
         submit_btn = gr.Button("Submit")
-        video_res = gr.Video(label="result")
+        #video_res = gr.Video(label="result")
+        video_res = gr.Textbox(label="result")
 
-    submit_btn.click(fn=infer, 
+    submit_btn.click(fn=run_inference, 
                      inputs=[prompt,
                              video_path,
                              condition,
-                             video_length,
-                             seed,
+                             video_length
                             ],
                     outputs=[video_res])