(path: str, scale: float = 4, dtype=torch.bfloat16, device='cuda')
| 85 | |
| 86 | |
| 87 | def prepare_input_tensor(path: str, scale: float = 4, dtype=torch.bfloat16, device='cuda'): |
| 88 | if os.path.isdir(path): |
| 89 | paths0 = list_images_natural(path) |
| 90 | if not paths0: |
| 91 | raise FileNotFoundError(f"No images in {path}") |
| 92 | |
| 93 | with Image.open(paths0[0]) as _img0: |
| 94 | w0, h0 = _img0.size |
| 95 | N0 = len(paths0) |
| 96 | print(f"[{os.path.basename(path)}] Original Resolution: {w0}x{h0} | Original Frames: {N0}") |
| 97 | |
| 98 | sW, sH, tW, tH = compute_scaled_and_target_dims(w0, h0, scale=scale, multiple=128) |
| 99 | print(f"[{os.path.basename(path)}] Scaled (x{scale:.2f}): {sW}x{sH} -> Target (128-multiple): {tW}x{tH}") |
| 100 | |
| 101 | paths = paths0 + [paths0[-1]] * 4 |
| 102 | F = largest_8n1_leq(len(paths)) |
| 103 | if F == 0: |
| 104 | raise RuntimeError(f"Not enough frames after padding in {path}. Got {len(paths)}.") |
| 105 | paths = paths[:F] |
| 106 | print(f"[{os.path.basename(path)}] Target Frames (8n-3): {F-4}") |
| 107 | |
| 108 | frames = [] |
| 109 | for p in paths: |
| 110 | with Image.open(p).convert('RGB') as img: |
| 111 | img_out = upscale_then_center_crop(img, scale=scale, tW=tW, tH=tH) |
| 112 | frames.append(pil_to_tensor_neg1_1(img_out, dtype, device)) |
| 113 | vid = torch.stack(frames, 0).permute(1,0,2,3).unsqueeze(0) # 1 C F H W |
| 114 | fps = 30 |
| 115 | return vid, tH, tW, F, fps |
| 116 | |
| 117 | if is_video(path): |
| 118 | rdr = imageio.get_reader(path) |
| 119 | first = Image.fromarray(rdr.get_data(0)).convert('RGB') |
| 120 | w0, h0 = first.size |
| 121 | |
| 122 | meta = {} |
| 123 | try: meta = rdr.get_meta_data() |
| 124 | except Exception: pass |
| 125 | fps_val = meta.get('fps', 30) |
| 126 | fps = int(round(fps_val)) if isinstance(fps_val, (int, float)) else 30 |
| 127 | |
| 128 | def count_frames(r): |
| 129 | try: |
| 130 | nf = meta.get('nframes', None) |
| 131 | if isinstance(nf,int) and nf>0: return nf |
| 132 | except Exception: pass |
| 133 | try: return r.count_frames() |
| 134 | except Exception: |
| 135 | n=0 |
| 136 | try: |
| 137 | while True: r.get_data(n); n+=1 |
| 138 | except Exception: |
| 139 | return n |
| 140 | |
| 141 | total = count_frames(rdr) |
| 142 | if total <= 0: |
| 143 | rdr.close() |
| 144 | raise RuntimeError(f"Cannot read frames from {path}") |
no test coverage detected