Coverage for wrapper/fluxupscaler/wrapper_fluxupscaler.py: 68%

1""" 

2Wrapper class for Flux Upscaler model generation. 

3""" 

4import logging 

5import sys 

6import random 

7import aiofiles 

8import asyncio 

9 

10from PIL import Image 

11 

12from typing import override 

13from typing import List 

14from typing import Union 

15from typing import Optional 

16from typing import Dict 

17from typing import Any 

18 

19import torch 

20import torch.distributed as dist 

21from torch import inference_mode 

22 

23from model_timing import GenTimer 

24from wrapper_flux import FluxGeneration 

25 

26from flux_xfuser import parallelize_transformer 

27 

28from diffusers import FluxPipeline 

29from diffusers import FluxControlNetModel 

30from diffusers.pipelines import FluxControlNetPipeline 

31 

32from image_utils import base64_to_img 

33from media_utils import save_video_frames 

34from media_utils import base64_to_video_frames 

35from file_utils import base64_to_binary 

36from media_utils import get_video_fps 

37 

38from xfuser.config import EngineConfig 

39from xfuser.core.distributed import get_runtime_state 

40from xfuser.core.distributed import initialize_runtime_state 

41from xfuser.core.distributed import get_pipeline_parallel_world_size 

42 

43 

44class FluxUpscalerGeneration(FluxGeneration): 

45 """Class for image and video upscaling using the Flux model with ControlNet.""" 

46 

47 def __init__( 

48 self, 

49 model_name: str = "fluxupscaler", 

50 engine_config: EngineConfig = None, 

51 param_dtype: torch.dtype = torch.bfloat16, 

52 ) -> None: 

53 super().__init__( 

54 model_name=model_name, 

55 engine_config=engine_config, 

56 param_dtype=param_dtype, 

57 ) 

58 

59 # Model components 

60 self.controlnet: Optional[FluxControlNetModel] = None 

61 self.pipeline: Optional[FluxControlNetPipeline] = None 

62 

63 def __del__(self) -> None: 

64 if self.pipeline is not None: 

65 del self.pipeline.transformer # type: ignore[attr-defined] 

66 self.pipeline = None 

67 if self.controlnet is not None: 

68 del self.controlnet 

69 self.controlnet = None 

70 super().__del__() 

71 

72 def load_model(self) -> None: 

73 assert torch.cuda.is_available() 

74 

75 self.load_timer.start("controlnet") 

76 self.CONTROL_NET_NAME = "jasperai/Flux.1-dev-Controlnet-Upscaler" 

77 self.controlnet = FluxControlNetModel.from_pretrained( 

78 self.CONTROL_NET_NAME, 

79 torch_dtype=self.param_dtype, 

80 ) 

81 self.load_timer.end("controlnet") 

82 

83 self.load_timer.start("pipeline") 

84 cache_args = None 

85 self.MODEL_NAME = "black-forest-labs/FLUX.1-dev" 

86 self.pipeline = FluxControlNetPipeline.from_pretrained( 

87 pretrained_model_name_or_path=self.MODEL_NAME, 

88 controlnet=self.controlnet, 

89 engine_config=self.engine_config, 

90 cache_args=cache_args, 

91 torch_dtype=self.param_dtype, 

92 ) 

93 if not self.pipeline: 

94 raise ValueError("Failed to load FluxControlNet pipeline.") 

95 assert isinstance(self.pipeline, FluxControlNetPipeline) 

96 self.pipeline = self.pipeline.to(self.device) # type: ignore[attr-defined] 

97 self.load_timer.end("pipeline") 

98 

99 logging.info( 

100 f"[{self.rank}] Loaded FluxUpscalerGeneration: {self.MODEL_NAME} and {self.CONTROL_NET_NAME} " 

101 f"device:{self.device} dtype:{self.param_dtype}.") 

102 

103 def init_model_parallelism(self) -> None: 

104 if not dist.is_initialized() or self.world_size <= 1: 

105 return 

106 

107 self.load_timer.start("dit_parallel") 

108 initialize_runtime_state(self.pipeline, self.engine_config) 

109 get_runtime_state().set_input_parameters( 

110 batch_size=1, 

111 # height=self.input_config.height, 

112 # width=self.input_config.width, 

113 # num_inference_steps=self.input_config.num_inference_steps, 

114 max_condition_sequence_length=512, 

115 split_text_embed_in_sp=get_pipeline_parallel_world_size() == 1, 

116 ) 

117 

118 parallelize_transformer(self.pipeline) 

119 self.load_timer.end("dit_parallel") 

120 

121 def model_compile(self) -> None: 

