fix: old SD3Transformer2DModel has no attributes

deepgen_pipeline.py

@@ -961,12 +961,31 @@ class DeepGenPipeline(DiffusionPipeline):
         with cond_latents support for image editing. No weight copying needed."""
         from diffusers.models.transformers.transformer_sd3 import SD3Transformer2DModel as _OrigSD3
         if isinstance(self.transformer, _OrigSD3) and not isinstance(self.transformer, SD3Transformer2DModel):
-            self.transformer.__class__ = SD3Transformer2DModel
-            for block in self.transformer.transformer_blocks:
-                block.__class__ = CustomJointTransformerBlock
-                block.attn.set_processor(CustomJointAttnProcessor2_0())
-                if block.attn2 is not None:
-                    block.attn2.set_processor(CustomJointAttnProcessor2_0())
+            state_dict = self.transformer.state_dict()
+            config = self.transformer.config
+            custom_config = {
+                'sample_size': config.sample_size,
+                'patch_size': config.patch_size,
+                'in_channels': config.in_channels,
+                'num_layers': config.num_layers,
+                'attention_head_dim': config.attention_head_dim,
+                'num_attention_heads': config.num_attention_heads,
+                'joint_attention_dim': config.joint_attention_dim,
+                'caption_projection_dim': config.caption_projection_dim,
+                'pooled_projection_dim': config.pooled_projection_dim,
+                'out_channels': config.out_channels,
+                'pos_embed_max_size': getattr(config, 'pos_embed_max_size', 96),
+                'dual_attention_layers': getattr(config, 'dual_attention_layers', ()),
+                'qk_norm': getattr(config, 'qk_norm', None),
+            }
+            device = self.transformer.device
+            dtype = self.transformer.dtype
+            self.transformer = SD3Transformer2DModel(**custom_config).to(device=device, dtype=dtype)
+            self.transformer.load_state_dict(state_dict, strict=False)
+            
+            # Set gradient checkpointing if it was enabled
+            if hasattr(self.transformer, 'gradient_checkpointing'):
+                self.transformer.gradient_checkpointing = self.transformer.gradient_checkpointing
 
     def _resolve_pretrained_path(self):
         path = self.config._name_or_path