brandstyle-flux-lora · Hugging Face Space

# app.py — brandstyle-flux-lora Space · CULTIVA IA
# Base: black-forest-labs/FLUX.1-schnell + LoRA: cultivaia/brandstyle-flux-lora
# Task: text-to-image · ZeroGPU · Gradio 6.x

import os
import torch
import gradio as gr
import spaces
from diffusers import FluxPipeline

# ── Constants ──────────────────────────────────────────────────────────────
LORA_REPO   = "cultivaia/brandstyle-flux-lora"
BASE_MODEL  = "black-forest-labs/FLUX.1-schnell"
TRIGGER     = "BRNDSTY"
STEPS       = 4        # schnell: few-step, no guidance
GUIDANCE    = 0.0
LORA_SCALE  = 0.85
MAX_SEQ_LEN = 256

ASPECT_RATIOS = {
    "1:1  (1024×1024)":  (1024, 1024),
    "4:3  (1024×768)":   (1024, 768),
    "3:2  (1024×688)":   (1024, 688),
    "16:9 (1024×576)":   (1024, 576),
    "9:16 (576×1024)":   (576,  1024),
    "2:3  (688×1024)":   (688,  1024),
}

FINISH_PROMPTS = {
    "Limpio / minimalista": "clean minimal, white space, editorial",
    "Oscuro / premium":    "dark matte premium, dramatic lighting",
    "Natural / orgánico":  "earthy natural textures, soft light",
}

# ── Module-level model load (ZeroGPU: must be at module level) ──────────────
pipe = FluxPipeline.from_pretrained(
    BASE_MODEL,
    torch_dtype=torch.bfloat16,
)
pipe.load_lora_weights(
    LORA_REPO,
    weight_name="pytorch_lora_weights.safetensors",
)
pipe.to("cuda")

# ── Inference ───────────────────────────────────────────────────────────────
@spaces.GPU(duration=60)
def generate(
    prompt: str,
    aspect_ratio: str,
    finish_style: str,
    seed: int,
    randomize_seed: bool,
    lora_scale: float,
    progress=gr.Progress(track_tqdm=True),
):
    if not prompt.strip():
        raise gr.Error("Escribe una descripción de la marca antes de generar.")

    if randomize_seed:
        seed = torch.randint(0, 2**32, ()).item()

    w, h = ASPECT_RATIOS[aspect_ratio]
    style_suffix = FINISH_PROMPTS.get(finish_style, "")
    full_prompt = f"{TRIGGER} {prompt.strip()}, {style_suffix}"

    pipe.set_adapters(["default"], adapter_weights=[lora_scale])

    generator = torch.Generator("cuda").manual_seed(seed)
    result = pipe(
        prompt=full_prompt,
        width=w, height=h,
        num_inference_steps=STEPS,
        guidance_scale=GUIDANCE,
        max_sequence_length=MAX_SEQ_LEN,
        generator=generator,
    ).images[0]

    return result, seed

# ── Gradio UI ────────────────────────────────────────────────────────────────
CSS = """
.generate-btn { font-size: 15px !important; }
.prompt-box textarea { font-size: 14px !important; }
"""

with gr.Blocks(title="BrandStyle LoRA — CULTIVA IA") as demo:
    gr.Markdown("""
# 🎨 BrandStyle LoRA · Imágenes de Marca de Lujo Sostenible
Genera identidades visuales con estética de lujo sostenible. Soportado por **CULTIVA IA**.
*El trigger `BRNDSTY` se añade automáticamente — escribe solo la descripción de la marca.*
""")

    with gr.Row(equal_height=True):
        with gr.Column(scale=1):
            prompt = gr.Textbox(
                label="Descripción de la marca",
                placeholder="luxury sustainable skincare brand identity, sage green and whites",
                lines=3, elem_classes=["prompt-box"],
            )
            aspect_ratio = gr.Dropdown(
                list(ASPECT_RATIOS.keys()),
                value="1:1  (1024×1024)",
                label="Ratio",
            )
            finish_style = gr.Radio(
                list(FINISH_PROMPTS.keys()),
                value="Limpio / minimalista",
                label="Estilo de acabado",
            )
            run_btn = gr.Button(
                "Generar imagen", variant="primary",
                size="lg", elem_classes=["generate-btn"],
            )
            with gr.Accordion("Avanzado", open=False):
                seed_slider = gr.Slider(0, 2**32-1, value=42, label="Seed", step=1)
                rand_seed   = gr.Checkbox(True, label="Aleatorizar seed")
                lora_scale  = gr.Slider(0.0, 1.0, value=LORA_SCALE, label="LoRA scale", step=0.05)

        with gr.Column(scale=1):
            output_img  = gr.Image(label="Resultado", type="pil")
            used_seed   = gr.Number(label="Seed usada (para reproducir)", interactive=False)

    gr.Examples(
        examples=[
            ["luxury sustainable skincare brand identity, sage green and whites",   "1:1  (1024×1024)", "Limpio / minimalista"],
            ["premium coffee packaging for specialty roastery, matte black and gold", "4:3  (1024×768)",  "Oscuro / premium"],
            ["eco-friendly fashion lookbook photography, natural textures",           "9:16 (576×1024)",  "Natural / orgánico"],
        ],
        inputs=[prompt, aspect_ratio, finish_style],
        cache_examples=True, cache_mode="lazy",
    )

    run_btn.click(
        generate,
        inputs=[prompt, aspect_ratio, finish_style, seed_slider, rand_seed, lora_scale],
        outputs=[output_img, used_seed],
    )
    prompt.submit(
        generate,
        inputs=[prompt, aspect_ratio, finish_style, seed_slider, rand_seed, lora_scale],
        outputs=[output_img, used_seed],
    )

if __name__ == "__main__":
    demo.launch(theme=gr.themes.Citrus(), css=CSS)