Upload folder using huggingface_hub

README.md

@@ -1,13 +1,12 @@
 ---
-title: Kroma Krea2 Lora Demo
-emoji: 🐨
-colorFrom: purple
-colorTo: yellow
+title: Kroma Krea 2 LoRA
+emoji: 🎨
+colorFrom: red
+colorTo: indigo
 sdk: gradio
 sdk_version: 6.22.0
-python_version: '3.12'
 app_file: app.py
-pinned: false
----
-
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+short_description: Kroma style LoRA for Krea 2 Turbo text-to-image
+python_version: "3.12"
+startup_duration_timeout: 30m
+---

app.py

@@ -0,0 +1,374 @@
+"""Kroma v0.1 — LoRA for Krea 2 Turbo.
+
+Loads the Kroma LoRA (which includes both rank-256 LoRA adapters and fused
+``.diff`` weight deltas for RMSNorm / modulation tensors) on top of the
+Krea-2-Turbo base pipeline and runs text-to-image inference.
+"""
+
+import os
+
+os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")
+
+import re
+
+import spaces  # MUST be before torch / diffusers
+import torch
+import gradio as gr
+import numpy as np
+from safetensors.torch import load_file
+from huggingface_hub import hf_hub_download
+from diffusers import Krea2Pipeline
+
+BASE_MODEL_ID = "krea/Krea-2-Turbo"
+LORA_REPO_ID = "lodestones/Kroma"
+LORA_FILENAME = "kroma-v0.1.safetensors"
+MAX_SEED = np.iinfo(np.int32).max
+MAX_IMAGE_SIZE = 1536
+
+
+# ---------------------------------------------------------------------------
+# Custom LoRA + .diff loading
+# ---------------------------------------------------------------------------
+# The Kroma safetensors contains two kinds of tensors:
+#   1. Standard LoRA adapters (lora_A / lora_B) — ComfyUI key naming, no ".weight" suffix
+#   2. ".diff" weight deltas for RMSNorm scales and modulation layers
+#
+# diffusers' built-in _convert_non_diffusers_krea2_lora_to_diffusers() expects
+# ".lora_A.weight" / ".lora_B.weight" suffixes and raises ValueError on any
+# remaining (non-LoRA) keys.  We therefore:
+#   - Separate LoRA keys from .diff keys
+#   - Convert LoRA keys to diffusers format (add .weight suffix, remap modules)
+#   - Load LoRA adapters via pipe.load_lora_weights()
+#   - Manually add .diff deltas to the corresponding model parameters
+
+_ATTN_MAP = {"wq": "to_q", "wk": "to_k", "wv": "to_v", "wo": "to_out.0", "gate": "to_gate"}
+_FF_MAP = {"gate": "ff.gate", "up": "ff.up", "down": "ff.down"}
+_STANDALONE_MAP = {
+    "first": "img_in",
+    "last.linear": "final_layer.linear",
+    "tmlp.0": "time_embed.linear_1",
+    "tmlp.2": "time_embed.linear_2",
+    "tproj.1": "time_mod_proj",
+    "txtmlp.1": "txt_in.linear_1",
+    "txtmlp.3": "txt_in.linear_2",
+    "txtfusion.projector": "text_fusion.projector",
+}
+
+
+def _convert_lora_module(module_path: str) -> str | None:
+    """Map a ComfyUI/Krea2 module path to its diffusers equivalent."""
+    m = re.match(r"blocks\.(\d+)\.(attn|mlp)\.(\w+)$", module_path)
+    if m:
+        idx, kind, sub = m.groups()
+        if kind == "attn" and sub in _ATTN_MAP:
+            return f"transformer_blocks.{idx}.attn.{_ATTN_MAP[sub]}"
+        if kind == "mlp" and sub in _FF_MAP:
+            return f"transformer_blocks.{idx}.{_FF_MAP[sub]}"
+        return None
+    m = re.match(r"txtfusion\.(layerwise_blocks|refiner_blocks)\.(\d+)\.(attn|mlp)\.(\w+)$", module_path)
+    if m:
+        block, idx, kind, sub = m.groups()
+        if kind == "attn" and sub in _ATTN_MAP:
+            return f"text_fusion.{block}.{idx}.attn.{_ATTN_MAP[sub]}"
+        if kind == "mlp" and sub in _FF_MAP:
+            return f"text_fusion.{block}.{idx}.{_FF_MAP[sub]}"
+        return None
+    return _STANDALONE_MAP.get(module_path)
+
+
+def _convert_diff_module(module_path: str) -> str | None:
+    """Map a ComfyUI .diff module path to a dotted attribute path in the transformer."""
+    # blocks.N.*  ->  transformer_blocks.N.*
+    m = re.match(r"blocks\.(\d+)\.(.+)$", module_path)
+    if m:
+        idx, rest = m.groups()
+        return f"transformer_blocks.{idx}.{rest}"
+
+    # txtfusion.*  ->  text_fusion.*
+    m = re.match(r"txtfusion\.(.+)$", module_path)
+    if m:
+        return f"text_fusion.{m.group(1)}"
+
+    # last.*  ->  final_layer.*
+    m = re.match(r"last\.(.+)$", module_path)
+    if m:
+        rest = m.group(1)
+        if rest == "norm.scale":
+            return "final_layer.norm.weight"
+        if rest == "modulation.lin":
+            return "final_layer.scale_shift_table"
+        return f"final_layer.{rest}"
+
+    # txtmlp.N  ->  txt_in.linear_{N//2+1} (0,1->linear_1  2,3->linear_2)
+    # Actually txtmlp.0 is txt_in.norm, txtmlp.1 is linear_1, txtmlp.3 is linear_2
+    # But .diff for txtmlp is a scale.diff -> txt_in.norm.weight
+    m = re.match(r"txtmlp\.(\d+)\.scale$", module_path)
+    if m:
+        return "txt_in.norm.weight"
+
+    return None
+
+
+def load_kroma_lora(pipe: Krea2Pipeline, lora_path: str, strength: float = 1.0):
+    """Load Kroma LoRA adapters + .diff deltas into the pipeline."""
+    state_dict = load_file(lora_path)
+
+    # Strip the "diffusion_model." prefix
+    state_dict = {
+        k.removeprefix("diffusion_model."): v for k, v in state_dict.items()
+    }
+
+    lora_state_dict = {}
+    diff_state_dict = {}
+
+    for key, value in state_dict.items():
+        # LoRA keys: end with .lora_A or .lora_B (no .weight suffix)
+        m = re.match(r"^(.+)\.(lora_[AB])$", key)
+        if m:
+            module_path, lora_type = m.groups()
+            diffusers_module = _convert_lora_module(module_path)
+            if diffusers_module is None:
+                print(f"  [Kroma LoRA] Skipping unmapped LoRA key: {key}")
+                continue
+            new_key = f"transformer.{diffusers_module}.{lora_type}.weight"
+            lora_state_dict[new_key] = value
+            continue
+
+        # .diff keys
+        m = re.match(r"^(.+)\.diff$", key)
+        if m:
+            module_path = m.group(1)
+            diffusers_path = _convert_diff_module(module_path)
+            if diffusers_path is None:
+                print(f"  [Kroma LoRA] Skipping unmapped .diff key: {key}")
+                continue
+            diff_state_dict[diffusers_path] = value
+            continue
+
+        print(f"  [Kroma LoRA] Skipping unknown key: {key}")
+
+    # Load LoRA adapters via diffusers' standard loader
+    print(f"  [Kroma LoRA] Loading {len(lora_state_dict)} LoRA tensors")
+    pipe.load_lora_weights(lora_state_dict, adapter_name="kroma")
+    pipe.set_adapters(["kroma"], adapter_weights=[strength])
+
+    # Apply .diff deltas manually by adding to existing parameters
+    transformer = pipe.transformer
+    print(f"  [Kroma LoRA] Applying {len(diff_state_dict)} .diff deltas")
+    for path, delta in diff_state_dict.items():
+        # Navigate the dotted path to the parameter
+        obj = transformer
+        parts = path.split(".")
+        for part in parts[:-1]:
+            idx = int(part) if part.isdigit() else None
+            obj = obj[idx] if idx is not None else getattr(obj, part)
+        param_name = parts[-1]
+        # For numeric indices at the end (e.g. to_out.0)
+        if param_name.isdigit():
+            idx = int(param_name)
+            obj = obj[idx]
+            param_name = "weight"
+
+        param = getattr(obj, param_name)
+        if param.shape != delta.shape:
+            print(
+                f"  [Kroma LoRA] Shape mismatch for {path}: "
+                f"param={param.shape} vs delta={delta.shape}"
+            )
+            continue
+        # Add the delta (scaled by strength) to the parameter
+        param.data.add_(delta.to(param.dtype) * strength)
+        print(f"  [Kroma LoRA] Applied .diff to {path} ({param.shape})")
+
+
+# ---------------------------------------------------------------------------
+# Model loading
+# ---------------------------------------------------------------------------
+print(f"Loading base pipeline from {BASE_MODEL_ID}...")
+pipe = Krea2Pipeline.from_pretrained(BASE_MODEL_ID, torch_dtype=torch.bfloat16)
+pipe.to("cuda")
+
+# Download and load the Kroma LoRA
+lora_path = hf_hub_download(LORA_REPO_ID, LORA_FILENAME)
+print(f"Loading Kroma LoRA from {lora_path}...")
+load_kroma_lora(pipe, lora_path, strength=1.0)
+print("Kroma LoRA loaded successfully.")
+
+
+# ---------------------------------------------------------------------------
+# Inference
+# ---------------------------------------------------------------------------
+@spaces.GPU(duration=60)
+def generate(
+    prompt: str,
+    negative_prompt: str = "",
+    seed: int = 0,
+    randomize_seed: bool = True,
+    width: int = 1024,
+    height: int = 1024,
+    num_inference_steps: int = 8,
+    guidance_scale: float = 0.0,
+    lora_scale: float = 1.0,
+    progress=gr.Progress(track_tqdm=True),
+):
+    """Generate an image from a text prompt using Krea 2 Turbo + Kroma LoRA.
+
+    Args:
+        prompt: What to generate.
+        negative_prompt: What to avoid generating.
+        seed: RNG seed for reproducibility.
+        randomize_seed: If True, pick a random seed each run.
+        width: Output image width in pixels.
+        height: Output image height in pixels.
+        num_inference_steps: Denoising steps (8 for Turbo).
+        guidance_scale: Classifier-free guidance (0.0 for Turbo).
+        lora_scale: Kroma LoRA strength (1.0 = full effect).
+    """
+    import random
+
+    if randomize_seed:
+        seed = random.randint(0, MAX_SEED)
+    seed = int(seed)
+
+    generator = torch.Generator(device="cuda").manual_seed(seed)
+
+    # Adjust LoRA strength at runtime
+    pipe.set_adapters(["kroma"], adapter_weights=[lora_scale])
+
+    image = pipe(
+        prompt=prompt,
+        negative_prompt=negative_prompt if negative_prompt else None,
+        width=width,
+        height=height,
+        num_inference_steps=num_inference_steps,
+        guidance_scale=guidance_scale,
+        generator=generator,
+    ).images[0]
+
+    return image, seed
+
+
+# ---------------------------------------------------------------------------
+# UI
+# ---------------------------------------------------------------------------
+examples = [
+    "A weathered fisherman mending nets on a misty dock at dawn, warm lantern light",
+    "A close-up portrait of a street violinist in the rain, neon reflections in puddles",
+    "Sun-drenched Mediterranean village alley with laundry lines and stray cats",
+]
+
+css = """
+#col-container { max-width: 1100px; margin: 0 auto; }
+.dark .gradio-container { color: var(--body-text-color); }
+"""
+
+with gr.Blocks(theme=gr.themes.Citrus(), css=css) as demo:
+    with gr.Column(elem_id="col-container"):
+        gr.Markdown(
+            """
+            # Kroma v0.1 — LoRA for Krea 2 Turbo
+            A style LoRA fine-tune for [Krea 2](https://huggingface.co/krea/Krea-2-Turbo)
+            by [lodestones](https://huggingface.co/lodestones/Kroma).
+            Runs on ZeroGPU with 8-step Turbo sampling.
+            """
+        )
+
+        with gr.Row():
+            prompt = gr.Text(
+                label="Prompt",
+                show_label=False,
+                max_lines=1,
+                placeholder="Describe the image you want to generate",
+                container=False,
+                scale=4,
+            )
+            run_button = gr.Button("Generate", variant="primary", scale=1)
+
+        result = gr.Image(label="Result", show_label=False)
+
+        with gr.Accordion("Advanced settings", open=False):
+            negative_prompt = gr.Text(
+                label="Negative prompt",
+                max_lines=1,
+                placeholder="Enter a negative prompt (optional)",
+            )
+
+            with gr.Row():
+                seed = gr.Slider(
+                    label="Seed",
+                    minimum=0,
+                    maximum=MAX_SEED,
+                    step=1,
+                    value=0,
+                )
+                randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
+
+            with gr.Row():
+                width = gr.Slider(
+                    label="Width",
+                    minimum=512,
+                    maximum=MAX_IMAGE_SIZE,
+                    step=64,
+                    value=1024,
+                )
+                height = gr.Slider(
+                    label="Height",
+                    minimum=512,
+                    maximum=MAX_IMAGE_SIZE,
+                    step=64,
+                    value=1024,
+                )
+
+            with gr.Row():
+                guidance_scale = gr.Slider(
+                    label="Guidance scale (CFG)",
+                    minimum=0.0,
+                    maximum=10.0,
+                    step=0.1,
+                    value=0.0,
+                )
+                num_inference_steps = gr.Slider(
+                    label="Inference steps",
+                    minimum=1,
+                    maximum=28,
+                    step=1,
+                    value=8,
+                )
+
+            lora_scale = gr.Slider(
+                label="LoRA strength (Kroma)",
+                minimum=0.0,
+                maximum=1.5,
+                step=0.05,
+                value=1.0,
+            )
+
+        gr.Examples(
+            examples=examples,
+            inputs=[prompt],
+            outputs=[result, seed],
+            fn=generate,
+            cache_examples=True,
+            cache_mode="lazy",
+        )
+
+    gr.on(
+        triggers=[run_button.click, prompt.submit],
+        fn=generate,
+        inputs=[
+            prompt,
+            negative_prompt,
+            seed,
+            randomize_seed,
+            width,
+            height,
+            num_inference_steps,
+            guidance_scale,
+            lora_scale,
+        ],
+        outputs=[result, seed],
+        api_name="generate",
+    )
+
+if __name__ == "__main__":
+    demo.launch(mcp_server=True)

requirements.txt

@@ -0,0 +1,5 @@
+diffusers @ git+https://github.com/huggingface/diffusers.git
+transformers
+accelerate
+safetensors
+sentencepiece