Update app.py

app.py

@@ -1,140 +1,132 @@
-# app.py
 import streamlit as st
-from transformers import BlipProcessor, BlipForConditionalGeneration
-from diffusers import DiffusionPipeline
 import torch
-import cv2
-import numpy as np
+from transformers import pipeline
 from PIL import Image
+from diffusers import LTXVideoProcessor, LTXVideoPipeline
+import numpy as np
+from moviepy.editor import ImageSequenceClip
 import tempfile
 import os
 
-# Configure page
-st.set_page_config(
-    page_title="Video Generator",
-    page_icon="🎥",
-    layout="wide"
-)
-
-@st.cache_resource
-def load_models():
-    # Load text-to-video model
-    pipeline = DiffusionPipeline.from_pretrained(
-        "cerspense/zeroscope_v2_576w",
-        torch_dtype=torch.float16
-    )
-    if torch.cuda.is_available():
-        pipeline.to("cuda")
-    else:
-        pipeline.to("cpu")
-    
-    # Load image captioning model
-    blip = BlipForConditionalGeneration.from_pretrained("Salesforce/blip-image-captioning-base")
-    blip_processor = BlipProcessor.from_pretrained("Salesforce/blip-image-captioning-base")
-    
-    if torch.cuda.is_available():
-        blip.to("cuda")
-    else:
-        blip.to("cpu")
-    
-    return pipeline, blip, blip_processor
-
-def enhance_image(image):
-    # Convert PIL Image to numpy array
-    img_array = np.array(image)
-    
-    # Basic enhancement: Increase contrast and brightness
-    enhanced = cv2.convertScaleAbs(img_array, alpha=1.2, beta=10)
-    
-    return Image.fromarray(enhanced)
-
-def get_description(image, blip, blip_processor):
-    # Process image for BLIP
-    inputs = blip_processor(image, return_tensors="pt")
-    
-    if torch.cuda.is_available():
-        inputs = {k: v.to("cuda") for k, v in inputs.items()}
-    
-    # Generate caption
-    with torch.no_grad():
-        generated_ids = blip.generate(pixel_values=inputs["pixel_values"], max_length=50)
-        description = blip_processor.decode(generated_ids[0], skip_special_tokens=True)
-    
-    return description
-
-def generate_video(pipeline, description):
-    # Generate video frames
-    video_frames = pipeline(
-        description,
-        num_inference_steps=30,
-        num_frames=16
-    ).frames
-    
-    # Create temporary directory and file path
-    temp_dir = tempfile.mkdtemp()
-    temp_path = os.path.join(temp_dir, "output.mp4")
-    
-    # Convert frames to video
-    height, width = video_frames[0].shape[:2]
-    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
-    video_writer = cv2.VideoWriter(temp_path, fourcc, 8, (width, height))
-    
-    for frame in video_frames:
-        video_writer.write(cv2.cvtColor(frame, cv2.COLOR_RGB2BGR))
-    
-    video_writer.release()
-    
-    return temp_path
-
-def main():
-    st.title("🎥 AI Video Generator")
-    st.write("Upload an image to generate a video based on its content!")
+def generate_video_from_image(image, duration_seconds=10, progress_bar=None):
+    """
+    Generate a video from an image using LTX-Video and image captioning.
     
+    Args:
+        image: PIL Image object
+        duration_seconds: Duration of output video in seconds
+        progress_bar: Streamlit progress bar object
+    """
     try:
-        # Load models
-        pipeline, blip, blip_processor = load_models()
+        if progress_bar:
+            progress_bar.progress(0.1, "Generating image caption...")
+            
+        # Setup image captioning pipeline
+        captioner = pipeline("image-to-text", model="Salesforce/blip-image-captioning-base")
         
-        # File uploader
-        image_file = st.file_uploader("Upload Image", type=['png', 'jpg', 'jpeg'])
+        # Generate caption
+        caption = captioner(image)[0]['generated_text']
+        st.write(f"Generated caption: *{caption}*")
         
