Spaces:
Runtime error
Runtime error
import random | |
import gradio as gr | |
import numpy as np | |
import spaces | |
import torch | |
import torchvision.transforms as transforms | |
from PIL import Image | |
from torchmetrics.functional.image import structural_similarity_index_measure as ssim | |
from transformers import CLIPModel, CLIPProcessor | |
from tools import synth | |
device = "cuda" if torch.cuda.is_available() else "cpu" | |
model_path = "runwayml/stable-diffusion-v1-5" | |
clip_model = CLIPModel.from_pretrained("openai/clip-vit-large-patch14").to(device) | |
clip_processor = CLIPProcessor.from_pretrained("openai/clip-vit-large-patch14") | |
if torch.cuda.is_available(): | |
torch.cuda.max_memory_allocated(device=device) | |
pipe = synth.pipe_img( | |
model_path=model_path, | |
device=device, | |
use_torchcompile=False, | |
) | |
else: | |
pipe = synth.pipe_img( | |
model_path=model_path, | |
device=device, | |
apply_optimization=False, | |
) | |
MAX_SEED = np.iinfo(np.int32).max | |
MAX_IMAGE_SIZE = 1024 | |
def infer( | |
input_image, | |
prompt1, | |
prompt2, | |
negative_prompt, | |
seed, | |
randomize_seed, | |
width, | |
height, | |
guidance_scale, | |
interpolation_step, | |
num_inference_steps, | |
num_interpolation_steps, | |
): | |
device = "cuda" if torch.cuda.is_available() else "cpu" | |
# Input Validation | |
try: | |
assert num_interpolation_steps % 2 == 0 | |
except AssertionError: | |
raise ValueError("num_interpolation_steps must be an even number") | |
if randomize_seed: | |
seed = random.randint(0, MAX_SEED) | |
prompts = [prompt1, prompt2] | |
generator = torch.Generator().manual_seed(seed) | |
sample_mid_interpolation = num_interpolation_steps | |
remove_n_middle = 0 | |
interpolated_prompt_embeds, prompt_metadata = synth.interpolatePrompts( | |
prompts, | |
pipe, | |
num_interpolation_steps, | |
sample_mid_interpolation, | |
remove_n_middle=remove_n_middle, | |
device=device, | |
) | |
negative_prompts = [negative_prompt, negative_prompt] | |
if negative_prompts != ["", ""]: | |
interpolated_negative_prompts_embeds, _ = synth.interpolatePrompts( | |
negative_prompts, | |
pipe, | |
num_interpolation_steps, | |
sample_mid_interpolation, | |
remove_n_middle=remove_n_middle, | |
device=device, | |
) | |
else: | |
interpolated_negative_prompts_embeds, _ = [None] * len( | |
interpolated_prompt_embeds | |
), None | |
latents = torch.randn( | |
(1, pipe.unet.config.in_channels, height // 8, width // 8), | |
generator=generator, | |
).to(device) | |
embed_pairs = zip(interpolated_prompt_embeds, interpolated_negative_prompts_embeds) | |
embed_pairs_list = list(embed_pairs) | |
# offset step by -1 | |
prompt_embeds, negative_prompt_embeds = embed_pairs_list[interpolation_step - 1] | |
preprocess_input = transforms.Compose( | |
[transforms.ToTensor(), transforms.Resize((512, 512))] | |
) | |
input_img_tensor = preprocess_input(input_image).unsqueeze(0) | |
if negative_prompt_embeds is not None: | |
npe = negative_prompt_embeds[None, ...] | |
else: | |
npe = None | |
images_list = pipe( | |
height=height, | |
width=width, | |
num_images_per_prompt=1, | |
prompt_embeds=prompt_embeds[None, ...], | |
negative_prompt_embeds=npe, | |
num_inference_steps=num_inference_steps, | |
guidance_scale=guidance_scale, | |
generator=generator, | |
latents=latents, | |
image=input_img_tensor, | |
) | |
if images_list["nsfw_content_detected"][0]: | |
image = Image.open("samples/unsafe.jpeg") | |
return image, seed, "Unsafe content detected", "Unsafe content detected" | |
else: | |
image = images_list.images[0] | |
pred_image = transforms.ToTensor()(image).unsqueeze(0) | |
ssim_score = ssim(pred_image, input_img_tensor).item() | |
real_inputs = clip_processor( | |
text=prompts, padding=True, images=input_image, return_tensors="pt" | |
).to(device) | |
real_output = clip_model(**real_inputs) | |
synth_inputs = clip_processor( | |
text=prompts, padding=True, images=image, return_tensors="pt" | |
).to(device) | |
synth_output = clip_model(**synth_inputs) | |
cos_sim = torch.nn.CosineSimilarity(dim=1) | |
cosine_sim = ( | |
cos_sim(real_output.image_embeds, synth_output.image_embeds) | |
.detach() | |
.cpu() | |
.numpy() | |
.squeeze() | |
* 100 | |
) | |
return image, seed, round(ssim_score, 4), round(cosine_sim, 2) | |
examples1 = [ | |
"A photo of a garbage truck, dustcart", | |
"A photo of a Shih-Tzu, a type of dog", | |
] | |
examples2 = [ | |
"A photo of a cassette player", | |
"A photo of a beagle, a type of dog", | |
] | |
def update_steps(total_steps, interpolation_step): | |
return gr.update(maximum=total_steps) | |
def update_format(image_format): | |
return gr.update(format=image_format) | |
if torch.cuda.is_available(): | |
power_device = "GPU" | |
else: | |
power_device = "CPU" | |
with gr.Blocks(title="Generative Date Augmentation Demo") as demo: | |
gr.Markdown( | |
""" | |
# Data Augmentation with Image-to-Image Diffusion Models via Prompt Interpolation. | |
Main GitHub Repo: [Generative Data Augmentation](https://github.com/zhulinchng/generative-data-augmentation) | Image Classification Demo: [Generative Augmented Classifiers](https://huggingface.co/spaces/czl/generative-augmented-classifiers). | |
""" | |
) | |
with gr.Row(): | |
with gr.Column(): | |
input_image = gr.Image(type="pil", label="Image to Augment") | |
with gr.Row(): | |
prompt1 = gr.Text( | |
label="Prompt for the image to synthesize. (Actual class)", | |
show_label=True, | |
max_lines=1, | |
placeholder="Enter Prompt for the image to synthesize. (Actual class)", | |
container=False, | |
) | |
with gr.Row(): | |
prompt2 = gr.Text( | |
label="Prompt to augment against. (Confusing class)", | |
show_label=True, | |
max_lines=1, | |
placeholder="Enter Prompt to augment against. (Confusing class)", | |
container=False, | |
) | |
with gr.Row(): | |
gr.Examples( | |
examples=[ | |
"samples/n03417042_5234.JPEG", | |
"samples/n02086240_2799.JPEG", | |
], | |
inputs=[input_image], | |
label="Example Images", | |
) | |
gr.Examples( | |
examples=examples1, | |
inputs=[prompt1], | |
label="Example for Prompt 1 (Actual class)", | |
) | |
gr.Examples( | |
examples=examples2, | |
inputs=[prompt2], | |
label="Example for Prompt 2 (Confusing class)", | |
) | |
with gr.Row(): | |
num_interpolation_steps = gr.Slider( | |
label="Total Interpolation Steps", | |
minimum=2, | |
maximum=128, | |
step=2, | |
value=16, | |
) | |
interpolation_step = gr.Slider( | |
label="Sample Interpolation Step", | |
minimum=1, | |
maximum=16, | |
step=1, | |
value=8, | |
) | |
num_interpolation_steps.change( | |
fn=update_steps, | |
inputs=[num_interpolation_steps, interpolation_step], | |
outputs=[interpolation_step], | |
) | |
run_button = gr.Button("Run", scale=0) | |
with gr.Accordion("Advanced Settings", open=True): | |
negative_prompt = gr.Text( | |
label="Negative prompt", | |
max_lines=1, | |
placeholder="Enter a negative prompt", | |
visible=False, | |
) | |
seed = gr.Slider( | |
label="Seed", | |
minimum=0, | |
maximum=MAX_SEED, | |
step=1, | |
value=0, | |
) | |
randomize_seed = gr.Checkbox(label="Randomize seed", value=True) | |
gr.Markdown("Negative Prompt: ") | |
with gr.Row(): | |
negative_prompt = gr.Text( | |
label="Negative Prompt", | |
show_label=True, | |
max_lines=1, | |
value="blurry image, disfigured, deformed, distorted, cartoon, drawings", | |
container=False, | |
) | |
with gr.Row(): | |
width = gr.Slider( | |
label="Width", | |
minimum=256, | |
maximum=MAX_IMAGE_SIZE, | |
step=32, | |
value=512, | |
) | |
height = gr.Slider( | |
label="Height", | |
minimum=256, | |
maximum=MAX_IMAGE_SIZE, | |
step=32, | |
value=512, | |
) | |
with gr.Row(): | |
guidance_scale = gr.Slider( | |
label="Guidance scale", | |
minimum=0.0, | |
maximum=10.0, | |
step=0.1, | |
value=8.0, | |
) | |
num_inference_steps = gr.Slider( | |
label="Number of inference steps", | |
minimum=1, | |
maximum=80, | |
step=1, | |
value=25, | |
) | |
with gr.Row(): | |
image_type = gr.Radio( | |
choices=[ | |
"webp", | |
"png", | |
"jpeg", | |
], | |
label="Download Image Format", | |
value="jpeg", | |
) | |
with gr.Column(): | |
result = gr.Image(label="Result", show_label=False, format="jpeg") | |
image_type.change( | |
fn=update_format, | |
inputs=[image_type], | |
outputs=[result], | |
) | |
gr.Markdown( | |
""" | |
Metadata: | |
""" | |
) | |
with gr.Row(): | |
show_seed = gr.Label(label="Seed:", value="Randomized seed") | |
ssim_score = gr.Label( | |
label="SSIM Score:", value="Generate to see score" | |
) | |
cos_sim = gr.Label(label="CLIP Score:", value="Generate to see score") | |
if power_device == "GPU": | |
gr.Markdown( | |
f""" | |
Currently running on {power_device}. | |
""" | |
) | |
else: | |
gr.Markdown( | |
f""" | |
Currently running on {power_device}. | |
Note: Running on CPU will take longer (approx. 6 minutes with default settings). | |
""" | |
) | |
gr.Markdown( | |
""" | |
This demo is created as part of the 'Investigating the Effectiveness of Generative Diffusion Models in Synthesizing Images for Data Augmentation in Image Classification' dissertation. | |
The user can augment an image by interpolating between two prompts, and specify the number of interpolation steps and the specific step to generate the image. | |
View the files used in this demo [here](https://huggingface.co/spaces/czl/generative-data-augmentation-demo/tree/main). | |
Usage Instructions & Documentation [here](https://huggingface.co/spaces/czl/generative-data-augmentation-demo/blob/main/README.md). | |
Note: Safety checker is enabled to prevent unsafe content from being displayed in this public demo. | |
""" | |
) | |
run_button.click( | |
fn=infer, | |
inputs=[ | |
input_image, | |
prompt1, | |
prompt2, | |
negative_prompt, | |
seed, | |
randomize_seed, | |
width, | |
height, | |
guidance_scale, | |
interpolation_step, | |
num_inference_steps, | |
num_interpolation_steps, | |
], | |
outputs=[result, show_seed, ssim_score, cos_sim], | |
) | |
demo.queue().launch(show_error=True) | |