Update app.py

app.py

@@ -2,55 +2,31 @@ import os
 import uuid
 from omegaconf import OmegaConf
 import spaces
-
 import random
-
 import imageio
 import torch
 import torchvision
 import gradio as gr
 import numpy as np
-
+from fastapi import FastAPI
+from fastapi.responses import FileResponse
 from gradio.components import Textbox, Video
 from huggingface_hub import hf_hub_download
-
 from utils.common_utils import load_model_checkpoint
 from utils.utils import instantiate_from_config
 from scheduler.t2v_turbo_scheduler import T2VTurboScheduler
 from pipeline.t2v_turbo_vc2_pipeline import T2VTurboVC2Pipeline
 
-DESCRIPTION = """# T2V-Turbo 🚀
-
-Our model is distilled from [VideoCrafter2](https://ailab-cvc.github.io/videocrafter2/).
-
-T2V-Turbo learns a LoRA on top of the base model by aligning to the reward feedback from [HPSv2.1](https://github.com/tgxs002/HPSv2/tree/master) and [InternVid2 Stage 2 Model](https://huggingface.co/OpenGVLab/InternVideo2-Stage2_1B-224p-f4).
-
-T2V-Turbo-v2 optimizes the training techniques by finetuning the full base model and further aligns to [CLIPScore](https://huggingface.co/laion/CLIP-ViT-H-14-laion2B-s32B-b79K)
-
-T2V-Turbo trains on pure WebVid-10M data, whereas T2V-Turbo-v2 carufully optimizes different learning objectives with a mixutre of VidGen-1M and WebVid-10M data.
-
-Moreover, T2V-Turbo-v2 supports to distill motion priors from the training videos. 
-
-[Project page for T2V-Turbo](https://t2v-turbo.github.io) 🥳
-
-[Project page for T2V-Turbo-v2](https://t2v-turbo-v2.github.io) 🤓
-"""
-if torch.cuda.is_available():
-    DESCRIPTION += "\n<p>Running on CUDA 😀</p>"
-elif hasattr(torch, "xpu") and torch.xpu.is_available():
-    DESCRIPTION += "\n<p>Running on XPU 🤓</p>"
-else:
-    DESCRIPTION += "\n<p>Running on CPU 🥶 This demo does not work on CPU.</p>"
+# Keep all your original constants and DESCRIPTION
 
 MAX_SEED = np.iinfo(np.int32).max
-
+app = FastAPI()
 
 def randomize_seed_fn(seed: int, randomize_seed: bool) -> int:
     if randomize_seed:
         seed = random.randint(0, MAX_SEED)
     return seed
 
-
 def save_video(video_array, video_save_path, fps: int = 16):
     video = video_array.detach().cpu()
     video = torch.clamp(video.float(), -1.0, 1.0)
@@ -62,17 +38,7 @@ def save_video(video_array, video_save_path, fps: int = 16):
         video_save_path, video, fps=fps, video_codec="h264", options={"crf": "10"}
     )
 
-example_txt = [
-    "An astronaut riding a horse.",
-    "Darth vader surfing in waves.",
-    "light wind, feathers moving, she moves her gaze, 4k",
-    "a girl floating underwater.",
-    "Pikachu snowboarding.",
-    "Self-portrait oil painting, a beautiful cyborg with golden hair, 8k",
-    "A musician strums his guitar, serenading the moonlit night.",
-]
-
-examples = [[i, 7.5, 0.5, 16, 16, 0, True, "bf16"] for i in example_txt]
+# Keep your original example_txt and examples
 
 @spaces.GPU(duration=120)
 @torch.inference_mode()
@@ -87,8 +53,8 @@ def generate(
     param_dtype="bf16",
     motion_gs: float = 0.05,
     fps: int = 8,
+    is_api: bool = False,  # New parameter to handle API calls
 ):
-
     seed = randomize_seed_fn(seed, randomize_seed)
     torch.manual_seed(seed)
 
@@ -123,26 +89,70 @@ def generate(
     )
 
     torch.cuda.empty_cache()
-    tmp_save_path = "tmp.mp4"
+    
+    # Generate unique filename for API calls
+    if is_api:
+        video_filename = f"{uuid.uuid4()}.mp4"
+    else:
+        video_filename = "tmp.mp4"
+        
     root_path = "./videos/"
     os.makedirs(root_path, exist_ok=True)
