Update app.py

app.py

@@ -1,47 +1,73 @@
 import streamlit as st
 from transformers import pipeline
 from PIL import Image
-import numpy as np
-import tempfile
 import os
+import pathlib
+from huggingface_hub import snapshot_download
 from modelscope.pipelines import pipeline as modelscope_pipeline
 from modelscope.outputs import OutputKeys
+import shutil
 
-def generate_video_from_image(image, duration_seconds=10, progress_bar=None):
-    """
-    Generate a video from an image using ModelScope's video generation.
-    """
+# Create a downloads directory if it doesn't exist
+if not os.path.exists('downloads'):
+    os.makedirs('downloads')
+
+def initialize_models():
+    """Initialize and cache the models to avoid reloading."""
+    if 'caption_pipeline' not in st.session_state:
+        st.session_state.caption_pipeline = pipeline("image-to-text", model="Salesforce/blip-image-captioning-base")
+    
+    if 'video_pipeline' not in st.session_state:
+        # Download and cache the model
+        model_dir = pathlib.Path('weights')
+        snapshot_download(
+            'damo-vilab/modelscope-damo-text-to-video-synthesis',
+            repo_type='model', 
+            local_dir=model_dir
+        )
+        st.session_state.video_pipeline = modelscope_pipeline(
+            'text-to-video-synthesis', 
+            model_dir.as_posix()
+        )
+
+def save_video(video_path, caption):
+    """Save video to downloads directory with a meaningful name."""
+    # Create a filename from the caption
+    safe_caption = "".join(x for x in caption[:30] if x.isalnum() or x in (' ','-','_')).strip()
+    save_name = f"video_{safe_caption}.mp4"
+    save_path = os.path.join('downloads', save_name)
+    
+    # Copy the video file
+    shutil.copy2(video_path, save_path)
+    return save_path
+
+def generate_video_from_image(image, progress_bar=None):
+    """Generate a video based on image caption using ModelScope's text-to-video model."""
     try:
         if progress_bar:
             progress_bar.progress(0.1, "Generating image caption...")
-            
-        # Setup image captioning
-        caption_pipe = pipeline("image-to-text", model="Salesforce/blip-image-captioning-base")
         
         # Generate caption
-        caption = caption_pipe(image)[0]['generated_text']
+        caption = st.session_state.caption_pipeline(image)[0]['generated_text']
         st.write(f"Generated caption: *{caption}*")
         
         if progress_bar:
-            progress_bar.progress(0.3, "Loading Video Generation model...")
-            
-        # Initialize video generation
-        video_pipe = modelscope_pipeline(
-            'text-to-video-synthesis',
-            model='damo/text-to-video-synthesis'
-        )
+            progress_bar.progress(0.3, "Generating video...")
+        
+        # Prepare input for video generation
+        input_text = {'text': caption}
         
-        if progress_bar:
-            progress_bar.progress(0.5, "Generating video...")
-            
         # Generate video
-        output = video_pipe(caption)
+        output = st.session_state.video_pipeline(input_text)
         video_path = output[OutputKeys.OUTPUT_VIDEO]
         
+        # Save video with meaningful name
+        final_path = save_video(video_path, caption)
+        
         if progress_bar:
             progress_bar.progress(1.0, "Video generation complete!")
-            
-        return video_path, caption
+        
+        return final_path, caption
         
     except Exception as e:
         st.error(f"Error generating video: {str(e)}")
@@ -50,13 +76,26 @@ def generate_video_from_image(image, duration_seconds=10, progress_bar=None):
 def main():
     st.set_page_config(page_title="AI Video Generator", page_icon="🎥")
     
-    st.title("🎥 AI Video Generator")
+    st.title("🎥 Text-to-Video Generator")
     st.write("""
-    Upload an image to generate a video with AI-powered motion and transitions.
-    The app will automatically generate a caption for your image and use it as inspiration for the video.
+    Upload an image to generate a video based on its content. The app will:
+    1. Generate a caption for your image
+    2. Create a video based on that caption
+    3. Provide options to view and download the video
+    """)
+    
+    # Display model limitations
+    st.warning("""
+    Model Limitations:
+    - Only English text is supported
+    - Cannot generate clear text in videos
+    - May have limitations with complex scenes
+    - Generation takes several minutes
     """)
     
-    st.info("Note: Video generation may take several minutes.")
+    # Initialize models
+    with st.spinner("Loading models... This may take a minute..."):
+        initialize_models()
     
     # File uploader
     uploaded_file = st.file_uploader("Choose an image", type=['png', 'jpg', 'jpeg'])
@@ -81,23 +120,47 @@ def main():
                     with open(video_path, 'rb') as video_file:
                         video_bytes = video_file.read()
                     
-                    # Create download button
-                    st.download_button(
-                        label="Download Video",
-                        data=video_bytes,
-                        file_name="generated_video.mp4",
-                        mime="video/mp4"
-                    )
+                    # Create a container for the video and download options
+                    st.success("Video generated successfully!")
+                    
+                    col1, col2 = st.columns(2)
+                    
+                    with col1:
+                        # Primary download button
+                        st.download_button(
+                            label="💾 Download Video",
+                            data=video_bytes,
+                            file_name=os.path.basename(video_path),
+                            mime="video/mp4",
+                            key="download1"
+                        )
+                    
+                    with col2:
+                        # Additional download button with caption
+                        st.download_button(
+                            label="📥 Download with Caption",
+                            data=video_bytes,
+                            file_name=f"{caption[:30]}.mp4",
+                            mime="video/mp4",
+                            key="download2"
+                        )
                     
                     # Display video
                     st.video(video_bytes)
+                    
+                    # Display additional information
+                    st.info(f"""
+                    Video Details:
+                    - Caption: {caption}
+                    - Filename: {os.path.basename(video_path)}
+                    - Size: {len(video_bytes)/1024/1024:.1f} MB
+                    """)
+                    
                 else:
                     st.error("Failed to generate video. Please try again.")
             
             except Exception as e:
                 st.error(f"An error occurred: {str(e)}")
-                st.error("Full error message for debugging:")
-                st.error(e)
 
 if __name__ == "__main__":
     main()