write complete correct code below 
[bold green]Original Model Summary[/bold green]
╭────────────╮
│ Sequential │
╰────────────╯
      Total Parameters Summary       
┏━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━┓
┃ Category                 ┃  Count ┃
┡━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━┩
│ Total Parameters         │ 40,298 │
│ Trainable Parameters     │ 40,298 │
│ Non-trainable Parameters │      0 │
└──────────────────────────┴────────┘
╭─────────────────────────────────────╮
│ Summary of model's total parameters │
╰─────────────────────────────────────╯
                Detailed Parameter Summary
┏━━━━━━━━━━━┳━━━━━━━━━━━━━━━━┳━━━━━━━━━━━┳━━━━━━━━━━━━━━━━┓
┃ Parameter ┃ Shape          ┃ Trainable ┃ Num Parameters ┃
┡━━━━━━━━━━━╇━━━━━━━━━━━━━━━━╇━━━━━━━━━━━╇━━━━━━━━━━━━━━━━┩
│ 0.weight  │ (256, 128)     │ ✓         │         32,768 │
├───────────┼────────────────┼───────────┼────────────────┤
│ 0.bias    │ (256,)         │ ✓         │            256 │
├───────────┼────────────────┼───────────┼────────────────┤
│ 2.weight  │ (32, 16, 3, 3) │ ✓         │          4,608 │
├───────────┼────────────────┼───────────┼────────────────┤
│ 2.bias    │ (32,)          │ ✓         │             32 │
├───────────┼────────────────┼───────────┼────────────────┤
│ 3.weight  │ (32,)          │ ✓         │             32 │
├───────────┼────────────────┼───────────┼────────────────┤
│ 3.bias    │ (32,)          │ ✓         │             32 │
├───────────┼────────────────┼───────────┼────────────────┤
│ 5.weight  │ (10, 256)      │ ✓         │          2,560 │
├───────────┼────────────────┼───────────┼────────────────┤
│ 5.bias    │ (10,)          │ ✓         │             10 │
└───────────┴────────────────┴───────────┴────────────────┘
╭──────────────────────────────────────────╮
│ Detailed breakdown of model's parameters │
╰──────────────────────────────────────────╯

Model Architecture
Sequential
├── 0
│   ├── weight: (256, 128) (Trainable)
│   └── bias: (256,) (Trainable)
├── 1
├── 2
│   ├── weight: (32, 16, 3, 3) (Trainable)
│   └── bias: (32,) (Trainable)
├── 3
│   ├── weight: (32,) (Trainable)
│   └── bias: (32,) (Trainable)
├── 4
└── 5
    ├── weight: (10, 256) (Trainable)
    └── bias: (10,) (Trainable)
╭───────────────────────────────────────────────────╮
│ Visual representation of the model's architecture │
╰───────────────────────────────────────────────────╯
ERROR:__main__:Error quantizing module Conv2d(16, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)): Conv2d.__init__() got multiple values for argument 'bias'
ERROR:__main__:Error during model quantization: Error quantizing module Conv2d(16, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)): Conv2d.__init__() got multiple values for argument 'bias'  
ERROR:__main__:Quantization failed: Error during model quantization: Error quantizing module Conv2d(16, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)): Conv2d.__init__() got multiple values for argument 'bias'
import copy
import logging
from typing import Callable, Optional, Type, Union

import torch
import torch.nn as nn
from torch import Tensor
from rich.console import Console
from rich.panel import Panel
from rich.table import Table
from rich.tree import Tree

# Configure logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# ----------------------------- Exception Classes -----------------------------

class QuantizationError(Exception):
    """Custom exception for quantization errors."""
    pass

# ----------------------------- Quantization Functions ------------------------

def round_clip(x: Tensor, a: float, b: float) -> Tensor:
    """
    Rounds the input tensor and clips its values to be within [a, b].

    Args:
        x (Tensor): Input tensor.
        a (float): Minimum clipping value.
        b (float): Maximum clipping value.

    Returns:
        Tensor: Rounded and clipped tensor.
    """
    return torch.clamp(torch.round(x), min=a, max=b)

