Spaces:
				
			
			
	
			
			
		Running
		
			on 
			
			Zero
	
	
	
			
			
	
	
	
	
		
		
		Running
		
			on 
			
			Zero
	add AoTI + FA3 (#1)
Browse files- add AoTI + FA3 (c2755f2573574227608b23ec3a1c25dcd9105755)
- Upload 3 files (748e90cea657a59721c266fb05672765c7c10945)
- Upload pipeline_qwenimage_edit_plus.py (7efe1326340efcff90bd2dba65d16fa8170cedfa)
- Create optimization.py (1ea3fe6942b427214bc78f9793a8260e2cc89d48)
- Update app.py (81a6058b6f7d88da31f234ec152ed4c043111907)
- app.py +12 -1
- optimization.py +70 -0
- qwenimage/__init__.py +0 -0
- qwenimage/pipeline_qwenimage_edit_plus.py +891 -0
- qwenimage/qwen_fa3_processor.py +142 -0
- qwenimage/transformer_qwenimage.py +642 -0
- requirements.txt +2 -1
    	
        app.py
    CHANGED
    
    | @@ -5,7 +5,12 @@ import torch | |
| 5 | 
             
            import spaces
         | 
| 6 |  | 
| 7 | 
             
            from PIL import Image
         | 
| 8 | 
            -
            from diffusers import  | 
|  | |
|  | |
|  | |
|  | |
|  | |
| 9 | 
             
            from huggingface_hub import InferenceClient
         | 
| 10 | 
             
            import math
         | 
| 11 |  | 
| @@ -193,6 +198,12 @@ pipe.load_lora_weights( | |
| 193 | 
             
                )
         | 
| 194 | 
             
            pipe.fuse_lora()
         | 
| 195 |  | 
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
| 196 | 
             
            # --- UI Constants and Helpers ---
         | 
| 197 | 
             
            MAX_SEED = np.iinfo(np.int32).max
         | 
| 198 |  | 
|  | |
| 5 | 
             
            import spaces
         | 
| 6 |  | 
| 7 | 
             
            from PIL import Image
         | 
| 8 | 
            +
            from diffusers import FlowMatchEulerDiscreteScheduler
         | 
| 9 | 
            +
            from optimization import optimize_pipeline_
         | 
| 10 | 
            +
            from qwenimage.pipeline_qwenimage_edit_plus import QwenImageEditPlusPipeline
         | 
| 11 | 
            +
            from qwenimage.transformer_qwenimage import QwenImageTransformer2DModel
         | 
| 12 | 
            +
            from qwenimage.qwen_fa3_processor import QwenDoubleStreamAttnProcessorFA3
         | 
| 13 | 
            +
             | 
| 14 | 
             
            from huggingface_hub import InferenceClient
         | 
| 15 | 
             
            import math
         | 
| 16 |  | 
|  | |
| 198 | 
             
                )
         | 
| 199 | 
             
            pipe.fuse_lora()
         | 
| 200 |  | 
| 201 | 
            +
            # Apply the same optimizations from the first version
         | 
| 202 | 
            +
            pipe.transformer.__class__ = QwenImageTransformer2DModel
         | 
| 203 | 
            +
            pipe.transformer.set_attn_processor(QwenDoubleStreamAttnProcessorFA3())
         | 
| 204 | 
            +
             | 
| 205 | 
            +
            # --- Ahead-of-time compilation ---
         | 
| 206 | 
            +
             | 
| 207 | 
             
            # --- UI Constants and Helpers ---
         | 
| 208 | 
             
            MAX_SEED = np.iinfo(np.int32).max
         | 
| 209 |  | 
    	
        optimization.py
    ADDED
    
    | @@ -0,0 +1,70 @@ | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | 
|  | |
| 1 | 
            +
            """
         | 
| 2 | 
            +
            """
         | 
| 3 | 
            +
             | 
| 4 | 
            +
            from typing import Any
         | 
| 5 | 
            +
            from typing import Callable
         | 
| 6 | 
            +
            from typing import ParamSpec
         | 
| 7 | 
            +
            from torchao.quantization import quantize_
         | 
| 8 | 
            +
            from torchao.quantization import Float8DynamicActivationFloat8WeightConfig
         | 
| 9 | 
            +
            import spaces
         | 
| 10 | 
            +
            import torch
         | 
| 11 | 
            +
            from torch.utils._pytree import tree_map
         | 
| 12 | 
            +
             | 
| 13 | 
            +
             | 
| 14 | 
            +
            P = ParamSpec('P')
         | 
| 15 | 
            +
             | 
| 16 | 
            +
             | 
| 17 | 
            +
            TRANSFORMER_IMAGE_SEQ_LENGTH_DIM = torch.export.Dim('image_seq_length')
         | 
| 18 | 
            +
            TRANSFORMER_TEXT_SEQ_LENGTH_DIM = torch.export.Dim('text_seq_length')
         | 
| 19 | 
            +
             | 
| 20 | 
            +
            TRANSFORMER_DYNAMIC_SHAPES = {
         | 
| 21 | 
            +
                'hidden_states': {
         | 
| 22 | 
            +
                    1: TRANSFORMER_IMAGE_SEQ_LENGTH_DIM,
         | 
| 23 | 
            +
                },
         | 
| 24 | 
            +
                'encoder_hidden_states': {
         | 
| 25 | 
            +
                    1: TRANSFORMER_TEXT_SEQ_LENGTH_DIM,
         | 
| 26 | 
            +
                },
         | 
| 27 | 
            +
                'encoder_hidden_states_mask': {
         | 
| 28 | 
            +
                    1: TRANSFORMER_TEXT_SEQ_LENGTH_DIM,
         | 
| 29 | 
            +
                },
         | 
| 30 | 
            +
                'image_rotary_emb': ({
         | 
| 31 | 
            +
                    0: TRANSFORMER_IMAGE_SEQ_LENGTH_DIM,
         | 
| 32 | 
            +
                }, {
         | 
| 33 | 
            +
                    0: TRANSFORMER_TEXT_SEQ_LENGTH_DIM,
         | 
| 34 | 
            +
                }),
         | 
| 35 | 
            +
            }
         | 
| 36 | 
            +
             | 
| 37 | 
            +
             | 
| 38 | 
            +
            INDUCTOR_CONFIGS = {
         | 
| 39 | 
            +
                'conv_1x1_as_mm': True,
         | 
| 40 | 
            +
                'epilogue_fusion': False,
         | 
| 41 | 
            +
                'coordinate_descent_tuning': True,
         | 
| 42 | 
            +
                'coordinate_descent_check_all_directions': True,
         | 
| 43 | 
            +
                'max_autotune': True,
         | 
| 44 | 
            +
                'triton.cudagraphs': True,
         | 
| 45 | 
            +
            }
         | 
| 46 | 
            +
             | 
| 47 | 
            +
             | 
| 48 | 
            +
            def optimize_pipeline_(pipeline: Callable[P, Any], *args: P.args, **kwargs: P.kwargs):
         | 
| 49 | 
            +
             | 
| 50 | 
            +
                @spaces.GPU(duration=1500)
         | 
| 51 | 
            +
                def compile_transformer():
         | 
| 52 | 
            +
             | 
| 53 | 
            +
                    with spaces.aoti_capture(pipeline.transformer) as call:
         | 
| 54 | 
            +
                        pipeline(*args, **kwargs)
         | 
| 55 | 
            +
             | 
| 56 | 
            +
                    dynamic_shapes = tree_map(lambda t: None, call.kwargs)
         | 
| 57 | 
            +
                    dynamic_shapes |= TRANSFORMER_DYNAMIC_SHAPES
         | 
| 58 | 
            +
             | 
| 59 | 
            +
                    # quantize_(pipeline.transformer, Float8DynamicActivationFloat8WeightConfig())
         | 
| 60 | 
            +
                    
         | 
| 61 | 
            +
                    exported = torch.export.export(
         | 
| 62 | 
            +
                        mod=pipeline.transformer,
         | 
| 63 | 
            +
                        args=call.args,
         | 
| 64 | 
            +
                        kwargs=call.kwargs,
         | 
| 65 | 
            +
                        dynamic_shapes=dynamic_shapes,
         | 
| 66 | 
            +
                    )
         | 
| 67 | 
            +
             | 
| 68 | 
            +
                    return spaces.aoti_compile(exported, INDUCTOR_CONFIGS)
         | 
| 69 | 
            +
             | 
| 70 | 
            +
                spaces.aoti_apply(compile_transformer(), pipeline.transformer)
         | 
    	
        qwenimage/__init__.py
    ADDED
    
    | 
            File without changes
         | 
    	
        qwenimage/pipeline_qwenimage_edit_plus.py
    ADDED
    
    | @@ -0,0 +1,891 @@ | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | 
|  | |
| 1 | 
            +
            # Copyright 2025 Qwen-Image Team and The HuggingFace Team. All rights reserved.
         | 
| 2 | 
            +
            #
         | 
| 3 | 
            +
            # Licensed under the Apache License, Version 2.0 (the "License");
         | 
| 4 | 
            +
            # you may not use this file except in compliance with the License.
         | 
| 5 | 
            +
            # You may obtain a copy of the License at
         | 
| 6 | 
            +
            #
         | 
| 7 | 
            +
            #     http://www.apache.org/licenses/LICENSE-2.0
         | 
| 8 | 
            +
            #
         | 
| 9 | 
            +
            # Unless required by applicable law or agreed to in writing, software
         | 
| 10 | 
            +
            # distributed under the License is distributed on an "AS IS" BASIS,
         | 
| 11 | 
            +
            # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
         | 
| 12 | 
            +
            # See the License for the specific language governing permissions and
         | 
| 13 | 
            +
            # limitations under the License.
         | 
| 14 | 
            +
             | 
| 15 | 
            +
            import inspect
         | 
| 16 | 
            +
            import math
         | 
| 17 | 
            +
            from typing import Any, Callable, Dict, List, Optional, Union
         | 
| 18 | 
            +
             | 
| 19 | 
            +
            import numpy as np
         | 
| 20 | 
            +
            import torch
         | 
| 21 | 
            +
            from transformers import Qwen2_5_VLForConditionalGeneration, Qwen2Tokenizer, Qwen2VLProcessor
         | 
| 22 | 
            +
             | 
| 23 | 
            +
            from diffusers.image_processor import PipelineImageInput, VaeImageProcessor
         | 
| 24 | 
            +
            from diffusers.loaders import QwenImageLoraLoaderMixin
         | 
| 25 | 
            +
            from diffusers.models import AutoencoderKLQwenImage, QwenImageTransformer2DModel
         | 
| 26 | 
            +
            from diffusers.schedulers import FlowMatchEulerDiscreteScheduler
         | 
| 27 | 
            +
            from diffusers.utils import is_torch_xla_available, logging, replace_example_docstring
         | 
| 28 | 
            +
            from diffusers.utils.torch_utils import randn_tensor
         | 
| 29 | 
            +
            from diffusers.pipelines.pipeline_utils import DiffusionPipeline
         | 
| 30 | 
            +
            from diffusers.pipelines.qwenimage.pipeline_output import QwenImagePipelineOutput
         | 
| 31 | 
            +
             | 
| 32 | 
            +
             | 
| 33 | 
            +
            if is_torch_xla_available():
         | 
| 34 | 
            +
                import torch_xla.core.xla_model as xm
         | 
| 35 | 
            +
             | 
| 36 | 
            +
                XLA_AVAILABLE = True
         | 
| 37 | 
            +
            else:
         | 
| 38 | 
            +
                XLA_AVAILABLE = False
         | 
| 39 | 
            +
             | 
| 40 | 
            +
             | 
| 41 | 
            +
            logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
         | 
| 42 | 
            +
             | 
| 43 | 
            +
            EXAMPLE_DOC_STRING = """
         | 
| 44 | 
            +
                Examples:
         | 
| 45 | 
            +
                    ```py
         | 
| 46 | 
            +
                    >>> import torch
         | 
| 47 | 
            +
                    >>> from PIL import Image
         | 
| 48 | 
            +
                    >>> from diffusers import QwenImageEditPlusPipeline
         | 
| 49 | 
            +
                    >>> from diffusers.utils import load_image
         | 
| 50 | 
            +
             | 
| 51 | 
            +
                    >>> pipe = QwenImageEditPlusPipeline.from_pretrained("Qwen/Qwen-Image-Edit-2509", torch_dtype=torch.bfloat16)
         | 
| 52 | 
            +
                    >>> pipe.to("cuda")
         | 
| 53 | 
            +
                    >>> image = load_image(
         | 
| 54 | 
            +
                    ...     "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/yarn-art-pikachu.png"
         | 
| 55 | 
            +
                    ... ).convert("RGB")
         | 
| 56 | 
            +
                    >>> prompt = (
         | 
| 57 | 
            +
                    ...     "Make Pikachu hold a sign that says 'Qwen Edit is awesome', yarn art style, detailed, vibrant colors"
         | 
| 58 | 
            +
                    ... )
         | 
| 59 | 
            +
                    >>> # Depending on the variant being used, the pipeline call will slightly vary.
         | 
| 60 | 
            +
                    >>> # Refer to the pipeline documentation for more details.
         | 
| 61 | 
            +
                    >>> image = pipe(image, prompt, num_inference_steps=50).images[0]
         | 
| 62 | 
            +
                    >>> image.save("qwenimage_edit_plus.png")
         | 
| 63 | 
            +
                    ```
         | 
| 64 | 
            +
            """
         | 
| 65 | 
            +
             | 
| 66 | 
            +
            CONDITION_IMAGE_SIZE = 384 * 384
         | 
| 67 | 
            +
            VAE_IMAGE_SIZE = 1024 * 1024
         | 
| 68 | 
            +
             | 
| 69 | 
            +
             | 
| 70 | 
            +
            # Copied from diffusers.pipelines.qwenimage.pipeline_qwenimage.calculate_shift
         | 
| 71 | 
            +
            def calculate_shift(
         | 
| 72 | 
            +
                image_seq_len,
         | 
| 73 | 
            +
                base_seq_len: int = 256,
         | 
| 74 | 
            +
                max_seq_len: int = 4096,
         | 
| 75 | 
            +
                base_shift: float = 0.5,
         | 
| 76 | 
            +
                max_shift: float = 1.15,
         | 
| 77 | 
            +
            ):
         | 
| 78 | 
            +
                m = (max_shift - base_shift) / (max_seq_len - base_seq_len)
         | 
| 79 | 
            +
                b = base_shift - m * base_seq_len
         | 
| 80 | 
            +
                mu = image_seq_len * m + b
         | 
| 81 | 
            +
                return mu
         | 
| 82 | 
            +
             | 
| 83 | 
            +
             | 
| 84 | 
            +
            # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps
         | 
| 85 | 
            +
            def retrieve_timesteps(
         | 
| 86 | 
            +
                scheduler,
         | 
| 87 | 
            +
                num_inference_steps: Optional[int] = None,
         | 
| 88 | 
            +
                device: Optional[Union[str, torch.device]] = None,
         | 
| 89 | 
            +
                timesteps: Optional[List[int]] = None,
         | 
| 90 | 
            +
                sigmas: Optional[List[float]] = None,
         | 
| 91 | 
            +
                **kwargs,
         | 
| 92 | 
            +
            ):
         | 
| 93 | 
            +
                r"""
         | 
| 94 | 
            +
                Calls the scheduler's `set_timesteps` method and retrieves timesteps from the scheduler after the call. Handles
         | 
| 95 | 
            +
                custom timesteps. Any kwargs will be supplied to `scheduler.set_timesteps`.
         | 
| 96 | 
            +
             | 
| 97 | 
            +
                Args:
         | 
| 98 | 
            +
                    scheduler (`SchedulerMixin`):
         | 
| 99 | 
            +
                        The scheduler to get timesteps from.
         | 
| 100 | 
            +
                    num_inference_steps (`int`):
         | 
| 101 | 
            +
                        The number of diffusion steps used when generating samples with a pre-trained model. If used, `timesteps`
         | 
| 102 | 
            +
                        must be `None`.
         | 
| 103 | 
            +
                    device (`str` or `torch.device`, *optional*):
         | 
| 104 | 
            +
                        The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
         | 
| 105 | 
            +
                    timesteps (`List[int]`, *optional*):
         | 
| 106 | 
            +
                        Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed,
         | 
| 107 | 
            +
                        `num_inference_steps` and `sigmas` must be `None`.
         | 
| 108 | 
            +
                    sigmas (`List[float]`, *optional*):
         | 
| 109 | 
            +
                        Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed,
         | 
| 110 | 
            +
                        `num_inference_steps` and `timesteps` must be `None`.
         | 
| 111 | 
            +
             | 
| 112 | 
            +
                Returns:
         | 
| 113 | 
            +
                    `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
         | 
| 114 | 
            +
                    second element is the number of inference steps.
         | 
| 115 | 
            +
                """
         | 
| 116 | 
            +
                if timesteps is not None and sigmas is not None:
         | 
| 117 | 
            +
                    raise ValueError("Only one of `timesteps` or `sigmas` can be passed. Please choose one to set custom values")
         | 
| 118 | 
            +
                if timesteps is not None:
         | 
| 119 | 
            +
                    accepts_timesteps = "timesteps" in set(inspect.signature(scheduler.set_timesteps).parameters.keys())
         | 
| 120 | 
            +
                    if not accepts_timesteps:
         | 
