Tramac's picture
Update app.py
4c7447e verified
raw
history blame
No virus
1.47 kB
import torch
import spaces
import gradio as gr
from diffusers import StableDiffusionPipeline
repo = "Tramac/style-portrait-v1-5"
# Ensure model and scheduler are initialized in GPU-enabled function
if torch.cuda.is_available():
pipeline = StableDiffusionPipeline.from_pretrained(
repo,
torch_dtype=torch.float16,
variant="fp16"
).to("cuda")
@spaces.GPU(enable_queue=True)
def generate_image(prompt, neg_prompt):
results = pipe(
prompt=prompt,
negative_prompt=neg_prompt,
height=832,
width=512,
num_inference_steps=20,
)
return results.images[0]
with gr.Blocks() as demo:
gr.HTML("<h1><center>Style Portrait</center></h1>")
gr.HTML("<p><center>text-to-image generation</center></p><p><center><a href='https://huggingface.co/Tramac/style-portrait-v1-5'></a></center></p>")
with gr.Row():
with gr.Column():
prompt = gr.Textbox(label='Enter your prompt', scale=8)
neg_prompt = gr.Textbox(label='Enter your negative prompt', scale=8)
submit = gr.Button(scale=1, variant='primary')
with gr.Column():
img = gr.Image(label='Style-Portrait Generated Image')
prompt.submit(
fn=generate_image,
inputs=[prompt, neg_prompt],
outputs=img,
)
submit.click(
fn=generate_image,
inputs=[prompt, neg_prompt],
outputs=img,
)
demo.queue().launch(share=True)