def compute_gamma(w: Tensor, epsilon: float = 1e-8) -> Tensor:
    """
    Computes the scaling factor gamma for the weight matrix.

    Args:
        w (Tensor): Weight tensor.
        epsilon (float, optional): Small constant to prevent division by zero. Defaults to 1e-8.

    Returns:
        Tensor: Scaling factor gamma.
    """
    try:
        gamma = torch.mean(torch.abs(w)) + epsilon
        return gamma
    except Exception as e:
        logger.error(f"Error computing gamma: {e}")
        raise QuantizationError(f"Error computing gamma: {e}")

def bitnet_b158_quantize_weights(w: Tensor, epsilon: float = 1e-8) -> Tensor:
    """
    Quantizes the weights using the BitNet b1.58 quantization method.

    Args:
        w (Tensor): Original weight tensor.
        epsilon (float, optional): Small constant to prevent division by zero. Defaults to 1e-8.

    Returns:
        Tensor: Quantized weight tensor.
    """
    gamma = compute_gamma(w, epsilon)
    scaled_w = w / gamma
    quantized_w = round_clip(scaled_w, -1.0, 1.0)
    return quantized_w * gamma

# ----------------------------- Quantizer Classes -----------------------------

class WeightQuantizer:
    """Class to handle weight quantization."""

    def __init__(self, epsilon: float = 1e-8) -> None:
        """
        Initializes the WeightQuantizer.

        Args:
            epsilon (float, optional): Small constant to prevent division by zero. Defaults to 1e-8.
        """
        self.epsilon = epsilon

    def quantize(self, w: Tensor) -> Tensor:
        """
        Quantizes the input weight tensor.

        Args:
            w (Tensor): Weight tensor to quantize.

        Returns:
            Tensor: Quantized weight tensor.
        """
        return bitnet_b158_quantize_weights(w, self.epsilon)

class ActivationQuantizer:
    """Class to handle activation quantization to a specified number of bits."""

    def __init__(self, num_bits: int = 8) -> None:
        """
        Initializes the ActivationQuantizer.

        Args:
            num_bits (int, optional): Number of bits for quantization. Defaults to 8.

        Raises:
            ValueError: If num_bits is not within [1, 16].
        """
        if not 1 <= num_bits <= 16:
            raise ValueError("num_bits must be between 1 and 16.")
        self.num_bits = num_bits

    def quantize(self, x: Tensor) -> Tensor:
        """
        Quantizes the input activation tensor to the specified number of bits.

        Args:
            x (Tensor): Activation tensor to quantize.

        Returns:
            Tensor: Quantized activation tensor.
        """
        try:
            qmin = 0
            qmax = 2 ** self.num_bits - 1
            min_val = x.min()
            max_val = x.max()
            scale = (max_val - min_val) / (qmax - qmin)
            scale = max(scale, 1e-8)  # Prevent division by zero
            zero_point = qmin - torch.round(min_val / scale)
            zero_point = zero_point.clamp(qmin, qmax).int()
            x_quant = torch.clamp(torch.round(x / scale) + zero_point, qmin, qmax)
            x_dequant = (x_quant - zero_point) * scale
            return x_dequant
        except Exception as e:
            logger.error(f"Error during activation quantization: {e}")
            raise QuantizationError(f"Error during activation quantization: {e}")

# ----------------------------- Quantized Module ------------------------------

class QuantizedModule(nn.Module):
    """
    Generic wrapper for quantized modules. Applies weight and activation quantization 
    to any module that contains a 'weight' attribute.
    """

    def __init__(
        self,
        module: nn.Module,
        weight_quantizer: WeightQuantizer,
        activation_quantizer: ActivationQuantizer,
        module_type: Type[nn.Module]
    ) -> None:
        """
        Initializes the QuantizedModule.

        Args:
            module (nn.Module): Original PyTorch module to quantize.
            weight_quantizer (WeightQuantizer): Quantizer for the weights.
            activation_quantizer (ActivationQuantizer): Quantizer for the activations.
            module_type (Type[nn.Module]): The type of the module (e.g., nn.Linear, nn.Conv2d).
        """
        super(QuantizedModule, self).__init__()
        self.activation_quantizer = activation_quantizer

        # Initialize the quantized module with quantized weights
        self.module = self._quantize_module(module, module_type, weight_quantizer)

    def _quantize_module(
        self,
        module: nn.Module,
        module_type: Type[nn.Module],
        weight_quantizer: WeightQuantizer
    ) -> nn.Module:
        """
        Quantizes the weights of the given module and initializes a new module with quantized weights.

        Args:
            module (nn.Module): The module to quantize.
            module_type (Type[nn.Module]): The type of the module.
            weight_quantizer (WeightQuantizer): The weight quantizer.

        Returns:
            nn.Module: Module with quantized weights.
        """
        try:
            quantized_weight = weight_quantizer.quantize(module.weight.data.clone())
            # Create a new instance of the module with quantized weights
            quantized_module = module_type(
                *module_args(module, module_type),
                **module_kwargs(module, module_type)
            )
            quantized_module.weight = nn.Parameter(quantized_weight)
            if module.bias is not None:
                quantized_module.bias = nn.Parameter(module.bias.data.clone())
            return quantized_module
        except Exception as e:
            logger.error(f"Error quantizing module {module}: {e}")
            raise QuantizationError(f"Error quantizing module {module}: {e}")

    def forward(self, input: Tensor) -> Tensor:
        """
        Forward pass with quantized weights and activations.

        Args:
            input (Tensor): Input tensor.

        Returns:
            Tensor: Output tensor after module operation and activation quantization.
        """
        try:
            output = self.module(input)
            output = self.activation_quantizer.quantize(output)
            return output
        except Exception as e:
            logger.error(f"Error in QuantizedModule forward pass: {e}")
            raise QuantizationError(f"Error in QuantizedModule forward pass: {e}")

