remove unuse code

.gitignore

@@ -0,0 +1,3 @@
+*.un~
+*.pyc
+__pycache__

demo.py

@@ -1,11 +1,12 @@
 from lyraSD import LyraSD
 
 t2imodel = LyraSD("text2img", "./sd1.5-engine")
-t2imodel.inference(prompt="a red ballon flying in the sky")
+t2imodel.inference(prompt="A fantasy landscape, trending on artstation", use_super=True)
 
 
 from PIL import Image
 i2imodel = LyraSD("img2img", "./sd1.5-engine")
 demo_img = Image.open("output/text2img_demo.jpg")
-i2imodel.inference(prompt="comic style", image=demo_img)
+i2imodel.inference(prompt="A fantasy landscape, trending on artstation", 
+                   image=demo_img)
 

lyraSD/muse_trt/models.py

@@ -259,44 +259,6 @@ class VAEEncoder(BaseModel):
             batch_size, image_height, image_width)
         return torch.randn(batch_size, 3, image_height, image_width, dtype=torch.float32, device=self.device)
 
-    def optimize(self, onnx_graph, minimal_optimization=False):
-        enable_optimization = not minimal_optimization
-
-        # Decompose InstanceNormalization into primitive Ops
-        bRemoveInstanceNorm = enable_optimization
-        # Remove Cast Node to optimize Attention block
-        bRemoveCastNode = enable_optimization
-        # Insert GroupNormalization Plugin
-        bGroupNormPlugin = enable_optimization
-
-        opt = Optimizer(onnx_graph, verbose=self.verbose)
-        opt.info('VAE Encoder: original')
-
-        if bRemoveInstanceNorm:
-            num_instancenorm_replaced = opt.decompose_instancenorms()
-            opt.info('VAE Encoder: replaced ' +
-                     str(num_instancenorm_replaced)+' InstanceNorms')
-
-        if bRemoveCastNode:
-            num_casts_removed = opt.remove_casts()
-            opt.info('VAE Encoder: removed '+str(num_casts_removed)+' casts')
-
-        opt.cleanup()
-        opt.info('VAE Encoder: cleanup')
-        opt.fold_constants()
-        opt.info('VAE Encoder: fold constants')
-        opt.infer_shapes()
-        opt.info('VAE Encoder: shape inference')
-
-        if bGroupNormPlugin:
-            num_groupnorm_inserted = opt.insert_groupnorm_plugin()
-            opt.info('VAE Encoder: inserted '+str(num_groupnorm_inserted) +
-                     ' GroupNorm plugins')
-
-        onnx_opt_graph = opt.cleanup(return_onnx=True)
-        opt.info('VAE Encoder: final')
-        return onnx_opt_graph
-
 
 class VAEDecoder(BaseModel):
     def get_model(self):
@@ -345,471 +307,4 @@ class VAEDecoder(BaseModel):
     def get_sample_input(self, batch_size, image_height, image_width):
         latent_height, latent_width = self.check_dims(
             batch_size, image_height, image_width)
