Spaces:
Running on Zero
Running on Zero
Replace demo entry point with gr.Workflow app (node canvas + styled App view)
Browse files
app.py
CHANGED
|
@@ -1,47 +1,587 @@
|
|
| 1 |
-
"""
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 2 |
|
| 3 |
-
import argparse
|
| 4 |
-
import logging
|
| 5 |
import os
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 6 |
|
| 7 |
-
|
|
|
|
|
|
|
| 8 |
|
| 9 |
-
|
| 10 |
-
|
| 11 |
-
|
| 12 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 13 |
|
|
|
|
| 14 |
|
| 15 |
-
|
| 16 |
-
|
| 17 |
-
|
| 18 |
-
|
| 19 |
-
|
| 20 |
-
|
| 21 |
-
|
| 22 |
-
|
| 23 |
-
|
| 24 |
-
|
| 25 |
-
|
| 26 |
-
|
| 27 |
-
|
| 28 |
-
|
| 29 |
-
|
| 30 |
-
|
| 31 |
-
|
| 32 |
-
|
| 33 |
-
|
| 34 |
-
|
| 35 |
-
|
| 36 |
-
|
| 37 |
-
|
| 38 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 39 |
)
|
| 40 |
-
|
| 41 |
-
|
| 42 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 43 |
)
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 44 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 45 |
|
| 46 |
if __name__ == "__main__":
|
| 47 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""MOSS-Image as a gr.Workflow app — node-canvas text-to-image pipeline.
|
| 2 |
+
|
| 3 |
+
A reimagining of https://huggingface.co/spaces/OpenMOSS-Team/MOSS-Image as a
|
| 4 |
+
Gradio Workflow (gradio 6.30.0): generation is a bound Python-function node on
|
| 5 |
+
a drag-and-drop canvas, fed by reference nodes (prompt, size, and every
|
| 6 |
+
sampling knob of the original demo) and writing into an output-gallery subject.
|
| 7 |
+
Edit the graph on the canvas or in workflow.json; the App view and REST API
|
| 8 |
+
stay in sync automatically.
|
| 9 |
+
|
| 10 |
+
Both checkpoints are loaded once at module scope (ZeroGPU packs and transfers
|
| 11 |
+
them when a GPU is actually allocated) and shared by every run of the node.
|
| 12 |
+
"""
|
| 13 |
|
|
|
|
|
|
|
| 14 |
import os
|
| 15 |
+
import secrets
|
| 16 |
+
import tempfile
|
| 17 |
+
import time
|
| 18 |
+
from collections import OrderedDict
|
| 19 |
+
from pathlib import Path
|
| 20 |
|
| 21 |
+
# `spaces` must be imported before anything that initializes CUDA.
|
| 22 |
+
import spaces # ty: ignore[unresolved-import]
|
| 23 |
+
from gradio_client import utils as client_utils
|
| 24 |
|
| 25 |
+
import gradio as gr
|
| 26 |
+
from gradio import workflow_api as _wf_api
|
| 27 |
+
|
| 28 |
+
# Workaround for gradio 6.30.0: the App view hands a subject's value to its
|
| 29 |
+
# output component via `_from_output`, which unwraps a top-level file dict to
|
| 30 |
+
# its path but leaves dicts nested in a LIST untouched — and gr.Gallery only
|
| 31 |
+
# accepts path strings, not file dicts. A gallery subject fed by an fn node
|
| 32 |
+
# (which must return file dicts for the canvas) therefore crashes the App
|
| 33 |
+
# view. Unwrap list items the same way so both views work.
|
| 34 |
+
_orig_from_output = _wf_api._from_output
|
| 35 |
+
|
| 36 |
+
|
| 37 |
+
def _from_output_unwrapping_lists(value):
|
| 38 |
+
if isinstance(value, list):
|
| 39 |
+
return [_orig_from_output(item) for item in value]
|
| 40 |
+
return _orig_from_output(value)
|
| 41 |
+
|
| 42 |
+
|
| 43 |
+
_wf_api._from_output = _from_output_unwrapping_lists
|
| 44 |
+
|
| 45 |
+
|
| 46 |
+
def _install_component_server_keepalive():
|
| 47 |
+
"""Keep `/gradio_api/component_server` connections alive during long fn calls.
|
| 48 |
+
|
| 49 |
+
Canvas fn nodes call back through the component_server route — a plain
|
| 50 |
+
POST that awaits the bound function synchronously, unlike App-view
|
| 51 |
+
endpoints, which run through the queue with SSE heartbeats. A MOSS-Image
|
| 52 |
+
generation takes 45-130s, and the silent connection is dropped by the
|
| 53 |
+
proxy chain, surfacing on the canvas as "Connection errored out" (the
|
| 54 |
+
App view is unaffected). Wrap the route so it streams whitespace
|
| 55 |
+
heartbeat bytes — legal leading whitespace in JSON — until the real
|
| 56 |
+
result is ready.
|
| 57 |
+
"""
|
| 58 |
+
import asyncio
|
| 59 |
+
import logging
|
| 60 |
+
|
| 61 |
+
from fastapi import Request
|
| 62 |
+
from fastapi.dependencies.utils import get_dependant
|
| 63 |
+
from fastapi.routing import APIRoute
|
| 64 |
+
from starlette.responses import StreamingResponse
|
| 65 |
+
from starlette.routing import request_response
|
| 66 |
+
from gradio.routes import App as _RoutesApp
|
| 67 |
+
|
| 68 |
+
original_create_app = _RoutesApp.create_app
|
| 69 |
+
|
| 70 |
+
def _wrap(endpoint):
|
| 71 |
+
async def component_server_with_keepalive(request: Request):
|
| 72 |
+
task = asyncio.create_task(endpoint(request))
|
| 73 |
+
|
| 74 |
+
async def body():
|
| 75 |
+
try:
|
| 76 |
+
while not task.done():
|
| 77 |
+
yield b" "
|
| 78 |
+
await asyncio.sleep(10)
|
| 79 |
+
payload = json.dumps(task.result(), default=str).encode()
|
| 80 |
+
except Exception as exc: # keep the canvas error path readable
|
| 81 |
+
payload = json.dumps({"error": str(exc)}).encode()
|
| 82 |
+
yield payload
|
| 83 |
+
|
| 84 |
+
return StreamingResponse(body(), media_type="application/json")
|
| 85 |
|
| 86 |
+
return component_server_with_keepalive
|
| 87 |
|
| 88 |
+
@staticmethod
|
| 89 |
+
def create_app(blocks, *args, **kwargs):
|
| 90 |
+
app = original_create_app(blocks, *args, **kwargs)
|
| 91 |
+
for route in app.routes:
|
| 92 |
+
if isinstance(route, APIRoute) and route.path.rstrip("/").endswith(
|
| 93 |
+
"/component_server"
|
| 94 |
+
):
|
| 95 |
+
route.endpoint = _wrap(route.endpoint)
|
| 96 |
+
dependant = get_dependant(
|
| 97 |
+
path=route.path_format, endpoint=route.endpoint
|
| 98 |
+
)
|
| 99 |
+
# Keep the route-level dependencies (e.g. Depends(login_check)).
|
| 100 |
+
dependant.dependencies = (
|
| 101 |
+
route.dependant.dependencies + dependant.dependencies
|
| 102 |
+
)
|
| 103 |
+
route.dependant = dependant
|
| 104 |
+
route.app = request_response(route.get_route_handler())
|
| 105 |
+
logging.getLogger(__name__).info(
|
| 106 |
+
"component_server keep-alive wrapper installed"
|
| 107 |
+
)
|
| 108 |
+
return app
|
| 109 |
+
|
| 110 |
+
_RoutesApp.create_app = create_app
|
| 111 |
+
|
| 112 |
+
|
| 113 |
+
_install_component_server_keepalive()
|
| 114 |
+
|
| 115 |
+
MODEL_IDS = {
|
| 116 |
+
"base": "OpenMOSS-Team/MOSS-Image-Base",
|
| 117 |
+
"sft": "OpenMOSS-Team/MOSS-Image-SFT",
|
| 118 |
+
}
|
| 119 |
+
|
| 120 |
+
# Set MOSS_WORKFLOW_PREVIEW=1 to open the canvas without downloading weights
|
| 121 |
+
# (useful for local UI work on a CPU-only machine).
|
| 122 |
+
PREVIEW = os.environ.get("MOSS_WORKFLOW_PREVIEW") == "1"
|
| 123 |
+
|
| 124 |
+
|
| 125 |
+
def _load_pipeline(model_id, shared_pipe=None):
|
| 126 |
+
"""Load one MOSS-Image checkpoint with the bundled inference classes.
|
| 127 |
+
|
| 128 |
+
Text encoder, tokenizer and VAE are shared with `shared_pipe` when their
|
| 129 |
+
Hub file hashes match — this is what keeps both checkpoints inside a
|
| 130 |
+
ZeroGPU worker's GPU memory (two full copies do not fit).
|
| 131 |
+
"""
|
| 132 |
+
import torch
|
| 133 |
+
from huggingface_hub import HfApi, snapshot_download
|
| 134 |
+
from pipeline_moss_image import (
|
| 135 |
+
MOSSImagePipeline,
|
| 136 |
+
MOSSImageScheduler,
|
| 137 |
+
MOSSImageTransformer,
|
| 138 |
)
|
| 139 |
+
|
| 140 |
+
info = HfApi().model_info(model_id, files_metadata=True)
|
| 141 |
+
revision = info.sha
|
| 142 |
+
shared_hashes = {
|
| 143 |
+
item.rfilename: item.lfs.sha256 if item.lfs else item.blob_id
|
| 144 |
+
for item in info.siblings
|
| 145 |
+
if item.rfilename.startswith(("text_encoder/", "tokenizer/", "vae/"))
|
| 146 |
+
}
|
| 147 |
+
components = {}
|
| 148 |
+
if shared_pipe is not None:
|
| 149 |
+
for name in ("text_encoder", "tokenizer", "vae"):
|
| 150 |
+
current = {k: v for k, v in shared_hashes.items() if k.startswith(f"{name}/")}
|
| 151 |
+
previous = {
|
| 152 |
+
k: v for k, v in shared_pipe._shared_component_hashes.items()
|
| 153 |
+
if k.startswith(f"{name}/")
|
| 154 |
+
}
|
| 155 |
+
if current and all(current.values()) and current == previous:
|
| 156 |
+
components[name] = getattr(shared_pipe, name)
|
| 157 |
+
snapshot = snapshot_download(
|
| 158 |
+
model_id, revision=revision, ignore_patterns=[f"{name}/*" for name in components]
|
| 159 |
)
|
| 160 |
+
pipe = MOSSImagePipeline.from_pretrained(
|
| 161 |
+
snapshot,
|
| 162 |
+
custom_pipeline=str(Path(__file__).with_name("pipeline_moss_image")),
|
| 163 |
+
trust_remote_code=True,
|
| 164 |
+
transformer=MOSSImageTransformer.from_pretrained(
|
| 165 |
+
snapshot, subfolder="transformer", torch_dtype=torch.bfloat16
|
| 166 |
+
),
|
| 167 |
+
scheduler=MOSSImageScheduler.from_pretrained(snapshot, subfolder="scheduler"),
|
| 168 |
+
local_files_only=True,
|
| 169 |
+
torch_dtype={"default": torch.bfloat16, "vae": torch.float32},
|
| 170 |
+
**components,
|
| 171 |
+
)
|
| 172 |
+
# On ZeroGPU the `spaces` package intercepts this move, packs the weights,
|
| 173 |
+
# and transfers them when a GPU is allocated — module scope costs no GPU time.
|
| 174 |
+
pipe.to("cuda")
|
| 175 |
+
pipe._prompt_cache = OrderedDict()
|
| 176 |
+
pipe._shared_component_hashes = shared_hashes
|
| 177 |
+
return pipe
|
| 178 |
+
|
| 179 |
+
|
| 180 |
+
PIPELINES = {}
|
| 181 |
+
if not PREVIEW:
|
| 182 |
+
PIPELINES["base"] = _load_pipeline(MODEL_IDS["base"])
|
| 183 |
+
PIPELINES["sft"] = _load_pipeline(MODEL_IDS["sft"], shared_pipe=PIPELINES["base"])
|
| 184 |
+
|
| 185 |
+
|
| 186 |
+
def _as_media(path: str) -> dict:
|
| 187 |
+
"""The dict shape the workflow canvas renders media from."""
|
| 188 |
+
return {
|
| 189 |
+
"path": path,
|
| 190 |
+
"url": f"/gradio_api/file={client_utils.encode_file_path(path)}",
|
| 191 |
+
"is_file": True,
|
| 192 |
+
}
|
| 193 |
+
|
| 194 |
|
| 195 |
+
def _parse_layers(value: str, name: str) -> tuple:
|
| 196 |
+
"""'22' -> (22,); '' -> (). Same rules as the original demo."""
|
| 197 |
+
if not (value or "").strip():
|
| 198 |
+
return ()
|
| 199 |
+
try:
|
| 200 |
+
layers = tuple(int(item.strip()) for item in value.split(","))
|
| 201 |
+
except ValueError:
|
| 202 |
+
raise gr.Error(
|
| 203 |
+
f"{name} must be comma-separated integer indices, e.g. 22."
|
| 204 |
+
) from None
|
| 205 |
+
if any(layer < 0 for layer in layers) or len(set(layers)) != len(layers):
|
| 206 |
+
raise gr.Error(f"{name} indices must be nonnegative and unique.")
|
| 207 |
+
return layers
|
| 208 |
+
|
| 209 |
+
|
| 210 |
+
@spaces.GPU(duration=80)
|
| 211 |
+
def generate(
|
| 212 |
+
prompt: str,
|
| 213 |
+
negative_prompt: str,
|
| 214 |
+
model: str,
|
| 215 |
+
width: float,
|
| 216 |
+
height: float,
|
| 217 |
+
steps: float,
|
| 218 |
+
cfg: float,
|
| 219 |
+
shift: float,
|
| 220 |
+
seed: float,
|
| 221 |
+
num_images: float,
|
| 222 |
+
max_sequence_length: float,
|
| 223 |
+
uncond_skip_layers: str,
|
| 224 |
+
slg_layers: str,
|
| 225 |
+
slg_scale: float,
|
| 226 |
+
slg_start: float,
|
| 227 |
+
slg_stop: float,
|
| 228 |
+
) -> list:
|
| 229 |
+
"""Generate images with MOSS-Image (Base or SFT). Full sampling controls."""
|
| 230 |
+
import torch
|
| 231 |
+
|
| 232 |
+
if PREVIEW:
|
| 233 |
+
raise gr.Error("Preview mode: models are not loaded (MOSS_WORKFLOW_PREVIEW=1).")
|
| 234 |
+
if not (prompt or "").strip():
|
| 235 |
+
raise gr.Error("Enter a prompt.")
|
| 236 |
+
key = (model or "").strip().lower().replace("moss-image-", "")
|
| 237 |
+
if key not in PIPELINES:
|
| 238 |
+
raise gr.Error(f"Unknown model {model!r}. Choose 'base' or 'sft'.")
|
| 239 |
+
width, height, steps, num_images, max_sequence_length = (
|
| 240 |
+
round(width), round(height), round(steps), round(num_images),
|
| 241 |
+
round(max_sequence_length),
|
| 242 |
+
)
|
| 243 |
+
seed = round(seed)
|
| 244 |
+
if width <= 0 or height <= 0 or width % 16 or height % 16:
|
| 245 |
+
raise gr.Error("Width and height must be positive multiples of 16.")
|
| 246 |
+
if steps < 1 or num_images < 1:
|
| 247 |
+
raise gr.Error("Steps and image count must be positive.")
|
| 248 |
+
if not 0 <= cfg or shift <= 0:
|
| 249 |
+
raise gr.Error("CFG must be nonnegative; shift must be positive.")
|
| 250 |
+
if not 1 <= max_sequence_length <= 4096:
|
| 251 |
+
raise gr.Error("Maximum text length must be between 1 and 4096.")
|
| 252 |
+
if slg_scale < 0 or not 0 <= slg_start <= slg_stop <= 1:
|
| 253 |
+
raise gr.Error("Require SLG scale >= 0 and 0 <= SLG start <= SLG stop <= 1.")
|
| 254 |
+
if seed == -1:
|
| 255 |
+
seed = secrets.randbelow(2**32)
|
| 256 |
+
if not 0 <= seed < 2**53:
|
| 257 |
+
raise gr.Error("Seed must be -1 (random) or an integer from 0 to 2^53 - 1.")
|
| 258 |
+
|
| 259 |
+
pipe = PIPELINES[key]
|
| 260 |
+
depth = pipe.transformer.config.num_layers
|
| 261 |
+
skip = _parse_layers(uncond_skip_layers, "Skip layers")
|
| 262 |
+
slg = _parse_layers(slg_layers, "SLG layers")
|
| 263 |
+
if any(layer >= depth for layer in skip + slg):
|
| 264 |
+
raise gr.Error(f"Skip-layer indices must be below {depth}.")
|
| 265 |
+
|
| 266 |
+
pipe.scheduler = type(pipe.scheduler).from_config(pipe.scheduler.config, shift=shift)
|
| 267 |
+
start = time.monotonic()
|
| 268 |
+
try:
|
| 269 |
+
images = pipe(
|
| 270 |
+
prompt=prompt,
|
| 271 |
+
negative_prompt=negative_prompt or "",
|
| 272 |
+
width=width,
|
| 273 |
+
height=height,
|
| 274 |
+
num_inference_steps=steps,
|
| 275 |
+
guidance_scale=float(cfg),
|
| 276 |
+
num_images_per_prompt=num_images,
|
| 277 |
+
max_sequence_length=max_sequence_length,
|
| 278 |
+
uncond_skip_layers=skip,
|
| 279 |
+
slg_layers=slg,
|
| 280 |
+
slg_scale=float(slg_scale),
|
| 281 |
+
slg_start=float(slg_start),
|
| 282 |
+
slg_stop=float(slg_stop),
|
| 283 |
+
generator=torch.Generator("cpu").manual_seed(seed),
|
| 284 |
+
).images
|
| 285 |
+
except torch.cuda.OutOfMemoryError as exc:
|
| 286 |
+
torch.cuda.empty_cache()
|
| 287 |
+
raise gr.Error("GPU memory exhausted. Try fewer or smaller images.") from exc
|
| 288 |
+
|
| 289 |
+
results = []
|
| 290 |
+
batch = os.urandom(8).hex()
|
| 291 |
+
for i, image in enumerate(images):
|
| 292 |
+
path = os.path.join(tempfile.gettempdir(), f"moss_{batch}_{i + 1:02d}.png")
|
| 293 |
+
image.save(path)
|
| 294 |
+
results.append(_as_media(path))
|
| 295 |
+
elapsed = time.monotonic() - start
|
| 296 |
+
status = (
|
| 297 |
+
f"{len(results)} image(s) · {MODEL_IDS[key].rsplit('/', 1)[-1]} "
|
| 298 |
+
f"· seed {seed} · {width}×{height} · {steps} steps "
|
| 299 |
+
f"· CFG {cfg:g} · shift {shift:g} · {elapsed:.1f}s"
|
| 300 |
+
)
|
| 301 |
+
# call_fn maps a fn's returned list/tuple onto the node's output PORTS
|
| 302 |
+
# (one element per port), so the gallery — one port holding many images —
|
| 303 |
+
# must come back as a single element containing the list: [results].
|
| 304 |
+
# Returning `results` bare would hand the gallery port only image #1.
|
| 305 |
+
return [results, status]
|
| 306 |
+
|
| 307 |
+
|
| 308 |
+
demo = gr.Workflow(
|
| 309 |
+
graph=str(Path(__file__).with_name("workflow.json")),
|
| 310 |
+
bind={"generate": generate},
|
| 311 |
+
)
|
| 312 |
+
demo.title = "MOSS-Image"
|
| 313 |
+
|
| 314 |
+
# --- Landing-page styling for the auto-generated App view -------------------
|
| 315 |
+
# gr.Workflow builds its App view inside #workflow-app; launch() threads
|
| 316 |
+
# theme/css/head straight through to Blocks.launch(). The hero uses the
|
| 317 |
+
# official lockup SVGs from the original Space (hot-linked, Apache-2.0).
|
| 318 |
+
|
| 319 |
+
_LOCKUP_LIGHT = (
|
| 320 |
+
"https://huggingface.co/spaces/OpenMOSS-Team/MOSS-Image"
|
| 321 |
+
"/resolve/main/assets/lockup-light.svg"
|
| 322 |
+
)
|
| 323 |
+
_LOCKUP_DARK = (
|
| 324 |
+
"https://huggingface.co/spaces/OpenMOSS-Team/MOSS-Image"
|
| 325 |
+
"/resolve/main/assets/lockup-dark.svg"
|
| 326 |
+
)
|
| 327 |
+
|
| 328 |
+
HEAD = """
|
| 329 |
+
<link rel="preconnect" href="https://fonts.googleapis.com">
|
| 330 |
+
<link rel="preconnect" href="https://fonts.gstatic.com" crossorigin>
|
| 331 |
+
<link href="https://fonts.googleapis.com/css2?family=Inter:wght@400;500;600;700;800&display=swap" rel="stylesheet">
|
| 332 |
+
"""
|
| 333 |
+
|
| 334 |
+
# Custom JS for launch(js=...): the gradio frontend executes this function once
|
| 335 |
+
# the page is up. A <script> tag in head= is unreliable for this under gradio 6
|
| 336 |
+
# SSR — the frontend can inject head HTML after the window load event, and
|
| 337 |
+
# innerHTML-inserted <script> tags never execute.
|
| 338 |
+
#
|
| 339 |
+
# Two pieces: the top nav, and the examples strip — the original Space's
|
| 340 |
+
# showcase, hot-linked. Clicking a thumbnail fetches its parameters.json and
|
| 341 |
+
# restores every field of the generated App-view form, replicating the
|
| 342 |
+
# original's click-to-restore.
|
| 343 |
+
JS = """
|
| 344 |
+
() => {
|
| 345 |
+
const ORIGIN = "https://huggingface.co/spaces/OpenMOSS-Team/MOSS-Image/resolve/main/assets/images";
|
| 346 |
+
const IDS = Array.from({length: 20}, (_, i) => String(i + 1).padStart(3, "0"));
|
| 347 |
+
|
| 348 |
+
function setField(prefix, value) {
|
| 349 |
+
if (value === undefined || value === null) return;
|
| 350 |
+
const labels = Array.from(document.querySelectorAll('#workflow-app label'));
|
| 351 |
+
const match = labels.find(l => (l.textContent || "").trim().toLowerCase().startsWith(prefix));
|
| 352 |
+
if (!match) return;
|
| 353 |
+
let root = match;
|
| 354 |
+
for (let i = 0; i < 6 && root; i++) {
|
| 355 |
+
const input = root.querySelector('input:not([type="hidden"]), textarea');
|
| 356 |
+
if (input) {
|
| 357 |
+
input.value = value;
|
| 358 |
+
input.dispatchEvent(new Event("input", {bubbles: true}));
|
| 359 |
+
input.dispatchEvent(new Event("change", {bubbles: true}));
|
| 360 |
+
return;
|
| 361 |
+
}
|
| 362 |
+
root = root.parentElement;
|
| 363 |
+
}
|
| 364 |
+
}
|
| 365 |
+
|
| 366 |
+
function applyRecord(rec) {
|
| 367 |
+
const p = rec.parameters || {};
|
| 368 |
+
let uncond = p.uncond_skip_layers ?? "";
|
| 369 |
+
let slgLayers = p.slg_layers ?? "";
|
| 370 |
+
let slgScale = p.slg_scale ?? 0;
|
| 371 |
+
if ((rec.schema_version ?? 2) === 1) {
|
| 372 |
+
// v1 records called negative-branch skipping "slg_layers".
|
| 373 |
+
uncond = (p.slg ?? true) ? (p.slg_layers ?? "") : "";
|
| 374 |
+
slgLayers = "";
|
| 375 |
+
slgScale = 0;
|
| 376 |
+
}
|
| 377 |
+
const model = (rec.model_repo || rec.model || "").toLowerCase().includes("sft") ? "sft" : "base";
|
| 378 |
+
setField("negative", p.negative_prompt ?? "");
|
| 379 |
+
setField("prompt", p.prompt ?? "");
|
| 380 |
+
setField("model", model);
|
| 381 |
+
setField("width", p.width);
|
| 382 |
+
setField("height", p.height);
|
| 383 |
+
setField("steps", p.steps);
|
| 384 |
+
setField("cfg", p.cfg);
|
| 385 |
+
setField("shift", p.shift);
|
| 386 |
+
setField("seed", p.seed ?? -1);
|
| 387 |
+
setField("images", p.num_images ?? 1);
|
| 388 |
+
setField("max text", p.max_sequence_length ?? 4096);
|
| 389 |
+
setField("skip start", p.slg_start ?? 0);
|
| 390 |
+
setField("skip stop", p.slg_stop ?? 1);
|
| 391 |
+
setField("skip layers", uncond);
|
| 392 |
+
setField("slg layers", slgLayers);
|
| 393 |
+
setField("slg scale", slgScale);
|
| 394 |
+
}
|
| 395 |
+
|
| 396 |
+
const recordCache = {};
|
| 397 |
+
function fetchRecord(id) {
|
| 398 |
+
if (!recordCache[id]) {
|
| 399 |
+
recordCache[id] = fetch(`${ORIGIN}/${id}/parameters.json`)
|
| 400 |
+
.then(r => (r.ok ? r.json() : null))
|
| 401 |
+
.catch(() => null);
|
| 402 |
+
}
|
| 403 |
+
return recordCache[id];
|
| 404 |
+
}
|
| 405 |
+
|
| 406 |
+
function buildStrip(app) {
|
| 407 |
+
if (document.getElementById('moss-examples')) return true;
|
| 408 |
+
const wrap = document.createElement('div');
|
| 409 |
+
wrap.id = 'moss-examples';
|
| 410 |
+
const title = document.createElement('div');
|
| 411 |
+
title.className = 'mx-title';
|
| 412 |
+
title.textContent = 'Examples · click to restore prompt & settings';
|
| 413 |
+
wrap.appendChild(title);
|
| 414 |
+
const row = document.createElement('div');
|
| 415 |
+
row.className = 'mx-row';
|
| 416 |
+
IDS.forEach((id, i) => {
|
| 417 |
+
const fig = document.createElement('button');
|
| 418 |
+
fig.className = 'mx-item';
|
| 419 |
+
fig.type = 'button';
|
| 420 |
+
fig.title = 'Load this example into the form';
|
| 421 |
+
const img = document.createElement('img');
|
| 422 |
+
img.loading = 'lazy';
|
| 423 |
+
img.alt = `Example ${i + 1}`;
|
| 424 |
+
img.src = `${ORIGIN}/${id}/image.webp`;
|
| 425 |
+
img.onerror = () => fig.remove();
|
| 426 |
+
fig.appendChild(img);
|
| 427 |
+
fig.addEventListener('click', () => {
|
| 428 |
+
fetchRecord(id).then(rec => { if (rec) applyRecord(rec); });
|
| 429 |
+
});
|
| 430 |
+
row.appendChild(fig);
|
| 431 |
+
});
|
| 432 |
+
wrap.appendChild(row);
|
| 433 |
+
app.prepend(wrap);
|
| 434 |
+
return true;
|
| 435 |
+
}
|
| 436 |
+
|
| 437 |
+
if (!document.getElementById('moss-nav')) {
|
| 438 |
+
const nav = document.createElement('nav');
|
| 439 |
+
nav.id = 'moss-nav';
|
| 440 |
+
nav.innerHTML = `
|
| 441 |
+
<span class="moss-brand">MOSS-Image</span>
|
| 442 |
+
<span class="moss-links">
|
| 443 |
+
<a href="https://huggingface.co/OpenMOSS-Team/MOSS-Image-Base" target="_blank" rel="noopener">Base</a>
|
| 444 |
+
<a href="https://huggingface.co/OpenMOSS-Team/MOSS-Image-SFT" target="_blank" rel="noopener">SFT</a>
|
| 445 |
+
<a href="https://huggingface.co/spaces/OpenMOSS-Team/MOSS-Image" target="_blank" rel="noopener">Original Demo</a>
|
| 446 |
+
<a href="https://github.com/OpenMOSS/MOSS-Image" target="_blank" rel="noopener">GitHub ↗</a>
|
| 447 |
+
</span>`;
|
| 448 |
+
document.body.prepend(nav);
|
| 449 |
+
}
|
| 450 |
+
|
| 451 |
+
// #workflow-app appears when the frontend mounts; poll briefly for it.
|
| 452 |
+
let tries = 0;
|
| 453 |
+
const timer = setInterval(() => {
|
| 454 |
+
const app = document.getElementById('workflow-app');
|
| 455 |
+
if ((app && buildStrip(app)) || ++tries > 75) clearInterval(timer);
|
| 456 |
+
}, 400);
|
| 457 |
+
}
|
| 458 |
+
"""
|
| 459 |
+
|
| 460 |
+
CSS = f"""
|
| 461 |
+
body, .gradio-container {{
|
| 462 |
+
font-family: 'Inter', ui-sans-serif, system-ui, sans-serif !important;
|
| 463 |
+
}}
|
| 464 |
+
.gradio-container {{
|
| 465 |
+
background: linear-gradient(180deg, #eef3fc 0%, #ffffff 460px) !important;
|
| 466 |
+
}}
|
| 467 |
+
.dark .gradio-container {{
|
| 468 |
+
background: linear-gradient(180deg, #0d1526 0%, #0b0f19 460px) !important;
|
| 469 |
+
}}
|
| 470 |
+
|
| 471 |
+
#moss-nav {{
|
| 472 |
+
display: flex;
|
| 473 |
+
align-items: center;
|
| 474 |
+
justify-content: space-between;
|
| 475 |
+
padding: 14px 24px 0;
|
| 476 |
+
max-width: 1200px;
|
| 477 |
+
margin: 0 auto;
|
| 478 |
+
font-size: 14px;
|
| 479 |
+
}}
|
| 480 |
+
#moss-nav .moss-brand {{
|
| 481 |
+
font-weight: 800;
|
| 482 |
+
font-size: 16px;
|
| 483 |
+
letter-spacing: -0.01em;
|
| 484 |
+
background: linear-gradient(90deg, #3b7fe0, #4aa5f2, #7fe6fb);
|
| 485 |
+
-webkit-background-clip: text;
|
| 486 |
+
background-clip: text;
|
| 487 |
+
color: transparent;
|
| 488 |
+
}}
|
| 489 |
+
#moss-nav .moss-links a {{
|
| 490 |
+
color: #475569;
|
| 491 |
+
text-decoration: none;
|
| 492 |
+
font-weight: 500;
|
| 493 |
+
margin-left: 22px;
|
| 494 |
+
}}
|
| 495 |
+
#moss-nav .moss-links a:hover {{ color: #2563eb; }}
|
| 496 |
+
.dark #moss-nav .moss-links a {{ color: #94a3b8; }}
|
| 497 |
+
.dark #moss-nav .moss-links a:hover {{ color: #7fe6fb; }}
|
| 498 |
+
|
| 499 |
+
#workflow-app::before {{
|
| 500 |
+
content: "MOSS-Image: Text-to-Image Generation\\Awith Native Layout Control";
|
| 501 |
+
white-space: pre-wrap;
|
| 502 |
+
display: block;
|
| 503 |
+
text-align: center;
|
| 504 |
+
font-weight: 800;
|
| 505 |
+
font-size: clamp(24px, 4.2vw, 40px);
|
| 506 |
+
line-height: 1.15;
|
| 507 |
+
letter-spacing: -0.02em;
|
| 508 |
+
color: #0b1220;
|
| 509 |
+
max-width: 860px;
|
| 510 |
+
margin: 0 auto 12px;
|
| 511 |
+
padding: 158px 16px 26px;
|
| 512 |
+
background: url("{_LOCKUP_LIGHT}") no-repeat top 26px center / min(480px, 78%) auto;
|
| 513 |
+
}}
|
| 514 |
+
.dark #workflow-app::before {{
|
| 515 |
+
color: #f1f5f9;
|
| 516 |
+
background-image: url("{_LOCKUP_DARK}");
|
| 517 |
+
}}
|
| 518 |
+
|
| 519 |
+
#workflow-app::after {{
|
| 520 |
+
content: "Models by the OpenMOSS Team · MOSS-Image-Base & MOSS-Image-SFT (Apache-2.0) · Rebuilt as a Gradio Workflow";
|
| 521 |
+
display: block;
|
| 522 |
+
text-align: center;
|
| 523 |
+
font-size: 13px;
|
| 524 |
+
color: #64748b;
|
| 525 |
+
padding: 8px 16px 36px;
|
| 526 |
+
}}
|
| 527 |
+
.dark #workflow-app::after {{ color: #94a3b8; }}
|
| 528 |
+
|
| 529 |
+
/* Examples strip (injected by the head script; hot-linked from the original Space). */
|
| 530 |
+
#moss-examples {{
|
| 531 |
+
max-width: 960px;
|
| 532 |
+
margin: 0 auto 18px;
|
| 533 |
+
padding: 0 16px;
|
| 534 |
+
}}
|
| 535 |
+
#moss-examples .mx-title {{
|
| 536 |
+
text-align: center;
|
| 537 |
+
font-size: 13px;
|
| 538 |
+
font-weight: 600;
|
| 539 |
+
letter-spacing: 0.02em;
|
| 540 |
+
color: #64748b;
|
| 541 |
+
margin-bottom: 10px;
|
| 542 |
+
}}
|
| 543 |
+
.dark #moss-examples .mx-title {{ color: #94a3b8; }}
|
| 544 |
+
#moss-examples .mx-row {{
|
| 545 |
+
display: flex;
|
| 546 |
+
gap: 10px;
|
| 547 |
+
overflow-x: auto;
|
| 548 |
+
padding-bottom: 6px;
|
| 549 |
+
scrollbar-width: thin;
|
| 550 |
+
}}
|
| 551 |
+
#moss-examples .mx-item {{
|
| 552 |
+
flex: 0 0 auto;
|
| 553 |
+
height: 128px;
|
| 554 |
+
padding: 0;
|
| 555 |
+
border: 2px solid transparent;
|
| 556 |
+
border-radius: 12px;
|
| 557 |
+
overflow: hidden;
|
| 558 |
+
cursor: pointer;
|
| 559 |
+
background: #e2e8f0;
|
| 560 |
+
transition: border-color 0.15s ease, transform 0.15s ease;
|
| 561 |
+
}}
|
| 562 |
+
#moss-examples .mx-item:hover {{
|
| 563 |
+
border-color: #2563eb;
|
| 564 |
+
transform: translateY(-2px);
|
| 565 |
+
}}
|
| 566 |
+
#moss-examples .mx-item img {{
|
| 567 |
+
height: 128px;
|
| 568 |
+
width: auto;
|
| 569 |
+
display: block;
|
| 570 |
+
}}
|
| 571 |
+
|
| 572 |
+
/* Soften the generated form into landing-page cards. */
|
| 573 |
+
#workflow-app .form, #workflow-app .panel, #workflow-app .block {{
|
| 574 |
+
border-radius: 14px !important;
|
| 575 |
+
}}
|
| 576 |
+
#workflow-app button.primary, #workflow-app button.lg {{
|
| 577 |
+
border-radius: 10px !important;
|
| 578 |
+
}}
|
| 579 |
+
"""
|
| 580 |
|
| 581 |
if __name__ == "__main__":
|
| 582 |
+
demo.launch(
|
| 583 |
+
theme=gr.themes.Soft(primary_hue="blue", neutral_hue="slate"),
|
| 584 |
+
css=CSS,
|
| 585 |
+
head=HEAD,
|
| 586 |
+
js=JS,
|
| 587 |
+
)
|