# ----------------------------- Helper Functions -------------------------------

def module_args(module: nn.Module, module_type: Type[nn.Module]) -> tuple:
    """
    Extracts the positional arguments required to instantiate a module of type `module_type` 
    using the attributes from the original `module`.

    Args:
        module (nn.Module): The original module.
        module_type (Type[nn.Module]): The type of the module to instantiate.

    Returns:
        tuple: Positional arguments for the module constructor.
    """
    if isinstance(module, nn.Linear):
        return (module.in_features, module.out_features)
    elif isinstance(module, nn.Conv1d):
        return (module.in_channels, module.out_channels, module.kernel_size)
    elif isinstance(module, nn.Conv2d):
        return (module.in_channels, module.out_channels, module.kernel_size,
                module.stride, module.padding, module.dilation, module.groups, module.bias is not None, module.padding_mode)
    elif isinstance(module, nn.Conv3d):
        return (module.in_channels, module.out_channels, module.kernel_size,
                module.stride, module.padding, module.dilation, module.groups, module.bias is not None, module.padding_mode)
    elif isinstance(module, nn.Embedding):
        return (module.num_embeddings, module.embedding_dim)
    else:
        # For other modules, you can add more conditions or raise an error
        return ()

def module_kwargs(module: nn.Module, module_type: Type[nn.Module]) -> dict:
    """
    Extracts the keyword arguments required to instantiate a module of type `module_type`
    using the attributes from the original `module`.

    Args:
        module (nn.Module): The original module.
        module_type (Type[nn.Module]): The type of the module to instantiate.

    Returns:
        dict: Keyword arguments for the module constructor.
    """
    kwargs = {}
    if isinstance(module, (nn.Conv1d, nn.Conv2d, nn.Conv3d)):
        kwargs = {
            'bias': module.bias is not None,
            'padding_mode': module.padding_mode if hasattr(module, 'padding_mode') else 'zeros'
        }
    return kwargs

# ----------------------------- Module Replacement -----------------------------

def replace_modules(
    model: nn.Module,
    quantized_module_factory: Callable[[nn.Module], nn.Module],
    exclude_modules: Optional[tuple] = None
) -> nn.Module:
    """
    Recursively replaces specified module types in the model with quantized modules.

    Args:
        model (nn.Module): Original PyTorch model.
        quantized_module_factory (Callable[[nn.Module], nn.Module]): Factory function to create quantized modules.
        exclude_modules (Optional[tuple]): Module types to exclude from quantization.

    Returns:
        nn.Module: Model with specified modules replaced by quantized modules.
    """
    if exclude_modules is None:
        exclude_modules = ()

    for name, module in model.named_children():
        if isinstance(module, exclude_modules):
            continue
        if hasattr(module, 'weight') and isinstance(module, nn.Module):
            quantized_module = quantized_module_factory(module)
            setattr(model, name, quantized_module)
        else:
            replace_modules(module, quantized_module_factory, exclude_modules)
    return model

# ----------------------------- BitNet Pipeline --------------------------------

