app.py

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)