Convert stable diffusion model as inference only

ai-edge-bot · copybara-github · commit 25db678984e7 · 2025-01-24T12:23:10.000-08:00
PiperOrigin-RevId: 719393588
diff --git a/ai_edge_torch/generative/examples/stable_diffusion/convert_to_tflite.py b/ai_edge_torch/generative/examples/stable_diffusion/convert_to_tflite.py
@@ -85,6 +85,7 @@ def convert_stable_diffusion_to_tflite(
       clip.TENSOR_NAMES,
   )
   loader.load(clip_model, strict=False)
+  clip_model.eval()
 
   diffusion_model = diffusion.Diffusion(
       diffusion.get_model_config(batch_size=2, device_type=_DEVICE_TYPE.value)
@@ -93,6 +94,7 @@ def convert_stable_diffusion_to_tflite(
       diffusion_ckpt_path, diffusion.TENSOR_NAMES
   )
   diffusion_loader.load(diffusion_model, strict=False)
+  diffusion_model.eval()
 
   decoder_model = decoder.Decoder(
       decoder.get_model_config(device_type=_DEVICE_TYPE.value)
@@ -101,6 +103,7 @@ def convert_stable_diffusion_to_tflite(
       decoder_ckpt_path, decoder.TENSOR_NAMES
   )
   decoder_loader.load(decoder_model, strict=False)
+  decoder_model.eval()
 
   # TODO(yichunk): enable image encoder conversion
   # if encoder_ckpt_path is not None:

Original file line number	Diff line number	Diff line change
`@@ -85,6 +85,7 @@ def convert_stable_diffusion_to_tflite(`
`85`	`85`	`clip.TENSOR_NAMES,`
`86`	`86`	`)`
`87`	`87`	`loader.load(clip_model, strict=False)`
	`88`	`+ clip_model.eval()`
`88`	`89`
`89`	`90`	`diffusion_model = diffusion.Diffusion(`
`90`	`91`	`diffusion.get_model_config(batch_size=2, device_type=_DEVICE_TYPE.value)`
`@@ -93,6 +94,7 @@ def convert_stable_diffusion_to_tflite(`
`93`	`94`	`diffusion_ckpt_path, diffusion.TENSOR_NAMES`
`94`	`95`	`)`
`95`	`96`	`diffusion_loader.load(diffusion_model, strict=False)`
	`97`	`+ diffusion_model.eval()`
`96`	`98`
`97`	`99`	`decoder_model = decoder.Decoder(`
`98`	`100`	`decoder.get_model_config(device_type=_DEVICE_TYPE.value)`
`@@ -101,6 +103,7 @@ def convert_stable_diffusion_to_tflite(`
`101`	`103`	`decoder_ckpt_path, decoder.TENSOR_NAMES`
`102`	`104`	`)`
`103`	`105`	`decoder_loader.load(decoder_model, strict=False)`
	`106`	`+ decoder_model.eval()`
`104`	`107`
`105`	`108`	`# TODO(yichunk): enable image encoder conversion`
`106`	`109`	`# if encoder_ckpt_path is not None:`