Coverage for wrapper/qwenimageedit/wrapper_qwenimageedit.py: 88%

1import logging 

2import os 

3import sys 

4import random 

5 

6from typing import override 

7from typing import Optional 

8from typing import Dict 

9from typing import Any 

10 

11from PIL import Image 

12 

13import torch 

14from torch import inference_mode 

15 

16from image_utils import base64_to_img 

17from wrapper_model import ModelGeneration 

18 

19from diffusers import QwenImageEditPipeline 

20 

21from xfuser.config import EngineConfig 

22 

23 

24class QwenImageEditGeneration(ModelGeneration): 

25 """Handle image generation using the QwenImageEdit model.""" 

26 HF_MODEL_NAME = "Qwen/Qwen-Image-Edit" 

27 

28 def __init__( 

29 self, 

30 model_name: str = "qwenimageedit", 

31 engine_config: EngineConfig = None, 

32 param_dtype: torch.dtype = torch.bfloat16, 

33 ) -> None: 

34 super().__init__(model_name) 

35 

36 self.engine_config = engine_config 

37 if self.engine_config is not None: 

38 self.torch_compile = self.engine_config.runtime_config.use_torch_compile 

39 self.param_dtype = param_dtype 

40 

41 # Model components 

42 self.pipeline: Optional[QwenImageEditPipeline] = None 

43 

44 def __del__(self) -> None: 

45 """Clean models.""" 

46 if self.pipeline is not None: 

47 self.pipeline = None 

48 

49 def init_parallelism(self) -> None: 

50 """Initialize distributed settings for multi-GPU setups.""" 

51 self.load_timer.start("torch_dist") 

52 

53 self.rank = int(os.getenv("RANK", 0)) 

54 self.local_rank = int(os.getenv("LOCAL_RANK", 0)) 

55 self.world_size = int(os.getenv("WORLD_SIZE", 1)) 

56 

57 self.device_id = self.local_rank 

58 self.device = torch.device(f"cuda:{self.device_id}") 

59 

60 torch.cuda.set_device(self.local_rank) 

61 

62 if self.world_size > 1: 

63 logging.warning("Qwen Image Edit is not optimized for multi-GPU setups (yet).") 

64 self.world_size = 1 

65 

66 self.load_timer.end("torch_dist") 

67 

68 def init_model_parallelism(self) -> None: 

69 """Qwen Image Edit does not support parallelism yet.""" 

70 if self.world_size > 1: 

71 logging.warning("Parallelism not supported.") 

72 

73 def load_model(self) -> None: 

74 """Load the Qwen Image Edit model from Hugging Face.""" 

75 assert torch.cuda.is_available() 

76 

77 self.load_timer.start("pipeline") 

78 self.pipeline = QwenImageEditPipeline.from_pretrained( 

79 pretrained_model_name_or_path=self.HF_MODEL_NAME, 

80 torch_dtype=self.param_dtype, 

81 ) 

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

83 self.load_timer.end("pipeline") 

84 

85 def model_compile(self) -> None: 

86 """Compile the model using torch.compile if enabled.""" 

87 if not self.torch_compile: 

88 return 

89 

90 self.load_timer.start("dit_compile") 

91 torch._inductor.config.reorder_for_compute_comm_overlap = True 

92 assert self.pipeline is not None 

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

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

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

96 ) 

97 self.load_timer.end("dit_compile") 

98 

99 def _assert_model_init(self) -> None: 

100 """Check if the model has been initialized.""" 

101 super()._assert_model_init() 

102 assert self.pipeline is not None 

103 

104 def _get_vae_scale_factor(self) -> int: 

105 """Return the VAE scale factor from the pipeline, raising if unavailable.""" 

106 if not self.pipeline: 

107 raise ValueError("Pipeline not initialized.") 

108 vae_scale_factor = getattr(self.pipeline, "vae_scale_factor", None) 

109 if vae_scale_factor is None: 

