diff --git a/android/image.go b/android/image.go
index bc6b8cd..9cad056 100644
--- a/android/image.go
+++ b/android/image.go
@@ -44,9 +44,8 @@
 	ExtraImageVariations(ctx BaseModuleContext) []string
 
 	// SetImageVariation is called for each newly created image variant. The receiver is the original
-	// module, "variation" is the name of the newly created variant and "module" is the newly created
-	// variant itself.
-	SetImageVariation(ctx BaseModuleContext, variation string, module Module)
+	// module, "variation" is the name of the newly created variant. "variation" is set on the receiver.
+	SetImageVariation(ctx BaseModuleContext, variation string)
 }
 
 const (
@@ -106,7 +105,7 @@
 		mod := ctx.CreateVariations(variations...)
 		for i, v := range variations {
 			mod[i].base().setImageVariation(v)
-			m.SetImageVariation(ctx, v, mod[i])
+			mod[i].(ImageInterface).SetImageVariation(ctx, v)
 		}
 	}
 }
diff --git a/bpf/bpf.go b/bpf/bpf.go
index 2eb869e..ce00b5b 100644
--- a/bpf/bpf.go
+++ b/bpf/bpf.go
@@ -131,7 +131,7 @@
 	return nil
 }
 
-func (bpf *bpf) SetImageVariation(ctx android.BaseModuleContext, variation string, module android.Module) {
+func (bpf *bpf) SetImageVariation(ctx android.BaseModuleContext, variation string) {
 	bpf.properties.VendorInternal = variation == "vendor"
 }
 
diff --git a/cc/genrule.go b/cc/genrule.go
index 431a01c..cd06818 100644
--- a/cc/genrule.go
+++ b/cc/genrule.go
@@ -114,5 +114,5 @@
 	return variants
 }
 
-func (g *GenruleExtraProperties) SetImageVariation(ctx android.BaseModuleContext, variation string, module android.Module) {
+func (g *GenruleExtraProperties) SetImageVariation(ctx android.BaseModuleContext, variation string) {
 }
diff --git a/cc/image.go b/cc/image.go
index f8c5ca5..28692d8 100644
--- a/cc/image.go
+++ b/cc/image.go
@@ -628,30 +628,29 @@
 	}
 }
 