-    video_save_path = os.path.join(root_path, tmp_save_path)
+    video_save_path = os.path.join(root_path, video_filename)
 
     save_video(result[0], video_save_path, fps=fps)
     display_model_info = f"Video size: {num_frames}x320x512, Sampling Step: {num_inference_steps}, Guidance Scale: {guidance_scale}"
+    
+    if is_api:
+        return {
+            "video_path": video_save_path,
+            "prompt": prompt,
+            "model_info": display_model_info,
+            "seed": seed
+        }
     return video_save_path, prompt, display_model_info, seed
 
-
-block_css = """
-#buttons button {
-    min-width: min(120px,100%);
-}
-"""
-
+# API endpoint
+@app.post("/generate")
+async def generate_api(
+    prompt: str,
+    guidance_scale: float = 7.5,
+    percentage: float = 0.5,
+    num_inference_steps: int = 4,
+    num_frames: int = 16,
+    seed: int = 0,
+    randomize_seed: bool = False,
+    param_dtype: str = "bf16",
+    motion_gs: float = 0.05,
+    fps: int = 8,
+):
+    result = generate(
+        prompt=prompt,
+        guidance_scale=guidance_scale,
+        percentage=percentage,
+        num_inference_steps=num_inference_steps,
+        num_frames=num_frames,
+        seed=seed,
+        randomize_seed=randomize_seed,
+        param_dtype=param_dtype,
+        motion_gs=motion_gs,
+        fps=fps,
+        is_api=True
+    )
+    
+    return FileResponse(
+        result["video_path"],
+        media_type="video/mp4",
+        headers={
+            "X-Model-Info": result["model_info"],
+            "X-Seed": str(result["seed"])
+        }
+    )
 
 if __name__ == "__main__":
     device = torch.device("cuda:0")
 
+    # Keep all your original model initialization code
     config = OmegaConf.load("configs/inference_t2v_512_v2.0.yaml")
     model_config = config.pop("model", OmegaConf.create())
     pretrained_t2v = instantiate_from_config(model_config)
@@ -169,54 +179,18 @@ if __name__ == "__main__":
     pipeline = T2VTurboVC2Pipeline(pretrained_t2v, scheduler, model_config)
     pipeline.to(device)
 
+    # Mount both Gradio and FastAPI
     demo = gr.Interface(
-        fn=generate,
+        fn=lambda *args: generate(*args, is_api=False),
         inputs=[
             Textbox(label="", placeholder="Please enter your prompt. \n"),
-            gr.Slider(
-                label="Guidance scale",
-                minimum=2,
-                maximum=14,
-                step=0.1,
-                value=7.5,
-            ),
-            gr.Slider(
-                label="Percentage of steps to apply motion guidance (v2 w/ MG only)",
-                minimum=0.0,
-                maximum=0.5,
-                step=0.05,
-                value=0.5,
-            ),
-            gr.Slider(
-                label="Number of inference steps",
-                minimum=4,
-                maximum=50,
-                step=1,
-                value=16,
-            ),
-            gr.Slider(
-                label="Number of Video Frames",
-                minimum=16,
-                maximum=48,
-                step=8,
-                value=16,
-            ),
-            gr.Slider(
-                label="Seed",
-                minimum=0,
-                maximum=MAX_SEED,
-                step=1,
-                value=0,
-                randomize=True,
-            ),
+            gr.Slider(label="Guidance scale", minimum=2, maximum=14, step=0.1, value=7.5),
+            gr.Slider(label="Percentage of steps to apply motion guidance", minimum=0.0, maximum=0.5, step=0.05, value=0.5),
+            gr.Slider(label="Number of inference steps", minimum=4, maximum=50, step=1, value=16),
+            gr.Slider(label="Number of Video Frames", minimum=16, maximum=48, step=8, value=16),
+            gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=0, randomize=True),
             gr.Checkbox(label="Randomize seed", value=True),
-            gr.Radio(
-                ["bf16", "fp16", "fp32"],
-                label="torch.dtype",
-                value="bf16",
-                interactive=True,
-                info="Dtype for inference. Default is bf16.",
-            )
+            gr.Radio(["bf16", "fp16", "fp32"], label="torch.dtype", value="bf16", interactive=True),
         ],
         outputs=[
             gr.Video(label="Generated Video", width=512, height=320, interactive=False, autoplay=True),
@@ -231,4 +205,9 @@ if __name__ == "__main__":
         cache_examples=False,
         concurrency_limit=10,
     )
-    demo.launch()
+    
+    app = gr.mount_gradio_app(app, demo, path="/")
+    
+    # Run both servers
+    import uvicorn
+    uvicorn.run(app, host="0.0.0.0", port=7860)