Update app.py

app.py

@@ -811,21 +811,28 @@ def get_video_html(video_path, width="100%"):
 
 
 # *********
-
-def validate_and_preprocess_image(image_data, target_size=(576, 1024)):
-    """Validate and preprocess image for video generation with detailed logging"""
+def validate_and_preprocess_image(file_data, target_size=(576, 1024)):
+    """Validate and preprocess image for video generation with improved BytesIO handling"""
     try:
         st.write("Starting image preprocessing...")
         
-        # Convert bytes to PIL Image if needed
-        if isinstance(image_data, bytes):
-            img = Image.open(io.BytesIO(image_data))
-        elif isinstance(image_data, Image.Image):
-            img = image_data
+        # Handle different input types
+        if isinstance(file_data, bytes):
+            st.write("Processing bytes input")
+            img = Image.open(io.BytesIO(file_data))
+        elif hasattr(file_data, 'read'):
+            st.write("Processing file-like object")
+            # Reset file pointer if possible
+            if hasattr(file_data, 'seek'):
+                file_data.seek(0)
+            img = Image.open(file_data)
+        elif isinstance(file_data, Image.Image):
+            st.write("Processing PIL Image input")
+            img = file_data
         else:
-            raise ValueError(f"Unsupported image data type: {type(image_data)}")
+            raise ValueError(f"Unsupported input type: {type(file_data)}")
         
-        st.write(f"Original image size: {img.size}, mode: {img.mode}")
+        st.write(f"Successfully loaded image: {img.format}, size={img.size}, mode={img.mode}")
         
         # Convert to RGB if necessary
         if img.mode != 'RGB':
@@ -864,118 +871,20 @@ def validate_and_preprocess_image(image_data, target_size=(576, 1024)):
         final_img.paste(resized_img, (paste_x, paste_y))
         
         st.write(f"Final image size: {final_img.size}")
-        
-        # Validate final image
-        if final_img.size != target_size:
-            raise ValueError(f"Final image size {final_img.size} doesn't match target size {target_size}")
-            
         return final_img
         
     except Exception as e:
-        st.error(f"Error in image preprocessing: {str(e)}")
+        st.error(f"Error in image preprocessing: {str(e)}\nType of input: {type(file_data)}")
         return None
 
-def generate_video_from_image(image_data, seed=None, motion_bucket_id=127, fps_id=6, max_retries=3):
-    """Generate video from image with improved preprocessing and error handling"""
-    temp_files = []
-    try:
-        # Set up progress tracking
-        progress_bar = st.progress(0)
-        status_text = st.empty()
-        
-        # Preprocess image
-        status_text.text("Preprocessing image...")
-        progress_bar.progress(10)
-        
-        processed_img = validate_and_preprocess_image(image_data)
-        if processed_img is None:
-            st.error("Image preprocessing failed")
-            return None, None
-            
-        # Show preprocessed image
-        st.write("Preprocessed image preview:")
-        st.image(processed_img, caption="Preprocessed image", use_column_width=True)
-        
-        # Save processed image
-        with tempfile.NamedTemporaryFile(suffix='.png', delete=False) as temp_img:
-            temp_files.append(temp_img.name)
-            processed_img.save(temp_img.name, format='PNG', optimize=True)
-            st.write(f"Saved preprocessed image to: {temp_img.name}")
-            
-            # Verify file size
-            file_size = os.path.getsize(temp_img.name)
-            st.write(f"Preprocessed image file size: {file_size/1024:.2f}KB")
-        
-        status_text.text("Connecting to video generation service...")
-        progress_bar.progress(30)
-        
-        # Initialize client with debug flags
-        client = Client(
-            "awacke1/stable-video-diffusion",
-            hf_token=os.environ.get("HUGGINGFACE_TOKEN"),
-        )
-        
-        if seed is None:
-            seed = int(time.time() * 1000)  # Use millisecond timestamp as seed
-        
-        status_text.text("Starting video generation...")
-        progress_bar.progress(40)
-        
-        for attempt in range(max_retries):
-            try:
-                status_text.text(f"Generating video (attempt {attempt + 1}/{max_retries})...")
-                progress_bar.progress(40 + (attempt * 20))
-                
-                # Call video generation API
-                result = client.predict(
-                    image=temp_img.name,
-                    seed=seed,
-                    randomize_seed=False,  # Set to False for reproducibility
-                    motion_bucket_id=motion_bucket_id,
-                    fps_id=fps_id,
-                    api_name="/video"
-                )
-                
-                # Validate result
-                if result and isinstance(result, tuple) and len(result) >= 1:
-                    if isinstance(result[0], dict) and 'video' in result[0]:
-                        video_path = result[0]['video']
-                        if os.path.exists(video_path):
-                            status_text.text("Video generated successfully!")
-                            progress_bar.progress(100)
-                            return video_path, seed
-                
-                st.warning(f"Invalid result format on attempt {attempt + 1}: {result}")
-                time.sleep(2 ** attempt)  # Exponential backoff
-                
-            except Exception as e:
-                st.warning(f"Attempt {attempt + 1} failed: {str(e)}")
-                time.sleep(2 ** attempt)
-        
-        raise Exception(f"Failed to generate video after {max_retries} attempts")
-        
-    except Exception as e:
-        st.error(f"Error in video generation: {str(e)}")
-        return None, None
-        
-    finally:
-        # Cleanup
-        for temp_file in temp_files:
-            try:
-                if os.path.exists(temp_file):
-                    os.unlink(temp_file)
-                    st.write(f"Cleaned up temporary file: {temp_file}")
-            except Exception as e:
-                st.warning(f"Error cleaning up {temp_file}: {str(e)}")
-
 def add_video_generation_ui(container):
