app.py

import gradio as gr
import requests
import os
import json
import time
from pathlib import Path
from datetime import datetime, timedelta
import io
import base64

# Configuration
API_URL = os.getenv("MIRAGIC_API_URL")
MAX_FREE_TRIALS = 2
ATTEMPTS_FILE = "user_attempts.json"

# Debug: Print API URL to check if it's loaded
print(f"API URL loaded: {API_URL is not None}")

class IPTracker:
    def __init__(self):
        self.attempts_file = ATTEMPTS_FILE
        self.attempts = self.load_attempts()
        
    def load_attempts(self):
        try:
            if Path(self.attempts_file).exists():
                with open(self.attempts_file, "r") as f:
                    data = json.load(f)
                    # Clean old entries (older than 24 hours)
                    now = datetime.now()
                    cleaned_data = {}
                    for ip, info in data.items():
                        try:
                            last_attempt = datetime.fromisoformat(info.get("last_attempt", "2000-01-01"))
                            if now - last_attempt < timedelta(hours=24):
                                cleaned_data[ip] = info
                        except:
                            continue
                    return cleaned_data
        except Exception as e:
            print(f"Error loading attempts: {e}")
        return {}
    
    def save_attempts(self):
        try:
            with open(self.attempts_file, "w") as f:
                json.dump(self.attempts, f, indent=2)
        except Exception as e:
            print(f"Error saving attempts: {e}")
    
    def get_attempts(self, ip):
        if ip not in self.attempts:
            return 0
        
        try:
            # Check if 24 hours have passed
            last_attempt = datetime.fromisoformat(self.attempts[ip]["last_attempt"])
            if datetime.now() - last_attempt > timedelta(hours=24):
                self.attempts[ip] = {"count": 0, "last_attempt": datetime.now().isoformat()}
                self.save_attempts()
                return 0
            
            return self.attempts[ip]["count"]
        except:
            return 0
    
    def increment_attempts(self, ip):
        if ip not in self.attempts:
            self.attempts[ip] = {"count": 0, "last_attempt": datetime.now().isoformat()}
        
        self.attempts[ip]["count"] += 1
        self.attempts[ip]["last_attempt"] = datetime.now().isoformat()
        self.save_attempts()
        return self.attempts[ip]["count"]

# Initialize IP tracker
ip_tracker = IPTracker()

def get_client_ip(request: gr.Request):
    """Get client IP address from request"""
    if not request:
        return "unknown"
    
    # Try different headers for getting real IP
    headers_to_check = [
        "X-Forwarded-For",
        "X-Real-IP", 
        "CF-Connecting-IP",
        "X-Forwarded-Proto"
    ]
    
    for header in headers_to_check:
        value = request.headers.get(header)
        if value:
            # For X-Forwarded-For, take the first IP
            if header == "X-Forwarded-For":
                return value.split(",")[0].strip()
            return value
    
    # Fallback to client host
    try:
        if hasattr(request, 'client') and request.client:
            return request.client.host
    except:
        pass
    
    return "unknown"

