app.py

import gradio as gr
import numpy as np
import random
import time

import spaces
from diffusers import DiffusionPipeline
import torch
from PIL import Image

from gradio.themes.base import Base
from gradio.themes.utils import colors, fonts, sizes
from typing import Iterable

class Seafoam(Base):
    def __init__(
        self,
        *,
        primary_hue: colors.Color | str = colors.emerald,
        secondary_hue: colors.Color | str = colors.blue,
        neutral_hue: colors.Color | str = colors.gray,
        spacing_size: sizes.Size | str = sizes.spacing_md,
        radius_size: sizes.Size | str = sizes.radius_md,
        text_size: sizes.Size | str = sizes.text_lg,
        font: fonts.Font
        | str
        | Iterable[fonts.Font | str] = (
            fonts.GoogleFont("Quicksand"),
            "ui-sans-serif",
            "sans-serif",
        ),
        font_mono: fonts.Font
        | str
        | Iterable[fonts.Font | str] = (
            fonts.GoogleFont("IBM Plex Mono"),
            "ui-monospace",
            "monospace",
        ),
    ):
        super().__init__(
            primary_hue=primary_hue,
            secondary_hue=secondary_hue,
            neutral_hue=neutral_hue,
            spacing_size=spacing_size,
            radius_size=radius_size,
            text_size=text_size,
            font=font,
            font_mono=font_mono,
        )

seafoam = Seafoam()

# Device and model setup
device = "cuda" if torch.cuda.is_available() else "cpu"
model_repo_id = "stabilityai/stable-diffusion-3.5-large-turbo"

torch_dtype = torch.bfloat16 if torch.cuda.is_available() else torch.float32
pipe = DiffusionPipeline.from_pretrained(model_repo_id, torch_dtype=torch_dtype)
pipe = pipe.to(device)

pipe.load_lora_weights("prithivMLmods/SD3.5-Turbo-Realism-2.0-LoRA", weight_name="SD3.5-Turbo-Realism-2.0-LoRA.safetensors")
trigger_word = "Turbo Realism"  
pipe.fuse_lora(lora_scale=1.0)

# Constants
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 = [s["name"] for s in style_list]
DEFAULT_STYLE_NAME = STYLE_NAMES[0]

# Define grid layouts
grid_sizes = {
    "2x1": (2, 1),
    "1x2": (1, 2),
    "2x2": (2, 2),
    "2x3": (2, 3),
    "3x2": (3, 2),
    "1x1": (1, 1),
}

@spaces.GPU
def infer(
     prompt,
     negative_prompt="",
     seed=42,
     randomize_seed=False,
     width=1024,
     height=1024,
     guidance_scale=7.5,
     num_inference_steps=10,
     style="Style Zero",
     grid_size="1x1",
     progress=gr.Progress(track_tqdm=True),
 ):
     if randomize_seed:
         seed = random.randint(0, MAX_SEED)
     generator = torch.Generator().manual_seed(seed)

     selected_style = next(s for s in style_list if s["name"] == style)
     styled_prompt = selected_style["prompt"].format(prompt=prompt)
     styled_negative = selected_style["negative_prompt"] or negative_prompt

     grid_x, grid_y = grid_sizes.get(grid_size, (1, 1))
     num_images = grid_x * grid_y

     output = pipe(
         prompt=styled_prompt,
         negative_prompt=styled_negative,
         width=width,
         height=height,
         guidance_scale=guidance_scale,
         num_inference_steps=num_inference_steps,
         generator=generator,
         num_images_per_prompt=num_images,
     )

     grid_img = Image.new('RGB', (width * grid_x, height * grid_y))
     for i, img in enumerate(output.images[:num_images]):
         x = (i % grid_x) * width
         y = (i // grid_x) * height
         grid_img.paste(img, (x, y))

     return grid_img, seed

examples = [
    "A tiny, cute astronaut emerging from a cracked eggshell on the surface of the moon, surrounded by cosmic dust and distant planets in the starry sky. The scene is illuminated by soft lunar light, showcasing the texture of the moon's surface. Shot in ultra-detailed 4K resolution, with a sci-fi fantasy atmosphere and planetary background",
    "A cute, fluffy cat standing upright on its hind legs, holding a hand-drawn sign that says 'Hello World' in bold, playful letters. The background is softly blurred, emphasizing the cat’s detailed fur texture and the colorful sign. Captured in portrait aspect ratio --ar 85:128, using --v 6.0 and --style raw for a semi-realistic, endearing look",
    "A photorealistic image of a cold coffee beverage in a glass cup, condensation on the surface, sitting on a wooden café table with shallow depth of field. The background features beautiful bokeh lighting, creating a cozy, blurred café ambiance. Shot in portrait mode --ar 85:128, 4K ultra-resolution, using --style raw5 for authentic textures, --v 6.0.",
    "An anime-style food illustration of a golden-brown wiener schnitzel, perfectly fried and crispy, served on a white ceramic plate. Accompanied by fresh parsley garnish and thin lemon slices, with artistic shading and stylized highlights. Captured in an anime illustration format --style raw5, high color saturation, food-themed aesthetic."
]

css = '''
.gradio-container {
    max-width: 100%;
    margin: 0 auto;
}
h1 { text-align: center; }
footer { visibility: hidden; }
'''

with gr.Blocks(theme=seafoam, css=css) as demo:
    gr.Markdown("## T2I SD3.5")

    with gr.Row():
        with gr.Column(scale=1):
            with gr.Row():
                prompt = gr.Text(
                    show_label=False,
                    max_lines=1,
                    placeholder="Enter your prompt",
                    container=False,
                )
                run_button = gr.Button("Run", scale=0, variant="primary")

            result = gr.Image(show_label=False)

            grid_size_selection = gr.Dropdown(
                choices=list(grid_sizes.keys()),
                value="1x1",
                label="Grid Size"
            )

            with gr.Accordion("Advanced Settings", open=False):
                negative_prompt = gr.Text(
                    label="Negative prompt",
                    max_lines=1,
                    placeholder="Enter a negative prompt",
                    value="(deformed, distorted, disfigured:1.3), poorly drawn, bad anatomy, wrong anatomy, extra limb, missing limb, floating limbs, (mutated hands and fingers:1.4), disconnected limbs, mutation, mutated, ugly, disgusting, blurry, amputation",
                )
                seed = gr.Slider(0, MAX_SEED, value=0, label="Seed")
                randomize_seed = gr.Checkbox(label="Randomize seed", value=True)

                with gr.Row():
                    width = gr.Slider(512, MAX_IMAGE_SIZE, step=32, value=1024, label="Width")
                    height = gr.Slider(512, MAX_IMAGE_SIZE, step=32, value=1024, label="Height")

                with gr.Row():
                    guidance_scale = gr.Slider(0.0, 7.5, step=0.1, value=0.0, label="Guidance scale")
                    num_inference_steps = gr.Slider(1, 50, step=1, value=10, label="Number of inference steps")

                style_selection = gr.Radio(
                    choices=STYLE_NAMES,
                    value=DEFAULT_STYLE_NAME,
                    label="Quality Style",
                )

        with gr.Column(scale=1):
            gr.Examples(
                examples=examples,
                inputs=[prompt],
                outputs=[result, seed],
                fn=infer,
                cache_examples=False,
                label="Prompt Examples"
            )

    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(ssr_mode=False)