rename

app.py

@@ -1,19 +1,11 @@
-import time
 import os
-from huggingface_hub import hf_hub_download, login
+import time
 
 import gradio as gr
-
-# Authenticate with Hugging Face (token is auto-available in HF Spaces)
-try:
-    login(token=os.environ.get("HF_TOKEN"))
-    print("Successfully authenticated with Hugging Face")
-except Exception as e:
-    print(f"Warning: Could not authenticate with HF: {e}")
-import spaces
 import torch
 from diffusers import AutoencoderKLCogVideoX, CogVideoXDDIMScheduler
 from diffusers.utils import export_to_video
+from huggingface_hub import hf_hub_download, login
 from PIL import Image
 from transformers import T5EncoderModel, T5Tokenizer
 
@@ -21,22 +13,23 @@ from cogvideo_transformer import CustomCogVideoXTransformer3DModel
 from EF_Net import EF_Net
 from Sci_Fi_inbetweening_pipeline import CogVideoXEFNetInbetweeningPipeline
 
-# Global variables for the pipeline
-pipe = None
+# Authenticate with Hugging Face
+try:
+    login(token=os.environ.get("HF_TOKEN"))
+    print("Successfully authenticated with Hugging Face")
+except Exception as e:
+    print(f"Warning: Could not authenticate with HF: {e}")
+
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
 
-def _load_pipeline_internal(
+def load_pipeline(
     pretrained_model_path="LiuhanChen/Sci-Fi",
     ef_net_path="weights/EF_Net.pth",
     dtype_str="bfloat16",
 ):
-    """Internal function to load the Sci-Fi pipeline"""
-    global pipe
-
-    # Return early if pipeline is already loaded
-    if pipe is not None:
-        return "Pipeline already loaded!"
+    """Load the Sci-Fi pipeline at startup"""
+    print("Loading Sci-Fi pipeline...")
 
     dtype = torch.float16 if dtype_str == "float16" else torch.bfloat16
 
@@ -48,29 +41,37 @@ def _load_pipeline_internal(
             repo_id="LiuhanChen/Sci-Fi",
             subfolder="EF_Net",
             filename="EF_Net.pth",
-            local_dir="weights"
+            local_dir="weights",
         )
         ef_net_path = "weights/EF_Net/EF_Net.pth"
         print(f"EF-Net weights downloaded to {ef_net_path}")
 
     # Load models from Hugging Face
+    print("Loading tokenizer and text encoder...")
     tokenizer = T5Tokenizer.from_pretrained(
         pretrained_model_path, subfolder="CogVideoX-5b-I2V/tokenizer"
     )
     text_encoder = T5EncoderModel.from_pretrained(
         pretrained_model_path, subfolder="CogVideoX-5b-I2V/text_encoder"
     )
+
+    print("Loading transformer...")
     transformer = CustomCogVideoXTransformer3DModel.from_pretrained(
         pretrained_model_path, subfolder="CogVideoX-5b-I2V/transformer"
     )
+
+    print("Loading VAE...")
     vae = AutoencoderKLCogVideoX.from_pretrained(
         pretrained_model_path, subfolder="CogVideoX-5b-I2V/vae"
     )
+
+    print("Loading scheduler...")
     scheduler = CogVideoXDDIMScheduler.from_pretrained(
         pretrained_model_path, subfolder="CogVideoX-5b-I2V/scheduler"
     )
 
     # Load EF-Net
+    print("Loading EF-Net...")
     EF_Net_model = (
         EF_Net(num_layers=4, downscale_coef=8, in_channels=2, num_attention_heads=48)
         .requires_grad_(False)
@@ -83,7 +84,8 @@ def _load_pipeline_internal(
     print(f"[EF-Net loaded] Missing: {len(m)} | Unexpected: {len(u)}")
 
     # Create pipeline
-    pipe = CogVideoXEFNetInbetweeningPipeline(
+    print("Creating pipeline...")
+    pipeline = CogVideoXEFNetInbetweeningPipeline(
         tokenizer=tokenizer,
         text_encoder=text_encoder,
         transformer=transformer,
@@ -91,20 +93,26 @@ def _load_pipeline_internal(
         EF_Net_model=EF_Net_model,
         scheduler=scheduler,
     )
-    pipe.scheduler = CogVideoXDDIMScheduler.from_config(
-        pipe.scheduler.config, timestep_spacing="trailing"
+    pipeline.scheduler = CogVideoXDDIMScheduler.from_config(
+        pipeline.scheduler.config, timestep_spacing="trailing"
     )
 
-    pipe.to(device)
-    pipe = pipe.to(dtype=dtype)
+    print(f"Moving pipeline to {device}...")
+    pipeline.to(device)
+    pipeline = pipeline.to(dtype=dtype)
+
+    pipeline.vae.enable_slicing()
+    pipeline.vae.enable_tiling()
 
-    pipe.vae.enable_slicing()
-    pipe.vae.enable_tiling()
+    print("Pipeline loaded successfully!")
+    return pipeline
 
-    return "Pipeline loaded successfully!"
+
+# Load pipeline at startup
+print("Initializing Sci-Fi pipeline at startup...")
+pipe = load_pipeline()
 
 
-@spaces.GPU(duration=500)
 def generate_inbetweening(
     first_image: Image.Image,
     last_image: Image.Image,
@@ -118,16 +126,6 @@ def generate_inbetweening(
     progress=gr.Progress(),
 ):
     """Generate frame inbetweening video"""
-    global pipe
-
-    # Load pipeline on first use (lazy loading with GPU access)
-    if pipe is None:
-        progress(0, desc="Loading pipeline (first run)...")
-        try:
-            _load_pipeline_internal()
-        except Exception as e:
-            return None, f"ERROR: Failed to load pipeline: {str(e)}"
-
     if first_image is None or last_image is None:
         return None, "Please upload both start and end frames!"
 
@@ -177,7 +175,7 @@ with gr.Blocks(title="Sci-Fi: Frame Inbetweening") as demo:
 
     Upload start and end frames to generate smooth inbetweening video.
 
-    **Note:** Pipeline loads on first generation (this may take 1-2 minutes).
+    **Model is pre-loaded and ready to use!**
     """
     )
 
@@ -280,5 +278,5 @@ with gr.Blocks(title="Sci-Fi: Frame Inbetweening") as demo:
         )
 
 if __name__ == "__main__":
-    print("App starting - pipeline will load on first generation request")
+    print("App ready - pipeline is loaded and ready for inference!")
     demo.launch()