Views
No views yet
1import torch
2from diffusers import AutoencoderKL, SD3Transformer2DModel
3from transformers import CLIPTokenizer, T5TokenizerFast
4from PIL import Image
5import numpy as np
6
7# Load base model components
8base_model = "stabilityai/stable-diffusion-3.5-medium"
9vae = AutoencoderKL.from_pretrained(base_model, subfolder="vae")
10
11# Load trained transformer
12transformer = SD3Transformer2DModel.from_pretrained(
13 "{model_id}",
14 subfolder="transformer",
15 torch_dtype=torch.bfloat16
16)
17
18# Load your custom pipeline (from training repo)
19from pipelines.pipeline_stable_diffusion_3 import StableDiffusion3Pipeline
20
21pipeline = StableDiffusion3Pipeline.from_pretrained(
22 base_model,
23 transformer=transformer,
24 vae=vae,
25 torch_dtype=torch.bfloat16,
26)
27pipeline.to("cuda")
28
29# Load and prepare albedo image
30albedo_image = Image.open("path/to/albedo.png").convert("RGB")
31albedo_image = albedo_image.resize((512, 512))
32
33# Convert to tensor and normalize
34albedo_np = np.array(albedo_image).astype(np.float32) / 255.0
35albedo_tensor = torch.from_numpy(albedo_np).permute(2, 0, 1) * 2.0 - 1.0
36albedo_tensor = albedo_tensor.unsqueeze(0).unsqueeze(0).to("cuda", dtype=torch.bfloat16)
37
38# Encode albedo to control latents
39from light_utils import encode_intrinsics
40control_latents = encode_intrinsics(albedo_tensor, vae, torch.bfloat16)
41
42# Generate
43prompt = "A beautiful landscape, soft golden hour lighting"
44image = pipeline(
45 prompt=prompt,
46 control_image=control_latents,
47 num_inference_steps=50,
48 guidance_scale=7.5,
49 height=512,
50 width=512,
51).images[0]
52
53image.save("output.png")1# Different lighting conditions
2prompts = [
3 "A forest scene, at sunrise",
4 "A forest scene, with fluorescent blue lighting",
5]
6
7for prompt in prompts:
8 image = pipeline(
9 prompt=prompt,
10 control_image=control_latents,
11 num_inference_steps=50,
12 ).images[0]
13 # Each will have different lighting/mood