app.py

import os
import requests
import gradio as gr
import uuid
import datetime
from supabase import create_client, Client
from supabase.lib.client_options import ClientOptions
import dotenv
from google.cloud import storage
import json
from pathlib import Path
import mimetypes
from workflow_handler import WanVideoWorkflow
from video_config import MODEL_FRAME_RATES, calculate_frames
import asyncio

dotenv.load_dotenv()

SCRIPT_DIR = Path(__file__).parent
CONFIG_PATH = SCRIPT_DIR / "config.json"
WORKFLOW_PATH = SCRIPT_DIR / "wani2v.json"

loras = [
  {
      #I suggest it to be a gif instead of an mp4!
      "image": "https://huggingface.co/Remade-AI/Squish/resolve/main/example_videos/tank_squish.mp4",
      #This is an id you can send to your backend, obviously you can change it
      "id": "06ce6840-f976-4963-9644-b6cf7f323f90",
      #This is the title that is shown on the front end
      "title": "Squish",

      "example_prompt": "In the video, a miniature rodent is presented. The rodent is held in a person's hands. The person then presses on the rodent, causing a sq41sh squish effect. The person keeps pressing down on the rodent, further showing the sq41sh squish effect.",
  },  
  {
      "image": "https://huggingface.co/Remade-AI/Rotate/resolve/main/example_videos/man_rotate.mp4",
      "id": "4ac08cfa-841e-4aa9-9022-c3fc80fb6ef4",
      "title": "Rotate",
      "example_prompt": "The video shows an elderly Asian man's head and shoulders with blurred background, performing a r0t4tion 360 degrees rotation.",
    },
    {
      "image": "https://huggingface.co/Remade-AI/Cakeify/resolve/main/example_videos/timberland_cakeify.mp4",
      "id": "b05c1dc7-a71c-4d24-b512-4877a12dea7e",
      "title": "Cakeify",
      "example_prompt": "The video opens on a woman. A knife, held by a hand, is coming into frame and hovering over the woman. The knife then begins cutting into the woman to c4k3 cakeify it. As the knife slices the woman open, the inside of the woman is revealed to be cake with chocolate layers. The knife cuts through and the contents of the woman are revealed."
    },
  
    
    
    
    
    
]

# Initialize Supabase client with async support
supabase: Client = create_client(
    os.getenv('SUPABASE_URL'),
    os.getenv('SUPABASE_KEY'),
   
)

def initialize_gcs():
    """Initialize Google Cloud Storage client with credentials from environment"""
    try:
        # Parse service account JSON from environment variable
        service_account_json = os.getenv('SERVICE_ACCOUNT_JSON')
        if not service_account_json:
            raise ValueError("SERVICE_ACCOUNT_JSON environment variable not found")
        
        credentials_info = json.loads(service_account_json)
        
        # Initialize storage client
        storage_client = storage.Client.from_service_account_info(credentials_info)
        print("Successfully initialized Google Cloud Storage client")
        return storage_client
    except Exception as e:
        print(f"Error initializing Google Cloud Storage: {e}")
        raise

def upload_to_gcs(file_path, content_type=None, folder='user_uploads'):
    """
    Uploads a file to Google Cloud Storage
    Args:
        file_path: Path to the file to upload
        content_type: MIME type of the file (optional)
        folder: Folder path in bucket (default: 'user_uploads')
    Returns:
        str: Public URL of the uploaded file
    """
    try:
        bucket_name = 'remade-v2'
        storage_client = initialize_gcs()
        bucket = storage_client.bucket(bucket_name)

        # Get file extension and generate unique filename
        file_extension = Path(file_path).suffix
        if not content_type:
            content_type = mimetypes.guess_type(file_path)[0] or 'application/octet-stream'
        
        # Validate file type
        valid_types = ['image/jpeg', 'image/png', 'image/gif']
        if content_type not in valid_types:
            raise ValueError("Invalid file type. Please upload a JPG, PNG or GIF image.")

        # Generate unique filename with proper path structure
        filename = f"{str(uuid.uuid4())}{file_extension}"
        file_path_in_gcs = f"{folder}/{filename}"
        
        # Create blob and set metadata
        blob = bucket.blob(file_path_in_gcs)
        blob.content_type = content_type
        blob.cache_control = 'public, max-age=31536000'
        
        print(f'Uploading file to GCS: {file_path_in_gcs}')
        
        # Upload the file
        blob.upload_from_filename(
            file_path,
            timeout=120  # 2 minute timeout
        )
        
        # Generate public URL with correct path format
        image_url = f"https://storage.googleapis.com/{bucket_name}/{file_path_in_gcs}"
        print(f"Successfully uploaded to GCS: {image_url}")
        return image_url

    except Exception as e:
        print(f"Error uploading to GCS: {e}")
        raise ValueError(f"Failed to upload image to storage: {str(e)}")

