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'' else: HERO_LOGO_HTML = '' 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"""
✦
{title}
{subtitle}
""" 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"""
Download
Generated image
""" 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"""
LIVE
{HERO_LOGO_HTML}

Generate Images

Preview ComfyUI workflows instantly with Solricks.

""" ) gr.HTML('
PROMPT
') 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('
STYLE
') 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('
RESOLUTION
') 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( '
✦Result Preview
' ) 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, )