app.py

import os
import time
import base64
from pathlib import Path
import gradio as gr
import requests

API_BASE_URL = os.getenv("SOLRICKS_API_BASE_URL", "https://api.solricks.com").rstrip("/")
API_KEY = os.getenv("SOLRICKS_SPACE_API_KEY", "")
API_ENDPOINT = f"{API_BASE_URL}/api/workflows/visual-ai/run"

REQUEST_TIMEOUT = int(os.getenv("SOLRICKS_REQUEST_TIMEOUT", "600"))
FETCH_TIMEOUT = int(os.getenv("SOLRICKS_FETCH_TIMEOUT", "180"))

ALLOWED_STYLES = ["Realistic", "Portrait", "Cinematic", "Product"]
ALLOWED_RESOLUTIONS = ["1024x1024", "768x1344", "1344x768"]
MAX_PROMPT_LENGTH = 500

APP_DIR = Path(__file__).resolve().parent
CSS_PATH = APP_DIR / "image-generator.css"
LOGO_PATH = APP_DIR / "solricks-logo.png"
custom_css = CSS_PATH.read_text(encoding="utf-8") if CSS_PATH.exists() else ""


def file_to_data_uri(file_path: Path, mime_type: str) -> str:
    encoded = base64.b64encode(file_path.read_bytes()).decode("utf-8")
    return f"data:{mime_type};base64,{encoded}"


if LOGO_PATH.exists():
    LOGO_DATA_URI = file_to_data_uri(LOGO_PATH, "image/png")
    HERO_LOGO_HTML = f'<img class="sol-logo" src="{LOGO_DATA_URI}" alt="SOLRICKS" />'
else:
    HERO_LOGO_HTML = '<div class="sol-text-logo">SOLRICKS</div>'


def render_state_html(
    title: str = "Result Preview",
    subtitle: str = "Your result will appear here.",
    processing: bool = False
) -> str:
    processing_class = " is-processing" if processing else ""
    return f"""
<div class="sol-preview-state{processing_class}">
    <div class="sol-preview-icon-box"><div class="sol-preview-icon">✦</div></div>
    <div class="sol-preview-title">{title}</div>
    <div class="sol-preview-subtitle">{subtitle}</div>
</div>
"""


EMPTY_IMAGE_HTML = render_state_html()
PROCESSING_IMAGE_HTML = render_state_html(subtitle="Processing...", processing=True)


def normalize_image_url(image_url: str) -> str:
    image_url = (image_url or "").strip()

    if not image_url:
        raise ValueError("Empty image URL")

    if image_url.startswith("data:image/"):
        return image_url

    if image_url.startswith("/"):
        return f"{API_BASE_URL}{image_url}"

    return image_url


def fetch_image_as_data_uri(image_url: str) -> tuple[str, str, str]:
    image_url = normalize_image_url(image_url)

    if image_url.startswith("data:image/"):
        mime_type = image_url.split(";", 1)[0].replace("data:", "") or "image/png"
        extension = "png"

        if "jpeg" in mime_type or "jpg" in mime_type:
            extension = "jpg"
        elif "webp" in mime_type:
            extension = "webp"

        return image_url, mime_type, extension

    headers = {}

    if API_KEY:
        headers["Authorization"] = f"Bearer {API_KEY}"

    last_error = None

    for attempt in range(6):
        try:
            response = requests.get(
                image_url,
                headers=headers,
                timeout=(10, FETCH_TIMEOUT)
            )

            if response.status_code in [404, 409, 425, 429, 500, 502, 503, 504]:
                last_error = f"Image not ready yet. HTTP {response.status_code}"
                print(f"IMAGE FETCH RETRY {attempt + 1}/6:", last_error, flush=True)
                time.sleep(2 + attempt * 3)
                continue

            response.raise_for_status()

            mime_type = (
                response.headers.get("content-type") or "image/png"
            ).split(";", 1)[0].strip().lower()

            if not mime_type.startswith("image/"):
                raise RuntimeError(f"Returned file is not an image. Content-Type: {mime_type}")

            extension = "png"

            if "jpeg" in mime_type or "jpg" in mime_type:
                extension = "jpg"
            elif "webp" in mime_type:
                extension = "webp"

            encoded = base64.b64encode(response.content).decode("utf-8")

            return f"data:{mime_type};base64,{encoded}", mime_type, extension

        except Exception as error:
            last_error = error
            print(f"IMAGE FETCH ERROR {attempt + 1}/6:", error, flush=True)
            time.sleep(2 + attempt * 3)

    raise RuntimeError(f"Could not fetch generated image after retries: {last_error}")


def render_image_html(image_data_uri: str | None = None, extension: str = "png") -> str:
    if not image_data_uri:
        return EMPTY_IMAGE_HTML

    filename = f"solricks-generated-image.{extension}"

    return f"""
<div class="sol-image-wrap">
    <div class="sol-image-actions"><a href="{image_data_uri}" download="{filename}">Download</a></div>
    <img src="{image_data_uri}" alt="Generated image" />
</div>
"""


