Update app.py

app.py

@@ -2,7 +2,6 @@ import os
 import copy
 import torch
 import random
-import spaces
 import gradio as gr
 from glob import glob
 from omegaconf import OmegaConf
@@ -15,6 +14,7 @@ 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,6 +62,7 @@ examples = [
 print(f"### Cleaning cached examples ...")
 os.system(f"rm -rf gradio_cached_examples/")
 
+device = torch.device('cuda:0')
 
 class MagicTimeController:
     def __init__(self):
@@ -87,9 +88,9 @@ class MagicTimeController:
         self.inference_config      = OmegaConf.load(inference_config_path)[1]
 
         self.tokenizer             = CLIPTokenizer.from_pretrained(pretrained_model_path, subfolder="tokenizer")
-        self.text_encoder          = CLIPTextModel.from_pretrained(pretrained_model_path, subfolder="text_encoder").cuda()
-        self.vae                   = AutoencoderKL.from_pretrained(pretrained_model_path, subfolder="vae").cuda()
-        self.unet                  = UNet3DConditionModel.from_pretrained_2d(pretrained_model_path, subfolder="unet", unet_additional_kwargs=OmegaConf.to_container(self.inference_config.unet_additional_kwargs)).cuda()
+        self.text_encoder          = CLIPTextModel.from_pretrained(pretrained_model_path, subfolder="text_encoder").to(device)
+        self.vae                   = AutoencoderKL.from_pretrained(pretrained_model_path, subfolder="vae").to(device)
+        self.unet                  = UNet3DConditionModel.from_pretrained_2d(pretrained_model_path, subfolder="unet", unet_additional_kwargs=OmegaConf.to_container(self.inference_config.unet_additional_kwargs)).to(device)
         self.text_model            = CLIPTextModel.from_pretrained("openai/clip-vit-large-patch14")
         self.unet_model            = UNet3DConditionModel.from_pretrained_2d(pretrained_model_path, subfolder="unet", unet_additional_kwargs=OmegaConf.to_container(self.inference_config.unet_additional_kwargs))
     
@@ -153,7 +154,8 @@ class MagicTimeController:
         _, unexpected = self.unet_model.load_state_dict(motion_module_state_dict, strict=False)
         assert len(unexpected) == 0
         return gr.Dropdown()
-    
+
+    @spaces.GPU(duration=300)
     def magictime(
         self,
         dreambooth_dropdown,
@@ -173,7 +175,7 @@ class MagicTimeController:
         pipeline = MagicTimePipeline(
             vae=self.vae, text_encoder=self.text_encoder, tokenizer=self.tokenizer, unet=self.unet,
             scheduler=DDIMScheduler(**OmegaConf.to_container(self.inference_config.noise_scheduler_kwargs))
-        ).to("cuda")
+        ).to(device)
         
         if int(seed_textbox) > 0: seed = int(seed_textbox)
         else: seed = random.randint(1, 1e16)
@@ -182,7 +184,7 @@ class MagicTimeController:
         assert seed == torch.initial_seed()
         print(f"### seed: {seed}")
         
-        generator = torch.Generator(device="cuda")
+        generator = torch.Generator(device=device)
         generator.manual_seed(seed)
         
         sample = pipeline(
@@ -208,51 +210,72 @@ class MagicTimeController:
             "dreambooth": dreambooth_dropdown,
         }
         return gr.Video(value=save_sample_path), gr.Json(value=json_config)
-        
-controller = MagicTimeController()
 
-@spaces.GPU(duration=300)
-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,
-    )
+# 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()
 
-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"),
-]
+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 )
 
-outputs = [
-    gr.Video(label="Generated Animation"),
-    gr.Json(label="Config")
-]
+                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' )
+
+            with gr.Column():
+                result_video = gr.Video( label="Generated Animation", interactive=False )
+                json_config  = gr.Json( label="Config", value=None )
+
+            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
 
-iface = gr.Interface(
-    fn=magictime_interface, 
-    inputs=inputs, 
-    outputs=outputs, 
-    title="MagicTime Controller",
-    examples=examples
-)
 
 if __name__ == "__main__":
-    iface.launch()
+    demo = ui()
+    demo.queue(max_size=20)
+    demo.launch()