110 raise ValueError("Pipeline does not have vae_scale_factor.") 

111 return vae_scale_factor 

112 

113 def _assert_args( 

114 self, 

115 height: int, 

116 width: int, 

117 ) -> None: 

118 """Check if the image size is supported for the current parallelism setting.""" 

119 vae_scale_factor = self._get_vae_scale_factor() 

120 height_latent = height // vae_scale_factor 

121 width_latent = width // vae_scale_factor 

122 img_latent_shape = (height_latent // 2) * (width_latent // 2) 

123 if img_latent_shape % self.world_size != 0: 

124 raise ValueError(f"{height}x{width} not supported for {self.world_size} GPUs.") 

125 

126 @inference_mode() 

127 async def warmup(self) -> None: 

128 """Warmup the model with a dummy generation.""" 

129 logging.info(f"[{self.rank}] Warmup for Qwen Image Edit generation.") 

130 empty_img = Image.new("RGB", (512, 512), (255, 255, 255)) 

131 await self.generate( 

132 empty_img, 

133 width=1280, 

134 height=800, 

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

136 neg_prompt="", 

137 sampling_steps=5) # It needs at least 5 steps to warm up properly 

138 

139 @override 

140 @inference_mode() 

141 async def generate( 

142 self, 

143 img: Image.Image, 

144 height: int, 

145 width: int, 

146 prompt: str, 

147 neg_prompt: str = "", 

148 sampling_steps: int = 25, # 10 

149 seed: Optional[int] = None, 

150 job_id: Optional[str] = None, 

151 ) -> Image.Image: 

152 """Generate an image using QwenImage.""" 

153 gen_timer = self._new_gen_timer(job_id) 

154 

155 self._assert_model_init() 

156 self._assert_args(height, width) 

157 assert self.pipeline is not None 

158 

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

160 

161 try: 

162 if seed is None or seed < 0: 

163 seed = random.randint(0, sys.maxsize) 

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

165 seed_g.manual_seed(seed) 

166 

167 def callback_gen_timer( 

168 pipeline: QwenImageEditPipeline, 

169 step: int, 

170 timestep: int, 

171 callback_kwargs: dict 

172 ) -> dict: 

173 gen_timer.end(f"step_{step:03d}") 

174 if step < sampling_steps - 1: 

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

176 self.check_interrupted() 

177 return callback_kwargs 

178 

179 gen_timer.start(f"step_{0:03d}") 

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

181 image=img, 

182 height=height, 

183 width=width, 

184 prompt=prompt, 

185 negative_prompt=neg_prompt, 

186 num_inference_steps=sampling_steps, 

187 output_type="pil", 

188 generator=seed_g, 

189 callback_on_step_end=callback_gen_timer, 

190 ) 

191 

192 if not output or len(output.images) != 1: 

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

194 image = output.images[0] 

195 return image 

196 finally: 

197 self.running = False 

198 gen_timer.end("total") 

199 

200 async def get_rest_args(self, data_json: Dict[str, str]) -> Dict[str, Any]: 

201 """Parse and validate REST API arguments for Qwen Image Edit generation.""" 

202 if data_json is None: 

203 raise ValueError("Missing JSON body") 

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

205 if img_base64 is None: 

206 raise ValueError("Missing 'img' parameter") 

207 img = base64_to_img(img_base64) 

208 prompt = data_json.get("prompt", None) 

209 if prompt is None: 

210 raise ValueError("Missing 'prompt' parameter") 

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

212 height = int(data_json.get("height", 480)) 

213 width = int(data_json.get("width", 640)) 

214 steps = int(data_json.get("sampling_steps", 20)) 

215 return { 

216 "task": self.model_name, 

217 "args": { 

218 "img": img, 

219 "prompt": prompt, 

220 "neg_prompt": neg_prompt, 

221 "height": height, 

222 "width": width, 

223 "sampling_steps": steps, 

224 } 

225 }