README.md

---
tags:
- image-to-image
- lora
- diffusers
- template:diffusion-lora
widget:
- output:
   url: 'd (1).jpg'
  text: '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.
'
- output:
   url: 'd (14).jpg'
  text: '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.
'
- output:
   url: 'd (16).jpg'
  text: '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.
'
base_model: black-forest-labs/FLUX.2-klein-9B
instance_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."
license: apache-2.0
---
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](/paom/texture2albedo-v2/tree/main) 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)
```