122 if not self.torch_compile: 

123 return 

124 if self.pipeline is None: 

125 return 

126 

127 self.load_timer.start("dit_compile") 

128 torch._inductor.config.reorder_for_compute_comm_overlap = True 

129 self.pipeline.transformer = torch.compile( # type: ignore[attr-defined] 

130 self.pipeline.transformer, # type: ignore[attr-defined] 

131 mode="max-autotune-no-cudagraphs" 

132 ) 

133 self.load_timer.end("dit_compile") 

134 

135 def _assert_model_init(self) -> None: 

136 super()._assert_model_init() 

137 assert self.pipeline is not None 

138 assert self.controlnet is not None 

139 

140 @inference_mode() 

141 async def warmup(self) -> None: 

142 logging.info(f"[{self.rank}] Warmup for Flux Upscaler generation.") 

143 await self.generate( 

144 img=Image.new("RGB", (640, 480), color=(255, 255, 255)), 

145 width=1280, 

146 height=720, 

147 prompt="A warmup image to initialize the model.", 

148 neg_prompt="", 

149 sampling_steps=2) 

150 

151 @override 

152 @inference_mode() 

153 async def generate( 

154 self, 

155 img: Optional[Image.Image] = None, 

156 video: Optional[List[Image.Image]] = None, 

157 height: int = 960, 

158 width: int = 1280, 

159 prompt: str = "", 

160 neg_prompt: str = "", 

161 sampling_steps: int = 28, 

162 controlnet_conditioning_scale: float = 0.6, 

163 guidance_scale: float = 3.5, 

164 video_fps: int = 30, 

165 job_id: Optional[str] = None, 

166 output_type: str = "pil", # "pil", "video_binary", "video_path" 

167 ) -> Any: # returns Image for images, or List[Image]/str/bytes for video 

168 gen_timer = self._new_gen_timer(job_id) 

169 

170 self._assert_model_init() 

171 self._assert_args(height, width) 

172 

173 self.running = True # Mark running to avoid concurrent calls 

174 

175 try: 

176 # Video upscaling 

177 if video is not None: 

178 ret: List[Optional[Image.Image]] = [] 

179 video_frames = video 

180 for it, video_frame in enumerate(video_frames): 

181 if video_frame is None: 

182 ret.append(None) 

183 else: 

184 gen_timer.start(f"frame_{it:03d}") 

185 resized_video_frame = await asyncio.to_thread( 

186 self.generate_image, 

187 gen_timer=gen_timer, 

188 img_id=it, 

189 img=video_frame, 

190 height=height, 

191 width=width, 

192 prompt=prompt, 

193 neg_prompt=neg_prompt, 

194 sampling_steps=sampling_steps, 

195 controlnet_conditioning_scale=controlnet_conditioning_scale, 

196 guidance_scale=guidance_scale) 

197 ret.append(resized_video_frame) 

198 gen_timer.end(f"frame_{it:03d}") 

199 if self.rank == 0: 

200 logging.info(f"[{self.rank}] Generated {len(video)}->{len(ret)} video frames.") 

201 return await self._output_video( 

202 job_id, 

203 gen_timer, 

204 ret, 

205 video_fps, 

206 output_type) 

207 

208 # Image upscaling 

209 if img is not None: 

210 out_image = await asyncio.to_thread( 

211 self.generate_image, 

212 gen_timer=gen_timer, 

213 img=img, 

214 height=height, 

215 width=width, 

216 prompt=prompt, 

217 neg_prompt=neg_prompt, 

218 sampling_steps=sampling_steps, 

219 controlnet_conditioning_scale=controlnet_conditioning_scale, 

220 guidance_scale=guidance_scale) 

221 return out_image 

222 

223 # Missing inputs 

224 raise ValueError("Image or video required for Flux Upscaling generation.") 

225 finally: 

226 self.running = False 

227 gen_timer.end("total") 

228 

229 async def _output_video( 

230 self, 

231 job_id: Optional[str], 

232 gen_timer: GenTimer, 

233 video_frames: List[Optional[Image.Image]], 

234 video_fps: int = 30, 

235 output_type: str = "pil", # "pil", "video_binary", "video_path" 

236 ) -> Union[List[Optional[Image.Image]], str, bytes, None]: 

237 gen_timer.start("output") 

238 try: 

239 if output_type == "pil": 

240 return video_frames 

241 

242 if output_type in ("video_binary", "video_path"): 

243 video_path = None 

244 if job_id: 

245 video_path = f"/tmp/{job_id}.mp4" 

246 video_path = await save_video_frames( 

247 [f for f in video_frames if f is not None], 

248 out_video_path=video_path, 

249 fps=video_fps) 