-func (c *Module) SetImageVariation(ctx android.BaseModuleContext, variant string, module android.Module) {
-	m := module.(*Module)
+func (c *Module) SetImageVariation(ctx android.BaseModuleContext, variant string) {
 	if variant == android.RamdiskVariation {
-		m.MakeAsPlatform()
-		squashRamdiskSrcs(m)
+		c.MakeAsPlatform()
+		squashRamdiskSrcs(c)
 	} else if variant == android.VendorRamdiskVariation {
-		m.MakeAsPlatform()
-		squashVendorRamdiskSrcs(m)
+		c.MakeAsPlatform()
+		squashVendorRamdiskSrcs(c)
 	} else if variant == android.RecoveryVariation {
-		m.MakeAsPlatform()
-		squashRecoverySrcs(m)
+		c.MakeAsPlatform()
+		squashRecoverySrcs(c)
 	} else if strings.HasPrefix(variant, VendorVariation) {
-		m.Properties.ImageVariation = VendorVariation
+		c.Properties.ImageVariation = VendorVariation
 
 		if strings.HasPrefix(variant, VendorVariationPrefix) {
-			m.Properties.VndkVersion = strings.TrimPrefix(variant, VendorVariationPrefix)
+			c.Properties.VndkVersion = strings.TrimPrefix(variant, VendorVariationPrefix)
 		}
-		squashVendorSrcs(m)
+		squashVendorSrcs(c)
 	} else if strings.HasPrefix(variant, ProductVariation) {
-		m.Properties.ImageVariation = ProductVariation
+		c.Properties.ImageVariation = ProductVariation
 		if strings.HasPrefix(variant, ProductVariationPrefix) {
-			m.Properties.VndkVersion = strings.TrimPrefix(variant, ProductVariationPrefix)
+			c.Properties.VndkVersion = strings.TrimPrefix(variant, ProductVariationPrefix)
 		}
-		squashProductSrcs(m)
+		squashProductSrcs(c)
 	}
 
 	if c.NeedsVendorPublicLibraryVariants() &&
diff --git a/cc/library_stub.go b/cc/library_stub.go
index 746b951..9643ec2 100644
--- a/cc/library_stub.go
+++ b/cc/library_stub.go
@@ -511,5 +511,5 @@
 
 	return variations
 }
-func (v *CcApiVariant) SetImageVariation(ctx android.BaseModuleContext, variation string, module android.Module) {
+func (v *CcApiVariant) SetImageVariation(ctx android.BaseModuleContext, variation string) {
 }
diff --git a/etc/prebuilt_etc.go b/etc/prebuilt_etc.go
index c1a0b9c..8c6d162 100644
--- a/etc/prebuilt_etc.go
+++ b/etc/prebuilt_etc.go
@@ -241,7 +241,7 @@
 	return nil
 }
 
-func (p *PrebuiltEtc) SetImageVariation(ctx android.BaseModuleContext, variation string, module android.Module) {
+func (p *PrebuiltEtc) SetImageVariation(ctx android.BaseModuleContext, variation string) {
 }
 
 func (p *PrebuiltEtc) SourceFilePath(ctx android.ModuleContext) android.Path {
diff --git a/genrule/genrule.go b/genrule/genrule.go
index 26dad01..06a7e18 100644
--- a/genrule/genrule.go
+++ b/genrule/genrule.go
@@ -650,7 +650,7 @@
 func (x noopImageInterface) DebugRamdiskVariantNeeded(android.BaseModuleContext) bool    { return false }
 func (x noopImageInterface) RecoveryVariantNeeded(android.BaseModuleContext) bool        { return false }
 func (x noopImageInterface) ExtraImageVariations(ctx android.BaseModuleContext) []string { return nil }
-func (x noopImageInterface) SetImageVariation(ctx android.BaseModuleContext, variation string, module android.Module) {
+func (x noopImageInterface) SetImageVariation(ctx android.BaseModuleContext, variation string) {
 }
 
 func NewGenSrcs() *Module {
diff --git a/rust/image.go b/rust/image.go
index e0d267d..fec6d92 100644
--- a/rust/image.go
+++ b/rust/image.go
@@ -197,21 +197,20 @@
 	return mod.InVendor() || mod.InProduct()
 }
 
-func (mod *Module) SetImageVariation(ctx android.BaseModuleContext, variant string, module android.Module) {
-	m := module.(*Module)
+func (mod *Module) SetImageVariation(ctx android.BaseModuleContext, variant string) {
 	if variant == android.VendorRamdiskVariation {
-		m.MakeAsPlatform()
+		mod.MakeAsPlatform()
 	} else if variant == android.RecoveryVariation {
-		m.MakeAsPlatform()
+		mod.MakeAsPlatform()
 	} else if strings.HasPrefix(variant, cc.VendorVariation) {
-		m.Properties.ImageVariation = cc.VendorVariation
+		mod.Properties.ImageVariation = cc.VendorVariation
 		if strings.HasPrefix(variant, cc.VendorVariationPrefix) {
-			m.Properties.VndkVersion = strings.TrimPrefix(variant, cc.VendorVariationPrefix)
+			mod.Properties.VndkVersion = strings.TrimPrefix(variant, cc.VendorVariationPrefix)
 		}
 	} else if strings.HasPrefix(variant, cc.ProductVariation) {
-		m.Properties.ImageVariation = cc.ProductVariation
+		mod.Properties.ImageVariation = cc.ProductVariation
 		if strings.HasPrefix(variant, cc.ProductVariationPrefix) {
-			m.Properties.VndkVersion = strings.TrimPrefix(variant, cc.ProductVariationPrefix)
+			mod.Properties.VndkVersion = strings.TrimPrefix(variant, cc.ProductVariationPrefix)
 		}
 	}
 }
diff --git a/sh/sh_binary.go b/sh/sh_binary.go
index 3f82560..48a442d 100644
--- a/sh/sh_binary.go
+++ b/sh/sh_binary.go
@@ -243,10 +243,8 @@
 	return extraVariations
 }
 
-func (s *ShBinary) SetImageVariation(ctx android.BaseModuleContext, variation string, module android.Module) {
-	if m, ok := module.(*ShBinary); ok {
-		m.properties.ImageVariation = variation
-	}
+func (s *ShBinary) SetImageVariation(ctx android.BaseModuleContext, variation string) {
+	s.properties.ImageVariation = variation
 }
 
 // Overrides ModuleBase.InstallInRamdisk() so that the install rule respects