from __future__ import annotations from __future__ import annotations from typing import Iterable import gradio as gr from gradio.themes.base import Base from gradio.themes.utils import colors, fonts, sizes import time class Seafoam(Base): def __init__( self, *, primary_hue: colors.Color | str = colors.emerald, secondary_hue: colors.Color | str = colors.blue, neutral_hue: colors.Color | str = colors.gray, spacing_size: sizes.Size | str = sizes.spacing_md, radius_size: sizes.Size | str = sizes.radius_md, text_size: sizes.Size | str = sizes.text_lg, font: fonts.Font | str | Iterable[fonts.Font | str] = ( fonts.GoogleFont("Quicksand"), "ui-sans-serif", "sans-serif", ), font_mono: fonts.Font | str | Iterable[fonts.Font | str] = ( fonts.GoogleFont("IBM Plex Mono"), "ui-monospace", "monospace", ), ): super().__init__( primary_hue=primary_hue, secondary_hue=secondary_hue, neutral_hue=neutral_hue, spacing_size=spacing_size, radius_size=radius_size, text_size=text_size, font=font, font_mono=font_mono, ) seafoam = Seafoam() import spaces import os import random import gradio as gr import numpy as np import torch from huggingface_hub import hf_hub_download from PIL import Image from transformers import AutoConfig, AutoModel, AutoTokenizer from safetensors.torch import safe_open import sensenova_u1 from sensenova_u1 import check_checkpoint_compatibility from sensenova_u1.models.neo_unify.utils import smart_resize from sensenova_u1.utils.lora import build_lora_names from huggingface_hub import snapshot_download import os MODEL_ID=snapshot_download(repo_id="Tele-AI/TeleStyleV2",allow_patterns="TeleStyle_SenseNova/*") MODEL_ID=os.path.join(MODEL_ID,"TeleStyle_SenseNova") LORA_REPO = "sensenova/SenseNova-U1-8B-MoT-LoRAs" LORA_FILE = "SenseNova-U1-8B-MoT-LoRA-8step-V1.0.safetensors" FAST_MODE = "Fast (8-step LoRA)" QUALITY_MODE = "Quality (50-step base)" NORM_MEAN = (0.5, 0.5, 0.5) NORM_STD = (0.5, 0.5, 0.5) # Trained T2I aspect-ratio buckets: aspect_label -> (W, H) T2I_RESOLUTIONS: dict[str, tuple[int, int]] = { "1:1": (2048, 2048), "16:9": (2720, 1536), "9:16": (1536, 2720), "3:2": (2496, 1664), "2:3": (1664, 2496), "4:3": (2368, 1760), "3:4": (1760, 2368), } # Editing output grid factor (= patch_size * merge_size = 32). EDIT_GRID_FACTOR = 32 EDIT_TARGET_PIXELS = 1024*1024#2048 * 2048 EDIT_INPUT_MAX_PIXELS = 2048 * 2048 MAX_SEED = 2**31 - 1 def _denorm(x: torch.Tensor) -> torch.Tensor: mean = torch.tensor(NORM_MEAN, device=x.device, dtype=x.dtype).view(1, 3, 1, 1) std = torch.tensor(NORM_STD, device=x.device, dtype=x.dtype).view(1, 3, 1, 1) return (x * std + mean).clamp(0, 1) def _to_pil(batch: torch.Tensor) -> list[Image.Image]: arr = _denorm(batch.float()).permute(0, 2, 3, 1).cpu().numpy() arr = (arr * 255.0).round().astype(np.uint8) return [Image.fromarray(a) for a in arr] def _coerce_pil(img) -> Image.Image: if isinstance(img, Image.Image): return img if isinstance(img, tuple): img = img[0] if isinstance(img, str): return Image.open(img) return img def _prep_input_image(img: Image.Image, max_pixels: int) -> Image.Image: if img.mode == "RGBA": bg = Image.new("RGB", img.size, (255, 255, 255)) bg.paste(img, mask=img.split()[3]) img = bg img = img.convert("RGB") h, w = smart_resize( height=img.height, width=img.width, factor=EDIT_GRID_FACTOR, min_pixels=max_pixels, max_pixels=max_pixels, ) if (w, h) != img.size: img = img.resize((w, h), Image.LANCZOS) print(f"input_img.size={img.size}") return img def _editing_output_size(input_img: Image.Image, target_pixels: int) -> tuple[int, int]: ''' h, w = smart_resize( height=input_img.height, width=input_img.width, factor=EDIT_GRID_FACTOR, min_pixels=target_pixels, max_pixels=target_pixels, ) ''' w,h=input_img.size print(f"output_img.size={w,h}") return w, h print("[startup] loading SenseNova-U1-8B-MoT (this may take a few minutes)...") sensenova_u1.set_attn_backend("auto") print(f"[startup] attn backend: {sensenova_u1.effective_attn_backend()!r}") print(f"Model_id={MODEL_ID}") print(os.listdir(MODEL_ID)) config = AutoConfig.from_pretrained(MODEL_ID) check_checkpoint_compatibility(config) tokenizer = AutoTokenizer.from_pretrained(MODEL_ID) model = AutoModel.from_pretrained(MODEL_ID, config=config, torch_dtype=torch.bfloat16).to("cuda").eval() print(f"[startup] downloading 8-step LoRA from {LORA_REPO}/{LORA_FILE}") lora_path = hf_hub_download( repo_id=LORA_REPO, filename=LORA_FILE, cache_dir=os.environ.get("HF_HOME"), ) def _build_lora_deltas(model, lora_weight_path: str) -> dict[str, torch.Tensor]: """Precompute delta_W = (alpha/rank) * (up @ down) for each affected param. Stored on CPU in fp32 so the toggle is a symmetric add/subtract that round-trips cleanly through bf16 quantization. """ lora_state_dict: dict[str, torch.Tensor] = {} with safe_open(lora_weight_path, framework="pt", device="cpu") as f: for key in f.keys(): lora_state_dict[key] = f.get_tensor(key) is_native = any("diffusion_model." in k for k in lora_state_dict) deltas: dict[str, torch.Tensor] = {} for name, param in model.named_parameters(): down_n, up_n, alpha_n = build_lora_names( name, ".lora_down.weight", ".lora_up.weight", is_native ) if down_n not in lora_state_dict: continue down = lora_state_dict[down_n].to(param.device, dtype=torch.float32) up = lora_state_dict[up_n].to(param.device, dtype=torch.float32) alpha = float(lora_state_dict[alpha_n]) scaling = alpha / down.shape[0] deltas[name] = (scaling * (up @ down)).detach().to("cpu", dtype=torch.float32) return deltas print(f"[startup] precomputing LoRA deltas from {lora_path}") LORA_DELTAS = _build_lora_deltas(model, lora_path) _LORA_ACTIVE = False def _set_lora_active(active: bool) -> None: global _LORA_ACTIVE if active == _LORA_ACTIVE: print(f"[lora] already {'ON' if active else 'OFF'} — no-op") return sign = 1.0 if active else -1.0 touched = 0 for name, param in model.named_parameters(): delta = LORA_DELTAS.get(name) if delta is None: continue d = delta.to(param.device, dtype=torch.float32, non_blocking=True) param.data = (param.data.float() + sign * d).to(param.dtype) touched += 1 _LORA_ACTIVE = active print(f"[lora] toggled to {'ON' if active else 'OFF'} ({touched} params updated)") # Default to Fast mode so the first call is snappy. _set_lora_active(True) # Pick one LoRA-affected param to fingerprint per generate call. Its sum changes # by exactly the delta when LoRA is toggled, so the per-call log makes the # active/inactive state unambiguous in the Space logs. _LORA_PROBE_NAME = next(iter(LORA_DELTAS), None) print(f"[startup] LoRA probe param: {_LORA_PROBE_NAME!r} ({len(LORA_DELTAS)} params have LoRA deltas)") print("[startup] model ready.") def _probe_weight() -> float: if _LORA_PROBE_NAME is None: return float("nan") p = dict(model.named_parameters())[_LORA_PROBE_NAME] return float(p.detach().float().sum().item()) def get_duration(images, prompt, mode, aspect_ratio, seed, randomize_seed): return 45 if mode == FAST_MODE else 95 @spaces.GPU(size="xlarge") def generate( content_ref,style_ref, prompt: str, mode: str, aspect_ratio: str, seed: int, randomize_seed: bool, resolution: int, progress=gr.Progress(track_tqdm=True), ): EDIT_TARGET_PIXELS= resolution*resolution if not prompt or not prompt.strip(): raise gr.Error("Please enter a prompt.") if randomize_seed: seed = random.randint(0, MAX_SEED) use_lora = mode == FAST_MODE print(f"[generate] mode={mode!r}, use_lora={use_lora}, _LORA_ACTIVE(before)={_LORA_ACTIVE}, probe(before)={_probe_weight():.6f}") _set_lora_active(use_lora) num_steps = 8 if use_lora else 50 cfg_scale = 1.0 if use_lora else 4.0 print(f"[generate] _LORA_ACTIVE(after)={_LORA_ACTIVE}, probe(after)={_probe_weight():.6f}, num_steps={num_steps}, cfg_scale={cfg_scale}") images=[] if content_ref is not None: images.append(Image.fromarray(content_ref)) if style_ref is not None: images.append(Image.fromarray(style_ref)) #images=[Image.fromarray(content_ref),Image.fromarray(style_ref)] has_input = images is not None and len(images) > 0 with torch.inference_mode(): if not has_input: width, height = T2I_RESOLUTIONS[aspect_ratio] tensor = model.t2i_generate( tokenizer, prompt, image_size=(width, height), cfg_scale=cfg_scale, cfg_norm="none", timestep_shift=3.0, cfg_interval=(0.0, 1.0), num_steps=num_steps, batch_size=1, seed=int(seed), think_mode=False, ) else: pil_inputs = [_prep_input_image(_coerce_pil(item), EDIT_TARGET_PIXELS) for item in images] out_w, out_h = _editing_output_size(pil_inputs[0], EDIT_TARGET_PIXELS) tensor = model.it2i_generate( tokenizer, prompt, pil_inputs, image_size=(out_w, out_h), cfg_scale=cfg_scale, img_cfg_scale=1.0, cfg_norm="none", timestep_shift=3.0, cfg_interval=(0.0, 1.0), num_steps=num_steps, batch_size=1, think_mode=False, seed=int(seed), ) images_out = _to_pil(tensor) return images_out[0], seed EXAMPLES = [] CSS = """ .fillable { max-width: 960px !important; } """ with gr.Blocks(title="SenseNova-U1-8B-MoT (8-step LoRA)") as demo: gr.Markdown( """ # TeleStyle-SenseNova-U1-8B-MoT This model reinforces SenseNova U1 for Content-Preserving Style Transfer and preserves its general image editing capability. This experimental model is trained in pixel space, making the sft quite hard and still having much space to improve. The model supports 1MP to 4MP. Paper:[TeleStyle V2: Beyond Content-Preserving StyleTransfer with Self-Distillation and Distribution-Matching-Distillation](https://arxiv.org/abs/2606.20709) | Codes: [Github](https://github.com/Tele-AI/TeleStyleV2) If you find the model and demo useful, please light a star for the [project](https://github.com/Tele-AI/TeleStyleV2), thanks """ ) with gr.Row(): with gr.Column(scale=1): with gr.Column(): with gr.Row(): content_ref = gr.Image(label="content ref", type="numpy", ) style_ref = gr.Image(label="style ref", type="numpy", ) # #print(f"type(content_ref)={type(content_ref)}") # image_input_gallery = gr.Gallery( # label="Upload one or more images here. Leave empty for text-to-image.", # file_types=["image"], # height="auto", # columns=4, # ) prompt_input = gr.Textbox( label="Prompt", placeholder="Describe the image to generate, or how to edit your input.", value='style transfer the style of Image 2 to Image 1, and keep the content of Image 1, ignore the content of Image 2', #'style transfer the style of Image 2 to Image 1, and keep the content of Image 1', lines=3, ) mode = gr.Radio( label="Mode", choices=[FAST_MODE, QUALITY_MODE], value=FAST_MODE, ) resolution = gr.Slider( label="The total pixel number of the generated image is the square of this value, default is 1536x1536", minimum=256, maximum=2048, step=8, value=1536, ) aspect_ratio = gr.Dropdown( label="Aspect ratio (text-to-image only — editing keeps input ratio)", choices=list(T2I_RESOLUTIONS.keys()), value="1:1", ) with gr.Row(): seed = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=42) randomize_seed = gr.Checkbox(label="Randomize seed", value=True) generate_button = gr.Button("Generate", variant="primary") with gr.Column(scale=1): output_image = gr.Image(label="Output", type="pil", format="png", interactive=False) used_seed = gr.Number(label="Seed used", interactive=False) gr.Examples( examples=EXAMPLES, inputs=[content_ref,style_ref, prompt_input, mode, aspect_ratio], cache_examples=False, ) generate_button.click( fn=generate, inputs=[content_ref,style_ref, prompt_input, mode, aspect_ratio, seed, randomize_seed,resolution], outputs=[output_image, used_seed], ) prompt_input.submit( fn=generate, inputs=[content_ref,style_ref, prompt_input, mode, aspect_ratio, seed, randomize_seed,resolution], outputs=[output_image, used_seed], ) if __name__ == "__main__": demo.launch(server_name='0.0.0.0',theme=seafoam, css=CSS)