def call_speedpainting_api(image):
    """Call the Miragic API with proper error handling for Hugging Face Spaces"""
    if not API_URL:
        raise gr.Error("API URL not configured. Please check your secrets in Space settings.")
    
    try:
        # Convert PIL image to bytes
        img_byte_arr = io.BytesIO()
        image.save(img_byte_arr, format='PNG')
        img_byte_arr.seek(0)
        
        # Prepare files for upload
        files = {
            'file': ('image.png', img_byte_arr, 'image/png')
        }
        
        headers = {
            'User-Agent': 'HuggingFace-Spaces/1.0'
        }
        
        print(f"Calling API: {API_URL}")
        
        # Make the API request
        response = requests.post(
            API_URL,
            files=files,
            headers=headers,
            timeout=120
        )
        
        print(f"API Response Status: {response.status_code}")
        
        if response.status_code == 200:
            try:
                result = response.json()
                print(f"API Response: {result}")
                
                if result.get("status") == "success":
                    video_url = result.get("link")
                    if video_url:
                        # Handle different URL formats
                        if "dl=0" in video_url:
                            return video_url.replace("dl=0", "raw=1")
                        elif "?dl=" in video_url:
                            return video_url.split("?dl=")[0] + "?raw=1"
                        return video_url
                    else:
                        raise gr.Error("No video URL in API response")
                else:
                    error_msg = result.get('message', 'Unknown error from API')
                    raise gr.Error(f"API error: {error_msg}")
                    
            except json.JSONDecodeError as e:
                print(f"JSON decode error: {e}")
                print(f"Response text: {response.text}")
                raise gr.Error("Invalid response from API")
        else:
            print(f"API Error Response: {response.text}")
            if response.status_code == 413:
                raise gr.Error("Image file too large. Please try a smaller image.")
            elif response.status_code == 429:
                raise gr.Error("API rate limit exceeded. Please try again later.")
            else:
                raise gr.Error(f"API request failed (Status: {response.status_code})")
                
    except requests.exceptions.Timeout:
        raise gr.Error("Request timeout. The API is taking too long to respond.")
    except requests.exceptions.ConnectionError:
        raise gr.Error("Cannot connect to API. Please check your internet connection.")
    except requests.exceptions.RequestException as e:
        print(f"Request Exception: {str(e)}")
        raise gr.Error(f"Network error: {str(e)}")
    except Exception as e:
        print(f"Unexpected Error: {str(e)}")
        raise gr.Error(f"Unexpected error: {str(e)}")

def generate_speedpainting(image, request: gr.Request):
    """Generate speedpainting with IP-based limiting"""
    if not image:
        raise gr.Error("Please upload an image first!")
    
    client_ip = get_client_ip(request)
    print(f"Client IP: {client_ip}")
    
    current_attempts = ip_tracker.get_attempts(client_ip)
    print(f"Current attempts for {client_ip}: {current_attempts}")
    
    if current_attempts >= MAX_FREE_TRIALS:
        raise gr.Error(
            f"You've used {MAX_FREE_TRIALS} free generations today. "
            f"Please visit https://miragic.ai/ to sign up for unlimited access!"
        )
    
    try:
        # Process the image
        video_url = call_speedpainting_api(image)
        
        # Increment attempts after successful generation
        new_count = ip_tracker.increment_attempts(client_ip)
        print(f"New attempt count for {client_ip}: {new_count}")
        
        return video_url
    
    except Exception as e:
        print(f"Error in generate_speedpainting: {str(e)}")
        raise e

def get_remaining_attempts(request: gr.Request):
    """Get remaining attempts for current IP"""
    client_ip = get_client_ip(request)
    current_attempts = ip_tracker.get_attempts(client_ip)
    remaining = MAX_FREE_TRIALS - current_attempts
    
    print(f"Getting remaining attempts for {client_ip}: {remaining}")
    
    if remaining <= 0:
        return f"Daily limit reached ({MAX_FREE_TRIALS}/{MAX_FREE_TRIALS}). Sign up for unlimited access!"
    else:
        return f"Remaining free generations: {remaining}/{MAX_FREE_TRIALS}"

# Example images
EXAMPLE_IMAGES = [
    ["static/assets/1.jpg"],
    ["static/assets/2.jpg"],
    ["static/assets/3.jpg"],
    ["static/assets/4.jpg"],
    ["static/assets/5.jpg"],
    ["static/assets/6.jpg"],
]

# Company information
COMPANY_INFO = """
## About US
### Miragic is a cutting-edge platform to serve AI-powered tools like Virtual Try-on, Speed Painting, Background Remover and Sales Pilot.
### [Visit our website](https://miragic.ai) to explore more innovative generative AI tools!
"""