-        return torch.randn(batch_size, 4, latent_height, latent_width, dtype=torch.float32, device=self.device)
-
-    def optimize(self, onnx_graph, minimal_optimization=False):
-        enable_optimization = not minimal_optimization
-
-        # Decompose InstanceNormalization into primitive Ops
-        bRemoveInstanceNorm = enable_optimization
-        # Remove Cast Node to optimize Attention block
-        bRemoveCastNode = enable_optimization
-        # Insert GroupNormalization Plugin
-        bGroupNormPlugin = enable_optimization
-
-        opt = Optimizer(onnx_graph, verbose=self.verbose)
-        opt.info('VAE Decoder: original')
-
-        if bRemoveInstanceNorm:
-            num_instancenorm_replaced = opt.decompose_instancenorms()
-            opt.info('VAE Decoder: replaced ' +
-                     str(num_instancenorm_replaced)+' InstanceNorms')
-
-        if bRemoveCastNode:
-            num_casts_removed = opt.remove_casts()
-            opt.info('VAE Decoder: removed '+str(num_casts_removed)+' casts')
-
-        opt.cleanup()
-        opt.info('VAE Decoder: cleanup')
-        opt.fold_constants()
-        opt.info('VAE Decoder: fold constants')
-        opt.infer_shapes()
-        opt.info('VAE Decoder: shape inference')
-
-        if bGroupNormPlugin:
-            num_groupnorm_inserted = opt.insert_groupnorm_plugin()
-            opt.info('VAE Decoder: inserted '+str(num_groupnorm_inserted) +
-                     ' GroupNorm plugins')
-
-        onnx_opt_graph = opt.cleanup(return_onnx=True)
-        opt.info('VAE Decoder: final')
-        return onnx_opt_graph
-
-
-class SuperModelX4(nn.Module):
-    def __init__(self, model_dir, scale=4, pre_pad=0):
-        super().__init__()
-        self.scale = scale
-        self.pre_pad = pre_pad
-
-        rrdb = RealESRGAN(model_dir=model_dir,
-                          model_name="RealESRGAN_x4plus_anime_6B").upsampler.model
-        self.rrdb = rrdb.eval()
-
-    def forward(self, x):
-        x = x / 255.
-        x = F.pad(x, (0, self.pre_pad, 0, self.pre_pad), 'reflect')
-        x = self.rrdb(x)
-        _, _, h, w = x.size()
-        x = x[:, :, 0:h-self.pre_pad * self.scale, 0:w-self.pre_pad*self.scale]
-        x = x.clamp(0, 1)
-        x = (x * 255).round()
-        return x
-
-
-class SuperResX4():
-    def __init__(
-        self,
-        local_model_path=None,
-        fp16=True,
-        device='cuda',
-        verbose=True,
-        max_batch_size=8
-    ):
-        self.fp16 = fp16
-        self.device = device
-        self.verbose = verbose
-        self.local_model_path = local_model_path
-
-        # Defaults
-        self.min_batch = 1
-        self.max_batch = max_batch_size
-        self.min_height = 64
-        self.max_height = 640
-        self.min_width = 64
-        self.max_width = 640
-
-    def get_model(self):
-        model = SuperModelX4(self.local_model_path, scale=4, pre_pad=0).to(device=self.device)
-        if self.fp16:
-            model = model.half()
-        return model
-
-    def get_input_names(self):
-        return ['input_image']
-
-    def get_output_names(self):
-        return ['output_image']
-
-    def get_dynamic_axes(self):
-        return {
-            'input_image': {0: 'B', },
-            'output_image': {0: 'B', }
-        }
-
-    def check_dims(self, batch_size, image_height, image_width):
-        assert batch_size >= self.min_batch and batch_size <= self.max_batch
-        return (image_height, image_width)
-
-    def get_minmax_dims(self, batch_size, image_height, image_width, static_batch, static_shape):
-        min_batch = batch_size if static_batch else self.min_batch
-        max_batch = batch_size if static_batch else self.max_batch
-        min_image_height = image_height if static_shape else self.min_height
-        max_image_height = image_height if static_shape else self.max_height
-        min_image_width = image_width if static_shape else self.min_width
-        max_image_width = image_width if static_shape else self.max_width
-        return (min_batch, max_batch, min_image_height, max_image_height, min_image_width, max_image_width)
-
-    def get_input_profile(self, batch_size, image_height, image_width, static_batch, static_shape):
-        image_height, image_width = self.check_dims(
-            batch_size, image_height, image_width)
-        min_batch, max_batch, min_image_height, max_image_height, min_image_width, max_image_width = \
-            self.get_minmax_dims(batch_size, image_height,
-                                 image_width, static_batch, static_shape)
-        return {
-            'input_image': [(min_batch, 3, min_image_height, min_image_width), (batch_size, 3, image_height, image_width), (max_batch, 3, max_image_height, max_image_width)]
-        }
-
-    def get_shape_dict(self, batch_size, image_height, image_width):
-        image_height, image_width = self.check_dims(
-            batch_size, image_height, image_width)
-        return {
-            'input_image': (batch_size, 3, image_height, image_width),
-            'output_image': (batch_size, 3, image_height*4, image_width*4),
-        }
-
-    def get_sample_input(self, batch_size, image_height, image_width):
-        dtype = torch.float16 if self.fp16 else torch.float32
-        image_height, image_width = self.check_dims(
-            batch_size, image_height, image_width)
-        return torch.randn(batch_size, 3, image_height, image_width, dtype=dtype, device=self.device)
-
-    def optimize(self, onnx_graph, minimal_optimization=False):
-        enable_optimization = not minimal_optimization
-
-        # Decompose InstanceNormalization into primitive Ops
-        bRemoveInstanceNorm = enable_optimization
-        # Remove Cast Node to optimize Attention block
-        bRemoveCastNode = enable_optimization
-        # Insert GroupNormalization Plugin
-        bGroupNormPlugin = enable_optimization
-
-        opt = Optimizer(onnx_graph, verbose=self.verbose)
-        opt.info('SuperX4: original')
-
-        if bRemoveInstanceNorm:
-            num_instancenorm_replaced = opt.decompose_instancenorms()
-            opt.info('SuperX4: replaced ' +
-                     str(num_instancenorm_replaced)+' InstanceNorms')
-
-        if bRemoveCastNode:
-            num_casts_removed = opt.remove_casts()
-            opt.info('SuperX4: removed '+str(num_casts_removed)+' casts')
-
-        opt.cleanup()
-        opt.info('SuperX4: cleanup')
-        opt.fold_constants()
-        opt.info('SuperX4: fold constants')
-        opt.infer_shapes()
-        opt.info('SuperX4: shape inference')
-
-        if bGroupNormPlugin:
-            num_groupnorm_inserted = opt.insert_groupnorm_plugin()
-            opt.info('SuperX4: inserted '+str(num_groupnorm_inserted) +
-                     ' GroupNorm plugins')
-
-        onnx_opt_graph = opt.cleanup(return_onnx=True)
-        opt.info('SuperX4: final')
-        return onnx_opt_graph
-
-
-class FusedControlNetModule(nn.Module):
-    def __init__(self, base_model_dir, control_model_dir, fp16=True) -> None:
-        super().__init__()
-        self.device = 'cuda:0'
-        self.fp16 = fp16
-        model_opts = {'revision': 'fp16',
-                      'torch_dtype': torch.float16} if self.fp16 else {}
-        self.base = UNet2DConditionModel.from_pretrained(
-            base_model_dir, subfolder="unet",
-            **model_opts
-        ).eval().to(self.device)
-        self.control = ControlNetModel.from_pretrained(
-            control_model_dir,
-            **model_opts
-        ).eval().to(self.device)
-
-    def forward(self, sample, timestep, encoder_hidden_states, controlnet_cond):
-        controlnet_conditioning_scale: float = 1.0
-        down_block_res_samples, mid_block_res_sample = self.control(
-            sample,
-            timestep,
-            encoder_hidden_states=encoder_hidden_states,
-            controlnet_cond=controlnet_cond,
-            return_dict=False,
-        )
-
-        down_block_res_samples = [
-            down_block_res_sample * controlnet_conditioning_scale
-            for down_block_res_sample in down_block_res_samples
-        ]
-        mid_block_res_sample *= controlnet_conditioning_scale
-
-        # predict the noise residual
-        noise_pred = self.base(
-            sample,
-            timestep,
-            encoder_hidden_states=encoder_hidden_states,
-            down_block_additional_residuals=down_block_res_samples,
-            mid_block_additional_residual=mid_block_res_sample,
-        ).sample
-
-        return noise_pred
-
-
-class FusedControlNet(BaseModel):
-    def __init__(self, local_model_path=None, controlnet_model_path=None, hf_token=None, text_maxlen=77,
-                 embedding_dim=768, fp16=False, device='cuda', verbose=True, max_batch_size=16):
-        super().__init__(local_model_path, hf_token, text_maxlen, embedding_dim, fp16, device, verbose, max_batch_size)
-        # if controlnet_model_path is None:
-        #     raise ValueError("Must give controlnet_model_path for FusedControlNet to load control net")
-        self.controlnet_model_path = controlnet_model_path
-        self.min_height = 256
-        self.max_height = 1024
-        self.min_width = 256
-        self.max_width = 1024
-
-    def get_minmax_dims(self, batch_size, image_height, image_width, static_batch, static_shape):
-        r = list(super().get_minmax_dims(batch_size, image_height, image_width, static_batch, static_shape))
-        min_height = image_height if static_shape else self.min_height
-        max_height = image_height if static_shape else self.max_height
-        min_width = image_width if static_shape else self.min_width
-        max_width = image_width if static_shape else self.max_width
-        r.extend([min_height, max_height, min_width, max_width])
-        return r
-
-    def get_model(self):
-        model = FusedControlNetModule(
-            base_model_dir=self.local_model_path,
-            control_model_dir=self.controlnet_model_path,
-            fp16=self.fp16
-        )
-        return model
-
-    def get_input_names(self):
-        return ['sample', 'timestep', 'encoder_hidden_states', 'controlnet_cond']
-
-    def get_output_names(self):
-        return ['latent']
-
-    def get_dynamic_axes(self):
-        return {
-            'sample': {0: '2B', 2: 'H', 3: 'W'},
-            'encoder_hidden_states': {0: '2B'},
-            'controlnet_cond': {0: '2B', 2: '8H', 3: '8W'},  # controlnet_cond is 8X sample and lantent
-            'latent': {0: '2B', 2: 'H', 3: 'W'}
-        }
-
-    def get_input_profile(self, batch_size, image_height, image_width, static_batch, static_shape):
-        latent_height, latent_width = self.check_dims(
-            batch_size, image_height, image_width)
-        min_batch, max_batch, min_latent_height, max_latent_height, min_latent_width, max_latent_width, min_height, max_height, min_width, max_width = \
-            self.get_minmax_dims(batch_size, image_height,
-                                 image_width, static_batch, static_shape)
-        return {
-            'sample': [(2*min_batch, 4, min_latent_height, min_latent_width), (2*batch_size, 4, latent_height, latent_width), (2*max_batch, 4, max_latent_height, max_latent_width)],
-            'encoder_hidden_states': [(2*min_batch, self.text_maxlen, self.embedding_dim), (2*batch_size, self.text_maxlen, self.embedding_dim), (2*max_batch, self.text_maxlen, self.embedding_dim)],
-            'controlnet_cond': [(2*min_batch, 3, min_height, min_width), (2*batch_size, 3, image_height, image_width), (2*max_batch, 3, max_height, max_width)]
-        }
-
-    def get_shape_dict(self, batch_size, image_height, image_width):
-        latent_height, latent_width = self.check_dims(
-            batch_size, image_height, image_width)
-        return {
-            'sample': (2*batch_size, 4, latent_height, latent_width),
-            'encoder_hidden_states': (2*batch_size, self.text_maxlen, self.embedding_dim),
-            'controlnet_cond': (2*batch_size, 3, image_height, image_width),
-            'latent': (2*batch_size, 4, latent_height, latent_width)
-        }
-
-    def get_sample_input(self, batch_size, image_height, image_width):
-        latent_height, latent_width = self.check_dims(
-            batch_size, image_height, image_width)
-        dtype = torch.float16 if self.fp16 else torch.float32
-        return (
-            torch.randn(2*batch_size, 4, latent_height, latent_width,
-                        dtype=torch.float32, device=self.device),   # sample
-            torch.tensor([1.], dtype=torch.float32, device=self.device),    # timestep
-            torch.randn(2*batch_size, self.text_maxlen,  # encoder_hidden_states
-                        self.embedding_dim, dtype=dtype, device=self.device),
-            torch.randn(2*batch_size, 3, image_height, image_width,
-                        dtype=torch.float32, device=self.device)    # controlnet_cond
-        )
-
-    def optimize(self, onnx_graph, minimal_optimization=False):
-        class_name = self.__class__.__name__
-
-        enable_optimization = not minimal_optimization
-
-        # Decompose InstanceNormalization into primitive Ops
-        bRemoveInstanceNorm = enable_optimization
-        # Remove Cast Node to optimize Attention block
-        bRemoveCastNode = enable_optimization
-        # Remove parallel Swish ops
-        bRemoveParallelSwish = enable_optimization
-        # Adjust the bias to be the second input to the Add ops
-        bAdjustAddNode = enable_optimization
-        # Change Resize node to take size instead of scale
-        bResizeFix = enable_optimization
-
-        # Common override for disabling all plugins below
-        bDisablePlugins = minimal_optimization
-        # Use multi-head attention Plugin
-        bMHAPlugin = True
-        # Use multi-head cross attention Plugin
-        bMHCAPlugin = True
-        # Insert GroupNormalization Plugin
-        bGroupNormPlugin = True
-        # Insert LayerNormalization Plugin
-        bLayerNormPlugin = True
-        # Insert Split+GeLU Plugin
-        bSplitGeLUPlugin = True
-        # Replace BiasAdd+ResidualAdd+SeqLen2Spatial with plugin
-        bSeqLen2SpatialPlugin = True
-
-        opt = Optimizer(onnx_graph, verbose=self.verbose)
-        opt.info(f'{class_name}: original')
-
-        if bRemoveInstanceNorm:
-            num_instancenorm_replaced = opt.decompose_instancenorms()
-            opt.info(f'{class_name}: replaced ' +
-                     str(num_instancenorm_replaced)+' InstanceNorms')
-
-        if bRemoveCastNode:
-            num_casts_removed = opt.remove_casts()
-            opt.info(f'{class_name}: removed '+str(num_casts_removed)+' casts')
-
-        if bRemoveParallelSwish:
-            num_parallel_swish_removed = opt.remove_parallel_swish()
-            opt.info(f'{class_name}: removed ' +
-                     str(num_parallel_swish_removed)+' parallel swish ops')
-
-        if bAdjustAddNode:
-            num_adjust_add = opt.adjustAddNode()
-            opt.info(f'{class_name}: adjusted '+str(num_adjust_add)+' adds')
-
-        if bResizeFix:
-            num_resize_fix = opt.resize_fix()
-            opt.info(f'{class_name}: fixed '+str(num_resize_fix)+' resizes')
-
-        opt.cleanup()
-        opt.info(f'{class_name}: cleanup')
-        opt.fold_constants()
-        opt.info(f'{class_name}: fold constants')
-        opt.infer_shapes()
-        opt.info(f'{class_name}: shape inference')
-
-        num_heads = 8
-        if bMHAPlugin and not bDisablePlugins:
-            num_fmha_inserted = opt.insert_fmha_plugin(num_heads)
-            opt.info(f'{class_name}: inserted '+str(num_fmha_inserted)+' fMHA plugins')
-
-        if bMHCAPlugin and not bDisablePlugins:
-            props = cudart.cudaGetDeviceProperties(0)[1]
-            sm = props.major * 10 + props.minor
-            num_fmhca_inserted = opt.insert_fmhca_plugin(num_heads, sm)
-            opt.info(f'{class_name}: inserted '+str(num_fmhca_inserted)+' fMHCA plugins')
-
-        if bGroupNormPlugin and not bDisablePlugins:
-            num_groupnorm_inserted = opt.insert_groupnorm_plugin()
-            opt.info(f'{class_name}: inserted '+str(num_groupnorm_inserted) +
-                     ' GroupNorm plugins')
-
-        if bLayerNormPlugin and not bDisablePlugins:
-            num_layernorm_inserted = opt.insert_layernorm_plugin()
-            opt.info(f'{class_name}: inserted '+str(num_layernorm_inserted) +
-                     ' LayerNorm plugins')
-
-        if bSplitGeLUPlugin and not bDisablePlugins:
-            num_splitgelu_inserted = opt.insert_splitgelu_plugin()
-            opt.info(f'{class_name}: inserted '+str(num_splitgelu_inserted) +
-                     ' SplitGeLU plugins')
-
-        if bSeqLen2SpatialPlugin and not bDisablePlugins:
-            num_seq2spatial_inserted = opt.insert_seq2spatial_plugin()
-            opt.info(f'{class_name}: inserted '+str(num_seq2spatial_inserted) +
-                     ' SeqLen2Spatial plugins')
-
-        onnx_opt_graph = opt.cleanup(return_onnx=True)
-        opt.info(f'{class_name}: final')
-        return onnx_opt_graph
-
-
-class ControlNetModule(nn.Module):
-    def __init__(self, control_model_dir, fp16=True) -> None:
-        super().__init__()
-        self.device = 'cuda:0'
-        self.fp16 = fp16
-        model_opts = {'revision': 'fp16',
-                      'torch_dtype': torch.float16} if self.fp16 else {}
-        self.control = ControlNetModel.from_pretrained(
-            control_model_dir,
-            **model_opts
-        ).eval().to(self.device)
-
-    def forward(self, sample, timestep, encoder_hidden_states, controlnet_cond):
-        controlnet_conditioning_scale: float = 1.0
-        down_block_res_samples, mid_block_res_sample = self.control(
-            sample,
-            timestep,
-            encoder_hidden_states=encoder_hidden_states,
-            controlnet_cond=controlnet_cond,
-            return_dict=False,
-        )
-        down_block_res_samples = [
-            down_block_res_sample * controlnet_conditioning_scale
-            for down_block_res_sample in down_block_res_samples
-        ]
-        mid_block_res_sample *= controlnet_conditioning_scale
-        # @vane: currently, only retun mid_blocks_res_sample: (B, 1280, height//8//8, width//8//8)
-        # down_block_res_samples is a tensor tuple that length is 12.
-        # it will be flatten to 12 nodes if we return the down_block_res_samples
-        return mid_block_res_sample
-
-
-class ControlNet(FusedControlNet):
-    def __init__(self, local_model_path=None, controlnet_model_path=None, hf_token=None, text_maxlen=77,
-                 embedding_dim=768, fp16=False, device='cuda', verbose=True, max_batch_size=16):
-        super().__init__(local_model_path, controlnet_model_path, hf_token,
-                         text_maxlen, embedding_dim, fp16, device, verbose, max_batch_size)
-
-    def get_model(self):
-        model = ControlNetModule(
-            control_model_dir=self.controlnet_model_path,
-            fp16=self.fp16
-        )
-        return model
-
-    def get_input_names(self):
-        return ['sample', 'timestep', 'encoder_hidden_states', 'controlnet_cond']
-
-    def get_output_names(self):
-        return ['mids']
-
-    def get_dynamic_axes(self):
-        return {
-            'sample': {0: '2B', 2: '8H', 3: '8W'},
-            'encoder_hidden_states': {0: '2B'},
-            'controlnet_cond': {0: '2B', 2: '16H', 3: '16W'},
-            'mids': {0: '2B', 2: 'H', 3: 'W'}
-        }
-
-    def get_shape_dict(self, batch_size, image_height, image_width):
-        latent_height, latent_width = self.check_dims(
-            batch_size, image_height, image_width)
-        return {
-            'sample': (2*batch_size, 4, latent_height, latent_width),
-            'encoder_hidden_states': (2*batch_size, self.text_maxlen, self.embedding_dim),
-            'controlnet_cond': (2*batch_size, 3, image_height, image_width),
-            'mids': (2*batch_size, 1280, latent_height//8, latent_width//8)
-        }
+        return torch.randn(batch_size, 4, latent_height, latent_width, dtype=torch.float32, device=self.device)