def generate_image(prompt: str, style: str, resolution: str):
    prompt = (prompt or "").strip()

    if not prompt:
        return render_state_html(subtitle="Please enter a prompt.")

    if len(prompt) > MAX_PROMPT_LENGTH:
        return render_state_html(
            subtitle=f"Prompt is too long. Keep it under {MAX_PROMPT_LENGTH} characters."
        )

    if style not in ALLOWED_STYLES:
        return render_state_html(subtitle="Invalid style selection.")

    if resolution not in ALLOWED_RESOLUTIONS:
        return render_state_html(subtitle="Invalid resolution selection.")

    payload = {
        "prompt": prompt,
        "style": style,
        "resolution": resolution,
        "enhancePrompt": False,
        "seedMode": "Random",
        "source": "huggingface-space",
    }

    headers = {"Content-Type": "application/json"}

    if API_KEY:
        headers["Authorization"] = f"Bearer {API_KEY}"

    try:
        response = requests.post(
            API_ENDPOINT,
            json=payload,
            headers=headers,
            timeout=(10, REQUEST_TIMEOUT)
        )

        try:
            data = response.json()
        except Exception:
            data = {}

        if response.status_code == 429:
            return render_state_html(
                subtitle="Too many requests. Please wait a little and try again."
            )

        if not response.ok:
            message = (
                data.get("error")
                or data.get("detail")
                or data.get("message")
                or f"Backend HTTP {response.status_code}: {response.text[:300]}"
            )

            print("BACKEND ERROR:", response.status_code, response.text[:1000], flush=True)

            return render_state_html(subtitle=message)

        image_url = data.get("image_url") or data.get("imageUrl") or data.get("url")

        if not image_url:
            return render_state_html(subtitle="The backend did not return an image.")

        image_data_uri, _mime_type, extension = fetch_image_as_data_uri(image_url)

        return render_image_html(image_data_uri, extension)

    except requests.exceptions.Timeout:
        return render_state_html(subtitle="Generation timed out. Please try again.")

    except Exception as error:
        print("SPACE ERROR:", error, flush=True)
        return render_state_html(subtitle=f"Error: {error}")


def clear_prompt() -> str:
    return ""


def begin_processing() -> str:
    return PROCESSING_IMAGE_HTML


sol_theme = gr.themes.Soft(
    primary_hue="blue",
    secondary_hue="indigo",
    neutral_hue="slate"
)


def build_demo() -> gr.Blocks:
    with gr.Blocks(
        title="SOLRICKS Image Generator",
        theme=sol_theme,
        css=custom_css
    ) as demo:
        with gr.Column(elem_id="sol-shell"):
            with gr.Row(equal_height=True, elem_id="sol-main-grid"):
                with gr.Column(scale=11, min_width=380, elem_id="sol-left-col"):
                    with gr.Group(elem_id="sol-left-card", elem_classes=["sol-panel"]):
                        gr.HTML(
                            f"""
                            <div class="sol-top-strip">
                              <span class="sol-mini-live"><span class="sol-live-dot"></span> LIVE</span>
                            </div>
                            <div class="sol-header-area">
                              <div class="sol-logo-holder">{HERO_LOGO_HTML}</div>
                              <h1 class="sol-title">Generate Images</h1>
                              <p class="sol-subtitle">Preview ComfyUI workflows instantly with Solricks.</p>
                            </div>
                            """
                        )

                        gr.HTML('<div class="sol-field-label">PROMPT</div>')

                        prompt = gr.Textbox(
                            label=None,
                            placeholder="Describe the scene, subject, and visual style...",
                            lines=5,
                            max_lines=6,
                            max_length=MAX_PROMPT_LENGTH,
                            show_label=False,
                            container=False,
                            elem_id="sol-prompt-box",
                            elem_classes=["sol-plain-input"],
                        )

                        with gr.Row(elem_id="sol-two-col-row"):
                            with gr.Column(scale=1, min_width=180):
                                gr.HTML('<div class="sol-field-label">STYLE</div>')

                                style = gr.Dropdown(
                                    label=None,
                                    choices=ALLOWED_STYLES,
                                    value="Cinematic",
                                    show_label=False,
                                    container=False,
                                    filterable=False,
                                    elem_id="sol-style-dropdown",
                                    elem_classes=["sol-plain-select"],
                                )

                            with gr.Column(scale=1, min_width=180):
                                gr.HTML('<div class="sol-field-label">RESOLUTION</div>')

                                resolution = gr.Dropdown(
                                    label=None,
                                    choices=ALLOWED_RESOLUTIONS,
                                    value="1024x1024",
                                    show_label=False,
                                    container=False,
                                    filterable=False,
                                    elem_id="sol-resolution-dropdown",
                                    elem_classes=["sol-plain-select"],
                                )

                        with gr.Row(elem_id="sol-action-row"):
                            clear_button = gr.Button(
                                "Clear",
                                variant="secondary",
                                elem_id="sol-clear-button",
                            )

                            generate_button = gr.Button(
                                "Generate",
                                variant="primary",
                                elem_id="sol-generate-button",
                            )

                with gr.Column(scale=11, min_width=380, elem_id="sol-right-col"):
                    with gr.Group(elem_id="sol-right-card", elem_classes=["sol-panel"]):
                        gr.HTML(
                            '<div id="sol-output-header"><div id="sol-output-title"><span>✦</span><span>Result Preview</span></div></div>'
                        )

                        output_image = gr.HTML(
                            value=render_image_html(),
                            elem_id="sol-output-image"
                        )

            clear_button.click(
                fn=clear_prompt,
                inputs=None,
                outputs=prompt,
                queue=False,
                show_progress="hidden",
            )

            generate_button.click(
                fn=begin_processing,
                inputs=None,
                outputs=output_image,
                queue=False,
                show_progress="hidden",
            ).then(
                fn=generate_image,
                inputs=[prompt, style, resolution],
                outputs=output_image,
                queue=True,
                show_progress="hidden",
            )

    return demo


if __name__ == "__main__":
    demo = build_demo()

    demo.queue(
        default_concurrency_limit=1,
        max_size=10
    )

    demo.launch(
        server_name="0.0.0.0",
        server_port=int(os.environ.get("PORT", 7860)),
        show_error=True,
        ssr_mode=False,
    )