Overrides _post_quantize to reset generate_function graph after quantization (#2436)

* Overrides _post_quantize to reset generate_function graph

* added _post_quantize override to image_to_image

Author: JyotinderSingh

keras_hub/src/models/causal_lm.py

@@ -424,3 +424,8 @@ def export_to_transformers(self, path):
         )
 
         export_to_safetensors(self, path)
+
+    def _post_quantize(self, mode, **kwargs):
+        super()._post_quantize(mode, **kwargs)
+        # Reset the compiled generate function.
+        self.generate_function = None

keras_hub/src/models/image_to_image.py

@@ -415,3 +415,8 @@ def generate(images, x):
         # Image-to-image.
         outputs = [generate(*x) for x in inputs]
         return self._normalize_generate_outputs(outputs, input_is_scalar)
+
+    def _post_quantize(self, mode, **kwargs):
+        super()._post_quantize(mode, **kwargs)
+        # Reset the compiled generate function.
+        self.generate_function = None

keras_hub/src/models/inpaint.py

@@ -518,3 +518,8 @@ def generate(images, masks, x):
         # Inpaint.
         outputs = [generate(*x) for x in inputs]
         return self._normalize_generate_outputs(outputs, input_is_scalar)
+
+    def _post_quantize(self, mode, **kwargs):
+        super()._post_quantize(mode, **kwargs)
+        # Reset the compiled generate function.
+        self.generate_function = None

keras_hub/src/models/text_to_image.py

@@ -345,3 +345,8 @@ def generate(x):
         # Text-to-image.
         outputs = [generate(x) for x in inputs]
         return self._normalize_generate_outputs(outputs, input_is_scalar)
+
+    def _post_quantize(self, mode, **kwargs):
+        super()._post_quantize(mode, **kwargs)
+        # Reset the compiled generate function.
+        self.generate_function = None