def build_lora_prompt(subject, lora_id):
    """
    Builds a standardized prompt based on the selected LoRA and subject
    """
    # Get LoRA config
    lora_config = next((lora for lora in loras if lora["id"] == lora_id), None)
    if not lora_config:
        raise ValueError(f"Invalid LoRA ID: {lora_id}")
    
    if lora_id == "06ce6840-f976-4963-9644-b6cf7f323f90":  # Squish
        return (
            f"In the video, a miniature {subject} is presented. "
            f"The {subject} is held in a person's hands. "
            f"The person then presses on the {subject}, causing a sq41sh squish effect. "
            f"The person keeps pressing down on the {subject}, further showing the sq41sh squish effect."
        )
    
    elif lora_id == "4ac08cfa-841e-4aa9-9022-c3fc80fb6ef4":  # Rotate
        return (
            f"The video shows a {subject} performing a r0t4tion 360 degrees rotation."
        )
    
    elif lora_id == "b05c1dc7-a71c-4d24-b512-4877a12dea7e":  # Cakeify
        return (
            f"The video opens on a {subject}. A knife, held by a hand, is coming into frame "
            f"and hovering over the {subject}. The knife then begins cutting into the {subject} "
            f"to c4k3 cakeify it. As the knife slices the {subject} open, the inside of the "
            f"{subject} is revealed to be cake with chocolate layers. The knife cuts through "
            f"and the contents of the {subject} are revealed."
        )
    
    else:
        raise ValueError(f"Unknown LoRA ID: {lora_id}")

def poll_generation_status(generation_id):
    """Poll generation status from database"""
    try:
        # Query the database for the current status
        response = supabase.table('generations') \
            .select('*') \
            .eq('generation_id', generation_id) \
            .execute()
        
        if not response.data:
            return None
        
        return response.data[0]
    except Exception as e:
        print(f"Error polling generation status: {e}")
        raise e

async def generate_video(input_image, subject, duration, selected_index, progress=gr.Progress()):
    try:
        # Initialize workflow handler with explicit paths
        workflow_handler = WanVideoWorkflow(
            supabase,
            config_path=str(CONFIG_PATH),
            workflow_path=str(WORKFLOW_PATH)
        )
        
        # Upload image to GCS and get public URL
        image_url = upload_to_gcs(input_image)
        
        # Map duration selection to actual seconds
        duration_mapping = {
            "Short (3s)": 3,
            "Long (5s)": 5
        }
        video_duration = duration_mapping[duration]
        
        # Get LoRA config
        lora_config = next((lora for lora in loras if lora["id"] == selected_index), None)
        if not lora_config:
            raise ValueError(f"Invalid LoRA ID: {selected_index}")

        # Generate unique ID
        generation_id = str(uuid.uuid4())

        # Update workflow
        prompt = build_lora_prompt(subject, selected_index)
        workflow_handler.update_prompt(prompt)
        workflow_handler.update_input_image(image_url)
        await workflow_handler.update_lora(lora_config)
        workflow_handler.update_length(video_duration)
        workflow_handler.update_output_name(generation_id)

        # Get final workflow
        workflow = workflow_handler.get_workflow()

        # Store generation data in Supabase
        generation_data = {
            "generation_id": generation_id,
            "user_id": "anonymous",
            "status": "queued",
            "progress": 0,
            "worker_id": None,
            "created_at": datetime.datetime.utcnow().isoformat(),
            "message": {
                "generationId": generation_id,
                "workflow": {
                    "prompt": workflow
                }
            },
            "metadata": {
                "prompt": {
                    "original": subject,
                    "enhanced": subject
                },
                "lora": {
                    "id": selected_index,
                    "strength": 1.0,
                    "name": lora_config["title"]
                },
                "workflow": "img2vid",
                "dimensions": None,
                "input_image_url": image_url,
                "video_length": {"duration": video_duration},
                "platform": "huggingface"
            },
            "error": None,
            "output_url": None,
            "batch_id": None
        }
        
        # Remove await - the execute() method returns the response directly
        response = supabase.table('generations').insert(generation_data).execute()
        print(f"Stored generation data with ID: {generation_id}")
        
        # Return generation ID for tracking
        return generation_id
        
    except Exception as e:
        print(f"Error in generate_video: {e}")
        raise e

