app.py

import os
import re
from pathlib import Path
import json
import base64
from datetime import datetime

import gradio as gr
from fastapi import FastAPI, Request, Response
from fastapi.staticfiles import StaticFiles
import uvicorn
import httpx

import spaces
from spaces.zero.client import _get_token

# Set environment variables so Gradio builds its URLs.
os.environ["GRADIO_SSR_MODE"] = "True"
os.environ["GRADIO_SERVER_PORT"] = "7860"
os.environ["GRADIO_SERVER_NAME"] = "0.0.0.0"
# Tell Gradio what the node server is (host and port)
os.environ["GRADIO_NODE_SERVER_NAME"] = "127.0.0.1:7861"
os.environ["GRADIO_ROOT_PATH"] = "/"

# Create FastAPI app
app = FastAPI()

# (Optional) Mount static files
static_dir = Path("./static")
static_dir.mkdir(parents=True, exist_ok=True)
app.mount("/static", StaticFiles(directory="static"), name="static")
os.environ["GRADIO_ALLOWED_PATHS"] = str(static_dir.resolve())

# RESPONSE MIDDLEWARE TO FIX SSR ASSET URLS
@app.middleware("http")
async def fix_asset_urls(request: Request, call_next):
    response = await call_next(request)
    content_type = response.headers.get("content-type", "")
    # Process only HTML responses
    if "text/html" in content_type:
        # Read the entire response body
        body = b""
        async for chunk in response.body_iterator:
            body += chunk
        # Step 1: Insert a slash after ':7861' if not already present.
        fixed_body = re.sub(rb'(:7861)(?!/)', rb'\1/', body)
        # Step 2: Fix the _appimmutable segment:
        # Replace "_appimmutable" with "_app/immutable"
        fixed_body = re.sub(rb'(_app)(immutable)', rb'\1/immutable', fixed_body)
        # Step 3: Ensure that after "_app/immutable" there is a slash.
        fixed_body = re.sub(rb'(_app/immutable)(?!/)', rb'\1/', fixed_body)
        # Step 4: If "assets" is immediately followed by an uppercase letter, insert a slash.
        fixed_body = re.sub(rb'(assets)(?=[A-Z])', rb'\1/', fixed_body)
        # Build a new HTML response using the fixed body.
        response = Response(content=fixed_body,
                            status_code=response.status_code,
                            headers=dict(response.headers),
                            media_type="text/html")
    return response

@spaces.GPU(duration=4*60)
def process_text(text):
    return text.upper()

def process_and_save(request: gr.Request, text):
    token = _get_token(request)
    payload = token.split('.')[1]
    payload = f"{payload}{'=' * ((4 - len(payload) % 4) % 4)}"
    payload = json.loads(base64.urlsafe_b64decode(payload).decode())
    print(f"Token payload: {payload}")
    processed_text = process_text(text)
    return processed_text

process_and_save.zerogpu = True

with gr.Blocks() as demo:
    text_input = gr.Textbox(label="Enter some text")
    submit_btn = gr.Button("Process and Download")
    output = gr.Textbox(label="Output")
    submit_btn.click(fn=process_and_save, inputs=[text_input], outputs=output)

# Mount the Gradio app into FastAPI using SSR mode (node_port 7861).
app = gr.mount_gradio_app(app, demo, path="/", ssr_mode=True, node_port=7861)

if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=7860)