mswhite's picture
Update app.py
92920e1 verified
# ==============================
# Main Code
# ==============================
import spaces
import os
import gradio as gr
import torch
from diffusers import DiffusionPipeline, StableDiffusionPipeline, StableDiffusionImg2ImgPipeline
# Nvidia GPU config
pipe = StableDiffusionPipeline.from_pretrained("mswhite/manga_model", torch_dtype=torch.float16) #, safety_checker=None)
pipe = pipe.to("cuda")
pipe.safety_checker = lambda images, **kwargs: (images, [False] * len(images))
@spaces.GPU # decorator to use ZeroGPU instance - comment out to use other GPUs
def infer_diffusion(text):
prompt = "nik_mag" + text
return pipe(prompt).images[0]
# Mobile Version of Blocks
with gr.Blocks() as demo:
gr.Markdown("<img src=https://huggingface.co/spaces/mswhite/Nik/resolve/main/logo.jpg width=720px>")
output = gr.Image(label="Output",width=720)
art_to_draw = gr.Textbox(label="Prompt to Draw: e.g. ")
flashb_btn = gr.Button("Generate A New Manga Character")
flashb_btn.click(fn=infer_diffusion, inputs=[art_to_draw], outputs=[output])
demo.launch()