Update app.py

app.py

@@ -1,140 +1,193 @@
-import gradio as gr
 import os
+import time
 import tempfile
 import shutil
+from pathlib import Path
 from typing import Optional, Tuple, Union
+
+import gradio as gr
 from huggingface_hub import InferenceClient, whoami
-from pathlib import Path
 
-# Initialize Hugging Face Inference Client with fal-ai provider
+# =========================
+# Inference client (fal-ai)
+# =========================
 client = InferenceClient(
     provider="fal-ai",
     api_key=os.environ.get("HF_TOKEN"),
     bill_to="huggingface",
 )
 
+# =========================
+# Auth / PRO helpers
+# =========================
 def verify_pro_status(token: Optional[Union[gr.OAuthToken, str]]) -> bool:
     """Verifies if the user is a Hugging Face PRO user or part of an enterprise org."""
     if not token:
         return False
-    
+
     if isinstance(token, gr.OAuthToken):
         token_str = token.token
     elif isinstance(token, str):
         token_str = token
     else:
         return False
-    
+
     try:
         user_info = whoami(token=token_str)
         return (
-            user_info.get("isPro", False) or
-            any(org.get("isEnterprise", False) for org in user_info.get("orgs", []))
+            user_info.get("isPro", False)
+            or any(org.get("isEnterprise", False) for org in user_info.get("orgs", []))
         )
     except Exception as e:
         print(f"Could not verify user's PRO/Enterprise status: {e}")
         return False
 
+# =========================
+# Storage hygiene
+# =========================
 def cleanup_temp_files():
-    """Clean up old temporary video files to prevent storage overflow."""
+    """Clean up old temporary .mp4 files to prevent storage overflow."""
     try:
         temp_dir = tempfile.gettempdir()
-        # Clean up old .mp4 files in temp directory
         for file_path in Path(temp_dir).glob("*.mp4"):
             try:
                 # Remove files older than 5 minutes
-                if file_path.stat().st_mtime < (os.time.time() - 300):
+                if file_path.stat().st_mtime < (time.time() - 300):
                     file_path.unlink(missing_ok=True)
             except Exception:
                 pass  # Ignore errors for individual files
     except Exception as e:
         print(f"Cleanup error: {e}")
 
+def _write_video_bytes_to_tempfile(video_bytes: bytes) -> str:
+    temp_file = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False)
+    try:
+        temp_file.write(video_bytes)
+        temp_file.flush()
+        return temp_file.name
+    finally:
+        temp_file.close()
+
+# =========================
+# Generation (Text → Video)
+# =========================
 def generate_video(
     prompt: str,
     duration: int = 8,
     size: str = "1280x720",
-    api_key: Optional[str] = None
+    api_key: Optional[str] = None,
 ) -> Tuple[Optional[str], str]:
     """
-    Generate video using Sora-2 through Hugging Face Inference API with fal-ai provider.
-    Returns tuple of (video_path, status_message).
+    Generate video using Sora-2 via Hugging Face Inference API (fal-ai provider).
+    Returns (video_path, status_message).
     """
-    # Clean up old files before generating new ones
     cleanup_temp_files()
-    
+
     try:
         # Use provided API key or environment variable