-        if image_file:
-            # Display original and enhanced image
-            col1, col2 = st.columns(2)
+        if progress_bar:
+            progress_bar.progress(0.3, "Loading LTX-Video model...")
             
-            with col1:
-                image = Image.open(image_file)
-                st.image(image, caption="Original Image")
+        # Initialize LTX-Video pipeline
+        processor = LTXVideoProcessor()
+        pipeline = LTXVideoPipeline.from_pretrained("Lightricks/ltx-video")
+        
+        if progress_bar:
+            progress_bar.progress(0.4, "Processing image...")
+            
+        # Process image for video generation
+        processed_image = processor(image).pixel_values
+        processed_image = torch.from_numpy(processed_image).unsqueeze(0)
+        
+        if progress_bar:
+            progress_bar.progress(0.5, "Generating video frames...")
             
-            with col2:
-                enhanced_image = enhance_image(image)
-                st.image(enhanced_image, caption="Enhanced Image")
+        # Generate video frames
+        num_frames = duration_seconds * 30  # 30 FPS
+        video_frames = pipeline(
+            processed_image,
+            num_inference_steps=50,
+            num_frames=num_frames,
+            guidance_scale=7.5,
+            prompt=caption,
+        ).videos
+        
+        if progress_bar:
+            progress_bar.progress(0.8, "Creating final video...")
             
-            # Get and display description
-            description = get_description(enhanced_image, blip, blip_processor)
-            st.write("📝 Generated Description:", description)
+        # Convert frames to format suitable for moviepy
+        frames = [np.array(frame) for frame in video_frames[0]]
+        
+        # Create temporary file for video
+        with tempfile.NamedTemporaryFile(delete=False, suffix='.mp4') as tmp_file:
+            output_path = tmp_file.name
             
-            # Allow user to edit description
-            modified_description = st.text_area("Edit description if needed:", description)
+        # Create and save video
+        clip = ImageSequenceClip(frames, fps=30)
+        clip.write_videofile(output_path, codec='libx264', audio=False)
+        
+        if progress_bar:
+            progress_bar.progress(1.0, "Video generation complete!")
             
-            # Generate video button
-            if st.button("🎬 Generate Video"):
-                with st.spinner("Generating video... This may take a few minutes."):
-                    video_path = generate_video(pipeline, modified_description)
-                    st.success("Video generated successfully!")
-                    st.video(video_path)
-                    
-                    # Add download button
-                    with open(video_path, 'rb') as f:
-                        st.download_button(
-                            label="Download Video",
-                            data=f,
-                            file_name="generated_video.mp4",
-                            mime="video/mp4"
-                        )
-    
+        return output_path, caption
+        
     except Exception as e:
-        st.error(f"An error occurred: {str(e)}")
-        st.error("Please try again or contact support if the error persists.")
+        st.error(f"Error generating video: {str(e)}")
+        return None, None
+
+def main():
+    st.set_page_config(page_title="Video Generator", page_icon="🎥")
+    
+    st.title("🎥 AI 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.
+    """)
+    
+    # File uploader
+    uploaded_file = st.file_uploader("Choose an image", type=['png', 'jpg', 'jpeg'])
+    
+    # Duration selector
+    duration = st.slider("Video duration (seconds)", min_value=1, max_value=30, value=10)
+    
+    if uploaded_file is not None:
+        # Display uploaded image
+        image = Image.open(uploaded_file)
+        st.image(image, caption="Uploaded Image", use_column_width=True)
+        
+        # Generate button
+        if st.button("Generate Video"):
+            # Create a progress bar
+            progress_text = "Operation in progress. Please wait..."
+            my_bar = st.progress(0, text=progress_text)
+            
+            # Generate video
+            video_path, caption = generate_video_from_image(image, duration, my_bar)
+            
+            if video_path and os.path.exists(video_path):
+                # Read the video file
+                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"
+                )
+                
+                # Display video
+                st.video(video_bytes)
+                
+                # Clean up temporary file
+                os.unlink(video_path)
+            else:
+                st.error("Failed to generate video. Please try again.")
 
 if __name__ == "__main__":
     main()