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 video_config import MODEL_FRAME_RATES, calculate_frames
import asyncio
from openai import OpenAI
import base64
from google.cloud import vision
from google.oauth2 import service_account
import time
from collections import defaultdict, deque

dotenv.load_dotenv()

SCRIPT_DIR = Path(__file__).parent

# Modal configuration
MODAL_ENDPOINT = os.getenv('FAL_MODAL_ENDPOINT')
MODAL_AUTH_TOKEN = os.getenv('MODAL_AUTH_TOKEN')

# Rate limiting configuration
RATE_LIMIT_GENERATIONS = int(os.getenv('RATE_LIMIT_GENERATIONS', '5'))  # Default 5 generations per hour
RATE_LIMIT_WINDOW = int(os.getenv('RATE_LIMIT_WINDOW', '3600'))  # Default 1 hour in seconds

# In-memory rate limiting storage (for production, consider Redis)
user_generations = defaultdict(deque)

loras = [
   {
      "image": "https://huggingface.co/Remade-AI/Crash-zoom-out/resolve/main/example_videos/1.gif",
      "id": "44c05ca1-422d-4cd4-8508-acadb6d0248c",
      "title": "Crash Zoom Out ",
      "example_prompt": "The video shows a man with a slight smile, then the j432mpscare jumpscare occurs, revealing a distorted and monstrous face with glowing red eyes, filling the frame and accompanied by a loud scream."
    },
         {
      "image": "https://huggingface.co/Remade-AI/Crash-zoom-in/resolve/main/example_videos/1.gif",
      "id": "34a80641-4702-4c1c-91bf-c436a59c79cb",
      "title": "Crash Zoom In ",
      "example_prompt": "The video shows a man with a slight smile, then the j432mpscare jumpscare occurs, revealing a distorted and monstrous face with glowing red eyes, filling the frame and accompanied by a loud scream."
    },
     {
      "image": "https://huggingface.co/Remade-AI/Car-chase/resolve/main/example_videos/2.gif",
      "id": "8b36b7fe-0a0b-4849-b0ed-d9a51ff0cc85",
      "title": "Car Chase",
      "example_prompt": "The video shows a man with a slight smile, then the j432mpscare jumpscare occurs, revealing a distorted and monstrous face with glowing red eyes, filling the frame and accompanied by a loud scream."
    },
     {
      "image": "https://huggingface.co/Remade-AI/Crane-down/resolve/main/example_videos/2.gif",
      "id": "f26db0b7-1c26-4587-b2b5-1cfd0c51c5b3",
      "title": "Crane Down ",
      "example_prompt": "The video shows a man with a slight smile, then the j432mpscare jumpscare occurs, revealing a distorted and monstrous face with glowing red eyes, filling the frame and accompanied by a loud scream."
    },
     {
      "image": "https://huggingface.co/Remade-AI/Crane_up/resolve/main/example_videos/1.gif",
      "id": "07c5e22b-7028-437c-9479-6eb9a50cf993",
      "title": "Crane Up ",
      "example_prompt": "The video shows a man with a slight smile, then the j432mpscare jumpscare occurs, revealing a distorted and monstrous face with glowing red eyes, filling the frame and accompanied by a loud scream."
    },

     {
      "image": "https://huggingface.co/Remade-AI/Crane_over_the_head/resolve/main/example_videos/1.gif",
      "id": "9393f8f4-abe6-4aa7-ba01-0b62e1507feb",
      "title": "Crane Overhead ",
      "example_prompt": "The video shows a man with a slight smile, then the j432mpscare jumpscare occurs, revealing a distorted and monstrous face with glowing red eyes, filling the frame and accompanied by a loud scream."
    },
    
     {
      "image": "https://huggingface.co/Remade-AI/matrix-shot/resolve/main/example_videos/1.gif",
      "id": "219ad5ad-8f23-48dc-b098-b8e6d9fbe6c0",
      "title": "Matrix Shot ",
      "example_prompt": "The video shows a man with a slight smile, then the j432mpscare jumpscare occurs, revealing a distorted and monstrous face with glowing red eyes, filling the frame and accompanied by a loud scream."
    },
     {
      "image": "https://huggingface.co/Remade-AI/360-Orbit/resolve/main/example_videos/1.gif",
      "id": "aaa3e820-5d94-4612-9488-0c9a1b2f5843",
      "title": "360 Orbit ",
      "example_prompt": "The video shows a man with a slight smile, then the j432mpscare jumpscare occurs, revealing a distorted and monstrous face with glowing red eyes, filling the frame and accompanied by a loud scream."
    },
         {
      "image": "https://huggingface.co/Remade-AI/Arc_shot/resolve/main/example_videos/1.gif",
      "id": "a5949ee3-61ea-4a18-bd4d-54c855f5401c",
      "title": "Arc Shot ",
      "example_prompt": "The video shows a man with a slight smile, then the j432mpscare jumpscare occurs, revealing a distorted and monstrous face with glowing red eyes, filling the frame and accompanied by a loud scream."
    },
         {
      "image": "https://huggingface.co/Remade-AI/Hero-run/resolve/main/example_videos/1.gif",
      "id": "36b9edf7-31d7-47d3-ad3b-e166fb3a9842",
      "title": "Hero Run ",
      "example_prompt": "The video shows a man with a slight smile, then the j432mpscare jumpscare occurs, revealing a distorted and monstrous face with glowing red eyes, filling the frame and accompanied by a loud scream."
    },


]

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

