File size: 1,164 Bytes
850467f 22ac5f3 850467f 22ac5f3 853bdb4 22ac5f3 853bdb4 850467f 853bdb4 22ac5f3 853bdb4 22ac5f3 853bdb4 22ac5f3 853bdb4 22ac5f3 853bdb4 22ac5f3 | 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 | import gradio as gr
from diffusers import PNDMScheduler, DDIMScheduler, LMSDiscreteScheduler
import torch
from torch import autocast
from diffusers import StableDiffusionPipeline
#funtion to call model and show images
#function to make API call writing
def show_image(prompt):
num_images = 2
scheduler = PNDMScheduler(beta_start=0.00085, beta_end=0.012, beta_schedule="scaled_linear", skip_prk_steps=True)
model_id = "runwayml/stable-diffusion-v1-5"
device = "cuda"
remove_safety = False
pipe = StableDiffusionPipeline.from_pretrained(model_id, scheduler=scheduler, torch_dtype=torch.float16, revision="fp16", use_auth_token=True)
if remove_safety:
pipe.safety_checker = lambda images, clip_input: (images, False)
pipe = pipe.to(device)
prompts = [ prompt ] * num_images
with autocast("cuda"):
images = pipe(prompts, guidance_scale=7.5, num_inference_steps=50).images
images[0].save("output.jpg")
print(type(images[0]))
return images[0]
demo = gr.Interface(fn=show_image, inputs="textbox", outputs=gr.Image(label = "Output Image Component"))
if __name__ == "__main__":
demo.launch() |