texture2albedo-v2

use this prompt:

Unlit flat-shaded albedo map. Remove all shadows, reflections, highlights, and specularity. Maintain absolute pixel-per-pixel structural identity, shape, and spatial alignment with the original image, displaying only raw base color.

texture2albedo

<Gallery />

Model description

This will return pure albedo from your texture.

Download model

Download them in the Files & versions tab.

'''

Python script for inference in gradio (install gradio in python with 'pip install gradio')

import os
import torch
import gradio as gr
from PIL import Image
from diffusers import Flux2KleinPipeline

# --- Configuration & Initialization ---

BASE_MODEL_FILE = "black-forest-labs/FLUX.2-klein-9B"
LORA_REPO = "paom/texture2albedo-v2"

print("Initializing device and pipeline...")
device = "cuda" if torch.cuda.is_available() else "cpu"
dtype = torch.bfloat16 if torch.cuda.is_available() else torch.float32

try:
    print(f"Loading transformer component from single file: {BASE_MODEL_FILE}")
 

    pipe = Flux2KleinPipeline.from_pretrained(
        BASE_MODEL_FILE, 
        torch_dtype=dtype
    )
    
   
    pipe.load_lora_weights(
        LORA_REPO, 
        weight_name="pytorch_lora_weights.safetensors", 
        adapter_name="albedo"
    )
    
    
    if device == "cuda":
        print("Enabling smart CPU offload...")
        pipe.enable_model_cpu_offload()
    else:
        pipe.to(device)
        
    print("Pipeline and LoRA weights loaded successfully.")
except Exception as e:
    import traceback
    print("!!! DETAILED INITIALIZATION ERROR !!!")
    traceback.print_exc()
    pipe = None

# --- Prompt Presets ---
PROMPT_PRESETS = {
    "Strict Unlit Flat (Default)": (
        "Unlit flat-shaded albedo map. Remove all shadows, reflections, highlights, and specularity. "
        "Maintain absolute pixel-per-pixel structural identity, shape, and spatial alignment with the "
        "original image, displaying only raw base color."
    )
}

# --- Core Inference Function ---
def generate_albedo(input_image, prompt_selection, custom_prompt, steps, guidance_scale, seed):
    if pipe is None:
        raise gr.Error("Model pipeline failed to initialize. Check your hardware compatibility.")
    
    if input_image is None:
        return None

   
    prompt = custom_prompt if custom_prompt.strip() else PROMPT_PRESETS[prompt_selection]
    
   
    orig_width, orig_height = input_image.size
    
   
    processed_input = input_image.resize((1024, 1024))
    
   
    generator = torch.manual_seed(seed) if seed >= 0 else None

    try:
       
        with torch.inference_mode():
            output_image = pipe(
                prompt=prompt,
                image=processed_input,
                guidance_scale=guidance_scale,
                num_inference_steps=int(steps),
                generator=generator
            ).images[0]
        
        
        albedo_map = output_image.resize((orig_width, orig_height))
        
        return albedo_map

    except Exception as e:
        raise gr.Error(f"Inference error occurred: {str(e)}")


with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown(
        """
        # Texture-to-Albedo Studio (Flux.2 Klein)
        Extract clean, flat, completely shadowless base color **Albedo maps** from textures and photos for your 3D/PBR pipelines.
        """
    )
    
    with gr.Row():
  
        with gr.Column(scale=1):
            input_img = gr.Image(label="Input Texture / Photo", type="pil")
            
            prompt_dropdown = gr.Dropdown(
                choices=list(PROMPT_PRESETS.keys()),
                value="Strict Unlit Flat (Default)",
                label="Prompt Style Preset"
            )
            
            custom_prompt_box = gr.Textbox(
                label="Custom Prompt Override", 
                placeholder="Leave blank to use chosen preset above...", 
                lines=2
            )
            
            with gr.Accordion("Advanced Parameters", open=False):
                inference_steps = gr.Slider(minimum=1, maximum=12, value=4, step=1, label="Inference Steps")
                guidance = gr.Slider(minimum=0.0, maximum=4.0, value=1.0, step=0.1, label="Guidance Scale")
                seed_input = gr.Number(value=0, label="Seed (-1 for random)", precision=0)
                
            submit_btn = gr.Button("Generate Albedo Map", variant="primary")
            
       
        with gr.Column(scale=1):
            albedo_out = gr.Image(label="Clean Albedo Texture Map", type="pil")
                
    
    submit_btn.click(
        fn=generate_albedo,
        inputs=[input_img, prompt_dropdown, custom_prompt_box, inference_steps, guidance, seed_input],
        outputs=[albedo_out]
    )

if __name__ == "__main__":
    demo.queue().launch(server_name="0.0.0.0", server_port=7860, share=False)