# Initialize OpenAI client
openai_client = OpenAI(api_key=os.getenv('OPENAI_API_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 == "c8972c6d-ab8a-4988-9a9d-38082264ef22":  # Jumpscare
        return (
            f"The video shows the {subject} with a slight smile, then the j432mpscare jumpscare occurs, "
            f"revealing a distorted and monstrous face with glowing red eyes, filling the frame and accompanied by a loud scream."
        )
    elif lora_id == "d7cbf9b4-82cd-4a94-ba2f-040e809635fa":  # Angry
        return (
            f"The video starts with the {subject} looking at the camera with a neutral face. "
            f"Then the facial expression of the {subject} changes to 4ngr23 angry face, and begins to yell with clenched fists."

        )
    elif lora_id == "e17959c4-9fa5-4e5b-8f69-d1fb01bbe4fa":  # Cartoon Jaw Drop
        return (
            f"The video shows {subject} smiling wide, "
            f"then {subject} mouth transforms into a dr0p_j88 comical jaw drop, extending down in a long, rectangular shape, and revealing his tongue and teeth."
        )
    elif lora_id == "687255bb-959e-4422-bdbb-5aba93c7c180":  # Kissing
        return (
            f"A {subject} is shown smiling. A man/woman comes into the scene and starts passionately k144ing kissing the {subject}."
        )
    elif lora_id == "4ac2fb4e-5ca2-4338-a59c-549167f5b6d0":  # Laughing
        return (
            f"A {subject} is smiling at the camera. He/she then begins l4a6ing laughing."
        )
    elif lora_id == "bcc4163d-ebda-4cdc-b153-7136cdbf563a":  # Crying
        return (
            f"The video starts with a {ubject} with a solemn expression. Then a tear rolls down his/her cheek, as he/she is cr471ng crying."
        )
    elif lora_id == "13093298-652c-4df8-ba28-62d9d5924754":  # Take a selfie with your younger self
        return (
            f"The video starts with the {subject} smiling at the camera, then s31lf13 taking a selfie with their younger self, "
            f"and the younger self appears next to the {subject} with similar facial features and eye color. "
            f"The younger self wears a white t-shirt and has a cream white jacket. The younger self is smiling slightly."
        )
        
    elif 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:
        # Fallback to using the example prompt from the LoRA config
        if "example_prompt" in lora_config:
            # Replace any specific subject in the example with the user's subject
            return lora_config["example_prompt"].replace("rodent", subject).replace("woman", subject).replace("man", subject)
        else:
            raise ValueError(f"Unknown LoRA ID: {lora_id} and no example prompt available")

def poll_generation_status(generation_id):
    """Poll generation status from Modal backend or database"""
    try:
        # First try to get status from Modal backend if available
        if MODAL_ENDPOINT:
            try:
                response = requests.get(
                    f"{MODAL_ENDPOINT}/fal-effects/status?generation_id={generation_id}",
                    headers=get_modal_auth_headers()
                )
        
            except Exception as e:
                print(f"Error polling Modal backend: {e}")

         

        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 moderate_prompt(prompt: str) -> dict:
    """
    Check if a text prompt contains NSFW content with strict rules against inappropriate content
    """
    try:
        # First check with OpenAI moderation
        response = openai_client.moderations.create(input=prompt)
        result = response.results[0]
        
        if result.flagged:
            # Find which categories were flagged
            flagged_categories = [
                category for category, flagged in result.categories.model_dump().items() 
                if flagged
            ]
            
            return {
                "isNSFW": True,
                "reason": f"Content flagged for: {', '.join(flagged_categories)}"
            }
        
        # Additional checks for keywords related to minors or inappropriate content
        keywords = [
            "child", "kid", "minor", "teen", "baby", "infant", "underage",
            "naked", "nude", "nsfw", "porn", "xxx", "sex", "explicit",
            "inappropriate", "adult content"
        ]
        
        lower_prompt = prompt.lower()
        found_keywords = [word for word in keywords if word in lower_prompt]
        
        if found_keywords:
            return {
                "isNSFW": True,
                "reason": f"Content contains inappropriate keywords: {', '.join(found_keywords)}"
            }
        
        return {"isNSFW": False, "reason": None}
    except Exception as e:
        print(f"Error during prompt moderation: {e}")
        # If there's an error, reject the prompt to be safe
        return {
            "isNSFW": True,
            "reason": "Failed to verify prompt safety - please try again"
        }

async def moderate_image(image_path: str) -> dict:
    """
    Check if an image contains NSFW content using both Google Cloud Vision API's SafeSearch detection
    and OpenAI's vision model for double verification
    """
    try:
        # Convert image to base64 for OpenAI
        with open(image_path, "rb") as image_file:
            base64_image = base64.b64encode(image_file.read()).decode('utf-8')
        
        # 1. Google Cloud Vision API Check using proper client library
        try:
            # Get service account info from environment
            service_account_info = json.loads(os.getenv('SERVICE_ACCOUNT_JSON'))
            
            # Initialize Vision client with credentials
            credentials = service_account.Credentials.from_service_account_info(service_account_info)
            vision_client = vision.ImageAnnotatorClient(credentials=credentials)
            
            # Load image content
            with open(image_path, "rb") as image_file:
                content = image_file.read()
            
            # Create image object
            image = vision.Image(content=content)
            
            # Perform safe search detection
            response = vision_client.safe_search_detection(image=image)
            safe_search = response.safe_search_annotation
            
            # Map likelihood values
            likelihood_values = {
                vision.Likelihood.VERY_LIKELY: 4,
                vision.Likelihood.LIKELY: 3,
                vision.Likelihood.POSSIBLE: 2,
                vision.Likelihood.UNLIKELY: 1,
                vision.Likelihood.VERY_UNLIKELY: 0,
                vision.Likelihood.UNKNOWN: 0
            }
            
            # Get likelihood scores
            adult_score = likelihood_values[safe_search.adult]
  
            
            # Determine if content is NSFW according to Vision API
            vision_reasons = []
            if adult_score >= 3:  # LIKELY or VERY_LIKELY
                vision_reasons.append("adult content")
   
                
            # Print Vision API results
            print("Google Cloud Vision API Results:")
            print(f"Adult: {vision.Likelihood(safe_search.adult).name}")
  
        except Exception as vision_error:
            print(f"Error with Vision API: {vision_error}")
            vision_reasons = []  # Continue with OpenAI check if Vision API fails
        
        # 2. OpenAI Vision Check
        openai_response = openai_client.chat.completions.create(
            model="gpt-4o",
            messages=[
                {
                    "role": "system",
                    "content": "You are a content review system. Your task is to analyze the image and determine if it contains any inappropriate content that would make it unsuitable for general audiences. Especially don't allow any content involving minors (children). Please respond with 'SAFE' if the image is appropriate, or 'NSFW: <reason>' if you detect any concerning content."
                },
                {
                    "role": "user",
                    "content": [
                        {"type": "text", "text": "Is this image appropriate for general audiences?"},
                        {"type": "image_url", "image_url": {
                            "url": f"data:image/jpeg;base64,{base64_image}"
                        }}
                    ]
                }
            ],
            max_tokens=150
        )
        
        openai_result = openai_response.choices[0].message.content.strip()
        print("\nOpenAI Vision Check Result:")
        print(openai_result)
        
        # Combine results from both services
        is_nsfw_vision = bool(vision_reasons)
        is_nsfw_openai = openai_result.startswith("NSFW:")
        
        # If either service flags the content, reject it
        if is_nsfw_vision or is_nsfw_openai:
            reasons = []
            if vision_reasons:
                reasons.append(f"Google Vision: {', '.join(vision_reasons)}")
            if is_nsfw_openai:
                reasons.append(f"OpenAI: {openai_result[5:].strip()}")
                
            return {
                "isNSFW": True,
                "reason": f"Image contains inappropriate content: {'; '.join(reasons)}"
            }
        
        return {"isNSFW": False, "reason": None}
        
    except Exception as e:
        print(f"Error during image moderation: {e}")
        # If there's an error, reject the image to be safe
        return {
            "isNSFW": True,
            "reason": "Failed to verify image safety - please try again"
        }

async def moderate_combined(prompt: str, image_path: str) -> dict:
    """
    Check if both a prompt and image contain NSFW content together with strict rules
    """
    try:
        # Convert image to base64
        with open(image_path, "rb") as image_file:
            base64_image = base64.b64encode(image_file.read()).decode('utf-8')
        
        response = openai_client.chat.completions.create(
            model="gpt-4o",
            messages=[
                {
                    "role": "system",
                    "content": """You are a strict content moderation system. Your task is to determine if an image and prompt combination would result in inappropriate content.

STRICT RULES:
1. NO content involving minors in any way
2. NO nudity or sexually suggestive content
3. NO extreme violence or gore
4. NO hate speech or discriminatory content
5. NO illegal activities

Respond with 'NSFW: <reason>' if ANY of these rules are violated, or 'SAFE' if appropriate.
Be extremely cautious - if there's any doubt, mark it as NSFW."""
                },
                {
                    "role": "user",
                    "content": [
                        {
                            "type": "text",
                            "text": f'Please moderate this image and prompt combination for an image-to-video generation:\n\nPrompt: "{prompt}"\n\nEnsure NO inappropriate content, especially involving minors.'
                        },
                        {
                            "type": "image_url",
                            "image_url": {
                                "url": f"data:image/jpeg;base64,{base64_image}"
                            }
                        }
                    ]
                }
            ],
            max_tokens=150
        )
        
        result = response.choices[0].message.content.strip()
        if result.startswith("NSFW:"):
            return {
                "isNSFW": True,
                "reason": result[5:].strip()
            }
        return {
            "isNSFW": False,
            "reason": None
        }
    except Exception as e:
        print(f"Error during combined moderation: {e}")
        # If there's an error, reject to be safe
        return {
            "isNSFW": True,
            "reason": "Failed to verify content safety - please try again"
        }

async def generate_video(input_image, subject, selected_index, progress=gr.Progress()):
    try:
        # Check if the input is a URL (example image) or a file path (user upload)
        if input_image.startswith('http'):
            # It's already a URL, use it directly
            image_url = input_image
        else:
            # It's a file path, upload to GCS
            image_url = upload_to_gcs(input_image)
        
        # Hardcode duration to 3 seconds
        video_duration = 5
        
        # 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())

        # Build prompt for the LoRA
        prompt = subject
        
        # Check if Modal endpoint is configured
        if not MODAL_ENDPOINT:
            raise ValueError("Modal endpoint not configured - FAL_MODAL_ENDPOINT environment variable not found")

        # Calculate frames based on duration and frame rate
        frame_rate = 16  # WanVideo frame rate
        num_frames = calculate_frames(video_duration, frame_rate)
        
        print(f"Sending request to Modal backend: {MODAL_ENDPOINT}/fal-effects")
        
        # Make POST request to the modal backend
        response = requests.post(f"{MODAL_ENDPOINT}/fal-effects", 
            headers=get_modal_auth_headers(),
            json={
                "user_id": "anonymous",  # Since we don't have user auth in this app
                "image_url": image_url,
                "subject": prompt,  # Use the built prompt as subject
                "aspect_ratio": "16:9",  # Default aspect ratio for effects
                "num_frames": 81,
                "frames_per_second": frame_rate,
                "length": str(5),
                "enhance_prompt": False,
                "lora_scale": 1.0,
                "turbo_mode": False,
                "lora_id": selected_index,
                "lora_strength": 1.0,
                "generation_ids": [generation_id]
            }
        )

        if not response.ok:
            error_text = response.text
            try:
                error_json = response.json()
                error_message = error_json.get('detail') or error_json.get('error') or 'Failed to create generation'
            except:
                error_message = f'Failed to create generation: {error_text}'
            raise ValueError(error_message)

        result = response.json()
        print(f"Modal backend response: {result}")
        
        # Extract generation ID from response
        if 'generation_id' in result:
            return result['generation_id']
        elif 'id' in result:
            return result['id']
        else:
            # Fallback to our generated ID if the response doesn't contain one
            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(image_input, subject, selected_index, request: gr.Request, progress=gr.Progress(track_tqdm=True)):
    try:
        if selected_index is None:
            raise gr.Error("You must select a LoRA before proceeding.")
        
        # Check rate limit first
        user_identifier = get_user_identifier(request)
        is_allowed, remaining, reset_time = check_rate_limit(user_identifier)
        
        if not is_allowed:
            minutes = reset_time // 60
            seconds = reset_time % 60
            time_str = f"{minutes}m {seconds}s" if minutes > 0 else f"{seconds}s"
            # Re-enable button on rate limit
            yield None, None, gr.update(visible=False), gr.update(value="Generate", interactive=True)
            raise gr.Error(f"Rate limit exceeded. Go to https://app.remade.ai for more generations and effects. Otherwise, you can generate {RATE_LIMIT_GENERATIONS} videos per hour. Try again in {time_str}.")
        
        # Record this generation attempt
        record_generation(user_identifier)
        
        # Show remaining generations to user
        if remaining > 0:
            print(f"User {user_identifier} has {remaining} generations remaining this hour")
        
        # First, moderate the prompt
        prompt_moderation = await moderate_prompt(subject)
        print(f"Prompt moderation result: {prompt_moderation}")
        if prompt_moderation["isNSFW"]:
            # Re-enable button on error
            yield None, None, gr.update(visible=False), gr.update(value="Generate", interactive=True)
            raise gr.Error(f"Content moderation failed: {prompt_moderation['reason']}")
        
        # Then, moderate the image
        image_moderation = await moderate_image(image_input)
        print(f"Image moderation result: {image_moderation}")
        if image_moderation["isNSFW"]:
            # Re-enable button on error
            yield None, None, gr.update(visible=False), gr.update(value="Generate", interactive=True)
            raise gr.Error(f"Content moderation failed: {image_moderation['reason']}")
        
        # Finally, check the combination
        combined_moderation = await moderate_combined(subject, image_input)
        print(f"Combined moderation result: {combined_moderation}")
        if combined_moderation["isNSFW"]:
            # Re-enable button on error
            yield None, None, gr.update(visible=False), gr.update(value="Generate", interactive=True)
            raise gr.Error(f"Content moderation failed: {combined_moderation['reason']}")
            
        # Generate the video and get generation ID
        generation_id = await generate_video(image_input, subject, selected_index)
        
        # Poll for status updates
        while True:
            generation = poll_generation_status(generation_id)
            
            if not generation:
                # Re-enable button on error
                yield None, None, gr.update(visible=False), gr.update(value="Generate", interactive=True)
                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;"><span class="progress-text">Processing: {progress_value}%</span></div></div><div class="refresh-warning">Please do not refresh this page while processing</div>'
                
                # Check status
                if generation['status'] == 'completed':
                    # Final yield with completed video and re-enabled button
                    yield generation['output_url'], generation_id, gr.update(visible=False), gr.update(value="Generate", interactive=True)
                    break  # Exit the loop
                elif generation['status'] == 'error':
                    # Re-enable button on error
                    yield None, None, gr.update(visible=False), gr.update(value="Generate", interactive=True)
                    raise ValueError(f"Generation failed: {generation.get('error')}")
                else:
                    # Yield progress update with button still disabled
                    yield None, generation_id, gr.update(value=progress_bar, visible=True), gr.update(value="Generating...", interactive=False)
            
            # Wait before next poll
            await asyncio.sleep(2)
            
    except Exception as e:
        print(f"Error in handle_generation: {e}")
        # Re-enable button on any error
        yield None, None, gr.update(visible=False), gr.update(value="Generate", interactive=True)
        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: auto; min-height: 350px}
#gallery .gallery-item {height: 100%; width: 100%; object-fit: cover}
#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: #2a2a2a;border-radius: 15px;overflow: hidden;margin-bottom: 20px;position: relative;}
.progress-bar {height: 100%;background-color: #7289DA;width: calc(var(--current) / var(--total) * 100%);transition: width 0.5s ease-in-out}
.progress-text {position: absolute;width: 100%;text-align: center;top: 50%;left: 0;transform: translateY(-50%);color: #ffffff;font-weight: bold;}
.refresh-warning {color: #ff7675;font-weight: bold;text-align: center;margin-top: 5px;}

/* Dark mode Discord styling */
.discord-banner {
    background: linear-gradient(135deg, #7289DA 0%, #5865F2 100%);
    color: #ffffff;
    padding: 20px;
    border-radius: 12px;
    margin: 15px 0;
    text-align: center;
    box-shadow: 0 4px 8px rgba(0,0,0,0.3);
}
.discord-banner h3 {
    margin-top: 0;
    font-size: 1.5em;
    text-shadow: 0 2px 4px rgba(0,0,0,0.3);
    color: #ffffff;
}
.discord-banner p {
    color: #ffffff;
    margin-bottom: 15px;
}
.discord-banner a {
    display: inline-block;
    background-color: #ffffff;
    color: #5865F2;
    text-decoration: none;
    font-weight: bold;
    padding: 10px 20px;
    border-radius: 24px;
    margin-top: 10px;
    transition: all 0.3s ease;
    box-shadow: 0 2px 8px rgba(0,0,0,0.3);
    border: none;
}
.discord-banner a:hover {
    transform: translateY(-3px);
    box-shadow: 0 6px 12px rgba(0,0,0,0.4);
    background-color: #f2f2f2;
}
.discord-banner .discord-community-btn {
    background-color: #ffffff !important;
    color: #5865F2 !important;
    opacity: 1 !important;
    font-weight: bold;
    font-size: 0.9em;
    padding: 8px 16px;
    border-radius: 20px;
    text-decoration: none;
    display: inline-block;
    transition: all 0.3s ease;
    box-shadow: 0 2px 6px rgba(0,0,0,0.2);
}
.discord-banner .discord-community-btn:hover {
    background-color: #f8f8f8 !important;
    transform: translateY(-2px);
    box-shadow: 0 4px 10px rgba(0,0,0,0.3);
}
.discord-feature {
    background-color: #2a2a2a;
    border-left: 4px solid #7289DA;
    padding: 12px 15px;
    margin: 10px 0;
    border-radius: 0 8px 8px 0;
    box-shadow: 0 2px 4px rgba(0,0,0,0.2);
    color: #e0e0e0;
}
.discord-feature-title {
    font-weight: bold;
    color: #7289DA;
}
.discord-locked {
    opacity: 0.7;
    position: relative;
    pointer-events: none;
}
.discord-locked::after {
    content: "🔒 Remade Canvas exclusive";
    position: absolute;
    top: 50%;
    left: 50%;
    transform: translate(-50%, -50%);
    background: rgba(114,137,218,0.9);
    color: white;
    padding: 5px 10px;
    border-radius: 20px;
    white-space: nowrap;
    font-size: 0.9em;
    font-weight: bold;
    box-shadow: 0 2px 4px rgba(0,0,0,0.3);
}
.discord-benefits-list {
    text-align: left;
    display: inline-block;
    margin: 10px 0;
    color: #ffffff;
}
.discord-benefits-list li {
    margin: 10px 0;
    position: relative;
    padding-left: 28px;
    color: #ffffff;
    font-weight: 500;
    text-shadow: 0 1px 2px rgba(0,0,0,0.2);
}
.discord-benefits-list li::before {
    content: "✨";
    position: absolute;
    left: 0;
    color: #FFD700;
}
.locked-option {
    opacity: 0.6;
    cursor: not-allowed;
}

/* Warning message styling */
.warning-message {
    background-color: #2a2a2a;
    border-left: 4px solid #ff7675;
    padding: 12px 15px;
    margin: 10px 0;
    border-radius: 0 8px 8px 0;
    box-shadow: 0 2px 4px rgba(0,0,0,0.2);
    color: #e0e0e0;
    font-weight: bold;
}

/* Example images and upload section styling */
.upload-section {
  display: flex;
  gap: 20px;
  margin: 20px 0;
}

.example-images-container {
  flex: 1;
}

.upload-container {
  flex: 1;
  display: flex;
  flex-direction: column;
  justify-content: center;
}

.section-title {
  font-weight: bold;
  margin-bottom: 10px;
  color: #7289DA;
}

.example-images-grid {
  display: grid;
  grid-template-columns: repeat(3, 1fr);
  gap: 10px;
}

.example-image-item {
  border-radius: 8px;
  overflow: hidden;
  cursor: pointer;
  transition: all 0.2s ease;
  border: 2px solid transparent;
}

.example-image-item:hover {
  transform: scale(1.05);
  box-shadow: 0 4px 8px rgba(0, 0, 0, 0.2);
}

.example-image-item.selected {
  border-color: #7289DA;
}

.upload-button {
  margin-top: 15px;
}
'''

def get_user_identifier(request: gr.Request) -> str:
    """Get user identifier from request (IP address)"""
    if request and hasattr(request, 'client') and hasattr(request.client, 'host'):
        return request.client.host
    return "unknown"

def get_rate_limit_status(request: gr.Request) -> str:
    """Get current rate limit status for display to user"""
    try:
        user_identifier = get_user_identifier(request)
        is_allowed, remaining, reset_time = check_rate_limit(user_identifier)
        
        if remaining == 0 and reset_time > 0:
            minutes = reset_time // 60
            seconds = reset_time % 60
            time_str = f"{minutes}m {seconds}s" if minutes > 0 else f"{seconds}s"
            return f"⚠️ Rate limit reached. Try again in {time_str}"
        elif remaining <= 2:
            return f"⚡ {remaining} generations remaining this hour"
        else:
            return f"✅ {remaining} generations remaining this hour"
    except:
        return "✅ Ready to generate"

with gr.Blocks(css=css, theme=gr.themes.Soft(primary_hue="indigo", neutral_hue="slate", text_size="lg")) as demo:
    selected_index = gr.State(None)
    current_generation_id = gr.State(None)
    
    # Updated title with Remade Canvas theme
    gr.Markdown("# Remade AI - Open Source Camera Controls")
    
    # Updated Remade Canvas callout
    gr.HTML(
        """
        <div class="discord-banner">
            <h3>🚀 Unlock 100s of AI Video Effects! 🎬</h3>
            <p>Access Remade Canvas with Veo, Kling, and hundreds of professional video effects. Create cinematic content with the most advanced AI video models!</p>
            <a href="https://app.remade.ai?utm_source=Huggingface&utm_medium=Social&utm_campaign=hugginface_space&utm_content=canvas_effects" target="_blank">Try Remade Canvas</a>
            <div style="margin-top: 15px; padding-top: 15px; border-top: 1px solid rgba(255,255,255,0.7);">
                <p style="font-size: 0.9em; margin-bottom: 10px;">Join our community for updates and tips:</p>
                <a href="https://remade.ai/join-discord?utm_source=Huggingface&utm_medium=Social&utm_campaign=hugginface_space&utm_content=canvas_effects" target="_blank" class="discord-community-btn">Discord Community</a>
            </div>
        </div>
        """
    )
    
    selected_info = gr.HTML("")

    with gr.Row():
        with gr.Column(scale=1):
            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="650px",
                object_fit="contain"
            )
            
            # Updated Discord/camera controls callout
            gr.HTML(
                """
                <div class="discord-feature">
                    <span class="discord-feature-title">🎬 Remade Canvas:</span> Access 100s of effects including Veo, Kling, and advanced camera controls beyond these samples!
                </div>
                """
            )
            
            gr.HTML('<div class="section-description">Click an example image or upload your own</div>')
            
            with gr.Row():
                with gr.Column(scale=1):
                    example_gallery = gr.Gallery(
                        [
                            ("https://storage.googleapis.com/remade-v2/huggingface_assets/image_fx%20(22).jpg", "Man with angel wings"),
                            ("https://storage.googleapis.com/remade-v2/huggingface_assets/image_fx%20(26).jpg", "Motorcyclist on the road"),
                            ("https://storage.googleapis.com/remade-v2/huggingface_assets/image_fx%20(27).jpg", "Superhero facing away in a tunnel"),
                            ("https://storage.googleapis.com/remade-v2/huggingface_assets/image_fx%20(75).jpg", "Girl with half her face underwater, staring at the camera"),
                            ("https://storage.googleapis.com/remade-v2/huggingface_assets/empire_state.jpg", "Workers sitting on construction at the top of Empire State Building"),
                            ("https://storage.googleapis.com/remade-v2/huggingface_assets/uploads_e6472106-4e9d-4620-b41b-a9bbe4893415.png", "Cartoon boy on bike")
                        ],
                        columns=3,
                        height="300px",
                        object_fit="cover"
                    )
                
                with gr.Column(scale=1):
                    image_input = gr.Image(type="filepath", label="")

            subject = gr.Textbox(label="Describe your subject", placeholder="Cat toy")
            
            # Rate limit status display
            rate_limit_status = gr.Markdown("✅ Ready to generate", elem_id="rate_limit_status")
            
            with gr.Row():
                button = gr.Button("Generate", variant="primary", elem_id="gen_btn")
                audio_button = gr.Button("Add Audio 🔒", interactive=False)
            
        with gr.Column(scale=1):
            warning_message = gr.HTML(
                """
                <div class="warning-message">
                    ⚠️ Please DO NOT refresh the page during generation. Processing camera controls takes time for best quality!
                </div>
                """,
                visible=True
            )
            
            gr.HTML(
                """
                <div class="discord-feature">
                    <span class="discord-feature-title">⚡ Remade Canvas:</span> Get faster generation speeds and access to Veo, Kling, and 100s of premium effects!
                </div>
                """
            )
            
            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]
    )
    
    # Modified function to handle example image selection
    def select_example_image(evt: gr.SelectData):
        """Handle example image selection and return image URL, description, and update image source"""
        example_images = [
            {
                "url": "https://storage.googleapis.com/remade-v2/huggingface_assets/image_fx%20(22).jpg",
                "description": "Man with angel wings"
            },
            {
                "url": "https://storage.googleapis.com/remade-v2/huggingface_assets/image_fx%20(26).jpg",
                "description": "Motorcyclist on the road"
            },
            {
                "url": "https://storage.googleapis.com/remade-v2/huggingface_assets/image_fx%20(27).jpg",
                "description": "Superhero facing away in a tunnel"
            },
            {
                "url": "https://storage.googleapis.com/remade-v2/huggingface_assets/image_fx%20(75).jpg",
                "description": "Girl with half her face underwater, staring at the camera"
            },
            {
                "url": "https://storage.googleapis.com/remade-v2/huggingface_assets/empire_state.jpg",
                "description": "Workers sitting on construction at the top of Empire State Building"
            },
            {
                "url": "https://storage.googleapis.com/remade-v2/huggingface_assets/uploads_e6472106-4e9d-4620-b41b-a9bbe4893415.png",
                "description": "Cartoon boy on bike"
            }
        ]
        
        selected = example_images[evt.index]
        
        # Return the URL, description, and update image source to "example"
        return selected["url"], selected["description"], "example"
    
    # Connect example gallery selection to image_input and subject
    example_gallery.select(
        fn=select_example_image,
        outputs=[image_input, subject]
    )
    
    # Add a custom handler to check if inputs are valid
    def check_inputs(subject, image_input, selected_index):
        if not selected_index:
            raise gr.Error("You must select a LoRA before proceeding.")
        if not subject.strip():
            raise gr.Error("Please describe your subject.")
        if image_input is None:
            raise gr.Error("Please upload an image or select an example image.")
    
    # Function to immediately disable button
    def start_generation():
        return gr.update(value="Generating...", interactive=False)
    
    # Use gr.on for the button click with validation
    button.click(
        fn=check_inputs,
        inputs=[subject, image_input, selected_index],
        outputs=None,
    ).success(
        fn=start_generation,
        inputs=None,
        outputs=[button]
    ).success(
        fn=handle_generation,
        inputs=[image_input, subject, selected_index],
        outputs=[output, current_generation_id, progress_bar, button]
    )
    
    # Add a click handler for the disabled audio button
    audio_button.click(
        fn=lambda: gr.Info("Try Remade Canvas to unlock audio generation and 100s of other effects!"),
        inputs=None,
        outputs=None
    )
    
    # Update rate limit status on page load
    demo.load(
        fn=get_rate_limit_status,
        inputs=None,
        outputs=[rate_limit_status]
    )

def get_modal_auth_headers():
    """Get authentication headers for Modal API requests"""
    if not MODAL_AUTH_TOKEN:
        raise ValueError("MODAL_AUTH_TOKEN environment variable not found")
    
    return {
        'Authorization': f'Bearer {MODAL_AUTH_TOKEN}',
        'Content-Type': 'application/json'
    }

def check_rate_limit(user_identifier: str) -> tuple[bool, int, int]:
    """
    Check if user has exceeded rate limit
    Returns: (is_allowed, remaining_generations, reset_time_seconds)
    """
    current_time = time.time()
    user_queue = user_generations[user_identifier]
    
    # Remove old entries outside the time window
    while user_queue and current_time - user_queue[0] > RATE_LIMIT_WINDOW:
        user_queue.popleft()
    
    # Check if user has exceeded limit
    if len(user_queue) >= RATE_LIMIT_GENERATIONS:
        # Calculate when the oldest entry will expire
        reset_time = int(user_queue[0] + RATE_LIMIT_WINDOW - current_time)
        return False, 0, reset_time
    
    remaining = RATE_LIMIT_GENERATIONS - len(user_queue)
    return True, remaining, 0

def record_generation(user_identifier: str):
    """Record a new generation for the user"""
    current_time = time.time()
    user_generations[user_identifier].append(current_time)

if __name__ == "__main__":
    demo.queue(default_concurrency_limit=20)
    demo.launch(ssr_mode=False, share=True)