app.py

# app.py

import os
import gradio as gr
import requests
import json
import time
import base64
import google.auth
import google.auth.transport.requests
from huggingface_hub import login

# --- 1. Configuration and Authentication ---

# IMPORTANT: Replace with your Google Cloud Project ID
GCP_PROJECT_ID = "YOUR_GCP_PROJECT_ID" # Make sure this is replaced!
GCP_LOCATION = "us-central1"
MODEL_ID = "veo-3.0-generate-preview"
API_ENDPOINT = f"{GCP_LOCATION}-aiplatform.googleapis.com"
PREDICT_URL = f"https://{API_ENDPOINT}/v1/projects/{GCP_PROJECT_ID}/locations/{GCP_LOCATION}/publishers/google/models/{MODEL_ID}:predictLongRunning"
FETCH_URL = f"https://{API_ENDPOINT}/v1/projects/{GCP_PROJECT_ID}/locations/{GCP_LOCATION}/publishers/google/models/{MODEL_ID}:fetchPredictOperation"


# --- Authentication Block ---

# Part A: Hugging Face Hub Authentication
hf_token = os.environ.get("HF_TOKEN")
if hf_token:
    print("Hugging Face token found. Logging in.")
    login(token=hf_token)
else:
    print("WARNING: Hugging Face token ('HF_TOKEN') not found. Hub-related features may be disabled.")

# Part B: Google Cloud Authentication
creds_json_str = os.environ.get("GOOGLE_APPLICATION_CREDENTIALS_JSON")
if not creds_json_str:
    print("FATAL: 'GOOGLE_APPLICATION_CREDENTIALS_JSON' secret not found. App cannot authenticate with Google Cloud.")
    def generate_video(prompt):
        raise gr.Error("Authentication failed. Server is missing Google Cloud credentials. Please check the Hugging Face Space secrets.")
else:
    with open("gcp_creds.json", "w") as f:
        f.write(creds_json_str)
    
    SCOPES = ["https://www.googleapis.com/auth/cloud-platform"]
    credentials, _ = google.auth.load_credentials_from_file("gcp_creds.json", scopes=SCOPES)
    print("GCP credentials loaded successfully.")


    def get_access_token():
        auth_req = google.auth.transport.requests.Request()
        credentials.refresh(auth_req)
        return credentials.token

    # --- 2. Core Video Generation Logic (WITH FIX) ---

    def generate_video(prompt: str):
        if not prompt:
            raise gr.Error("Prompt cannot be empty.")
            
        yield "Status: Authenticating and submitting job...", None

        try:
            access_token = get_access_token()
            headers = {
                "Authorization": f"Bearer {access_token}",
                "Content-Type": "application/json",
            }

            # --- Step A: Submit the long-running prediction job ---
            payload = {
                "instances": [{"prompt": prompt}],
                "parameters": {
                    "aspectRatio": "16:9",
                    "sampleCount": 1,
                    "durationSeconds": 8,
                    "personGeneration": "allow_all",
                    "addWatermark": True,
                    "includeRaiReason": True,
                    "generateAudio": True,
                }
            }
            
            response = requests.post(PREDICT_URL, headers=headers, json=payload)
            response.raise_for_status() 

            operation_name = response.json()["name"]
            print(f"Successfully submitted job. Operation Name: {operation_name}")

            # --- Step B: Poll for the result ---
            MAX_POLL_ATTEMPTS = 60 
            for i in range(MAX_POLL_ATTEMPTS):
                status_message = f"Status: Job submitted. Polling for result (Attempt {i+1}/{MAX_POLL_ATTEMPTS})... Please wait."
                yield status_message, None

                access_token = get_access_token()
                headers["Authorization"] = f"Bearer {access_token}"
                
                fetch_payload = {"operationName": operation_name}
                poll_response = requests.post(FETCH_URL, headers=headers, json=fetch_payload)
                poll_response.raise_for_status()
                
                poll_result = poll_response.json()

                if poll_result.get("done"):
                    print("Job finished successfully.")
                    # --- START: ROBUST RESPONSE HANDLING FIX ---
                    
                    # For debugging, print the entire final response from Google
                    print(f"Full successful response payload: {json.dumps(poll_result, indent=2)}")

                    response_data = poll_result.get("response", {})
                    
                    # Check if the 'predictions' key exists and is not empty
                    if "predictions" in response_data and response_data["predictions"]:
                        video_base64 = response_data["predictions"][0]["bytesBase64Encoded"]
                        video_bytes = base64.b64decode(video_base64)
                        
                        temp_video_path = "generated_video.mp4"
                        with open(temp_video_path, "wb") as f:
                            f.write(video_bytes)
                        
                        yield "Status: Done! Video generated successfully.", temp_video_path
                        return
                    else:
                        # This branch handles cases where the job is "done" but no video was returned,
                        # most likely due to safety filters.
                        error_message = "Video generation finished, but the content was blocked by safety filters or another issue prevented video creation."
                        print(f"ERROR: {error_message}")
                        raise gr.Error(error_message)

                    # --- END: ROBUST RESPONSE HANDLING FIX ---
                
                time.sleep(10)

            raise gr.Error("Operation timed out after several minutes. The job may have failed or is taking too long.")

        except requests.exceptions.HTTPError as e:
            print(f"HTTP Error: {e.response.text}")
            raise gr.Error(f"API Error: {e.response.status_code}. Details: {e.response.text}")
        except Exception as e:
            # Catch the specific Gradio error we raised, and any other unexpected errors.
            print(f"An error occurred in the generation process: {e}")
            raise gr.Error(str(e))


# --- 3. Gradio User Interface (Unchanged) ---

with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown("# 🎬 Vertex AI VEO Video Generator")
    gr.Markdown(
        "Generate short videos from a text prompt using Google's VEO model. "
        "Generation can take several minutes. Please be patient."
    )
    
    with gr.Row():
        with gr.Column(scale=1):
            prompt_input = gr.Textbox(
                label="Prompt",
                placeholder="A majestic lion roaming the savanna at sunrise, cinematic 4K.",
                lines=3
            )
            submit_button = gr.Button("Generate Video", variant="primary")

        with gr.Column(scale=1):
            status_output = gr.Markdown("Status: Ready")
            video_output = gr.Video(label="Generated Video", interactive=False)
            
    gr.Examples(
        examples=[
            "A high-speed drone shot flying through a futuristic city with flying vehicles.",
            "A raccoon happily eating popcorn in a movie theater, cinematic lighting.",
            "A beautiful time-lapse of a flower blooming, from bud to full blossom, ultra-realistic.",
        ],
        inputs=prompt_input,
    )

    submit_button.click(
        fn=generate_video,
        inputs=prompt_input,
        outputs=[status_output, video_output]
    )

demo.launch()