Spaces:
Sleeping
Sleeping
| from flax.jax_utils import replicate | |
| from jax import pmap | |
| from flax.training.common_utils import shard | |
| import jax | |
| import jax.numpy as jnp | |
| import gradio as gr | |
| from pathlib import Path | |
| from PIL import Image | |
| import numpy as np | |
| from diffusers import FlaxStableDiffusionPipeline | |
| import os | |
| if 'TPU_NAME' in os.environ: | |
| import requests | |
| if 'TPU_DRIVER_MODE' not in globals(): | |
| url = 'http:' + os.environ['TPU_NAME'].split(':')[1] + ':8475/requestversion/tpu_driver_nightly' | |
| resp = requests.post(url) | |
| TPU_DRIVER_MODE = 1 | |
| from jax.config import config | |
| config.FLAGS.jax_xla_backend = "tpu_driver" | |
| config.FLAGS.jax_backend_target = os.environ['TPU_NAME'] | |
| print('Registered TPU:', config.FLAGS.jax_backend_target) | |
| else: | |
| print('No TPU detected. Can be changed under "Runtime/Change runtime type".') | |
| import jax | |
| jax.local_devices() | |
| num_devices = jax.device_count() | |
| device_type = jax.devices()[0].device_kind | |
| print(f"Found {num_devices} JAX devices of type {device_type}.") | |
| def sd2_inference(pipeline, prompts, params, seed = 42, num_inference_steps = 50 ): | |
| prng_seed = jax.random.PRNGKey(seed) | |
| prompt_ids = pipeline.prepare_inputs(prompts) | |
| params = replicate(params) | |
| prng_seed = jax.random.split(prng_seed, jax.device_count()) | |
| prompt_ids = shard(prompt_ids) | |
| images = pipeline(prompt_ids, params, prng_seed, num_inference_steps, jit=True).images | |
| images = images.reshape((images.shape[0] * images.shape[1], ) + images.shape[-3:]) | |
| images = pipeline.numpy_to_pil(images) | |
| return images | |
| HF_ACCESS_TOKEN = os.environ["HFAUTH"] | |
| # Load Model | |
| # - Reference: https://github.com/huggingface/diffusers/blob/main/README.md | |
| pipeline, params = FlaxStableDiffusionPipeline.from_pretrained( | |
| "CompVis/stable-diffusion-v1-4", | |
| use_auth_token = HF_ACCESS_TOKEN, | |
| revision="bf16", | |
| dtype=jnp.bfloat16, | |
| ) | |
| def text_to_image(text): | |
| images = sd2_inference(pipeline, [text], params, seed = 42, num_inference_steps = 5 ) | |
| img = images[0] | |
| return img | |
| examples = ["apple", | |
| "banana", | |
| "chocolate"] | |
| if __name__ == '__main__': | |
| interFace = gr.Interface(fn=text_to_image, | |
| inputs=gr.inputs.Textbox(placeholder="Enter the text to Encode to an image", label="Text " | |
| "query", | |
| lines=1), | |
| outputs=gr.outputs.Image(type="auto", label="Generated Image"), | |
| verbose=True, | |
| examples=examples, | |
| title="Generate Image from Text", | |
| description="", | |
| theme="huggingface") | |
| interFace.launch() |