class BitNetPipeline:
    """Class to handle the full BitNet b1.58 quantization pipeline."""

    def __init__(self, model: nn.Module) -> None:
        """
        Initializes the BitNetPipeline.

        Args:
            model (nn.Module): The PyTorch model to quantize.
        """
        if not isinstance(model, nn.Module):
            raise TypeError("Input model must be an instance of torch.nn.Module.")
        self.original_model = model
        self.quantized_model = copy.deepcopy(model)
        self.weight_quantizer = WeightQuantizer()
        self.activation_quantizer = ActivationQuantizer()

    def quantize_module(self, module: nn.Module) -> nn.Module:
        """
        Factory method to create a quantized version of the given module.

        Args:
            module (nn.Module): The module to quantize.

        Returns:
            nn.Module: Quantized module.
        """
        return QuantizedModule(
            module=module,
            weight_quantizer=self.weight_quantizer,
            activation_quantizer=self.activation_quantizer,
            module_type=type(module)
        )

    def quantize(self, exclude_modules: Optional[tuple] = (
        nn.ReLU, nn.Sigmoid, nn.Tanh, nn.LeakyReLU, nn.ELU, nn.Softmax,
        nn.BatchNorm1d, nn.BatchNorm2d, nn.BatchNorm3d, nn.LayerNorm
    )) -> nn.Module:
        """
        Applies the quantization pipeline to the model.

        Args:
            exclude_modules (Optional[tuple], optional): Module types to exclude from quantization.
                Defaults to common activation and normalization layers.

        Returns:
            nn.Module: Quantized PyTorch model.
        """
        try:
            self.quantized_model = replace_modules(
                self.quantized_model,
                self.quantize_module,
                exclude_modules=exclude_modules
            )
            logger.info("Model quantization completed successfully.")
            return self.quantized_model
        except Exception as e:
            logger.error(f"Error during model quantization: {e}")
            raise QuantizationError(f"Error during model quantization: {e}")

# ----------------------------- Model Summary ----------------------------------

def model_summary(model: nn.Module, console: Optional[Console] = None) -> None:
    """
    Generate a detailed summary of a PyTorch model, including the number of total,
    trainable, and non-trainable parameters. Also, visually display the model architecture
    using a tree structure.

    Args:
        model (nn.Module): The PyTorch model to summarize.
        console (Optional[Console], optional): Rich console for output. Defaults to None.

    Raises:
        TypeError: If the input is not a PyTorch nn.Module.
    """
    if not isinstance(model, nn.Module):
        raise TypeError("Input must be a PyTorch nn.Module")
    
    if console is None:
        console = Console()

    # Compute parameters' stats
    total_params = sum(p.numel() for p in model.parameters())
    trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
    non_trainable_params = total_params - trainable_params

    # Create a caption for the model
    model_caption = f"[bold cyan]{model.__class__.__name__}[/bold cyan]"
    console.print(Panel(model_caption, expand=False))

    # Create a summary table for the total parameters
    total_params_table = Table(
        title="Total Parameters Summary",
        show_header=True,
        header_style="bold magenta"
    )
    total_params_table.add_column("Category", style="cyan")
    total_params_table.add_column("Count", style="green", justify="right")

    total_params_table.add_row("Total Parameters", f"{total_params:,}")
    total_params_table.add_row("Trainable Parameters", f"{trainable_params:,}")
    total_params_table.add_row("Non-trainable Parameters", f"{non_trainable_params:,}")

    console.print(total_params_table)
    console.print(Panel("[italic]Summary of model's total parameters[/italic]", expand=False))

    # Create a detailed table for the parameters
    detailed_table = Table(
        title="Detailed Parameter Summary",
        show_lines=True
    )
    detailed_table.add_column("Parameter", style="cyan", no_wrap=True)
    detailed_table.add_column("Shape", style="magenta")
    detailed_table.add_column("Trainable", style="green")
    detailed_table.add_column("Num Parameters", style="yellow", justify="right")

    def add_module_to_table(module: nn.Module, prefix: Optional[str] = "") -> None:
        """
        Recursively add model parameters to the detailed table.

        Args:
            module (nn.Module): A PyTorch module.
            prefix (Optional[str]): Prefix for parameter names, used for recursion.
        """
        for name, param in module.named_parameters(recurse=False):
            full_name = f"{prefix}.{name}" if prefix else name
            shape = tuple(param.shape)
            trainable = param.requires_grad
            num_params = param.numel()

            detailed_table.add_row(
                full_name,
                str(shape),
                "[green]✓[/green]" if trainable else "[red]✗[/red]",
                f"{num_params:,}"
            )

        for name, child in module.named_children():
            child_prefix = f"{prefix}.{name}" if prefix else name
            add_module_to_table(child, child_prefix)

    add_module_to_table(model)
    console.print(detailed_table)
    console.print(Panel("[italic]Detailed breakdown of model's parameters[/italic]", expand=False))

    # Create a rich Tree for the model architecture
    def build_tree(module: nn.Module, prefix: Optional[str] = "") -> Tree:
        """
        Recursively build a tree structure for the model architecture.

        Args:
            module (nn.Module): The module to represent in the tree.
            prefix (Optional[str]): The prefix for the current module.

        Returns:
            Tree: A rich Tree representing the module architecture.
        """
        root_name = prefix if prefix else module.__class__.__name__
        tree = Tree(f"[bold]{root_name}[/bold]")

        for name, child in module.named_children():
            child_tree = build_tree(child, prefix=name)
            tree.add(child_tree)

        for name, param in module.named_parameters(recurse=False):
            param_info = (
                f"{name}: {tuple(param.shape)} "
                f"{'(Trainable)' if param.requires_grad else '(Non-trainable)'}"
            )
            tree.add(f"[cyan]{param_info}[/cyan]")

        return tree

    architecture_tree = build_tree(model)
    console.print("\n[bold yellow]Model Architecture[/bold yellow]")
    console.print(architecture_tree)
    console.print(Panel("[italic]Visual representation of the model's architecture[/italic]", expand=False))

