diff --git a/drivers/leds/leds-qti-flash.c b/drivers/leds/leds-qti-flash.c
index 319b09ab71c5..790a2dc81f2b 100644
--- a/drivers/leds/leds-qti-flash.c
+++ b/drivers/leds/leds-qti-flash.c
@@ -162,7 +162,8 @@ struct flash_switch_data {
  * @revision		: Revision of the flash LED module
  * @subtype		: Peripheral subtype of the flash LED module
  * @max_channels	: Maximum number of channels supported by flash module
- * @ref_count		: Reference count used to enable/disable flash LED
+ * @chan_en_map		: Bit map of individual channel enable
+ * @module_en		: Flag used to enable/disable flash LED module
  * @trigger_lmh		: Flag to enable lmh mitigation
  */
 struct qti_flash_led {
@@ -184,7 +185,8 @@ struct qti_flash_led {
 	u8		revision;
 	u8		subtype;
 	u8		max_channels;
-	u8		ref_count;
+	u8		chan_en_map;
+	bool		module_en;
 	bool		trigger_lmh;
 };
 
@@ -289,25 +291,24 @@ static int qti_flash_led_module_control(struct qti_flash_led *led,
 	u8 val;
 
 	if (enable) {
-		if (!led->ref_count) {
+		if (!led->module_en && led->chan_en_map) {
 			val = FLASH_MODULE_ENABLE;
 			rc = qti_flash_led_write(led, FLASH_ENABLE_CONTROL,
 						&val, 1);
 			if (rc < 0)
 				return rc;
+
+			led->module_en = true;
 		}
-
-		led->ref_count++;
 	} else {
-		if (led->ref_count)
-			led->ref_count--;
-
-		if (!led->ref_count) {
+		if (led->module_en && !led->chan_en_map) {
 			val = FLASH_MODULE_DISABLE;
 			rc = qti_flash_led_write(led, FLASH_ENABLE_CONTROL,
 						&val, 1);
 			if (rc < 0)
 				return rc;
+
+			led->module_en = false;
 		}
 	}
 
@@ -318,12 +319,16 @@ static int qti_flash_led_strobe(struct qti_flash_led *led,
 				struct flash_switch_data *snode,
 				u8 mask, u8 value)
 {
-	int rc;
+	int rc, i;
 	bool enable = mask & value;
 
 	spin_lock(&led->lock);
 
 	if (enable) {
+		for (i = 0; i < led->max_channels; i++)
+			if ((mask & BIT(i)) && (value & BIT(i)))
+				led->chan_en_map |= BIT(i);
+
 		rc = qti_flash_led_module_control(led, enable);
 		if (rc < 0)
 			goto error;
@@ -341,6 +346,11 @@ static int qti_flash_led_strobe(struct qti_flash_led *led,
 		if (rc < 0)
 			goto error;
 	} else {
+		for (i = 0; i < led->max_channels; i++)
+			if ((led->chan_en_map & BIT(i)) &&
+			    (mask & BIT(i)) && !(value & BIT(i)))
+				led->chan_en_map &= ~(BIT(i));
+
 		rc = qti_flash_led_masked_write(led, FLASH_EN_LED_CTRL,
 				mask, value);
 		if (rc < 0)
@@ -411,6 +421,9 @@ static int qti_flash_led_disable(struct flash_node_data *fnode)
 	struct qti_flash_led *led = fnode->led;
 	int rc;
 
+	if (!fnode->configured)
+		return -EINVAL;
+
 	spin_lock(&led->lock);
 	if ((fnode->strobe_sel == HW_STROBE) &&
 		gpio_is_valid(led->hw_strobe_gpio[fnode->id]))
@@ -427,6 +440,7 @@ static int qti_flash_led_disable(struct flash_node_data *fnode)
 	if (rc < 0)
 		goto out;
 
+	fnode->configured = false;
 	fnode->current_ma = 0;
 
 out:
@@ -454,11 +468,18 @@ static void qti_flash_led_brightness_set(struct led_classdev *led_cdev,
 	fnode = container_of(fdev, struct flash_node_data, fdev);
 	led = fnode->led;
 
-	if (brightness <= 0) {
+	if (!brightness) {
+		rc = qti_flash_led_strobe(fnode->led, NULL,
+			FLASH_LED_ENABLE(fnode->id), 0);
+		if (rc < 0) {
+			pr_err("Failed to destrobe LED, rc=%d\n", rc);
+			return;
+		}
+
 		rc = qti_flash_led_disable(fnode);
 		if (rc < 0)
-			pr_err("Failed to set brightness %d to LED\n",
-				brightness);
+			pr_err("Failed to disable LED\n");
+
 		return;
 	}
 
@@ -578,6 +599,21 @@ static int qti_flash_switch_disable(struct flash_switch_data *snode)
 	int rc = 0, i;
 	u8 led_dis = 0;
 
+	for (i = 0; i < led->num_fnodes; i++) {
+		if (!(snode->led_mask & BIT(led->fnode[i].id)) ||
+				!led->fnode[i].configured)
+			continue;
+
+		led_dis |= BIT(led->fnode[i].id);
+	}
+
+	rc = qti_flash_led_strobe(led, NULL, led_dis, ~led_dis);
+	if (rc < 0) {
+		pr_err("Failed to destrobe LEDs under with switch, rc=%d\n",
+			rc);
+		return rc;
+	}
+
 	for (i = 0; i < led->num_fnodes; i++) {
 		/*
 		 * Do not turn OFF flash/torch device if
@@ -585,7 +621,7 @@ static int qti_flash_switch_disable(struct flash_switch_data *snode)
 		 * ii. brightness is not configured for device under this switch
 		 */
 		if (!(snode->led_mask & BIT(led->fnode[i].id)) ||
-			!led->fnode[i].configured)
+				!led->fnode[i].configured)
 			continue;
 
 		rc = qti_flash_led_disable(&led->fnode[i]);
@@ -594,15 +630,12 @@ static int qti_flash_switch_disable(struct flash_switch_data *snode)
 				&led->fnode[i].id);
 			break;
 		}
-
-		led_dis |= (1 << led->fnode[i].id);
-		led->fnode[i].configured = false;
 	}
 
 	snode->on_time_ms = 0;
 	snode->off_time_ms = 0;
 
-	return qti_flash_led_strobe(led, NULL, led_dis, ~led_dis);
+	return rc;
 }
 
 static void qti_flash_led_switch_brightness_set(
@@ -1087,20 +1120,25 @@ static int qti_flash_strobe_set(struct led_classdev_flash *fdev,
 	if (fnode->enabled == state)
 		return 0;
 
-	if (!state) {
-		rc = qti_flash_led_disable(fnode);
-		if (rc < 0) {
-			pr_err("Failed to disable LED %u\n", fnode->id);
-			return rc;
-		}
-	}
+	if (state && !fnode->configured)
+		return -EINVAL;
 
 	mask = FLASH_LED_ENABLE(fnode->id);
 	value = state ? FLASH_LED_ENABLE(fnode->id) : 0;
 
 	rc = qti_flash_led_strobe(fnode->led, NULL, mask, value);
-	if (!rc)
-		fnode->enabled = state;
+	if (rc < 0) {
+		pr_err("Failed to %s LED, rc=%d\n",
+			state ? "strobe" : "desrobe", rc);
+		return rc;
+	}
+	fnode->enabled = state;
+
+	if (!state) {
+		rc = qti_flash_led_disable(fnode);
+		if (rc < 0)
+			pr_err("Failed to disable LED %u\n", fnode->id);
+	}
 
 	return rc;
 }