-        if api_key:
-            temp_client = InferenceClient(
-                provider="fal-ai",
-                api_key=api_key,
-                bill_to="huggingface",
-            )
-        else:
-            temp_client = client
-            if not os.environ.get("HF_TOKEN") and not api_key:
-                return None, "❌ Please set HF_TOKEN environment variable."
-        
-        # Call Sora-2 through Hugging Face Inference API
+        temp_client = (
+            InferenceClient(provider="fal-ai", api_key=api_key, bill_to="huggingface")
+            if api_key
+            else client
+        )
+        if not (api_key or os.environ.get("HF_TOKEN")):
+            return None, "❌ Please set HF_TOKEN environment variable."
+
+        # Call text-to-video
         video_bytes = temp_client.text_to_video(
             prompt,
             model="akhaliq/sora-2",
+            # If your backend supports these, you can forward them as kwargs:
+            # duration=duration,
+            # size=size,
+        )
+
+        video_path = _write_video_bytes_to_tempfile(video_bytes)
+        return video_path, "✅ Video generated successfully!"
+    except Exception as e:
+        return None, f"❌ Error generating video: {str(e)}"
+
+# =========================
+# Generation (Image → Video)
+# =========================
+def generate_video_from_image(
+    prompt: str,
+    image_path: str,
+    api_key: Optional[str] = None,
+) -> Tuple[Optional[str], str]:
+    """
+    Generate video from a single input image using Sora-2 image-to-video.
+    Returns (video_path, status_message).
+    """
+    cleanup_temp_files()
+
+    if not image_path or not Path(image_path).exists():
+        return None, "❌ Please upload an image."
+
+    try:
+        temp_client = (
+            InferenceClient(provider="fal-ai", api_key=api_key, bill_to="huggingface")
+            if api_key
+            else client
+        )
+        if not (api_key or os.environ.get("HF_TOKEN")):
+            return None, "❌ Please set HF_TOKEN environment variable."
+
+        with open(image_path, "rb") as f:
+            input_image = f.read()
+
+        video_bytes = temp_client.image_to_video(
+            input_image,
+            prompt=prompt or "",
+            model="akhaliq/sora-2-image-to-video",
         )
-        
-        # Save to temporary file with proper cleanup
-        # Use NamedTemporaryFile with delete=True but keep reference
-        temp_file = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False)
-        try:
-            temp_file.write(video_bytes)
-            temp_file.flush()
-            video_path = temp_file.name
-        finally:
-            temp_file.close()
-        
-        status_message = f"✅ Video generated successfully!"
-        return video_path, status_message
-        
+
+        video_path = _write_video_bytes_to_tempfile(video_bytes)
+        return video_path, "✅ Video generated successfully from image!"
     except Exception as e:
-        error_msg = f"❌ Error generating video: {str(e)}"
-        return None, error_msg
+        return None, f"❌ Error generating video from image: {str(e)}"
 
+# =========================
+# PRO wrapper (uses request)
+# =========================
 def generate_with_pro_auth(
-    prompt: str, 
-    oauth_token: Optional[gr.OAuthToken] = None
+    mode: str,
+    prompt: str,
+    image_path: Optional[str],
+    request: gr.Request,
 ) -> Tuple[Optional[str], str]:
     """
-    Wrapper function that checks if user is PRO before generating video.
+    Check PRO status from the request's OAuth token, then route to the
+    appropriate generation function based on mode.
     """
-    # Check if user is PRO
+    oauth_token = getattr(request, "oauth_token", None)
     if not verify_pro_status(oauth_token):
-        raise gr.Error("Access Denied. This app is exclusively for Hugging Face PRO users. Please subscribe to PRO to use this app.")
-    
-    if not prompt or prompt.strip() == "":
-        return None, "❌ Please enter a prompt"
-    
-    # Use the environment token for API calls (with bill_to="huggingface")
-    # Don't use the user's OAuth token for the API call
-    video_path, status = generate_video(
-        prompt, 
-        duration=8, 
-        size="1280x720", 
-        api_key=None  # This will use the environment HF_TOKEN
-    )
-    
-    return video_path, status
+        raise gr.Error(
+            "Access Denied. This app is exclusively for Hugging Face PRO users. Please subscribe to PRO to use this app."
+        )
+
+    if mode == "Text → Video":
+        if not prompt or not prompt.strip():
+            return None, "❌ Please enter a prompt."
+        return generate_video(prompt, duration=8, size="1280x720", api_key=None)
+
+    # Image → Video
+    if not image_path:
+        return None, "❌ Please upload an image."
+    # Prompt is optional for image→video; pass empty string if not provided
+    return generate_video_from_image(prompt or "", image_path, api_key=None)
 
 def simple_generate(prompt: str) -> Optional[str]:
-    """Simplified wrapper for examples that only returns video."""
-    if not prompt or prompt.strip() == "":
+    """Examples: only return the video path (text→video)."""
+    if not prompt or not prompt.strip():
         return None
-    
     video_path, _ = generate_video(prompt, duration=8, size="1280x720", api_key=None)
     return video_path
 
+# =========================
+# UI
+# =========================
 def create_ui():
-    """Create the Gradio interface with PRO verification."""
-    
     css = '''
     .logo-dark{display: none}
     .dark .logo-dark{display: block !important}
@@ -151,7 +204,7 @@ def create_ui():
         margin-left: 8px;
     }
     '''
-    
+
     with gr.Blocks(title="Sora-2 Text-to-Video Generator", theme=gr.themes.Soft(), css=css) as demo:
         gr.HTML("""
             <div style="text-align: center; max-width: 800px; margin: 0 auto;">
@@ -169,52 +222,62 @@ def create_ui():
                 </p>
             </div>
         """)
-        
-        # Login button for OAuth
+
+        # HF OAuth (Spaces must have hf_oauth: true)
         gr.LoginButton()
-        
-        # PRO message for non-PRO users
+
+        # Hidden by default; we’ll toggle based on PRO status
         pro_message = gr.Markdown(visible=False)
-        
-        # Main interface (hidden by default)
         main_interface = gr.Column(visible=False)
-        
+
         with main_interface:
             gr.HTML("""
                 <div style="text-align: center; margin: 20px 0;">
                     <p style="color: #28a745; font-weight: bold;">✨ Welcome PRO User! You have full access to Sora-2.</p>
                 </div>
             """)
-            
+
             with gr.Row():
                 with gr.Column(scale=1):
+                    mode_radio = gr.Radio(
+                        choices=["Text → Video", "Image → Video"],
+                        value="Text → Video",
+                        label="Mode",
+                    )
                     prompt_input = gr.Textbox(
-                        label="Enter your prompt",
-                        placeholder="Describe the video you want to create...",
-                        lines=4
+                        label="Prompt",
+                        placeholder="Describe the video you want to create (optional for image→video)...",
+                        lines=4,
                     )
-                    
+
+                    image_group = gr.Group(visible=False)
+                    with image_group:
+                        image_input = gr.Image(
+                            label="Input Image (for Image → Video)",
+                            type="filepath",
+                            sources=["upload", "clipboard"],
+                            image_mode="RGB",
+                        )
+
                     with gr.Accordion("Advanced Settings", open=False):
                         gr.Markdown("*Coming soon: Duration and resolution controls*")
-                    
+
                     generate_btn = gr.Button("🎥 Generate Video", variant="primary", size="lg")
-                
+
                 with gr.Column(scale=1):
                     video_output = gr.Video(
                         label="Generated Video",
                         height=400,
                         interactive=False,
-                        show_download_button=True
+                        show_download_button=True,
                     )
                     status_output = gr.Textbox(
                         label="Status",
                         interactive=False,
-                        visible=True
+                        visible=True,
                     )
-            
-            # Hidden manual token input removed - not needed anymore
-            
-            # Examples section with queue disabled
+
+            # Examples (text→video only)
             gr.Examples(
                 examples=[
                     "A serene beach at sunset with waves gently rolling onto the shore",
@@ -222,60 +285,72 @@ def create_ui():
                     "Northern lights dancing across a starry night sky",
                     "A bustling city street transitioning from day to night in timelapse",
                     "A close-up of coffee being poured into a cup with steam rising",
-                    "Cherry blossoms falling in slow motion in a Japanese garden"
+                    "Cherry blossoms falling in slow motion in a Japanese garden",
                 ],
                 inputs=prompt_input,
                 outputs=video_output,
-                fn=simple_generate,  # Examples use simplified function
+                fn=simple_generate,
                 cache_examples=False,
                 api_name=False,
                 show_api=False,
             )
