DarwinAnim8or's picture
Create app.py
312c9c4 verified
import gradio as gr
import time
from openvino.runtime import Core
from optimum.intel import OVStableDiffusionPipeline
from PIL import Image as PImg
import numpy as np
model_id = "NoCrypt/SomethingV2_2"
core = Core()
# Check available devices
devices = core.available_devices
print("Available devices:", devices)
# Use 'GPU' if available, otherwise fall back to 'CPU'
device = "GPU" if "GPU" in devices else "CPU"
print(f"Using device: {device}")
ov_pipe_bf16 = OVStableDiffusionPipeline.from_pretrained(
model_id,
export=True,
device=device # Specify the device here
)
# The compile step is not needed as it's handled internally
def generate_image(prompt, num_inference_steps):
# Generate an image from the prompt using the pipeline
start = time.time()
output = ov_pipe_bf16(prompt, num_inference_steps=num_inference_steps, output_type="np")
end = time.time()
print("Inference time: ", end - start)
# Convert the image to PIL Image object
image_data = output.images[0]
image_data = (image_data * 255).clip(0, 255).astype(np.uint8)
image = PImg.fromarray(image_data)
# Calculate the target size based on the scaling factor
target_resolution = 1.2
width = int(image.width * target_resolution)
height = int(image.height * target_resolution)
target_size = (width, height)
# Upscale the image to the target resolution
upscaled_image = image.resize(target_size, resample=PImg.BICUBIC)
return upscaled_image
examples = [
["masterpiece, best quality, 1girl, blonde, colorful, clouds, outdoors, falling leaves, smiling, whimsical"],
["masterpiece, best quality, landscape"],
["masterpiece, best quality, 1girl, aqua eyes, baseball cap, blonde hair, looking at viewer, shirt, short hair, simple background, solo, upper body, yellow shirt"]
]
iface = gr.Interface(
fn=generate_image,
inputs=[
gr.Textbox(label="Enter a prompt"),
gr.Slider(minimum=1, maximum=20, value=8, step=1, label="Number of inference steps")
],
outputs=gr.Image(label="Generated image"),
title="OpenVINO Anime Diffusion",
description="A gradio app that generates an image from a text prompt using the stable diffusion pipeline using the OpenVINO library for speed!",
examples=examples,
)
iface.launch()