Spaces:
Running on Zero
Running on Zero
| """MOSS-Image as a gr.Workflow app — node-canvas text-to-image pipeline. | |
| A reimagining of https://huggingface.co/spaces/OpenMOSS-Team/MOSS-Image as a | |
| Gradio Workflow (gradio 6.30.0): generation is a bound Python-function node on | |
| a drag-and-drop canvas, fed by reference nodes (prompt, size, and every | |
| sampling knob of the original demo) and writing into an output-gallery subject. | |
| Edit the graph on the canvas or in workflow.json; the App view and REST API | |
| stay in sync automatically. | |
| Both checkpoints are loaded once at module scope (ZeroGPU packs and transfers | |
| them when a GPU is actually allocated) and shared by every run of the node. | |
| """ | |
| import os | |
| import secrets | |
| import tempfile | |
| import time | |
| from collections import OrderedDict | |
| from pathlib import Path | |
| # `spaces` must be imported before anything that initializes CUDA. | |
| import spaces # ty: ignore[unresolved-import] | |
| from gradio_client import utils as client_utils | |
| import gradio as gr | |
| from gradio import workflow_api as _wf_api | |
| # Workaround for gradio 6.30.0: the App view hands a subject's value to its | |
| # output component via `_from_output`, which unwraps a top-level file dict to | |
| # its path but leaves dicts nested in a LIST untouched — and gr.Gallery only | |
| # accepts path strings, not file dicts. A gallery subject fed by an fn node | |
| # (which must return file dicts for the canvas) therefore crashes the App | |
| # view. Unwrap list items the same way so both views work. | |
| _orig_from_output = _wf_api._from_output | |
| def _from_output_unwrapping_lists(value): | |
| if isinstance(value, list): | |
| return [_orig_from_output(item) for item in value] | |
| return _orig_from_output(value) | |
| _wf_api._from_output = _from_output_unwrapping_lists | |
| def _install_component_server_keepalive(): | |
| """Keep `/gradio_api/component_server` connections alive during long fn calls. | |
| Canvas fn nodes call back through the component_server route — a plain | |
| POST that awaits the bound function synchronously, unlike App-view | |
| endpoints, which run through the queue with SSE heartbeats. A MOSS-Image | |
| generation takes 45-130s, and the silent connection is dropped by the | |
| proxy chain, surfacing on the canvas as "Connection errored out" (the | |
| App view is unaffected). Wrap the route so it streams whitespace | |
| heartbeat bytes — legal leading whitespace in JSON — until the real | |
| result is ready. | |
| """ | |
| import asyncio | |
| import logging | |
| from fastapi import Request | |
| from fastapi.dependencies.utils import get_dependant | |
| from fastapi.routing import APIRoute | |
| from starlette.responses import StreamingResponse | |
| from starlette.routing import request_response | |
| from gradio.routes import App as _RoutesApp | |
| original_create_app = _RoutesApp.create_app | |
| def _wrap(endpoint): | |
| async def component_server_with_keepalive(request: Request): | |
| task = asyncio.create_task(endpoint(request)) | |
| async def body(): | |
| try: | |
| while not task.done(): | |
| yield b" " | |
| await asyncio.sleep(10) | |
| payload = json.dumps(task.result(), default=str).encode() | |
| except Exception as exc: # keep the canvas error path readable | |
| payload = json.dumps({"error": str(exc)}).encode() | |
| yield payload | |
| return StreamingResponse(body(), media_type="application/json") | |
| return component_server_with_keepalive | |
| def create_app(blocks, *args, **kwargs): | |
| app = original_create_app(blocks, *args, **kwargs) | |
| for route in app.routes: | |
| if isinstance(route, APIRoute) and route.path.rstrip("/").endswith( | |
| "/component_server" | |
| ): | |
| route.endpoint = _wrap(route.endpoint) | |
| dependant = get_dependant( | |
| path=route.path_format, endpoint=route.endpoint | |
| ) | |
| # Keep the route-level dependencies (e.g. Depends(login_check)). | |
| dependant.dependencies = ( | |
| route.dependant.dependencies + dependant.dependencies | |
| ) | |
| route.dependant = dependant | |
| route.app = request_response(route.get_route_handler()) | |
| logging.getLogger(__name__).info( | |
| "component_server keep-alive wrapper installed" | |
| ) | |
| return app | |
| _RoutesApp.create_app = create_app | |
| _install_component_server_keepalive() | |
| MODEL_IDS = { | |
| "base": "OpenMOSS-Team/MOSS-Image-Base", | |
| "sft": "OpenMOSS-Team/MOSS-Image-SFT", | |
| } | |
| # Set MOSS_WORKFLOW_PREVIEW=1 to open the canvas without downloading weights | |
| # (useful for local UI work on a CPU-only machine). | |
| PREVIEW = os.environ.get("MOSS_WORKFLOW_PREVIEW") == "1" | |
| def _load_pipeline(model_id, shared_pipe=None): | |
| """Load one MOSS-Image checkpoint with the bundled inference classes. | |
| Text encoder, tokenizer and VAE are shared with `shared_pipe` when their | |
| Hub file hashes match — this is what keeps both checkpoints inside a | |
| ZeroGPU worker's GPU memory (two full copies do not fit). | |
| """ | |
| import torch | |
| from huggingface_hub import HfApi, snapshot_download | |
| from pipeline_moss_image import ( | |
| MOSSImagePipeline, | |
| MOSSImageScheduler, | |
| MOSSImageTransformer, | |
| ) | |
| info = HfApi().model_info(model_id, files_metadata=True) | |
| revision = info.sha | |
| shared_hashes = { | |
| item.rfilename: item.lfs.sha256 if item.lfs else item.blob_id | |
| for item in info.siblings | |
| if item.rfilename.startswith(("text_encoder/", "tokenizer/", "vae/")) | |
| } | |
| components = {} | |
| if shared_pipe is not None: | |
| for name in ("text_encoder", "tokenizer", "vae"): | |
| current = {k: v for k, v in shared_hashes.items() if k.startswith(f"{name}/")} | |
| previous = { | |
| k: v for k, v in shared_pipe._shared_component_hashes.items() | |
| if k.startswith(f"{name}/") | |
| } | |
| if current and all(current.values()) and current == previous: | |
| components[name] = getattr(shared_pipe, name) | |
| snapshot = snapshot_download( | |
| model_id, revision=revision, ignore_patterns=[f"{name}/*" for name in components] | |
| ) | |
| pipe = MOSSImagePipeline.from_pretrained( | |
| snapshot, | |
| custom_pipeline=str(Path(__file__).with_name("pipeline_moss_image")), | |
| trust_remote_code=True, | |
| transformer=MOSSImageTransformer.from_pretrained( | |
| snapshot, subfolder="transformer", torch_dtype=torch.bfloat16 | |
| ), | |
| scheduler=MOSSImageScheduler.from_pretrained(snapshot, subfolder="scheduler"), | |
| local_files_only=True, | |
| torch_dtype={"default": torch.bfloat16, "vae": torch.float32}, | |
| **components, | |
| ) | |
| # On ZeroGPU the `spaces` package intercepts this move, packs the weights, | |
| # and transfers them when a GPU is allocated — module scope costs no GPU time. | |
| pipe.to("cuda") | |
| pipe._prompt_cache = OrderedDict() | |
| pipe._shared_component_hashes = shared_hashes | |
| return pipe | |
| PIPELINES = {} | |
| if not PREVIEW: | |
| PIPELINES["base"] = _load_pipeline(MODEL_IDS["base"]) | |
| PIPELINES["sft"] = _load_pipeline(MODEL_IDS["sft"], shared_pipe=PIPELINES["base"]) | |
| def _as_media(path: str) -> dict: | |
| """The dict shape the workflow canvas renders media from.""" | |
| return { | |
| "path": path, | |
| "url": f"/gradio_api/file={client_utils.encode_file_path(path)}", | |
| "is_file": True, | |
| } | |
| def _parse_layers(value: str, name: str) -> tuple: | |
| """'22' -> (22,); '' -> (). Same rules as the original demo.""" | |
| if not (value or "").strip(): | |
| return () | |
| try: | |
| layers = tuple(int(item.strip()) for item in value.split(",")) | |
| except ValueError: | |
| raise gr.Error( | |
| f"{name} must be comma-separated integer indices, e.g. 22." | |
| ) from None | |
| if any(layer < 0 for layer in layers) or len(set(layers)) != len(layers): | |
| raise gr.Error(f"{name} indices must be nonnegative and unique.") | |
| return layers | |
| def generate( | |
| prompt: str, | |
| negative_prompt: str, | |
| model: str, | |
| width: float, | |
| height: float, | |
| steps: float, | |
| cfg: float, | |
| shift: float, | |
| seed: float, | |
| num_images: float, | |
| max_sequence_length: float, | |
| uncond_skip_layers: str, | |
| slg_layers: str, | |
| slg_scale: float, | |
| slg_start: float, | |
| slg_stop: float, | |
| ) -> list: | |
| """Generate images with MOSS-Image (Base or SFT). Full sampling controls.""" | |
| import torch | |
| if PREVIEW: | |
| raise gr.Error("Preview mode: models are not loaded (MOSS_WORKFLOW_PREVIEW=1).") | |
| if not (prompt or "").strip(): | |
| raise gr.Error("Enter a prompt.") | |
| key = (model or "").strip().lower().replace("moss-image-", "") | |
| if key not in PIPELINES: | |
| raise gr.Error(f"Unknown model {model!r}. Choose 'base' or 'sft'.") | |
| width, height, steps, num_images, max_sequence_length = ( | |
| round(width), round(height), round(steps), round(num_images), | |
| round(max_sequence_length), | |
| ) | |
| seed = round(seed) | |
| if width <= 0 or height <= 0 or width % 16 or height % 16: | |
| raise gr.Error("Width and height must be positive multiples of 16.") | |
| if steps < 1 or num_images < 1: | |
| raise gr.Error("Steps and image count must be positive.") | |
| if not 0 <= cfg or shift <= 0: | |
| raise gr.Error("CFG must be nonnegative; shift must be positive.") | |
| if not 1 <= max_sequence_length <= 4096: | |
| raise gr.Error("Maximum text length must be between 1 and 4096.") | |
| if slg_scale < 0 or not 0 <= slg_start <= slg_stop <= 1: | |
| raise gr.Error("Require SLG scale >= 0 and 0 <= SLG start <= SLG stop <= 1.") | |
| if seed == -1: | |
| seed = secrets.randbelow(2**32) | |
| if not 0 <= seed < 2**53: | |
| raise gr.Error("Seed must be -1 (random) or an integer from 0 to 2^53 - 1.") | |
| pipe = PIPELINES[key] | |
| depth = pipe.transformer.config.num_layers | |
| skip = _parse_layers(uncond_skip_layers, "Skip layers") | |
| slg = _parse_layers(slg_layers, "SLG layers") | |
| if any(layer >= depth for layer in skip + slg): | |
| raise gr.Error(f"Skip-layer indices must be below {depth}.") | |
| pipe.scheduler = type(pipe.scheduler).from_config(pipe.scheduler.config, shift=shift) | |
| start = time.monotonic() | |
| try: | |
| images = pipe( | |
| prompt=prompt, | |
| negative_prompt=negative_prompt or "", | |
| width=width, | |
| height=height, | |
| num_inference_steps=steps, | |
| guidance_scale=float(cfg), | |
| num_images_per_prompt=num_images, | |
| max_sequence_length=max_sequence_length, | |
| uncond_skip_layers=skip, | |
| slg_layers=slg, | |
| slg_scale=float(slg_scale), | |
| slg_start=float(slg_start), | |
| slg_stop=float(slg_stop), | |
| generator=torch.Generator("cpu").manual_seed(seed), | |
| ).images | |
| except torch.cuda.OutOfMemoryError as exc: | |
| torch.cuda.empty_cache() | |
| raise gr.Error("GPU memory exhausted. Try fewer or smaller images.") from exc | |
| results = [] | |
| batch = os.urandom(8).hex() | |
| for i, image in enumerate(images): | |
| path = os.path.join(tempfile.gettempdir(), f"moss_{batch}_{i + 1:02d}.png") | |
| image.save(path) | |
| results.append(_as_media(path)) | |
| elapsed = time.monotonic() - start | |
| status = ( | |
| f"{len(results)} image(s) · {MODEL_IDS[key].rsplit('/', 1)[-1]} " | |
| f"· seed {seed} · {width}×{height} · {steps} steps " | |
| f"· CFG {cfg:g} · shift {shift:g} · {elapsed:.1f}s" | |
| ) | |
| # call_fn maps a fn's returned list/tuple onto the node's output PORTS | |
| # (one element per port), so the gallery — one port holding many images — | |
| # must come back as a single element containing the list: [results]. | |
| # Returning `results` bare would hand the gallery port only image #1. | |
| return [results, status] | |
| demo = gr.Workflow( | |
| graph=str(Path(__file__).with_name("workflow.json")), | |
| bind={"generate": generate}, | |
| ) | |
| demo.title = "MOSS-Image" | |
| # --- Landing-page styling for the auto-generated App view ------------------- | |
| # gr.Workflow builds its App view inside #workflow-app; launch() threads | |
| # theme/css/head straight through to Blocks.launch(). The hero uses the | |
| # official lockup SVGs from the original Space (hot-linked, Apache-2.0). | |
| _LOCKUP_LIGHT = ( | |
| "https://huggingface.co/spaces/OpenMOSS-Team/MOSS-Image" | |
| "/resolve/main/assets/lockup-light.svg" | |
| ) | |
| _LOCKUP_DARK = ( | |
| "https://huggingface.co/spaces/OpenMOSS-Team/MOSS-Image" | |
| "/resolve/main/assets/lockup-dark.svg" | |
| ) | |
| HEAD = """ | |
| <link rel="preconnect" href="https://fonts.googleapis.com"> | |
| <link rel="preconnect" href="https://fonts.gstatic.com" crossorigin> | |
| <link href="https://fonts.googleapis.com/css2?family=Inter:wght@400;500;600;700;800&display=swap" rel="stylesheet"> | |
| """ | |
| # Custom JS for launch(js=...): the gradio frontend executes this function once | |
| # the page is up. A <script> tag in head= is unreliable for this under gradio 6 | |
| # SSR — the frontend can inject head HTML after the window load event, and | |
| # innerHTML-inserted <script> tags never execute. | |
| # | |
| # Two pieces: the top nav, and the examples strip — the original Space's | |
| # showcase, hot-linked. Clicking a thumbnail fetches its parameters.json and | |
| # restores every field of the generated App-view form, replicating the | |
| # original's click-to-restore. | |
| JS = """ | |
| () => { | |
| const ORIGIN = "https://huggingface.co/spaces/OpenMOSS-Team/MOSS-Image/resolve/main/assets/images"; | |
| const IDS = Array.from({length: 20}, (_, i) => String(i + 1).padStart(3, "0")); | |
| function setField(prefix, value) { | |
| if (value === undefined || value === null) return; | |
| const labels = Array.from(document.querySelectorAll('#workflow-app label')); | |
| const match = labels.find(l => (l.textContent || "").trim().toLowerCase().startsWith(prefix)); | |
| if (!match) return; | |
| let root = match; | |
| for (let i = 0; i < 6 && root; i++) { | |
| const input = root.querySelector('input:not([type="hidden"]), textarea'); | |
| if (input) { | |
| input.value = value; | |
| input.dispatchEvent(new Event("input", {bubbles: true})); | |
| input.dispatchEvent(new Event("change", {bubbles: true})); | |
| return; | |
| } | |
| root = root.parentElement; | |
| } | |
| } | |
| function applyRecord(rec) { | |
| const p = rec.parameters || {}; | |
| let uncond = p.uncond_skip_layers ?? ""; | |
| let slgLayers = p.slg_layers ?? ""; | |
| let slgScale = p.slg_scale ?? 0; | |
| if ((rec.schema_version ?? 2) === 1) { | |
| // v1 records called negative-branch skipping "slg_layers". | |
| uncond = (p.slg ?? true) ? (p.slg_layers ?? "") : ""; | |
| slgLayers = ""; | |
| slgScale = 0; | |
| } | |
| const model = (rec.model_repo || rec.model || "").toLowerCase().includes("sft") ? "sft" : "base"; | |
| setField("negative", p.negative_prompt ?? ""); | |
| setField("prompt", p.prompt ?? ""); | |
| setField("model", model); | |
| setField("width", p.width); | |
| setField("height", p.height); | |
| setField("steps", p.steps); | |
| setField("cfg", p.cfg); | |
| setField("shift", p.shift); | |
| setField("seed", p.seed ?? -1); | |
| setField("images", p.num_images ?? 1); | |
| setField("max text", p.max_sequence_length ?? 4096); | |
| setField("skip start", p.slg_start ?? 0); | |
| setField("skip stop", p.slg_stop ?? 1); | |
| setField("skip layers", uncond); | |
| setField("slg layers", slgLayers); | |
| setField("slg scale", slgScale); | |
| } | |
| const recordCache = {}; | |
| function fetchRecord(id) { | |
| if (!recordCache[id]) { | |
| recordCache[id] = fetch(`${ORIGIN}/${id}/parameters.json`) | |
| .then(r => (r.ok ? r.json() : null)) | |
| .catch(() => null); | |
| } | |
| return recordCache[id]; | |
| } | |
| function buildStrip(app) { | |
| if (document.getElementById('moss-examples')) return true; | |
| const wrap = document.createElement('div'); | |
| wrap.id = 'moss-examples'; | |
| const title = document.createElement('div'); | |
| title.className = 'mx-title'; | |
| title.textContent = 'Examples · click to restore prompt & settings'; | |
| wrap.appendChild(title); | |
| const row = document.createElement('div'); | |
| row.className = 'mx-row'; | |
| IDS.forEach((id, i) => { | |
| const fig = document.createElement('button'); | |
| fig.className = 'mx-item'; | |
| fig.type = 'button'; | |
| fig.title = 'Load this example into the form'; | |
| const img = document.createElement('img'); | |
| img.loading = 'lazy'; | |
| img.alt = `Example ${i + 1}`; | |
| img.src = `${ORIGIN}/${id}/image.webp`; | |
| img.onerror = () => fig.remove(); | |
| fig.appendChild(img); | |
| fig.addEventListener('click', () => { | |
| fetchRecord(id).then(rec => { if (rec) applyRecord(rec); }); | |
| }); | |
| row.appendChild(fig); | |
| }); | |
| wrap.appendChild(row); | |
| app.prepend(wrap); | |
| return true; | |
| } | |
| if (!document.getElementById('moss-nav')) { | |
| const nav = document.createElement('nav'); | |
| nav.id = 'moss-nav'; | |
| nav.innerHTML = ` | |
| <span class="moss-brand">MOSS-Image</span> | |
| <span class="moss-links"> | |
| <a href="https://huggingface.co/OpenMOSS-Team/MOSS-Image-Base" target="_blank" rel="noopener">Base</a> | |
| <a href="https://huggingface.co/OpenMOSS-Team/MOSS-Image-SFT" target="_blank" rel="noopener">SFT</a> | |
| <a href="https://huggingface.co/spaces/OpenMOSS-Team/MOSS-Image" target="_blank" rel="noopener">Original Demo</a> | |
| <a href="https://github.com/OpenMOSS/MOSS-Image" target="_blank" rel="noopener">GitHub ↗</a> | |
| </span>`; | |
| document.body.prepend(nav); | |
| } | |
| // #workflow-app appears when the frontend mounts; poll briefly for it. | |
| let tries = 0; | |
| const timer = setInterval(() => { | |
| const app = document.getElementById('workflow-app'); | |
| if ((app && buildStrip(app)) || ++tries > 75) clearInterval(timer); | |
| }, 400); | |
| } | |
| """ | |
| CSS = f""" | |
| body, .gradio-container {{ | |
| font-family: 'Inter', ui-sans-serif, system-ui, sans-serif !important; | |
| }} | |
| .gradio-container {{ | |
| background: linear-gradient(180deg, #eef3fc 0%, #ffffff 460px) !important; | |
| }} | |
| .dark .gradio-container {{ | |
| background: linear-gradient(180deg, #0d1526 0%, #0b0f19 460px) !important; | |
| }} | |
| #moss-nav {{ | |
| display: flex; | |
| align-items: center; | |
| justify-content: space-between; | |
| padding: 14px 24px 0; | |
| max-width: 1200px; | |
| margin: 0 auto; | |
| font-size: 14px; | |
| }} | |
| #moss-nav .moss-brand {{ | |
| font-weight: 800; | |
| font-size: 16px; | |
| letter-spacing: -0.01em; | |
| background: linear-gradient(90deg, #3b7fe0, #4aa5f2, #7fe6fb); | |
| -webkit-background-clip: text; | |
| background-clip: text; | |
| color: transparent; | |
| }} | |
| #moss-nav .moss-links a {{ | |
| color: #475569; | |
| text-decoration: none; | |
| font-weight: 500; | |
| margin-left: 22px; | |
| }} | |
| #moss-nav .moss-links a:hover {{ color: #2563eb; }} | |
| .dark #moss-nav .moss-links a {{ color: #94a3b8; }} | |
| .dark #moss-nav .moss-links a:hover {{ color: #7fe6fb; }} | |
| #workflow-app::before {{ | |
| content: "MOSS-Image: Text-to-Image Generation\\Awith Native Layout Control"; | |
| white-space: pre-wrap; | |
| display: block; | |
| text-align: center; | |
| font-weight: 800; | |
| font-size: clamp(24px, 4.2vw, 40px); | |
| line-height: 1.15; | |
| letter-spacing: -0.02em; | |
| color: #0b1220; | |
| max-width: 860px; | |
| margin: 0 auto 12px; | |
| padding: 158px 16px 26px; | |
| background: url("{_LOCKUP_LIGHT}") no-repeat top 26px center / min(480px, 78%) auto; | |
| }} | |
| .dark #workflow-app::before {{ | |
| color: #f1f5f9; | |
| background-image: url("{_LOCKUP_DARK}"); | |
| }} | |
| #workflow-app::after {{ | |
| content: "Models by the OpenMOSS Team · MOSS-Image-Base & MOSS-Image-SFT (Apache-2.0) · Rebuilt as a Gradio Workflow"; | |
| display: block; | |
| text-align: center; | |
| font-size: 13px; | |
| color: #64748b; | |
| padding: 8px 16px 36px; | |
| }} | |
| .dark #workflow-app::after {{ color: #94a3b8; }} | |
| /* Examples strip (injected by the head script; hot-linked from the original Space). */ | |
| #moss-examples {{ | |
| max-width: 960px; | |
| margin: 0 auto 18px; | |
| padding: 0 16px; | |
| }} | |
| #moss-examples .mx-title {{ | |
| text-align: center; | |
| font-size: 13px; | |
| font-weight: 600; | |
| letter-spacing: 0.02em; | |
| color: #64748b; | |
| margin-bottom: 10px; | |
| }} | |
| .dark #moss-examples .mx-title {{ color: #94a3b8; }} | |
| #moss-examples .mx-row {{ | |
| display: flex; | |
| gap: 10px; | |
| overflow-x: auto; | |
| padding-bottom: 6px; | |
| scrollbar-width: thin; | |
| }} | |
| #moss-examples .mx-item {{ | |
| flex: 0 0 auto; | |
| height: 128px; | |
| padding: 0; | |
| border: 2px solid transparent; | |
| border-radius: 12px; | |
| overflow: hidden; | |
| cursor: pointer; | |
| background: #e2e8f0; | |
| transition: border-color 0.15s ease, transform 0.15s ease; | |
| }} | |
| #moss-examples .mx-item:hover {{ | |
| border-color: #2563eb; | |
| transform: translateY(-2px); | |
| }} | |
| #moss-examples .mx-item img {{ | |
| height: 128px; | |
| width: auto; | |
| display: block; | |
| }} | |
| /* Soften the generated form into landing-page cards. */ | |
| #workflow-app .form, #workflow-app .panel, #workflow-app .block {{ | |
| border-radius: 14px !important; | |
| }} | |
| #workflow-app button.primary, #workflow-app button.lg {{ | |
| border-radius: 10px !important; | |
| }} | |
| """ | |
| if __name__ == "__main__": | |
| demo.launch( | |
| theme=gr.themes.Soft(primary_hue="blue", neutral_hue="slate"), | |
| css=CSS, | |
| head=HEAD, | |
| js=JS, | |
| ) | |