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 = "gen-lang-client-0193353123"
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 (NEW)
# This section looks for a secret named 'HF_TOKEN' to log into the Hub.
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 (Unchanged)
# This section expects a secret named 'GOOGLE_APPLICATION_CREDENTIALS_JSON'
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.")
    # Define a dummy function to show an error in the UI
    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():
        """Generates a fresh short-lived access token for Google Cloud."""
        auth_req = google.auth.transport.requests.Request()
        credentials.refresh(auth_req)
        return credentials.token

    # --- 2. Core Video Generation Logic (Unchanged) ---

    def generate_video(prompt: str):
        """
        The main function to generate a video. It submits, polls, and returns the result.
        """
        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.")
                    video_base64 = poll_result["response"]["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!", temp_video_path
                    return
                
                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:
            print(f"An unexpected error occurred: {e}")
            raise gr.Error(f"An unexpected error occurred: {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()