Update opensora/serve/gradio_web_server.py

opensora/serve/gradio_web_server.py

@@ -22,32 +22,38 @@ from opensora.models.ae import ae_stride_config, getae, getae_wrapper
 from opensora.models.ae.videobase import CausalVQVAEModelWrapper, CausalVAEModelWrapper
 from opensora.models.diffusion.latte.modeling_latte import LatteT2V
 from opensora.sample.pipeline_videogen import VideoGenPipeline
-from opensora.serve.gradio_utils import block_css, title_markdown, randomize_seed_fn, set_env, examples, DESCRIPTION
-
+from opensora.serve.gradio_utils import block_css, title_markdown, randomize_seed_fn, set_env, DESCRIPTION
+from opensora.serve.gradio_utils import examples_txt, examples
 
 @spaces.GPU(duration=300)
+@torch.inference_mode()
+
 @torch.inference_mode()
 def generate_img(prompt, sample_steps, scale, seed=0, randomize_seed=False, force_images=False):
-    seed = int(randomize_seed_fn(seed, randomize_seed))
-    set_env(seed)
     video_length = transformer_model.config.video_length if not force_images else 1
     height, width = int(args.version.split('x')[1]), int(args.version.split('x')[2])
     num_frames = 1 if video_length == 1 else int(args.version.split('x')[0])
-    videos = videogen_pipeline(prompt,
-                               num_frames=num_frames,
-                               height=height,
-                               width=width,
-                               num_inference_steps=sample_steps,
-                               guidance_scale=scale,
-                               enable_temporal_attentions=not force_images,
-                               num_images_per_prompt=1,
-                               mask_feature=True,
-                               ).video
+    if not force_images and prompt in examples_txt:
+        idx = examples_txt.index(prompt)
+        tmp_save_path = f'demo65-221/f65/{idx+1}.mp4'
+    else:
+        seed = int(randomize_seed_fn(seed, randomize_seed))
+        set_env(seed)
+        videos = videogen_pipeline(prompt,
+                                num_frames=num_frames,
+                                height=height,
+                                width=width,
+                                num_inference_steps=sample_steps,
+                                guidance_scale=scale,
+                                enable_temporal_attentions=not force_images,
+                                num_images_per_prompt=1,
+                                mask_feature=True,
+                                ).video
 
-    torch.cuda.empty_cache()
-    videos = videos[0]
-    tmp_save_path = 'tmp.mp4'
-    imageio.mimwrite(tmp_save_path, videos, fps=24, quality=9)  # highest quality is 10, lowest is 0
+        torch.cuda.empty_cache()
+        videos = videos[0]
+        tmp_save_path = 'tmp.mp4'
+        imageio.mimwrite(tmp_save_path, videos, fps=24, quality=9)  # highest quality is 10, lowest is 0
     display_model_info = f"Video size: {num_frames}×{height}×{width}, \nSampling Step: {sample_steps}, \nGuidance Scale: {scale}"
     return tmp_save_path, prompt, display_model_info, seed