File size: 2,962 Bytes
2b2a90e dff845d 021cffb 2101ee0 c68da04 2b2a90e ddef426 021cffb dff845d ddef426 021cffb ddef426 021cffb ddef426 021cffb dff845d ddef426 2101ee0 73cfa52 021cffb 2101ee0 06da073 73cfa52 06da073 73cfa52 06da073 73cfa52 ddef426 021cffb ddef426 021cffb ddef426 021cffb ddef426 021cffb ddef426 021cffb ddef426 dff845d 2101ee0 18b2b4b 2101ee0 021cffb ddef426 dff845d ddef426 c68da04 ddef426 dff845d ddef426 021cffb ddef426 021cffb ddef426 021cffb 2b2a90e 2101ee0 021cffb ddef426 3a64eb8 021cffb 18b2b4b 021cffb 2b2a90e 021cffb |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 |
import gradio as gr
import torch
import numpy as np
from diffusers import StableDiffusionXLPipeline
from transformers import DPTFeatureExtractor, DPTForDepthEstimation
from PIL import Image, ImageEnhance, ImageOps
device = "cpu" # or "cuda" if you have a GPU
torch_dtype = torch.float32
print("Loading SDXL Base model...")
pipe = StableDiffusionXLPipeline.from_pretrained(
"stabilityai/stable-diffusion-xl-base-1.0",
torch_dtype=torch_dtype
).to(device)
print("Loading bas-relief LoRA weights with PEFT...")
pipe.load_lora_weights(
"KappaNeuro/bas-relief", # The HF repo with BAS-RELIEF.safetensors
weight_name="BAS-RELIEF.safetensors",
peft_backend="peft" # This is crucial
)
print("Loading DPT Depth Model...")
feature_extractor = DPTFeatureExtractor.from_pretrained("Intel/dpt-large")
depth_model = DPTForDepthEstimation.from_pretrained("Intel/dpt-large").to(device)
def enhance_depth_map(depth_arr: np.ndarray) -> Image.Image:
d_min, d_max = depth_arr.min(), depth_arr.max()
depth_stretched = (depth_arr - d_min) / (d_max - d_min + 1e-8)
depth_stretched = (depth_stretched * 255).astype(np.uint8)
depth_pil = Image.fromarray(depth_stretched)
depth_pil = ImageOps.autocontrast(depth_pil)
enhancer = ImageEnhance.Sharpness(depth_pil)
depth_pil = enhancer.enhance(2.0)
return depth_pil
def generate_bas_relief_and_depth(prompt):
# Use the token "BAS-RELIEF" so the LoRA triggers
full_prompt = f"BAS-RELIEF {prompt}"
print("Generating image with LoRA style...")
result = pipe(
prompt=full_prompt,
num_inference_steps=15, # reduce if too slow
guidance_scale=7.5,
height=512, # reduce if you still get timeouts
width=512
)
image = result.images[0]
print("Running DPT Depth Estimation...")
inputs = feature_extractor(image, return_tensors="pt").to(device)
with torch.no_grad():
outputs = depth_model(**inputs)
predicted_depth = outputs.predicted_depth
prediction = torch.nn.functional.interpolate(
predicted_depth.unsqueeze(1),
size=image.size[::-1],
mode="bicubic",
align_corners=False
).squeeze()
depth_map_pil = enhance_depth_map(prediction.cpu().numpy())
return image, depth_map_pil
title = "Bas-Relief (SDXL + LoRA) + Depth Map (with PEFT)"
description = (
"Loads stable-diffusion-xl-base-1.0 on CPU, merges LoRA from 'KappaNeuro/bas-relief'. "
"Use 'BAS-RELIEF' token in your prompt to trigger the style, then compute a depth map."
)
iface = gr.Interface(
fn=generate_bas_relief_and_depth,
inputs=gr.Textbox(
label="Description",
placeholder="bas-relief with roman soldier, marble relief, intricately carved"
),
outputs=[gr.Image(label="Bas-Relief Image"), gr.Image(label="Depth Map")],
title=title,
description=description
)
if __name__ == "__main__":
iface.launch()
|