optimize

app.py

@@ -1,56 +1,67 @@
-import torch
-import gradio as gr
-from diffusers import DiffusionPipeline
+import torch  
+import gradio as gr  
+from diffusers import DiffusionPipeline  
+import time  # For monitoring time  
 
-def load_video_model():
-    try:
-        # Ensure all necessary libraries are imported
-        import sentencepiece
+def load_video_model():  
+    try:  
+        # Ensure all necessary libraries are imported  
+        import sentencepiece  
         
-        # Load the model with specific error handling
-        pipe = DiffusionPipeline.from_pretrained("Lightricks/LTX-Video")
-        return pipe
-    except ImportError as e:
-        print(f"Dependency Error: {e}")
-        print("Please install required libraries: pip install sentencepiece diffusers transformers torch gradio imageio")
-        return None
-    except Exception as e:
-        print(f"Error loading model: {e}")
-        return None
+        # Load the model with specific error handling  
+        print("Loading model...")  
+        pipe = DiffusionPipeline.from_pretrained("Lightricks/LTX-Video")  
+        print("Model loaded successfully!")  
+        return pipe  
+    except ImportError as e:  
+        print(f"Dependency Error: {e}")  
+        print("Please install required libraries: pip install sentencepiece diffusers transformers torch gradio imageio")  
+        return None  
+    except Exception as e:  
+        print(f"Error loading model: {e}")  
+        return None  
 
-# Load the model when the script starts
-pipe = load_video_model()
+# Load the model when the script starts  
+pipe = load_video_model()  
 
-def generate_video(prompt):
-    if pipe is None:
-        return "Error: Model could not be loaded. Check your dependencies."
+def generate_video(prompt):  
+    if pipe is None:  
+        return "Error: Model could not be loaded. Check your dependencies."  
     
-    try:
-        # Generate video frames
-        images = pipe(prompt).images
+    try:  
+        print("Starting video generation...")  
+        start_time = time.time()  
         
-        # Save the generated images as a video
-        output_path = "generated_video.mp4"
+        # Generate video frames with a limit of 24 frames (approx. 5 seconds at 5 FPS)  
+        images = pipe(prompt, num_inference_steps=24).images  
+        print("Video frames generated successfully!")  
         
-        # Convert images to video
-        import imageio
-        imageio.mimsave(output_path, images, fps=5)
+        # Save the generated images as a video  
+        output_path = "generated_video.mp4"  
         
-        return output_path
-    except Exception as e:
-        return f"Error generating video: {str(e)}"
+        # Convert images to video  
+        import imageio  
+        imageio.mimsave(output_path, images, fps=5)  
+        
+        end_time = time.time()  
+        print(f"Video saved to {output_path} in {end_time - start_time:.2f} seconds")  
+        
+        return output_path  
+    except Exception as e:  
+        print(f"Error during video generation: {str(e)}")  
+        return f"Error generating video: {str(e)}"  
 
-# Create Gradio interface
-demo = gr.Interface(
-    fn=generate_video,
-    inputs=gr.Textbox(label="Enter Video Generation Prompt"),
-    outputs=gr.Video(label="Generated Video"),
-    title="LTX-Video Generation",
-    description="Generate a video using Lightricks Video Diffusion Model"
-)
+# Create Gradio interface  
+demo = gr.Interface(  
+    fn=generate_video,  
+    inputs=gr.Textbox(label="Enter Video Generation Prompt"),  
+    outputs=gr.Video(label="Generated Video"),  
+    title="LTX-Video Generation",  
+    description="Generate a video using Lightricks Video Diffusion Model"  
+)  
 
-if __name__ == "__main__":
-    if pipe is not None:
-        demo.launch()
-    else:
+if __name__ == "__main__":  
+    if pipe is not None:  
+        demo.launch(timeout=300)  # Increase timeout to 300 seconds (5 minutes)  
+    else:  
         print("Could not launch app due to model loading failure.")