prithivMLmods's picture
Update app.py
57709d9 verified
raw
history blame
6.44 kB
import gradio as gr
import numpy as np
import random
import uuid
from PIL import Image
import spaces
from diffusers import DiffusionPipeline
import torch
DESCRIPTIONx = """## SD-3.5 LARGE TURBO """
device = "cuda" if torch.cuda.is_available() else "cpu"
model_repo_id = "stabilityai/stable-diffusion-3.5-large-turbo"
if torch.cuda.is_available():
torch_dtype = torch.bfloat16
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
# Define styles
style_list = [
{
"name": "3840 x 2160",
"prompt": "hyper-realistic 8K image of {prompt}. ultra-detailed, lifelike, high-resolution, sharp, vibrant colors, photorealistic",
"negative_prompt": "cartoonish, low resolution, blurry, simplistic, abstract, deformed, ugly",
},
{
"name": "2560 x 1440",
"prompt": "hyper-realistic 4K image of {prompt}. ultra-detailed, lifelike, high-resolution, sharp, vibrant colors, photorealistic",
"negative_prompt": "cartoonish, low resolution, blurry, simplistic, abstract, deformed, ugly",
},
{
"name": "HD+",
"prompt": "hyper-realistic 2K image of {prompt}. ultra-detailed, lifelike, high-resolution, sharp, vibrant colors, photorealistic",
"negative_prompt": "cartoonish, low resolution, blurry, simplistic, abstract, deformed, ugly",
},
{
"name": "Style Zero",
"prompt": "{prompt}",
"negative_prompt": "",
},
]
STYLE_NAMES = [style["name"] for style in style_list]
DEFAULT_STYLE_NAME = STYLE_NAMES[0]
grid_sizes = {
"2x1": (2, 1),
"1x2": (1, 2),
"2x2": (2, 2),
"2x3": (2, 3),
"3x2": (3, 2),
"1x1": (1, 1)
}
@spaces.GPU(duration=60, enable_queue=True)
def infer(
prompt,
negative_prompt="",
seed=42,
randomize_seed=False,
width=1024,
height=1024,
guidance_scale=0.0,
num_inference_steps=4,
style="Style Zero",
grid_size="1x1",
progress=gr.Progress(track_tqdm=True),
):
selected_style = next(s for s in style_list if s["name"] == style)
styled_prompt = selected_style["prompt"].format(prompt=prompt)
styled_negative_prompt = selected_style["negative_prompt"]
if randomize_seed:
seed = random.randint(0, MAX_SEED)
generator = torch.Generator().manual_seed(seed)
grid_size_x, grid_size_y = grid_sizes.get(grid_size, (2, 2))
num_images = grid_size_x * grid_size_y
images = []
for _ in range(num_images):
image = pipe(
prompt=styled_prompt,
negative_prompt=styled_negative_prompt,
guidance_scale=guidance_scale,
num_inference_steps=num_inference_steps,
width=width,
height=height,
generator=generator,
).images[0]
images.append(image)
# Create a grid image
grid_img = Image.new('RGB', (width * grid_size_x, height * grid_size_y))
for i, img in enumerate(images[:num_images]):
grid_img.paste(img, (i % grid_size_x * width, i // grid_size_x * height))
# Save the grid image
unique_name = str(uuid.uuid4()) + ".png"
grid_img.save(unique_name)
return unique_name, seed
examples = [
"A capybara wearing a suit holding a sign that reads Hello World",
]
css = '''
.gradio-container{max-width: 585px !important}
h1{text-align:center}
footer {
visibility: hidden
}
'''
with gr.Blocks(css=css, theme="prithivMLmods/Minecraft-Theme") as demo:
gr.Markdown(DESCRIPTIONx)
with gr.Row():
prompt = gr.Text(
label="Prompt",
show_label=False,
max_lines=1,
placeholder="Enter your prompt",
container=False,
)
run_button = gr.Button("Run", scale=0, variant="primary")
result = gr.Image(label="Result", show_label=False)
with gr.Row(visible=True):
grid_size_selection = gr.Dropdown(
choices=["2x1", "1x2", "2x2", "2x3", "3x2", "1x1"],
value="1x1",
label="Grid Size"
)
with gr.Row(visible=True):
style_selection = gr.Radio(
show_label=True,
container=True,
interactive=True,
choices=STYLE_NAMES,
value=DEFAULT_STYLE_NAME,
label="Quality Style",
)
with gr.Accordion("Advanced Settings", open=False):
negative_prompt = gr.Text(
label="Negative prompt",
max_lines=1,
placeholder="Enter a negative prompt",
visible=False,
)
seed = gr.Slider(
label="Seed",
minimum=0,
maximum=MAX_SEED,
step=1,
value=0,
)
randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
with gr.Row():
width = gr.Slider(
label="Width",
minimum=512,
maximum=MAX_IMAGE_SIZE,
step=32,
value=1024,
)
height = gr.Slider(
label="Height",
minimum=512,
maximum=MAX_IMAGE_SIZE,
step=32,
value=1024,
)
with gr.Row():
guidance_scale = gr.Slider(
label="Guidance scale",
minimum=0.0,
maximum=7.5,
step=0.1,
value=0.0,
)
num_inference_steps = gr.Slider(
label="Number of inference steps",
minimum=1,
maximum=50,
step=1,
value=4,
)
gr.Examples(
examples=examples,
inputs=[prompt],
outputs=[result, seed],
fn=infer,
cache_examples=True,
cache_mode="lazy"
)
gr.on(
triggers=[run_button.click, prompt.submit],
fn=infer,
inputs=[
prompt,
negative_prompt,
seed,
randomize_seed,
width,
height,
guidance_scale,
num_inference_steps,
style_selection,
grid_size_selection,
],
outputs=[result, seed],
)
if __name__ == "__main__":
demo.launch()