diff --git a/packages/ui/components/avatar/avatar.stories.mdx b/packages/ui/components/avatar/avatar.stories.mdx
index 5e7d5fbe31..671a780a62 100644
--- a/packages/ui/components/avatar/avatar.stories.mdx
+++ b/packages/ui/components/avatar/avatar.stories.mdx
@@ -125,14 +125,29 @@ Avatar group can be composed differently based on the number of user profile.
 
 <Canvas>
   <Story
-    name="Avatar">
-    <VariantsTable titles={['Default']} columnMinWidth={150}>
-      <VariantRow variant="Sm">
-        <Avatar size="sm" alt="Avatar Story" gravatarFallbackMd5="Ui@CAL.com" />
+    name="Avatar"
+    args={{
+      size: 'sm',
+      alt: 'Avatar Story',
+      gravatarFallbackMd5: 'Ui@CAL.com',
+    }}
+    argTypes={{
+      size: {
+        control: {
+          type: 'inline-radio',
+          options: ['sm', 'lg'],
+        },
+      },
+      alt: { control: 'text' },
+      gravatarFallbackMd5: { control: 'text' },
+    }}
+  >
+    {({ size, alt, gravatarFallbackMd5 }) => (
+  <VariantsTable titles={['Default']} columnMinWidth={150}>
+      <VariantRow variant={size}>
+        <Avatar size={size} alt={alt} gravatarFallbackMd5={gravatarFallbackMd5} />
       </VariantRow>
-      <VariantRow variant="Lg">
-        <Avatar size="lg" alt="Avatar Story" gravatarFallbackMd5="Ui@CAL.com" />
-      </VariantRow>
-    </VariantsTable>
+  </VariantsTable>
+    )}
   </Story>
 </Canvas>
diff --git a/packages/ui/components/badge/badge.stories.mdx b/packages/ui/components/badge/badge.stories.mdx
index 817d6dcfac..d3233ca2d4 100644
--- a/packages/ui/components/badge/badge.stories.mdx
+++ b/packages/ui/components/badge/badge.stories.mdx
@@ -56,28 +56,32 @@ Alet badge is used in conjunction with an item, profile or label to indicate num
 <Title offset title="Badge" suffix="Variants" />
 
 <Canvas>
-  <Story name="All Variants">
-    <VariantsTable titles={['Default','With Dot', 'With Icon']} columnMinWidth={150}>
-      <VariantRow variant="Default">
-        <Badge variant="default" >Badge text</Badge>
-        <Badge variant="default" withDot>Badge text</Badge>
-        <Badge variant="default" startIcon={Plus}>Badge text</Badge>
+  <Story name="All Variants"
+  args={{
+    severity: 'default',
+    label: 'Badge text',
+  }}
+  argTypes={{
+    severity: {
+      control: {
+        type: 'inline-radio',
+        options: ['default', 'success', 'gray', 'error'],
+      },
+    },
+    label: {
+      control: {
+        type: 'text',
+      },
+    },
+  }}>
+    {({ severity, label }) => (
+  <VariantsTable titles={['Default','With Dot', 'With Icon']} columnMinWidth={150}>
+      <VariantRow variant={severity}>
+        <Badge variant={severity}>{label}</Badge>
+        <Badge variant={severity} withDot>{label}</Badge>
+        <Badge variant={severity} startIcon={Plus}>{label}</Badge>
       </VariantRow>
-      <VariantRow variant="Success">
-        <Badge variant="success" >Badge text</Badge>
-        <Badge variant="success" withDot>Badge text</Badge>
-        <Badge variant="success" startIcon={Plus}>Badge text</Badge>
-      </VariantRow>
-      <VariantRow variant="gray">
-        <Badge variant="gray" >Badge text</Badge>
-        <Badge variant="gray" withDot>Badge text</Badge>
-        <Badge variant="gray" startIcon={Plus}>Badge text</Badge>
-      </VariantRow>
-      <VariantRow variant="error">
-        <Badge variant="error" >Badge text</Badge>
-        <Badge variant="error" withDot>Badge text</Badge>
-        <Badge variant="error" startIcon={Plus}>Badge text</Badge>
-      </VariantRow>
-    </VariantsTable>
+  </VariantsTable>
+    )}
   </Story>
 </Canvas>