| 121 | 
            +
                        raise ValueError(
         | 
| 122 | 
            +
                            f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"
         | 
| 123 | 
            +
                            f" timestep schedules. Please check whether you are using the correct scheduler."
         | 
| 124 | 
            +
                        )
         | 
| 125 | 
            +
                    scheduler.set_timesteps(timesteps=timesteps, device=device, **kwargs)
         | 
| 126 | 
            +
                    timesteps = scheduler.timesteps
         | 
| 127 | 
            +
                    num_inference_steps = len(timesteps)
         | 
| 128 | 
            +
                elif sigmas is not None:
         | 
| 129 | 
            +
                    accept_sigmas = "sigmas" in set(inspect.signature(scheduler.set_timesteps).parameters.keys())
         | 
| 130 | 
            +
                    if not accept_sigmas:
         | 
| 131 | 
            +
                        raise ValueError(
         | 
| 132 | 
            +
                            f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"
         | 
| 133 | 
            +
                            f" sigmas schedules. Please check whether you are using the correct scheduler."
         | 
| 134 | 
            +
                        )
         | 
| 135 | 
            +
                    scheduler.set_timesteps(sigmas=sigmas, device=device, **kwargs)
         | 
| 136 | 
            +
                    timesteps = scheduler.timesteps
         | 
| 137 | 
            +
                    num_inference_steps = len(timesteps)
         | 
| 138 | 
            +
                else:
         | 
| 139 | 
            +
                    scheduler.set_timesteps(num_inference_steps, device=device, **kwargs)
         | 
| 140 | 
            +
                    timesteps = scheduler.timesteps
         | 
| 141 | 
            +
                return timesteps, num_inference_steps
         | 
| 142 | 
            +
             | 
| 143 | 
            +
             | 
| 144 | 
            +
            # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents
         | 
| 145 | 
            +
            def retrieve_latents(
         | 
| 146 | 
            +
                encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample"
         | 
| 147 | 
            +
            ):
         | 
| 148 | 
            +
                if hasattr(encoder_output, "latent_dist") and sample_mode == "sample":
         | 
| 149 | 
            +
                    return encoder_output.latent_dist.sample(generator)
         | 
| 150 | 
            +
                elif hasattr(encoder_output, "latent_dist") and sample_mode == "argmax":
         | 
| 151 | 
            +
                    return encoder_output.latent_dist.mode()
         | 
| 152 | 
            +
                elif hasattr(encoder_output, "latents"):
         | 
| 153 | 
            +
                    return encoder_output.latents
         | 
| 154 | 
            +
                else:
         | 
| 155 | 
            +
                    raise AttributeError("Could not access latents of provided encoder_output")
         | 
| 156 | 
            +
             | 
| 157 | 
            +
             | 
| 158 | 
            +
            def calculate_dimensions(target_area, ratio):
         | 
| 159 | 
            +
                width = math.sqrt(target_area * ratio)
         | 
| 160 | 
            +
                height = width / ratio
         | 
| 161 | 
            +
             | 
| 162 | 
            +
                width = round(width / 32) * 32
         | 
| 163 | 
            +
                height = round(height / 32) * 32
         | 
| 164 | 
            +
             | 
| 165 | 
            +
                return width, height
         | 
| 166 | 
            +
             | 
| 167 | 
            +
             | 
| 168 | 
            +
            class QwenImageEditPlusPipeline(DiffusionPipeline, QwenImageLoraLoaderMixin):
         | 
| 169 | 
            +
                r"""
         | 
| 170 | 
            +
                The Qwen-Image-Edit pipeline for image editing.
         | 
| 171 | 
            +
             | 
| 172 | 
            +
                Args:
         | 
| 173 | 
            +
                    transformer ([`QwenImageTransformer2DModel`]):
         | 
| 174 | 
            +
                        Conditional Transformer (MMDiT) architecture to denoise the encoded image latents.
         | 
| 175 | 
            +
                    scheduler ([`FlowMatchEulerDiscreteScheduler`]):
         | 
| 176 | 
            +
                        A scheduler to be used in combination with `transformer` to denoise the encoded image latents.
         | 
| 177 | 
            +
                    vae ([`AutoencoderKL`]):
         | 
| 178 | 
            +
                        Variational Auto-Encoder (VAE) Model to encode and decode images to and from latent representations.
         | 
| 179 | 
            +
                    text_encoder ([`Qwen2.5-VL-7B-Instruct`]):
         | 
| 180 | 
            +
                        [Qwen2.5-VL-7B-Instruct](https://huggingface.co/Qwen/Qwen2.5-VL-7B-Instruct), specifically the
         | 
| 181 | 
            +
                        [Qwen2.5-VL-7B-Instruct](https://huggingface.co/Qwen/Qwen2.5-VL-7B-Instruct) variant.
         | 
| 182 | 
            +
                    tokenizer (`QwenTokenizer`):
         | 
| 183 | 
            +
                        Tokenizer of class
         | 
| 184 | 
            +
                        [CLIPTokenizer](https://huggingface.co/docs/transformers/en/model_doc/clip#transformers.CLIPTokenizer).
         | 
| 185 | 
            +
                """
         | 
| 186 | 
            +
             | 
| 187 | 
            +
                model_cpu_offload_seq = "text_encoder->transformer->vae"
         | 
| 188 | 
            +
                _callback_tensor_inputs = ["latents", "prompt_embeds"]
         | 
| 189 | 
            +
             | 
| 190 | 
            +
                def __init__(
         | 
| 191 | 
            +
                    self,
         | 
| 192 | 
            +
                    scheduler: FlowMatchEulerDiscreteScheduler,
         | 
| 193 | 
            +
                    vae: AutoencoderKLQwenImage,
         | 
| 194 | 
            +
                    text_encoder: Qwen2_5_VLForConditionalGeneration,
         | 
| 195 | 
            +
                    tokenizer: Qwen2Tokenizer,
         | 
| 196 | 
            +
                    processor: Qwen2VLProcessor,
         | 
| 197 | 
            +
                    transformer: QwenImageTransformer2DModel,
         | 
| 198 | 
            +
                ):
         | 
| 199 | 
            +
                    super().__init__()
         | 
| 200 | 
            +
             | 
| 201 | 
            +
                    self.register_modules(
         | 
| 202 | 
            +
                        vae=vae,
         | 
| 203 | 
            +
                        text_encoder=text_encoder,
         | 
| 204 | 
            +
                        tokenizer=tokenizer,
         | 
| 205 | 
            +
                        processor=processor,
         | 
| 206 | 
            +
                        transformer=transformer,
         | 
| 207 | 
            +
                        scheduler=scheduler,
         | 
| 208 | 
            +
                    )
         | 
| 209 | 
            +
                    self.vae_scale_factor = 2 ** len(self.vae.temperal_downsample) if getattr(self, "vae", None) else 8
         | 
| 210 | 
            +
                    self.latent_channels = self.vae.config.z_dim if getattr(self, "vae", None) else 16
         | 
| 211 | 
            +
                    # QwenImage latents are turned into 2x2 patches and packed. This means the latent width and height has to be divisible
         | 
| 212 | 
            +
                    # by the patch size. So the vae scale factor is multiplied by the patch size to account for this
         | 
| 213 | 
            +
                    self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor * 2)
         | 
| 214 | 
            +
                    self.tokenizer_max_length = 1024
         | 
| 215 | 
            +
             | 
| 216 | 
            +
                    self.prompt_template_encode = "<|im_start|>system\nDescribe the key features of the input image (color, shape, size, texture, objects, background), then explain how the user's text instruction should alter or modify the image. Generate a new image that meets the user's requirements while maintaining consistency with the original input where appropriate.<|im_end|>\n<|im_start|>user\n{}<|im_end|>\n<|im_start|>assistant\n"
         | 
| 217 | 
            +
                    self.prompt_template_encode_start_idx = 64
         | 
| 218 | 
            +
                    self.default_sample_size = 128
         | 
| 219 | 
            +
             | 
| 220 | 
            +
                # Copied from diffusers.pipelines.qwenimage.pipeline_qwenimage.QwenImagePipeline._extract_masked_hidden
         | 
| 221 | 
            +
                def _extract_masked_hidden(self, hidden_states: torch.Tensor, mask: torch.Tensor):
         | 
| 222 | 
            +
                    bool_mask = mask.bool()
         | 
| 223 | 
            +
                    valid_lengths = bool_mask.sum(dim=1)
         | 
| 224 | 
            +
                    selected = hidden_states[bool_mask]
         | 
| 225 | 
            +
                    split_result = torch.split(selected, valid_lengths.tolist(), dim=0)
         | 
| 226 | 
            +
             | 
| 227 | 
            +
                    return split_result
         | 
| 228 | 
            +
             | 
| 229 | 
            +
                def _get_qwen_prompt_embeds(
         | 
| 230 | 
            +
                    self,
         | 
| 231 | 
            +
                    prompt: Union[str, List[str]] = None,
         | 
| 232 | 
            +
                    image: Optional[torch.Tensor] = None,
         | 
| 233 | 
            +
                    device: Optional[torch.device] = None,
         | 
| 234 | 
            +
                    dtype: Optional[torch.dtype] = None,
         | 
| 235 | 
            +
                ):
         | 
| 236 | 
            +
                    device = device or self._execution_device
         | 
| 237 | 
            +
                    dtype = dtype or self.text_encoder.dtype
         | 
| 238 | 
            +
             | 
| 239 | 
            +
                    prompt = [prompt] if isinstance(prompt, str) else prompt
         | 
| 240 | 
            +
                    img_prompt_template = "Picture {}: <|vision_start|><|image_pad|><|vision_end|>"
         | 
| 241 | 
            +
                    if isinstance(image, list):
         | 
| 242 | 
            +
                        base_img_prompt = ""
         | 
| 243 | 
            +
                        for i, img in enumerate(image):
         | 
| 244 | 
            +
                            base_img_prompt += img_prompt_template.format(i + 1)
         | 
| 245 | 
            +
                    elif image is not None:
         | 
| 246 | 
            +
                        base_img_prompt = img_prompt_template.format(1)
         | 
| 247 | 
            +
                    else:
         | 
| 248 | 
            +
                        base_img_prompt = ""
         | 
| 249 | 
            +
             | 
| 250 | 
            +
                    template = self.prompt_template_encode
         | 
| 251 | 
            +
             | 
| 252 | 
            +
                    drop_idx = self.prompt_template_encode_start_idx
         | 
| 253 | 
            +
                    txt = [template.format(base_img_prompt + e) for e in prompt]
         | 
| 254 | 
            +
             | 
| 255 | 
            +
                    model_inputs = self.processor(
         | 
| 256 | 
            +
                        text=txt,
         | 
| 257 | 
            +
                        images=image,
         | 
| 258 | 
            +
                        padding=True,
         | 
| 259 | 
            +
                        return_tensors="pt",
         | 
| 260 | 
            +
                    ).to(device)
         | 
| 261 | 
            +
             | 
| 262 | 
            +
                    outputs = self.text_encoder(
         | 
| 263 | 
            +
                        input_ids=model_inputs.input_ids,
         | 
| 264 | 
            +
                        attention_mask=model_inputs.attention_mask,
         | 
| 265 | 
            +
                        pixel_values=model_inputs.pixel_values,
         | 
| 266 | 
            +
                        image_grid_thw=model_inputs.image_grid_thw,
         | 
| 267 | 
            +
                        output_hidden_states=True,
         | 
| 268 | 
            +
                    )
         | 
| 269 | 
            +
             | 
| 270 | 
            +
                    hidden_states = outputs.hidden_states[-1]
         | 
| 271 | 
            +
                    split_hidden_states = self._extract_masked_hidden(hidden_states, model_inputs.attention_mask)
         | 
| 272 | 
            +
                    split_hidden_states = [e[drop_idx:] for e in split_hidden_states]
         | 
| 273 | 
            +
                    attn_mask_list = [torch.ones(e.size(0), dtype=torch.long, device=e.device) for e in split_hidden_states]
         | 
| 274 | 
            +
                    max_seq_len = max([e.size(0) for e in split_hidden_states])
         | 
| 275 | 
            +
                    prompt_embeds = torch.stack(
         | 
| 276 | 
            +
                        [torch.cat([u, u.new_zeros(max_seq_len - u.size(0), u.size(1))]) for u in split_hidden_states]
         | 
| 277 | 
            +
                    )
         | 
| 278 | 
            +
                    encoder_attention_mask = torch.stack(
         | 
| 279 | 
            +
                        [torch.cat([u, u.new_zeros(max_seq_len - u.size(0))]) for u in attn_mask_list]
         | 
| 280 | 
            +
                    )
         | 
| 281 | 
            +
             | 
| 282 | 
            +
                    prompt_embeds = prompt_embeds.to(dtype=dtype, device=device)
         | 
| 283 | 
            +
             | 
| 284 | 
            +
                    return prompt_embeds, encoder_attention_mask
         | 
| 285 | 
            +
             | 
| 286 | 
            +
                # Copied from diffusers.pipelines.qwenimage.pipeline_qwenimage_edit.QwenImageEditPipeline.encode_prompt
         | 
| 287 | 
            +
                def encode_prompt(
         | 
| 288 | 
            +
                    self,
         | 
| 289 | 
            +
                    prompt: Union[str, List[str]],
         | 
| 290 | 
            +
                    image: Optional[torch.Tensor] = None,
         | 
| 291 | 
            +
                    device: Optional[torch.device] = None,
         | 
| 292 | 
            +
                    num_images_per_prompt: int = 1,
         | 
| 293 | 
            +
                    prompt_embeds: Optional[torch.Tensor] = None,
         | 
| 294 | 
            +
                    prompt_embeds_mask: Optional[torch.Tensor] = None,
         | 
| 295 | 
            +
                    max_sequence_length: int = 1024,
         | 
| 296 | 
            +
                ):
         | 
| 297 | 
            +
                    r"""
         | 
| 298 | 
            +
             | 
| 299 | 
            +
                    Args:
         | 
| 300 | 
            +
                        prompt (`str` or `List[str]`, *optional*):
         | 
| 301 | 
            +
                            prompt to be encoded
         | 
| 302 | 
            +
                        image (`torch.Tensor`, *optional*):
         | 
| 303 | 
            +
                            image to be encoded
         | 
| 304 | 
            +
                        device: (`torch.device`):
         | 
| 305 | 
            +
                            torch device
         | 
| 306 | 
            +
                        num_images_per_prompt (`int`):
         | 
| 307 | 
            +
                            number of images that should be generated per prompt
         | 
| 308 | 
            +
                        prompt_embeds (`torch.Tensor`, *optional*):
         | 
| 309 | 
            +
                            Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not
         | 
| 310 | 
            +
                            provided, text embeddings will be generated from `prompt` input argument.
         | 
| 311 | 
            +
                    """
         | 
| 312 | 
            +
                    device = device or self._execution_device
         | 
| 313 | 
            +
             | 
| 314 | 
            +
                    prompt = [prompt] if isinstance(prompt, str) else prompt
         | 
| 315 | 
            +
                    batch_size = len(prompt) if prompt_embeds is None else prompt_embeds.shape[0]
         | 
| 316 | 
            +
             | 
| 317 | 
            +
                    if prompt_embeds is None:
         | 
| 318 | 
            +
                        prompt_embeds, prompt_embeds_mask = self._get_qwen_prompt_embeds(prompt, image, device)
         | 
| 319 | 
            +
             | 
| 320 | 
            +
                    _, seq_len, _ = prompt_embeds.shape
         | 
| 321 | 
            +
                    prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1)
         | 
| 322 | 
            +
                    prompt_embeds = prompt_embeds.view(batch_size * num_images_per_prompt, seq_len, -1)
         | 
| 323 | 
            +
                    prompt_embeds_mask = prompt_embeds_mask.repeat(1, num_images_per_prompt, 1)
         | 
| 324 | 
            +
                    prompt_embeds_mask = prompt_embeds_mask.view(batch_size * num_images_per_prompt, seq_len)
         | 
| 325 | 
            +
             | 
| 326 | 
            +
                    return prompt_embeds, prompt_embeds_mask
         | 
| 327 | 
            +
             | 
| 328 | 
            +
                # Copied from diffusers.pipelines.qwenimage.pipeline_qwenimage_edit.QwenImageEditPipeline.check_inputs
         | 
| 329 | 
            +
                def check_inputs(
         | 
| 330 | 
            +
                    self,
         | 
| 331 | 
            +
                    prompt,
         | 
| 332 | 
            +
                    height,
         | 
| 333 | 
            +
                    width,
         | 
| 334 | 
            +
                    negative_prompt=None,
         | 
| 335 | 
            +
                    prompt_embeds=None,
         | 
| 336 | 
            +
                    negative_prompt_embeds=None,
         | 
| 337 | 
            +
                    prompt_embeds_mask=None,
         | 
| 338 | 
            +
                    negative_prompt_embeds_mask=None,
         | 
| 339 | 
            +
                    callback_on_step_end_tensor_inputs=None,
         | 
| 340 | 
            +
                    max_sequence_length=None,
         | 
| 341 | 
            +
                ):
         | 
| 342 | 
            +
                    if height % (self.vae_scale_factor * 2) != 0 or width % (self.vae_scale_factor * 2) != 0:
         | 
| 343 | 
            +
                        logger.warning(
         | 
| 344 | 
            +
                            f"`height` and `width` have to be divisible by {self.vae_scale_factor * 2} but are {height} and {width}. Dimensions will be resized accordingly"
         | 
| 345 | 
            +
                        )
         | 
