Loading model with 'from_pretrained_keras'

app.py

@@ -5,20 +5,20 @@ import gradio as gr
 import keras_cv
 import numpy as np
 import tensorflow as tf
+from huggingface_hub import from_pretrained_keras
 
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__file__)
-prompt_token = os.environ.get("TOKEN", "<token>")
-text_encoder_path = os.environ.get(
-    "TEXT_ENCODER", "./models/example/text_encoder/keras"
-)
+prompt_token = "<token>"
+text_encoder_url = "Dimitre/stablediffusion-canarinho_pistola"
 
 logger.info(f'Inversed token used: "{prompt_token}"')
-logger.info(f'Loading text encoder from: "{text_encoder_path}"')
+logger.info(f'Loading text encoder from: "{text_encoder_url}"')
 
 stable_diffusion = keras_cv.models.StableDiffusion()
 stable_diffusion.tokenizer.add_tokens(prompt_token)
-loaded_text_encoder_ = tf.keras.models.load_model(text_encoder_path)
+# loaded_text_encoder_ = tf.keras.models.load_model(text_encoder_path)
+loaded_text_encoder_ = from_pretrained_keras(text_encoder_url)
 stable_diffusion._text_encoder = loaded_text_encoder_
 stable_diffusion._text_encoder.compile(jit_compile=True)