Lifeinhockey's picture
Update app.py
28f39ff verified
raw
history blame
4.73 kB
!pip install rembg
import gradio as gr
import numpy as np
import random
from diffusers import DiffusionPipeline
import torch
from rembg import remove
from PIL import Image
device = "cuda" if torch.cuda.is_available() else "cpu"
model_repo_id = "stabilityai/sdxl-turbo" # Replace to the model you would like to use
if torch.cuda.is_available():
torch_dtype = torch.float16
else:
torch_dtype = torch.float32
pipe = DiffusionPipeline.from_pretrained(model_repo_id, torch_dtype=torch_dtype)
pipe = pipe.to(device)
MAX_SEED = np.iinfo(np.int32).max
MAX_IMAGE_SIZE = 1024
# @spaces.GPU #[uncomment to use ZeroGPU]
def infer(
model,
prompt,
negative_prompt,
seed,
width,
height,
guidance_scale,
num_inference_steps,
progress=gr.Progress(track_tqdm=True),
):
global model_repo_id
if model != model_repo_id:
print(model, model_repo_id)
pipe = DiffusionPipeline.from_pretrained(model, torch_dtype=torch_dtype)
pipe = pipe.to(device)
generator = torch.Generator().manual_seed(seed)
image = pipe(
prompt=prompt,
negative_prompt=negative_prompt,
guidance_scale=guidance_scale,
num_inference_steps=num_inference_steps,
width=width,
height=height,
generator=generator,
).images[0]
return image, seed
examples = [
"Young man in anime style. The image is of high sharpness and resolution. A handsome, thoughtful man. The man is depicted in the foreground, close-up or middle plan. The background is blurry, not sharp. The play of light and shadow is visible on the face and clothes."
]
examples_negative = [
"blurred details, low resolution, poor image of a man's face, poor quality, artifacts, black and white image"
]
css = """
#col-container {
margin: 0 auto;
max-width: 640px;
}
"""
available_models = [
"stable-diffusion-v1-5/stable-diffusion-v1-5",
"CompVis/stable-diffusion-v1-4",
]
with gr.Blocks(css=css) as demo:
with gr.Column(elem_id="col-container"):
gr.Markdown(" # Text-to-Image Gradio Template from V. Gorsky")
model = gr.Dropdown(
label="Model Selection",
choices=available_models,
value="stable-diffusion-v1-5/stable-diffusion-v1-5",
interactive=True
)
prompt = gr.Text(
label="Prompt",
show_label=False,
max_lines=1,
placeholder="Enter your prompt",
container=False,
)
negative_prompt = gr.Text(
label="Negative prompt",
max_lines=1,
placeholder="Enter a negative prompt",
visible=True,
)
seed = gr.Slider(
label="Seed",
minimum=0,
maximum=MAX_SEED,
step=1,
value=0,
)
guidance_scale = gr.Slider(
label="Guidance scale",
minimum=0.0,
maximum=10.0,
step=0.1,
value=7.5, # Replace with defaults that work for your model
)
num_inference_steps = gr.Slider(
label="Number of inference steps",
minimum=1,
maximum=100,
step=1,
value=30, # Replace with defaults that work for your model
)
with gr.Accordion("Advanced Settings", open=False):
with gr.Row():
width = gr.Slider(
label="Width",
minimum=256,
maximum=MAX_IMAGE_SIZE,
step=32,
value=512, # Replace with defaults that work for your model
)
height = gr.Slider(
label="Height",
minimum=256,
maximum=MAX_IMAGE_SIZE,
step=32,
value=512, # Replace with defaults that work for your model
)
gr.Examples(examples=examples, inputs=[prompt])
gr.Examples(examples=examples_negative, inputs=[negative_prompt])
run_button = gr.Button("Run", scale=0, variant="primary")
result = gr.Image(label="Result", show_label=False)
gr.on(
triggers=[run_button.click, prompt.submit],
fn=infer,
inputs=[
model,
prompt,
negative_prompt,
seed,
width,
height,
guidance_scale,
num_inference_steps,
],
outputs=[result, seed],
)
if __name__ == "__main__":
demo.launch()