Spaces:
Running on Zero
Running on Zero
| """MOSS-Image Gradio interface for Hugging Face Spaces.""" | |
| import hashlib | |
| import html | |
| import importlib.metadata | |
| import json | |
| import logging | |
| import math | |
| import secrets | |
| import shutil | |
| import tempfile | |
| import threading | |
| import time | |
| from collections import OrderedDict | |
| import spaces | |
| from spaces.config import Config as SpacesConfig | |
| from dataclasses import asdict, dataclass, replace | |
| from functools import partial | |
| from importlib.resources import files | |
| from pathlib import Path | |
| from secrets import randbits | |
| from uuid import uuid4 | |
| from inference_runtime import generate_with_progress, set_scheduler | |
| from gradio_presets import ( | |
| PIXELS, | |
| fit_size, | |
| pixel_label, | |
| ratio_value, | |
| size_selection, | |
| ) | |
| logger = logging.getLogger(__name__) | |
| logger.setLevel(logging.INFO) | |
| MODEL_IDS = ("OpenMOSS-Team/MOSS-Image-Base", "OpenMOSS-Team/MOSS-Image-SFT") | |
| DEFAULT_MODEL = MODEL_IDS[0] | |
| class GenerationConfig: | |
| prompt: str = "" | |
| negative_prompt: str = "" | |
| height: int = 1024 | |
| width: int = 1024 | |
| steps: int = 50 | |
| cfg: float = 3.0 | |
| shift: float = 5.0 | |
| seed: int = 42 | |
| num_images: int = 1 | |
| max_sequence_length: int = 4096 | |
| uncond_skip_layers: str = "22" | |
| slg_layers: str = "" | |
| slg_scale: float = 0.0 | |
| slg_start: float = 0.0 | |
| slg_stop: float = 1.0 | |
| offload: str = "none" | |
| compile_transformer: bool = False | |
| def layers(self, name): | |
| value = getattr(self, name) | |
| if not value.strip(): | |
| return () | |
| try: | |
| layers = tuple(int(item.strip()) for item in value.split(",")) | |
| except ValueError as exc: | |
| raise ValueError("SLG layers must be comma-separated integer indices, e.g. 22.") from exc | |
| if any(layer < 0 for layer in layers) or len(set(layers)) != len(layers): | |
| raise ValueError("SLG indices must be nonnegative and unique.") | |
| return layers | |
| def validate(self): | |
| if not isinstance(self.prompt, str) or not self.prompt.strip(): | |
| raise ValueError("Enter a prompt.") | |
| for name in ("height", "width", "steps", "seed", "num_images", "max_sequence_length"): | |
| value = getattr(self, name) | |
| if not isinstance(value, int) or isinstance(value, bool): | |
| raise ValueError(f"{name} must be an integer.") | |
| if self.height <= 0 or self.width <= 0 or self.height % 16 or self.width % 16: | |
| raise ValueError("Height and width must be positive multiples of 16.") | |
| if self.steps < 1 or self.num_images < 1: | |
| raise ValueError("Steps and image count must be positive.") | |
| if not math.isfinite(self.cfg) or self.cfg < 0 or not math.isfinite(self.shift) or self.shift <= 0: | |
| raise ValueError("CFG must be finite and nonnegative; shift must be finite and positive.") | |
| if not -1 <= self.seed < 2**53: | |
| raise ValueError("Seed must be -1 (random) or an integer from 0 to 2⁵³−1 (browser-safe range).") | |
| if not 1 <= self.max_sequence_length <= 4096: | |
| raise ValueError("Maximum text length must be between 1 and 4096.") | |
| if self.offload not in ("none", "model", "sequential"): | |
| raise ValueError("Unknown offload mode.") | |
| if not isinstance(self.compile_transformer, bool): | |
| raise ValueError("Compile transformer must be true or false.") | |
| if self.offload == "sequential" and self.compile_transformer: | |
| raise ValueError("Compile transformer is incompatible with sequential offload.") | |
| if not math.isfinite(self.slg_scale) or self.slg_scale < 0: | |
| raise ValueError("SLG scale must be finite and nonnegative.") | |
| if not 0 <= self.slg_start <= self.slg_stop <= 1: | |
| raise ValueError("Require 0 <= SLG start <= SLG stop <= 1.") | |
| self.layers("uncond_skip_layers") | |
| self.layers("slg_layers") | |
| def pipeline_kwargs(self): | |
| return dict( | |
| prompt=self.prompt, | |
| negative_prompt=self.negative_prompt, | |
| height=self.height, | |
| width=self.width, | |
| num_inference_steps=self.steps, | |
| guidance_scale=self.cfg, | |
| num_images_per_prompt=self.num_images, | |
| max_sequence_length=self.max_sequence_length, | |
| uncond_skip_layers=self.layers("uncond_skip_layers"), | |
| slg_layers=self.layers("slg_layers"), | |
| slg_scale=self.slg_scale, | |
| slg_start=self.slg_start, | |
| slg_stop=self.slg_stop, | |
| ) | |
| def read_prompt_defaults(path): | |
| headers = {"[prompt_prefix]": "prompt_prefix", "[negative_prompt]": "negative_prompt"} | |
| sections = {} | |
| current = None | |
| for line in Path(path).read_text(encoding="utf-8").splitlines(): | |
| header = line.strip() | |
| if header in headers: | |
| current = headers[header] | |
| if current in sections: | |
| raise ValueError(f"Duplicate section {header} in {path}.") | |
| sections[current] = [] | |
| continue | |
| if current is None: | |
| if header and not header.startswith("#"): | |
| raise ValueError(f"Expected [prompt_prefix] or [negative_prompt] in {path}.") | |
| continue | |
| sections[current].append(line) | |
| if sections.keys() != set(headers.values()): | |
| raise ValueError(f"Both [prompt_prefix] and [negative_prompt] are required in {path}.") | |
| return {name: "\n".join(lines).strip() for name, lines in sections.items()} | |
| def parse_parameters(parameters, version): | |
| if version not in (1, 2): | |
| raise ValueError("Unsupported generation metadata version.") | |
| parameters = dict(parameters) | |
| if version == 1: | |
| # Version 1 called negative-branch skipping "slg_layers". | |
| layers = parameters.pop("slg_layers", "") | |
| parameters["uncond_skip_layers"] = layers if parameters.pop("slg", True) else "" | |
| if parameters.get("seed") is None: | |
| parameters["seed"] = -1 | |
| parameters.setdefault("compile_transformer", False) | |
| if parameters.get("offload") == "staged": | |
| parameters["offload"] = "model" | |
| config = GenerationConfig(**parameters) | |
| config.validate() | |
| return config | |
| def load_pipeline(model, shared_pipe=None): | |
| import torch | |
| from huggingface_hub import HfApi, snapshot_download | |
| from pipeline_moss_image import MOSSImagePipeline, MOSSImageScheduler, MOSSImageTransformer | |
| if model not in MODEL_IDS: | |
| raise ValueError("Choose an official MOSS-Image model repository.") | |
| info = HfApi().model_info(model, 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 = {key: value for key, value in shared_hashes.items() if key.startswith(f"{name}/")} | |
| previous = {key: value for key, value in shared_pipe._shared_component_hashes.items() | |
| if key.startswith(f"{name}/")} | |
| if current and all(current.values()) and current == previous: | |
| components[name] = getattr(shared_pipe, name) | |
| # Download unshared components and keep the Space's bundled inference classes. | |
| snapshot = snapshot_download( | |
| model, 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, | |
| ) | |
| pipe.to("cuda") | |
| pipe._prompt_cache = OrderedDict() | |
| pipe._shared_component_hashes = shared_hashes | |
| pipe._moss_revision = revision | |
| return pipe | |
| class ModelRunner: | |
| def __init__(self, output_dir): | |
| self.output_dir = Path(output_dir).expanduser().resolve() | |
| self.pipelines = {} | |
| def generate(self, model, config, progress=None): | |
| import torch | |
| logger.info( | |
| "Generation request: %s", | |
| json.dumps({"model": model, "parameters": asdict(config)}, ensure_ascii=False), | |
| ) | |
| config.validate() | |
| if model not in MODEL_IDS: | |
| raise ValueError("Choose an official MOSS-Image model repository.") | |
| if config.offload != "none" or config.compile_transformer: | |
| raise ValueError("This Space uses fixed GPU residency; runtime torch.compile is disabled.") | |
| pipe = self.pipelines.get(model) | |
| if pipe is None: | |
| raise ValueError("The selected model is not loaded. UI preview does not run inference.") | |
| if any(layer >= pipe.transformer.config.num_layers for layer in config.layers("uncond_skip_layers") + config.layers("slg_layers")): | |
| raise ValueError(f"SLG layer index must be below {pipe.transformer.config.num_layers}.") | |
| if config.seed == -1: | |
| config = replace(config, seed=secrets.randbelow(2**32)) | |
| logger.info("Generation seed resolved: model=%s seed=%d", model, config.seed) | |
| start = time.monotonic() | |
| total_steps = config.steps | |
| stage_seconds = {} | |
| current_stage = None | |
| stage_started = start | |
| def finish_stage(stage): | |
| nonlocal current_stage, stage_started | |
| if torch.cuda.is_available(): | |
| torch.cuda.synchronize() | |
| now = time.monotonic() | |
| if current_stage is not None: | |
| stage_seconds[current_stage] = round(now - stage_started, 3) | |
| current_stage, stage_started = stage, now | |
| def report(value, **kwargs): | |
| if progress is not None: | |
| progress(value, **kwargs) | |
| def report_stage(stage): | |
| nonlocal total_steps | |
| finish_stage(stage) | |
| if stage == "denoise": | |
| timesteps = getattr(getattr(pipe, "scheduler", None), "timesteps", None) | |
| if timesteps is not None and len(timesteps): | |
| total_steps = len(timesteps) | |
| descriptions = { | |
| "encode": "Encoding prompt", | |
| "denoise": f"Sampling · step 0/{total_steps}", | |
| "decode": "Decoding image", | |
| } | |
| report((0, total_steps) if stage == "denoise" else 0, desc=descriptions[stage], unit="steps") | |
| def report_step(pipe, step_index, timestep, callback_kwargs): | |
| completed = step_index + 1 | |
| report( | |
| (completed, total_steps), desc=f"Sampling · step {completed}/{total_steps}", unit="steps" | |
| ) | |
| return callback_kwargs | |
| try: | |
| set_scheduler(pipe, config.shift) | |
| images = generate_with_progress( | |
| pipe, report_stage, report_step, | |
| **config.pipeline_kwargs(), | |
| generator=torch.Generator("cpu").manual_seed(config.seed), | |
| ).images | |
| except torch.cuda.OutOfMemoryError as exc: | |
| torch.cuda.empty_cache() | |
| raise ValueError( | |
| "GPU memory exhausted. Try fewer or smaller images." | |
| ) from exc | |
| finish_stage(None) | |
| report(0, desc="Saving images and parameters") | |
| folder = self.output_dir / uuid4().hex | |
| folder.mkdir(parents=True, exist_ok=True) | |
| paths = [] | |
| for i, image in enumerate(images): | |
| path = folder / f"{i + 1:02d}.png" | |
| image.save(path) | |
| paths.append(str(path)) | |
| configs = {} | |
| for name in ("transformer", "text_encoder", "vae", "scheduler"): | |
| component_config = getattr(pipe, name).config | |
| configs[name] = ( | |
| component_config.to_dict() | |
| if hasattr(component_config, "to_dict") else dict(component_config) | |
| ) | |
| record = { | |
| "schema_version": 2, | |
| "model": model.rsplit("/", 1)[-1], | |
| "model_repo": model, | |
| "model_revision": pipe._moss_revision, | |
| "parameters": asdict(config), | |
| "sampler": "Euler", | |
| "weight_dtype": "bfloat16", | |
| "vae_dtype": "float32", | |
| "pipeline_sha256": hashlib.sha256( | |
| Path(__file__).with_name("pipeline_moss_image.py").read_bytes() | |
| ).hexdigest(), | |
| "model_config_sha256": { | |
| name: hashlib.sha256(json.dumps(config, sort_keys=True, default=str).encode()).hexdigest() | |
| for name, config in configs.items() | |
| }, | |
| "versions": { | |
| name: importlib.metadata.version(name) | |
| for name in ("torch", "diffusers", "transformers", "accelerate") | |
| }, | |
| "stage_seconds": stage_seconds, | |
| "execution_backend": getattr(pipe, "_execution_backend", "eager"), | |
| "gpu": torch.cuda.get_device_name(0) if torch.cuda.is_available() else None, | |
| "elapsed_seconds": round(time.monotonic() - start, 3), | |
| "images": [Path(p).name for p in paths], | |
| } | |
| metadata = folder / "parameters.json" | |
| metadata.write_text(json.dumps(record, indent=2, ensure_ascii=False) + "\n") | |
| report(1, desc="Complete") | |
| return ( | |
| paths, | |
| str(metadata), | |
| f"Saved {len(images)} image(s) · seed {config.seed} · {record['elapsed_seconds']:.1f}s", | |
| ) | |
| def save_showcase(metadata, destination): | |
| """Append the completed batch with its recorded parameters to the examples.""" | |
| if not metadata: | |
| raise ValueError("Generate an image successfully before adding it to showcase.") | |
| metadata = Path(metadata) | |
| record = json.loads(metadata.read_text()) | |
| destination = Path(destination) | |
| destination.mkdir(parents=True, exist_ok=True) | |
| saved = [] | |
| # Stage outside the watched directory so the gallery sees complete examples. | |
| with tempfile.TemporaryDirectory(dir=destination.parent) as temporary: | |
| for index, name in enumerate(record["images"]): | |
| folder = Path(temporary) / f"showcase-{uuid4().hex}" | |
| folder.mkdir() | |
| shutil.copyfile(metadata.parent / name, folder / "image.png") | |
| entry = {**record, "image": "image.png", "images": ["image.png"], | |
| "image_index": index, "source_images": record["images"]} | |
| (folder / "parameters.json").write_text( | |
| json.dumps(entry, ensure_ascii=False, indent=2) + "\n" | |
| ) | |
| saved.append(destination / folder.name) | |
| for target in saved: | |
| (Path(temporary) / target.name).rename(target) | |
| return saved | |
| def delete_showcase(metadata, destination): | |
| """Remove only a selected example directory within the showcase.""" | |
| if not metadata: | |
| raise ValueError("Select a showcase image to delete first.") | |
| metadata = Path(metadata).resolve() | |
| destination = Path(destination).resolve() | |
| if metadata.name != "parameters.json" or metadata.parent.parent != destination: | |
| raise ValueError("Selected example must be inside the showcase directory.") | |
| if not metadata.is_file(): | |
| raise ValueError("This showcase image has already been removed.") | |
| shutil.rmtree(metadata.parent) | |
| UI_ASSETS = Path(str(files("assets"))) | |
| EXAMPLES = UI_ASSETS / "images" | |
| RATIO_PRESETS = ("9:16", "2:3", "3:4", "1:1", "4:3", "3:2", "16:9") | |
| def build_app( | |
| model=DEFAULT_MODEL, | |
| runner=None, | |
| output_dir="outputs/gradio", | |
| prompt_defaults=UI_ASSETS / "prompt_defaults.txt", | |
| ): | |
| import gradio as gr | |
| runner = runner or ModelRunner(output_dir) | |
| choices = MODEL_IDS | |
| initial = model | |
| height_default, width_default = 1024, 1024 | |
| pixels_default, ratio_default = size_selection(width_default, height_default) | |
| min_budget, max_budget = PIXELS[0] / 1048576, PIXELS[-1] / 1048576 | |
| defaults = GenerationConfig() | |
| texts = read_prompt_defaults(prompt_defaults) | |
| asset_lock = threading.RLock() | |
| def asset_snapshot(): | |
| with asset_lock: | |
| return read_asset_snapshot() | |
| def read_asset_snapshot(): | |
| paths = sorted(EXAMPLES.glob("*/parameters.json")) | |
| images = [] | |
| stamps = [] | |
| for path in paths: | |
| record = json.loads(path.read_text()) | |
| image = (path.parent / record.get("image", "image.webp")).resolve() | |
| if not image.is_relative_to(EXAMPLES.resolve()): | |
| raise ValueError("Example images must be inside the assets directory.") | |
| images.append(str(image)) | |
| stamps.append([str(path), path.stat().st_mtime_ns, image.stat().st_mtime_ns]) | |
| return {"paths": [str(p) for p in paths], "images": images, "stamps": stamps, | |
| "defaults": Path(prompt_defaults).stat().st_mtime_ns, "selected": None} | |
| initial_assets = asset_snapshot() | |
| with gr.Blocks(title="MOSS-Image", analytics_enabled=False) as app: | |
| light_logo = app.serve_static_file(UI_ASSETS / "lockup-light.svg") | |
| dark_logo = app.serve_static_file(UI_ASSETS / "lockup-dark.svg") | |
| report = app.serve_static_file(UI_ASSETS / "MOSS_Image_Technical_Report.pdf") | |
| modelscope_base = app.serve_static_file(UI_ASSETS / "modelscope-base-badge.svg") | |
| modelscope_sft = app.serve_static_file(UI_ASSETS / "modelscope-sft-badge.svg") | |
| feishu_badge = app.serve_static_file(UI_ASSETS / "feishu-builders-badge.svg") | |
| gr.HTML( | |
| '<header style="text-align:center;padding:20px 0 24px">' | |
| f'<img class="moss-logo moss-logo-light" ' | |
| f'src="{html.escape(light_logo["url"].lstrip("/"), quote=True)}" alt="MOSS-Image">' | |
| f'<img class="moss-logo moss-logo-dark" ' | |
| f'src="{html.escape(dark_logo["url"].lstrip("/"), quote=True)}" alt="MOSS-Image">' | |
| '<nav aria-label="Project links" style="display:flex;justify-content:center;' | |
| 'align-items:center;flex-wrap:wrap;gap:8px;margin-top:16px">' | |
| '<a target="_blank" rel="noopener noreferrer" href="https://github.com/OpenMOSS/MOSS-Image"><img alt="GitHub" src="https://img.shields.io/badge/GitHub-Code-181717?style=flat-square&logo=github&logoColor=white"></a>' | |
| '<a target="_blank" rel="noopener noreferrer" href="https://huggingface.co/OpenMOSS-Team/MOSS-Image-Base"><img alt="Hugging Face · MOSS-Image-Base" src="https://img.shields.io/badge/🤗%20Hugging%20Face-MOSS--Image--Base-ffc107?style=flat-square"></a>' | |
| '<a target="_blank" rel="noopener noreferrer" href="https://huggingface.co/OpenMOSS-Team/MOSS-Image-SFT"><img alt="Hugging Face · MOSS-Image-SFT" src="https://img.shields.io/badge/🤗%20Hugging%20Face-MOSS--Image--SFT-ffc107?style=flat-square"></a>' | |
| '<a target="_blank" rel="noopener noreferrer" href="https://modelscope.cn/models/openmoss/MOSS-Image-Base">' | |
| f'<img alt="ModelScope · MOSS-Image-Base" src="{html.escape(modelscope_base["url"].lstrip("/"), quote=True)}"></a>' | |
| '<a target="_blank" rel="noopener noreferrer" href="https://modelscope.cn/models/openmoss/MOSS-Image-SFT">' | |
| f'<img alt="ModelScope · MOSS-Image-SFT" src="{html.escape(modelscope_sft["url"].lstrip("/"), quote=True)}"></a>' | |
| '<a target="_blank" rel="noopener noreferrer" href="https://x.com/Open_MOSS"><img alt="Follow OpenMOSS on X" src="https://img.shields.io/badge/Follow-OpenMOSS-000000?style=flat-square&logo=x&logoColor=white"></a>' | |
| '<a target="_blank" rel="noopener noreferrer" href="https://qm.qq.com/q/hphM7VhaDu"><img alt="MOSS-Image QQ 交流群：755535143" src="https://img.shields.io/badge/QQ-755535143-12B7F5?style=flat-square&logo=qq&logoColor=white"></a>' | |
| '<a target="_blank" rel="noopener noreferrer" href="https://applink.feishu.cn/client/chat/chatter/add_by_link?link_token=7a3m2007-815f-409b-9c3d-a4129f5c058e">' | |
| f'<img alt="MOSS-Builders Feishu Community" src="{html.escape(feishu_badge["url"].lstrip("/"), quote=True)}"></a>' | |
| '<span title="Homepage URL coming soon"><img alt="Homepage" src="https://img.shields.io/badge/🌐%20Project-Homepage-2563eb?style=flat-square"></span>' | |
| f'<a target="_blank" rel="noopener noreferrer" href="{html.escape(report["url"].lstrip("/"), quote=True)}">' | |
| '<img alt="Technical Report" src="https://img.shields.io/badge/📄%20Paper-Technical_Report-b31b1b?style=flat-square"></a>' | |
| '</nav></header>', | |
| css_template=( | |
| '.moss-logo {width:520px;max-width:100%;height:auto;margin:auto}' | |
| '.moss-logo-light {display:block}.moss-logo-dark {display:none}' | |
| '.moss-logo-light:is(.dark *) {display:none}.moss-logo-dark:is(.dark *) {display:block}' | |
| ), | |
| ) | |
| with gr.Row(equal_height=True): | |
| with gr.Column(scale=3, min_width=640): | |
| with gr.Row(equal_height=False, scale=0, elem_id="model-toolbar"): | |
| checkpoint = gr.Dropdown( | |
| choices=[(Path(p).name, p) for p in choices], | |
| value=initial, | |
| interactive=True, | |
| label="Model", | |
| scale=3, | |
| min_width=200, | |
| ) | |
| with gr.Group(): | |
| prompt = gr.Textbox( | |
| value=texts["prompt_prefix"], | |
| label="Prompt", | |
| lines=4, | |
| max_lines=8, | |
| placeholder="Describe your image — subject, light, composition, style…", | |
| ) | |
| negative = gr.Textbox( | |
| value=texts["negative_prompt"], | |
| label="Negative Prompt", | |
| lines=1, | |
| max_lines=3, | |
| placeholder="What would you like to leave out?", | |
| ) | |
| with gr.Column(): | |
| resolution = gr.State(pixels_default) | |
| ratio = gr.State(ratio_default) | |
| with gr.Row(): | |
| resolution_slider = gr.Slider( | |
| minimum=min_budget, | |
| maximum=max_budget, | |
| step=1 / 16, | |
| value=pixels_default / 1048576, | |
| label="Resolution (M pixels)", | |
| scale=3, | |
| min_width=170, | |
| ) | |
| aspect_slider = gr.Slider( | |
| minimum=0.125, | |
| maximum=8, | |
| step=0.001, | |
| value=width_default / height_default, | |
| label="Aspect ratio (W/H)", | |
| scale=3, | |
| min_width=170, | |
| ) | |
| width = gr.Number( | |
| value=width_default, | |
| precision=0, | |
| label="Width", | |
| minimum=16, | |
| step=16, | |
| scale=1, | |
| min_width=80, | |
| ) | |
| height = gr.Number( | |
| value=height_default, | |
| precision=0, | |
| label="Height", | |
| minimum=16, | |
| step=16, | |
| scale=1, | |
| min_width=80, | |
| ) | |
| with gr.Row(): | |
| resolution_choices = gr.Radio( | |
| choices=[(pixel_label(n), n) for n in PIXELS], | |
| value=pixels_default, | |
| label="Resolution presets", | |
| scale=5, | |
| min_width=240, | |
| ) | |
| aspect_choices = gr.Radio( | |
| choices=list(RATIO_PRESETS), | |
| value=ratio_default if ratio_default in RATIO_PRESETS else None, | |
| label="Aspect ratio presets", | |
| scale=7, | |
| min_width=320, | |
| ) | |
| with gr.Row(equal_height=True): | |
| with gr.Column(scale=1, min_width=100): | |
| count = gr.Number(value=1, precision=0, label="Images", minimum=1) | |
| with gr.Column(scale=1, min_width=100): | |
| random_seed = gr.Button("Random", size="sm", render=False) | |
| seed = gr.Number( | |
| value=42, | |
| precision=0, | |
| label="Seed", | |
| minimum=-1, | |
| maximum=2**53 - 1, | |
| buttons=[random_seed], | |
| ) | |
| with gr.Column(scale=1, min_width=100): | |
| steps = gr.Number(value=defaults.steps, precision=0, minimum=1, label="Steps") | |
| with gr.Accordion("Advanced settings", open=False): | |
| with gr.Row(): | |
| cfg = gr.Number(value=defaults.cfg, label="CFG", minimum=0) | |
| shift = gr.Number(value=defaults.shift, label="Shift", minimum=0.01) | |
| # ZeroGPU fixes residency; AOT compilation runs before serving. | |
| offload = gr.State("none") | |
| compile_transformer = gr.State(False) | |
| skip_mode = gr.Radio( | |
| choices=[("Off", "off"), ("Negative-branch skipping", "negative"), | |
| ("Standard SLG", "standard")], | |
| value="negative", | |
| label="Skip layer mode", | |
| ) | |
| layers = gr.Textbox( | |
| value=defaults.uncond_skip_layers, | |
| label="Skip layers", | |
| info="Zero-based layer indices, separated by commas. Empty disables skipping.", | |
| ) | |
| with gr.Row(visible=False) as standard_settings: | |
| slg_layers = gr.Textbox( | |
| value=defaults.slg_layers, label="SLG layers", | |
| info="Zero-based layer indices, separated by commas.", | |
| ) | |
| slg_scale = gr.Number(value=defaults.slg_scale, minimum=0, label="SLG scale") | |
| with gr.Row() as skip_interval: | |
| slg_start = gr.Number(value=defaults.slg_start, minimum=0, maximum=1, label="Skip start") | |
| slg_stop = gr.Number(value=defaults.slg_stop, minimum=0, maximum=1, label="Skip stop") | |
| with gr.Column(scale=2): | |
| gallery = gr.Gallery(label="生成结果", columns=2, scale=1, object_fit="contain") | |
| status = gr.Textbox(label="Status", interactive=False, scale=0) | |
| add_showcase = gr.Button("Add to showcase", size="sm", visible=False) | |
| last_generation = gr.State(None) | |
| with gr.Row(): | |
| with gr.Column(scale=3, min_width=640): | |
| generate = gr.Button("Generate", variant="primary") | |
| gr.Column(scale=2) | |
| gr.Markdown("## Examples") | |
| with gr.Row(equal_height=False): | |
| with gr.Column(scale=3, min_width=360): | |
| examples = gr.Gallery( | |
| value=initial_assets["images"], | |
| label="Examples", | |
| columns=4, | |
| object_fit="contain", | |
| interactive=False, | |
| ) | |
| delete_example = gr.Button("Delete from showcase", size="sm", visible=False) | |
| with gr.Column(scale=2, min_width=300): | |
| example_prompt = gr.Textbox( | |
| label="Original Prompt", | |
| lines=5, | |
| max_lines=8, | |
| interactive=False, | |
| placeholder="Click an image to view its original prompt.", | |
| ) | |
| example_negative = gr.Textbox( | |
| label="Negative Prompt", | |
| lines=1, | |
| max_lines=3, | |
| interactive=False, | |
| ) | |
| example_settings = gr.Textbox( | |
| label="Original settings", | |
| lines=2, | |
| max_lines=3, | |
| interactive=False, | |
| ) | |
| fields = [ | |
| prompt, | |
| negative, | |
| height, | |
| width, | |
| steps, | |
| cfg, | |
| shift, | |
| seed, | |
| count, | |
| layers, | |
| slg_layers, | |
| slg_scale, | |
| slg_start, | |
| slg_stop, | |
| offload, | |
| compile_transformer, | |
| ] | |
| field_names = ( | |
| "prompt", | |
| "negative_prompt", | |
| "height", | |
| "width", | |
| "steps", | |
| "cfg", | |
| "shift", | |
| "seed", | |
| "num_images", | |
| "uncond_skip_layers", | |
| "slg_layers", | |
| "slg_scale", | |
| "slg_start", | |
| "slg_stop", | |
| "offload", | |
| "compile_transformer", | |
| ) | |
| def run(selected, progress=gr.Progress(), *values): | |
| try: | |
| parameters = dict(zip(field_names, values)) | |
| if parameters.get("negative_prompt") is None: | |
| parameters["negative_prompt"] = "" | |
| config = GenerationConfig(**parameters) | |
| images, metadata, message = runner.generate(selected or "", config, progress=progress) | |
| return images, message, metadata | |
| except (ValueError, OSError) as exc: | |
| logger.exception("Generation failed") | |
| raise gr.Error(str(exc)) from exc | |
| def add_to_showcase(metadata): | |
| try: | |
| saved = save_showcase(metadata, EXAMPLES) | |
| return f"Added {len(saved)} image(s) to showcase. Examples refresh automatically." | |
| except (ValueError, OSError) as exc: | |
| logger.exception("Saving showcase failed") | |
| raise gr.Error(str(exc)) from exc | |
| def load_example(path): | |
| path = Path(path) | |
| record = json.loads(path.read_text()) | |
| # Examples supply image settings; keep the user's runtime configuration. | |
| parameters = { | |
| key: value | |
| for key, value in record["parameters"].items() | |
| if key not in ("offload", "compile_transformer") | |
| } | |
| config = parse_parameters(parameters, record.get("schema_version", 1)) | |
| model_id = record.get("model_repo") or record.get("model") | |
| if model_id not in MODEL_IDS: | |
| model_id = next((item for item in MODEL_IDS | |
| if item.rsplit("/", 1)[-1] == model_id), None) | |
| # Legacy examples without model metadata keep the current selection. | |
| return ( | |
| config.prompt, | |
| config.negative_prompt, | |
| config.height, | |
| config.width, | |
| config.steps, | |
| config.cfg, | |
| config.shift, | |
| config.num_images, | |
| config.seed, | |
| gr.update(value=config.uncond_skip_layers, visible=bool(config.uncond_skip_layers)), | |
| config.slg_layers, | |
| config.slg_scale, | |
| config.slg_start, | |
| config.slg_stop, | |
| "standard" if config.slg_layers and config.slg_scale > 0 else | |
| "negative" if config.uncond_skip_layers else "off", | |
| gr.update(visible=bool(config.slg_layers and config.slg_scale > 0)), | |
| gr.update(visible=bool(config.uncond_skip_layers or config.slg_layers and config.slg_scale > 0)), | |
| model_id if model_id else gr.skip(), | |
| config.prompt, | |
| config.negative_prompt, | |
| ( | |
| f"Model: {model_id or 'unknown (keep current selection)'}\n" | |
| f"{config.width} × {config.height} · {config.steps} steps · CFG {config.cfg:g}\n" | |
| f"Euler · Seed {parameters.get('seed') if parameters.get('seed') is not None else 'unknown (random on generation)'}\n" | |
| f"Shift {config.shift:g} · " | |
| f"Uncond skip {config.uncond_skip_layers or 'off'} · " | |
| f"Standard SLG {config.slg_layers or 'off'} (scale {config.slg_scale:g}) · " | |
| f"Interval [{config.slg_start:g}, {config.slg_stop:g})" | |
| ), | |
| ) | |
| assets = gr.State(initial_assets) | |
| def choose_example(state, evt: gr.SelectData): | |
| with asset_lock: | |
| if evt.index < 0 or evt.index >= len(state["paths"]): | |
| raise gr.Error("Choose an example image.") | |
| path = state["paths"][evt.index] | |
| if not Path(path).is_file(): | |
| raise gr.Error("This showcase image has already been removed.") | |
| return (*load_example(path), {**state, "selected": path}) | |
| def refresh_assets(state): | |
| with asset_lock: | |
| fresh = asset_snapshot() | |
| changed = fresh["paths"] != state["paths"] or fresh["stamps"] != state["stamps"] | |
| selected = state["selected"] | |
| fresh["selected"] = selected if selected in fresh["paths"] else None | |
| values = [gr.skip() for _ in example_fields] | |
| if fresh["selected"]: | |
| before = state["stamps"][state["paths"].index(selected)] | |
| after = fresh["stamps"][fresh["paths"].index(selected)] | |
| if before != after: | |
| values = list(load_example(selected)) | |
| if selected and not fresh["selected"]: | |
| values[-3:] = ["", "", "Example removed from assets."] | |
| if fresh["defaults"] != state["defaults"]: | |
| defaults = read_prompt_defaults(prompt_defaults) | |
| values[0] = defaults["prompt_prefix"] | |
| values[1] = defaults["negative_prompt"] | |
| selection = fresh["paths"].index(selected) if fresh["selected"] else None | |
| gallery = gr.update(value=fresh["images"], selected_index=selection) if changed else gr.skip() | |
| return (gallery, fresh, *values) | |
| def remove_from_showcase(state): | |
| with asset_lock: | |
| try: | |
| selected = state.get("selected") | |
| if selected not in state["paths"]: | |
| raise ValueError("Select a showcase image to delete first.") | |
| delete_showcase(selected, EXAMPLES) | |
| return (*refresh_assets(state), "Deleted selected image from showcase.") | |
| except (ValueError, OSError) as exc: | |
| raise gr.Error(str(exc)) from exc | |
| def update_skip(mode): | |
| negative = mode == "negative" | |
| standard = mode == "standard" | |
| return ( | |
| gr.update(value="22" if negative else "", visible=negative), | |
| "", 3.0 if standard else 0.0, 0.01 if standard else 0.0, 1.0, | |
| gr.update(visible=standard), gr.update(visible=mode != "off"), | |
| ) | |
| def apply_ratio(selected, pixels, w, h): | |
| if selected == "custom" or (not pixels and (not w or not h)): | |
| return gr.skip(), gr.skip() | |
| return fit_size(pixels or w * h, selected) | |
| def apply_resolution(pixels, selected, w, h): | |
| if not pixels or (selected == "custom" and (not w or not h)): | |
| return gr.skip(), gr.skip() | |
| return fit_size(pixels, f"{w}:{h}" if selected == "custom" else selected) | |
| def apply_budget(value, selected, w, h): | |
| return apply_resolution(value * 1048576, selected, w, h) | |
| def apply_aspect(value, pixels, w, h): | |
| return apply_ratio(f"{value}:1", pixels, w, h) | |
| def sync_size(w, h): | |
| pixels, selected = size_selection(w, h) | |
| choices = [pixels, selected, pixels or None, selected if selected in RATIO_PRESETS else None] | |
| if not w or not h or w <= 0 or h <= 0: | |
| return choices + [gr.skip(), gr.skip()] | |
| budget = (pixels or w * h) / 1048576 | |
| aspect = ratio_value(selected) if selected != "custom" else w / h | |
| return choices + [min(max_budget, max(min_budget, budget)), min(8, max(0.125, aspect))] | |
| resolution_slider.release( | |
| apply_budget, [resolution_slider, ratio, width, height], [width, height], queue=False | |
| ) | |
| aspect_slider.release( | |
| apply_aspect, [aspect_slider, resolution, width, height], [width, height], queue=False | |
| ) | |
| resolution_choices.input( | |
| apply_resolution, | |
| [resolution_choices, ratio, width, height], | |
| [width, height], | |
| queue=False, | |
| ) | |
| aspect_choices.input( | |
| apply_ratio, | |
| [aspect_choices, resolution, width, height], | |
| [width, height], | |
| queue=False, | |
| ) | |
| for field in (width, height): | |
| field.change( | |
| sync_size, | |
| [width, height], | |
| [ | |
| resolution, | |
| ratio, | |
| resolution_choices, | |
| aspect_choices, | |
| resolution_slider, | |
| aspect_slider, | |
| ], | |
| queue=False, | |
| api_name=False, | |
| ) | |
| random_seed.click(partial(randbits, 32), outputs=seed, queue=False, api_name=False) | |
| example_fields = [ | |
| prompt, negative, height, width, steps, cfg, shift, count, seed, layers, | |
| slg_layers, slg_scale, slg_start, slg_stop, | |
| skip_mode, standard_settings, skip_interval, | |
| checkpoint, example_prompt, example_negative, example_settings, | |
| ] | |
| examples.select( | |
| choose_example, inputs=assets, outputs=[*example_fields, assets], | |
| queue=False, api_name=False, | |
| ) | |
| gr.Timer(2).tick( | |
| refresh_assets, inputs=assets, outputs=[examples, assets, *example_fields], | |
| queue=False, api_name=False, | |
| ) | |
| skip_mode.input( | |
| update_skip, inputs=skip_mode, | |
| outputs=[layers, slg_layers, slg_scale, slg_start, slg_stop, standard_settings, skip_interval], | |
| queue=False, api_name=False, | |
| ) | |
| delete_example.click( | |
| remove_from_showcase, inputs=assets, | |
| outputs=[examples, assets, *example_fields, status], | |
| concurrency_limit=1, concurrency_id="showcase", api_name=False, | |
| ) | |
| add_showcase.click( | |
| add_to_showcase, inputs=last_generation, outputs=status, | |
| concurrency_limit=1, concurrency_id="showcase", api_name=False, | |
| ) | |
| generate.click( | |
| run, | |
| [checkpoint] + fields, | |
| [gallery, status, last_generation], | |
| # ZeroGPU isolates GPU tasks in workers; local runs share the pipeline. | |
| concurrency_limit=32 if SpacesConfig.zero_gpu else 1, | |
| concurrency_id="model", | |
| api_name="generate", | |
| ) | |
| return app.queue(default_concurrency_limit=1, max_size=32) | |