Spaces:
Sleeping
Sleeping
| import shutil | |
| import tempfile | |
| from pathlib import Path | |
| import streamlit as st | |
| from sorawm.core import SoraWM | |
| from sorawm.schemas import CleanerType | |
| def main(): | |
| st.set_page_config( | |
| page_title="Sora Watermark Cleaner", page_icon="🎬", layout="centered" | |
| ) | |
| # Header section with improved layout | |
| st.markdown( | |
| """ | |
| <div style='text-align: center; padding: 1rem 0;'> | |
| <h1 style='margin-bottom: 0.5rem;'> | |
| 🎬 Sora Watermark Cleaner | |
| </h1> | |
| <p style='font-size: 1.2rem; color: #666; margin-bottom: 1rem;'> | |
| Remove watermarks from Sora-generated videos with AI-powered precision | |
| </p> | |
| </div> | |
| """, | |
| unsafe_allow_html=True, | |
| ) | |
| # # Feature badges | |
| # col1, col2, col3 = st.columns(3) | |
| # with col1: | |
| # st.markdown( | |
| # """ | |
| # <div style='text-align: center; padding: 0.8rem; background: linear-gradient(135deg, #667eea 0%, #764ba2 100%); | |
| # border-radius: 10px; color: white;'> | |
| # <div style='font-size: 1.5rem;'>⚡</div> | |
| # <div style='font-weight: bold;'>Fast Processing</div> | |
| # <div style='font-size: 0.85rem; opacity: 0.9;'>GPU Accelerated</div> | |
| # </div> | |
| # """, | |
| # unsafe_allow_html=True, | |
| # ) | |
| # with col2: | |
| # st.markdown( | |
| # """ | |
| # <div style='text-align: center; padding: 0.8rem; background: linear-gradient(135deg, #f093fb 0%, #f5576c 100%); | |
| # border-radius: 10px; color: white;'> | |
| # <div style='font-size: 1.5rem;'>🎯</div> | |
| # <div style='font-weight: bold;'>High Precision</div> | |
| # <div style='font-size: 0.85rem; opacity: 0.9;'>AI-Powered</div> | |
| # </div> | |
| # """, | |
| # unsafe_allow_html=True, | |
| # ) | |
| # with col3: | |
| # st.markdown( | |
| # """ | |
| # <div style='text-align: center; padding: 0.8rem; background: linear-gradient(135deg, #4facfe 0%, #00f2fe 100%); | |
| # border-radius: 10px; color: white;'> | |
| # <div style='font-size: 1.5rem;'>📦</div> | |
| # <div style='font-weight: bold;'>Batch Support</div> | |
| # <div style='font-size: 0.85rem; opacity: 0.9;'>Process Multiple</div> | |
| # </div> | |
| # """, | |
| # unsafe_allow_html=True, | |
| # ) | |
| # Promo banner | |
| st.markdown( | |
| """ | |
| <div style='text-align: center; padding: 0.8rem 1rem; margin-top: 0.5rem; | |
| background: linear-gradient(135deg, #667eea 0%, #764ba2 100%); | |
| border-radius: 10px; color: white;'> | |
| <p style='margin: 0; font-size: 1rem;'> | |
| 🌐 Also try <strong><a href='https://cleanmark.org' target='_blank' | |
| style='color: #fff; text-decoration: underline;'>CleanMark.org</a></strong> | |
| — Free online tools to remove AI watermarks from images & videos (Gemini, Doubao, Sora) | |
| </p> | |
| </div> | |
| """, | |
| unsafe_allow_html=True, | |
| ) | |
| # Footer info | |
| st.markdown( | |
| """ | |
| <div style='text-align: center; padding: 1rem 0; margin-top: 1rem;'> | |
| <p style='color: #888; font-size: 0.9rem;'> | |
| Built with ❤️ using Streamlit and AI | | |
| <a href='https://cleanmark.org' target='_blank' style='color: #667eea; text-decoration: none;'> | |
| 🌐 CleanMark.org | |
| </a> | | |
| <a href='https://github.com/sunzhenyu/CleanMark' | |
| target='_blank' style='color: #667eea; text-decoration: none;'> | |
| ⭐ Star on GitHub | |
| </a> | |
| </p> | |
| </div> | |
| """, | |
| unsafe_allow_html=True, | |
| ) | |
| st.markdown("---") | |
| # Model selection | |
| st.markdown("### ⚙️ Model Settings") | |
| col1, col2 = st.columns([2, 3]) | |
| with col1: | |
| model_type = st.selectbox( | |
| "Select Cleaner Model:", | |
| options=[CleanerType.LAMA, CleanerType.E2FGVI_HQ], | |
| format_func=lambda x: { | |
| CleanerType.LAMA: "🚀 LAMA (Fast, Good Quality)", | |
| CleanerType.E2FGVI_HQ: "💎 E2FGVI-HQ (Slower when not on GPU, Best Quality with time consistency)", | |
| }[x], | |
| help="LAMA: Fast processing with good quality. E2FGVI-HQ: Slower when not on GPU but highest quality results.", | |
| ) | |
| with col2: | |
| model_info = { | |
| CleanerType.LAMA: "⚡ **Fast processing** - Recommended for most videos. Uses LaMa (Large Mask Inpainting) for quick watermark removal.", | |
| CleanerType.E2FGVI_HQ: "🎯 **Highest quality** - Uses temporal flow-based video inpainting. Best for professional results. Slower when not on GPU. Time consistency is guaranteed.", | |
| } | |
| st.info(model_info[model_type]) | |
| # Initialize or reinitialize SoraWM if model changed | |
| if ( | |
| "sora_wm" not in st.session_state | |
| or st.session_state.get("current_model") != model_type | |
| ): | |
| with st.spinner(f"Loading {model_type.value.upper()} model..."): | |
| st.session_state.sora_wm = SoraWM(cleaner_type=model_type) | |
| st.session_state.current_model = model_type | |
| st.success(f"✅ {model_type.value.upper()} model loaded!") | |
| st.markdown("---") | |
| # Mode selection | |
| mode = st.radio( | |
| "Select input mode:", | |
| ["📁 Upload Video File", "🗂️ Process Folder"], | |
| horizontal=True, | |
| ) | |
| if mode == "📁 Upload Video File": | |
| # File uploader | |
| uploaded_file = st.file_uploader( | |
| "Upload your video", | |
| type=["mp4", "avi", "mov", "mkv"], | |
| accept_multiple_files=False, | |
| help="Select a video file to remove watermark", | |
| ) | |
| if uploaded_file: | |
| # Clear previous processed video if a new file is uploaded | |
| if ( | |
| "current_file_name" not in st.session_state | |
| or st.session_state.current_file_name != uploaded_file.name | |
| ): | |
| st.session_state.current_file_name = uploaded_file.name | |
| if "processed_video_data" in st.session_state: | |
| del st.session_state.processed_video_data | |
| if "processed_video_path" in st.session_state: | |
| del st.session_state.processed_video_path | |
| if "processed_video_name" in st.session_state: | |
| del st.session_state.processed_video_name | |
| # Display video info | |
| st.success(f"✅ Uploaded: {uploaded_file.name}") | |
| # Create two columns for before/after comparison | |
| col_left, col_right = st.columns(2) | |
| with col_left: | |
| st.markdown("### 📥 Original Video") | |
| st.video(uploaded_file) | |
| with col_right: | |
| st.markdown("### 🎬 Processed Video") | |
| # Placeholder for processed video | |
| if "processed_video_data" not in st.session_state: | |
| st.info("Click 'Remove Watermark' to process the video") | |
| else: | |
| st.video(st.session_state.processed_video_data) | |
| # Process button | |
| if st.button( | |
| "🚀 Remove Watermark", type="primary", use_container_width=True | |
| ): | |
| with tempfile.TemporaryDirectory() as tmp_dir: | |
| tmp_path = Path(tmp_dir) | |
| try: | |
| # Create progress bar and status text | |
| progress_bar = st.progress(0) | |
| status_text = st.empty() | |
| def update_progress(progress: int): | |
| progress_bar.progress(progress / 100) | |
| if progress < 50: | |
| status_text.text( | |
| f"🔍 Detecting watermarks... {progress}%" | |
| ) | |
| elif progress < 95: | |
| status_text.text( | |
| f"🧹 Removing watermarks... {progress}%" | |
| ) | |
| else: | |
| status_text.text(f"🎵 Merging audio... {progress}%") | |
| # Single file processing | |
| input_path = tmp_path / uploaded_file.name | |
| with open(input_path, "wb") as f: | |
| f.write(uploaded_file.read()) | |
| output_path = tmp_path / f"cleaned_{uploaded_file.name}" | |
| st.session_state.sora_wm.run( | |
| input_path, output_path, progress_callback=update_progress | |
| ) | |
| progress_bar.progress(100) | |
| status_text.text("✅ Processing complete!") | |
| st.success("✅ Watermark removed successfully!") | |
| # Store processed video path and read video data | |
| with open(output_path, "rb") as f: | |
| video_data = f.read() | |
| st.session_state.processed_video_path = output_path | |
| st.session_state.processed_video_data = video_data | |
| st.session_state.processed_video_name = ( | |
| f"cleaned_{uploaded_file.name}" | |
| ) | |
| # Rerun to show the video in the right column | |
| st.rerun() | |
| except Exception as e: | |
| st.error(f"❌ Error processing video: {str(e)}") | |
| # Download button (show only if video is processed) | |
| if "processed_video_data" in st.session_state: | |
| st.download_button( | |
| label="⬇️ Download Cleaned Video", | |
| data=st.session_state.processed_video_data, | |
| file_name=st.session_state.processed_video_name, | |
| mime="video/mp4", | |
| use_container_width=True, | |
| ) | |
| else: # Folder mode | |
| st.info( | |
| "💡 Drag and drop your video folder here, or click to browse and select multiple video files" | |
| ) | |
| # File uploader for multiple files (supports folder drag & drop) | |
| uploaded_files = st.file_uploader( | |
| "Upload videos from folder", | |
| type=["mp4", "avi", "mov", "mkv"], | |
| accept_multiple_files=True, | |
| help="You can drag & drop an entire folder here, or select multiple video files", | |
| key="folder_uploader", | |
| ) | |
| if uploaded_files: | |
| # Display uploaded files info | |
| video_count = len(uploaded_files) | |
| st.success(f"✅ {video_count} video file(s) uploaded") | |
| # Show file list in an expander | |
| with st.expander("📋 View uploaded files", expanded=False): | |
| for i, file in enumerate(uploaded_files, 1): | |
| file_size_mb = file.size / (1024 * 1024) | |
| st.text(f"{i}. {file.name} ({file_size_mb:.2f} MB)") | |
| # Process button | |
| if st.button( | |
| "🚀 Process All Videos", type="primary", use_container_width=True | |
| ): | |
| with tempfile.TemporaryDirectory() as tmp_dir: | |
| tmp_path = Path(tmp_dir) | |
| input_folder = tmp_path / "input" | |
| output_folder = tmp_path / "output" | |
| input_folder.mkdir(exist_ok=True) | |
| output_folder.mkdir(exist_ok=True) | |
| try: | |
| # Save all uploaded files to temp folder | |
| status_text = st.empty() | |
| status_text.text("📥 Saving uploaded files...") | |
| for uploaded_file in uploaded_files: | |
| # Preserve folder structure if file.name contains subdirectories | |
| file_path = input_folder / uploaded_file.name | |
| file_path.parent.mkdir(parents=True, exist_ok=True) | |
| with open(file_path, "wb") as f: | |
| f.write(uploaded_file.read()) | |
| # Create progress tracking | |
| progress_bar = st.progress(0) | |
| current_file_text = st.empty() | |
| processed_count = 0 | |
| def update_progress(progress: int): | |
| # Calculate overall progress | |
| overall_progress = ( | |
| (processed_count * 100 + progress) / video_count / 100 | |
| ) | |
| progress_bar.progress(overall_progress) | |
| if progress < 50: | |
| current_file_text.text( | |
| f"🔍 Processing file {processed_count + 1}/{video_count}: Detecting watermarks... {progress}%" | |
| ) | |
| elif progress < 95: | |
| current_file_text.text( | |
| f"🧹 Processing file {processed_count + 1}/{video_count}: Removing watermarks... {progress}%" | |
| ) | |
| else: | |
| current_file_text.text( | |
| f"🎵 Processing file {processed_count + 1}/{video_count}: Merging audio... {progress}%" | |
| ) | |
| # Process each video file | |
| for video_file in input_folder.rglob("*"): | |
| if video_file.is_file() and video_file.suffix.lower() in [ | |
| ".mp4", | |
| ".avi", | |
| ".mov", | |
| ".mkv", | |
| ]: | |
| # Determine output path maintaining folder structure | |
| rel_path = video_file.relative_to(input_folder) | |
| output_path = ( | |
| output_folder | |
| / rel_path.parent | |
| / f"cleaned_{rel_path.name}" | |
| ) | |
| output_path.parent.mkdir(parents=True, exist_ok=True) | |
| # Process the video | |
| st.session_state.sora_wm.run( | |
| video_file, | |
| output_path, | |
| progress_callback=update_progress, | |
| ) | |
| processed_count += 1 | |
| progress_bar.progress(100) | |
| current_file_text.text("✅ All videos processed!") | |
| st.success(f"✅ {video_count} video(s) processed successfully!") | |
| # Create download option for processed videos | |
| st.markdown("### 📦 Download Processed Videos") | |
| # Store processed files info in session state | |
| if "batch_processed_files" not in st.session_state: | |
| st.session_state.batch_processed_files = [] | |
| st.session_state.batch_processed_files.clear() | |
| for processed_file in output_folder.rglob("*"): | |
| if processed_file.is_file(): | |
| with open(processed_file, "rb") as f: | |
| video_data = f.read() | |
| rel_path = processed_file.relative_to(output_folder) | |
| st.session_state.batch_processed_files.append( | |
| {"name": str(rel_path), "data": video_data} | |
| ) | |
| st.rerun() | |
| except Exception as e: | |
| st.error(f"❌ Error processing videos: {str(e)}") | |
| import traceback | |
| st.error(f"Details: {traceback.format_exc()}") | |
| # Show download buttons for processed files | |
| if ( | |
| "batch_processed_files" in st.session_state | |
| and st.session_state.batch_processed_files | |
| ): | |
| st.markdown("---") | |
| st.markdown("### ⬇️ Download Processed Videos") | |
| for file_info in st.session_state.batch_processed_files: | |
| col1, col2 = st.columns([3, 1]) | |
| with col1: | |
| st.text(f"📹 {file_info['name']}") | |
| with col2: | |
| st.download_button( | |
| label="⬇️ Download", | |
| data=file_info["data"], | |
| file_name=file_info["name"], | |
| mime="video/mp4", | |
| key=f"download_{file_info['name']}", | |
| use_container_width=True, | |
| ) | |
| if __name__ == "__main__": | |
| main() | |