CIawevy/SD3.5M-TextPecker-SQPA

Model Card for Model ID

This model is trained using Flow-GRPO with LoRA. We provide only the LoRA weights here, so you will need to download the SD 3.5 Medium base model first.

Model Details

Model Sources

• Repository: https://github.com/CIawevy/TextPecker/tree/main
• Paper: https://www.arxiv.org/pdf/2602.XXXXX

Uses

import os
import torch
from diffusers import StableDiffusion3Pipeline
from peft import PeftModel

# Environment variable configuration (consistent with inference code)
os.environ["PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION"] = "python"

def load_model(model_path, lora_path=None):
    """Load SD3.5 pipeline with optional LoRA weights (aligned with inference code)"""
    torch_dtype = torch.bfloat16
    device = "cuda"
    
    # Initialize SD3.5 Pipeline (SD3/3.5 dedicated pipeline)
    pipe = StableDiffusion3Pipeline.from_pretrained(
        model_path,
        torch_dtype=torch_dtype,
    ).to(device)
    
    # Disable safety checker (core configuration from inference code)
    pipe.safety_checker = None
    # Optimize progress bar display (consistent with inference code)
    pipe.set_progress_bar_config(
        position=1,
        disable=False,
        leave=False,
        desc="Timestep",
        dynamic_ncols=True,
    )

    # Load LoRA weights (standard SD3 LoRA loading method)
    if lora_path is not None and os.path.exists(lora_path):
        pipe.transformer = PeftModel.from_pretrained(pipe.transformer, lora_path)
        pipe.transformer.eval()  # Set to inference mode
        print(f"Successfully loaded LoRA weights from: {lora_path}")
    
    return pipe

# Core configuration (aligned with SD3.5 parameters in inference code)
model_id = "stabilityai/stable-diffusion-3.5-medium"
lora_ckpt_path = "CIawevy/SD3.5M-TextPecker-SQPA"
device = "cuda"

# Inference parameters (SD3 standard parameters)
negative_prompt = " "
width, height = 1024, 1024  # Standard resolution for SD3/3.5
num_inference_steps = 50     # Inference steps from reference code
guidance_scale = 3.5         # Guidance scale from reference code

# Load model (simplified logic without redundant parameters)
pipe = load_model(model_id, lora_ckpt_path)

# Generate image (aligned with inference code parameter format)
prompt = 'a weathered cave explorers journal page, with the phrase "TextPecker" prominently written in faded ink, surrounded by sketches of ancient ruins and cryptic symbols, under a dim, mystical light.'
image = pipe(
    prompt=prompt,
    negative_prompt=negative_prompt,
    width=width,
    height=height,
    num_inference_steps=num_inference_steps,
    guidance_scale=guidance_scale,
    generator=torch.Generator(device=device).manual_seed(42)
).images[0]

# Save result (SD3/3.5 naming convention)
image.save("TextPecker_sd35_demo.png")
print("Image saved as: TextPecker_sd35_demo.png")