Update app.py

app.py

@@ -7,6 +7,7 @@ import numpy as np
 from PIL import Image
 import tempfile
 import os
+import base64
 
 @st.cache_resource
 def load_models():
@@ -22,44 +23,32 @@ def load_models():
    
    return pipeline, blip, blip_processor
 
-def enhance_image(image):
-   img = np.array(image)
-   denoised = cv2.fastNlMeansDenoisingColored(img)
-   lab = cv2.cvtColor(denoised, cv2.COLOR_RGB2LAB)
-   l, a, b = cv2.split(lab)
-   clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8,8))
-   l = clahe.apply(l)
-   enhanced = cv2.cvtColor(cv2.merge([l,a,b]), cv2.COLOR_LAB2RGB)
-   return Image.fromarray(enhanced)
-
-def get_description(image, blip_model, blip_processor):
-   inputs = blip_processor(images=image, return_tensors="pt")
-   output = blip_model.generate(**inputs, max_length=50)
-   return blip_processor.decode(output[0], skip_special_tokens=True)
-
-def save_video_frames(frames, fps=8):
+def generate_video(pipeline, description):
+   video_frames = pipeline(
+       description, 
+       num_inference_steps=30,  # Reduced from 50
+       num_frames=16  # Reduced from 24
+   ).frames
+   
    temp_dir = tempfile.mkdtemp()
    temp_path = os.path.join(temp_dir, "output.mp4")
    
-   height, width = frames[0].shape[:2]
+   height, width = video_frames[0].shape[:2]
    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
-   video_writer = cv2.VideoWriter(temp_path, fourcc, fps, (width, height))
+   video_writer = cv2.VideoWriter(temp_path, fourcc, 8, (width, height))
    
-   for frame in frames:
+   for frame in video_frames:
        video_writer.write(frame)
    video_writer.release()
    
    return temp_path
 
-def generate_video(pipeline, description):
-   video_frames = pipeline(
-       description, 
-       num_inference_steps=50,
-       num_frames=24
-   ).frames
-   
-   video_path = save_video_frames(video_frames)
-   return video_path
+def get_binary_file_downloader_html(bin_file, file_label='File'):
+   with open(bin_file, 'rb') as f:
+       data = f.read()
+   bin_str = base64.b64encode(data).decode()
+   href = f'<a href="data:application/octet-stream;base64,{bin_str}" download="{file_label}">Download {file_label}</a>'
+   return href
 
 def main():
    st.title("Video Generator")
@@ -80,6 +69,7 @@ def main():
            with st.spinner("Generating video..."):
                video_path = generate_video(pipeline, description)
                st.video(video_path)
+               st.markdown(get_binary_file_downloader_html(video_path, 'video.mp4'), unsafe_allow_html=True)
 
 if __name__ == "__main__":
    main()