@@ -155,7 +155,11 @@ impl FluxEngine {
155155 // Load VAE
156156 tracing:: info!( path = %self . paths. vae. display( ) , "loading VAE..." ) ;
157157 let vae_vb = unsafe {
158- VarBuilder :: from_mmaped_safetensors ( std:: slice:: from_ref ( & self . paths . vae ) , dtype, & device) ?
158+ VarBuilder :: from_mmaped_safetensors (
159+ std:: slice:: from_ref ( & self . paths . vae ) ,
160+ dtype,
161+ & device,
162+ ) ?
159163 } ;
160164 let vae_cfg = if is_schnell {
161165 flux:: autoencoder:: Config :: schnell ( )
@@ -214,8 +218,7 @@ impl InferenceEngine for FluxEngine {
214218 . get_ids ( )
215219 . to_vec ( ) ;
216220 tokens. resize ( 256 , 0 ) ;
217- let input_ids =
218- Tensor :: new ( & tokens[ ..] , & loaded. device ) ?. unsqueeze ( 0 ) ?;
221+ let input_ids = Tensor :: new ( & tokens[ ..] , & loaded. device ) ?. unsqueeze ( 0 ) ?;
219222 loaded. t5_model . forward ( & input_ids) ?
220223 } ;
221224 tracing:: info!( "T5 encoding complete" ) ;
@@ -228,15 +231,14 @@ impl InferenceEngine for FluxEngine {
228231 . map_err ( |e| anyhow:: anyhow!( "CLIP tokenization failed: {e}" ) ) ?
229232 . get_ids ( )
230233 . to_vec ( ) ;
231- let input_ids =
232- Tensor :: new ( & tokens[ ..] , & loaded. device ) ?. unsqueeze ( 0 ) ?;
234+ let input_ids = Tensor :: new ( & tokens[ ..] , & loaded. device ) ?. unsqueeze ( 0 ) ?;
233235 loaded. clip_model . forward ( & input_ids) ?
234236 } ;
235237 tracing:: info!( "CLIP encoding complete" ) ;
236238
237239 // 3. Generate initial noise
238- let img = flux :: sampling :: get_noise ( 1 , height , width , & loaded . device ) ?
239- . to_dtype ( loaded. dtype ) ?;
240+ let img =
241+ flux :: sampling :: get_noise ( 1 , height , width , & loaded . device ) ? . to_dtype ( loaded. dtype ) ?;
240242
241243 // 4. Build sampling state
242244 let state = flux:: sampling:: State :: new ( & t5_emb, & clip_emb, & img) ?;
@@ -245,10 +247,7 @@ impl InferenceEngine for FluxEngine {
245247 let timesteps = if loaded. is_schnell {
246248 flux:: sampling:: get_schedule ( req. steps as usize , None )
247249 } else {
248- flux:: sampling:: get_schedule (
249- req. steps as usize ,
250- Some ( ( state. img . dim ( 1 ) ?, 0.5 , 1.15 ) ) ,
251- )
250+ flux:: sampling:: get_schedule ( req. steps as usize , Some ( ( state. img . dim ( 1 ) ?, 0.5 , 1.15 ) ) )
252251 } ;
253252
254253 tracing:: info!( steps = timesteps. len( ) , "running denoising loop..." ) ;
@@ -273,8 +272,7 @@ impl InferenceEngine for FluxEngine {
273272 let img = loaded. vae . decode ( & img) ?;
274273
275274 // 9. Convert to u8 image: clamp to [-1, 1], map to [0, 255]
276- let img = ( ( img. clamp ( -1f32 , 1f32 ) ? + 1.0 ) ? * 127.5 ) ?
277- . to_dtype ( DType :: U8 ) ?;
275+ let img = ( ( img. clamp ( -1f32 , 1f32 ) ? + 1.0 ) ? * 127.5 ) ?. to_dtype ( DType :: U8 ) ?;
278276 let img = img. i ( 0 ) ?; // remove batch dim: [3, H, W]
279277
280278 tracing:: info!( "VAE decode complete, encoding output image..." ) ;
0 commit comments