app.py

import logging
import os

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 = "<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_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_ = from_pretrained_keras(text_encoder_url)
stable_diffusion._text_encoder = loaded_text_encoder_
stable_diffusion._text_encoder.compile(jit_compile=True)


def generate_fn(input_prompt: str) -> np.ndarray:
    """Generates images from a text prompt

    Args:
        input_prompt (str): Text input prompt

    Returns:
        np.ndarray: Generated image
    """
    generated = stable_diffusion.text_to_image(
        prompt=input_prompt, batch_size=1, num_steps=1
    )
    return generated[0]


iface = gr.Interface(
    fn=generate_fn,
    title="Textual Inversion",
    description="Textual Inversion Demo",
    article="Note: Keras-cv uses lazy intialization, so the first use will be slower while the model is initialized.",
    inputs=gr.Textbox(
        label="Prompt",
        show_label=False,
        max_lines=2,
        placeholder="Enter your prompt",
        elem_id="input-prompt",
    ),
    outputs=gr.Image(),
)

if __name__ == "__main__":
    app, local_url, share_url = iface.launch()