Spaces:
Runtime error
Runtime error
| """ | |
| Lumina 2.0 ZeroGPU demo – Hugging Face Spaces | |
| Run locally with: python app.py | |
| Push to HF and pick the “ZeroGPU” hardware tier. | |
| """ | |
| import os, gc, json, random, time | |
| import numpy as np | |
| import torch | |
| import gradio as gr | |
| import spaces | |
| from torchvision.transforms.functional import to_pil_image | |
| from tqdm import tqdm | |
| from diffusers.models import AutoencoderKL | |
| from transformers import AutoModel, AutoTokenizer | |
| from huggingface_hub import hf_hub_download | |
| import functools | |
| import models # local | |
| from transport import Sampler, create_transport | |
| # ------------------------------------------------------------------ | |
| # CONSTANTS / ENV CONFIG | |
| # ------------------------------------------------------------------ | |
| REPO_ID = "OnomaAIResearch/Illustrious-Lumina-v0.03" | |
| HF_TOKEN = os.getenv("HF_TOKEN") | |
| CKPT_FILE = "consolidated_ema.00-of-01.pth" | |
| ARGS_FILE = "model_args.pth" | |
| VAE_TYPE = os.getenv("VAE_TYPE", "flux") # flux | ema | mse | sdxl | |
| PRECISION = os.getenv("PRECISION", "bf16") # bf16 | fp16 | fp32 | |
| TEXT_ENCODER_MODEL = os.getenv("TEXT_ENCODER_MODEL", "google/gemma-2-2B") | |
| DTYPE = {"bf16": torch.bfloat16, "fp16": torch.float16, "fp32": torch.float32}[PRECISION] | |
| # ------------------------------------------------------------------ | |
| # GLOBALS (kept on **CPU**) | |
| # ------------------------------------------------------------------ | |
| tokenizer = text_encoder = vae = model = sampler = transport = None | |
| train_args = cap_feat_dim = None | |
| # ------------------------------------------------------------------ | |
| # HELPERS | |
| # ------------------------------------------------------------------ | |
| def encode_prompt(batch, enc, tok, dev, dtype): | |
| """Temporarily moves the text‑encoder to *dev* to extract embeddings.""" | |
| captions = [(c if isinstance(c, str) else c[0]) for c in batch] | |
| enc.to(dev) | |
| with torch.no_grad(), torch.autocast(device_type=dev.split(":")[0], dtype=dtype): | |
| inputs = tok( | |
| captions, | |
| padding=True, | |
| pad_to_multiple_of=8, | |
| max_length=256, | |
| truncation=True, | |
| return_tensors="pt", | |
| ).to(dev) | |
| out = enc(**inputs, output_hidden_states=True).hidden_states[-2] | |
| enc.to("cpu"); gc.collect() | |
| return out.cpu(), inputs.attention_mask.cpu() | |
| def none_or_str(v): # used by create_transport in original repo | |
| return None if v in (None, "None") else str(v) | |
| def load_models(): | |
| global tokenizer, text_encoder, vae, model, sampler, transport, train_args, cap_feat_dim | |
| torch.set_grad_enabled(False) | |
| # ---------------- Hub downloads (cached after first run) ------ | |
| ckpt_path = hf_hub_download(REPO_ID, filename=CKPT_FILE, token=HF_TOKEN) | |
| args_path = hf_hub_download(REPO_ID, filename=ARGS_FILE, token=HF_TOKEN) | |
| # --------------- training args -------------------------------- | |
| train_args = torch.load(args_path, map_location="cpu", weights_only=False) | |
| # --------------- tokenizer / text encoder --------------------- | |
| tokenizer = AutoTokenizer.from_pretrained("google/gemma-2-2b", token=HF_TOKEN, padding_side="right") | |
| text_encoder = AutoModel.from_pretrained( | |
| "google/gemma-2-2b", torch_dtype=DTYPE, token=HF_TOKEN | |
| ).eval().cpu() | |
| cap_feat_dim = text_encoder.config.hidden_size | |
| # --------------- VAE ------------------------------------------ | |
| vae = AutoencoderKL.from_pretrained( | |
| "black-forest-labs/FLUX.1-dev", | |
| subfolder="vae", | |
| token=HF_TOKEN, | |
| torch_dtype=DTYPE | |
| ).eval().cpu() | |
| # --------------- DiT backbone -------------------------------- | |
| dit_cls = getattr(models, train_args.model) | |
| model = dit_cls(in_channels=16, qk_norm=getattr(train_args, "qk_norm", True), | |
| cap_feat_dim=cap_feat_dim).eval().cpu() | |
| state = torch.load(ckpt_path, map_location="cpu") | |
| state = {k[len("module."):] if k.startswith("module.") else k: v for k, v in state.items()} | |
| model.load_state_dict(state, strict=False) | |
| # --------------- Sampler -------------------------------------- | |
| transport = create_transport("Linear", "velocity", None, None, None) | |
| sampler = Sampler(transport) | |
| print("🔄 Loading models to CPU …") | |
| load_models() | |
| print("✅ Models loaded on CPU") | |
| # ------------------------------------------------------------------ | |
| # INFERENCE (GPU on‑demand) | |
| # ------------------------------------------------------------------ | |
| # << ZeroGPU magic | |
| def generate_image( | |
| prompt, negative_prompt, system_type, solver, resolution_str, | |
| guidance_scale, num_steps, seed, time_shifting_factor, t_shift, | |
| atol=1e-6, rtol=1e-3, | |
| ): | |
| """Runs every time a user clicks “Generate”. | |
| ZeroGPU will attach an A100; GPU is released on return.""" | |
| # ---- pick GPU if present (after decorator) ---- | |
| device = "cuda" if torch.cuda.is_available() else "cpu" | |
| dtype = DTYPE if device == "cuda" else torch.float32 | |
| # ---- seed ---- | |
| seed = random.randint(0, 2**32 - 1) if int(seed) == -1 else int(seed) | |
| torch.manual_seed(seed); np.random.seed(seed); random.seed(seed) | |
| # ---- prompts ---- | |
| sys_prompts = { | |
| "align": "You are an assistant designed to generate high-quality images with the highest degree of image‑text alignment based on textual prompts. <Prompt Start> ", | |
| "base": "You are an assistant designed to generate high-quality images based on user prompts. <Prompt Start> ", | |
| "aesthetics": "You are an assistant designed to generate high-quality images with highest degree of aesthetics based on user prompts. <Prompt Start> ", | |
| "real": "You are an assistant designed to generate superior images with the superior degree of image‑text alignment based on textual prompts or user prompts. <Prompt Start> ", | |
| "4grid": "You are an assistant designed to generate four high-quality images with highest degree of aesthetics arranged in 2x2 grids based on user prompts. <Prompt Start> ", | |
| "tags": "You are an assistant designed to generate high-quality images based on user prompts based on danbooru tags. <Prompt Start> ", | |
| "empty": "", | |
| } | |
| full_prompt = sys_prompts.get(system_type, sys_prompts["base"]) + prompt | |
| full_neg = (sys_prompts.get(system_type, "") + negative_prompt) if negative_prompt else "" | |
| # ---- resolution ---- | |
| w, h = map(int, resolution_str.split("x")); lat_w, lat_h = w // 8, h // 8 | |
| # ---- encode prompts ---- | |
| cap_feats_cpu, cap_mask_cpu = encode_prompt( | |
| [full_prompt, full_neg], text_encoder, tokenizer, device, dtype | |
| ) | |
| # ---------------- SAMPLING ---------------- | |
| model.to(device) | |
| z = torch.randn([1, 16, lat_h, lat_w], device=device, dtype=dtype).repeat(2,1,1,1) | |
| model_kwargs = dict( | |
| cap_feats = cap_feats_cpu.to(device, dtype=dtype), | |
| cap_mask = cap_mask_cpu.to(device), | |
| cfg_scale = guidance_scale, | |
| ); del cap_feats_cpu, cap_mask_cpu | |
| with torch.no_grad(), torch.autocast(device_type="cuda", dtype=dtype): | |
| if solver == "dpm": | |
| _sampler = Sampler(create_transport("Linear", "velocity")) | |
| samples = _sampler.sample_dpm( | |
| model.forward_with_cfg, model_kwargs=model_kwargs | |
| )(z, steps=num_steps, order=2, | |
| skip_type="time_uniform_flow", method="multistep", | |
| flow_shift=time_shifting_factor) | |
| else: | |
| samples = sampler.sample_ode( | |
| sampling_method=solver, num_steps=num_steps, | |
| atol=atol, rtol=rtol, time_shifting_factor=t_shift | |
| )(z, model.forward_with_cfg, **model_kwargs)[-1] | |
| samples = samples[:1] # keep positive branch | |
| # ---------------- DECODE ---------------- | |
| vae.to(device) | |
| with torch.no_grad(), torch.autocast(device_type="cuda", dtype=dtype): | |
| sf, sh = vae.config.scaling_factor, vae.config.shift_factor | |
| img = vae.decode(samples / sf + sh)[0] | |
| vae.to("cpu"); model.to("cpu"); torch.cuda.empty_cache(); gc.collect() | |
| img = ((img.cpu() + 1) / 2).clamp(0,1) | |
| pil = to_pil_image(img[0].float()) | |
| return pil, f"Seed {seed}", seed | |
| # ------------------------------------------------------------------ | |
| # GRADIO UI | |
| # ------------------------------------------------------------------ | |
| with gr.Blocks() as demo: | |
| gr.Markdown("# Lumina 2.0 (ZeroGPU)") | |
| with gr.Row(): | |
| with gr.Column(scale=2): | |
| prompt = gr.Textbox(label="Prompt", value="1girl, kita ikuyo, bocchi the rock!, solo, backlighting, blurry, depth of field, bloom, light particles, transparent, blurry foreground, indoors, upper body, red hair, yellow flower, school uniform, white shirt, hair between eyes, lily \(flower\), floating hair, serafuku, chromatic aberration, white lily, green eyes, looking at viewer, red neckerchief, flower, pink flower, sunlight, day, neckerchief, sailor collar, grey sailor collar, white flower, lens flare abuse, medium hair, holding, closed mouth, one side up, long sleeves, holding bouquet, arms at sides, light smile, shirt, blurry background, bouquet") | |
| negative_prompt = gr.Textbox(label="Negative prompt", value="mutated, worst quality, blurry, bad anatomy, bad hands") | |
| system_type = gr.Dropdown(choices=["align","base","aesthetics","real","4grid","tags","empty"], | |
| value="tags", label="System prompt") | |
| resolution = gr.Dropdown(choices=["1024x1024","1280x768","768x1280","1536x1024","1024x1536"], | |
| value="1024x1024", label="Resolution") | |
| solver = gr.Dropdown(choices=["dpm","euler","midpoint","heun","rk4"], | |
| value="euler", label="Solver") | |
| run_btn = gr.Button("Generate", variant="primary") | |
| with gr.Column(scale=1): | |
| guidance_scale = gr.Slider(1.0, 15.0, step=0.5, value=4.0, label="CFG scale") | |
| num_steps = gr.Slider(10, 200, step=1, value=50, label="Steps") | |
| seed = gr.Number(value=-1, precision=0, label="Seed (-1=random)") | |
| time_shifting_factor = gr.Slider(0.0, 10.0, step=0.1, value=1.0, label="Time‑shift (DPM)") | |
| t_shift = gr.Slider(0, 10, step=1, value=4, label="T‑shift (ODE)") | |
| out_img = gr.Image(label="Output") | |
| out_txt = gr.Textbox(label="Status") | |
| out_seed = gr.Number(label="Seed used", interactive=False) | |
| run_btn.click( | |
| generate_image, | |
| inputs=[prompt,negative_prompt,system_type,solver,resolution, | |
| guidance_scale,num_steps,seed,time_shifting_factor,t_shift], | |
| outputs=[out_img,out_txt,out_seed], | |
| concurrency_limit=1 # helps stay within 1 GPU / user | |
| ) | |
| if __name__ == "__main__": | |
| demo.launch() | |