# Custom CSS
css = """
footer {visibility: hidden}
.banner {
    background-color: #f8f9fa;
    padding: 10px;
    border-radius: 5px;
    margin-bottom: 20px;
    text-align: center;
}
.button-gradient {
  background: linear-gradient(45deg, #ff416c, #ff4b2b, #ff9b00, #ff416c);
  background-size: 400% 400%;
  border: none;
  padding: 14px 28px;
  font-size: 16px;
  font-weight: bold;
  color: white;
  border-radius: 10px;
  cursor: pointer;
  transition: 0.3s ease-in-out;
  animation: gradientAnimation 2s infinite linear;
  box-shadow: 0 4px 10px rgba(255, 65, 108, 0.6);
}
@keyframes gradientAnimation {
  0% { background-position: 0% 50%; }
  100% { background-position: 100% 50%; }
}
.button-gradient:hover {
  transform: scale(1.05);
  box-shadow: 0 6px 15px rgba(255, 75, 43, 0.8);
}
.signup-container {
    text-align: center;
    padding: 20px;
    background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
    border-radius: 8px;
    margin-top: 20px;
    color: white;
    box-shadow: 0 4px 15px rgba(0,0,0,0.2);
}
.signup-container h3 {
    margin-bottom: 10px;
    color: white;
}
.signup-container p {
    margin-bottom: 15px;
    color: #f0f0f0;
}
.signup-button {
  background: linear-gradient(45deg, #ff416c, #ff4b2b);
  border: none;
  padding: 12px 25px;
  font-size: 16px;
  font-weight: bold;
  color: white;
  border-radius: 8px;
  text-decoration: none;
  display: inline-block;
  transition: all 0.3s ease;
  box-shadow: 0 4px 15px rgba(0, 0, 0, 0.2);
}
.signup-button:hover {
  transform: translateY(-2px);
  box-shadow: 0 6px 20px rgba(0, 0, 0, 0.3);
}
.attempts-counter {
    background: #e3f2fd;
    padding: 10px;
    border-radius: 5px;
    margin: 10px 0;
    text-align: center;
    font-weight: bold;
    color: #1976d2;
}
"""

# Create Gradio interface
with gr.Blocks(title="Miragic Speed-Painting", theme=gr.themes.Ocean(), css=css) as demo:
    gr.Markdown("""
    <div style="display: flex; align-items: center;">
        <img src="https://avatars.githubusercontent.com/u/211682198?s=200&v=4" style="width: 80px; margin-right: 20px;"/>
        <div>
            <h1 style="margin-bottom: 0;">Miragic Speed-Painting 🎨</h1>
            <p>Upload an image to see AI create speedpainting animations!</p>
        </div>
    </div>
    """)
    
    gr.Markdown(COMPANY_INFO)
    
    # Usage counter
    usage_display = gr.HTML(elem_classes="attempts-counter")
    
    with gr.Row():
        with gr.Column():
            image_input = gr.Image(
                label="Upload Image",
                type="pil",
                sources=["upload", "clipboard"],
                height=300
            )
            
            gr.Examples(
                examples=EXAMPLE_IMAGES,
                inputs=image_input,
                label="Try these examples!",
                examples_per_page=6
            )
            
            submit_btn = gr.Button("Generate Speedpainting 🚀", elem_classes="button-gradient")
        
        with gr.Column():
            video_output = gr.Video(
                label="Speedpainting Result",
                autoplay=True,
                height=300
            )
            
            signup_prompt = gr.HTML(
                visible=True,
                value="""<div class="signup-container">
                    <h3>🚀 Want unlimited generations?</h3>
                    <p>Please sign up at Miragic.ai for unlimited access to all our AI tools!</p>
                    <a href='https://miragic.ai/products/speed-painting' target='_blank' class="signup-button">
                        SignUp for Free 🚀
                    </a>
                </div>"""
            )
    
    # Update usage display on page load
    demo.load(
        fn=get_remaining_attempts,
        outputs=usage_display
    )
    
    # Handle generation
    submit_btn.click(
        fn=generate_speedpainting,
        inputs=[image_input],
        outputs=video_output
    ).then(
        fn=get_remaining_attempts,
        outputs=usage_display
    )

    gr.HTML('<a href="https://visitorbadge.io/status?path=https%3A%2F%2Fhuggingface.co%2Fspaces%2FMiragic-AI%2FMiragic-Speed-Painting"><img src="https://api.visitorbadge.io/api/combined?path=https%3A%2F%2Fhuggingface.co%2Fspaces%2FMiragic-AI%2FMiragic-Speed-Painting&label=VISITORS&labelColor=%2337d67a&countColor=%23f47373&style=plastic&labelStyle=upper" /></a>')

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