def update_selection(evt: gr.SelectData):
  selected_lora = loras[evt.index]
  sentence = f"Selected LoRA: {selected_lora['title']}"
  return selected_lora['id'], sentence

async def handle_generation(input_image, subject, duration, selected_index, progress=gr.Progress(track_tqdm=True)):
    try:
        if selected_index is None:
            raise gr.Error("You must select a LoRA before proceeding.")
            
        # Generate the video and get generation ID
        generation_id = await generate_video(input_image, subject, duration, selected_index)
        
        # Poll for status updates
        while True:
            generation = poll_generation_status(generation_id)
            
            if not generation:
                raise ValueError(f"Generation {generation_id} not found")
                
            # Update progress
            if 'progress' in generation:
                progress_value = generation['progress']
                progress_bar = f'<div class="progress-container"><div class="progress-bar" style="--current: {progress_value}; --total: 100;"></div></div>'
                
                # Check status
                if generation['status'] == 'completed':
                    # Final yield with completed video
                    yield generation['output_url'], generation_id, gr.update(visible=False)
                    break  # Exit the loop
                elif generation['status'] == 'error':
                    raise ValueError(f"Generation failed: {generation.get('error')}")
                else:
                    # Yield progress update
                    yield None, generation_id, gr.update(value=progress_bar, visible=True)
            
            # Wait before next poll
            await asyncio.sleep(2)
            
    except Exception as e:
        print(f"Error in handle_generation: {e}")
        raise e

css = '''
#gen_btn{height: 100%}
#gen_column{align-self: stretch}
#title{text-align: center}
#title h1{font-size: 3em; display:inline-flex; align-items:center}
#title img{width: 100px; margin-right: 0.5em}
#gallery .grid-wrap{height: 10vh}
#lora_list{background: var(--block-background-fill);padding: 0 1em .3em; font-size: 90%}
.card_internal{display: flex;height: 100px;margin-top: .5em}
.card_internal img{margin-right: 1em}
.styler{--form-gap-width: 0px !important}
#progress{height:30px}
#progress .generating{display:none}
.progress-container {width: 100%;height: 30px;background-color: #f0f0f0;border-radius: 15px;overflow: hidden;margin-bottom: 20px}
.progress-bar {height: 100%;background-color: #4f46e5;width: calc(var(--current) / var(--total) * 100%);transition: width 0.5s ease-in-out}
'''

with gr.Blocks(css=css) as demo:
    selected_index = gr.State(None)
    current_generation_id = gr.State(None)
    
    gr.Markdown("# Remade AI - Wan 2.1 I2V effects LoRAs ")
    selected_info = gr.HTML("")
    
    with gr.Row():
        with gr.Column():
            gallery = gr.Gallery(
                [(item["image"], item["title"]) for item in loras],
                label="Select LoRA",
                allow_preview=False,
                columns=4,
                elem_id="gallery",
                show_share_button=False,
                height=350
            )
            input_image = gr.Image(type="filepath")
            subject = gr.Textbox(label="Describe your subject", placeholder="Cat toy")
            duration = gr.Radio(["Short (3s)", "Long (5s)"], label="Duration", value="Short (3s)")
            button = gr.Button("Generate", variant="primary", elem_id="gen_btn")
            
        with gr.Column():
            progress_bar = gr.Markdown(elem_id="progress", visible=False)
            output = gr.Video(interactive=False, label="Output video")
    
    gallery.select(
        update_selection,
        outputs=[selected_index, selected_info]
    )
    
    # Use gr.on for the button click to match the example
    gr.on(
        triggers=[button.click],
        fn=handle_generation,
        inputs=[input_image, subject, duration, selected_index],
        outputs=[output, current_generation_id, progress_bar]
    )

if __name__ == "__main__":
    demo.queue()
    demo.launch()