ysn-rfd/text-dataset-tiny-code-script-py-format

1"""
2sd_directml_tiled.py
3Tiled diffusion + VAE tiling + partial offload for torch-directml
4Requires:
5  pip install torch-directml diffusers transformers accelerate safetensors pillow
6Usage: just run. Edit MODEL_PATH and PROMPT below.
7"""
8import os
9import time
10import math
11import torch
12import torch_directml
13from diffusers import StableDiffusionPipeline, LCMScheduler
14from PIL import Image
15import numpy as np
16
17# -------------------------
18# User config (edit these)
19# -------------------------
20MODEL_PATH = r"C:\Users\ysnrfd\Desktop\torch_directml\ds_lcm.safetensors"
21OUT_IMAGE = "out_tiled_directml.png"
22PROMPT = "A cinematic landscape, sunrise, ultra-detailed"
23WIDTH = 768   # try 512/384/256 for lower VRAM; tiled helps but base size still matters
24HEIGHT = 512
25STEPS = 20
26GUIDANCE = 1
27SEED = 1234
28
29# Tile config (image-space pixels). We will compute latent tile size by dividing by vae_scale.
30# Keep tile_px reasonably small for lower memory (e.g., 256 or 128)
31TILE_PX = 256
32TILE_OVERLAP_PX = 64  # overlap to blend tiles and avoid seams
33
34# Precision and devices
35DML_DEVICE_NAME = torch_directml.default_device()
36DML_DEVICE = torch_directml.device(DML_DEVICE_NAME)
37CPU_DEVICE = torch.device("cpu")
38USE_FP16 = True  # try True first on DirectML
39
40# -------------------------
41# Helpers
42# -------------------------
43def get_latent_downscale(pipe):
44    """
45    Determine the downscale factor from image to latents (usually 8 for SD: image/8 = latent spatial dims)
46    We'll infer by comparing vae configs if possible, otherwise default 8.
47    """
48    try:
49        # many diffusers VAE have config with 'scaling_factor' or rely on encoder/decoder shapes; fallback to 8
50        return 8
51    except Exception:
52        return 8
53
54def to_device_partial(pipe, unet_device):
55    """
56    Move UNet to GPU (DirectML), move vae and text_encoder to CPU.
57    """
58    try:
59        pipe.unet.to(unet_device)
60    except Exception as e:
61        print("⚠️ couldn't move UNet to DML via .to():", e)
62    # put heavy but tolerable parts to CPU
63    try:
64        pipe.vae.to(CPU_DEVICE)
65    except Exception:
66        pass
67    try:
68        if hasattr(pipe, "text_encoder"):
69            pipe.text_encoder.to(CPU_DEVICE)
70    except Exception:
71        pass
72
73def apply_optimizations(pipe):
74    """
75    common diffusers optimizations
76    """
77    try:
78        if hasattr(pipe, "enable_attention_slicing"):
79            pipe.enable_attention_slicing()
80            print("🔧 attention slicing enabled")
81        if hasattr(pipe, "vae") and hasattr(pipe.vae, "enable_slicing"):
82            pipe.vae.enable_slicing()
83            print("🔧 VAE slicing enabled")
84        if hasattr(pipe, "set_progress_bar_config"):
85            pipe.set_progress_bar_config(disable=True)
86        # channels-last attempt
87        def ch_last(mod):
88            try:
89                for p in mod.parameters():
90                    if p.ndim == 4:
91                        p.data = p.data.contiguous(memory_format=torch.channels_last)
92            except Exception:
93                pass
94        for sub in [pipe.unet, getattr(pipe, "vae", None), getattr(pipe, "text_encoder", None)]:
95            if sub is not None:
96                ch_last(sub)
97    except Exception as e:
98        print("⚠️ apply_optimizations error:", e)
99
100def make_prompt_embedding(pipe, prompt, device):
101    """
102    Encode prompt into text embeddings on CPU (since text_encoder is on CPU).
103    Returns encoder_hidden_states on CPU or moved to DML only when UNet needs it.
104    """
105    # Use pipeline's tokenizer + text_encoder
106    tokenizer = pipe.tokenizer
107    text_encoder = pipe.text_encoder
108    inputs = tokenizer(prompt, padding="max_length", truncation=True, max_length=tokenizer.model_max_length, return_tensors="pt")
109    input_ids = inputs.input_ids.to(CPU_DEVICE)
110    with torch.no_grad():
111        # text_encoder is on CPU_device
112        encoder_hidden_states = text_encoder(input_ids)[0].to(CPU_DEVICE)
113    return encoder_hidden_states
114
115# -------------------------
116# tiled UNet inference helpers
117# -------------------------
118def slice_latent_to_tiles(latents, tile_px, overlap_px, vae_scale):
119    """
120    latents: tensor (b, c, H_lat, W_lat) in latent resolution
121    tile_px, overlap_px: in image pixels. Convert to latent pixels by dividing by vae_scale.
122    Returns list of tile boxes in latent coordinates: (y0,y1,x0,x1)
123    """
124    _, _, H_lat, W_lat = latents.shape
125    tile_lat = max(1, tile_px // vae_scale)
126    overlap_lat = max(0, overlap_px // vae_scale)
127    step = tile_lat - overlap_lat
128    tiles = []
129    y = 0
130    while y < H_lat:
131        y1 = min(y + tile_lat, H_lat)
132        y0 = max(0, y1 - tile_lat)
133        x = 0
134        while x < W_lat:
135            x1 = min(x + tile_lat, W_lat)
136            x0 = max(0, x1 - tile_lat)
137            tiles.append((y0, y1, x0, x1))
138            x += step
139        y += step
140    return tiles
141
142def merge_tile_predictions(pred_full, pred_tile, box):
143    """
144    Merge predicted noise for a tile into the full prediction tensor.
145    Use simple overwrite in this implementation (we can average in overlaps).
146    box = (y0,y1,x0,x1)
147    """
148    y0, y1, x0, x1 = box
149    pred_full[..., y0:y1, x0:x1] = pred_tile
150    return pred_full
151
152@torch.no_grad()
153def denoise_with_tiled_unet(pipe, latents, scheduler, encoder_hidden_states, guidance_scale, device_unet, tile_px, overlap_px, vae_scale):
154    """
155    Run denoising loop but at each UNet call split latents into tiles to reduce peak memory.
156    latents: initial latents (b, c, H_lat, W_lat) on CPU or DML? We'll ensure UNet gets tile tensor moved to DML.
157    scheduler: pipeline.scheduler (already configured)
158    encoder_hidden_states: text embeddings on CPU
159    device_unet: DML device object
160    """
161    pipe_unet = pipe.unet
162    timesteps = scheduler.timesteps
163    # prepare latent dtype / device: keep latents on CPU to save VRAM, move tiles to DML on demand
164    # but UNet expects latents on device_unet
165    for t in timesteps:
166        # prepare model_input as in diffusers: scale latents etc.
167        # mimic pipeline.prepare_latents behavior: here latents are already batched
168        # predict noise for each tile then merge
169        predicted_noise_full = torch.zeros_like(latents)  # keep on CPU to save VRAM
170        tiles = slice_latent_to_tiles(latents, tile_px, overlap_px, vae_scale)
171        for box in tiles:
172            y0, y1, x0, x1 = box
173            latent_tile = latent_tile.to(device_unet, dtype=pipe.unet.dtype)
174            # timestep tensor on device_unet
175            t_tensor = t_tensor.to(device_unet, dtype=pipe.unet.dtype)
176            # encoder_hidden_states may be on CPU: move to device_unet temporarily
177            enc = enc.to(device_unet, dtype=pipe.unet.dtype)
178            # classifier-free guidance split
179            latent_model_input = torch.cat([latent_tile] * 2) if guidance_scale > 1.0 else latent_tile
180            # scale model input if required by scheduler (some schedulers expect it)
181            # call unet
182            with torch.no_grad():
183                noise_pred = pipe_unet(latent_model_input, t_tensor, encoder_hidden_states=enc).sample
184            # guidance
185            if guidance_scale > 1.0:
186                eps_uncond, eps_cond = noise_pred.chunk(2)
187                noise_pred = eps_uncond + guidance_scale * (eps_cond - eps_uncond)
188            # move noise_pred back to CPU and merge
189            noise_pred = noise_pred.to(latents.device)
190            predicted_noise_full = merge_tile_predictions(predicted_noise_full, noise_pred, box)
191            # free DML memory by deleting temp vars
192            del latent_tile, t_tensor, enc, noise_pred, latent_model_input
193            torch.directml.empty_cache()  # try free DML memory (best-effort)
194        # scheduler step: update latents based on predicted noise (on CPU)
195        latents = scheduler.step(predicted_noise_full, t, latents)["prev_sample"]
196    return latents
197
198@torch.no_grad()
199def decode_latents_tiled(pipe, latents, tile_px, overlap_px, vae_scale):
200    """
201    Decode latents to image via VAE in tiles.
202    vae is on CPU (so keep tiles on CPU). We decode each tile and paste into final image with blending in overlaps.
203    Returns PIL Image.
204    """
205    vae = pipe.vae  # on CPU
206    # latents shape: (b, c, H_lat, W_lat)
207    b, c, H_lat, W_lat = latents.shape
208    assert b == 1, "batch>1 not supported in tiled decode"
209    tile_lat = max(1, tile_px // vae_scale)
210    overlap_lat = max(0, overlap_px // vae_scale)
211    step = tile_lat - overlap_lat
212    # compute final image size
213    img_h = H_lat * vae_scale
214    img_w = W_lat * vae_scale
215    # We'll decode tile-by-tile and paste into numpy array, using a weight map for blending
216    final_acc = np.zeros((img_h, img_w, 3), dtype=np.float32)
217    weight_acc = np.zeros((img_h, img_w, 1), dtype=np.float32)
218    tiles = slice_latent_to_tiles(latents, tile_px, overlap_px, vae_scale)
219    for box in tiles:
220        y0, y1, x0, x1 = box
221        latent_tile = latents[..., y0:y1, x0:x1].to(CPU_DEVICE)
222        # VAE decode expects latents scaled by vae.config.scaling maybe; use pipeline.decode_latents if present
223        try:
224            # use pipeline's decode_latents if exists (handles scaling)
225            decoded = pipe.decode_latents(latent_tile)
226            # decode_latents returns list or tensor; normalize to numpy HWC
227            if isinstance(decoded, list):
228                decoded = decoded[0]
229        except Exception:
230            # fallback using vae directly
231            with torch.no_grad():
232                decoded = vae.decode(latent_tile).sample
233        img_tile = decoded.cpu().permute(0,2,3,1).numpy()[0]  # HxWxC, usually in [-1,1]
234        # convert to 0..1
235        img_tile = (img_tile + 1.0) / 2.0
236        # compute pixel coords in final image
237        py0 = int(y0 * vae_scale)
238        py1 = int(y1 * vae_scale)
239        px0 = int(x0 * vae_scale)
240        px1 = int(x1 * vae_scale)
241        h_tile = py1 - py0
242        w_tile = px1 - px0
243        # ensure tile shape matches expected (sometimes rounding)
244        img_tile = Image.fromarray((np.clip(img_tile*255,0,255)).astype(np.uint8)).resize((w_tile,h_tile), resample=Image.LANCZOS)
245        img_tile = np.array(img_tile).astype(np.float32) / 255.0
246        # weight (rectangular) and blend
247        weight = np.ones((h_tile,w_tile,1), dtype=np.float32)
248        final_acc[py0:py1, px0:px1] += img_tile * weight
249        weight_acc[py0:py1, px0:px1] += weight
250        # free mem
251        del latent_tile, decoded, img_tile, weight
252    # normalize by weights
253    weight_acc = np.maximum(weight_acc, 1e-6)
254    final = final_acc / weight_acc
255    final = (np.clip(final,0,1)*255).astype(np.uint8)
256    pil = Image.fromarray(final)
257    return pil
258
259# -------------------------
260# Main flow
261# -------------------------
262def main():
263    print("DirectML device:", DML_DEVICE_NAME)
264    dtype = torch.float16 if USE_FP16 else torch.float32
265    print("Loading pipeline (dtype=%s)..." % dtype)
266    # load pipeline on CPU first
267    pipe = StableDiffusionPipeline.from_single_file(MODEL_PATH, torch_dtype=dtype, use_safetensors=True, safety_checker=None)
268    # apply scheduler LCM if available
269    try:
270        pipe.scheduler = LCMScheduler.from_config(pipe.scheduler.config)
271        print("LCMScheduler active")
272    except Exception:
273        print("LCMScheduler not available or failed to set (falling back to current scheduler)")
274
275    apply_optimizations(pipe)
276    # partial device placement
277    to_device_partial(pipe, DML_DEVICE)
278    # prepare prompt embeddings (on CPU)
279    encoder_hidden_states = make_prompt_embedding(pipe, PROMPT, CPU_DEVICE)
280    # prepare initial latents on CPU (normal pipeline uses random latents)
281    vae_scale = get_latent_downscale(pipe)
282    # latent dims
283    H_lat = HEIGHT // vae_scale
284    W_lat = WIDTH // vae_scale
285    latents = torch.randn((1, pipe.unet.in_channels, H_lat, W_lat), device=CPU_DEVICE, generator=torch.Generator(device=CPU_DEVICE).manual_seed(SEED))
286    # scheduler prepare (use pipeline methods if exist)
287    scheduler = pipe.scheduler
288    scheduler.set_timesteps(STEPS)
289    # scale initial noise if needed by scheduler
290    latents = latents * scheduler.init_noise_sigma
291    start = time.time()
292    print("Starting tiled denoising (this will be slower but memory-efficient)...")
293    latents = denoise_with_tiled_unet(pipe, latents, scheduler, encoder_hidden_states, GUIDANCE, DML_DEVICE, TILE_PX, TILE_OVERLAP_PX, vae_scale)
294    print("Denoising done. Decoding latents via tiled VAE (CPU)...")
295    img = decode_latents_tiled(pipe, latents, TILE_PX, TILE_OVERLAP_PX, vae_scale)
296    img.save(OUT_IMAGE)
297    elapsed = time.time() - start
298    print(f"Saved {OUT_IMAGE} (time: {elapsed:.1f}s)")
299
300if __name__ == "__main__":
301    main()
302