-    """Enhanced video generation UI with better error handling and feedback"""
+    """Enhanced video generation UI with improved file handling"""
     st.markdown("### 🎥 Video Generation")
     
     col1, col2 = st.columns([2, 1])
     
     with col1:
-        uploaded_image = st.file_uploader(
+        uploaded_file = st.file_uploader(
             "Upload Image for Video Generation 🖼️",
             type=['png', 'jpg', 'jpeg'],
             help="Upload a clear, well-lit image. Recommended size: 576x1024 pixels."
@@ -998,7 +907,6 @@ def add_video_generation_ui(container):
             help="Higher values create smoother but potentially less stable videos"
         )
         
-        # Add advanced options in an expander
         with st.expander("Advanced Options"):
             use_custom_seed = st.checkbox("Use Custom Seed")
             if use_custom_seed:
@@ -1006,71 +914,103 @@ def add_video_generation_ui(container):
             else:
                 seed = None
 
-    if uploaded_image:
+    if uploaded_file is not None:
         try:
+            # Read file data
+            file_data = uploaded_file.read()
+            
             # Preview original image
             preview_col1, preview_col2 = st.columns(2)
             with preview_col1:
                 st.write("Original Image:")
-                st.image(uploaded_image, caption="Original", use_column_width=True)
+                original_img = Image.open(io.BytesIO(file_data))
+                st.image(original_img, caption="Original", use_column_width=True)
             
             # Preview preprocessed image
             with preview_col2:
-                preprocessed = validate_and_preprocess_image(uploaded_image.read())
+                # Create a new BytesIO object with the file data
+                preprocessed = validate_and_preprocess_image(io.BytesIO(file_data))
                 if preprocessed:
                     st.write("Preprocessed Image:")
                     st.image(preprocessed, caption="Preprocessed", use_column_width=True)
-        except Exception as e:
-            st.error(f"Error previewing image: {str(e)}")
-        
-        if st.button("🎥 Generate Video", help="Start video generation process"):
-            try:
+                else:
+                    st.error("Failed to preprocess image")
+                    return
+            
+            if st.button("🎥 Generate Video", help="Start video generation process"):
                 with st.spinner("Processing your video... This may take a few minutes 🎬"):
-                    video_path, used_seed = generate_video_from_image(
-                        uploaded_image.read(),
-                        seed=seed,
-                        motion_bucket_id=motion_bucket_id,
-                        fps_id=fps_id
-                    )
-                    
-                    if video_path and os.path.exists(video_path):
-                        # Save video locally
-                        video_filename = f"generated_video_{datetime.now().strftime('%Y%m%d_%H%M%S')}.mp4"
+                    # Save preprocessed image to temporary file
+                    with tempfile.NamedTemporaryFile(suffix='.png', delete=False) as temp_file:
+                        preprocessed.save(temp_file.name, format='PNG')
+                        st.write(f"Saved preprocessed image to temporary file: {temp_file.name}")
+                        
                         try:
-                            shutil.copy(video_path, video_filename)
-                            
-                            # Display success and video
-                            st.success(f"""
-                            Video generated successfully! 🎉
-                            - Seed: {used_seed}
-                            - Motion Intensity: {motion_bucket_id}
-                            - FPS: {fps_id}
-                            """)
+                            # Initialize the Gradio client
+                            client = Client(
+                                "awacke1/stable-video-diffusion",
+                                hf_token=os.environ.get("HUGGINGFACE_TOKEN")
+                            )
                             
-                            st.video(video_filename)
+                            # Generate video
+                            result = client.predict(
+                                image=temp_file.name,
+                                seed=seed if seed is not None else int(time.time() * 1000),
+                                randomize_seed=seed is None,
+                                motion_bucket_id=motion_bucket_id,
+                                fps_id=fps_id,
+                                api_name="/video"
+                            )
                             
-                            # Save to Cosmos DB
-                            if container:
-                                video_record = {
-                                    "id": generate_unique_id(),
-                                    "type": "generated_video",
-                                    "filename": video_filename,
-                                    "seed": used_seed,
-                                    "motion_bucket_id": motion_bucket_id,
-                                    "fps_id": fps_id,
-                                    "timestamp": datetime.now().isoformat()
-                                }
-                                success, message = insert_record(container, video_record)
-                                if success:
-                                    st.success("Video record saved to database!")
+                            if result and isinstance(result, tuple) and len(result) >= 1:
+                                video_path = result[0].get('video') if isinstance(result[0], dict) else None
+                                if video_path and os.path.exists(video_path):
+                                    # Save video locally
+                                    video_filename = f"generated_video_{datetime.now().strftime('%Y%m%d_%H%M%S')}.mp4"
+                                    shutil.copy(video_path, video_filename)
+                                    
+                                    st.success(f"""
+                                    Video generated successfully! 🎉
+                                    - Seed: {seed if seed is not None else 'Random'}
+                                    - Motion Intensity: {motion_bucket_id}
+                                    - FPS: {fps_id}
+                                    """)
+                                    
+                                    st.video(video_filename)
+                                    
+                                    # Save to Cosmos DB
+                                    if container:
+                                        video_record = {
+                                            "id": generate_unique_id(),
+                                            "type": "generated_video",
+                                            "filename": video_filename,
+                                            "seed": seed if seed is not None else "random",
+                                            "motion_bucket_id": motion_bucket_id,
+                                            "fps_id": fps_id,
+                                            "timestamp": datetime.now().isoformat()
+                                        }
+                                        success, message = insert_record(container, video_record)
+                                        if success:
+                                            st.success("Video record saved to database!")
+                                        else:
+                                            st.error(f"Error saving video record: {message}")
                                 else:
-                                    st.error(f"Error saving video record: {message}")
+                                    st.error("Failed to generate video: Invalid result format")
+                            else:
+                                st.error("Failed to generate video: No result returned")
+                            
                         except Exception as e:
-                            st.error(f"Error saving video: {str(e)}")
-                    else:
-                        st.error("Video generation failed. Please try again with different parameters.")
-            except Exception as e:
-                st.error(f"Error during video generation process: {str(e)}")
+                            st.error(f"Error generating video: {str(e)}")
+                        finally:
+                            # Cleanup temporary file
+                            try:
+                                os.unlink(temp_file.name)
+                                st.write("Cleaned up temporary file")
+                            except Exception as e:
+                                st.warning(f"Error cleaning up temporary file: {str(e)}")
+                
+        except Exception as e:
+            st.error(f"Error processing uploaded file: {str(e)}")
+
 
 # ******************************************