Duplicate from bguisard/stable-diffusion-nano

Co-authored-by: Bruno Guisard <[email protected]>

.gitattributes

@@ -0,0 +1,34 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,16 @@
+---
+title: Stable Diffusion Nano
+emoji: 📊
+colorFrom: indigo
+colorTo: pink
+sdk: gradio
+sdk_version: 3.27.0
+app_file: app.py
+pinned: false
+license: creativeml-openrail-m
+tags:
+- jax-diffusers-event
+duplicated_from: bguisard/stable-diffusion-nano
+---
+
+Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

app.py

@@ -0,0 +1,85 @@
+import gradio as gr
+import jax
+import jax.numpy as jnp
+from diffusers import FlaxPNDMScheduler, FlaxStableDiffusionPipeline
+from flax.jax_utils import replicate
+from flax.training.common_utils import shard
+
+DTYPE = jnp.float16
+
+pipeline, pipeline_params = FlaxStableDiffusionPipeline.from_pretrained(
+    "bguisard/stable-diffusion-nano-2-1",
+    revision="flax",
+    dtype=DTYPE,
+)
+if DTYPE != jnp.float32:
+    # There is a known issue with schedulers when loading from a pre trained
+    # pipeline. We need the schedulers to always use float32.
+    # See: https://github.com/huggingface/diffusers/issues/2155
+    scheduler, scheduler_params = FlaxPNDMScheduler.from_pretrained(
+        pretrained_model_name_or_path="bguisard/stable-diffusion-nano-2-1",
+        subfolder="scheduler",
+        dtype=jnp.float32,
+    )
+    pipeline_params["scheduler"] = scheduler_params
+    pipeline.scheduler = scheduler
+
+
+def generate_image(prompt: str, inference_steps: int = 30, prng_seed: int = 0):
+    rng = jax.random.PRNGKey(int(prng_seed))
+    rng = jax.random.split(rng, jax.device_count())
+    p_params = replicate(pipeline_params)
+
+    num_samples = 1
+    prompt_ids = pipeline.prepare_inputs([prompt] * num_samples)
+    prompt_ids = shard(prompt_ids)
+
+    images = pipeline(
+        prompt_ids=prompt_ids,
+        params=p_params,
+        prng_seed=rng,
+        height=128,
+        width=128,
+        num_inference_steps=int(inference_steps),
+        jit=True,
+    ).images
+
+    images = images.reshape((num_samples,) + images.shape[-3:])
+    images = pipeline.numpy_to_pil(images)
+    return images[0]
+
+
+prompt_input = gr.inputs.Textbox(
+    label="Prompt", placeholder="A watercolor painting of a bird"
+)
+inf_steps_input = gr.inputs.Slider(
+    minimum=1, maximum=100, default=30, step=1, label="Inference Steps"
+)
+seed_input = gr.inputs.Number(default=0, label="Seed")
+
+
+app = gr.Interface(
+    fn=generate_image,
+    inputs=[prompt_input, inf_steps_input, seed_input],
+    outputs="image",
+    title="🤗 Stable Diffusion Nano 🧨",
+    description=(
+        "Based on stable diffusion and fine-tuned on 128x128 images, "
+        "[Stable Diffusion Nano](https://huggingface.co/bguisard/stable-diffusion-nano-2-1) allows "
+        "for fast prototyping of diffusion models, enabling quick experimentation "
+        "with easily available hardware."
+        " It performs reasonably well on several tasks, but it struggles "
+        "with small details such as faces."
+    ),
+    css="h1 { text-align: center }",
+    # Some examples were copied from hf.co/spaces/stabilityai/stable-diffusion
+    examples=[
+        ["A watercolor painting of a bird", 50, 0],
+        ["A small cabin on top of a snowy mountain in the style of Disney, artstation", 50, 232190380],
+        ["A mecha robot in a favela in expressionist style", 50, 827198341273],
+        ["Marvel MCU deadpool, red mask, red shirt, red gloves, black shoulders, black elbow pads, black legs, gold buckle, black belt, black mask, white eyes, black boots, fuji low light color 35mm film, downtown Osaka alley at night out of focus in background, neon lights", 50, 23856839],
+    ],
+)
+
+app.launch()
+

requirements.txt

@@ -0,0 +1,6 @@
+-f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
+jax[cuda11_cudnn82]
+jaxlib
+flax
+transformers
+git+https://github.com/huggingface/diffusers@main