Added experimental support for Ideogram 4
This commit is contained in:
Jaret Burkett
@@ -52,6 +52,7 @@ AI Toolkit is an easy to use all in one training suite for diffusion models. I t
|
||||
|
||||
### Experimental
|
||||
- [lodestones/Zeta-Chroma](https://huggingface.co/lodestones/Zeta-Chroma) (Zeta Chroma)
|
||||
- [ideogram-ai/ideogram-4-fp8](https://huggingface.co/ideogram-ai/ideogram-4-fp8) (Ideogram 4 FP8)
|
||||
|
||||
## Installation
|
||||
|
||||
|
||||
@@ -13,6 +13,7 @@ from .ernie_image import ErnieImageModel
|
||||
from .nucleus_image import NucleusImageModel
|
||||
from .hidream.hidream_o1_model import HidreamO1Model
|
||||
from .z_image.z_image_l2p_model import ZImageL2PModel
|
||||
from .ideogram4 import Ideogram4Model
|
||||
|
||||
AI_TOOLKIT_MODELS = [
|
||||
# put a list of models here
|
||||
@@ -40,4 +41,5 @@ AI_TOOLKIT_MODELS = [
|
||||
NucleusImageModel,
|
||||
HidreamO1Model,
|
||||
ZImageL2PModel,
|
||||
Ideogram4Model,
|
||||
]
|
||||
|
||||
@@ -0,0 +1 @@
|
||||
from .ideogram4 import Ideogram4Model
|
||||
534
extensions_built_in/diffusion_models/ideogram4/ideogram4.py
Normal file
534
extensions_built_in/diffusion_models/ideogram4/ideogram4.py
Normal file
@@ -0,0 +1,534 @@
|
||||
import os
|
||||
from typing import List, Optional
|
||||
|
||||
import torch
|
||||
import yaml
|
||||
from safetensors.torch import load_file, save_file
|
||||
|
||||
from toolkit.config_modules import GenerateImageConfig, ModelConfig
|
||||
from toolkit.models.base_model import BaseModel
|
||||
from toolkit.basic import flush
|
||||
from toolkit.print import print_acc
|
||||
from toolkit.prompt_utils import PromptEmbeds
|
||||
from toolkit.samplers.custom_flowmatch_sampler import (
|
||||
CustomFlowMatchEulerDiscreteScheduler,
|
||||
)
|
||||
from toolkit.accelerator import unwrap_model
|
||||
from toolkit.metadata import get_meta_for_safetensors
|
||||
from toolkit.memory_management import MemoryManager
|
||||
from toolkit.util.quantize import quantize, get_qtype, quantize_model
|
||||
from optimum.quanto import freeze, QTensor
|
||||
|
||||
import huggingface_hub
|
||||
from huggingface_hub.errors import EntryNotFoundError
|
||||
from transformers import AutoModel, AutoTokenizer
|
||||
|
||||
from .src.transformer import Ideogram4Config, Ideogram4Transformer2DModel
|
||||
from .src.vae import AutoEncoder, AutoEncoderParams, convert_diffusers_state_dict
|
||||
from .src.latent_norm import get_latent_norm
|
||||
from .src.pipeline import (
|
||||
Ideogram4Pipeline,
|
||||
get_qwen3_vl_features,
|
||||
patchify_latents,
|
||||
predict_velocity,
|
||||
unpatchify_latents,
|
||||
)
|
||||
|
||||
|
||||
scheduler_config = {
|
||||
"base_image_seq_len": 256,
|
||||
"base_shift": 0.5,
|
||||
"invert_sigmas": False,
|
||||
"max_image_seq_len": 4096,
|
||||
"max_shift": 1.15,
|
||||
"num_train_timesteps": 1000,
|
||||
"shift": 1.0,
|
||||
"shift_terminal": None,
|
||||
"stochastic_sampling": False,
|
||||
"time_shift_type": "exponential",
|
||||
"use_beta_sigmas": False,
|
||||
"use_dynamic_shifting": False,
|
||||
"use_exponential_sigmas": False,
|
||||
"use_karras_sigmas": False,
|
||||
}
|
||||
|
||||
# Weight-only FP8 (e4m3) Linear weights carry a per-output-channel float32 scale
|
||||
# saved alongside as ``<name>.weight_scale``. Folding it back gives bf16 weights.
|
||||
FP8_SCALE_SUFFIX = ".weight_scale"
|
||||
|
||||
# The text encoder is frozen, stock Qwen3-VL-8B-Instruct.
|
||||
QWEN3_VL_PATH = "Qwen/Qwen3-VL-8B-Instruct"
|
||||
|
||||
HF_TOKEN = os.getenv("HF_TOKEN", None)
|
||||
|
||||
|
||||
def _dequantize_fp8_state_dict(
|
||||
state_dict: dict,
|
||||
dtype: torch.dtype,
|
||||
device: torch.device,
|
||||
low_vram: bool,
|
||||
) -> dict:
|
||||
"""Fold weight-only FP8 scales back into the weights, casting to ``dtype``.
|
||||
|
||||
Linear weights stored as float8 with a sibling ``.weight_scale`` are
|
||||
reconstructed as ``weight_fp8.to(float32) * scale[:, None]``. Everything else
|
||||
is simply cast to ``dtype`` (non-floating tensors are left untouched). If the
|
||||
checkpoint isn't quantized this is just a dtype cast.
|
||||
|
||||
The fold/cast runs on ``device`` (GPU is much faster than CPU). With
|
||||
``low_vram=True`` each tensor is moved to ``device``, processed, then moved
|
||||
back to CPU so the whole bf16 model never sits on the GPU at once; otherwise
|
||||
the dequantized tensors are left on ``device`` ready to load.
|
||||
"""
|
||||
work_device = torch.device(device)
|
||||
|
||||
def _finish(t: torch.Tensor) -> torch.Tensor:
|
||||
return t.to("cpu") if low_vram else t
|
||||
|
||||
num_fp8 = sum(1 for k in state_dict if k.endswith(FP8_SCALE_SUFFIX))
|
||||
if num_fp8 > 0:
|
||||
print_acc(f" dequantizing {num_fp8} fp8 weights -> {dtype} on {work_device}")
|
||||
else:
|
||||
print_acc(f" casting weights -> {dtype} on {work_device}")
|
||||
|
||||
out = {}
|
||||
for key, tensor in state_dict.items():
|
||||
if key.endswith(FP8_SCALE_SUFFIX):
|
||||
continue
|
||||
scale_key = key + "_scale"
|
||||
if key.endswith(".weight") and scale_key in state_dict:
|
||||
w = tensor.to(work_device, torch.float32)
|
||||
scale = state_dict[scale_key].to(work_device, torch.float32)
|
||||
out[key] = _finish((w * scale.unsqueeze(1)).to(dtype))
|
||||
elif tensor.is_floating_point():
|
||||
out[key] = _finish(tensor.to(work_device, dtype))
|
||||
else:
|
||||
out[key] = tensor
|
||||
return out
|
||||
|
||||
|
||||
def _load_component_state_dict(base: str, subfolder: str, basename: str) -> dict:
|
||||
"""Load a component's weights whether local or on the hub, sharded or single."""
|
||||
index_name = f"{basename}.safetensors.index.json"
|
||||
single_name = f"{basename}.safetensors"
|
||||
|
||||
# Local directory layout: <base>/<subfolder>/<file>
|
||||
local_dir = os.path.join(base, subfolder)
|
||||
if os.path.isdir(local_dir):
|
||||
index_path = os.path.join(local_dir, index_name)
|
||||
if os.path.exists(index_path):
|
||||
return _load_sharded(local_dir, index_path, is_local=True)
|
||||
return load_file(os.path.join(local_dir, single_name))
|
||||
|
||||
# Hub repo layout: <subfolder>/<file>
|
||||
prefix = f"{subfolder}/" if subfolder else ""
|
||||
try:
|
||||
index_path = huggingface_hub.hf_hub_download(
|
||||
repo_id=base, filename=f"{prefix}{index_name}", token=HF_TOKEN
|
||||
)
|
||||
return _load_sharded(base, index_path, is_local=False, prefix=prefix)
|
||||
except EntryNotFoundError:
|
||||
single_path = huggingface_hub.hf_hub_download(
|
||||
repo_id=base, filename=f"{prefix}{single_name}", token=HF_TOKEN
|
||||
)
|
||||
return load_file(single_path)
|
||||
|
||||
|
||||
def _load_sharded(base, index_path, is_local, prefix="") -> dict:
|
||||
import json
|
||||
|
||||
with open(index_path) as f:
|
||||
index = json.load(f)
|
||||
shard_files = sorted(set(index["weight_map"].values()))
|
||||
state_dict = {}
|
||||
num_shards = len(shard_files)
|
||||
for i, shard in enumerate(shard_files):
|
||||
if is_local:
|
||||
shard_path = os.path.join(base, shard)
|
||||
else:
|
||||
print_acc(f" downloading shard {i + 1}/{num_shards}: {shard}")
|
||||
shard_path = huggingface_hub.hf_hub_download(
|
||||
repo_id=base, filename=f"{prefix}{shard}", token=HF_TOKEN
|
||||
)
|
||||
print_acc(f" loading shard {i + 1}/{num_shards}: {shard}")
|
||||
state_dict.update(load_file(shard_path))
|
||||
return state_dict
|
||||
|
||||
|
||||
class Ideogram4Model(BaseModel):
|
||||
arch = "ideogram4"
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
device,
|
||||
model_config: ModelConfig,
|
||||
dtype="bf16",
|
||||
custom_pipeline=None,
|
||||
noise_scheduler=None,
|
||||
**kwargs,
|
||||
):
|
||||
super().__init__(
|
||||
device, model_config, dtype, custom_pipeline, noise_scheduler, **kwargs
|
||||
)
|
||||
self.is_flow_matching = True
|
||||
self.is_transformer = True
|
||||
self.target_lora_modules = ["Ideogram4Transformer2DModel"]
|
||||
|
||||
self.patch_size = 2
|
||||
self.vae_scale_factor = 8
|
||||
# Length the Qwen3-VL caption is padded/truncated to. Fixed so cached and
|
||||
# batched embeddings always share a sequence length.
|
||||
self.max_text_length = int(
|
||||
self.model_config.model_kwargs.get("max_text_length", 512)
|
||||
)
|
||||
|
||||
self._latent_shift = None
|
||||
self._latent_scale = None
|
||||
|
||||
@staticmethod
|
||||
def get_train_scheduler():
|
||||
return CustomFlowMatchEulerDiscreteScheduler(**scheduler_config)
|
||||
|
||||
def get_bucket_divisibility(self):
|
||||
# 8 for the VAE downsample, 2 for the patch size.
|
||||
return self.vae_scale_factor * self.patch_size
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Loading
|
||||
# ------------------------------------------------------------------
|
||||
def _load_text_encoder(self, base: str):
|
||||
dtype = self.torch_dtype
|
||||
# The text encoder is frozen, stock Qwen3-VL-8B-Instruct. The ideogram repo
|
||||
# only ships an fp8 copy of it, so load the public bf16 model directly --
|
||||
# faster and higher precision than dequantizing the fp8 weights.
|
||||
te_path = self.model_config.model_kwargs.get("text_encoder_path", QWEN3_VL_PATH)
|
||||
self.print_and_status_update(f"Loading Qwen3-VL text encoder from {te_path}")
|
||||
|
||||
tokenizer = AutoTokenizer.from_pretrained(te_path, token=HF_TOKEN)
|
||||
text_encoder = AutoModel.from_pretrained(
|
||||
te_path, torch_dtype=dtype, token=HF_TOKEN
|
||||
)
|
||||
flush()
|
||||
|
||||
text_encoder.eval()
|
||||
text_encoder.requires_grad_(False)
|
||||
return tokenizer, text_encoder
|
||||
|
||||
def _load_transformer(self, base: str):
|
||||
dtype = self.torch_dtype
|
||||
self.print_and_status_update("Loading transformer")
|
||||
|
||||
transformer_config = Ideogram4Config()
|
||||
with torch.device("meta"):
|
||||
transformer = Ideogram4Transformer2DModel(transformer_config)
|
||||
|
||||
self.print_and_status_update(" - fetching transformer weights")
|
||||
state_dict = _load_component_state_dict(
|
||||
base, "transformer", "diffusion_pytorch_model"
|
||||
)
|
||||
self.print_and_status_update(" - dequantizing transformer weights")
|
||||
state_dict = _dequantize_fp8_state_dict(
|
||||
state_dict, dtype, self.device_torch, self.model_config.low_vram
|
||||
)
|
||||
self.print_and_status_update(" - loading transformer state dict")
|
||||
transformer.load_state_dict(state_dict, assign=True)
|
||||
del state_dict
|
||||
flush()
|
||||
|
||||
# inv_freq is a non-persistent buffer absent from the checkpoint; rebuild
|
||||
# it now that the module is off the meta device.
|
||||
head_dim = transformer_config.emb_dim // transformer_config.num_heads
|
||||
inv_freq = 1.0 / (
|
||||
transformer_config.rope_theta
|
||||
** (torch.arange(0, head_dim, 2, dtype=torch.float32) / head_dim)
|
||||
)
|
||||
transformer.rotary_emb.register_buffer("inv_freq", inv_freq, persistent=False)
|
||||
return transformer
|
||||
|
||||
def _load_vae(self, base: str):
|
||||
dtype = self.torch_dtype
|
||||
self.print_and_status_update("Loading VAE")
|
||||
vae_sd = _load_component_state_dict(base, "vae", "diffusion_pytorch_model")
|
||||
vae_sd = convert_diffusers_state_dict(vae_sd)
|
||||
vae = AutoEncoder(AutoEncoderParams())
|
||||
vae.load_state_dict(vae_sd)
|
||||
del vae_sd
|
||||
vae.to(self.vae_device_torch, dtype=dtype)
|
||||
vae.eval()
|
||||
vae.requires_grad_(False)
|
||||
return vae
|
||||
|
||||
def load_model(self):
|
||||
dtype = self.torch_dtype
|
||||
self.print_and_status_update("Loading Ideogram4 model")
|
||||
base = self.model_config.name_or_path
|
||||
|
||||
transformer = self._load_transformer(base)
|
||||
|
||||
if self.model_config.quantize:
|
||||
self.print_and_status_update("Quantizing Transformer")
|
||||
quantize_model(self, transformer)
|
||||
flush()
|
||||
else:
|
||||
transformer.to(self.device_torch, dtype=dtype)
|
||||
flush()
|
||||
|
||||
if (
|
||||
self.model_config.layer_offloading
|
||||
and self.model_config.layer_offloading_transformer_percent > 0
|
||||
):
|
||||
MemoryManager.attach(
|
||||
transformer,
|
||||
self.device_torch,
|
||||
offload_percent=self.model_config.layer_offloading_transformer_percent,
|
||||
ignore_modules=[transformer.input_proj, transformer.llm_cond_proj],
|
||||
)
|
||||
elif self.model_config.low_vram:
|
||||
self.print_and_status_update("Moving transformer to CPU")
|
||||
transformer.to("cpu")
|
||||
else:
|
||||
# quantize_model leaves the model on CPU; make sure it lands on device.
|
||||
transformer.to(self.device_torch)
|
||||
flush()
|
||||
|
||||
tokenizer, text_encoder = self._load_text_encoder(base)
|
||||
if self.model_config.quantize_te:
|
||||
self.print_and_status_update("Quantizing Text Encoder")
|
||||
text_encoder.to(self.device_torch)
|
||||
quantize(text_encoder, weights=get_qtype(self.model_config.qtype_te))
|
||||
freeze(text_encoder)
|
||||
flush()
|
||||
if (
|
||||
self.model_config.layer_offloading
|
||||
and self.model_config.layer_offloading_text_encoder_percent > 0
|
||||
):
|
||||
MemoryManager.attach(
|
||||
text_encoder,
|
||||
self.device_torch,
|
||||
offload_percent=self.model_config.layer_offloading_text_encoder_percent,
|
||||
)
|
||||
elif self.model_config.low_vram:
|
||||
self.print_and_status_update("Moving text encoder to CPU")
|
||||
text_encoder.to("cpu")
|
||||
else:
|
||||
self.print_and_status_update("Moving text encoder to device")
|
||||
text_encoder.to(self.device_torch)
|
||||
flush()
|
||||
|
||||
vae = self._load_vae(base)
|
||||
|
||||
self.noise_scheduler = Ideogram4Model.get_train_scheduler()
|
||||
|
||||
shift, scale = get_latent_norm()
|
||||
self._latent_shift = shift.view(1, -1, 1, 1)
|
||||
self._latent_scale = scale.view(1, -1, 1, 1)
|
||||
|
||||
self.vae = vae
|
||||
self.text_encoder = text_encoder
|
||||
self.tokenizer = tokenizer
|
||||
self.model = transformer
|
||||
self.pipeline = Ideogram4Pipeline(self)
|
||||
self.print_and_status_update("Model Loaded")
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Generation
|
||||
# ------------------------------------------------------------------
|
||||
def get_generation_pipeline(self):
|
||||
return Ideogram4Pipeline(self)
|
||||
|
||||
def generate_single_image(
|
||||
self,
|
||||
pipeline: Ideogram4Pipeline,
|
||||
gen_config: GenerateImageConfig,
|
||||
conditional_embeds: PromptEmbeds,
|
||||
unconditional_embeds: PromptEmbeds,
|
||||
generator: torch.Generator,
|
||||
extra: dict,
|
||||
):
|
||||
if self.model.device == torch.device("cpu"):
|
||||
self.model.to(self.device_torch)
|
||||
|
||||
sc = self.get_bucket_divisibility()
|
||||
gen_config.width = int(gen_config.width // sc * sc)
|
||||
gen_config.height = int(gen_config.height // sc * sc)
|
||||
|
||||
img = pipeline(
|
||||
conditional_embeds=conditional_embeds,
|
||||
unconditional_embeds=unconditional_embeds,
|
||||
height=gen_config.height,
|
||||
width=gen_config.width,
|
||||
num_inference_steps=gen_config.num_inference_steps,
|
||||
guidance_scale=gen_config.guidance_scale,
|
||||
latents=gen_config.latents,
|
||||
generator=generator,
|
||||
)[0]
|
||||
return img
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Training hooks
|
||||
# ------------------------------------------------------------------
|
||||
def get_noise_prediction(
|
||||
self,
|
||||
latent_model_input: torch.Tensor, # (B, 128, gh, gw)
|
||||
timestep: torch.Tensor, # 0 to 1000 scale
|
||||
text_embeddings: PromptEmbeds,
|
||||
**kwargs,
|
||||
):
|
||||
if self.model.device == torch.device("cpu"):
|
||||
self.model.to(self.device_torch)
|
||||
|
||||
t01 = timestep.to(self.device_torch, dtype=torch.float32) / 1000.0
|
||||
if t01.dim() == 0:
|
||||
t01 = t01.unsqueeze(0)
|
||||
if t01.shape[0] != latent_model_input.shape[0]:
|
||||
t01 = t01.expand(latent_model_input.shape[0])
|
||||
|
||||
llm_features = text_embeddings.text_embeds.to(self.device_torch)
|
||||
text_mask = text_embeddings.attention_mask.to(self.device_torch)
|
||||
|
||||
pred = predict_velocity(
|
||||
self.transformer,
|
||||
latent_model_input.to(self.device_torch),
|
||||
t01,
|
||||
llm_features,
|
||||
text_mask,
|
||||
)
|
||||
return pred
|
||||
|
||||
def get_prompt_embeds(self, prompt) -> PromptEmbeds:
|
||||
if isinstance(prompt, str):
|
||||
prompt = [prompt]
|
||||
|
||||
if self.text_encoder.device == torch.device("cpu"):
|
||||
self.text_encoder.to(self.device_torch)
|
||||
device = self.text_encoder.device
|
||||
|
||||
pad_id = self.tokenizer.pad_token_id
|
||||
if pad_id is None:
|
||||
pad_id = self.tokenizer.eos_token_id or 0
|
||||
|
||||
token_id_list = []
|
||||
for p in prompt:
|
||||
messages = [{"role": "user", "content": [{"type": "text", "text": p}]}]
|
||||
text = self.tokenizer.apply_chat_template(
|
||||
messages, add_generation_prompt=True, tokenize=False
|
||||
)
|
||||
ids = self.tokenizer(
|
||||
text,
|
||||
add_special_tokens=False,
|
||||
truncation=True,
|
||||
max_length=self.max_text_length,
|
||||
)["input_ids"]
|
||||
token_id_list.append(ids)
|
||||
|
||||
seq_len = self.max_text_length
|
||||
batch_size = len(token_id_list)
|
||||
token_ids = torch.full((batch_size, seq_len), pad_id, dtype=torch.long)
|
||||
attention_mask = torch.zeros((batch_size, seq_len), dtype=torch.long)
|
||||
for b, ids in enumerate(token_id_list):
|
||||
n = min(len(ids), seq_len)
|
||||
# left pad
|
||||
token_ids[b, seq_len - n :] = torch.tensor(ids[:n], dtype=torch.long)
|
||||
attention_mask[b, seq_len - n :] = 1
|
||||
|
||||
token_ids = token_ids.to(device)
|
||||
attention_mask = attention_mask.to(device)
|
||||
pos_2d = (attention_mask.cumsum(dim=-1) - 1).clamp(min=0).to(torch.long)
|
||||
|
||||
features = get_qwen3_vl_features(
|
||||
self.text_encoder, token_ids, attention_mask, pos_2d
|
||||
)
|
||||
|
||||
pe = PromptEmbeds(features.to(self.torch_dtype))
|
||||
pe.attention_mask = attention_mask
|
||||
return pe
|
||||
|
||||
def get_model_has_grad(self):
|
||||
return False
|
||||
|
||||
def get_te_has_grad(self):
|
||||
return False
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# VAE
|
||||
# ------------------------------------------------------------------
|
||||
def encode_images(self, image_list: List[torch.Tensor], device=None, dtype=None):
|
||||
if device is None:
|
||||
device = self.vae_device_torch
|
||||
if dtype is None:
|
||||
dtype = self.vae_torch_dtype
|
||||
if self.vae.device == torch.device("cpu"):
|
||||
self.vae.to(self.vae_device_torch)
|
||||
|
||||
if isinstance(image_list, list):
|
||||
images = torch.stack(image_list, dim=0)
|
||||
else:
|
||||
images = image_list
|
||||
images = images.to(device, dtype=dtype)
|
||||
|
||||
ae_channels = self.vae.params.z_channels
|
||||
moments = self.vae.encoder(images)
|
||||
mean = moments[:, :ae_channels]
|
||||
|
||||
patched = patchify_latents(mean, self.patch_size)
|
||||
shift = self._latent_shift.to(patched.device, patched.dtype)
|
||||
scale = self._latent_scale.to(patched.device, patched.dtype)
|
||||
latents = (patched - shift) / scale
|
||||
return latents.to(device, dtype=dtype)
|
||||
|
||||
def decode_latents(self, latents: torch.Tensor, device=None, dtype=None):
|
||||
if device is None:
|
||||
device = self.vae_device_torch
|
||||
if dtype is None:
|
||||
dtype = self.vae_torch_dtype
|
||||
if self.vae.device == torch.device("cpu"):
|
||||
self.vae.to(self.vae_device_torch)
|
||||
|
||||
latents = latents.to(device, dtype=dtype)
|
||||
shift = self._latent_shift.to(device, dtype)
|
||||
scale = self._latent_scale.to(device, dtype)
|
||||
patched = latents * scale + shift
|
||||
z = unpatchify_latents(patched, self.patch_size)
|
||||
images = self.vae.decoder(z)
|
||||
return images
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Saving / misc
|
||||
# ------------------------------------------------------------------
|
||||
def get_loss_target(self, *args, **kwargs):
|
||||
noise = kwargs.get("noise")
|
||||
batch = kwargs.get("batch")
|
||||
return (noise - batch.latents).detach()
|
||||
|
||||
def save_model(self, output_path, meta, save_dtype):
|
||||
if not output_path.endswith(".safetensors"):
|
||||
output_path = output_path + ".safetensors"
|
||||
transformer: Ideogram4Transformer2DModel = unwrap_model(self.model)
|
||||
state_dict = transformer.state_dict()
|
||||
save_dict = {}
|
||||
for k, v in state_dict.items():
|
||||
if isinstance(v, QTensor):
|
||||
v = v.dequantize()
|
||||
save_dict[k] = v.clone().to("cpu", dtype=save_dtype)
|
||||
meta = get_meta_for_safetensors(meta, name="ideogram4")
|
||||
save_file(save_dict, output_path, metadata=meta)
|
||||
|
||||
def get_base_model_version(self):
|
||||
return "ideogram4"
|
||||
|
||||
def get_transformer_block_names(self) -> Optional[List[str]]:
|
||||
return ["layers"]
|
||||
|
||||
def convert_lora_weights_before_save(self, state_dict):
|
||||
new_sd = {}
|
||||
for key, value in state_dict.items():
|
||||
new_key = key.replace("transformer.", "diffusion_model.")
|
||||
new_sd[new_key] = value
|
||||
return new_sd
|
||||
|
||||
def convert_lora_weights_before_load(self, state_dict):
|
||||
new_sd = {}
|
||||
for key, value in state_dict.items():
|
||||
new_key = key.replace("diffusion_model.", "transformer.")
|
||||
new_sd[new_key] = value
|
||||
return new_sd
|
||||
@@ -0,0 +1,272 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import torch
|
||||
|
||||
LATENT_SHIFT: tuple[float, ...] = (
|
||||
0.01984364,
|
||||
0.10149707,
|
||||
0.29689495,
|
||||
0.27188619,
|
||||
-0.21445648,
|
||||
-0.15979549,
|
||||
0.05021099,
|
||||
-0.15083604,
|
||||
-0.15360136,
|
||||
-0.20131799,
|
||||
0.01922352,
|
||||
0.0622626,
|
||||
0.10140969,
|
||||
-0.06739428,
|
||||
0.3758261,
|
||||
-0.233712,
|
||||
0.35164491,
|
||||
-0.02590912,
|
||||
-0.0271935,
|
||||
-0.10833897,
|
||||
-0.1476848,
|
||||
-0.01130957,
|
||||
-0.2298372,
|
||||
0.23526423,
|
||||
-0.10893522,
|
||||
0.11957631,
|
||||
0.04047799,
|
||||
0.3134589,
|
||||
-0.17225064,
|
||||
-0.18646109,
|
||||
-0.34691978,
|
||||
-0.03571246,
|
||||
0.02583857,
|
||||
0.10190072,
|
||||
0.28402294,
|
||||
0.26952152,
|
||||
-0.21634675,
|
||||
-0.17938656,
|
||||
0.04358909,
|
||||
-0.15007621,
|
||||
-0.1548502,
|
||||
-0.18971131,
|
||||
0.02710861,
|
||||
0.05609494,
|
||||
0.10697846,
|
||||
-0.06854968,
|
||||
0.38167698,
|
||||
-0.24269937,
|
||||
0.35705471,
|
||||
-0.03063305,
|
||||
-0.02946109,
|
||||
-0.11244286,
|
||||
-0.14336038,
|
||||
-0.01362137,
|
||||
-0.21863696,
|
||||
0.23228983,
|
||||
-0.11739769,
|
||||
0.11693044,
|
||||
0.02563311,
|
||||
0.31356594,
|
||||
-0.17420591,
|
||||
-0.19006285,
|
||||
-0.34905377,
|
||||
-0.04025005,
|
||||
0.01924137,
|
||||
0.07652984,
|
||||
0.2995608,
|
||||
0.2628057,
|
||||
-0.22011674,
|
||||
-0.12715361,
|
||||
0.04879879,
|
||||
-0.14075719,
|
||||
-0.15935895,
|
||||
-0.2123584,
|
||||
0.01974813,
|
||||
0.05523547,
|
||||
0.10011992,
|
||||
-0.06428964,
|
||||
0.37781868,
|
||||
-0.21491644,
|
||||
0.34254215,
|
||||
-0.03153528,
|
||||
-0.0310082,
|
||||
-0.10761415,
|
||||
-0.14730405,
|
||||
-0.02475182,
|
||||
-0.2285588,
|
||||
0.2515081,
|
||||
-0.10445128,
|
||||
0.12446,
|
||||
0.07062869,
|
||||
0.30880162,
|
||||
-0.18016875,
|
||||
-0.18869164,
|
||||
-0.34533499,
|
||||
-0.0129177,
|
||||
0.02578168,
|
||||
0.07993659,
|
||||
0.28642181,
|
||||
0.26038408,
|
||||
-0.22459419,
|
||||
-0.14820155,
|
||||
0.04059549,
|
||||
-0.14043529,
|
||||
-0.16111187,
|
||||
-0.2020305,
|
||||
0.02602069,
|
||||
0.04852717,
|
||||
0.10432153,
|
||||
-0.06309942,
|
||||
0.38402443,
|
||||
-0.22397003,
|
||||
0.34814481,
|
||||
-0.03774432,
|
||||
-0.03381438,
|
||||
-0.11245691,
|
||||
-0.14128767,
|
||||
-0.02853208,
|
||||
-0.21752016,
|
||||
0.24872463,
|
||||
-0.11399775,
|
||||
0.1222687,
|
||||
0.05620835,
|
||||
0.309178,
|
||||
-0.18065738,
|
||||
-0.19401479,
|
||||
-0.34495114,
|
||||
-0.01760592,
|
||||
)
|
||||
|
||||
LATENT_SCALE: tuple[float, ...] = (
|
||||
1.63933691,
|
||||
1.70204478,
|
||||
1.73642566,
|
||||
1.90004803,
|
||||
1.6675316,
|
||||
1.69059584,
|
||||
1.56853198,
|
||||
1.62314944,
|
||||
1.89106626,
|
||||
1.58086668,
|
||||
1.60822129,
|
||||
1.60962993,
|
||||
1.63322129,
|
||||
1.56074359,
|
||||
1.73419528,
|
||||
1.7919265,
|
||||
1.64040632,
|
||||
1.66802808,
|
||||
1.60390303,
|
||||
1.75480492,
|
||||
1.63187587,
|
||||
1.64334594,
|
||||
1.61722884,
|
||||
1.60146046,
|
||||
1.63459219,
|
||||
1.55291476,
|
||||
1.68771497,
|
||||
1.68415657,
|
||||
1.78966054,
|
||||
1.66631641,
|
||||
1.65626686,
|
||||
1.65976433,
|
||||
1.63487607,
|
||||
1.69513249,
|
||||
1.72933756,
|
||||
1.91310663,
|
||||
1.67035057,
|
||||
1.72286863,
|
||||
1.56719251,
|
||||
1.61934825,
|
||||
1.88628859,
|
||||
1.56911539,
|
||||
1.59455129,
|
||||
1.60829869,
|
||||
1.62470611,
|
||||
1.56052853,
|
||||
1.73677003,
|
||||
1.77563606,
|
||||
1.63732541,
|
||||
1.66370527,
|
||||
1.59508952,
|
||||
1.75153949,
|
||||
1.63029275,
|
||||
1.64517667,
|
||||
1.61659342,
|
||||
1.59722044,
|
||||
1.64103121,
|
||||
1.5408531,
|
||||
1.68610394,
|
||||
1.67772755,
|
||||
1.78998563,
|
||||
1.66621713,
|
||||
1.65458955,
|
||||
1.66041308,
|
||||
1.64710857,
|
||||
1.68163503,
|
||||
1.74000294,
|
||||
1.92784786,
|
||||
1.67411194,
|
||||
1.67395548,
|
||||
1.57406532,
|
||||
1.62199356,
|
||||
1.87618195,
|
||||
1.5584375,
|
||||
1.57438785,
|
||||
1.61711053,
|
||||
1.63094305,
|
||||
1.55644029,
|
||||
1.73124302,
|
||||
1.80666627,
|
||||
1.6463621,
|
||||
1.65932006,
|
||||
1.60816188,
|
||||
1.75682671,
|
||||
1.64695873,
|
||||
1.63121722,
|
||||
1.61380832,
|
||||
1.60478651,
|
||||
1.63396035,
|
||||
1.53505068,
|
||||
1.65534289,
|
||||
1.67132281,
|
||||
1.80317197,
|
||||
1.6767314,
|
||||
1.65700938,
|
||||
1.68426259,
|
||||
1.65339716,
|
||||
1.67540638,
|
||||
1.73298504,
|
||||
1.94067348,
|
||||
1.67893609,
|
||||
1.70635117,
|
||||
1.5730906,
|
||||
1.61928553,
|
||||
1.87148809,
|
||||
1.56244866,
|
||||
1.56697152,
|
||||
1.61584394,
|
||||
1.62759496,
|
||||
1.55480378,
|
||||
1.73484107,
|
||||
1.79055143,
|
||||
1.64688773,
|
||||
1.66121492,
|
||||
1.60135887,
|
||||
1.75254572,
|
||||
1.64798332,
|
||||
1.62989921,
|
||||
1.61381592,
|
||||
1.60792883,
|
||||
1.63939668,
|
||||
1.53075757,
|
||||
1.65371318,
|
||||
1.66801185,
|
||||
1.80029087,
|
||||
1.67591476,
|
||||
1.65655173,
|
||||
1.68533454,
|
||||
)
|
||||
|
||||
|
||||
def get_latent_norm() -> tuple[torch.Tensor, torch.Tensor]:
|
||||
shift = torch.tensor(LATENT_SHIFT, dtype=torch.float32)
|
||||
scale = torch.tensor(LATENT_SCALE, dtype=torch.float32)
|
||||
assert shift.shape == (128,) and scale.shape == (128,)
|
||||
return shift, scale
|
||||
310
extensions_built_in/diffusion_models/ideogram4/src/pipeline.py
Normal file
310
extensions_built_in/diffusion_models/ideogram4/src/pipeline.py
Normal file
@@ -0,0 +1,310 @@
|
||||
"""Packing / sampling helpers for Ideogram 4.
|
||||
|
||||
This module holds the glue that turns image latents + Qwen3-VL text features into
|
||||
the single packed sequence the transformer consumes, plus a minimal flow-matching
|
||||
sampling pipeline used to render preview images during training.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import List, Optional
|
||||
|
||||
import torch
|
||||
from PIL import Image
|
||||
from diffusers.utils.torch_utils import randn_tensor
|
||||
|
||||
from transformers.masking_utils import create_causal_mask
|
||||
|
||||
from .transformer import (
|
||||
IMAGE_POSITION_OFFSET,
|
||||
LLM_TOKEN_INDICATOR,
|
||||
OUTPUT_IMAGE_INDICATOR,
|
||||
QWEN3_VL_ACTIVATION_LAYERS,
|
||||
SEQUENCE_PADDING_INDICATOR,
|
||||
Ideogram4Transformer2DModel,
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Latent (un)patchification.
|
||||
#
|
||||
# The VAE produces (B, ae_ch=32, H/8, W/8) latents. The transformer works on
|
||||
# tokens of dim ae_ch * patch**2 = 128. We store the patchified latent in a 4-D
|
||||
# (B, 128, gh, gw) layout so the rest of ai-toolkit (noise, add_noise, loss) can
|
||||
# treat it like an ordinary image latent. The channel ordering here matches the
|
||||
# reference Ideogram 4 decode exactly: 128 = (patch_h, patch_w, ae_ch) with ae_ch
|
||||
# the fastest-varying axis.
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def patchify_latents(z: torch.Tensor, patch_size: int = 2) -> torch.Tensor:
|
||||
"""(B, ae_ch, H8, W8) -> (B, ae_ch * patch**2, gh, gw)."""
|
||||
b, ae_ch, h8, w8 = z.shape
|
||||
ph = pw = patch_size
|
||||
gh, gw = h8 // ph, w8 // pw
|
||||
z = z.view(b, ae_ch, gh, ph, gw, pw)
|
||||
# -> (B, ph, pw, ae_ch, gh, gw) then merge (ph, pw, ae_ch) -> channels
|
||||
z = z.permute(0, 3, 5, 1, 2, 4).reshape(b, ph * pw * ae_ch, gh, gw)
|
||||
return z
|
||||
|
||||
|
||||
def unpatchify_latents(z: torch.Tensor, patch_size: int = 2) -> torch.Tensor:
|
||||
"""(B, ae_ch * patch**2, gh, gw) -> (B, ae_ch, H8, W8)."""
|
||||
b, c, gh, gw = z.shape
|
||||
ph = pw = patch_size
|
||||
ae_ch = c // (ph * pw)
|
||||
z = z.view(b, ph, pw, ae_ch, gh, gw)
|
||||
# -> (B, ae_ch, gh, ph, gw, pw) then merge spatial
|
||||
z = z.permute(0, 3, 4, 1, 5, 2).reshape(b, ae_ch, gh * ph, gw * pw)
|
||||
return z
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Qwen3-VL hidden-state extraction.
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@torch.no_grad()
|
||||
def get_qwen3_vl_features(
|
||||
text_encoder,
|
||||
token_ids: torch.Tensor,
|
||||
attention_mask: torch.Tensor,
|
||||
pos_2d: torch.Tensor,
|
||||
) -> torch.Tensor:
|
||||
"""Run Qwen3-VL and concat the hidden states from the activation layers.
|
||||
|
||||
Returns a (B, L, hidden_size * num_layers) tensor (in the encoder's dtype),
|
||||
zeroed at non-text (padding) positions.
|
||||
"""
|
||||
language_model = text_encoder.language_model
|
||||
|
||||
inputs_embeds = language_model.embed_tokens(token_ids)
|
||||
|
||||
position_ids_4d = pos_2d[None, ...].expand(4, pos_2d.shape[0], -1)
|
||||
text_position_ids = position_ids_4d[0]
|
||||
mrope_position_ids = position_ids_4d[1:]
|
||||
|
||||
causal_mask = create_causal_mask(
|
||||
config=language_model.config,
|
||||
inputs_embeds=inputs_embeds,
|
||||
attention_mask=attention_mask,
|
||||
past_key_values=None,
|
||||
position_ids=text_position_ids,
|
||||
)
|
||||
position_embeddings = language_model.rotary_emb(inputs_embeds, mrope_position_ids)
|
||||
|
||||
tap_set = set(QWEN3_VL_ACTIVATION_LAYERS)
|
||||
captured: dict[int, torch.Tensor] = {}
|
||||
hidden_states = inputs_embeds
|
||||
for layer_idx, decoder_layer in enumerate(language_model.layers):
|
||||
hidden_states = decoder_layer(
|
||||
hidden_states,
|
||||
attention_mask=causal_mask,
|
||||
position_ids=text_position_ids,
|
||||
past_key_values=None,
|
||||
position_embeddings=position_embeddings,
|
||||
)
|
||||
if layer_idx in tap_set:
|
||||
captured[layer_idx] = hidden_states
|
||||
|
||||
selected = [captured[i] for i in QWEN3_VL_ACTIVATION_LAYERS]
|
||||
batch_size, seq_len = token_ids.shape
|
||||
stacked = torch.stack(selected, dim=0) # (num_taps, B, L, H)
|
||||
stacked = torch.permute(stacked, (1, 2, 3, 0)) # (B, L, H, num_taps)
|
||||
stacked = stacked.reshape(batch_size, seq_len, -1)
|
||||
|
||||
text_mask = attention_mask.to(stacked.dtype).unsqueeze(-1)
|
||||
stacked = stacked * text_mask
|
||||
return stacked
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Packing + velocity prediction.
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def predict_velocity(
|
||||
transformer: Ideogram4Transformer2DModel,
|
||||
latents: torch.Tensor, # (B, 128, gh, gw)
|
||||
t: torch.Tensor, # (B,) toolkit flow time in [0, 1] (1 = pure noise)
|
||||
llm_features: torch.Tensor, # (B, Lt, llm_dim)
|
||||
text_mask: torch.Tensor, # (B, Lt) 1 for real text tokens
|
||||
) -> torch.Tensor:
|
||||
"""Run the transformer on the packed [text | image] sequence.
|
||||
|
||||
``t`` is in the ai-toolkit flow-matching convention: ``t=1`` is pure noise,
|
||||
``t=0`` is clean, and the returned velocity is ``noise - clean`` (matching the
|
||||
toolkit scheduler / loss target).
|
||||
|
||||
Ideogram's transformer uses the opposite convention internally (``t=1`` is
|
||||
clean) and predicts ``clean - noise``, so we feed it ``1 - t`` and negate its
|
||||
output. Returns the velocity reshaped to the (B, 128, gh, gw) latent layout.
|
||||
"""
|
||||
device = latents.device
|
||||
b, c, gh, gw = latents.shape
|
||||
num_image_tokens = gh * gw
|
||||
num_text_tokens = llm_features.shape[1]
|
||||
seq_len = num_text_tokens + num_image_tokens
|
||||
|
||||
# image latents -> tokens (row-major: h outer, w inner)
|
||||
image_tokens = latents.permute(0, 2, 3, 1).reshape(b, num_image_tokens, c)
|
||||
|
||||
# The mask may arrive as a float (PromptEmbeds.to casts it to the embed
|
||||
# dtype); work in long so cumsum positions stay exact for long prompts.
|
||||
text_mask_bool = text_mask.to(device) > 0
|
||||
text_mask_long = text_mask_bool.long()
|
||||
|
||||
# noise tokens: text region is zeroed (masked out anyway)
|
||||
x = torch.cat(
|
||||
[
|
||||
torch.zeros(b, num_text_tokens, c, device=device, dtype=image_tokens.dtype),
|
||||
image_tokens,
|
||||
],
|
||||
dim=1,
|
||||
)
|
||||
|
||||
# llm features: image region is zero
|
||||
llm_full = torch.cat(
|
||||
[
|
||||
llm_features,
|
||||
torch.zeros(
|
||||
b,
|
||||
num_image_tokens,
|
||||
llm_features.shape[-1],
|
||||
device=device,
|
||||
dtype=llm_features.dtype,
|
||||
),
|
||||
],
|
||||
dim=1,
|
||||
)
|
||||
|
||||
# indicator: real text -> 3, image -> 2, text pad -> 0
|
||||
indicator = torch.zeros(b, seq_len, dtype=torch.long, device=device)
|
||||
indicator[:, :num_text_tokens] = text_mask_long * LLM_TOKEN_INDICATOR
|
||||
indicator[:, num_text_tokens:] = OUTPUT_IMAGE_INDICATOR
|
||||
|
||||
# segment ids: real text + image -> 1, text pad -> -1 (its own padding segment)
|
||||
segment_ids = torch.ones(b, seq_len, dtype=torch.long, device=device)
|
||||
segment_ids[:, :num_text_tokens] = torch.where(
|
||||
text_mask_bool,
|
||||
torch.ones_like(text_mask_long),
|
||||
torch.full_like(text_mask_long, SEQUENCE_PADDING_INDICATOR),
|
||||
)
|
||||
|
||||
# position ids (t, h, w)
|
||||
# text positions: 0..num_real-1 at the real slots (relative; pad -> 0)
|
||||
text_pos = (text_mask_long.cumsum(dim=-1) - 1).clamp(min=0) # (B, Lt)
|
||||
text_pos_3d = text_pos.unsqueeze(-1).expand(-1, -1, 3)
|
||||
|
||||
h_idx = torch.arange(gh, device=device).view(-1, 1).expand(gh, gw).reshape(-1)
|
||||
w_idx = torch.arange(gw, device=device).view(1, -1).expand(gh, gw).reshape(-1)
|
||||
t_idx = torch.zeros_like(h_idx)
|
||||
image_pos = torch.stack([t_idx, h_idx, w_idx], dim=1) + IMAGE_POSITION_OFFSET
|
||||
image_pos_3d = image_pos.unsqueeze(0).expand(b, -1, -1)
|
||||
|
||||
position_ids = torch.cat([text_pos_3d, image_pos_3d], dim=1)
|
||||
|
||||
# Flip into the model's time convention (t=1 -> clean).
|
||||
model_t = 1.0 - t
|
||||
|
||||
out = transformer(
|
||||
llm_features=llm_full,
|
||||
x=x,
|
||||
t=model_t,
|
||||
position_ids=position_ids,
|
||||
segment_ids=segment_ids,
|
||||
indicator=indicator,
|
||||
)
|
||||
|
||||
image_velocity = out[:, num_text_tokens:] # (B, Li, 128)
|
||||
image_velocity = image_velocity.reshape(b, gh, gw, c).permute(0, 3, 1, 2)
|
||||
# Model predicts clean - noise; negate to return toolkit velocity (noise - clean).
|
||||
return -image_velocity
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Minimal sampling pipeline (for training previews).
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class Ideogram4Pipeline:
|
||||
"""Lightweight flow-matching sampler used by ai-toolkit's preview generation."""
|
||||
|
||||
def __init__(self, model):
|
||||
# ``model`` is the Ideogram4Model so we can reuse its encode/decode and
|
||||
# latent helpers without duplicating state.
|
||||
self.model = model
|
||||
|
||||
@property
|
||||
def device(self):
|
||||
return self.model.device_torch
|
||||
|
||||
def to(self, *args, **kwargs):
|
||||
return self
|
||||
|
||||
@torch.no_grad()
|
||||
def __call__(
|
||||
self,
|
||||
conditional_embeds,
|
||||
unconditional_embeds,
|
||||
height: int = 1024,
|
||||
width: int = 1024,
|
||||
num_inference_steps: int = 30,
|
||||
guidance_scale: float = 7.0,
|
||||
latents: Optional[torch.Tensor] = None,
|
||||
generator: Optional[torch.Generator] = None,
|
||||
**kwargs,
|
||||
) -> List[Image.Image]:
|
||||
model = self.model
|
||||
device = model.device_torch
|
||||
dtype = model.torch_dtype
|
||||
transformer = model.transformer
|
||||
patch = model.patch_size
|
||||
|
||||
# Use a fresh scheduler so we never mutate the training scheduler's state.
|
||||
scheduler = model.get_train_scheduler()
|
||||
scheduler.set_timesteps(num_inference_steps, device=device)
|
||||
timesteps = scheduler.timesteps
|
||||
|
||||
ae_scale = model.vae_scale_factor # 8
|
||||
gh = height // (ae_scale * patch)
|
||||
gw = width // (ae_scale * patch)
|
||||
latent_channels = transformer.config.in_channels
|
||||
|
||||
do_cfg = unconditional_embeds is not None and guidance_scale != 1.0
|
||||
|
||||
if latents is None:
|
||||
shape = (1, latent_channels, gh, gw)
|
||||
latents = randn_tensor(
|
||||
shape, generator=generator, device=device, dtype=torch.float32
|
||||
)
|
||||
latents = latents.to(device, dtype=torch.float32)
|
||||
latents = latents * scheduler.init_noise_sigma
|
||||
|
||||
cond_feats = conditional_embeds.text_embeds.to(device, dtype=dtype)
|
||||
cond_mask = conditional_embeds.attention_mask.to(device)
|
||||
if do_cfg:
|
||||
uncond_feats = unconditional_embeds.text_embeds.to(device, dtype=dtype)
|
||||
uncond_mask = unconditional_embeds.attention_mask.to(device)
|
||||
|
||||
for t in timesteps:
|
||||
t01 = (t / 1000.0).to(device).expand(latents.shape[0])
|
||||
v_cond = predict_velocity(
|
||||
transformer, latents.to(dtype), t01, cond_feats, cond_mask
|
||||
)
|
||||
if do_cfg:
|
||||
v_uncond = predict_velocity(
|
||||
transformer, latents.to(dtype), t01, uncond_feats, uncond_mask
|
||||
)
|
||||
v = v_uncond + guidance_scale * (v_cond - v_uncond)
|
||||
else:
|
||||
v = v_cond
|
||||
latents = scheduler.step(
|
||||
v.to(torch.float32), t, latents, return_dict=False
|
||||
)[0]
|
||||
|
||||
images = model.decode_latents(latents, device=device, dtype=dtype)
|
||||
images = images.float().clamp(-1.0, 1.0)
|
||||
images = ((images + 1.0) * 127.5).round().to(torch.uint8)
|
||||
images = images.permute(0, 2, 3, 1).cpu().numpy()
|
||||
return [Image.fromarray(arr) for arr in images]
|
||||
@@ -0,0 +1,406 @@
|
||||
"""Ideogram4 transformer backbone.
|
||||
|
||||
Ported from the reference ``modeling_ideogram4.py`` for ai-toolkit. The
|
||||
transformer consumes Qwen3-VL hidden states and flow-matching noise tokens
|
||||
(packed into a single sequence) and produces velocity predictions on the image
|
||||
latent tokens.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import math
|
||||
from dataclasses import dataclass
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
from torch.utils.checkpoint import checkpoint
|
||||
|
||||
# Per-token role indicators used inside the packed sequence.
|
||||
SEQUENCE_PADDING_INDICATOR = -1
|
||||
OUTPUT_IMAGE_INDICATOR = 2
|
||||
LLM_TOKEN_INDICATOR = 3
|
||||
|
||||
# Image grid coordinates start at this offset so they never collide with text
|
||||
# token indices (text positions start at 0 and never exceed max_text_tokens).
|
||||
IMAGE_POSITION_OFFSET = 65536
|
||||
|
||||
# Layers of Qwen3-VL whose hidden states are concatenated and fed to the transformer.
|
||||
QWEN3_VL_ACTIVATION_LAYERS = (0, 3, 6, 9, 12, 15, 18, 21, 24, 27, 30, 33, 35)
|
||||
|
||||
|
||||
@dataclass
|
||||
class Ideogram4Config:
|
||||
emb_dim: int = 4608
|
||||
num_layers: int = 34
|
||||
num_heads: int = 18
|
||||
intermediate_size: int = 12288
|
||||
adanln_dim: int = 512
|
||||
|
||||
# Latent dimension after patchification: ae_channels (32) * patch_size**2 (4) = 128.
|
||||
in_channels: int = 128
|
||||
|
||||
# Hidden size of Qwen3-VL-8B-Instruct multiplied by the number of extracted layers.
|
||||
# Qwen3-VL hidden size = 4096
|
||||
llm_features_dim: int = 4096 * len(QWEN3_VL_ACTIVATION_LAYERS)
|
||||
|
||||
rope_theta: int = 5_000_000
|
||||
mrope_section: tuple[int, ...] = (24, 20, 20)
|
||||
|
||||
norm_eps: float = 1e-5
|
||||
|
||||
|
||||
def _rotate_half(x: torch.Tensor) -> torch.Tensor:
|
||||
half = x.shape[-1] // 2
|
||||
x1 = x[..., :half]
|
||||
x2 = x[..., half:]
|
||||
return torch.cat((-x2, x1), dim=-1)
|
||||
|
||||
|
||||
def _apply_rotary_pos_emb(
|
||||
q: torch.Tensor,
|
||||
k: torch.Tensor,
|
||||
cos: torch.Tensor,
|
||||
sin: torch.Tensor,
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
# q, k: (B, num_heads, L, head_dim); cos/sin: (B, L, head_dim).
|
||||
cos = cos.unsqueeze(1)
|
||||
sin = sin.unsqueeze(1)
|
||||
q_embed = (q * cos) + (_rotate_half(q) * sin)
|
||||
k_embed = (k * cos) + (_rotate_half(k) * sin)
|
||||
return q_embed, k_embed
|
||||
|
||||
|
||||
class Ideogram4MRoPE(nn.Module):
|
||||
inv_freq: torch.Tensor
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
head_dim: int,
|
||||
base: int,
|
||||
mrope_section: tuple[int, ...],
|
||||
) -> None:
|
||||
super().__init__()
|
||||
inv_freq = 1.0 / (
|
||||
base ** (torch.arange(0, head_dim, 2, dtype=torch.float32) / head_dim)
|
||||
)
|
||||
self.register_buffer("inv_freq", inv_freq, persistent=False)
|
||||
self.mrope_section = tuple(mrope_section)
|
||||
self.head_dim = head_dim
|
||||
|
||||
@torch.no_grad()
|
||||
def forward(self, position_ids: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
# position_ids: (B, L, 3) of int.
|
||||
assert position_ids.ndim == 3 and position_ids.shape[-1] == 3
|
||||
batch_size, seq_len, _ = position_ids.shape
|
||||
|
||||
# (3, B, inv_freq_size, L)
|
||||
pos = position_ids.permute(2, 0, 1).to(dtype=torch.float32)
|
||||
inv_freq = self.inv_freq.to(dtype=torch.float32)[None, None, :, None].expand(
|
||||
3, batch_size, -1, 1
|
||||
)
|
||||
freqs = inv_freq @ pos.unsqueeze(2)
|
||||
freqs = freqs.transpose(2, 3) # (3, B, L, inv_freq_size)
|
||||
|
||||
# interleaved mrope: pull H freqs into idx 1 mod 3, W freqs into idx 2 mod 3.
|
||||
freqs_t = freqs[0].clone()
|
||||
for axis, offset in ((1, 1), (2, 2)):
|
||||
length = self.mrope_section[axis] * 3
|
||||
idx = torch.arange(offset, length, 3, device=freqs_t.device)
|
||||
freqs_t[..., idx] = freqs[axis][..., idx]
|
||||
|
||||
emb = torch.cat((freqs_t, freqs_t), dim=-1)
|
||||
return emb.cos(), emb.sin()
|
||||
|
||||
|
||||
class Ideogram4RMSNorm(nn.Module):
|
||||
def __init__(self, dim: int, eps: float = 1e-6) -> None:
|
||||
super().__init__()
|
||||
self.weight = nn.Parameter(torch.ones(dim))
|
||||
self.eps = eps
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
return F.rms_norm(x, self.weight.shape, self.weight, self.eps)
|
||||
|
||||
|
||||
class Ideogram4Attention(nn.Module):
|
||||
def __init__(self, hidden_size: int, num_heads: int, eps: float = 1e-5) -> None:
|
||||
super().__init__()
|
||||
assert hidden_size % num_heads == 0
|
||||
self.hidden_size = hidden_size
|
||||
self.num_heads = num_heads
|
||||
self.head_dim = hidden_size // num_heads
|
||||
|
||||
self.qkv = nn.Linear(hidden_size, hidden_size * 3, bias=False)
|
||||
self.norm_q = Ideogram4RMSNorm(self.head_dim, eps=eps)
|
||||
self.norm_k = Ideogram4RMSNorm(self.head_dim, eps=eps)
|
||||
self.o = nn.Linear(hidden_size, hidden_size, bias=False)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
x: torch.Tensor,
|
||||
attn_mask: torch.Tensor,
|
||||
cos: torch.Tensor,
|
||||
sin: torch.Tensor,
|
||||
) -> torch.Tensor:
|
||||
batch_size, seq_len, _ = x.shape
|
||||
|
||||
qkv = self.qkv(x)
|
||||
qkv = qkv.view(batch_size, seq_len, 3, self.num_heads, self.head_dim)
|
||||
q, k, v = qkv.unbind(dim=2)
|
||||
|
||||
q = self.norm_q(q)
|
||||
k = self.norm_k(k)
|
||||
|
||||
# SDPA expects (B, num_heads, L, head_dim).
|
||||
q = q.transpose(1, 2)
|
||||
k = k.transpose(1, 2)
|
||||
v = v.transpose(1, 2)
|
||||
|
||||
q, k = _apply_rotary_pos_emb(q, k, cos, sin)
|
||||
|
||||
out = F.scaled_dot_product_attention(q, k, v, attn_mask=attn_mask)
|
||||
out = out.transpose(1, 2).reshape(batch_size, seq_len, self.hidden_size)
|
||||
return self.o(out)
|
||||
|
||||
|
||||
class Ideogram4MLP(nn.Module):
|
||||
def __init__(self, dim: int, hidden_dim: int) -> None:
|
||||
super().__init__()
|
||||
self.w1 = nn.Linear(dim, hidden_dim, bias=False)
|
||||
self.w2 = nn.Linear(hidden_dim, dim, bias=False)
|
||||
self.w3 = nn.Linear(dim, hidden_dim, bias=False)
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
return self.w2(F.silu(self.w1(x)) * self.w3(x))
|
||||
|
||||
|
||||
class Ideogram4TransformerBlock(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
hidden_size: int,
|
||||
intermediate_size: int,
|
||||
num_heads: int,
|
||||
norm_eps: float,
|
||||
adanln_dim: int,
|
||||
) -> None:
|
||||
super().__init__()
|
||||
self.attention = Ideogram4Attention(hidden_size, num_heads, eps=1e-5)
|
||||
self.feed_forward = Ideogram4MLP(hidden_size, intermediate_size)
|
||||
|
||||
self.attention_norm1 = Ideogram4RMSNorm(hidden_size, eps=norm_eps)
|
||||
self.ffn_norm1 = Ideogram4RMSNorm(hidden_size, eps=norm_eps)
|
||||
self.attention_norm2 = Ideogram4RMSNorm(hidden_size, eps=norm_eps)
|
||||
self.ffn_norm2 = Ideogram4RMSNorm(hidden_size, eps=norm_eps)
|
||||
|
||||
self.adaln_modulation = nn.Linear(adanln_dim, 4 * hidden_size, bias=True)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
x: torch.Tensor,
|
||||
attn_mask: torch.Tensor,
|
||||
cos: torch.Tensor,
|
||||
sin: torch.Tensor,
|
||||
adaln_input: torch.Tensor,
|
||||
) -> torch.Tensor:
|
||||
mod = self.adaln_modulation(adaln_input)
|
||||
scale_msa, gate_msa, scale_mlp, gate_mlp = mod.chunk(4, dim=-1)
|
||||
gate_msa = torch.tanh(gate_msa)
|
||||
gate_mlp = torch.tanh(gate_mlp)
|
||||
scale_msa = 1.0 + scale_msa
|
||||
scale_mlp = 1.0 + scale_mlp
|
||||
|
||||
attn_out = self.attention(
|
||||
self.attention_norm1(x) * scale_msa,
|
||||
attn_mask=attn_mask,
|
||||
cos=cos,
|
||||
sin=sin,
|
||||
)
|
||||
x = x + gate_msa * self.attention_norm2(attn_out)
|
||||
x = x + gate_mlp * self.ffn_norm2(
|
||||
self.feed_forward(self.ffn_norm1(x) * scale_mlp)
|
||||
)
|
||||
return x
|
||||
|
||||
|
||||
def _sinusoidal_embedding(
|
||||
t: torch.Tensor, dim: int, scale: float = 1e4
|
||||
) -> torch.Tensor:
|
||||
t = t.to(torch.float32)
|
||||
half = dim // 2
|
||||
freq = math.log(scale) / (half - 1)
|
||||
freq = torch.exp(torch.arange(half, dtype=torch.float32, device=t.device) * -freq)
|
||||
emb = t.unsqueeze(-1) * freq
|
||||
emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=-1)
|
||||
if dim % 2 == 1:
|
||||
emb = F.pad(emb, (0, 1))
|
||||
return emb
|
||||
|
||||
|
||||
class Ideogram4EmbedScalar(nn.Module):
|
||||
def __init__(self, dim: int, input_range: tuple[float, float]) -> None:
|
||||
super().__init__()
|
||||
self.dim = dim
|
||||
self.range_min, self.range_max = input_range
|
||||
assert self.range_max > self.range_min
|
||||
self.mlp_in = nn.Linear(dim, dim, bias=True)
|
||||
self.mlp_out = nn.Linear(dim, dim, bias=True)
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
# x is shape (..., 1) or (...,) holding a scalar per token.
|
||||
x = x.to(torch.float32)
|
||||
scaled = 1e4 * (x - self.range_min) / (self.range_max - self.range_min)
|
||||
emb = _sinusoidal_embedding(scaled, self.dim)
|
||||
emb = emb.to(self.mlp_in.weight.dtype)
|
||||
emb = F.silu(self.mlp_in(emb))
|
||||
return self.mlp_out(emb)
|
||||
|
||||
|
||||
class Ideogram4FinalLayer(nn.Module):
|
||||
def __init__(self, hidden_size: int, out_channels: int, adanln_dim: int) -> None:
|
||||
super().__init__()
|
||||
self.norm_final = nn.LayerNorm(hidden_size, eps=1e-6, elementwise_affine=False)
|
||||
self.linear = nn.Linear(hidden_size, out_channels, bias=True)
|
||||
self.adaln_modulation = nn.Linear(adanln_dim, hidden_size, bias=True)
|
||||
|
||||
def forward(self, x: torch.Tensor, c: torch.Tensor) -> torch.Tensor:
|
||||
scale = 1.0 + self.adaln_modulation(F.silu(c))
|
||||
return self.linear(self.norm_final(x) * scale)
|
||||
|
||||
|
||||
class Ideogram4Transformer2DModel(nn.Module):
|
||||
"""Ideogram 4 flow-matching transformer."""
|
||||
|
||||
def __init__(self, config: Ideogram4Config) -> None:
|
||||
super().__init__()
|
||||
self.config = config
|
||||
self.gradient_checkpointing = False
|
||||
|
||||
head_dim = config.emb_dim // config.num_heads
|
||||
|
||||
self.input_proj = nn.Linear(config.in_channels, config.emb_dim, bias=True)
|
||||
self.llm_cond_norm = Ideogram4RMSNorm(config.llm_features_dim, eps=1e-6)
|
||||
self.llm_cond_proj = nn.Linear(
|
||||
config.llm_features_dim, config.emb_dim, bias=True
|
||||
)
|
||||
self.t_embedding = Ideogram4EmbedScalar(config.emb_dim, input_range=(0.0, 1.0))
|
||||
self.adaln_proj = nn.Linear(config.emb_dim, config.adanln_dim, bias=True)
|
||||
|
||||
self.embed_image_indicator = nn.Embedding(2, config.emb_dim)
|
||||
|
||||
self.rotary_emb = Ideogram4MRoPE(
|
||||
head_dim=head_dim,
|
||||
base=config.rope_theta,
|
||||
mrope_section=config.mrope_section,
|
||||
)
|
||||
|
||||
self.layers = nn.ModuleList(
|
||||
[
|
||||
Ideogram4TransformerBlock(
|
||||
hidden_size=config.emb_dim,
|
||||
intermediate_size=config.intermediate_size,
|
||||
num_heads=config.num_heads,
|
||||
norm_eps=config.norm_eps,
|
||||
adanln_dim=config.adanln_dim,
|
||||
)
|
||||
for _ in range(config.num_layers)
|
||||
]
|
||||
)
|
||||
|
||||
self.final_layer = Ideogram4FinalLayer(
|
||||
hidden_size=config.emb_dim,
|
||||
out_channels=config.in_channels,
|
||||
adanln_dim=config.adanln_dim,
|
||||
)
|
||||
|
||||
@property
|
||||
def device(self) -> torch.device:
|
||||
return next(self.parameters()).device
|
||||
|
||||
@property
|
||||
def dtype(self) -> torch.dtype:
|
||||
return next(self.parameters()).dtype
|
||||
|
||||
def enable_gradient_checkpointing(self):
|
||||
self.gradient_checkpointing = True
|
||||
|
||||
def disable_gradient_checkpointing(self):
|
||||
self.gradient_checkpointing = False
|
||||
|
||||
def forward(
|
||||
self,
|
||||
*,
|
||||
llm_features: torch.Tensor,
|
||||
x: torch.Tensor,
|
||||
t: torch.Tensor,
|
||||
position_ids: torch.Tensor,
|
||||
segment_ids: torch.Tensor,
|
||||
indicator: torch.Tensor,
|
||||
) -> torch.Tensor:
|
||||
"""Velocity prediction.
|
||||
|
||||
Args:
|
||||
llm_features: (B, L, llm_features_dim) Qwen3-VL conditioning features.
|
||||
x: (B, L, in_channels) noise tokens.
|
||||
t: (B,) or (B, L) flow-matching time in [0, 1].
|
||||
position_ids: (B, L, 3) (t, h, w) positions for MRoPE.
|
||||
segment_ids: (B, L) sample id within a packed batch.
|
||||
indicator: (B, L) per-token role: LLM_TOKEN_INDICATOR or OUTPUT_IMAGE_INDICATOR.
|
||||
|
||||
Returns:
|
||||
(B, L, in_channels) velocity prediction in float32. Only the positions
|
||||
with ``indicator == OUTPUT_IMAGE_INDICATOR`` are meaningful.
|
||||
"""
|
||||
batch_size, seq_len, in_channels = x.shape
|
||||
assert in_channels == self.config.in_channels
|
||||
|
||||
param_dtype = self.input_proj.weight.dtype
|
||||
x = x.to(param_dtype)
|
||||
t = t.to(param_dtype)
|
||||
llm_features = llm_features.to(param_dtype)
|
||||
|
||||
indicator = indicator.to(torch.long)
|
||||
llm_token_mask = (indicator == LLM_TOKEN_INDICATOR).to(x.dtype).unsqueeze(-1)
|
||||
output_image_mask = (
|
||||
(indicator == OUTPUT_IMAGE_INDICATOR).to(x.dtype).unsqueeze(-1)
|
||||
)
|
||||
|
||||
llm_features = llm_features * llm_token_mask
|
||||
x = x * output_image_mask
|
||||
|
||||
x = self.input_proj(x) * output_image_mask
|
||||
|
||||
# Keep shape (B, 1, ...) when t is per-sample so downstream adaln_modulation
|
||||
# projections don't pay for L identical copies.
|
||||
t_cond = self.t_embedding(t)
|
||||
if t.dim() == 1:
|
||||
t_cond = t_cond.unsqueeze(1)
|
||||
adaln_input = F.silu(self.adaln_proj(t_cond))
|
||||
|
||||
llm_features = self.llm_cond_norm(llm_features)
|
||||
llm_features = self.llm_cond_proj(llm_features) * llm_token_mask
|
||||
|
||||
h = x + llm_features
|
||||
|
||||
image_indicator_embedding = self.embed_image_indicator(
|
||||
(indicator == OUTPUT_IMAGE_INDICATOR).to(torch.long)
|
||||
)
|
||||
h = h + image_indicator_embedding
|
||||
|
||||
cos, sin = self.rotary_emb(position_ids)
|
||||
cos = cos.to(h.dtype)
|
||||
sin = sin.to(h.dtype)
|
||||
|
||||
# Block-diagonal mask from segment ids: (B, 1, L, L), True = attend.
|
||||
attn_mask = (segment_ids.unsqueeze(2) == segment_ids.unsqueeze(1)).unsqueeze(1)
|
||||
|
||||
for layer in self.layers:
|
||||
if self.gradient_checkpointing and torch.is_grad_enabled():
|
||||
h = checkpoint(
|
||||
layer, h, attn_mask, cos, sin, adaln_input, use_reentrant=False
|
||||
)
|
||||
else:
|
||||
h = layer(h, attn_mask, cos, sin, adaln_input)
|
||||
|
||||
out = self.final_layer(h, c=adaln_input)
|
||||
return out.to(torch.float32)
|
||||
431
extensions_built_in/diffusion_models/ideogram4/src/vae.py
Normal file
431
extensions_built_in/diffusion_models/ideogram4/src/vae.py
Normal file
@@ -0,0 +1,431 @@
|
||||
"""Flux2 KL autoencoder."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import math
|
||||
import re
|
||||
from dataclasses import dataclass, field
|
||||
|
||||
import torch
|
||||
from einops import rearrange
|
||||
from torch import Tensor, nn
|
||||
|
||||
|
||||
@dataclass
|
||||
class AutoEncoderParams:
|
||||
resolution: int = 256
|
||||
in_channels: int = 3
|
||||
ch: int = 128
|
||||
out_ch: int = 3
|
||||
ch_mult: list[int] = field(default_factory=lambda: [1, 2, 4, 4])
|
||||
num_res_blocks: int = 2
|
||||
z_channels: int = 32
|
||||
|
||||
|
||||
def swish(x: Tensor) -> Tensor:
|
||||
return x * torch.sigmoid(x)
|
||||
|
||||
|
||||
class AttnBlock(nn.Module):
|
||||
def __init__(self, in_channels: int):
|
||||
super().__init__()
|
||||
self.in_channels = in_channels
|
||||
|
||||
self.norm = nn.GroupNorm(
|
||||
num_groups=32, num_channels=in_channels, eps=1e-6, affine=True
|
||||
)
|
||||
|
||||
self.q = nn.Conv2d(in_channels, in_channels, kernel_size=1)
|
||||
self.k = nn.Conv2d(in_channels, in_channels, kernel_size=1)
|
||||
self.v = nn.Conv2d(in_channels, in_channels, kernel_size=1)
|
||||
self.proj_out = nn.Conv2d(in_channels, in_channels, kernel_size=1)
|
||||
|
||||
def attention(self, h_: Tensor) -> Tensor:
|
||||
h_ = self.norm(h_)
|
||||
q = self.q(h_)
|
||||
k = self.k(h_)
|
||||
v = self.v(h_)
|
||||
|
||||
b, c, h, w = q.shape
|
||||
q = rearrange(q, "b c h w -> b 1 (h w) c").contiguous()
|
||||
k = rearrange(k, "b c h w -> b 1 (h w) c").contiguous()
|
||||
v = rearrange(v, "b c h w -> b 1 (h w) c").contiguous()
|
||||
h_ = nn.functional.scaled_dot_product_attention(q, k, v)
|
||||
|
||||
return rearrange(h_, "b 1 (h w) c -> b c h w", h=h, w=w, c=c, b=b)
|
||||
|
||||
def forward(self, x: Tensor) -> Tensor:
|
||||
return x + self.proj_out(self.attention(x))
|
||||
|
||||
|
||||
class ResnetBlock(nn.Module):
|
||||
def __init__(self, in_channels: int, out_channels: int):
|
||||
super().__init__()
|
||||
self.in_channels = in_channels
|
||||
out_channels = in_channels if out_channels is None else out_channels
|
||||
self.out_channels = out_channels
|
||||
|
||||
self.norm1 = nn.GroupNorm(
|
||||
num_groups=32, num_channels=in_channels, eps=1e-6, affine=True
|
||||
)
|
||||
self.conv1 = nn.Conv2d(
|
||||
in_channels, out_channels, kernel_size=3, stride=1, padding=1
|
||||
)
|
||||
self.norm2 = nn.GroupNorm(
|
||||
num_groups=32, num_channels=out_channels, eps=1e-6, affine=True
|
||||
)
|
||||
self.conv2 = nn.Conv2d(
|
||||
out_channels, out_channels, kernel_size=3, stride=1, padding=1
|
||||
)
|
||||
if self.in_channels != self.out_channels:
|
||||
self.nin_shortcut = nn.Conv2d(
|
||||
in_channels, out_channels, kernel_size=1, stride=1, padding=0
|
||||
)
|
||||
|
||||
def forward(self, x):
|
||||
h = x
|
||||
h = self.norm1(h)
|
||||
h = swish(h)
|
||||
h = self.conv1(h)
|
||||
|
||||
h = self.norm2(h)
|
||||
h = swish(h)
|
||||
h = self.conv2(h)
|
||||
|
||||
if self.in_channels != self.out_channels:
|
||||
x = self.nin_shortcut(x)
|
||||
|
||||
return x + h
|
||||
|
||||
|
||||
class Downsample(nn.Module):
|
||||
def __init__(self, in_channels: int):
|
||||
super().__init__()
|
||||
# no asymmetric padding in torch conv, must do it ourselves
|
||||
self.conv = nn.Conv2d(
|
||||
in_channels, in_channels, kernel_size=3, stride=2, padding=0
|
||||
)
|
||||
|
||||
def forward(self, x: Tensor):
|
||||
pad = (0, 1, 0, 1)
|
||||
x = nn.functional.pad(x, pad, mode="constant", value=0)
|
||||
x = self.conv(x)
|
||||
return x
|
||||
|
||||
|
||||
class Upsample(nn.Module):
|
||||
def __init__(self, in_channels: int):
|
||||
super().__init__()
|
||||
self.conv = nn.Conv2d(
|
||||
in_channels, in_channels, kernel_size=3, stride=1, padding=1
|
||||
)
|
||||
|
||||
def forward(self, x: Tensor):
|
||||
x = nn.functional.interpolate(x, scale_factor=2.0, mode="nearest")
|
||||
x = self.conv(x)
|
||||
return x
|
||||
|
||||
|
||||
class Encoder(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
resolution: int,
|
||||
in_channels: int,
|
||||
ch: int,
|
||||
ch_mult: list[int],
|
||||
num_res_blocks: int,
|
||||
z_channels: int,
|
||||
):
|
||||
super().__init__()
|
||||
self.quant_conv = torch.nn.Conv2d(2 * z_channels, 2 * z_channels, 1)
|
||||
self.ch = ch
|
||||
self.num_resolutions = len(ch_mult)
|
||||
self.num_res_blocks = num_res_blocks
|
||||
self.resolution = resolution
|
||||
self.in_channels = in_channels
|
||||
# downsampling
|
||||
self.conv_in = nn.Conv2d(
|
||||
in_channels, self.ch, kernel_size=3, stride=1, padding=1
|
||||
)
|
||||
|
||||
curr_res = resolution
|
||||
in_ch_mult = (1,) + tuple(ch_mult)
|
||||
self.in_ch_mult = in_ch_mult
|
||||
self.down = nn.ModuleList()
|
||||
block_in = self.ch
|
||||
for i_level in range(self.num_resolutions):
|
||||
block = nn.ModuleList()
|
||||
attn = nn.ModuleList()
|
||||
block_in = ch * in_ch_mult[i_level]
|
||||
block_out = ch * ch_mult[i_level]
|
||||
for _ in range(self.num_res_blocks):
|
||||
block.append(ResnetBlock(in_channels=block_in, out_channels=block_out))
|
||||
block_in = block_out
|
||||
down = nn.Module()
|
||||
down.block = block
|
||||
down.attn = attn
|
||||
if i_level != self.num_resolutions - 1:
|
||||
down.downsample = Downsample(block_in)
|
||||
curr_res = curr_res // 2
|
||||
self.down.append(down)
|
||||
|
||||
# middle
|
||||
self.mid = nn.Module()
|
||||
self.mid.block_1 = ResnetBlock(in_channels=block_in, out_channels=block_in)
|
||||
self.mid.attn_1 = AttnBlock(block_in)
|
||||
self.mid.block_2 = ResnetBlock(in_channels=block_in, out_channels=block_in)
|
||||
|
||||
# end
|
||||
self.norm_out = nn.GroupNorm(
|
||||
num_groups=32, num_channels=block_in, eps=1e-6, affine=True
|
||||
)
|
||||
self.conv_out = nn.Conv2d(
|
||||
block_in, 2 * z_channels, kernel_size=3, stride=1, padding=1
|
||||
)
|
||||
|
||||
def forward(self, x: Tensor) -> Tensor:
|
||||
# downsampling
|
||||
hs = [self.conv_in(x)]
|
||||
for i_level in range(self.num_resolutions):
|
||||
for i_block in range(self.num_res_blocks):
|
||||
h = self.down[i_level].block[i_block](hs[-1]) # type: ignore[index, operator]
|
||||
if len(self.down[i_level].attn) > 0: # type: ignore[arg-type]
|
||||
h = self.down[i_level].attn[i_block](h) # type: ignore[index, operator]
|
||||
hs.append(h)
|
||||
if i_level != self.num_resolutions - 1:
|
||||
hs.append(self.down[i_level].downsample(hs[-1])) # type: ignore[operator]
|
||||
|
||||
# middle
|
||||
h = hs[-1]
|
||||
h = self.mid.block_1(h) # type: ignore[operator]
|
||||
h = self.mid.attn_1(h) # type: ignore[operator]
|
||||
h = self.mid.block_2(h) # type: ignore[operator]
|
||||
# end
|
||||
h = self.norm_out(h)
|
||||
h = swish(h)
|
||||
h = self.conv_out(h)
|
||||
h = self.quant_conv(h)
|
||||
return h
|
||||
|
||||
|
||||
class Decoder(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
ch: int,
|
||||
out_ch: int,
|
||||
ch_mult: list[int],
|
||||
num_res_blocks: int,
|
||||
in_channels: int,
|
||||
resolution: int,
|
||||
z_channels: int,
|
||||
):
|
||||
super().__init__()
|
||||
self.post_quant_conv = torch.nn.Conv2d(z_channels, z_channels, 1)
|
||||
self.ch = ch
|
||||
self.num_resolutions = len(ch_mult)
|
||||
self.num_res_blocks = num_res_blocks
|
||||
self.resolution = resolution
|
||||
self.in_channels = in_channels
|
||||
self.ffactor = 2 ** (self.num_resolutions - 1)
|
||||
|
||||
# compute in_ch_mult, block_in and curr_res at lowest res
|
||||
block_in = ch * ch_mult[self.num_resolutions - 1]
|
||||
curr_res = resolution // 2 ** (self.num_resolutions - 1)
|
||||
self.z_shape = (1, z_channels, curr_res, curr_res)
|
||||
|
||||
# z to block_in
|
||||
self.conv_in = nn.Conv2d(
|
||||
z_channels, block_in, kernel_size=3, stride=1, padding=1
|
||||
)
|
||||
|
||||
# middle
|
||||
self.mid = nn.Module()
|
||||
self.mid.block_1 = ResnetBlock(in_channels=block_in, out_channels=block_in)
|
||||
self.mid.attn_1 = AttnBlock(block_in)
|
||||
self.mid.block_2 = ResnetBlock(in_channels=block_in, out_channels=block_in)
|
||||
|
||||
# upsampling
|
||||
self.up = nn.ModuleList()
|
||||
for i_level in reversed(range(self.num_resolutions)):
|
||||
block = nn.ModuleList()
|
||||
attn = nn.ModuleList()
|
||||
block_out = ch * ch_mult[i_level]
|
||||
for _ in range(self.num_res_blocks + 1):
|
||||
block.append(ResnetBlock(in_channels=block_in, out_channels=block_out))
|
||||
block_in = block_out
|
||||
up = nn.Module()
|
||||
up.block = block
|
||||
up.attn = attn
|
||||
if i_level != 0:
|
||||
up.upsample = Upsample(block_in)
|
||||
curr_res = curr_res * 2
|
||||
self.up.insert(0, up) # prepend to get consistent order
|
||||
|
||||
# end
|
||||
self.norm_out = nn.GroupNorm(
|
||||
num_groups=32, num_channels=block_in, eps=1e-6, affine=True
|
||||
)
|
||||
self.conv_out = nn.Conv2d(block_in, out_ch, kernel_size=3, stride=1, padding=1)
|
||||
|
||||
def forward(self, z: Tensor) -> Tensor:
|
||||
z = self.post_quant_conv(z)
|
||||
|
||||
# get dtype for proper tracing
|
||||
upscale_dtype = next(self.up.parameters()).dtype
|
||||
|
||||
# z to block_in
|
||||
h = self.conv_in(z)
|
||||
|
||||
# middle
|
||||
h = self.mid.block_1(h) # type: ignore[operator]
|
||||
h = self.mid.attn_1(h) # type: ignore[operator]
|
||||
h = self.mid.block_2(h) # type: ignore[operator]
|
||||
|
||||
# cast to proper dtype
|
||||
h = h.to(upscale_dtype)
|
||||
# upsampling
|
||||
for i_level in reversed(range(self.num_resolutions)):
|
||||
for i_block in range(self.num_res_blocks + 1):
|
||||
h = self.up[i_level].block[i_block](h) # type: ignore[index, operator]
|
||||
if len(self.up[i_level].attn) > 0: # type: ignore[arg-type]
|
||||
h = self.up[i_level].attn[i_block](h) # type: ignore[index, operator]
|
||||
if i_level != 0:
|
||||
h = self.up[i_level].upsample(h) # type: ignore[operator]
|
||||
|
||||
# end
|
||||
h = self.norm_out(h)
|
||||
h = swish(h)
|
||||
h = self.conv_out(h)
|
||||
return h
|
||||
|
||||
|
||||
class AutoEncoder(nn.Module):
|
||||
def __init__(self, params: AutoEncoderParams):
|
||||
super().__init__()
|
||||
self.params = params
|
||||
self.encoder = Encoder(
|
||||
resolution=params.resolution,
|
||||
in_channels=params.in_channels,
|
||||
ch=params.ch,
|
||||
ch_mult=params.ch_mult,
|
||||
num_res_blocks=params.num_res_blocks,
|
||||
z_channels=params.z_channels,
|
||||
)
|
||||
self.decoder = Decoder(
|
||||
resolution=params.resolution,
|
||||
in_channels=params.in_channels,
|
||||
ch=params.ch,
|
||||
out_ch=params.out_ch,
|
||||
ch_mult=params.ch_mult,
|
||||
num_res_blocks=params.num_res_blocks,
|
||||
z_channels=params.z_channels,
|
||||
)
|
||||
|
||||
self.bn_eps = 1e-4
|
||||
self.bn_momentum = 0.1
|
||||
self.ps = [2, 2]
|
||||
self.bn = torch.nn.BatchNorm2d(
|
||||
math.prod(self.ps) * params.z_channels,
|
||||
eps=self.bn_eps,
|
||||
momentum=self.bn_momentum,
|
||||
affine=False,
|
||||
track_running_stats=True,
|
||||
)
|
||||
|
||||
@property
|
||||
def device(self) -> torch.device:
|
||||
return next(self.parameters()).device
|
||||
|
||||
@property
|
||||
def dtype(self) -> torch.dtype:
|
||||
return next(self.parameters()).dtype
|
||||
|
||||
|
||||
_NUM_RESOLUTIONS = 4
|
||||
|
||||
|
||||
def convert_diffusers_state_dict(src: dict[str, Tensor]) -> dict[str, Tensor]:
|
||||
out: dict[str, Tensor] = {}
|
||||
attn_substrings = (".mid.attn_1.",)
|
||||
for src_key, tensor in src.items():
|
||||
dst_key = _rewrite_diffusers_key(src_key)
|
||||
if dst_key is None:
|
||||
raise KeyError(f"Unrecognized diffusers VAE state-dict key: {src_key}")
|
||||
if (
|
||||
any(s in dst_key for s in attn_substrings)
|
||||
and dst_key.endswith(".weight")
|
||||
and tensor.ndim == 2
|
||||
):
|
||||
tensor = tensor.unsqueeze(-1).unsqueeze(-1)
|
||||
out[dst_key] = tensor
|
||||
return out
|
||||
|
||||
|
||||
def _rewrite_diffusers_key(key: str) -> str | None:
|
||||
if key.startswith("bn."):
|
||||
return key
|
||||
|
||||
if key.startswith("quant_conv."):
|
||||
return key.replace("quant_conv.", "encoder.quant_conv.", 1)
|
||||
if key.startswith("post_quant_conv."):
|
||||
return key.replace("post_quant_conv.", "decoder.post_quant_conv.", 1)
|
||||
|
||||
if key == "encoder.conv_norm_out.weight":
|
||||
return "encoder.norm_out.weight"
|
||||
if key == "encoder.conv_norm_out.bias":
|
||||
return "encoder.norm_out.bias"
|
||||
if key == "decoder.conv_norm_out.weight":
|
||||
return "decoder.norm_out.weight"
|
||||
if key == "decoder.conv_norm_out.bias":
|
||||
return "decoder.norm_out.bias"
|
||||
|
||||
m = re.match(r"^(encoder|decoder)\.mid_block\.resnets\.(\d+)\.(.+)$", key)
|
||||
if m:
|
||||
side, idx, rest = m.group(1), int(m.group(2)), m.group(3)
|
||||
rest = rest.replace("conv_shortcut", "nin_shortcut")
|
||||
return f"{side}.mid.block_{idx + 1}.{rest}"
|
||||
m = re.match(r"^(encoder|decoder)\.mid_block\.attentions\.0\.(.+)$", key)
|
||||
if m:
|
||||
side, rest = m.group(1), m.group(2)
|
||||
rest = (
|
||||
rest.replace("group_norm.", "norm.")
|
||||
.replace("to_q.", "q.")
|
||||
.replace("to_k.", "k.")
|
||||
.replace("to_v.", "v.")
|
||||
.replace("to_out.0.", "proj_out.")
|
||||
)
|
||||
return f"{side}.mid.attn_1.{rest}"
|
||||
|
||||
m = re.match(r"^encoder\.down_blocks\.(\d+)\.resnets\.(\d+)\.(.+)$", key)
|
||||
if m:
|
||||
level, res_idx, rest = m.group(1), m.group(2), m.group(3)
|
||||
rest = rest.replace("conv_shortcut", "nin_shortcut")
|
||||
return f"encoder.down.{level}.block.{res_idx}.{rest}"
|
||||
m = re.match(r"^encoder\.down_blocks\.(\d+)\.downsamplers\.0\.conv\.(.+)$", key)
|
||||
if m:
|
||||
return f"encoder.down.{m.group(1)}.downsample.conv.{m.group(2)}"
|
||||
|
||||
m = re.match(r"^decoder\.up_blocks\.(\d+)\.resnets\.(\d+)\.(.+)$", key)
|
||||
if m:
|
||||
diffusers_idx = int(m.group(1))
|
||||
res_idx = m.group(2)
|
||||
rest = m.group(3).replace("conv_shortcut", "nin_shortcut")
|
||||
return (
|
||||
f"decoder.up.{_NUM_RESOLUTIONS - 1 - diffusers_idx}.block.{res_idx}.{rest}"
|
||||
)
|
||||
m = re.match(r"^decoder\.up_blocks\.(\d+)\.upsamplers\.0\.conv\.(.+)$", key)
|
||||
if m:
|
||||
diffusers_idx = int(m.group(1))
|
||||
return f"decoder.up.{_NUM_RESOLUTIONS - 1 - diffusers_idx}.upsample.conv.{m.group(2)}"
|
||||
|
||||
if key.startswith(
|
||||
(
|
||||
"encoder.conv_in.",
|
||||
"encoder.conv_out.",
|
||||
"decoder.conv_in.",
|
||||
"decoder.conv_out.",
|
||||
)
|
||||
):
|
||||
return key
|
||||
|
||||
return None
|
||||
@@ -1,6 +1,6 @@
|
||||
import { GroupedSelectOption, SelectOption, JobConfig } from '@/types';
|
||||
import { defaultSliderConfig } from './jobConfig';
|
||||
import { defaultAudioSampleConfig, defaultSampleConfig } from '@/helpers/defaultSamples';
|
||||
import { defaultAudioSampleConfig, defaultSampleConfig, defaultIdeogramSamplesConfig } from '@/helpers/defaultSamples';
|
||||
|
||||
type Control = 'depth' | 'line' | 'pose' | 'inpaint';
|
||||
|
||||
@@ -935,7 +935,7 @@ export const modelArchs: ModelArch[] = [
|
||||
'config.process[0].network.conv': [undefined, 16],
|
||||
'config.process[0].network.conv_alpha': [undefined, 16],
|
||||
'config.process[0].train.max_loss': [1.0, undefined],
|
||||
'config.process[0].network.network_kwargs.ignore_if_contains': [['lm_head','patch_embed', 'visual'], []],
|
||||
'config.process[0].network.network_kwargs.ignore_if_contains': [['lm_head', 'patch_embed', 'visual'], []],
|
||||
'config.process[0].network.transformer_only': [false, undefined],
|
||||
'config.process[0].sample.width': [2048, 1024],
|
||||
'config.process[0].sample.height': [2048, 1024],
|
||||
@@ -979,6 +979,29 @@ export const modelArchs: ModelArch[] = [
|
||||
'model.layer_offloading',
|
||||
],
|
||||
},
|
||||
{
|
||||
name: 'ideogram4',
|
||||
label: 'Ideogram4',
|
||||
group: 'experimental',
|
||||
defaults: {
|
||||
'config.process[0].model.name_or_path': ['ideogram-ai/ideogram-4-fp8', defaultNameOrPath],
|
||||
'config.process[0].model.quantize': [true, false],
|
||||
'config.process[0].model.quantize_te': [true, false],
|
||||
'config.process[0].train.timestep_type': ['linear', 'sigmoid'],
|
||||
'config.process[0].network.conv': [undefined, 16],
|
||||
'config.process[0].network.conv_alpha': [undefined, 16],
|
||||
'config.process[0].model.low_vram': [true, false],
|
||||
'config.process[0].sample': [defaultIdeogramSamplesConfig, defaultSampleConfig],
|
||||
},
|
||||
disableSections: [
|
||||
'network.conv',
|
||||
],
|
||||
additionalSections: [
|
||||
'model.low_vram',
|
||||
'model.layer_offloading',
|
||||
],
|
||||
hasMultiLinePrompts: true,
|
||||
},
|
||||
].sort((a, b) => {
|
||||
// Sort by label, case-insensitive
|
||||
return a.label.localeCompare(b.label, undefined, { sensitivity: 'base' });
|
||||
|
||||
@@ -401,3 +401,522 @@ It's Converging!
|
||||
num_frames: 1,
|
||||
fps: 1,
|
||||
};
|
||||
|
||||
|
||||
export const defaultIdeogramSamplesConfig: SampleConfig = {
|
||||
sampler: 'flowmatch',
|
||||
sample_every: 250,
|
||||
width: 1024,
|
||||
height: 1024,
|
||||
samples: [
|
||||
{
|
||||
prompt: `
|
||||
{
|
||||
"high_level_description": "A red-haired woman calmly playing chess at a wooden table in a sunlit public park during late afternoon, while a violent orange explosion erupts among distant trees behind her, creating a surreal contrast between serene focus and sudden chaos.",
|
||||
"style_description": {
|
||||
"aesthetics": "moody cinematic realism, sharp foreground focus with hazy atmospheric background, dramatic contrast",
|
||||
"lighting": "warm low-angle afternoon sun from the right, intense orange glow from the distant blast adding a secondary warm rim light, long soft shadows",
|
||||
"photo": "full-frame DSLR, 50mm lens, f/2.0, shallow depth of field with motion blur in the background blast",
|
||||
"medium": "photography",
|
||||
"color_palette": ["#1F3D1A", "#C9521E", "#E8B14C", "#9FB37A", "#3A2A1F", "#D9C9A8"]
|
||||
},
|
||||
"compositional_deconstruction": {
|
||||
"background": "An open green park in late afternoon, mown grass receding into the distance, a gravel path curving away to the right, scattered mature trees along the horizon under a hazy pale-blue sky. Atmospheric haze and drifting smoke soften the far distance, with sunlight filtering warmly across the open space to create depth between the near table and the far tree line.",
|
||||
"elements": [
|
||||
{
|
||||
"type": "obj",
|
||||
"bbox": [560, 120, 880, 430],
|
||||
"desc": "A large explosion erupting among the distant trees in the upper-right background, a billowing fireball of orange and yellow flame with a rising column of dark grey-black smoke and debris flung outward. Soft-focused and motion-blurred to read as far away, casting a warm glow that bleeds into the surrounding haze.",
|
||||
"color_palette": ["#C9521E", "#E8B14C", "#3A2A1F", "#7A6A55"]
|
||||
},
|
||||
{
|
||||
"type": "obj",
|
||||
"bbox": [40, 300, 360, 470],
|
||||
"desc": "A cluster of mid-distance park trees along the left horizon, full green canopies catching warm afternoon light on their right edges, partially obscured by drifting smoke from the blast. Slightly soft-focused to sit in the midground behind the table.",
|
||||
"color_palette": ["#1F3D1A", "#9FB37A", "#5C7A3A"]
|
||||
},
|
||||
{
|
||||
"type": "obj",
|
||||
"bbox": [180, 560, 820, 1000],
|
||||
"desc": "A weathered wooden park table occupying the lower portion of the frame, viewed slightly from the side, grain and knots visible on the planks, warm sunlight raking across the surface from the right and casting a long shadow to the left. Sharp focus as the anchoring foreground surface.",
|
||||
"color_palette": ["#5A3A22", "#8A6038", "#D9C9A8"]
|
||||
},
|
||||
{
|
||||
"type": "obj",
|
||||
"bbox": [340, 600, 680, 760],
|
||||
"desc": "A checkered chessboard resting flat on the table in the lower-center of the frame, mid-game with carved wooden pieces casting small crisp shadows in the afternoon light, a few captured pieces set off to one side. Tilted slightly toward the viewer, in sharp focus.",
|
||||
"color_palette": ["#E8E0CC", "#3A2A1F", "#8A6038"]
|
||||
},
|
||||
{
|
||||
"type": "obj",
|
||||
"bbox": [330, 250, 640, 760],
|
||||
"desc": "A woman with vivid red hair seated at the far side of the table, center-left, leaning forward with one hand resting near a chess piece, her gaze fixed calmly downward on the board, fully absorbed and unaware of the chaos behind her. Wavy shoulder-length hair catching warm light and a faint glow from the distant fire, wearing a simple olive jacket, rendered in sharp focus as the primary subject.",
|
||||
"color_palette": ["#B0301F", "#C9521E", "#6E7A4A", "#E2C2A0"]
|
||||
},
|
||||
{
|
||||
"type": "obj",
|
||||
"bbox": [680, 540, 880, 900],
|
||||
"desc": "An empty weathered wooden chair on the near right side of the table, angled slightly outward as if its occupant just left, sunlight warming its right edge and a long shadow falling across the grass beside it. Slightly softer focus than the woman, in the right foreground.",
|
||||
"color_palette": ["#5A3A22", "#8A6038", "#3A2A1F"]
|
||||
}
|
||||
]
|
||||
}
|
||||
}
|
||||
`,
|
||||
},
|
||||
{
|
||||
prompt: `
|
||||
{
|
||||
"high_level_description": "A young woman in a knit beanie sits at a cozy café table on an overcast morning, cradling a warm coffee cup in both hands and gazing thoughtfully toward the window light, evoking a calm, intimate, contemplative mood.",
|
||||
"style_description": {
|
||||
"aesthetics": "moody cinematic lifestyle photography, sharp focus on subject, shallow depth of field with creamy bokeh",
|
||||
"lighting": "soft diffused natural window light from the left, gentle cool daylight key with warm ambient fill from interior tungsten",
|
||||
"photo": "50mm prime lens, f/1.8, slight film grain, DSLR full-frame",
|
||||
"medium": "photography",
|
||||
"color_palette": ["#3E2C20", "#A9876B", "#D8C4A8", "#6E5A48", "#C24A2E", "#E9E2D6"]
|
||||
},
|
||||
"compositional_deconstruction": {
|
||||
"background": "The blurred interior of a rustic café: warm wooden wall paneling and a softly out-of-focus window on the left spilling pale daylight, with hanging pendant lights and shelves of jars receding into a hazy bokeh on the right, creating a sense of warm enclosed depth.",
|
||||
"elements": [
|
||||
{
|
||||
"type": "obj",
|
||||
"bbox": [40, 110, 360, 640],
|
||||
"desc": "A tall frosted café window in the midground left, glowing with diffuse overcast daylight, condensation faintly softening its surface and acting as the primary light source washing across the scene.",
|
||||
"color_palette": ["#E9E2D6", "#C9C0B2", "#9AA0A0"]
|
||||
},
|
||||
{
|
||||
"type": "obj",
|
||||
"bbox": [560, 230, 940, 600],
|
||||
"desc": "A weathered wooden shelf along the back wall on the right, holding a few out-of-focus glass jars and a small potted plant, rendered as soft warm bokeh that anchors the depth behind the subject.",
|
||||
"color_palette": ["#6E5A48", "#8B7152", "#4A3A2C"]
|
||||
},
|
||||
{
|
||||
"type": "obj",
|
||||
"bbox": [240, 250, 760, 1000],
|
||||
"desc": "A young woman seated at the table, occupying the central-lower frame and turned slightly toward the window at left, wearing a chunky cream-and-rust knit beanie pushed back over wavy hair, an oversized oatmeal sweater, her shoulders relaxed and her gaze directed softly off-frame left, catching the cool window light along her cheek and jaw.",
|
||||
"color_palette": ["#A9876B", "#D8C4A8", "#C24A2E", "#3E2C20"]
|
||||
},
|
||||
{
|
||||
"type": "obj",
|
||||
"bbox": [400, 600, 640, 850],
|
||||
"desc": "A rounded ceramic coffee cup held in both of the woman's hands near her chest, foreground center, thin curl of steam rising from the dark surface, the warm interior light catching the glazed rim and her fingers wrapped snugly around it.",
|
||||
"color_palette": ["#3E2C20", "#C9BBA6", "#7A5A40"]
|
||||
},
|
||||
{
|
||||
"type": "obj",
|
||||
"bbox": [180, 780, 880, 1000],
|
||||
"desc": "A rustic dark-wood café table filling the foreground edge, its grain catching a soft sheen, with a faint coffee ring and a folded napkin resting near the bottom corner, grounding the composition.",
|
||||
"color_palette": ["#4A3526", "#6E5238", "#2C1F16"]
|
||||
}
|
||||
]
|
||||
}
|
||||
}
|
||||
`,
|
||||
},
|
||||
{
|
||||
prompt: `
|
||||
{
|
||||
"high_level_description": "A confident horse works the decks as a DJ in a packed underground night club at peak hour, surrounded by laser beams and rolling fog, holding a martini glass in one hoof while the crowd glows in the haze.",
|
||||
"style_description": {
|
||||
"aesthetics": "high-energy nightlife photography, surreal and playful, deep shadows with neon punch, slight chromatic aberration",
|
||||
"lighting": "low ambient darkness cut by saturated laser beams and colored wash from above, strobing key light on the subject, cool magenta and cyan rim light with warm amber spill",
|
||||
"photo": "fisheye lens, ~8mm, exaggerated barrel distortion, f/2.8, slight motion blur in the haze, faint sensor grain",
|
||||
"medium": "photography",
|
||||
"color_palette": ["#0A0410", "#E0146E", "#16D6E8", "#7A1FB0", "#F2B705", "#1B0A2E"]
|
||||
},
|
||||
"compositional_deconstruction": {
|
||||
"background": "A dark, cavernous club interior with a low ceiling rig of moving heads and laser projectors, distant blurred crowd silhouettes packed toward a back bar, dense low-lying fog drifting across the floor, the whole space curved and stretched by the fisheye so vertical lines bow outward toward the frame edges.",
|
||||
"elements": [
|
||||
{
|
||||
"type": "obj",
|
||||
"bbox": [0, 30, 1000, 520],
|
||||
"desc": "Crisscrossing green and magenta laser beams fanning out from ceiling projectors at the top of the frame, splaying across the upper third and bending along the fisheye curve, catching in the fog to form bright volumetric sheets and dotted scatter planes.",
|
||||
"color_palette": ["#16D6E8", "#E0146E", "#3DF27A", "#7A1FB0"]
|
||||
},
|
||||
{
|
||||
"type": "obj",
|
||||
"bbox": [120, 380, 880, 820],
|
||||
"desc": "Thick rolling smoke-machine fog blanketing the midground and floor, lit through by colored beams, denser at center where it glows and thinning toward the warped edges, partially veiling the lower body of the subject.",
|
||||
"color_palette": ["#2A1140", "#7A1FB0", "#16D6E8", "#1B0A2E"]
|
||||
},
|
||||
{
|
||||
"type": "obj",
|
||||
"bbox": [310, 250, 760, 640],
|
||||
"desc": "A glossy black DJ booth and mixer console set across the lower-center foreground, twin turntables and a lit channel mixer with glowing faders and knobs, front edge bulging toward the viewer from the fisheye distortion, reflective surface catching pink and cyan highlights.",
|
||||
"color_palette": ["#0A0410", "#E0146E", "#16D6E8", "#F2B705"]
|
||||
},
|
||||
{
|
||||
"type": "obj",
|
||||
"bbox": [270, 90, 720, 560],
|
||||
"desc": "A chestnut-brown horse standing as the DJ behind the booth, head and torso filling the center frame, leaning into one turntable with a front hoof, mane catching colored rim light, ears forward and gaze toward the crowd, slight beam glare across its glossy coat.",
|
||||
"color_palette": ["#7A4B2A", "#3A2414", "#E0146E", "#16D6E8"]
|
||||
},
|
||||
{
|
||||
"type": "obj",
|
||||
"bbox": [640, 300, 820, 560],
|
||||
"desc": "A classic martini glass held aloft in the horse's raised right hoof on the mid-right, pale liquid with an olive on a pick, the glass rim and surface flashing cyan and magenta laser glints.",
|
||||
"color_palette": ["#E8E2C0", "#16D6E8", "#E0146E", "#9BAF3A"]
|
||||
},
|
||||
{
|
||||
"type": "obj",
|
||||
"bbox": [0, 560, 1000, 1000],
|
||||
"desc": "Foreground crowd of raised hands and blurred dancing silhouettes along the bottom edge, backlit into dark shapes rimmed by neon, hands reaching up into the fog and beams, heavily curved by the fisheye at the frame's lower corners.",
|
||||
"color_palette": ["#1B0A2E", "#E0146E", "#16D6E8", "#7A1FB0"]
|
||||
}
|
||||
]
|
||||
}
|
||||
}
|
||||
`,
|
||||
},
|
||||
{
|
||||
prompt: `
|
||||
{
|
||||
"high_level_description": "A cheerful man stands on a sunny beach proudly showing off his new graphic t-shirt while a large shark breaches the ocean surface behind him, blending a relaxed summer mood with a playful jolt of drama.",
|
||||
"style_description": {
|
||||
"aesthetics": "vibrant lifestyle photography, sharp focus on the subject with a softly blurred background, candid summer energy",
|
||||
"lighting": "bright natural midday sunlight from the upper right, warm and slightly hard, with strong highlights on skin and water and gentle fill from sand bounce",
|
||||
"photo": "DSLR, 50mm lens, f/4, fast shutter to freeze the breaching shark",
|
||||
"medium": "photography",
|
||||
"color_palette": ["#3FA9D6", "#F2E2B6", "#1C6E8C", "#E8B04B", "#2B3A42"]
|
||||
},
|
||||
"compositional_deconstruction": {
|
||||
"background": "A wide stretch of bright tropical beach under a clear blue sky with scattered wispy clouds near the horizon; turquoise ocean water fills the middle band of the frame, transitioning to pale golden wet sand in the lower foreground, with soft heat haze and shimmering reflections suggesting strong midday sun and open depth toward the horizon.",
|
||||
"elements": [
|
||||
{
|
||||
"type": "obj",
|
||||
"bbox": [560, 300, 800, 470],
|
||||
"desc": "A large gray shark breaching upward out of the ocean in the midground right, body angled diagonally with its head clearing the surface and tail still in the water, mouth slightly open, spraying a burst of white sea foam and droplets around it; smaller in scale due to its distance, creating a dramatic but slightly comedic focal accent behind the man.",
|
||||
"color_palette": ["#6B7C85", "#9FB3BC", "#FFFFFF", "#1C6E8C"]
|
||||
},
|
||||
{
|
||||
"type": "obj",
|
||||
"bbox": [230, 250, 560, 940],
|
||||
"desc": "A smiling man standing center-left in the foreground, facing the camera, tugging the hem of his t-shirt outward with both hands to show it off, weight on one leg in a relaxed casual stance; tanned skin catching warm sunlight, wind lightly moving his hair, gaze directed at the viewer with an enthusiastic expression, occupying the dominant vertical mass of the frame.",
|
||||
"color_palette": ["#D9A66C", "#3B3027", "#2B3A42", "#F2E2B6"]
|
||||
},
|
||||
{
|
||||
"type": "obj",
|
||||
"bbox": [290, 470, 510, 700],
|
||||
"desc": "The man's new graphic t-shirt, a vivid printed cotton tee stretched outward by his hands to display its bold front design, fabric slightly creased from being pulled, catching crisp highlights along the folds and reading as the proud centerpiece of his pose.",
|
||||
"color_palette": ["#E84C3D", "#FFFFFF", "#2B3A42", "#E8B04B"]
|
||||
},
|
||||
{
|
||||
"type": "obj",
|
||||
"bbox": [0, 760, 1000, 1000],
|
||||
"desc": "A band of damp golden beach sand sweeping across the lower foreground, dotted with faint footprints and tiny shell fragments, surface glinting where the wet sheen reflects the bright sky, anchoring the man and adding tactile depth at the bottom of the frame.",
|
||||
"color_palette": ["#E8D6A6", "#C9A86A", "#F5EBD0"]
|
||||
}
|
||||
]
|
||||
}
|
||||
}
|
||||
`,
|
||||
},
|
||||
{
|
||||
prompt: `
|
||||
{
|
||||
"high_level_description": "A large brown bear works on constructing a half-built log cabin in a snow-covered mountain valley on a crisp, overcast winter morning, conveying a cozy yet industrious, slightly whimsical wilderness mood.",
|
||||
"style_description": {
|
||||
"aesthetics": "warm storybook digital illustration, painterly textures, gentle stylization with soft edges and rich detail",
|
||||
"lighting": "soft diffused overcast daylight from the upper left, cool ambient shadows with subtle warm bounce from the wood",
|
||||
"photo": "",
|
||||
"medium": "digital illustration",
|
||||
"color_palette": ["#E8EEF2", "#A9C0D0", "#7A5230", "#3D5A4C", "#C8895A", "#2E3A44"]
|
||||
},
|
||||
"compositional_deconstruction": {
|
||||
"background": "Towering snow-capped mountain peaks recede into a pale grey-blue winter sky, partially veiled by drifting mist. A dense stand of frosted evergreen pines blankets the mid-slope, and an unbroken sheet of fresh snow covers the valley floor, broken only by faint drifts and gentle undulations that establish depth toward the distant ridgeline.",
|
||||
"elements": [
|
||||
{
|
||||
"type": "obj",
|
||||
"bbox": [60, 230, 420, 540],
|
||||
"desc": "A cluster of tall frosted pine trees standing in the midground left, their dark green boughs heavy with clumps of snow, slightly out of focus to push them back in depth and frame the construction scene.",
|
||||
"color_palette": ["#3D5A4C", "#2E3A44", "#E8EEF2"]
|
||||
},
|
||||
{
|
||||
"type": "obj",
|
||||
"bbox": [380, 420, 880, 760],
|
||||
"desc": "A half-finished log cabin occupying the center-right midground, built of stacked honey-brown timber logs about four courses high with notched corner joints, an open doorway gap on the left side, and a dusting of snow along the top logs; freshly hewn and inviting.",
|
||||
"color_palette": ["#7A5230", "#C8895A", "#5A3C22", "#E8EEF2"]
|
||||
},
|
||||
{
|
||||
"type": "obj",
|
||||
"bbox": [520, 280, 700, 480],
|
||||
"desc": "A large brown bear standing upright on its hind legs in the right midground beside the cabin, gripping a heavy log between its forepaws as it lifts it into place, fur thick and shaggy with snowflakes caught in it, head tilted toward its work with a focused expression.",
|
||||
"color_palette": ["#5A3C22", "#3A2818", "#8A6038", "#E8EEF2"]
|
||||
},
|
||||
{
|
||||
"type": "obj",
|
||||
"bbox": [150, 700, 470, 850],
|
||||
"desc": "A small stack of cut logs and a couple of fresh wood shavings resting on the snow in the foreground left, the sawn ends showing pale concentric rings, a few flecks of sawdust scattered across the surrounding snow.",
|
||||
"color_palette": ["#C8895A", "#7A5230", "#EDE3D2"]
|
||||
},
|
||||
{
|
||||
"type": "obj",
|
||||
"bbox": [600, 760, 760, 880],
|
||||
"desc": "A worn steel axe embedded blade-down in a chopping stump in the lower foreground right, its wooden handle angled up and to the right, light snow gathered on the stump's flat top.",
|
||||
"color_palette": ["#6B7178", "#4A3320", "#E8EEF2"]
|
||||
},
|
||||
{
|
||||
"type": "obj",
|
||||
"bbox": [0, 0, 1000, 220],
|
||||
"desc": "Faint drifting snowflakes and a soft mist layer across the upper portion of the frame, thinning the distant peaks and adding cold atmospheric depth without obscuring the foreground.",
|
||||
"color_palette": ["#E8EEF2", "#A9C0D0", "#C4D4DE"]
|
||||
}
|
||||
]
|
||||
}
|
||||
}
|
||||
`,
|
||||
},
|
||||
{
|
||||
prompt: `
|
||||
{
|
||||
"high_level_description": "A punk rock woman performs alone on a dark stage at night, mid-song with her guitar, surrounded by sweeping laser lights and haze that cut through the darkness, charged with raw, electric energy.",
|
||||
"style_description": {
|
||||
"aesthetics": "moody cinematic concert photography, high contrast, gritty and energetic with motion in the lights",
|
||||
"lighting": "dramatic stage lighting from above and behind, vivid colored laser beams cutting through atmospheric haze, hard rim light on the performer, deep shadows",
|
||||
"photo": "DSLR concert photograph, 85mm, f/2.0, slight motion blur on the lasers, ISO grain",
|
||||
"medium": "photography",
|
||||
"color_palette": ["#0B0B12", "#E0204A", "#1FB6C9", "#7A2BD4", "#F2E9D8", "#101830"]
|
||||
},
|
||||
"compositional_deconstruction": {
|
||||
"background": "A deep, near-black stage interior at night, dense with atmospheric haze that catches the light. Vertical truss rigging and dim silhouettes of speaker stacks recede into darkness at the edges. The far depth dissolves into a smoky void that gives the colored beams room to glow.",
|
||||
"elements": [
|
||||
{
|
||||
"type": "obj",
|
||||
"bbox": [0, 20, 1000, 620],
|
||||
"desc": "A web of vivid laser beams fanning out from rigging high in the frame, sweeping diagonally across the upper two-thirds of the scene. Sharp magenta, cyan, and violet lines pierce the haze, some converging toward the performer, others splaying outward, with slight motion blur conveying their sweep.",
|
||||
"color_palette": ["#E0204A", "#1FB6C9", "#7A2BD4", "#F2E9D8"]
|
||||
},
|
||||
{
|
||||
"type": "obj",
|
||||
"bbox": [120, 80, 360, 520],
|
||||
"desc": "A tall stage spotlight rig at midground left, its lens flaring a hot magenta glow that bleeds into the surrounding haze. Partially silhouetted metal housing, angled down toward center stage.",
|
||||
"color_palette": ["#E0204A", "#101830", "#0B0B12"]
|
||||
},
|
||||
{
|
||||
"type": "obj",
|
||||
"bbox": [650, 90, 880, 500],
|
||||
"desc": "A cluster of cyan-tinted backlights at midground right mounted on dark truss, throwing a cool teal wash and lens flares through the smoke, balancing the warm magenta on the opposite side.",
|
||||
"color_palette": ["#1FB6C9", "#101830", "#0B0B12"]
|
||||
},
|
||||
{
|
||||
"type": "obj",
|
||||
"bbox": [355, 180, 660, 940],
|
||||
"desc": "The punk rocker woman, full-bodied and slightly right of center, captured mid-song. She leans into a microphone on a stand, mouth open singing, eyes intense. Spiked, brightly dyed hair, ripped band tee, studded leather, and torn jeans. Her body is rim-lit by colored stage light against the dark haze, sweat catching highlights, in a dynamic forward-leaning stance.",
|
||||
"color_palette": ["#E0204A", "#F2E9D8", "#0B0B12", "#7A2BD4"]
|
||||
},
|
||||
{
|
||||
"type": "obj",
|
||||
"bbox": [330, 440, 690, 820],
|
||||
"desc": "An electric guitar slung across the woman's torso, held mid-strum with her right hand near the strings and left hand on the fretboard. Glossy body catching sharp magenta and cyan reflections, chrome hardware flaring, positioned diagonally across her lower body.",
|
||||
"color_palette": ["#101830", "#1FB6C9", "#E0204A", "#F2E9D8"]
|
||||
},
|
||||
{
|
||||
"type": "obj",
|
||||
"bbox": [430, 250, 545, 560],
|
||||
"desc": "A chrome microphone on a slim stand directly in front of the singer's face at center frame, catching a bright specular highlight from the stage lights, slightly haloed by the haze.",
|
||||
"color_palette": ["#F2E9D8", "#0B0B12", "#1FB6C9"]
|
||||
},
|
||||
{
|
||||
"type": "obj",
|
||||
"bbox": [0, 780, 1000, 1000],
|
||||
"desc": "The dark, glossy stage floor across the foreground, reflecting smeared streaks of magenta, cyan, and violet from the lasers and the performer above, fading to black at the front edge.",
|
||||
"color_palette": ["#0B0B12", "#E0204A", "#1FB6C9", "#7A2BD4"]
|
||||
}
|
||||
]
|
||||
}
|
||||
}
|
||||
`,
|
||||
},
|
||||
{
|
||||
prompt: `
|
||||
{
|
||||
"high_level_description": "A bearded hipster craftsman in a sunlit woodworking shop, focused intently as he assembles a wooden chair surrounded by tools and sawdust, evoking a warm, artisanal, hands-on mood.",
|
||||
"style_description": {
|
||||
"aesthetics": "warm documentary photography, sharp focus on subject, shallow depth of field, authentic artisanal feel",
|
||||
"lighting": "soft natural light streaming from a large window on the left, warm afternoon color temperature with gentle dust particles catching the light",
|
||||
"photo": "35mm DSLR, 50mm lens, f/2.0, slight film grain",
|
||||
"medium": "photography",
|
||||
"color_palette": ["#6B4A2E", "#C9A36B", "#E8D5B0", "#3A2C1F", "#8C7355", "#D9C2A0"]
|
||||
},
|
||||
"compositional_deconstruction": {
|
||||
"background": "A rustic woodworking shop interior with weathered plank walls, a large window on the left admitting bright afternoon light, blurred shelves of tools and lumber stacked in the soft-focus depth, faint sawdust haze drifting through the air, warm earthy tones throughout.",
|
||||
"elements": [
|
||||
{
|
||||
"type": "obj",
|
||||
"bbox": [40, 120, 320, 520],
|
||||
"desc": "A pegboard and tool rack mounted on the back wall, upper-left of frame, slightly out of focus, holding hand saws, chisels, and clamps hanging in rows, midground depth, muted in the soft background light.",
|
||||
"color_palette": ["#5A4630", "#8C7355", "#3A2C1F"]
|
||||
},
|
||||
{
|
||||
"type": "obj",
|
||||
"bbox": [620, 250, 980, 720],
|
||||
"desc": "A heavy wooden workbench occupying the right side of the frame, its surface scattered with hand tools, wood shavings, and a coffee mug, warm light grazing the worn timber grain, midground to foreground.",
|
||||
"color_palette": ["#6B4A2E", "#C9A36B", "#4A3622"]
|
||||
},
|
||||
{
|
||||
"type": "obj",
|
||||
"bbox": [300, 160, 640, 760],
|
||||
"desc": "A bearded hipster man in his early thirties, center-left of frame, wearing a rolled-sleeve flannel shirt and a leather apron, leaning forward with focused concentration as he works, hair tied back, forearms tensed, side-lit by the window so highlights catch his beard and the dust around him.",
|
||||
"color_palette": ["#8C5A3A", "#C9A36B", "#5A3B28", "#E8D5B0"]
|
||||
},
|
||||
{
|
||||
"type": "obj",
|
||||
"bbox": [380, 520, 760, 920],
|
||||
"desc": "A partially assembled solid-wood chair in the lower-center foreground, frame and three legs joined with a fourth resting nearby, freshly planed pale timber showing clean tool marks, the man's hands fitting a joint, sharply in focus.",
|
||||
"color_palette": ["#D9C2A0", "#C9A36B", "#9C7A4E"]
|
||||
},
|
||||
{
|
||||
"type": "obj",
|
||||
"bbox": [620, 780, 940, 980],
|
||||
"desc": "A scatter of wood shavings and sawdust across the shop floor in the immediate foreground, curled pale ribbons of planed wood catching warm light, slightly blurred at the bottom edge.",
|
||||
"color_palette": ["#E8D5B0", "#C9A36B", "#8C7355"]
|
||||
}
|
||||
]
|
||||
}
|
||||
}
|
||||
`,
|
||||
},
|
||||
{
|
||||
prompt: `
|
||||
{
|
||||
"high_level_description": "A whimsical low-light photograph of an anthropomorphic cat whose face and body are split vertically into a black half and an orange tabby half, wearing a blue tophat and holding a yarn-filled martini glass in one paw and a movie DVD case in the other, set inside a crowded, neon-lit nightclub where a giant mushroom man dances with a bear.",
|
||||
"style_description": {
|
||||
"aesthetics": "playful surreal portrait photography, sharp focus on the subject with a softly blurred background, shallow depth of field, vivid saturated color",
|
||||
"lighting": "moody nightclub ambiance with colorful neon key light, magenta and cyan rim accents, a warm spotlight catching the cat from front-left",
|
||||
"photo": "DSLR portrait, 50mm lens, f/2.0, slight motion blur in the background crowd",
|
||||
"medium": "photography",
|
||||
"color_palette": ["#101012", "#E8741C", "#2E63C9", "#E84B8A", "#3FB55C", "#C9A227"]
|
||||
},
|
||||
"compositional_deconstruction": {
|
||||
"background": "Interior of a busy, dimly lit nightclub seen at night, packed with a blurred dancing crowd, hazy atmosphere thick with colored light beams, glowing magenta and cyan neon signage along the walls, a distant illuminated bar and DJ glow, reflective dark floor scattering specular highlights into bokeh.",
|
||||
"elements": [
|
||||
{
|
||||
"type": "obj",
|
||||
"bbox": [110, 180, 400, 760],
|
||||
"desc": "A giant mushroom man towering in the midground left, with a broad spotted red-and-cream cap for a head and a plump pale stalk body, arms thrown up mid-dance, slightly motion-blurred and softened by depth of field, lit by shifting neon so its cap glows magenta and cyan at the edges.",
|
||||
"color_palette": ["#D24B4B", "#EDE3CF", "#B0457F", "#3A8FC4"]
|
||||
},
|
||||
{
|
||||
"type": "obj",
|
||||
"bbox": [360, 240, 650, 780],
|
||||
"desc": "A large shaggy brown bear standing on its hind legs in the midground right of center, dancing face-to-face with the mushroom man, one paw raised, fur catching cyan rim light, blurred by the shallow focus to read as a joyful background partner.",
|
||||
"color_palette": ["#5A3A22", "#7A4B2A", "#2E8FB0", "#1B1A1C"]
|
||||
},
|
||||
{
|
||||
"type": "obj",
|
||||
"bbox": [290, 250, 720, 950],
|
||||
"desc": "The main subject, an upright anthropomorphic cat filling the foreground center, its entire body split cleanly down the vertical midline: the left half solid glossy black, the right half warm orange tabby with darker stripes, the seam running straight through nose, chest and belly. It sits tall and forward-facing, both front paws raised to hold props, eyes catching the front spotlight with bright reflective pupils, fur crisply in focus against the soft crowd behind.",
|
||||
"color_palette": ["#141414", "#E8741C", "#B85416", "#F0E2C8"]
|
||||
},
|
||||
{
|
||||
"type": "obj",
|
||||
"bbox": [395, 175, 625, 360],
|
||||
"desc": "A tall royal-blue tophat perched on the cat's head in the upper-center of the frame, with a satin sheen, a slim darker band around the base and a stiff curled brim, neon highlights skating across its smooth surface.",
|
||||
"color_palette": ["#2E63C9", "#1A3C8A", "#5C8BE0", "#0E1F4A"]
|
||||
},
|
||||
{
|
||||
"type": "obj",
|
||||
"bbox": [225, 540, 430, 825],
|
||||
"desc": "A clear martini glass gripped in the cat's right paw (viewer's left, lower-left of frame), tilted slightly, holding a round ball of bright pink yarn instead of a drink, with two slender green knitting needles poking up and out at angles from the yarn; glass rim and stem flash sharp neon specular highlights.",
|
||||
"color_palette": ["#E84B8A", "#3FB55C", "#D8E4EC", "#A8285F"]
|
||||
},
|
||||
{
|
||||
"type": "obj",
|
||||
"bbox": [615, 560, 835, 835],
|
||||
"desc": "A glossy DVD case held flat-faced toward the camera in the cat's left paw (viewer's right, lower-right of frame), its cover showing a standing golden robot rendered in metallic sheen and the bold title text 'This is a test' across the top, plastic surface reflecting magenta and cyan club light.",
|
||||
"color_palette": ["#C9A227", "#E8C95A", "#1B1B22", "#E8741C"]
|
||||
}
|
||||
]
|
||||
}
|
||||
}
|
||||
`,
|
||||
},
|
||||
{
|
||||
prompt: `
|
||||
{
|
||||
"high_level_description": "A man stands in a relaxed pose holding up a plain handheld sign reading 'this is a sign' in clear hand-lettered text, photographed in a bright neutral studio setting with a candid, mildly amused mood.",
|
||||
"style_description": {
|
||||
"aesthetics": "clean editorial portrait photography, sharp focus on the subject, gentle background falloff",
|
||||
"lighting": "soft diffused key from the upper left with a subtle fill, neutral-to-slightly-warm color temperature, gentle shadows",
|
||||
"photo": "85mm portrait lens, f/4, full-frame DSLR, slight depth of field separating subject from backdrop",
|
||||
"medium": "photography",
|
||||
"color_palette": ["#E8E2D8", "#C9B79C", "#4A4038", "#8A6F52", "#F5F2EC"]
|
||||
},
|
||||
"compositional_deconstruction": {
|
||||
"background": "A smooth, evenly lit seamless studio backdrop in a warm off-white tone, fading to a marginally darker gradient toward the edges, with soft out-of-focus texture and no distinct objects, creating shallow depth behind the subject.",
|
||||
"elements": [
|
||||
{
|
||||
"type": "obj",
|
||||
"bbox": [300, 180, 720, 1000],
|
||||
"desc": "A man of average build standing centered-slightly-left and facing the camera, occupying most of the vertical frame from mid-thigh upward, wearing a casual crew-neck sweater and simple trousers; his left arm is raised to grip the bottom of the sign while his right hand steadies its side, head tilted with a faint knowing half-smile and direct eye contact toward the viewer, softly lit from the upper left with a clean catchlight in his eyes.",
|
||||
"color_palette": ["#5A4B3A", "#9C8266", "#E2D4C2", "#3A3128"]
|
||||
},
|
||||
{
|
||||
"type": "obj",
|
||||
"bbox": [360, 250, 700, 470],
|
||||
"desc": "A rectangular flat handheld sign held up at chest-to-shoulder height in the upper-middle of the frame, tilted just slightly off-level, made of plain matte white poster board with a thin neutral border; across its face in bold, evenly spaced hand-lettered black marker reads the phrase 'this is a sign', legible and centered, catching the soft key light with a faint even sheen and casting a gentle shadow onto the man's chest.",
|
||||
"color_palette": ["#F7F4EE", "#1E1B17", "#C8C0B2"]
|
||||
}
|
||||
]
|
||||
}
|
||||
}
|
||||
`,
|
||||
},
|
||||
{
|
||||
prompt: `
|
||||
{
|
||||
"high_level_description": "A grizzled bulldog clad in a battered leather jacket stands guard with a sawed-off shotgun beside his motorcycle in a sun-scorched post-apocalyptic desert at harsh late afternoon, exuding a gritty, defiant survivor mood.",
|
||||
"style_description": {
|
||||
"aesthetics": "moody cinematic, gritty post-apocalyptic realism, high detail with weathered textures, shallow depth of field on the subject",
|
||||
"lighting": "harsh low-angle late afternoon sun from the right, long shadows, warm golden key with dusty haze and faint cool fill in the shadows",
|
||||
"photo": "full-frame DSLR, 50mm lens, f/2.8, slight atmospheric grain",
|
||||
"medium": "photography",
|
||||
"color_palette": ["#C8923F", "#7A4B2A", "#3A2E22", "#D9C49A", "#5A6B5E", "#2B2520"]
|
||||
},
|
||||
"compositional_deconstruction": {
|
||||
"background": "A vast cracked desert wasteland stretching to a hazy horizon, dunes and dry parched earth underfoot, rusted debris and a skeletal collapsed structure faintly visible in the far distance, blown dust and heat shimmer thickening the air, pale washed-out sky tinged amber by airborne particulates and the low sun.",
|
||||
"elements": [
|
||||
{
|
||||
"type": "obj",
|
||||
"bbox": [40, 360, 360, 620],
|
||||
"desc": "Rusted skeletal remains of a collapsed structure in the far background midground-left, sun-bleached metal girders and corrugated sheeting half-buried in sand, small in the frame to establish scale and ruin, partially obscured by dust haze.",
|
||||
"color_palette": ["#8A6A45", "#5A4632", "#B8A079"]
|
||||
},
|
||||
{
|
||||
"type": "obj",
|
||||
"bbox": [480, 380, 880, 760],
|
||||
"desc": "A heavily weathered chopper-style motorcycle parked at a three-quarter angle on the right side of the frame just behind the dog, chrome dulled and pitted, leather seat cracked, dust-caked tank and rust streaks, kickstand down, long shadow cast leftward across the cracked earth.",
|
||||
"color_palette": ["#3A2E22", "#7A5A38", "#9A9088", "#2B2520"]
|
||||
},
|
||||
{
|
||||
"type": "obj",
|
||||
"bbox": [250, 300, 640, 920],
|
||||
"desc": "A muscular English bulldog standing upright and front-and-center slightly left, broad wrinkled jowls and stocky frame facing the viewer with a determined glaring gaze, wearing a scuffed brown leather jacket with frayed collar and metal studs, warm rim light catching the right edge of his face and shoulders, the dominant subject in sharp focus.",
|
||||
"color_palette": ["#C8923F", "#6B4528", "#3A2A1E", "#E0C9A0"]
|
||||
},
|
||||
{
|
||||
"type": "obj",
|
||||
"bbox": [360, 540, 700, 780],
|
||||
"desc": "A sawed-off double-barrel shotgun gripped across the bulldog's body in the lower midground, weathered steel barrels with surface rust and a worn wooden stock, held at a downward diagonal angle catching a glint of the warm sun on the metal.",
|
||||
"color_palette": ["#4A4038", "#6B4A2E", "#8A7460"]
|
||||
},
|
||||
{
|
||||
"type": "obj",
|
||||
"bbox": [120, 760, 900, 960],
|
||||
"desc": "Foreground cracked desert floor in the lower frame, parched fissured clay and scattered sand with small pebbles, dust kicked up faintly, raking warm light emphasizing the deep crack texture and the long shadows of the subject and motorcycle.",
|
||||
"color_palette": ["#D9C49A", "#A8895C", "#6B5A40"]
|
||||
}
|
||||
]
|
||||
}
|
||||
}
|
||||
`,
|
||||
},
|
||||
],
|
||||
neg: '',
|
||||
seed: 42,
|
||||
walk_seed: true,
|
||||
guidance_scale: 4,
|
||||
sample_steps: 30,
|
||||
num_frames: 1,
|
||||
fps: 1,
|
||||
}
|
||||
@@ -1 +1 @@
|
||||
VERSION = "0.9.15"
|
||||
VERSION = "0.10.0"
|
||||
|
||||
Reference in New Issue
Block a user