-            
-            # Event handler for generation with queue disabled
+
+            # Toggle image upload visibility with mode
+            def _toggle_image_group(mode: str):
+                return gr.update(visible=(mode == "Image → Video"))
+
+            mode_radio.change(
+                _toggle_image_group,
+                inputs=[mode_radio],
+                outputs=[image_group],
+                show_progress=False,
+            )
+
+            # Generation handler (uses request to read OAuth token)
             generate_btn.click(
                 fn=generate_with_pro_auth,
-                inputs=[prompt_input],
+                inputs=[mode_radio, prompt_input, image_input],
                 outputs=[video_output, status_output],
                 queue=False,
                 api_name=False,
                 show_api=False,
             )
-            
+
             # Footer
             gr.HTML("""
                 <div style="text-align: center; margin-top: 40px; padding: 20px; border-top: 1px solid #e0e0e0;">
                     <h3 style="color: #667eea;">Thank you for being a PRO user! 🤗</h3>
                 </div>
             """)
-        
-        def control_access(profile: Optional[gr.OAuthProfile] = None, oauth_token: Optional[gr.OAuthToken] = None):
-            """Control interface visibility based on PRO status."""
-            if not profile:
-                # User not logged in
+
+        # Use request to check access on load
+        def control_access(request: gr.Request):
+            oauth_profile = getattr(request, "oauth_profile", None)
+            oauth_token = getattr(request, "oauth_token", None)
+
+            if not oauth_profile:
+                # Not logged in
                 return gr.update(visible=False), gr.update(visible=False)
-            
+
             if verify_pro_status(oauth_token):
-                # User is PRO - show main interface
                 return gr.update(visible=True), gr.update(visible=False)
             else:
-                # User is not PRO - show upgrade message
                 message = """
                 ## ✨ Exclusive Access for PRO Users
-                
+
                 Thank you for your interest in the Sora-2 Text-to-Video Generator!
-                
+
                 This advanced AI video generation tool is available exclusively for Hugging Face **PRO** members.
-                
+
                 ### What you get with PRO:
                 - ✅ Unlimited access to Sora-2 video generation
                 - ✅ High-quality video outputs up to 1280x720
                 - ✅ Fast generation times with priority queue
                 - ✅ Access to other exclusive PRO Spaces
                 - ✅ Support the development of cutting-edge AI tools
-                
+
                 ### Ready to create amazing videos?
-                
+
                 <div style="text-align: center; margin: 30px 0;">
                     <a href="http://huggingface.co/subscribe/pro?source=sora2_video" target="_blank" style="
                         display: inline-block;
@@ -292,39 +367,37 @@ def create_ui():
                         🚀 Become a PRO Today!
                     </a>
                 </div>
-                
+
                 <p style="text-align: center; color: #666; margin-top: 20px;">
                     Join thousands of creators who are already using PRO tools to bring their ideas to life.
                 </p>
                 """
                 return gr.update(visible=False), gr.update(visible=True, value=message)
-        
-        # Check access on load
+
         demo.load(
             control_access,
             inputs=None,
-            outputs=[main_interface, pro_message]
+            outputs=[main_interface, pro_message],
         )
-    
+
     return demo
 
-# Launch the application
+# =========================
+# Entrypoint
+# =========================
 if __name__ == "__main__":
     # Clean up any leftover files on startup
     try:
         cleanup_temp_files()
-        # Also try to clear Gradio's cache
         if os.path.exists("gradio_cached_examples"):
             shutil.rmtree("gradio_cached_examples", ignore_errors=True)
     except Exception as e:
         print(f"Initial cleanup error: {e}")
-    
+
     app = create_ui()
-    # Launch without special auth parameters and no queue
-    # OAuth is enabled via Space metadata (hf_oauth: true in README.md)
     app.launch(
         show_api=False,
         enable_monitoring=False,
         quiet=True,
-        max_threads=10,  # Limit threads to prevent resource exhaustion
-    )
+        max_threads=10,
+    )