Upload app.py

app.py

@@ -0,0 +1,106 @@
+import gradio as gr
+
+import requests
+from PIL import Image
+from pathlib import Path
+from io import BytesIO
+
+# Diffusers
+from diffusers import StableDiffusionImg2ImgPipeline
+from diffusers import (
+    FlaxStableDiffusionControlNetPipeline,
+    FlaxControlNetModel,
+    FlaxStableDiffusionPipeline,
+)
+from diffusers import ControlNetModel
+from diffusers.utils import load_image
+
+# Pytorch
+import torch
+
+# Numpy
+import numpy as np
+
+# Jax
+import jax
+import jax.numpy as jnp
+from jax import pmap
+
+# Flax
+import flax
+from flax.jax_utils import replicate
+from flax.training.common_utils import shard
+
+
+def create_key(seed=0):
+    return jax.random.PRNGKey(seed)
+
+
+def image_grid(imgs, rows, cols):
+    w, h = imgs[0].size
+    grid = Image.new("RGB", size=(cols * w, rows * h))
+    for i, img in enumerate(imgs):
+        grid.paste(img, box=(i % cols * w, i // cols * h))
+    return grid
+
+
+# load control net and stable diffusion v1-5
+controlnet, controlnet_params = FlaxControlNetModel.from_pretrained(
+    "jax-diffuser-event/learner/trained_model_v0.1", from_flax=True, dtype=jnp.float32
+)
+
+pipe, params = FlaxStableDiffusionControlNetPipeline.from_pretrained(
+    "runwayml/stable-diffusion-v1-5",
+    controlnet=controlnet,
+    from_pt=True,
+    dtype=jnp.float32,
+    safety_checker=None,
+)
+
+
+# inference function takes prompt, negative prompt and image
+def infer(prompts, negative_prompts, image):
+    params["controlnet"] = controlnet_params
+
+    num_samples = 1  # jax.device_count()
+    rng = create_key(0)
+    rng = jax.random.split(rng, jax.device_count())
+    battlemap_image = load_image(image)
+
+    prompt_ids = pipe.prepare_text_inputs([prompts] * num_samples)
+    negative_prompt_ids = pipe.prepare_text_inputs([negative_prompts] * num_samples)
+    processed_image = pipe.prepare_image_inputs([battlemap_image] * num_samples)
+
+    p_params = replicate(params)
+    prompt_ids = shard(prompt_ids)
+    negative_prompt_ids = shard(negative_prompt_ids)
+    processed_image = shard(processed_image)
+
+    output = pipe(
+        prompt_ids=prompt_ids,
+        image=processed_image,
+        params=p_params,
+        # params = params,
+        prng_seed=rng,
+        num_inference_steps=50,
+        neg_prompt_ids=negative_prompt_ids,
+        jit=True,
+    ).images
+
+    output_image = pipe.numpy_to_pil(
+        np.asarray(output.reshape((num_samples,) + output.shape[-3:]))
+    )
+
+    return output_image
+
+
+title = "ControlNet on Battlemaps"
+description = "This is a demo on ControlNet based on Bettlemaps."
+# you need to pass inputs and outputs according to inference function
+gr.Interface(
+    fn=infer,
+    inputs=["text", "text", "image"],
+    outputs="gallery",
+    title=title,
+    description=description,
+).launch()