Update base/app.py

base/app.py

@@ -16,21 +16,12 @@ from diffusers.schedulers import DDIMScheduler, DDPMScheduler, PNDMScheduler, Eu
 from diffusers.models import AutoencoderKL
 from transformers import CLIPTokenizer, CLIPTextModel, CLIPTextModelWithProjection
 
+SECRET_TOKEN = os.getenv('SECRET_TOKEN', 'default_secret')
 
 config_path = "./base/configs/sample.yaml"
 args = OmegaConf.load("./base/configs/sample.yaml")
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
-css = """
-h1 {
-  text-align: center;
-}
-#component-0 {
-  max-width: 730px;
-  margin: auto;
-}
-"""
-
 sd_path = args.pretrained_path
 unet = get_models(args, sd_path).to(device, dtype=torch.float16)
 state_dict = find_model("./pretrained_models/lavie_base.pt")
@@ -42,7 +33,10 @@ unet.eval()
 vae.eval()
 text_encoder_one.eval()
 
-def infer(prompt, seed_inp, ddim_steps,cfg, infer_type):
+def infer(secret_token, prompt, seed_inp, ddim_steps,cfg, infer_type):
+    if secret_token != SECRET_TOKEN:
+        raise gr.Error(f'Invalid secret token. Please fork the original space if you want to use it for yourself.')
+       
     if seed_inp!=-1:
         setup_seed(seed_inp)
     else:
@@ -79,38 +73,21 @@ def infer(prompt, seed_inp, ddim_steps,cfg, infer_type):
     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;">
-        <div
-        style="
-            display: inline-flex;
-            align-items: center;
-            gap: 0.8rem;
-            font-size: 1.75rem;
-        "
-        >
-        <h1 style="font-weight: 900; margin-bottom: 7px; margin-top: 5px;">
-            Intern·Vchitect (Text-to-Video)
-        </h1>
-        </div>
-        <p style="margin-bottom: 10px; font-size: 94%">
-        Apply Intern·Vchitect to generate a video 
-        </p>
-    </div>
-"""
-
-with gr.Blocks(css='style.css') as demo:
-    gr.Markdown("<font color=red size=10><center>LaVie: Text-to-Video generation</center></font>")
-    gr.Markdown(
-        """<div style="text-align:center">
-        [<a href="https://arxiv.org/abs/2309.15103">Arxiv Report</a>] | [<a href="https://vchitect.github.io/LaVie-project/">Project Page</a>] | [<a href="https://github.com/Vchitect/LaVie">Github</a>]</div>
-        """
-    )
+
+with gr.Blocks() as demo:
+    gr.HTML("""
+            <div style="z-index: 100; position: fixed; top: 0px; right: 0px; left: 0px; bottom: 0px; width: 100%; height: 100%; background: white; display: flex; align-items: center; justify-content: center; color: black;">
+              <div style="text-align: center; color: black;">
+                <p style="color: black;">This space is a REST API to programmatically generate MP4 videos.</p>
+                <p style="color: black;">Interested in using it? Look no further than the <a href="https://huggingface.co/spaces/Vchitect/LaVie" target="_blank">original space</a>!</p>
+              </div>
+        </div>""")
+    secret_token = gr.Textbox(label="Secret token")
     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)
+                prompt = gr.Textbox(value="", 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')
                 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)
@@ -120,29 +97,9 @@ with gr.Blocks(css='style.css') as demo:
                 submit_btn = gr.Button("Generate video")
                 video_out = gr.Video(label="Video result", elem_id="video-output")
 
-            inputs = [prompt, seed_inp, ddim_steps, cfg, infer_type]
+            inputs = [secret_token, prompt, seed_inp, ddim_steps, cfg, infer_type]
             outputs = [video_out]
 
-        ex = gr.Examples(
-            examples = [['a corgi walking in the park at sunrise, oil painting style',400,50,7,'ddim'],
-                    ['a cut 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=False,
-        )
-        ex.dataset.headers = [""]
-         
     submit_btn.click(infer, inputs, outputs)
 
 demo.queue(max_size=12).launch()