| 346 | 
            +
             | 
| 347 | 
            +
                    if callback_on_step_end_tensor_inputs is not None and not all(
         | 
| 348 | 
            +
                        k in self._callback_tensor_inputs for k in callback_on_step_end_tensor_inputs
         | 
| 349 | 
            +
                    ):
         | 
| 350 | 
            +
                        raise ValueError(
         | 
| 351 | 
            +
                            f"`callback_on_step_end_tensor_inputs` has to be in {self._callback_tensor_inputs}, but found {[k for k in callback_on_step_end_tensor_inputs if k not in self._callback_tensor_inputs]}"
         | 
| 352 | 
            +
                        )
         | 
| 353 | 
            +
             | 
| 354 | 
            +
                    if prompt is not None and prompt_embeds is not None:
         | 
| 355 | 
            +
                        raise ValueError(
         | 
| 356 | 
            +
                            f"Cannot forward both `prompt`: {prompt} and `prompt_embeds`: {prompt_embeds}. Please make sure to"
         | 
| 357 | 
            +
                            " only forward one of the two."
         | 
| 358 | 
            +
                        )
         | 
| 359 | 
            +
                    elif prompt is None and prompt_embeds is None:
         | 
| 360 | 
            +
                        raise ValueError(
         | 
| 361 | 
            +
                            "Provide either `prompt` or `prompt_embeds`. Cannot leave both `prompt` and `prompt_embeds` undefined."
         | 
| 362 | 
            +
                        )
         | 
| 363 | 
            +
                    elif prompt is not None and (not isinstance(prompt, str) and not isinstance(prompt, list)):
         | 
| 364 | 
            +
                        raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}")
         | 
| 365 | 
            +
             | 
| 366 | 
            +
                    if negative_prompt is not None and negative_prompt_embeds is not None:
         | 
| 367 | 
            +
                        raise ValueError(
         | 
| 368 | 
            +
                            f"Cannot forward both `negative_prompt`: {negative_prompt} and `negative_prompt_embeds`:"
         | 
| 369 | 
            +
                            f" {negative_prompt_embeds}. Please make sure to only forward one of the two."
         | 
| 370 | 
            +
                        )
         | 
| 371 | 
            +
             | 
| 372 | 
            +
                    if prompt_embeds is not None and prompt_embeds_mask is None:
         | 
| 373 | 
            +
                        raise ValueError(
         | 
| 374 | 
            +
                            "If `prompt_embeds` are provided, `prompt_embeds_mask` also have to be passed. Make sure to generate `prompt_embeds_mask` from the same text encoder that was used to generate `prompt_embeds`."
         | 
| 375 | 
            +
                        )
         | 
| 376 | 
            +
                    if negative_prompt_embeds is not None and negative_prompt_embeds_mask is None:
         | 
| 377 | 
            +
                        raise ValueError(
         | 
| 378 | 
            +
                            "If `negative_prompt_embeds` are provided, `negative_prompt_embeds_mask` also have to be passed. Make sure to generate `negative_prompt_embeds_mask` from the same text encoder that was used to generate `negative_prompt_embeds`."
         | 
| 379 | 
            +
                        )
         | 
| 380 | 
            +
             | 
| 381 | 
            +
                    if max_sequence_length is not None and max_sequence_length > 1024:
         | 
| 382 | 
            +
                        raise ValueError(f"`max_sequence_length` cannot be greater than 1024 but is {max_sequence_length}")
         | 
| 383 | 
            +
             | 
| 384 | 
            +
                @staticmethod
         | 
| 385 | 
            +
                # Copied from diffusers.pipelines.qwenimage.pipeline_qwenimage.QwenImagePipeline._pack_latents
         | 
| 386 | 
            +
                def _pack_latents(latents, batch_size, num_channels_latents, height, width):
         | 