250 if output_type == "video_path": 

251 return video_path 

252 

253 # video_binary 

254 async with aiofiles.open(video_path, "rb") as f: 

255 video_binary = await f.read() 

256 return video_binary 

257 

258 logging.error(f"Unknown output type: {output_type}") 

259 return None 

260 finally: 

261 gen_timer.end("output") 

262 

263 @inference_mode() 

264 def generate_image( 

265 self, 

266 gen_timer: GenTimer, 

267 img_id: int = 0, 

268 img: Optional[Image.Image] = None, 

269 height: int = 960, 

270 width: int = 1280, 

271 prompt: str = "", 

272 neg_prompt: str = "", 

273 sampling_steps: int = 28, 

274 controlnet_conditioning_scale: float = 0.6, 

275 guidance_scale: float = 3.5, 

276 ) -> Image.Image: 

277 """ 

278 Upscale one image from a prompt using the Flux model. 

279 """ 

280 assert img is not None, "Image is required for Flux Upscaling." 

281 img = img.resize((width, height), Image.Resampling.LANCZOS) 

282 

283 seed = self.base_seed if self.base_seed >= 0 else random.randint(0, sys.maxsize) 

284 seed_g = torch.Generator(device=self.device) 

285 seed_g.manual_seed(seed) 

286 

287 def callback_gen_timer( 

288 pipeline: FluxPipeline, 

289 step: int, 

290 timestep: int, 

291 callback_kwargs: dict 

292 ) -> dict: 

293 gen_timer.end(f"step_{img_id:03d}_{step:03d}") 

294 if step < sampling_steps - 1: 

295 gen_timer.start(f"step_{img_id:03d}_{step + 1:03d}") 

296 self.check_interrupted() 

297 return callback_kwargs 

298 

299 assert self.pipeline is not None, "Flux pipeline not initialized." 

300 gen_timer.start(f"step_{img_id:03d}_{0:03d}") 

301 output = self.pipeline( # type: ignore[operator] 

302 control_image=img, 

303 height=height, 

304 width=width, 

305 prompt=prompt, 

306 negative_prompt=neg_prompt, 

307 num_inference_steps=sampling_steps, 

308 controlnet_conditioning_scale=controlnet_conditioning_scale, 

309 guidance_scale=guidance_scale, 

310 output_type="pil", 

311 generator=seed_g, 

312 callback_on_step_end=callback_gen_timer, 

313 ) 

314 

315 if len(output.images) != 1: 

316 raise ValueError(f"Expected 1 image, but got {len(output.images)} images.") 

317 

318 return output.images[0] 

319 

320 def get_health(self) -> Dict[str, Any]: 

321 ret = super().get_health() 

322 ret.update({ 

323 "device_map": getattr(self.pipeline, "hf_device_map", None) if self.pipeline else None, 

324 }) 

325 return ret 

326 

327 async def get_rest_args( 

328 self, 

329 data_json: Dict[str, Union[str, int, float]] 

330 ) -> Dict[str, Any]: 

331 if data_json is None or not isinstance(data_json, dict): 

332 raise ValueError("Missing JSON body") 

333 

334 job_id = data_json.get("job_id", None) 

335 

336 img_base64 = data_json.get("img", None) 

337 img = None 

338 if img_base64 is not None: 

339 img = base64_to_img(str(img_base64)) 

340 

341 video_base64 = data_json.get("video", None) 

342 video_frames = None 

343 video_fps: float = -1 

344 if video_base64 is not None: 

345 video_frames = base64_to_video_frames(str(video_base64)) 

346 video_binary = base64_to_binary(str(video_base64)) 

347 video_fps = get_video_fps(video_binary) 

348 

349 prompt = data_json.get("prompt", "") 

350 neg_prompt = data_json.get("neg_prompt", "") 

351 

352 rest_args: Dict[str, Any] = { 

353 "task": self.model_name, 

354 "args": { 

355 "job_id": job_id, 

356 "img": img, 

357 "video": video_frames, 

358 "prompt": prompt, 

359 "neg_prompt": neg_prompt, 

360 "width": int(data_json.get("width", 640)), 

361 "height": int(data_json.get("height", 480)), 

362 "sampling_steps": int(data_json.get("sampling_steps", 28)), 

363 "controlnet_conditioning_scale": data_json.get("controlnet_conditioning_scale", 0.6), 

364 "guidance_scale": data_json.get("guidance_scale", 3.5), 

365 } 

366 } 

367 if video_fps > 0: 

368 rest_args["args"]["video_fps"] = video_fps 

369 return rest_args