Refactor CogVideoX app with modular functions and improved code structure

app.py

@@ -5,7 +5,8 @@ import torch
 import tempfile
 import os
 import spaces
-# Lista de modelos disponibles
+
+# Available transformer models
 TRANSFORMER_MODELS = [
     "sayakpaul/pika-dissolve-v0",
     "finetrainers/crush-smol-v0", 
@@ -13,24 +14,42 @@ TRANSFORMER_MODELS = [
     "finetrainers/cakeify-v0"
 ]
 
-@spaces.GPU
-def generate_video(transformer_model, prompt, negative_prompt, num_frames, height, width, num_inference_steps):
-    # Cargar el modelo del transformer seleccionado
+def load_models(transformer_model):
+    """Load transformer and pipeline models"""
+    # Load the selected transformer model
+    print(f"Loading model: {transformer_model}")
     transformer = CogVideoXTransformer3DModel.from_pretrained(
         transformer_model, 
         torch_dtype=torch.bfloat16
     )
-
-    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
     
-    # Inicializar el pipeline
+    # Initialize the pipeline
+    print("Initializing pipeline")
     pipeline = DiffusionPipeline.from_pretrained(
         "THUDM/CogVideoX-5b", 
         transformer=transformer, 
         torch_dtype=torch.bfloat16
-    ).to(device)
+    )
+    
+    return pipeline
+
+def save_video(video_frames, fps=25):
+    """Save video frames to a temporary file"""
+    print("Saving video")
+    with tempfile.NamedTemporaryFile(suffix='.mp4', delete=False) as tmp_file:
+        export_to_video(video_frames, tmp_file.name, fps=fps)
+        return tmp_file.name
+
+@spaces.GPU
+def generate_video_pipeline(pipeline, prompt, negative_prompt, num_frames, height, width, num_inference_steps):
+    """Generate video using the pipeline"""
+    # Move to appropriate device
+    print("Moving to device")
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    pipeline = pipeline.to(device)
     
