Update README.md

README.md

@@ -38,3 +38,147 @@ This will return pure albedo from your texture.
 
 
 [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)
+```