Duplicate from jax-diffusers-event/canny_coyo1m

Co-authored-by: Merve Noyan <[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,14 @@
+---
+title: Canny Coyo1m
+emoji: 🌖
+colorFrom: indigo
+colorTo: green
+sdk: gradio
+sdk_version: 3.24.1
+app_file: app.py
+pinned: false
+license: apache-2.0
+duplicated_from: jax-diffusers-event/canny_coyo1m
+---
+
+Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

app.py

@@ -0,0 +1,59 @@
+import gradio as gr
+import jax
+import numpy as np
+import jax.numpy as jnp
+from flax.jax_utils import replicate
+from flax.training.common_utils import shard
+from PIL import Image
+from diffusers import FlaxStableDiffusionControlNetPipeline, FlaxControlNetModel
+import cv2
+
+def create_key(seed=0):
+    return jax.random.PRNGKey(seed)
+
+def canny_filter(image):
+    gray_image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
+    blurred_image = cv2.GaussianBlur(gray_image, (5, 5), 0)
+    edges_image = cv2.Canny(blurred_image, 50, 150)
+    return edges_image
+
+# load control net and stable diffusion v1-5
+controlnet, controlnet_params = FlaxControlNetModel.from_pretrained(
+    "jax-diffusers-event/canny-coyo1m", dtype=jnp.bfloat16
+)
+pipe, params = FlaxStableDiffusionControlNetPipeline.from_pretrained(
+    "runwayml/stable-diffusion-v1-5", controlnet=controlnet, revision="flax", dtype=jnp.bfloat16
+)
+
+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())
+    im = canny_filter(image)
+    canny_image = Image.fromarray(im)
+    
+    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([canny_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,
+        prng_seed=rng,
+        num_inference_steps=50,
+        neg_prompt_ids=negative_prompt_ids,
+        jit=True,
+    ).images
+    
+    output_images = pipe.numpy_to_pil(np.asarray(output.reshape((num_samples,) + output.shape[-3:])))
+    return output_images
+
+gr.Interface(infer, inputs=["text", "text", "image"], outputs="gallery").launch()

requirements.txt

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