Update app.py

app.py

@@ -2,6 +2,7 @@ import os
 import copy
 import torch
 import random
+import spaces
 import gradio as gr
 from glob import glob
 from omegaconf import OmegaConf
@@ -14,7 +15,6 @@ from transformers import CLIPTextModel, CLIPTokenizer
 from utils.unet import UNet3DConditionModel
 from utils.pipeline_magictime import MagicTimePipeline
 from utils.util import save_videos_grid, convert_ldm_unet_checkpoint, convert_ldm_clip_checkpoint, convert_ldm_vae_checkpoint, load_diffusers_lora_unet, convert_ldm_clip_text_model
-import spaces
 
 pretrained_model_path   = "./ckpts/Base_Model/stable-diffusion-v1-5"
 inference_config_path   = "./sample_configs/RealisticVision.yaml"
@@ -62,7 +62,7 @@ examples = [
 print(f"### Cleaning cached examples ...")
 os.system(f"rm -rf gradio_cached_examples/")
 
-@spaces.GPU(duration=300)
+
 class MagicTimeController:
     def __init__(self):
         
@@ -208,72 +208,51 @@ class MagicTimeController:
             "dreambooth": dreambooth_dropdown,
         }
         return gr.Video(value=save_sample_path), gr.Json(value=json_config)
+        
+controller = MagicTimeController()
 
-# inference_config = OmegaConf.load(inference_config_path)[1]
-# tokenizer        =  CLIPTokenizer.from_pretrained(pretrained_model_path, subfolder="tokenizer")
-# text_encoder     =  CLIPTextModel.from_pretrained(pretrained_model_path, subfolder="text_encoder").cuda()
-# vae              =  AutoencoderKL.from_pretrained(pretrained_model_path, subfolder="vae").cuda()
-# unet             =  UNet3DConditionModel.from_pretrained_2d(pretrained_model_path, subfolder="unet", unet_additional_kwargs=OmegaConf.to_container(inference_config.unet_additional_kwargs)).cuda()
-# text_model       =  CLIPTextModel.from_pretrained("openai/clip-vit-large-patch14")
-# controller       = MagicTimeController(tokenizer=tokenizer, text_encoder=text_encoder, vae=vae, unet=unet, text_model=text_model)
-controller       = MagicTimeController()
-
-def ui():
-    with gr.Blocks(css=css) as demo:
-        gr.Markdown(
-            """
-            <div style='display: flex; align-items: center; justify-content: center; text-align: center;'>
-                <img src='https://www.pnglog.com/48rWnj.png' style='width: 300px; height: auto; margin-right: 10px;' />
-            </div>
-            
-            <h2 align="center"> <a href="https://github.com/PKU-YuanGroup/MagicTime">MagicTime: Time-lapse Video Generation Models as Metamorphic Simulators</a></h2>
-            <h5 style="text-align:left;">If you like our project, please give us a star ⭐ on GitHub for the latest update.</h5>
-
-            [GitHub](https://github.com/PKU-YuanGroup/MagicTime) | [arXiv](https://arxiv.org/abs/2404.05014) | [Home Page](https://pku-yuangroup.github.io/MagicTime/) | [Dataset](https://drive.google.com/drive/folders/1WsomdkmSp3ql3ImcNsmzFuSQ9Qukuyr8?usp=sharing)
-            """
-        )
-        with gr.Row():
-            with gr.Column():
-                dreambooth_dropdown     = gr.Dropdown( label="DreamBooth Model", choices=controller.dreambooth_list,    value=controller.dreambooth_list[0],    interactive=True )
-                motion_module_dropdown  = gr.Dropdown( label="Motion Module",  choices=controller.motion_module_list, value=controller.motion_module_list[0], interactive=True )
-
-                dreambooth_dropdown.change(fn=controller.update_dreambooth,       inputs=[dreambooth_dropdown],    outputs=[dreambooth_dropdown])
-                motion_module_dropdown.change(fn=controller.update_motion_module, inputs=[motion_module_dropdown], outputs=[motion_module_dropdown])
-
-                prompt_textbox          = gr.Textbox( label="Prompt",          lines=3 )
-                negative_prompt_textbox = gr.Textbox( label="Negative Prompt", lines=3, value="worst quality, low quality, nsfw, logo")
-
-                with gr.Accordion("Advance", open=False):
-                    with gr.Row():
-                        width_slider  = gr.Slider(  label="Width",  value=512, minimum=256, maximum=1024, step=64 )
-                        height_slider = gr.Slider(  label="Height", value=512, minimum=256, maximum=1024, step=64 )
-                    with gr.Row():
-                        seed_textbox = gr.Textbox( label="Seed (-1 means random)",  value=-1)
-                        seed_button  = gr.Button(value="\U0001F3B2", elem_classes="toolbutton")
-                        seed_button.click(fn=lambda: gr.Textbox(value=random.randint(1, 1e16)), inputs=[], outputs=[seed_textbox])
-
-                generate_button = gr.Button( value="Generate", variant='primary' )
+@spaces.GPU
+def magictime_interface(
+    dreambooth_dropdown,
+    motion_module_dropdown,
+    prompt_textbox,
+    negative_prompt_textbox,
+    width_slider,
+    height_slider,
+    seed_textbox,
+):
+    return controller.magictime(
+        dreambooth_dropdown,
+        motion_module_dropdown,
+        prompt_textbox,
+        negative_prompt_textbox,
+        width_slider,
+        height_slider,
+        seed_textbox,
+    )
 
-            with gr.Column():
-                result_video = gr.Video( label="Generated Animation", interactive=False )
-                json_config  = gr.Json( label="Config", value=None )
+inputs = [
+    gr.Dropdown(label="DreamBooth Model", choices=controller.dreambooth_list, value=controller.dreambooth_list[0]),
+    gr.Dropdown(label="Motion Module", choices=controller.motion_module_list, value=controller.motion_module_list[0]),
+    gr.Textbox(label="Prompt", lines=3),
+    gr.Textbox(label="Negative Prompt", lines=3, value="worst quality, low quality, nsfw, logo"),
+    gr.Slider(label="Width", value=512, minimum=256, maximum=1024, step=64),
+    gr.Slider(label="Height", value=512, minimum=256, maximum=1024, step=64),
+    gr.Textbox(label="Seed", value="-1"),
+]
 
-            inputs  = [dreambooth_dropdown, motion_module_dropdown, prompt_textbox, negative_prompt_textbox, width_slider, height_slider, seed_textbox]
-            outputs = [result_video, json_config]
-            
-            generate_button.click( fn=controller.magictime, inputs=inputs, outputs=outputs )
-        
-        gr.Markdown(
-            """
-            <h5 style="text-align:left;">Warning: It is worth noting that even if we use the same seed and prompt but we change a machine, the results will be different. If you find a better seed and prompt, please tell me in a GitHub issue.</h5>
-            """
-        )        
-        gr.Examples( fn=controller.magictime, examples=examples, inputs=inputs, outputs=outputs, cache_examples=True )
-        
-    return demo
+outputs = [
+    gr.Video(label="Generated Animation"),
+    gr.Json(label="Config")
+]
 
+iface = gr.Interface(
+    fn=magictime_interface, 
+    inputs=inputs, 
+    outputs=outputs, 
+    title="MagicTime Controller",
+    examples=examples
+)
 
 if __name__ == "__main__":
-    demo = ui()
-    demo.queue(max_size=20)
-    demo.launch()
+    iface.launch()