Spaces:
Runtime error
Runtime error
File size: 5,704 Bytes
174ccbd 369d822 174ccbd 369d822 174ccbd 369d822 41c10f1 369d822 41c10f1 369d822 41c10f1 369d822 41c10f1 369d822 41c10f1 369d822 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 |
import gradio as gr
import numpy as np
import torch
import jax
import jax.numpy as jnp
from flax.jax_utils import replicate
from flax.training.common_utils import shard
from PIL import Image
from segment_anything import SamPredictor, sam_model_registry, SamAutomaticMaskGenerator
from diffusers import (
FlaxStableDiffusionControlNetPipeline,
FlaxControlNetModel,
)
from transformers import pipeline
import colorsys
sam_checkpoint = "sam_vit_h_4b8939.pth"
model_type = "vit_h"
device = "cuda" if torch.cuda.is_available() else "cpu"
#sam = sam_model_registry[model_type](checkpoint=sam_checkpoint)
#sam.to(device=device)
#predictor = SamPredictor(sam)
#mask_generator = SamAutomaticMaskGenerator(sam)
generator = pipeline(model="facebook/sam-vit-base", task="mask-generation", points_per_batch=256)
#image_url = "https://huggingface.co/ybelkada/segment-anything/resolve/main/assets/car.png"
# controlnet, controlnet_params = FlaxControlNetModel.from_pretrained(
# "SAMControlNet/sd-controlnet-sam-seg", dtype=jnp.float32
# )
# pipe, params = FlaxStableDiffusionControlNetPipeline.from_pretrained(
# "runwayml/stable-diffusion-v1-5",
# controlnet=controlnet,
# revision="flax",
# dtype=jnp.bfloat16,
# )
# params["controlnet"] = controlnet_params
# p_params = replicate(params)
with gr.Blocks() as demo:
gr.Markdown("# Ahsans version WildSynth: Synthetic Wildlife Data Generation")
gr.Markdown(
"""
## Work in Progress
### About
We have trained a JAX ControlNet model for semantic segmentation on Wildlife Animal Images.
For the training data creation we used the [Wildlife Animals Images](https://www.kaggle.com/datasets/anshulmehtakaggl/wildlife-animals-images) dataset.
We created segmentation masks with the help of [Grounded SAM](https://github.com/IDEA-Research/Grounded-Segment-Anything) where we used the animals names
as input prompts for detection and more accurate segmentation.
### How To Use
"""
)
with gr.Row():
input_img = gr.Image(label="Input", type="pil")
mask_img = gr.Image(label="Mask", interactive=False)
output_img = gr.Image(label="Output", interactive=False)
with gr.Row():
prompt_text = gr.Textbox(lines=1, label="Prompt")
negative_prompt_text = gr.Textbox(lines=1, label="Negative Prompt")
with gr.Row():
submit = gr.Button("Submit")
clear = gr.Button("Clear")
def generate_mask(image):
outputs = generator(image, points_per_batch=256)
mask_images = []
for mask in outputs["masks"]:
color = np.concatenate([np.random.random(3), np.array([1.0])], axis=0)
h, w = mask.shape[-2:]
mask_image = mask.reshape(h, w, 1) * color.reshape(1, 1, -1)
mask_images.append(mask_image)
return np.stack(mask_images)
# def infer(
# image, prompts, negative_prompts, num_inference_steps=50, seed=4, num_samples=4
# ):
# try:
# rng = jax.random.PRNGKey(int(seed))
# num_inference_steps = int(num_inference_steps)
# image = Image.fromarray(image, mode="RGB")
# num_samples = max(jax.device_count(), int(num_samples))
# p_rng = jax.random.split(rng, jax.device_count())
# 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([image] * num_samples)
# 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=p_rng,
# num_inference_steps=num_inference_steps,
# neg_prompt_ids=negative_prompt_ids,
# jit=True,
# ).images
# del negative_prompt_ids
# del processed_image
# del prompt_ids
# output = output.reshape((num_samples,) + output.shape[-3:])
# final_image = [np.array(x * 255, dtype=np.uint8) for x in output]
# print(output.shape)
# del output
# except Exception as e:
# print("Error: " + str(e))
# final_image = [np.zeros((512, 512, 3), dtype=np.uint8)] * num_samples
# finally:
# gc.collect()
# return final_image
# def _clear(sel_pix, img, mask, seg, out, prompt, neg_prompt, bg):
# img = None
# mask = None
# seg = None
# out = None
# prompt = ""
# neg_prompt = ""
# bg = False
# return img, mask, seg, out, prompt, neg_prompt, bg
input_img.change(
generate_mask,
inputs=[input_img],
outputs=[mask_img],
)
# submit.click(
# infer,
# inputs=[mask_img, prompt_text, negative_prompt_text],
# outputs=[output_img],
# )
# clear.click(
# _clear,
# inputs=[
# input_img,
# mask_img,
# output_img,
# prompt_text,
# negative_prompt_text,
# ],
# outputs=[
# input_img,
# mask_img,
# output_img,
# prompt_text,
# negative_prompt_text,
# ],
# )
if __name__ == "__main__":
demo.queue()
demo.launch() |