def model_summary_update(model: nn.Module, save_path: Optional[str] = "model_summary.html") -> None:
    """
    Generate a detailed summary of a PyTorch model, including the number of total,
    trainable, and non-trainable parameters. Also, visually display the model architecture
    using a tree structure. Optionally saves the summary as an HTML file.

    Args:
        model (nn.Module): The PyTorch model to summarize.
        save_path (Optional[str], optional): Path to save the summary as an HTML file. 
            Defaults to "model_summary.html".

    Raises:
        TypeError: If the input is not a PyTorch nn.Module.
        QuantizationError: If saving the summary fails.
    """
    if not isinstance(model, nn.Module):
        raise TypeError("Input must be a PyTorch nn.Module")

    console = Console(record=True)  # Enable recording of console output

    model_summary(model, console)

    if save_path:
        try:
            html = console.export_html(clear=False)
            with open(save_path, "w", encoding="utf-8") as f:
                f.write(html)
            console.print(f"Summary saved as [green]{save_path}[/green]")
        except Exception as e:
            logger.error(f"Failed to save summary to {save_path}: {e}")
            raise QuantizationError(f"Failed to save summary to {save_path}: {e}")

# ----------------------------- Quantization Application ----------------------

def apply_bitnet_b158_quantization(model: nn.Module) -> nn.Module:
    """
    Applies the BitNet b1.58 quantization pipeline to any PyTorch model.

    Args:
        model (nn.Module): The PyTorch model to quantize.

    Returns:
        nn.Module: Quantized PyTorch model.

    Raises:
        TypeError: If the input is not a PyTorch nn.Module.
        QuantizationError: If quantization fails.
    """
    if not isinstance(model, nn.Module):
        raise TypeError("Input model must be an instance of torch.nn.Module.")

    pipeline = BitNetPipeline(model)
    quantized_model = pipeline.quantize()
    return quantized_model

# ----------------------------- Example Usage -----------------------------------

if __name__ == "__main__":
    try:
        # Define a sample model
        sample_model = nn.Sequential(
            nn.Linear(128, 256),
            nn.ReLU(),
            nn.Conv2d(16, 32, kernel_size=3, stride=1, padding=1),
            nn.BatchNorm2d(32),
            nn.ReLU(),
            nn.Linear(256, 10)
        )

        print("\n[bold green]Original Model Summary[/bold green]")
        model_summary(sample_model)

        # Apply BitNet b1.58 quantization
        quantized_sample_model = apply_bitnet_b158_quantization(sample_model)

        print("\n[bold green]Quantized Model Summary[/bold green]")
        model_summary(quantized_sample_model)

    except QuantizationError as qe:
        logger.error(f"Quantization failed: {qe}")
    except Exception as ex:
        logger.error(f"An unexpected error occurred: {ex}")