--- 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 ## 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) ```