-    # Generar el video
+    # Generate video
+    print("Generating video")
     video_frames = pipeline(
         prompt=prompt,
         negative_prompt=negative_prompt,
@@ -40,128 +59,151 @@ def generate_video(transformer_model, prompt, negative_prompt, num_frames, heigh
         num_inference_steps=num_inference_steps
     ).frames[0]
     
-    # Guardar el video en un archivo temporal
-    with tempfile.NamedTemporaryFile(suffix='.mp4', delete=False) as tmp_file:
-        export_to_video(video_frames, tmp_file.name, fps=25)
-        return tmp_file.name
+    print("Video generated")
+    return video_frames
 
-# Crear la interfaz de Gradio
-with gr.Blocks() as demo:
-    gr.Markdown("# CogVideoX Video Generator")
+def generate_video(transformer_model, prompt, negative_prompt, num_frames, height, width, num_inference_steps):
+    """Main function to handle the video generation process"""
+    # Load models
+    pipeline = load_models(transformer_model)
     
-    with gr.Row():
-        with gr.Column():
-            # Entradas
-            model_dropdown = gr.Dropdown(
-                choices=TRANSFORMER_MODELS,
-                value=TRANSFORMER_MODELS[0],
-                label="Transformer Model"
-            )
-            prompt_input = gr.Textbox(
-                lines=5,
-                label="Prompt",
-                placeholder="Describe the video you want to generate..."
-            )
-            negative_prompt_input = gr.Textbox(
-                lines=2,
-                label="Negative Prompt",
-                value="inconsistent motion, blurry motion, worse quality, degenerate outputs, deformed outputs"
-            )
-            
-            with gr.Accordion("Advanced Parameters", open=False):
-                num_frames = gr.Slider(
-                    minimum=8,
-                    maximum=128,
-                    value=50,
-                    step=1,
-                    label="Number of Frames",
-                    info="Number of frames in the video"
-                )
-                height = gr.Slider(
-                    minimum=32,
-                    maximum=1024,
-                    value=256,
-                    step=64,
-                    label="Height",
-                    info="Video height in pixels"
+    # Generate video frames
+    video_frames = generate_video_pipeline(
+        pipeline,
+        prompt,
+        negative_prompt,
+        num_frames,
+        height,
+        width,
+        num_inference_steps
+    )
+    
+    # Save and return video path
+    print("Saving video")
+    return save_video(video_frames)
+
+def create_interface():
+    """Create and configure the Gradio interface"""
+    with gr.Blocks() as demo:
+        gr.Markdown("# CogVideoX Video Generator")
+        
+        with gr.Row():
+            with gr.Column():
+                # Inputs
+                model_dropdown = gr.Dropdown(
+                    choices=TRANSFORMER_MODELS,
+                    value=TRANSFORMER_MODELS[0],
+                    label="Transformer Model"
                 )
-                width = gr.Slider(
-                    minimum=32,
-                    maximum=1024,
-                    value=256,
-                    step=64,
-                    label="Width",
-                    info="Video width in pixels"
+                prompt_input = gr.Textbox(
+                    lines=5,
+                    label="Prompt",
+                    placeholder="Describe the video you want to generate..."
                 )
-                num_inference_steps = gr.Slider(
-                    minimum=10,
-                    maximum=100,
-                    value=50,
-                    step=1,
-                    label="Inference Steps",
-                    info="Higher number = better quality but slower"
+                negative_prompt_input = gr.Textbox(
+                    lines=2,
+                    label="Negative Prompt",
+                    value="inconsistent motion, blurry motion, worse quality, degenerate outputs, deformed outputs"
                 )
+                
+                with gr.Accordion("Advanced Parameters", open=False):
+                    num_frames = gr.Slider(
+                        minimum=8,
+                        maximum=128,
+                        value=50,
+                        step=1,
+                        label="Number of Frames",
+                        info="Number of frames in the video"
+                    )
+                    height = gr.Slider(
+                        minimum=32,
+                        maximum=1024,
+                        value=224,
+                        step=64,
+                        label="Height",
+                        info="Video height in pixels"
+                    )
+                    width = gr.Slider(
+                        minimum=32,
+                        maximum=1024,
+                        value=224,
+                        step=64,
+                        label="Width",
+                        info="Video width in pixels"
+                    )
+                    num_inference_steps = gr.Slider(
+                        minimum=10,
+                        maximum=100,
+                        value=50,
+                        step=1,
+                        label="Inference Steps",
+                        info="Higher number = better quality but slower"
+                    )
+                
+                generate_btn = gr.Button("Generate Video")
             
-            generate_btn = gr.Button("Generate Video")
+            with gr.Column():
+                # Output
+                video_output = gr.Video(label="Generated Video")
         
-        with gr.Column():
-            # Salida
-            video_output = gr.Video(label="Generated Video")
-    
-    # Agregar ejemplos
-    gr.Examples(
-        examples=[
-            [
-                "sayakpaul/pika-dissolve-v0",
-                "PIKA_DISSOLVE A slender glass vase, brimming with tiny white pebbles, stands centered on a polished ebony dais. Without warning, the glass begins to dissolve from the edges inward. Wisps of translucent dust swirl upward in an elegant spiral, illuminating each pebble as they drop onto the dais. The gently drifting dust eventually settles, leaving only the scattered stones and faint traces of shimmering powder on the stage.",
-                "inconsistent motion, blurry motion, worse quality, degenerate outputs, deformed outputs",
-                8, 32, 32, 10
+        # Add examples
+        gr.Examples(
+            examples=[
+                [
+                    "sayakpaul/pika-dissolve-v0",
+                    "PIKA_DISSOLVE A slender glass vase, brimming with tiny white pebbles, stands centered on a polished ebony dais. Without warning, the glass begins to dissolve from the edges inward. Wisps of translucent dust swirl upward in an elegant spiral, illuminating each pebble as they drop onto the dais. The gently drifting dust eventually settles, leaving only the scattered stones and faint traces of shimmering powder on the stage.",
+                    "inconsistent motion, blurry motion, worse quality, degenerate outputs, deformed outputs",
+                    8, 32, 32, 10
+                ],
+                [
+                    "finetrainers/crush-smol-v0",
+                    "DIFF_crush A thick burger is placed on a dining table, and a large metal cylinder descends from above, crushing the burger as if it were under a hydraulic press. The bulb is crushed, leaving a pile of debris around it.",
+                    "inconsistent motion, blurry motion, worse quality, degenerate outputs, deformed outputs",
+                    8, 32, 32, 10
+                ],
+                [
+                    "finetrainers/3dgs-v0",
+                    "3D_dissolve In a 3D appearance, a bookshelf filled with books is surrounded by a burst of red sparks, creating a dramatic and explosive effect against a black background.",
+                    "inconsistent motion, blurry motion, worse quality, degenerate outputs, deformed outputs",
+                    8, 32, 32, 10
+                ],
+                [
+                    "finetrainers/cakeify-v0",
+                    "PIKA_CAKEIFY On a gleaming glass display stand, a sleek black purse quietly commands attention. Suddenly, a knife appears and slices through the shoe, revealing a fluffy vanilla sponge at its core. Immediately, it turns into a hyper-realistic prop cake, delighting the senses with its playful juxtaposition of the everyday and the extraordinary.",
+                    "inconsistent motion, blurry motion, worse quality, degenerate outputs, deformed outputs",
+                    8, 32, 32, 10
+                ]
             ],
-            [
-                "finetrainers/crush-smol-v0",
-                "DIFF_crush A thick burger is placed on a dining table, and a large metal cylinder descends from above, crushing the burger as if it were under a hydraulic press. The bulb is crushed, leaving a pile of debris around it.",
-                "inconsistent motion, blurry motion, worse quality, degenerate outputs, deformed outputs",
-                8, 32, 32, 10
+            inputs=[
+                model_dropdown,
+                prompt_input,
+                negative_prompt_input,
+                num_frames,
+                height,
+                width,
+                num_inference_steps
             ],
-            [
-                "finetrainers/3dgs-v0",
-                "3D_dissolve In a 3D appearance, a bookshelf filled with books is surrounded by a burst of red sparks, creating a dramatic and explosive effect against a black background.",
-                "inconsistent motion, blurry motion, worse quality, degenerate outputs, deformed outputs",
-                8, 32, 32, 10
-            ],
-            [
-                "finetrainers/cakeify-v0",
-                "PIKA_CAKEIFY On a gleaming glass display stand, a sleek black purse quietly commands attention. Suddenly, a knife appears and slices through the shoe, revealing a fluffy vanilla sponge at its core. Immediately, it turns into a hyper-realistic prop cake, delighting the senses with its playful juxtaposition of the everyday and the extraordinary.",
-                "inconsistent motion, blurry motion, worse quality, degenerate outputs, deformed outputs",
-                8, 32, 32, 10
-            ]
-        ],
-        inputs=[
-            model_dropdown,
-            prompt_input,
-            negative_prompt_input,
-            num_frames,
-            height,
-            width,
-            num_inference_steps
-        ],
-        label="Prompt Examples"
-    )
+            label="Prompt Examples"
+        )
 
-    # Conectar la función
-    generate_btn.click(
-        fn=generate_video,
-        inputs=[
-            model_dropdown, 
-            prompt_input, 
-            negative_prompt_input,
-            num_frames,
-            height,
-            width,
-            num_inference_steps
-        ],
-        outputs=video_output
-    )
+        # Connect the function
+        generate_btn.click(
+            fn=generate_video,
+            inputs=[
+                model_dropdown, 
+                prompt_input, 
+                negative_prompt_input,
+                num_frames,
+                height,
+                width,
+                num_inference_steps
+            ],
+            outputs=video_output
+        )
+    
+    return demo
 
-# Lanzar la aplicación
-demo.launch()
+# Launch the application
+if __name__ == "__main__":
+    demo = create_interface()
+    demo.launch()