| 387 | 
            +
                    latents = latents.view(batch_size, num_channels_latents, height // 2, 2, width // 2, 2)
         | 
| 388 | 
            +
                    latents = latents.permute(0, 2, 4, 1, 3, 5)
         | 
| 389 | 
            +
                    latents = latents.reshape(batch_size, (height // 2) * (width // 2), num_channels_latents * 4)
         | 
| 390 | 
            +
             | 
| 391 | 
            +
                    return latents
         | 
| 392 | 
            +
             | 
| 393 | 
            +
                @staticmethod
         | 
| 394 | 
            +
                # Copied from diffusers.pipelines.qwenimage.pipeline_qwenimage.QwenImagePipeline._unpack_latents
         | 
| 395 | 
            +
                def _unpack_latents(latents, height, width, vae_scale_factor):
         | 
| 396 | 
            +
                    batch_size, num_patches, channels = latents.shape
         | 
| 397 | 
            +
             | 
| 398 | 
            +
                    # VAE applies 8x compression on images but we must also account for packing which requires
         | 
| 399 | 
            +
                    # latent height and width to be divisible by 2.
         | 
| 400 | 
            +
                    height = 2 * (int(height) // (vae_scale_factor * 2))
         | 
| 401 | 
            +
                    width = 2 * (int(width) // (vae_scale_factor * 2))
         | 
| 402 | 
            +
             | 
| 403 | 
            +
                    latents = latents.view(batch_size, height // 2, width // 2, channels // 4, 2, 2)
         | 
| 404 | 
            +
                    latents = latents.permute(0, 3, 1, 4, 2, 5)
         | 
| 405 | 
            +
             | 
| 406 | 
            +
                    latents = latents.reshape(batch_size, channels // (2 * 2), 1, height, width)
         | 
| 407 | 
            +
             | 
| 408 | 
            +
                    return latents
         | 
| 409 | 
            +
             | 
| 410 | 
            +
                # Copied from diffusers.pipelines.qwenimage.pipeline_qwenimage_edit.QwenImageEditPipeline._encode_vae_image
         | 
| 411 | 
            +
                def _encode_vae_image(self, image: torch.Tensor, generator: torch.Generator):
         | 
| 412 | 
            +
                    if isinstance(generator, list):
         | 
| 413 | 
            +
                        image_latents = [
         | 
| 414 | 
            +
                            retrieve_latents(self.vae.encode(image[i : i + 1]), generator=generator[i], sample_mode="argmax")
         | 
| 415 | 
            +
                            for i in range(image.shape[0])
         | 
| 416 | 
            +
                        ]
         | 
| 417 | 
            +
                        image_latents = torch.cat(image_latents, dim=0)
         | 
| 418 | 
            +
                    else:
         | 
| 419 | 
            +
                        image_latents = retrieve_latents(self.vae.encode(image), generator=generator, sample_mode="argmax")
         | 
| 420 | 
            +
                    latents_mean = (
         | 
| 421 | 
            +
                        torch.tensor(self.vae.config.latents_mean)
         | 
| 422 | 
            +
                        .view(1, self.latent_channels, 1, 1, 1)
         | 
| 423 | 
            +
                        .to(image_latents.device, image_latents.dtype)
         | 
| 424 | 
            +
                    )
         | 
| 425 | 
            +
                    latents_std = (
         | 
| 426 | 
            +
                        torch.tensor(self.vae.config.latents_std)
         | 
| 427 | 
            +
                        .view(1, self.latent_channels, 1, 1, 1)
         | 
| 428 | 
            +
                        .to(image_latents.device, image_latents.dtype)
         | 
| 429 | 
            +
                    )
         | 
| 430 | 
            +
                    image_latents = (image_latents - latents_mean) / latents_std
         | 
| 431 | 
            +
             | 
| 432 | 
            +
                    return image_latents
         | 
| 433 | 
            +
             | 
| 434 | 
            +
                def prepare_latents(
         | 
| 435 | 
            +
                    self,
         | 
| 436 | 
            +
                    images,
         | 
| 437 | 
            +
                    batch_size,
         | 
| 438 | 
            +
                    num_channels_latents,
         | 
| 439 | 
            +
                    height,
         | 
| 440 | 
            +
                    width,
         | 
| 441 | 
            +
                    dtype,
         | 
| 442 | 
            +
                    device,
         | 
| 443 | 
            +
                    generator,
         | 
| 444 | 
            +
                    latents=None,
         | 
| 445 | 
            +
                ):
         | 
| 446 | 
            +
                    # VAE applies 8x compression on images but we must also account for packing which requires
         | 
| 447 | 
            +
                    # latent height and width to be divisible by 2.
         | 
| 448 | 
            +
                    height = 2 * (int(height) // (self.vae_scale_factor * 2))
         | 
| 449 | 
            +
                    width = 2 * (int(width) // (self.vae_scale_factor * 2))
         | 
| 450 | 
            +
             | 
| 451 | 
            +
                    shape = (batch_size, 1, num_channels_latents, height, width)
         | 
| 452 | 
            +
             | 
| 453 | 
            +
                    image_latents = None
         | 
| 454 | 
            +
                    if images is not None:
         | 
| 455 | 
            +
                        if not isinstance(images, list):
         | 
| 456 | 
            +
                            images = [images]
         | 
| 457 | 
            +
                        all_image_latents = []
         | 
| 458 | 
            +
                        for image in images:
         | 
| 459 | 
            +
                            image = image.to(device=device, dtype=dtype)
         | 
| 460 | 
            +
                            if image.shape[1] != self.latent_channels:
         | 
| 461 | 
            +
                                image_latents = self._encode_vae_image(image=image, generator=generator)
         | 
| 462 | 
            +
                            else:
         | 
| 463 | 
            +
                                image_latents = image
         | 
| 464 | 
            +
                            if batch_size > image_latents.shape[0] and batch_size % image_latents.shape[0] == 0:
         | 
| 465 | 
            +
                                # expand init_latents for batch_size
         | 
| 466 | 
            +
                                additional_image_per_prompt = batch_size // image_latents.shape[0]
         | 
| 467 | 
            +
                                image_latents = torch.cat([image_latents] * additional_image_per_prompt, dim=0)
         | 
| 468 | 
            +
                            elif batch_size > image_latents.shape[0] and batch_size % image_latents.shape[0] != 0:
         | 
| 469 | 
            +
                                raise ValueError(
         | 
| 470 | 
            +
                                    f"Cannot duplicate `image` of batch size {image_latents.shape[0]} to {batch_size} text prompts."
         | 
| 471 | 
            +
                                )
         | 
| 472 | 
            +
                            else:
         | 
| 473 | 
            +
                                image_latents = torch.cat([image_latents], dim=0)
         | 
| 474 | 
            +
             | 
| 475 | 
            +
                            image_latent_height, image_latent_width = image_latents.shape[3:]
         | 
| 476 | 
            +
                            image_latents = self._pack_latents(
         | 
| 477 | 
            +
                                image_latents, batch_size, num_channels_latents, image_latent_height, image_latent_width
         | 
| 478 | 
            +
                            )
         | 
| 479 | 
            +
                            all_image_latents.append(image_latents)
         | 
| 480 | 
            +
                        image_latents = torch.cat(all_image_latents, dim=1)
         | 
| 481 | 
            +
             | 
| 482 | 
            +
                    if isinstance(generator, list) and len(generator) != batch_size:
         | 
| 483 | 
            +
                        raise ValueError(
         | 
| 484 | 
            +
                            f"You have passed a list of generators of length {len(generator)}, but requested an effective batch"
         | 
| 485 | 
            +
                            f" size of {batch_size}. Make sure the batch size matches the length of the generators."
         | 
| 486 | 
            +
                        )
         | 
| 487 | 
            +
                    if latents is None:
         | 
| 488 | 
            +
                        latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
         | 
| 489 | 
            +
                        latents = self._pack_latents(latents, batch_size, num_channels_latents, height, width)
         | 
| 490 | 
            +
                    else:
         | 
| 491 | 
            +
                        latents = latents.to(device=device, dtype=dtype)
         | 
| 492 | 
            +
             | 
| 493 | 
            +
                    return latents, image_latents
         | 
| 494 | 
            +
             | 
| 495 | 
            +
                @property
         | 
| 496 | 
            +
                def guidance_scale(self):
         | 
| 497 | 
            +
                    return self._guidance_scale
         | 
| 498 | 
            +
             | 
| 499 | 
            +
                @property
         | 
| 500 | 
            +
                def attention_kwargs(self):
         | 
| 501 | 
            +
                    return self._attention_kwargs
         | 
| 502 | 
            +
             | 
| 503 | 
            +
                @property
         | 
| 504 | 
            +
                def num_timesteps(self):
         | 
| 505 | 
            +
                    return self._num_timesteps
         | 
| 506 | 
            +
             | 
| 507 | 
            +
                @property
         | 
| 508 | 
            +
                def current_timestep(self):
         | 
| 509 | 
            +
                    return self._current_timestep
         | 
| 510 | 
            +
             | 
| 511 | 
            +
                @property
         | 
| 512 | 
            +
                def interrupt(self):
         | 
| 513 | 
            +
                    return self._interrupt
         | 
| 514 | 
            +
             | 
| 515 | 
            +
                @torch.no_grad()
         | 
| 516 | 
            +
                @replace_example_docstring(EXAMPLE_DOC_STRING)
         | 
| 517 | 
            +
                def __call__(
         | 
| 518 | 
            +
                    self,
         | 
| 519 | 
            +
                    image: Optional[PipelineImageInput] = None,
         | 
| 520 | 
            +
                    prompt: Union[str, List[str]] = None,
         | 
| 521 | 
            +
                    negative_prompt: Union[str, List[str]] = None,
         | 
| 522 | 
            +
                    true_cfg_scale: float = 4.0,
         | 
| 523 | 
            +
                    height: Optional[int] = None,
         | 
| 524 | 
            +
                    width: Optional[int] = None,
         | 
| 525 | 
            +
                    num_inference_steps: int = 50,
         | 
| 526 | 
            +
                    sigmas: Optional[List[float]] = None,
         | 
| 527 | 
            +
                    guidance_scale: Optional[float] = None,
         | 
| 528 | 
            +
                    num_images_per_prompt: int = 1,
         | 
| 529 | 
            +
                    generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
         | 
| 530 | 
            +
                    latents: Optional[torch.Tensor] = None,
         | 
| 531 | 
            +
                    prompt_embeds: Optional[torch.Tensor] = None,
         | 
| 532 | 
            +
                    prompt_embeds_mask: Optional[torch.Tensor] = None,
         | 
| 533 | 
            +
                    negative_prompt_embeds: Optional[torch.Tensor] = None,
         | 
| 534 | 
            +
                    negative_prompt_embeds_mask: Optional[torch.Tensor] = None,
         | 
| 535 | 
            +
                    output_type: Optional[str] = "pil",
         | 
| 536 | 
            +
                    return_dict: bool = True,
         | 
| 537 | 
            +
                    attention_kwargs: Optional[Dict[str, Any]] = None,
         | 
| 538 | 
            +
                    callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
         | 
| 539 | 
            +
                    callback_on_step_end_tensor_inputs: List[str] = ["latents"],
         | 
| 540 | 
            +
                    max_sequence_length: int = 512,
         | 
| 541 | 
            +
                ):
         | 
| 542 | 
            +
                    r"""
         | 
| 543 | 
            +
                    Function invoked when calling the pipeline for generation.
         | 
| 544 | 
            +
             | 
| 545 | 
            +
                    Args:
         | 
| 546 | 
            +
                        image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `List[torch.Tensor]`, `List[PIL.Image.Image]`, or `List[np.ndarray]`):
         | 
| 547 | 
            +
                            `Image`, numpy array or tensor representing an image batch to be used as the starting point. For both
         | 
| 548 | 
            +
                            numpy array and pytorch tensor, the expected value range is between `[0, 1]` If it's a tensor or a list
         | 
| 549 | 
            +
                            or tensors, the expected shape should be `(B, C, H, W)` or `(C, H, W)`. If it is a numpy array or a
         | 
| 550 | 
            +
                            list of arrays, the expected shape should be `(B, H, W, C)` or `(H, W, C)` It can also accept image
         | 
| 551 | 
            +
                            latents as `image`, but if passing latents directly it is not encoded again.
         | 
| 552 | 
            +
                        prompt (`str` or `List[str]`, *optional*):
         | 
| 553 | 
            +
                            The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.
         | 
| 554 | 
            +
                            instead.
         | 
| 555 | 
            +
                        negative_prompt (`str` or `List[str]`, *optional*):
         | 
| 556 | 
            +
                            The prompt or prompts not to guide the image generation. If not defined, one has to pass
         | 
| 557 | 
            +
                            `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `true_cfg_scale` is
         | 
| 558 | 
            +
                            not greater than `1`).
         | 
| 559 | 
            +
                        true_cfg_scale (`float`, *optional*, defaults to 1.0):
         | 
| 560 | 
            +
                            true_cfg_scale (`float`, *optional*, defaults to 1.0): Guidance scale as defined in [Classifier-Free
         | 
| 561 | 
            +
                            Diffusion Guidance](https://huggingface.co/papers/2207.12598). `true_cfg_scale` is defined as `w` of
         | 
| 562 | 
            +
                            equation 2. of [Imagen Paper](https://huggingface.co/papers/2205.11487). Classifier-free guidance is
         | 
| 563 | 
            +
                            enabled by setting `true_cfg_scale > 1` and a provided `negative_prompt`. Higher guidance scale
         | 
| 564 | 
            +
                            encourages to generate images that are closely linked to the text `prompt`, usually at the expense of
         | 
| 565 | 
            +
                            lower image quality.
         | 
| 566 | 
            +
                        height (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor):
         | 
| 567 | 
            +
                            The height in pixels of the generated image. This is set to 1024 by default for the best results.
         | 
| 568 | 
            +
                        width (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor):
         | 
| 569 | 
            +
                            The width in pixels of the generated image. This is set to 1024 by default for the best results.
         | 
| 570 | 
            +
                        num_inference_steps (`int`, *optional*, defaults to 50):
         | 
| 571 | 
            +
                            The number of denoising steps. More denoising steps usually lead to a higher quality image at the
         | 
| 572 | 
            +
                            expense of slower inference.
         | 
| 573 | 
            +
                        sigmas (`List[float]`, *optional*):
         | 
| 574 | 
            +
                            Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in
         | 
| 575 | 
            +
                            their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed
         | 
| 576 | 
            +
                            will be used.
         | 
| 577 | 
            +
                        guidance_scale (`float`, *optional*, defaults to None):
         | 
| 578 | 
            +
                            A guidance scale value for guidance distilled models. Unlike the traditional classifier-free guidance
         | 
| 579 | 
            +
                            where the guidance scale is applied during inference through noise prediction rescaling, guidance
         | 
| 580 | 
            +
                            distilled models take the guidance scale directly as an input parameter during forward pass. Guidance
         | 
| 581 | 
            +
                            scale is enabled by setting `guidance_scale > 1`. Higher guidance scale encourages to generate images
         | 
| 582 | 
            +
                            that are closely linked to the text `prompt`, usually at the expense of lower image quality. This
         | 
| 583 | 
            +
                            parameter in the pipeline is there to support future guidance-distilled models when they come up. It is
         | 
| 584 | 
            +
                            ignored when not using guidance distilled models. To enable traditional classifier-free guidance,
         | 
| 585 | 
            +
                            please pass `true_cfg_scale > 1.0` and `negative_prompt` (even an empty negative prompt like " " should
         | 
| 586 | 
            +
                            enable classifier-free guidance computations).
         | 
| 587 | 
            +
                        num_images_per_prompt (`int`, *optional*, defaults to 1):
         | 
| 588 | 
            +
                            The number of images to generate per prompt.
         | 
| 589 | 
            +
                        generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
         | 
| 590 | 
            +
                            One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
         | 
| 591 | 
            +
                            to make generation deterministic.
         | 
| 592 | 
            +
                        latents (`torch.Tensor`, *optional*):
         | 
| 593 | 
            +
                            Pre-generated noisy latents, sampled from a Gaussian distribution, to be used as inputs for image
         | 
| 594 | 
            +
                            generation. Can be used to tweak the same generation with different prompts. If not provided, a latents
         | 
| 595 | 
            +
                            tensor will be generated by sampling using the supplied random `generator`.
         | 
| 596 | 
            +
                        prompt_embeds (`torch.Tensor`, *optional*):
         | 
| 597 | 
            +
                            Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not
         | 
| 598 | 
            +
                            provided, text embeddings will be generated from `prompt` input argument.
         | 
| 599 | 
            +
                        negative_prompt_embeds (`torch.Tensor`, *optional*):
         | 
| 600 | 
            +
                            Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt
         | 
| 601 | 
            +
                            weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input
         | 
| 602 | 
            +
                            argument.
         | 
| 603 | 
            +
                        output_type (`str`, *optional*, defaults to `"pil"`):
         | 
| 604 | 
            +
                            The output format of the generate image. Choose between
         | 
| 605 | 
            +
                            [PIL](https://pillow.readthedocs.io/en/stable/): `PIL.Image.Image` or `np.array`.
         | 
| 606 | 
            +
                        return_dict (`bool`, *optional*, defaults to `True`):
         | 
| 607 | 
            +
                            Whether or not to return a [`~pipelines.qwenimage.QwenImagePipelineOutput`] instead of a plain tuple.
         | 
| 608 | 
            +
                        attention_kwargs (`dict`, *optional*):
         | 
| 609 | 
            +
                            A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under
         | 
| 610 | 
            +
                            `self.processor` in
         | 
| 611 | 
            +
                            [diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py).
         | 
| 612 | 
            +
                        callback_on_step_end (`Callable`, *optional*):
         | 
| 613 | 
            +
                            A function that calls at the end of each denoising steps during the inference. The function is called
         | 
| 614 | 
            +
                            with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int,
         | 
| 615 | 
            +
                            callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by
         | 
| 616 | 
            +
                            `callback_on_step_end_tensor_inputs`.
         | 
| 617 | 
            +
                        callback_on_step_end_tensor_inputs (`List`, *optional*):
         | 
| 618 | 
            +
                            The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list
         | 
| 619 | 
            +
                            will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the
         | 
| 620 | 
            +
                            `._callback_tensor_inputs` attribute of your pipeline class.
         | 
| 621 | 
            +
                        max_sequence_length (`int` defaults to 512): Maximum sequence length to use with the `prompt`.
         | 
| 622 | 
            +
             | 
| 623 | 
            +
                    Examples:
         | 
| 624 | 
            +
             | 
| 625 | 
            +
                    Returns:
         | 
| 626 | 
            +
                        [`~pipelines.qwenimage.QwenImagePipelineOutput`] or `tuple`:
         | 
| 627 | 
            +
                        [`~pipelines.qwenimage.QwenImagePipelineOutput`] if `return_dict` is True, otherwise a `tuple`. When
         | 
| 628 | 
            +
                        returning a tuple, the first element is a list with the generated images.
         | 
| 629 | 
            +
                    """
         | 
| 630 | 
            +
                    image_size = image[-1].size if isinstance(image, list) else image.size
         | 
| 631 | 
            +
                    calculated_width, calculated_height = calculate_dimensions(1024 * 1024, image_size[0] / image_size[1])
         | 
| 632 | 
            +
                    height = height or calculated_height
         | 
| 633 | 
            +
                    width = width or calculated_width
         | 
| 634 | 
            +
             | 
| 635 | 
            +
                    multiple_of = self.vae_scale_factor * 2
         | 
| 636 | 
            +
                    width = width // multiple_of * multiple_of
         | 
| 637 | 
            +
                    height = height // multiple_of * multiple_of
         | 
| 638 | 
            +
             | 
| 639 | 
            +
                    # 1. Check inputs. Raise error if not correct
         | 
| 640 | 
            +
                    self.check_inputs(
         | 
| 641 | 
            +
                        prompt,
         | 
| 642 | 
            +
                        height,
         | 
| 643 | 
            +
                        width,
         | 
| 644 | 
            +
                        negative_prompt=negative_prompt,
         | 
| 645 | 
            +
                        prompt_embeds=prompt_embeds,
         | 
| 646 | 
            +
                        negative_prompt_embeds=negative_prompt_embeds,
         | 
| 647 | 
            +
                        prompt_embeds_mask=prompt_embeds_mask,
         | 
| 648 | 
            +
                        negative_prompt_embeds_mask=negative_prompt_embeds_mask,
         | 
| 649 | 
            +
                        callback_on_step_end_tensor_inputs=callback_on_step_end_tensor_inputs,
         | 
| 650 | 
            +
                        max_sequence_length=max_sequence_length,
         | 
| 651 | 
            +
                    )
         | 
| 652 | 
            +
             | 
| 653 | 
            +
                    self._guidance_scale = guidance_scale
         | 
| 654 | 
            +
                    self._attention_kwargs = attention_kwargs
         | 
| 655 | 
            +
                    self._current_timestep = None
         | 
| 656 | 
            +
                    self._interrupt = False
         | 
| 657 | 
            +
             | 
| 658 | 
            +
                    # 2. Define call parameters
         | 
| 659 | 
            +
                    if prompt is not None and isinstance(prompt, str):
         | 
| 660 | 
            +
                        batch_size = 1
         | 
| 661 | 
            +
                    elif prompt is not None and isinstance(prompt, list):
         | 
| 662 | 
            +
                        batch_size = len(prompt)
         | 
| 663 | 
            +
                    else:
         | 
| 664 | 
            +
                        batch_size = prompt_embeds.shape[0]
         | 
| 665 | 
            +
             | 
| 666 | 
            +
                    device = self._execution_device
         | 
| 667 | 
            +
                    # 3. Preprocess image
         | 
| 668 | 
            +
                    if image is not None and not (isinstance(image, torch.Tensor) and image.size(1) == self.latent_channels):
         | 
| 669 | 
            +
                        if not isinstance(image, list):
         | 
| 670 | 
            +
                            image = [image]
         | 
| 671 | 
            +
                        condition_image_sizes = []
         | 
| 672 | 
            +
                        condition_images = []
         | 
| 673 | 
            +
                        vae_image_sizes = []
         | 
| 674 | 
            +
                        vae_images = []
         | 
| 675 | 
            +
                        for img in image:
         | 
| 676 | 
            +
                            image_width, image_height = img.size
         | 
| 677 | 
            +
                            condition_width, condition_height = calculate_dimensions(
         | 
| 678 | 
            +
                                CONDITION_IMAGE_SIZE, image_width / image_height
         | 
| 679 | 
            +
                            )
         | 
| 680 | 
            +
                            vae_width, vae_height = calculate_dimensions(VAE_IMAGE_SIZE, image_width / image_height)
         | 
| 681 | 
            +
                            condition_image_sizes.append((condition_width, condition_height))
         | 
| 682 | 
            +
                            vae_image_sizes.append((vae_width, vae_height))
         | 
| 683 | 
            +
                            condition_images.append(self.image_processor.resize(img, condition_height, condition_width))
         | 
| 684 | 
            +
                            vae_images.append(self.image_processor.preprocess(img, vae_height, vae_width).unsqueeze(2))
         | 
| 685 | 
            +
             | 
| 686 | 
            +
                    has_neg_prompt = negative_prompt is not None or (
         | 
| 687 | 
            +
                        negative_prompt_embeds is not None and negative_prompt_embeds_mask is not None
         | 
| 688 | 
            +
                    )
         | 
| 689 | 
            +
             | 
| 690 | 
            +
                    if true_cfg_scale > 1 and not has_neg_prompt:
         | 
| 691 | 
            +
                        logger.warning(
         | 
| 692 | 
            +
                            f"true_cfg_scale is passed as {true_cfg_scale}, but classifier-free guidance is not enabled since no negative_prompt is provided."
         | 
| 693 | 
            +
                        )
         | 
| 694 | 
            +
                    elif true_cfg_scale <= 1 and has_neg_prompt:
         | 
| 695 | 
            +
                        logger.warning(
         | 
| 696 | 
            +
                            " negative_prompt is passed but classifier-free guidance is not enabled since true_cfg_scale <= 1"
         | 
| 697 | 
            +
                        )
         | 
| 698 | 
            +
             | 
| 699 | 
            +
                    do_true_cfg = true_cfg_scale > 1 and has_neg_prompt
         | 
| 700 | 
            +
                    prompt_embeds, prompt_embeds_mask = self.encode_prompt(
         | 
| 701 | 
            +
                        image=condition_images,
         | 
| 702 | 
            +
                        prompt=prompt,
         | 
| 703 | 
            +
                        prompt_embeds=prompt_embeds,
         | 
| 704 | 
            +
                        prompt_embeds_mask=prompt_embeds_mask,
         | 
| 705 | 
            +
                        device=device,
         | 
| 706 | 
            +
                        num_images_per_prompt=num_images_per_prompt,
         | 
| 707 | 
            +
                        max_sequence_length=max_sequence_length,
         | 
| 708 | 
            +
                    )
         | 
| 709 | 
            +
                    if do_true_cfg:
         | 
| 710 | 
            +
                        negative_prompt_embeds, negative_prompt_embeds_mask = self.encode_prompt(
         | 
| 711 | 
            +
                            image=condition_images,
         | 
| 712 | 
            +
                            prompt=negative_prompt,
         | 
| 713 | 
            +
                            prompt_embeds=negative_prompt_embeds,
         | 
| 714 | 
            +
                            prompt_embeds_mask=negative_prompt_embeds_mask,
         | 
| 715 | 
            +
                            device=device,
         | 
| 716 | 
            +
                            num_images_per_prompt=num_images_per_prompt,
         | 
| 717 | 
            +
                            max_sequence_length=max_sequence_length,
         | 
| 718 | 
            +
                        )
         | 
| 719 | 
            +
             | 
| 720 | 
            +
                    # 4. Prepare latent variables
         | 
| 721 | 
            +
                    num_channels_latents = self.transformer.config.in_channels // 4
         | 
| 722 | 
            +
                    latents, image_latents = self.prepare_latents(
         | 
| 723 | 
            +
                        vae_images,
         | 
| 724 | 
            +
                        batch_size * num_images_per_prompt,
         | 
| 725 | 
            +
                        num_channels_latents,
         | 
| 726 | 
            +
                        height,
         | 
| 727 | 
            +
                        width,
         | 
| 728 | 
            +
                        prompt_embeds.dtype,
         | 
| 729 | 
            +
                        device,
         | 
| 730 | 
            +
                        generator,
         | 
| 731 | 
            +
                        latents,
         | 
| 732 | 
            +
                    )
         | 
| 733 | 
            +
                    img_shapes = [
         | 
| 734 | 
            +
                        [
         | 
| 735 | 
            +
                            (1, height // self.vae_scale_factor // 2, width // self.vae_scale_factor // 2),
         | 
| 736 | 
            +
                            *[
         | 
| 737 | 
            +
                                (1, vae_height // self.vae_scale_factor // 2, vae_width // self.vae_scale_factor // 2)
         | 
| 738 | 
            +
                                for vae_width, vae_height in vae_image_sizes
         | 
| 739 | 
            +
                            ],
         | 
| 740 | 
            +
                        ]
         | 
| 741 | 
            +
                    ] * batch_size
         | 
| 742 | 
            +
             | 
| 743 | 
            +
                    # 5. Prepare timesteps
         | 
| 744 | 
            +
                    sigmas = np.linspace(1.0, 1 / num_inference_steps, num_inference_steps) if sigmas is None else sigmas
         | 
| 745 | 
            +
                    image_seq_len = latents.shape[1]
         | 
| 746 | 
            +
                    mu = calculate_shift(
         | 
| 747 | 
            +
                        image_seq_len,
         | 
| 748 | 
            +
                        self.scheduler.config.get("base_image_seq_len", 256),
         | 
| 749 | 
            +
                        self.scheduler.config.get("max_image_seq_len", 4096),
         | 
| 750 | 
            +
                        self.scheduler.config.get("base_shift", 0.5),
         | 
| 751 | 
            +
                        self.scheduler.config.get("max_shift", 1.15),
         | 
| 752 | 
            +
                    )
         | 
| 753 | 
            +
                    timesteps, num_inference_steps = retrieve_timesteps(
         | 
| 754 | 
            +
                        self.scheduler,
         | 
| 755 | 
            +
                        num_inference_steps,
         | 
| 756 | 
            +
                        device,
         | 
| 757 | 
            +
                        sigmas=sigmas,
         | 
| 758 | 
            +
                        mu=mu,
         | 
| 759 | 
            +
                    )
         | 
| 760 | 
            +
                    num_warmup_steps = max(len(timesteps) - num_inference_steps * self.scheduler.order, 0)
         | 
| 761 | 
            +
                    self._num_timesteps = len(timesteps)
         | 
| 762 | 
            +
             | 
| 763 | 
            +
                    # handle guidance
         | 
| 764 | 
            +
                    if self.transformer.config.guidance_embeds and guidance_scale is None:
         | 
| 765 | 
            +
                        raise ValueError("guidance_scale is required for guidance-distilled model.")
         | 
| 766 | 
            +
                    elif self.transformer.config.guidance_embeds:
         | 
| 767 | 
            +
                        guidance = torch.full([1], guidance_scale, device=device, dtype=torch.float32)
         | 
| 768 | 
            +
                        guidance = guidance.expand(latents.shape[0])
         | 
| 769 | 
            +
                    elif not self.transformer.config.guidance_embeds and guidance_scale is not None:
         | 
| 770 | 
            +
                        logger.warning(
         | 
| 771 | 
            +
                            f"guidance_scale is passed as {guidance_scale}, but ignored since the model is not guidance-distilled."
         | 
| 772 | 
            +
                        )
         | 
| 773 | 
            +
                        guidance = None
         | 
| 774 | 
            +
                    elif not self.transformer.config.guidance_embeds and guidance_scale is None:
         | 
| 775 | 
            +
                        guidance = None
         | 
| 776 | 
            +
             | 
| 777 | 
            +
                    if self.attention_kwargs is None:
         | 
| 778 | 
            +
                        self._attention_kwargs = {}
         | 
| 779 | 
            +
             | 
| 780 | 
            +
                    txt_seq_lens = prompt_embeds_mask.sum(dim=1).tolist() if prompt_embeds_mask is not None else None
         | 
| 781 | 
            +
                    
         | 
| 782 | 
            +
                    image_rotary_emb = self.transformer.pos_embed(img_shapes, txt_seq_lens, device=latents.device)
         | 
| 783 | 
            +
                    if do_true_cfg:
         | 
| 784 | 
            +
                        negative_txt_seq_lens = (
         | 
| 785 | 
            +
                            negative_prompt_embeds_mask.sum(dim=1).tolist()
         | 
| 786 | 
            +
                            if negative_prompt_embeds_mask is not None
         | 
| 787 | 
            +
                            else None
         | 
| 788 | 
            +
                        )
         | 
| 789 | 
            +
                        uncond_image_rotary_emb = self.transformer.pos_embed(
         | 
| 790 | 
            +
                            img_shapes, negative_txt_seq_lens, device=latents.device
         | 
| 791 | 
            +
                        )
         | 
| 792 | 
            +
                    else:
         | 
| 793 | 
            +
                        uncond_image_rotary_emb = None
         | 
| 794 | 
            +
             | 
| 795 | 
            +
                    # 6. Denoising loop
         | 
| 796 | 
            +
                    self.scheduler.set_begin_index(0)
         | 
| 797 | 
            +
                    with self.progress_bar(total=num_inference_steps) as progress_bar:
         | 
| 798 | 
            +
                        for i, t in enumerate(timesteps):
         | 
| 799 | 
            +
                            if self.interrupt:
         | 
| 800 | 
            +
                                continue
         | 
| 801 | 
            +
             | 
| 802 | 
            +
                            self._current_timestep = t
         | 
| 803 | 
            +
             | 
| 804 | 
            +
                            latent_model_input = latents
         | 
| 805 | 
            +
                            if image_latents is not None:
         | 
| 806 | 
            +
                                latent_model_input = torch.cat([latents, image_latents], dim=1)
         | 
| 807 | 
            +
             | 
| 808 | 
            +
                            # broadcast to batch dimension in a way that's compatible with ONNX/Core ML
         | 
| 809 | 
            +
                            timestep = t.expand(latents.shape[0]).to(latents.dtype)
         | 
| 810 | 
            +
                            with self.transformer.cache_context("cond"):
         | 
| 811 | 
            +
                                noise_pred = self.transformer(
         | 
| 812 | 
            +
                                    hidden_states=latent_model_input,
         | 
| 813 | 
            +
                                    timestep=timestep / 1000,
         | 
| 814 | 
            +
                                    guidance=guidance,
         | 
| 815 | 
            +
                                    encoder_hidden_states_mask=prompt_embeds_mask,
         | 
| 816 | 
            +
                                    encoder_hidden_states=prompt_embeds,
         | 
| 817 | 
            +
                                    image_rotary_emb=image_rotary_emb,
         | 
| 818 | 
            +
                                    attention_kwargs=self.attention_kwargs,
         | 
| 819 | 
            +
                                    return_dict=False,
         | 
| 820 | 
            +
                                )[0]
         | 
| 821 | 
            +
                                noise_pred = noise_pred[:, : latents.size(1)]
         | 
| 822 | 
            +
             | 
| 823 | 
            +
                            if do_true_cfg:
         | 
| 824 | 
            +
                                with self.transformer.cache_context("uncond"):
         | 
| 825 | 
            +
                                    neg_noise_pred = self.transformer(
         | 
| 826 | 
            +
                                        hidden_states=latent_model_input,
         | 
| 827 | 
            +
                                        timestep=timestep / 1000,
         | 
| 828 | 
            +
                                        guidance=guidance,
         | 
| 829 | 
            +
                                        encoder_hidden_states_mask=negative_prompt_embeds_mask,
         | 
| 830 | 
            +
                                        encoder_hidden_states=negative_prompt_embeds,
         | 
| 831 | 
            +
                                        image_rotary_emb=uncond_image_rotary_emb,
         | 
| 832 | 
            +
                                        attention_kwargs=self.attention_kwargs,
         | 
| 833 | 
            +
                                        return_dict=False,
         | 
| 834 | 
            +
                                    )[0]
         | 
| 835 | 
            +
                                neg_noise_pred = neg_noise_pred[:, : latents.size(1)]
         | 
| 836 | 
            +
                                comb_pred = neg_noise_pred + true_cfg_scale * (noise_pred - neg_noise_pred)
         | 
| 837 | 
            +
             | 
| 838 | 
            +
                                cond_norm = torch.norm(noise_pred, dim=-1, keepdim=True)
         | 
| 839 | 
            +
                                noise_norm = torch.norm(comb_pred, dim=-1, keepdim=True)
         | 
| 840 | 
            +
                                noise_pred = comb_pred * (cond_norm / noise_norm)
         | 
| 841 | 
            +
             | 
| 842 | 
            +
                            # compute the previous noisy sample x_t -> x_t-1
         | 
| 843 | 
            +
                            latents_dtype = latents.dtype
         | 
| 844 | 
            +
                            latents = self.scheduler.step(noise_pred, t, latents, return_dict=False)[0]
         | 
| 845 | 
            +
             | 
| 846 | 
            +
                            if latents.dtype != latents_dtype:
         | 
| 847 | 
            +
                                if torch.backends.mps.is_available():
         | 
| 848 | 
            +
                                    # some platforms (eg. apple mps) misbehave due to a pytorch bug: https://github.com/pytorch/pytorch/pull/99272
         | 
| 849 | 
            +
                                    latents = latents.to(latents_dtype)
         | 
| 850 | 
            +
             | 
| 851 | 
            +
                            if callback_on_step_end is not None:
         | 
| 852 | 
            +
                                callback_kwargs = {}
         | 
| 853 | 
            +
                                for k in callback_on_step_end_tensor_inputs:
         | 
| 854 | 
            +
                                    callback_kwargs[k] = locals()[k]
         | 
| 855 | 
            +
                                callback_outputs = callback_on_step_end(self, i, t, callback_kwargs)
         | 
| 856 | 
            +
             | 
| 857 | 
            +
                                latents = callback_outputs.pop("latents", latents)
         | 
| 858 | 
            +
                                prompt_embeds = callback_outputs.pop("prompt_embeds", prompt_embeds)
         | 
| 859 | 
            +
             | 
| 860 | 
            +
                            # call the callback, if provided
         | 
| 861 | 
            +
                            if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0):
         | 
| 862 | 
            +
                                progress_bar.update()
         | 
| 863 | 
            +
             | 
| 864 | 
            +
                            if XLA_AVAILABLE:
         | 
| 865 | 
            +
                                xm.mark_step()
         | 
| 866 | 
            +
             | 
| 867 | 
            +
                    self._current_timestep = None
         | 
| 868 | 
            +
                    if output_type == "latent":
         | 
| 869 | 
            +
                        image = latents
         | 
| 870 | 
            +
                    else:
         | 
| 871 | 
            +
                        latents = self._unpack_latents(latents, height, width, self.vae_scale_factor)
         | 
| 872 | 
            +
                        latents = latents.to(self.vae.dtype)
         | 
| 873 | 
            +
                        latents_mean = (
         | 
| 874 | 
            +
                            torch.tensor(self.vae.config.latents_mean)
         | 
| 875 | 
            +
                            .view(1, self.vae.config.z_dim, 1, 1, 1)
         | 
| 876 | 
            +
                            .to(latents.device, latents.dtype)
         | 
| 877 | 
            +
                        )
         | 
| 878 | 
            +
                        latents_std = 1.0 / torch.tensor(self.vae.config.latents_std).view(1, self.vae.config.z_dim, 1, 1, 1).to(
         | 
| 879 | 
            +
                            latents.device, latents.dtype
         | 
| 880 | 
            +
                        )
         | 
| 881 | 
            +
                        latents = latents / latents_std + latents_mean
         | 
| 882 | 
            +
                        image = self.vae.decode(latents, return_dict=False)[0][:, :, 0]
         | 
| 883 | 
            +
                        image = self.image_processor.postprocess(image, output_type=output_type)
         | 
| 884 | 
            +
             | 
| 885 | 
            +
                    # Offload all models
         | 
| 886 | 
            +
                    self.maybe_free_model_hooks()
         | 
| 887 | 
            +
             | 
| 888 | 
            +
                    if not return_dict:
         | 
| 889 | 
            +
                        return (image,)
         | 
| 890 | 
            +
             | 
| 891 | 
            +
                    return QwenImagePipelineOutput(images=image)
         | 
    	
        qwenimage/qwen_fa3_processor.py
    ADDED
    
    | @@ -0,0 +1,142 @@ | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | 
|  | |
| 1 | 
            +
            """
         | 
| 2 | 
            +
            Paired with a good language model. Thanks!
         | 
| 3 | 
            +
            """
         | 
| 4 | 
            +
             | 
| 5 | 
            +
            import torch
         | 
| 6 | 
            +
            from typing import Optional, Tuple
         | 
| 7 | 
            +
            from diffusers.models.transformers.transformer_qwenimage import apply_rotary_emb_qwen
         | 
| 8 | 
            +
             | 
| 9 | 
            +
            try:
         | 
| 10 | 
            +
                from kernels import get_kernel
         | 
| 11 | 
            +
                _k = get_kernel("kernels-community/vllm-flash-attn3")
         | 
| 12 | 
            +
                _flash_attn_func = _k.flash_attn_func
         | 
| 13 | 
            +
            except Exception as e:
         | 
| 14 | 
            +
                _flash_attn_func = None
         | 
| 15 | 
            +
                _kernels_err = e
         | 
| 16 | 
            +
             | 
| 17 | 
            +
             | 
| 18 | 
            +
            def _ensure_fa3_available():
         | 
| 19 | 
            +
                if _flash_attn_func is None:
         | 
| 20 | 
            +
                    raise ImportError(
         | 
| 21 | 
            +
                        "FlashAttention-3 via Hugging Face `kernels` is required. "
         | 
| 22 | 
            +
                        "Tried `get_kernel('kernels-community/vllm-flash-attn3')` and failed with:\n"
         | 
| 23 | 
            +
                        f"{_kernels_err}"
         | 
| 24 | 
            +
                    )
         | 
| 25 | 
            +
             | 
| 26 | 
            +
            @torch.library.custom_op("flash::flash_attn_func", mutates_args=())
         | 
| 27 | 
            +
            def flash_attn_func(
         | 
| 28 | 
            +
                q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, causal: bool = False
         | 
| 29 | 
            +
            ) -> torch.Tensor:
         | 
| 30 | 
            +
                outputs, lse = _flash_attn_func(q, k, v, causal=causal)
         | 
| 31 | 
            +
                return outputs
         | 
| 32 | 
            +
             | 
| 33 | 
            +
            @flash_attn_func.register_fake
         | 
| 34 | 
            +
            def _(q, k, v, **kwargs):
         | 
| 35 | 
            +
                # two outputs:
         | 
| 36 | 
            +
                # 1. output: (batch, seq_len, num_heads, head_dim)
         | 
| 37 | 
            +
                # 2. softmax_lse: (batch, num_heads, seq_len) with dtype=torch.float32
         | 
| 38 | 
            +
                meta_q = torch.empty_like(q).contiguous()
         | 
| 39 | 
            +
                return meta_q #, q.new_empty((q.size(0), q.size(2), q.size(1)), dtype=torch.float32)
         | 
| 40 | 
            +
             | 
| 41 | 
            +
             | 
| 42 | 
            +
            class QwenDoubleStreamAttnProcessorFA3:
         | 
| 43 | 
            +
                """
         | 
| 44 | 
            +
                FA3-based attention processor for Qwen double-stream architecture.
         | 
| 45 | 
            +
                Computes joint attention over concatenated [text, image] streams using vLLM FlashAttention-3
         | 
| 46 | 
            +
                accessed via Hugging Face `kernels`.
         | 
| 47 | 
            +
             | 
| 48 | 
            +
                Notes / limitations:
         | 
| 49 | 
            +
                - General attention masks are not supported here (FA3 path). `is_causal=False` and no arbitrary mask.
         | 
| 50 | 
            +
                - Optional windowed attention / sink tokens / softcap can be plumbed through if you use those features.
         | 
| 51 | 
            +
                - Expects an available `apply_rotary_emb_qwen` in scope (same as your non-FA3 processor).
         | 
| 52 | 
            +
                """
         | 
| 53 | 
            +
             | 
| 54 | 
            +
                _attention_backend = "fa3"  # for parity with your other processors, not used internally
         | 
| 55 | 
            +
             | 
| 56 | 
            +
                def __init__(self):
         | 
| 57 | 
            +
                    _ensure_fa3_available()
         | 
| 58 | 
            +
             | 
| 59 | 
            +
                @torch.no_grad()
         | 
| 60 | 
            +
                def __call__(
         | 
| 61 | 
            +
                    self,
         | 
| 62 | 
            +
                    attn,  # Attention module with to_q/to_k/to_v/add_*_proj, norms, to_out, to_add_out, and .heads
         | 
| 63 | 
            +
                    hidden_states: torch.FloatTensor,                 # (B, S_img, D_model)  image stream
         | 
| 64 | 
            +
                    encoder_hidden_states: torch.FloatTensor = None,  # (B, S_txt, D_model)  text stream
         | 
| 65 | 
            +
                    encoder_hidden_states_mask: torch.FloatTensor = None,  # unused in FA3 path
         | 
| 66 | 
            +
                    attention_mask: Optional[torch.FloatTensor] = None,    # unused in FA3 path
         | 
| 67 | 
            +
                    image_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,  # (img_freqs, txt_freqs)
         | 
| 68 | 
            +
                ) -> Tuple[torch.FloatTensor, torch.FloatTensor]:
         | 
| 69 | 
            +
                    if encoder_hidden_states is None:
         | 
| 70 | 
            +
                        raise ValueError("QwenDoubleStreamAttnProcessorFA3 requires encoder_hidden_states (text stream).")
         | 
| 71 | 
            +
                    if attention_mask is not None:
         | 
| 72 | 
            +
                        # FA3 kernel path here does not consume arbitrary masks; fail fast to avoid silent correctness issues.
         | 
| 73 | 
            +
                        raise NotImplementedError("attention_mask is not supported in this FA3 implementation.")
         | 
| 74 | 
            +
             | 
| 75 | 
            +
                    _ensure_fa3_available()
         | 
| 76 | 
            +
             | 
| 77 | 
            +
                    B, S_img, _ = hidden_states.shape
         | 
| 78 | 
            +
                    S_txt = encoder_hidden_states.shape[1]
         | 
| 79 | 
            +
             | 
| 80 | 
            +
                    # ---- QKV projections (image/sample stream) ----
         | 
| 81 | 
            +
                    img_q = attn.to_q(hidden_states)   # (B, S_img, D)
         | 
| 82 | 
            +
                    img_k = attn.to_k(hidden_states)
         | 
| 83 | 
            +
                    img_v = attn.to_v(hidden_states)
         | 
| 84 | 
            +
             | 
| 85 | 
            +
                    # ---- QKV projections (text/context stream) ----
         | 
| 86 | 
            +
                    txt_q = attn.add_q_proj(encoder_hidden_states)  # (B, S_txt, D)
         | 
| 87 | 
            +
                    txt_k = attn.add_k_proj(encoder_hidden_states)
         | 
| 88 | 
            +
                    txt_v = attn.add_v_proj(encoder_hidden_states)
         | 
| 89 | 
            +
             | 
| 90 | 
            +
                    # ---- Reshape to (B, S, H, D_h) ----
         | 
| 91 | 
            +
                    H = attn.heads
         | 
| 92 | 
            +
                    img_q = img_q.unflatten(-1, (H, -1))
         | 
| 93 | 
            +
                    img_k = img_k.unflatten(-1, (H, -1))
         | 
| 94 | 
            +
                    img_v = img_v.unflatten(-1, (H, -1))
         | 
| 95 | 
            +
             | 
| 96 | 
            +
                    txt_q = txt_q.unflatten(-1, (H, -1))
         | 
| 97 | 
            +
                    txt_k = txt_k.unflatten(-1, (H, -1))
         | 
| 98 | 
            +
                    txt_v = txt_v.unflatten(-1, (H, -1))
         | 
| 99 | 
            +
             | 
| 100 | 
            +
                    # ---- Q/K normalization (per your module contract) ----
         | 
| 101 | 
            +
                    if getattr(attn, "norm_q", None) is not None:
         | 
| 102 | 
            +
                        img_q = attn.norm_q(img_q)
         | 
| 103 | 
            +
                    if getattr(attn, "norm_k", None) is not None:
         | 
| 104 | 
            +
                        img_k = attn.norm_k(img_k)
         | 
| 105 | 
            +
                    if getattr(attn, "norm_added_q", None) is not None:
         | 
| 106 | 
            +
                        txt_q = attn.norm_added_q(txt_q)
         | 
| 107 | 
            +
                    if getattr(attn, "norm_added_k", None) is not None:
         | 
| 108 | 
            +
                        txt_k = attn.norm_added_k(txt_k)
         | 
| 109 | 
            +
             | 
| 110 | 
            +
                    # ---- RoPE (Qwen variant) ----
         | 
| 111 | 
            +
                    if image_rotary_emb is not None:
         | 
| 112 | 
            +
                        img_freqs, txt_freqs = image_rotary_emb
         | 
| 113 | 
            +
                        # expects tensors shaped (B, S, H, D_h)
         | 
| 114 | 
            +
                        img_q = apply_rotary_emb_qwen(img_q, img_freqs, use_real=False)
         | 
| 115 | 
            +
                        img_k = apply_rotary_emb_qwen(img_k, img_freqs, use_real=False)
         | 
| 116 | 
            +
                        txt_q = apply_rotary_emb_qwen(txt_q, txt_freqs, use_real=False)
         | 
| 117 | 
            +
                        txt_k = apply_rotary_emb_qwen(txt_k, txt_freqs, use_real=False)
         | 
| 118 | 
            +
             | 
| 119 | 
            +
                    # ---- Joint attention over [text, image] along sequence axis ----
         | 
| 120 | 
            +
                    # Shapes: (B, S_total, H, D_h)
         | 
| 121 | 
            +
                    q = torch.cat([txt_q, img_q], dim=1)
         | 
| 122 | 
            +
                    k = torch.cat([txt_k, img_k], dim=1)
         | 
| 123 | 
            +
                    v = torch.cat([txt_v, img_v], dim=1)
         | 
| 124 | 
            +
             | 
| 125 | 
            +
                    # FlashAttention-3 path expects (B, S, H, D_h) and returns (out, softmax_lse)
         | 
| 126 | 
            +
                    out = flash_attn_func(q, k, v, causal=False)  # out: (B, S_total, H, D_h)
         | 
| 127 | 
            +
             | 
| 128 | 
            +
                    # ---- Back to (B, S, D_model) ----
         | 
| 129 | 
            +
                    out = out.flatten(2, 3).to(q.dtype)
         | 
| 130 | 
            +
             | 
| 131 | 
            +
                    # Split back to text / image segments
         | 
| 132 | 
            +
                    txt_attn_out = out[:, :S_txt, :]
         | 
| 133 | 
            +
                    img_attn_out = out[:, S_txt:, :]
         | 
| 134 | 
            +
             | 
| 135 | 
            +
                    # ---- Output projections ----
         | 
| 136 | 
            +
                    img_attn_out = attn.to_out[0](img_attn_out)
         | 
| 137 | 
            +
                    if len(attn.to_out) > 1:
         | 
| 138 | 
            +
                        img_attn_out = attn.to_out[1](img_attn_out)  # dropout if present
         | 
| 139 | 
            +
             | 
| 140 | 
            +
                    txt_attn_out = attn.to_add_out(txt_attn_out)
         | 
| 141 | 
            +
             | 
| 142 | 
            +
                    return img_attn_out, txt_attn_out
         | 
    	
        qwenimage/transformer_qwenimage.py
    ADDED
    
    | @@ -0,0 +1,642 @@ | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | 
|  | |
| 1 | 
            +
            # Copyright 2025 Qwen-Image Team, The HuggingFace Team. All rights reserved.
         | 
| 2 | 
            +
            #
         | 
| 3 | 
            +
            # Licensed under the Apache License, Version 2.0 (the "License");
         | 
| 4 | 
            +
            # you may not use this file except in compliance with the License.
         | 
| 5 | 
            +
            # You may obtain a copy of the License at
         | 
| 6 | 
            +
            #
         | 
| 7 | 
            +
            #     http://www.apache.org/licenses/LICENSE-2.0
         | 
| 8 | 
            +
            #
         | 
| 9 | 
            +
            # Unless required by applicable law or agreed to in writing, software
         | 
| 10 | 
            +
            # distributed under the License is distributed on an "AS IS" BASIS,
         | 
| 11 | 
            +
            # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
         | 
| 12 | 
            +
            # See the License for the specific language governing permissions and
         | 
| 13 | 
            +
            # limitations under the License.
         | 
| 14 | 
            +
             | 
| 15 | 
            +
            import functools
         | 
| 16 | 
            +
            import math
         | 
| 17 | 
            +
            from typing import Any, Dict, List, Optional, Tuple, Union
         | 
| 18 | 
            +
             | 
| 19 | 
            +
            import torch
         | 
| 20 | 
            +
            import torch.nn as nn
         | 
| 21 | 
            +
            import torch.nn.functional as F
         | 
| 22 | 
            +
             | 
| 23 | 
            +
            from diffusers.configuration_utils import ConfigMixin, register_to_config
         | 
| 24 | 
            +
            from diffusers.loaders import FromOriginalModelMixin, PeftAdapterMixin
         | 
| 25 | 
            +
            from diffusers.utils import USE_PEFT_BACKEND, logging, scale_lora_layers, unscale_lora_layers
         | 
| 26 | 
            +
            from diffusers.utils.torch_utils import maybe_allow_in_graph
         | 
| 27 | 
            +
            from diffusers.models.attention import FeedForward, AttentionMixin
         | 
| 28 | 
            +
            from diffusers.models.attention_dispatch import dispatch_attention_fn
         | 
| 29 | 
            +
            from diffusers.models.attention_processor import Attention
         | 
| 30 | 
            +
            from diffusers.models.cache_utils import CacheMixin
         | 
| 31 | 
            +
            from diffusers.models.embeddings import TimestepEmbedding, Timesteps
         | 
| 32 | 
            +
            from diffusers.models.modeling_outputs import Transformer2DModelOutput
         | 
| 33 | 
            +
            from diffusers.models.modeling_utils import ModelMixin
         | 
| 34 | 
            +
            from diffusers.models.normalization import AdaLayerNormContinuous, RMSNorm
         | 
| 35 | 
            +
             | 
| 36 | 
            +
             | 
| 37 | 
            +
            logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
         | 
| 38 | 
            +
             | 
| 39 | 
            +
             | 
| 40 | 
            +
            def get_timestep_embedding(
         | 
| 41 | 
            +
                timesteps: torch.Tensor,
         | 
| 42 | 
            +
                embedding_dim: int,
         | 
| 43 | 
            +
                flip_sin_to_cos: bool = False,
         | 
| 44 | 
            +
                downscale_freq_shift: float = 1,
         | 
| 45 | 
            +
                scale: float = 1,
         | 
| 46 | 
            +
                max_period: int = 10000,
         | 
| 47 | 
            +
            ) -> torch.Tensor:
         | 
| 48 | 
            +
                """
         | 
| 49 | 
            +
                This matches the implementation in Denoising Diffusion Probabilistic Models: Create sinusoidal timestep embeddings.
         | 
| 50 | 
            +
             | 
| 51 | 
            +
                Args
         | 
| 52 | 
            +
                    timesteps (torch.Tensor):
         | 
| 53 | 
            +
                        a 1-D Tensor of N indices, one per batch element. These may be fractional.
         | 
| 54 | 
            +
                    embedding_dim (int):
         | 
| 55 | 
            +
                        the dimension of the output.
         | 
| 56 | 
            +
                    flip_sin_to_cos (bool):
         | 
| 57 | 
            +
                        Whether the embedding order should be `cos, sin` (if True) or `sin, cos` (if False)
         | 
| 58 | 
            +
                    downscale_freq_shift (float):
         | 
| 59 | 
            +
                        Controls the delta between frequencies between dimensions
         | 
| 60 | 
            +
                    scale (float):
         | 
| 61 | 
            +
                        Scaling factor applied to the embeddings.
         | 
| 62 | 
            +
                    max_period (int):
         | 
| 63 | 
            +
                        Controls the maximum frequency of the embeddings
         | 
| 64 | 
            +
                Returns
         | 
| 65 | 
            +
                    torch.Tensor: an [N x dim] Tensor of positional embeddings.
         | 
| 66 | 
            +
                """
         | 
| 67 | 
            +
                assert len(timesteps.shape) == 1, "Timesteps should be a 1d-array"
         | 
| 68 | 
            +
             | 
| 69 | 
            +
                half_dim = embedding_dim // 2
         | 
| 70 | 
            +
                exponent = -math.log(max_period) * torch.arange(
         | 
| 71 | 
            +
                    start=0, end=half_dim, dtype=torch.float32, device=timesteps.device
         | 
| 72 | 
            +
                )
         | 
| 73 | 
            +
                exponent = exponent / (half_dim - downscale_freq_shift)
         | 
| 74 | 
            +
             | 
| 75 | 
            +
                emb = torch.exp(exponent).to(timesteps.dtype)
         | 
| 76 | 
            +
                emb = timesteps[:, None].float() * emb[None, :]
         | 
| 77 | 
            +
             | 
| 78 | 
            +
                # scale embeddings
         | 
| 79 | 
            +
                emb = scale * emb
         | 
| 80 | 
            +
             | 
| 81 | 
            +
                # concat sine and cosine embeddings
         | 
| 82 | 
            +
                emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=-1)
         | 
| 83 | 
            +
             | 
| 84 | 
            +
                # flip sine and cosine embeddings
         | 
| 85 | 
            +
                if flip_sin_to_cos:
         | 
| 86 | 
            +
                    emb = torch.cat([emb[:, half_dim:], emb[:, :half_dim]], dim=-1)
         | 
| 87 | 
            +
             | 
| 88 | 
            +
                # zero pad
         | 
| 89 | 
            +
                if embedding_dim % 2 == 1:
         | 
| 90 | 
            +
                    emb = torch.nn.functional.pad(emb, (0, 1, 0, 0))
         | 
| 91 | 
            +
                return emb
         | 
| 92 | 
            +
             | 
| 93 | 
            +
             | 
| 94 | 
            +
            def apply_rotary_emb_qwen(
         | 
| 95 | 
            +
                x: torch.Tensor,
         | 
| 96 | 
            +
                freqs_cis: Union[torch.Tensor, Tuple[torch.Tensor]],
         | 
| 97 | 
            +
                use_real: bool = True,
         | 
| 98 | 
            +
                use_real_unbind_dim: int = -1,
         | 
| 99 | 
            +
            ) -> Tuple[torch.Tensor, torch.Tensor]:
         | 
| 100 | 
            +
                """
         | 
| 101 | 
            +
                Apply rotary embeddings to input tensors using the given frequency tensor. This function applies rotary embeddings
         | 
| 102 | 
            +
                to the given query or key 'x' tensors using the provided frequency tensor 'freqs_cis'. The input tensors are
         | 
| 103 | 
            +
                reshaped as complex numbers, and the frequency tensor is reshaped for broadcasting compatibility. The resulting
         | 
| 104 | 
            +
                tensors contain rotary embeddings and are returned as real tensors.
         | 
| 105 | 
            +
             | 
| 106 | 
            +
                Args:
         | 
| 107 | 
            +
                    x (`torch.Tensor`):
         | 
| 108 | 
            +
                        Query or key tensor to apply rotary embeddings. [B, S, H, D] xk (torch.Tensor): Key tensor to apply
         | 
| 109 | 
            +
                    freqs_cis (`Tuple[torch.Tensor]`): Precomputed frequency tensor for complex exponentials. ([S, D], [S, D],)
         | 
| 110 | 
            +
             | 
| 111 | 
            +
                Returns:
         | 
| 112 | 
            +
                    Tuple[torch.Tensor, torch.Tensor]: Tuple of modified query tensor and key tensor with rotary embeddings.
         | 
| 113 | 
            +
                """
         | 
| 114 | 
            +
                if use_real:
         | 
| 115 | 
            +
                    cos, sin = freqs_cis  # [S, D]
         | 
| 116 | 
            +
                    cos = cos[None, None]
         | 
| 117 | 
            +
                    sin = sin[None, None]
         | 
| 118 | 
            +
                    cos, sin = cos.to(x.device), sin.to(x.device)
         | 
| 119 | 
            +
             | 
| 120 | 
            +
                    if use_real_unbind_dim == -1:
         | 
| 121 | 
            +
                        # Used for flux, cogvideox, hunyuan-dit
         | 
| 122 | 
            +
                        x_real, x_imag = x.reshape(*x.shape[:-1], -1, 2).unbind(-1)  # [B, S, H, D//2]
         | 
| 123 | 
            +
                        x_rotated = torch.stack([-x_imag, x_real], dim=-1).flatten(3)
         | 
| 124 | 
            +
                    elif use_real_unbind_dim == -2:
         | 
| 125 | 
            +
                        # Used for Stable Audio, OmniGen, CogView4 and Cosmos
         | 
| 126 | 
            +
                        x_real, x_imag = x.reshape(*x.shape[:-1], 2, -1).unbind(-2)  # [B, S, H, D//2]
         | 
| 127 | 
            +
                        x_rotated = torch.cat([-x_imag, x_real], dim=-1)
         | 
| 128 | 
            +
                    else:
         | 
| 129 | 
            +
                        raise ValueError(f"`use_real_unbind_dim={use_real_unbind_dim}` but should be -1 or -2.")
         | 
| 130 | 
            +
             | 
| 131 | 
            +
                    out = (x.float() * cos + x_rotated.float() * sin).to(x.dtype)
         | 
| 132 | 
            +
             | 
| 133 | 
            +
                    return out
         | 
| 134 | 
            +
                else:
         | 
| 135 | 
            +
                    x_rotated = torch.view_as_complex(x.float().reshape(*x.shape[:-1], -1, 2))
         | 
| 136 | 
            +
                    freqs_cis = freqs_cis.unsqueeze(1)
         | 
| 137 | 
            +
                    x_out = torch.view_as_real(x_rotated * freqs_cis).flatten(3)
         | 
| 138 | 
            +
             | 
| 139 | 
            +
                    return x_out.type_as(x)
         | 
| 140 | 
            +
             | 
| 141 | 
            +
             | 
| 142 | 
            +
            class QwenTimestepProjEmbeddings(nn.Module):
         | 
| 143 | 
            +
                def __init__(self, embedding_dim):
         | 
| 144 | 
            +
                    super().__init__()
         | 
| 145 | 
            +
             | 
| 146 | 
            +
                    self.time_proj = Timesteps(num_channels=256, flip_sin_to_cos=True, downscale_freq_shift=0, scale=1000)
         | 
| 147 | 
            +
                    self.timestep_embedder = TimestepEmbedding(in_channels=256, time_embed_dim=embedding_dim)
         | 
| 148 | 
            +
             | 
| 149 | 
            +
                def forward(self, timestep, hidden_states):
         | 
| 150 | 
            +
                    timesteps_proj = self.time_proj(timestep)
         | 
| 151 | 
            +
                    timesteps_emb = self.timestep_embedder(timesteps_proj.to(dtype=hidden_states.dtype))  # (N, D)
         | 
| 152 | 
            +
             | 
| 153 | 
            +
                    conditioning = timesteps_emb
         | 
| 154 | 
            +
             | 
| 155 | 
            +
                    return conditioning
         | 
| 156 | 
            +
             | 
| 157 | 
            +
             | 
| 158 | 
            +
            class QwenEmbedRope(nn.Module):
         | 
| 159 | 
            +
                def __init__(self, theta: int, axes_dim: List[int], scale_rope=False):
         | 
| 160 | 
            +
                    super().__init__()
         | 
| 161 | 
            +
                    self.theta = theta
         | 
| 162 | 
            +
                    self.axes_dim = axes_dim
         | 
| 163 | 
            +
                    pos_index = torch.arange(4096)
         | 
| 164 | 
            +
                    neg_index = torch.arange(4096).flip(0) * -1 - 1
         | 
| 165 | 
            +
                    self.pos_freqs = torch.cat(
         | 
| 166 | 
            +
                        [
         | 
| 167 | 
            +
                            self.rope_params(pos_index, self.axes_dim[0], self.theta),
         | 
| 168 | 
            +
                            self.rope_params(pos_index, self.axes_dim[1], self.theta),
         | 
| 169 | 
            +
                            self.rope_params(pos_index, self.axes_dim[2], self.theta),
         | 
| 170 | 
            +
                        ],
         | 
| 171 | 
            +
                        dim=1,
         | 
| 172 | 
            +
                    )
         | 
| 173 | 
            +
                    self.neg_freqs = torch.cat(
         | 
| 174 | 
            +
                        [
         | 
| 175 | 
            +
                            self.rope_params(neg_index, self.axes_dim[0], self.theta),
         | 
| 176 | 
            +
                            self.rope_params(neg_index, self.axes_dim[1], self.theta),
         | 
| 177 | 
            +
                            self.rope_params(neg_index, self.axes_dim[2], self.theta),
         | 
| 178 | 
            +
                        ],
         | 
| 179 | 
            +
                        dim=1,
         | 
| 180 | 
            +
                    )
         | 
| 181 | 
            +
                    self.rope_cache = {}
         | 
| 182 | 
            +
             | 
| 183 | 
            +
                    # DO NOT USING REGISTER BUFFER HERE, IT WILL CAUSE COMPLEX NUMBERS LOSE ITS IMAGINARY PART
         | 
| 184 | 
            +
                    self.scale_rope = scale_rope
         | 
| 185 | 
            +
             | 
| 186 | 
            +
                def rope_params(self, index, dim, theta=10000):
         | 
| 187 | 
            +
                    """
         | 
| 188 | 
            +
                    Args:
         | 
| 189 | 
            +
                        index: [0, 1, 2, 3] 1D Tensor representing the position index of the token
         | 
| 190 | 
            +
                    """
         | 
| 191 | 
            +
                    assert dim % 2 == 0
         | 
| 192 | 
            +
                    freqs = torch.outer(index, 1.0 / torch.pow(theta, torch.arange(0, dim, 2).to(torch.float32).div(dim)))
         | 
| 193 | 
            +
                    freqs = torch.polar(torch.ones_like(freqs), freqs)
         | 
| 194 | 
            +
                    return freqs
         | 
| 195 | 
            +
             | 
| 196 | 
            +
                def forward(self, video_fhw, txt_seq_lens, device):
         | 
| 197 | 
            +
                    """
         | 
| 198 | 
            +
                    Args: video_fhw: [frame, height, width] a list of 3 integers representing the shape of the video Args:
         | 
| 199 | 
            +
                    txt_length: [bs] a list of 1 integers representing the length of the text
         | 
| 200 | 
            +
                    """
         | 
| 201 | 
            +
                    if self.pos_freqs.device != device:
         | 
| 202 | 
            +
                        self.pos_freqs = self.pos_freqs.to(device)
         | 
| 203 | 
            +
                        self.neg_freqs = self.neg_freqs.to(device)
         | 
| 204 | 
            +
             | 
| 205 | 
            +
                    if isinstance(video_fhw, list):
         | 
| 206 | 
            +
                        video_fhw = video_fhw[0]
         | 
| 207 | 
            +
                    if not isinstance(video_fhw, list):
         | 
| 208 | 
            +
                        video_fhw = [video_fhw]
         | 
| 209 | 
            +
             | 
| 210 | 
            +
                    vid_freqs = []
         | 
| 211 | 
            +
                    max_vid_index = 0
         | 
| 212 | 
            +
                    for idx, fhw in enumerate(video_fhw):
         | 
| 213 | 
            +
                        frame, height, width = fhw
         | 
| 214 | 
            +
                        rope_key = f"{idx}_{height}_{width}"
         | 
| 215 | 
            +
             | 
| 216 | 
            +
                        if not torch.compiler.is_compiling():
         | 
| 217 | 
            +
                            if rope_key not in self.rope_cache:
         | 
| 218 | 
            +
                                self.rope_cache[rope_key] = self._compute_video_freqs(frame, height, width, idx)
         | 
| 219 | 
            +
                            video_freq = self.rope_cache[rope_key]
         | 
| 220 | 
            +
                        else:
         | 
| 221 | 
            +
                            video_freq = self._compute_video_freqs(frame, height, width, idx)
         | 
| 222 | 
            +
                        video_freq = video_freq.to(device)
         | 
| 223 | 
            +
                        vid_freqs.append(video_freq)
         | 
| 224 | 
            +
             | 
| 225 | 
            +
                        if self.scale_rope:
         | 
| 226 | 
            +
                            max_vid_index = max(height // 2, width // 2, max_vid_index)
         | 
| 227 | 
            +
                        else:
         | 
| 228 | 
            +
                            max_vid_index = max(height, width, max_vid_index)
         | 
| 229 | 
            +
             | 
| 230 | 
            +
                    max_len = max(txt_seq_lens)
         | 
| 231 | 
            +
                    txt_freqs = self.pos_freqs[max_vid_index : max_vid_index + max_len, ...]
         | 
| 232 | 
            +
                    vid_freqs = torch.cat(vid_freqs, dim=0)
         | 
| 233 | 
            +
             | 
| 234 | 
            +
                    return vid_freqs, txt_freqs
         | 
| 235 | 
            +
             | 
| 236 | 
            +
                @functools.lru_cache(maxsize=None)
         | 
| 237 | 
            +
                def _compute_video_freqs(self, frame, height, width, idx=0):
         | 
| 238 | 
            +
                    seq_lens = frame * height * width
         | 
| 239 | 
            +
                    freqs_pos = self.pos_freqs.split([x // 2 for x in self.axes_dim], dim=1)
         | 
| 240 | 
            +
                    freqs_neg = self.neg_freqs.split([x // 2 for x in self.axes_dim], dim=1)
         | 
| 241 | 
            +
             | 
| 242 | 
            +
                    freqs_frame = freqs_pos[0][idx : idx + frame].view(frame, 1, 1, -1).expand(frame, height, width, -1)
         | 
| 243 | 
            +
                    if self.scale_rope:
         | 
| 244 | 
            +
                        freqs_height = torch.cat([freqs_neg[1][-(height - height // 2) :], freqs_pos[1][: height // 2]], dim=0)
         | 
| 245 | 
            +
                        freqs_height = freqs_height.view(1, height, 1, -1).expand(frame, height, width, -1)
         | 
| 246 | 
            +
                        freqs_width = torch.cat([freqs_neg[2][-(width - width // 2) :], freqs_pos[2][: width // 2]], dim=0)
         | 
| 247 | 
            +
                        freqs_width = freqs_width.view(1, 1, width, -1).expand(frame, height, width, -1)
         | 
| 248 | 
            +
                    else:
         | 
| 249 | 
            +
                        freqs_height = freqs_pos[1][:height].view(1, height, 1, -1).expand(frame, height, width, -1)
         | 
| 250 | 
            +
                        freqs_width = freqs_pos[2][:width].view(1, 1, width, -1).expand(frame, height, width, -1)
         | 
| 251 | 
            +
             | 
| 252 | 
            +
                    freqs = torch.cat([freqs_frame, freqs_height, freqs_width], dim=-1).reshape(seq_lens, -1)
         | 
| 253 | 
            +
                    return freqs.clone().contiguous()
         | 
| 254 | 
            +
             | 
| 255 | 
            +
             | 
| 256 | 
            +
            class QwenDoubleStreamAttnProcessor2_0:
         | 
| 257 | 
            +
                """
         | 
| 258 | 
            +
                Attention processor for Qwen double-stream architecture, matching DoubleStreamLayerMegatron logic. This processor
         | 
| 259 | 
            +
                implements joint attention computation where text and image streams are processed together.
         | 
| 260 | 
            +
                """
         | 
| 261 | 
            +
             | 
| 262 | 
            +
                _attention_backend = None
         | 
| 263 | 
            +
             | 
| 264 | 
            +
                def __init__(self):
         | 
| 265 | 
            +
                    if not hasattr(F, "scaled_dot_product_attention"):
         | 
| 266 | 
            +
                        raise ImportError(
         | 
| 267 | 
            +
                            "QwenDoubleStreamAttnProcessor2_0 requires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0."
         | 
| 268 | 
            +
                        )
         | 
| 269 | 
            +
             | 
| 270 | 
            +
                def __call__(
         | 
| 271 | 
            +
                    self,
         | 
| 272 | 
            +
                    attn: Attention,
         | 
| 273 | 
            +
                    hidden_states: torch.FloatTensor,  # Image stream
         | 
| 274 | 
            +
                    encoder_hidden_states: torch.FloatTensor = None,  # Text stream
         | 
| 275 | 
            +
                    encoder_hidden_states_mask: torch.FloatTensor = None,
         | 
| 276 | 
            +
                    attention_mask: Optional[torch.FloatTensor] = None,
         | 
| 277 | 
            +
                    image_rotary_emb: Optional[torch.Tensor] = None,
         | 
| 278 | 
            +
                ) -> torch.FloatTensor:
         | 
| 279 | 
            +
                    if encoder_hidden_states is None:
         | 
| 280 | 
            +
                        raise ValueError("QwenDoubleStreamAttnProcessor2_0 requires encoder_hidden_states (text stream)")
         | 
| 281 | 
            +
             | 
| 282 | 
            +
                    seq_txt = encoder_hidden_states.shape[1]
         | 
| 283 | 
            +
             | 
| 284 | 
            +
                    # Compute QKV for image stream (sample projections)
         | 
| 285 | 
            +
                    img_query = attn.to_q(hidden_states)
         | 
| 286 | 
            +
                    img_key = attn.to_k(hidden_states)
         | 
| 287 | 
            +
                    img_value = attn.to_v(hidden_states)
         | 
| 288 | 
            +
             | 
| 289 | 
            +
                    # Compute QKV for text stream (context projections)
         | 
| 290 | 
            +
                    txt_query = attn.add_q_proj(encoder_hidden_states)
         | 
| 291 | 
            +
                    txt_key = attn.add_k_proj(encoder_hidden_states)
         | 
| 292 | 
            +
                    txt_value = attn.add_v_proj(encoder_hidden_states)
         | 
| 293 | 
            +
             | 
| 294 | 
            +
                    # Reshape for multi-head attention
         | 
| 295 | 
            +
                    img_query = img_query.unflatten(-1, (attn.heads, -1))
         | 
| 296 | 
            +
                    img_key = img_key.unflatten(-1, (attn.heads, -1))
         | 
| 297 | 
            +
                    img_value = img_value.unflatten(-1, (attn.heads, -1))
         | 
| 298 | 
            +
             | 
| 299 | 
            +
                    txt_query = txt_query.unflatten(-1, (attn.heads, -1))
         | 
| 300 | 
            +
                    txt_key = txt_key.unflatten(-1, (attn.heads, -1))
         | 
| 301 | 
            +
                    txt_value = txt_value.unflatten(-1, (attn.heads, -1))
         | 
| 302 | 
            +
             | 
| 303 | 
            +
                    # Apply QK normalization
         | 
| 304 | 
            +
                    if attn.norm_q is not None:
         | 
| 305 | 
            +
                        img_query = attn.norm_q(img_query)
         | 
| 306 | 
            +
                    if attn.norm_k is not None:
         | 
| 307 | 
            +
                        img_key = attn.norm_k(img_key)
         | 
| 308 | 
            +
                    if attn.norm_added_q is not None:
         | 
| 309 | 
            +
                        txt_query = attn.norm_added_q(txt_query)
         | 
| 310 | 
            +
                    if attn.norm_added_k is not None:
         | 
| 311 | 
            +
                        txt_key = attn.norm_added_k(txt_key)
         | 
| 312 | 
            +
             | 
| 313 | 
            +
                    # Apply RoPE
         | 
| 314 | 
            +
                    if image_rotary_emb is not None:
         | 
| 315 | 
            +
                        img_freqs, txt_freqs = image_rotary_emb
         | 
| 316 | 
            +
                        img_query = apply_rotary_emb_qwen(img_query, img_freqs, use_real=False)
         | 
| 317 | 
            +
                        img_key = apply_rotary_emb_qwen(img_key, img_freqs, use_real=False)
         | 
| 318 | 
            +
                        txt_query = apply_rotary_emb_qwen(txt_query, txt_freqs, use_real=False)
         | 
| 319 | 
            +
                        txt_key = apply_rotary_emb_qwen(txt_key, txt_freqs, use_real=False)
         | 
| 320 | 
            +
             | 
| 321 | 
            +
                    # Concatenate for joint attention
         | 
| 322 | 
            +
                    # Order: [text, image]
         | 
| 323 | 
            +
                    joint_query = torch.cat([txt_query, img_query], dim=1)
         | 
| 324 | 
            +
                    joint_key = torch.cat([txt_key, img_key], dim=1)
         | 
| 325 | 
            +
                    joint_value = torch.cat([txt_value, img_value], dim=1)
         | 
| 326 | 
            +
             | 
| 327 | 
            +
                    # Compute joint attention
         | 
| 328 | 
            +
                    joint_hidden_states = dispatch_attention_fn(
         | 
| 329 | 
            +
                        joint_query,
         | 
| 330 | 
            +
                        joint_key,
         | 
| 331 | 
            +
                        joint_value,
         | 
| 332 | 
            +
                        attn_mask=attention_mask,
         | 
| 333 | 
            +
                        dropout_p=0.0,
         | 
| 334 | 
            +
                        is_causal=False,
         | 
| 335 | 
            +
                        backend=self._attention_backend,
         | 
| 336 | 
            +
                    )
         | 
| 337 | 
            +
             | 
| 338 | 
            +
                    # Reshape back
         | 
| 339 | 
            +
                    joint_hidden_states = joint_hidden_states.flatten(2, 3)
         | 
| 340 | 
            +
                    joint_hidden_states = joint_hidden_states.to(joint_query.dtype)
         | 
| 341 | 
            +
             | 
| 342 | 
            +
                    # Split attention outputs back
         | 
| 343 | 
            +
                    txt_attn_output = joint_hidden_states[:, :seq_txt, :]  # Text part
         | 
| 344 | 
            +
                    img_attn_output = joint_hidden_states[:, seq_txt:, :]  # Image part
         | 
| 345 | 
            +
             | 
| 346 | 
            +
                    # Apply output projections
         | 
| 347 | 
            +
                    img_attn_output = attn.to_out[0](img_attn_output)
         | 
| 348 | 
            +
                    if len(attn.to_out) > 1:
         | 
| 349 | 
            +
                        img_attn_output = attn.to_out[1](img_attn_output)  # dropout
         | 
| 350 | 
            +
             | 
| 351 | 
            +
                    txt_attn_output = attn.to_add_out(txt_attn_output)
         | 
| 352 | 
            +
             | 
| 353 | 
            +
                    return img_attn_output, txt_attn_output
         | 
| 354 | 
            +
             | 
| 355 | 
            +
             | 
| 356 | 
            +
            @maybe_allow_in_graph
         | 
| 357 | 
            +
            class QwenImageTransformerBlock(nn.Module):
         | 
| 358 | 
            +
                def __init__(
         | 
| 359 | 
            +
                    self, dim: int, num_attention_heads: int, attention_head_dim: int, qk_norm: str = "rms_norm", eps: float = 1e-6
         | 
| 360 | 
            +
                ):
         | 
| 361 | 
            +
                    super().__init__()
         | 
| 362 | 
            +
             | 
| 363 | 
            +
                    self.dim = dim
         | 
| 364 | 
            +
                    self.num_attention_heads = num_attention_heads
         | 
| 365 | 
            +
                    self.attention_head_dim = attention_head_dim
         | 
| 366 | 
            +
             | 
| 367 | 
            +
                    # Image processing modules
         | 
| 368 | 
            +
                    self.img_mod = nn.Sequential(
         | 
| 369 | 
            +
                        nn.SiLU(),
         | 
| 370 | 
            +
                        nn.Linear(dim, 6 * dim, bias=True),  # For scale, shift, gate for norm1 and norm2
         | 
| 371 | 
            +
                    )
         | 
| 372 | 
            +
                    self.img_norm1 = nn.LayerNorm(dim, elementwise_affine=False, eps=eps)
         | 
| 373 | 
            +
                    self.attn = Attention(
         | 
| 374 | 
            +
                        query_dim=dim,
         | 
| 375 | 
            +
                        cross_attention_dim=None,  # Enable cross attention for joint computation
         | 
| 376 | 
            +
                        added_kv_proj_dim=dim,  # Enable added KV projections for text stream
         | 
| 377 | 
            +
                        dim_head=attention_head_dim,
         | 
| 378 | 
            +
                        heads=num_attention_heads,
         | 
| 379 | 
            +
                        out_dim=dim,
         | 
| 380 | 
            +
                        context_pre_only=False,
         | 
| 381 | 
            +
                        bias=True,
         | 
| 382 | 
            +
                        processor=QwenDoubleStreamAttnProcessor2_0(),
         | 
| 383 | 
            +
                        qk_norm=qk_norm,
         | 
| 384 | 
            +
                        eps=eps,
         | 
| 385 | 
            +
                    )
         | 
| 386 | 
            +
                    self.img_norm2 = nn.LayerNorm(dim, elementwise_affine=False, eps=eps)
         | 
| 387 | 
            +
                    self.img_mlp = FeedForward(dim=dim, dim_out=dim, activation_fn="gelu-approximate")
         | 
| 388 | 
            +
             | 
| 389 | 
            +
                    # Text processing modules
         | 
| 390 | 
            +
                    self.txt_mod = nn.Sequential(
         | 
| 391 | 
            +
                        nn.SiLU(),
         | 
| 392 | 
            +
                        nn.Linear(dim, 6 * dim, bias=True),  # For scale, shift, gate for norm1 and norm2
         | 
| 393 | 
            +
                    )
         | 
| 394 | 
            +
                    self.txt_norm1 = nn.LayerNorm(dim, elementwise_affine=False, eps=eps)
         | 
| 395 | 
            +
                    # Text doesn't need separate attention - it's handled by img_attn joint computation
         | 
| 396 | 
            +
                    self.txt_norm2 = nn.LayerNorm(dim, elementwise_affine=False, eps=eps)
         | 
| 397 | 
            +
                    self.txt_mlp = FeedForward(dim=dim, dim_out=dim, activation_fn="gelu-approximate")
         | 
| 398 | 
            +
             | 
| 399 | 
            +
                def _modulate(self, x, mod_params):
         | 
| 400 | 
            +
                    """Apply modulation to input tensor"""
         | 
| 401 | 
            +
                    shift, scale, gate = mod_params.chunk(3, dim=-1)
         | 
| 402 | 
            +
                    return x * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1), gate.unsqueeze(1)
         | 
| 403 | 
            +
             | 
| 404 | 
            +
                def forward(
         | 
| 405 | 
            +
                    self,
         | 
| 406 | 
            +
                    hidden_states: torch.Tensor,
         | 
| 407 | 
            +
                    encoder_hidden_states: torch.Tensor,
         | 
| 408 | 
            +
                    encoder_hidden_states_mask: torch.Tensor,
         | 
| 409 | 
            +
                    temb: torch.Tensor,
         | 
| 410 | 
            +
                    image_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
         | 
| 411 | 
            +
                    joint_attention_kwargs: Optional[Dict[str, Any]] = None,
         | 
| 412 | 
            +
                ) -> Tuple[torch.Tensor, torch.Tensor]:
         | 
| 413 | 
            +
                    # Get modulation parameters for both streams
         | 
| 414 | 
            +
                    img_mod_params = self.img_mod(temb)  # [B, 6*dim]
         | 
| 415 | 
            +
                    txt_mod_params = self.txt_mod(temb)  # [B, 6*dim]
         | 
| 416 | 
            +
             | 
| 417 | 
            +
                    # Split modulation parameters for norm1 and norm2
         | 
| 418 | 
            +
                    img_mod1, img_mod2 = img_mod_params.chunk(2, dim=-1)  # Each [B, 3*dim]
         | 
| 419 | 
            +
                    txt_mod1, txt_mod2 = txt_mod_params.chunk(2, dim=-1)  # Each [B, 3*dim]
         | 
| 420 | 
            +
             | 
| 421 | 
            +
                    # Process image stream - norm1 + modulation
         | 
| 422 | 
            +
                    img_normed = self.img_norm1(hidden_states)
         | 
| 423 | 
            +
                    img_modulated, img_gate1 = self._modulate(img_normed, img_mod1)
         | 
| 424 | 
            +
             | 
| 425 | 
            +
                    # Process text stream - norm1 + modulation
         | 
| 426 | 
            +
                    txt_normed = self.txt_norm1(encoder_hidden_states)
         | 
| 427 | 
            +
                    txt_modulated, txt_gate1 = self._modulate(txt_normed, txt_mod1)
         | 
| 428 | 
            +
             | 
| 429 | 
            +
                    # Use QwenAttnProcessor2_0 for joint attention computation
         | 
| 430 | 
            +
                    # This directly implements the DoubleStreamLayerMegatron logic:
         | 
| 431 | 
            +
                    # 1. Computes QKV for both streams
         | 
| 432 | 
            +
                    # 2. Applies QK normalization and RoPE
         | 
| 433 | 
            +
                    # 3. Concatenates and runs joint attention
         | 
| 434 | 
            +
                    # 4. Splits results back to separate streams
         | 
| 435 | 
            +
                    joint_attention_kwargs = joint_attention_kwargs or {}
         | 
| 436 | 
            +
                    attn_output = self.attn(
         | 
| 437 | 
            +
                        hidden_states=img_modulated,  # Image stream (will be processed as "sample")
         | 
| 438 | 
            +
                        encoder_hidden_states=txt_modulated,  # Text stream (will be processed as "context")
         | 
| 439 | 
            +
                        encoder_hidden_states_mask=encoder_hidden_states_mask,
         | 
| 440 | 
            +
                        image_rotary_emb=image_rotary_emb,
         | 
| 441 | 
            +
                        **joint_attention_kwargs,
         | 
| 442 | 
            +
                    )
         | 
| 443 | 
            +
             | 
| 444 | 
            +
                    # QwenAttnProcessor2_0 returns (img_output, txt_output) when encoder_hidden_states is provided
         | 
| 445 | 
            +
                    img_attn_output, txt_attn_output = attn_output
         | 
| 446 | 
            +
             | 
| 447 | 
            +
                    # Apply attention gates and add residual (like in Megatron)
         | 
| 448 | 
            +
                    hidden_states = hidden_states + img_gate1 * img_attn_output
         | 
| 449 | 
            +
                    encoder_hidden_states = encoder_hidden_states + txt_gate1 * txt_attn_output
         | 
| 450 | 
            +
             | 
| 451 | 
            +
                    # Process image stream - norm2 + MLP
         | 
| 452 | 
            +
                    img_normed2 = self.img_norm2(hidden_states)
         | 
| 453 | 
            +
                    img_modulated2, img_gate2 = self._modulate(img_normed2, img_mod2)
         | 
| 454 | 
            +
                    img_mlp_output = self.img_mlp(img_modulated2)
         | 
| 455 | 
            +
                    hidden_states = hidden_states + img_gate2 * img_mlp_output
         | 
| 456 | 
            +
             | 
| 457 | 
            +
                    # Process text stream - norm2 + MLP
         | 
| 458 | 
            +
                    txt_normed2 = self.txt_norm2(encoder_hidden_states)
         | 
| 459 | 
            +
                    txt_modulated2, txt_gate2 = self._modulate(txt_normed2, txt_mod2)
         | 
| 460 | 
            +
                    txt_mlp_output = self.txt_mlp(txt_modulated2)
         | 
| 461 | 
            +
                    encoder_hidden_states = encoder_hidden_states + txt_gate2 * txt_mlp_output
         | 
| 462 | 
            +
             | 
| 463 | 
            +
                    # Clip to prevent overflow for fp16
         | 
| 464 | 
            +
                    if encoder_hidden_states.dtype == torch.float16:
         | 
| 465 | 
            +
                        encoder_hidden_states = encoder_hidden_states.clip(-65504, 65504)
         | 
| 466 | 
            +
                    if hidden_states.dtype == torch.float16:
         | 
| 467 | 
            +
                        hidden_states = hidden_states.clip(-65504, 65504)
         | 
| 468 | 
            +
             | 
| 469 | 
            +
                    return encoder_hidden_states, hidden_states
         | 
| 470 | 
            +
             | 
| 471 | 
            +
             | 
| 472 | 
            +
            class QwenImageTransformer2DModel(ModelMixin, ConfigMixin, PeftAdapterMixin, FromOriginalModelMixin, CacheMixin, AttentionMixin):
         | 
| 473 | 
            +
                """
         | 
| 474 | 
            +
                The Transformer model introduced in Qwen.
         | 
| 475 | 
            +
             | 
| 476 | 
            +
                Args:
         | 
| 477 | 
            +
                    patch_size (`int`, defaults to `2`):
         | 
| 478 | 
            +
                        Patch size to turn the input data into small patches.
         | 
| 479 | 
            +
                    in_channels (`int`, defaults to `64`):
         | 
| 480 | 
            +
                        The number of channels in the input.
         | 
| 481 | 
            +
                    out_channels (`int`, *optional*, defaults to `None`):
         | 
| 482 | 
            +
                        The number of channels in the output. If not specified, it defaults to `in_channels`.
         | 
| 483 | 
            +
                    num_layers (`int`, defaults to `60`):
         | 
| 484 | 
            +
                        The number of layers of dual stream DiT blocks to use.
         | 
| 485 | 
            +
                    attention_head_dim (`int`, defaults to `128`):
         | 
| 486 | 
            +
                        The number of dimensions to use for each attention head.
         | 
| 487 | 
            +
                    num_attention_heads (`int`, defaults to `24`):
         | 
| 488 | 
            +
                        The number of attention heads to use.
         | 
| 489 | 
            +
                    joint_attention_dim (`int`, defaults to `3584`):
         | 
| 490 | 
            +
                        The number of dimensions to use for the joint attention (embedding/channel dimension of
         | 
| 491 | 
            +
                        `encoder_hidden_states`).
         | 
| 492 | 
            +
                    guidance_embeds (`bool`, defaults to `False`):
         | 
| 493 | 
            +
                        Whether to use guidance embeddings for guidance-distilled variant of the model.
         | 
| 494 | 
            +
                    axes_dims_rope (`Tuple[int]`, defaults to `(16, 56, 56)`):
         | 
| 495 | 
            +
                        The dimensions to use for the rotary positional embeddings.
         | 
| 496 | 
            +
                """
         | 
| 497 | 
            +
             | 
| 498 | 
            +
                _supports_gradient_checkpointing = True
         | 
| 499 | 
            +
                _no_split_modules = ["QwenImageTransformerBlock"]
         | 
| 500 | 
            +
                _skip_layerwise_casting_patterns = ["pos_embed", "norm"]
         | 
| 501 | 
            +
                _repeated_blocks = ["QwenImageTransformerBlock"]
         | 
| 502 | 
            +
             | 
| 503 | 
            +
                @register_to_config
         | 
| 504 | 
            +
                def __init__(
         | 
| 505 | 
            +
                    self,
         | 
| 506 | 
            +
                    patch_size: int = 2,
         | 
| 507 | 
            +
                    in_channels: int = 64,
         | 
| 508 | 
            +
                    out_channels: Optional[int] = 16,
         | 
| 509 | 
            +
                    num_layers: int = 60,
         | 
| 510 | 
            +
                    attention_head_dim: int = 128,
         | 
| 511 | 
            +
                    num_attention_heads: int = 24,
         | 
| 512 | 
            +
                    joint_attention_dim: int = 3584,
         | 
| 513 | 
            +
                    guidance_embeds: bool = False,  # TODO: this should probably be removed
         | 
| 514 | 
            +
                    axes_dims_rope: Tuple[int, int, int] = (16, 56, 56),
         | 
| 515 | 
            +
                ):
         | 
| 516 | 
            +
                    super().__init__()
         | 
| 517 | 
            +
                    self.out_channels = out_channels or in_channels
         | 
| 518 | 
            +
                    self.inner_dim = num_attention_heads * attention_head_dim
         | 
| 519 | 
            +
             | 
| 520 | 
            +
                    self.pos_embed = QwenEmbedRope(theta=10000, axes_dim=list(axes_dims_rope), scale_rope=True)
         | 
| 521 | 
            +
             | 
| 522 | 
            +
                    self.time_text_embed = QwenTimestepProjEmbeddings(embedding_dim=self.inner_dim)
         | 
| 523 | 
            +
             | 
| 524 | 
            +
                    self.txt_norm = RMSNorm(joint_attention_dim, eps=1e-6)
         | 
| 525 | 
            +
             | 
| 526 | 
            +
                    self.img_in = nn.Linear(in_channels, self.inner_dim)
         | 
| 527 | 
            +
                    self.txt_in = nn.Linear(joint_attention_dim, self.inner_dim)
         | 
| 528 | 
            +
             | 
| 529 | 
            +
                    self.transformer_blocks = nn.ModuleList(
         | 
| 530 | 
            +
                        [
         | 
| 531 | 
            +
                            QwenImageTransformerBlock(
         | 
| 532 | 
            +
                                dim=self.inner_dim,
         | 
| 533 | 
            +
                                num_attention_heads=num_attention_heads,
         | 
| 534 | 
            +
                                attention_head_dim=attention_head_dim,
         | 
| 535 | 
            +
                            )
         | 
| 536 | 
            +
                            for _ in range(num_layers)
         | 
| 537 | 
            +
                        ]
         | 
| 538 | 
            +
                    )
         | 
| 539 | 
            +
             | 
| 540 | 
            +
                    self.norm_out = AdaLayerNormContinuous(self.inner_dim, self.inner_dim, elementwise_affine=False, eps=1e-6)
         | 
| 541 | 
            +
                    self.proj_out = nn.Linear(self.inner_dim, patch_size * patch_size * self.out_channels, bias=True)
         | 
| 542 | 
            +
             | 
| 543 | 
            +
                    self.gradient_checkpointing = False
         | 
| 544 | 
            +
             | 
| 545 | 
            +
                def forward(
         | 
| 546 | 
            +
                    self,
         | 
| 547 | 
            +
                    hidden_states: torch.Tensor,
         | 
| 548 | 
            +
                    encoder_hidden_states: torch.Tensor = None,
         | 
| 549 | 
            +
                    encoder_hidden_states_mask: torch.Tensor = None,
         | 
| 550 | 
            +
                    timestep: torch.LongTensor = None,
         | 
| 551 | 
            +
                    image_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
         | 
| 552 | 
            +
                    guidance: torch.Tensor = None,  # TODO: this should probably be removed
         | 
| 553 | 
            +
                    attention_kwargs: Optional[Dict[str, Any]] = None,
         | 
| 554 | 
            +
                    return_dict: bool = True,
         | 
| 555 | 
            +
                ) -> Union[torch.Tensor, Transformer2DModelOutput]:
         | 
| 556 | 
            +
                    """
         | 
| 557 | 
            +
                    The [`QwenTransformer2DModel`] forward method.
         | 
| 558 | 
            +
             | 
| 559 | 
            +
                    Args:
         | 
| 560 | 
            +
                        hidden_states (`torch.Tensor` of shape `(batch_size, image_sequence_length, in_channels)`):
         | 
| 561 | 
            +
                            Input `hidden_states`.
         | 
| 562 | 
            +
                        encoder_hidden_states (`torch.Tensor` of shape `(batch_size, text_sequence_length, joint_attention_dim)`):
         | 
| 563 | 
            +
                            Conditional embeddings (embeddings computed from the input conditions such as prompts) to use.
         | 
| 564 | 
            +
                        encoder_hidden_states_mask (`torch.Tensor` of shape `(batch_size, text_sequence_length)`):
         | 
| 565 | 
            +
                            Mask of the input conditions.
         | 
| 566 | 
            +
                        timestep ( `torch.LongTensor`):
         | 
| 567 | 
            +
                            Used to indicate denoising step.
         | 
| 568 | 
            +
                        attention_kwargs (`dict`, *optional*):
         | 
| 569 | 
            +
                            A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under
         | 
| 570 | 
            +
                            `self.processor` in
         | 
| 571 | 
            +
                            [diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py).
         | 
| 572 | 
            +
                        return_dict (`bool`, *optional*, defaults to `True`):
         | 
| 573 | 
            +
                            Whether or not to return a [`~models.transformer_2d.Transformer2DModelOutput`] instead of a plain
         | 
| 574 | 
            +
                            tuple.
         | 
| 575 | 
            +
             | 
| 576 | 
            +
                    Returns:
         | 
| 577 | 
            +
                        If `return_dict` is True, an [`~models.transformer_2d.Transformer2DModelOutput`] is returned, otherwise a
         | 
| 578 | 
            +
                        `tuple` where the first element is the sample tensor.
         | 
| 579 | 
            +
                    """
         | 
| 580 | 
            +
                    if attention_kwargs is not None:
         | 
| 581 | 
            +
                        attention_kwargs = attention_kwargs.copy()
         | 
| 582 | 
            +
                        lora_scale = attention_kwargs.pop("scale", 1.0)
         | 
| 583 | 
            +
                    else:
         | 
| 584 | 
            +
                        lora_scale = 1.0
         | 
| 585 | 
            +
             | 
| 586 | 
            +
                    if USE_PEFT_BACKEND:
         | 
| 587 | 
            +
                        # weight the lora layers by setting `lora_scale` for each PEFT layer
         | 
| 588 | 
            +
                        scale_lora_layers(self, lora_scale)
         | 
| 589 | 
            +
                    else:
         | 
| 590 | 
            +
                        if attention_kwargs is not None and attention_kwargs.get("scale", None) is not None:
         | 
| 591 | 
            +
                            logger.warning(
         | 
| 592 | 
            +
                                "Passing `scale` via `joint_attention_kwargs` when not using the PEFT backend is ineffective."
         | 
| 593 | 
            +
                            )
         | 
| 594 | 
            +
             | 
| 595 | 
            +
                    hidden_states = self.img_in(hidden_states)
         | 
| 596 | 
            +
             | 
| 597 | 
            +
                    timestep = timestep.to(hidden_states.dtype)
         | 
| 598 | 
            +
                    encoder_hidden_states = self.txt_norm(encoder_hidden_states)
         | 
| 599 | 
            +
                    encoder_hidden_states = self.txt_in(encoder_hidden_states)
         | 
| 600 | 
            +
             | 
| 601 | 
            +
                    if guidance is not None:
         | 
| 602 | 
            +
                        guidance = guidance.to(hidden_states.dtype) * 1000
         | 
| 603 | 
            +
             | 
| 604 | 
            +
                    temb = (
         | 
| 605 | 
            +
                        self.time_text_embed(timestep, hidden_states)
         | 
| 606 | 
            +
                        if guidance is None
         | 
| 607 | 
            +
                        else self.time_text_embed(timestep, guidance, hidden_states)
         | 
| 608 | 
            +
                    )
         | 
| 609 | 
            +
             | 
| 610 | 
            +
                    for index_block, block in enumerate(self.transformer_blocks):
         | 
| 611 | 
            +
                        if torch.is_grad_enabled() and self.gradient_checkpointing:
         | 
| 612 | 
            +
                            encoder_hidden_states, hidden_states = self._gradient_checkpointing_func(
         | 
| 613 | 
            +
                                block,
         | 
| 614 | 
            +
                                hidden_states,
         | 
| 615 | 
            +
                                encoder_hidden_states,
         | 
| 616 | 
            +
                                encoder_hidden_states_mask,
         | 
| 617 | 
            +
                                temb,
         | 
| 618 | 
            +
                                image_rotary_emb,
         | 
| 619 | 
            +
                            )
         | 
| 620 | 
            +
             | 
| 621 | 
            +
                        else:
         | 
| 622 | 
            +
                            encoder_hidden_states, hidden_states = block(
         | 
| 623 | 
            +
                                hidden_states=hidden_states,
         | 
| 624 | 
            +
                                encoder_hidden_states=encoder_hidden_states,
         | 
| 625 | 
            +
                                encoder_hidden_states_mask=encoder_hidden_states_mask,
         | 
| 626 | 
            +
                                temb=temb,
         | 
| 627 | 
            +
                                image_rotary_emb=image_rotary_emb,
         | 
| 628 | 
            +
                                joint_attention_kwargs=attention_kwargs,
         | 
| 629 | 
            +
                            )
         | 
| 630 | 
            +
             | 
| 631 | 
            +
                    # Use only the image part (hidden_states) from the dual-stream blocks
         | 
| 632 | 
            +
                    hidden_states = self.norm_out(hidden_states, temb)
         | 
| 633 | 
            +
                    output = self.proj_out(hidden_states)
         | 
| 634 | 
            +
             | 
| 635 | 
            +
                    if USE_PEFT_BACKEND:
         | 
| 636 | 
            +
                        # remove `lora_scale` from each PEFT layer
         | 
| 637 | 
            +
                        unscale_lora_layers(self, lora_scale)
         | 
| 638 | 
            +
             | 
| 639 | 
            +
                    if not return_dict:
         | 
| 640 | 
            +
                        return (output,)
         | 
| 641 | 
            +
             | 
| 642 | 
            +
                    return Transformer2DModelOutput(sample=output)
         | 
    	
        requirements.txt
    CHANGED
    
    | @@ -7,4 +7,5 @@ sentencepiece | |
| 7 | 
             
            dashscope
         | 
| 8 | 
             
            kernels
         | 
| 9 | 
             
            torchvision
         | 
| 10 | 
            -
            peft
         | 
|  | 
|  | |
| 7 | 
             
            dashscope
         | 
| 8 | 
             
            kernels
         | 
| 9 | 
             
            torchvision
         | 
| 10 | 
            +
            peft
         | 
| 11 | 
            +
            torchao==0.11.0
         | 

