app.py

import os
import random
import uuid
import json

import gradio as gr
import numpy as np
from PIL import Image
import spaces
import torch

from diffusers import DiffusionPipeline
from diffusers import StableDiffusionXLPipeline, EulerAncestralDiscreteScheduler # EulerAncestralDiscreteScheduler not explicitly used but imported
from typing import Tuple

bad_words = json.loads(os.getenv('BAD_WORDS', "[]"))
bad_words_negative = json.loads(os.getenv('BAD_WORDS_NEGATIVE', "[]"))
default_negative = os.getenv("default_negative","")

def check_text(prompt, negative=""):
    for i in bad_words:
        if i in prompt:
            return True
    for i in bad_words_negative:
        if i in negative:
            return True
    return False

style_list = [
    {
        "name": "Photo",
        "prompt": "cinematic photo {prompt}. 35mm photograph, film, bokeh, professional, 4k, highly detailed",
        "negative_prompt": "drawing, painting, crayon, sketch, graphite, impressionist, noisy, blurry, soft, deformed, ugly",
    },
    {
        "name": "Cinematic",
        "prompt": "cinematic still {prompt}. emotional, harmonious, vignette, highly detailed, high budget, bokeh, cinemascope, moody, epic, gorgeous, film grain, grainy",
        "negative_prompt": "anime, cartoon, graphic, text, painting, crayon, graphite, abstract, glitch, deformed, mutated, ugly, disfigured",
    },
    {
        "name": "Anime",
        "prompt": "anime artwork {prompt}. anime style, key visual, vibrant, studio anime, highly detailed",
        "negative_prompt": "photo, deformed, black and white, realism, disfigured, low contrast",
    },
    {
        "name": "3D Model",
        "prompt": "professional 3d model {prompt}. octane render, highly detailed, volumetric, dramatic lighting",
        "negative_prompt": "ugly, deformed, noisy, low poly, blurry, painting",
    },
    {
        "name": "(No style)",
        "prompt": "{prompt}",
        "negative_prompt": "",
    },
]

styles = {k["name"]: (k["prompt"], k["negative_prompt"]) for k in style_list}
STYLE_NAMES = list(styles.keys())
DEFAULT_STYLE_NAME = "Photo"

def apply_style(style_name: str, positive: str, negative: str = "") -> Tuple[str, str]:
    p, n = styles.get(style_name, styles[DEFAULT_STYLE_NAME])
    if not negative:
        negative = ""
    return p.replace("{prompt}", positive), n + negative

DESCRIPTION = """## SDXL Image Generation

"""

if not torch.cuda.is_available():
    DESCRIPTION += "\n<p>⚠️Running on CPU, This may not work on CPU.</p>"

MAX_SEED = np.iinfo(np.int32).max
CACHE_EXAMPLES = torch.cuda.is_available() and os.getenv("CACHE_EXAMPLES", "0") == "1"
MAX_IMAGE_SIZE = int(os.getenv("MAX_IMAGE_SIZE", "2048"))
USE_TORCH_COMPILE = os.getenv("USE_TORCH_COMPILE", "0") == "1"
ENABLE_CPU_OFFLOAD = os.getenv("ENABLE_CPU_OFFLOAD", "0") == "1"

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

NUM_IMAGES_PER_PROMPT = 1

if torch.cuda.is_available():
    pipe = StableDiffusionXLPipeline.from_pretrained(
        "SG161222/RealVisXL_V5.0_Lightning",
        torch_dtype=torch.float16,
        use_safetensors=True,
        add_watermarker=False,
        variant="fp16"
    )
    pipe2 = StableDiffusionXLPipeline.from_pretrained(
        "SG161222/RealVisXL_V4.0_Lightning",
        torch_dtype=torch.float16,
        use_safetensors=True,
        add_watermarker=False,
        variant="fp16"
    )
    if ENABLE_CPU_OFFLOAD:
        pipe.enable_model_cpu_offload()
        pipe2.enable_model_cpu_offload()
    else:
        pipe.to(device)
        pipe2.to(device)
        print("Loaded on Device!")

    if USE_TORCH_COMPILE:
        pipe.unet = torch.compile(pipe.unet, mode="reduce-overhead", fullgraph=True)
        pipe2.unet = torch.compile(pipe2.unet, mode="reduce-overhead", fullgraph=True)
        print("Model Compiled!")

def save_image(img):
    unique_name = str(uuid.uuid4()) + ".png"
    img.save(unique_name)
    return unique_name

def randomize_seed_fn(seed: int, randomize_seed: bool) -> int:
    if randomize_seed:
        seed = random.randint(0, MAX_SEED)
    return seed

@spaces.GPU(duration=30)
@torch.no_grad()
def generate(
    prompt: str,
    negative_prompt: str = "",
    use_negative_prompt: bool = False,
    style: str = DEFAULT_STYLE_NAME,
    seed: int = 0,
    width: int = 1024,
    height: int = 1024,
    guidance_scale: float = 3,
    randomize_seed: bool = False,
    use_resolution_binning: bool = True, # This parameter is not exposed in the UI by default
    progress=gr.Progress(track_tqdm=True),
):
    if check_text(prompt, negative_prompt):
        raise ValueError("Prompt contains restricted words.")

    prompt, negative_prompt_from_style = apply_style(style, prompt, "") # Apply style positive first
    
    # Combine negative prompts
    if use_negative_prompt:
        final_negative_prompt = negative_prompt_from_style + " " + negative_prompt + " " + default_negative
    else:
        final_negative_prompt = negative_prompt_from_style + " " + default_negative
    final_negative_prompt = final_negative_prompt.strip()


    seed = int(randomize_seed_fn(seed, randomize_seed))
    generator = torch.Generator(device=device).manual_seed(seed) # Ensure generator is on the correct device

    options = {
        "prompt": prompt,
        "negative_prompt": final_negative_prompt,
        "width": width,
        "height": height,
        "guidance_scale": guidance_scale,
        "num_inference_steps": 25, # This is hardcoded, UI slider for steps is not connected
        "generator": generator,
        "num_images_per_prompt": NUM_IMAGES_PER_PROMPT, # UI slider for images is not connected to this
        # "use_resolution_binning": use_resolution_binning, # This was in original code, but not defined. Diffusers handles it.
        "output_type": "pil",
    }
    
    # If on CPU, ensure generator is for CPU
    if device.type == 'cpu':
        generator = torch.Generator(device='cpu').manual_seed(seed)
        options["generator"] = generator

    images = []
    if 'pipe' in globals(): # Check if pipes are loaded (i.e. on GPU)
         images.extend(pipe(**options).images)
         images.extend(pipe2(**options).images)
    else: # Fallback for CPU or if pipes are not loaded (though the DESCRIPTION warns about CPU)
        # This part would need a CPU-compatible pipeline if one isn't loaded.
        # For now, it will likely error if pipe/pipe2 aren't available.
        # Or, we can return a placeholder or raise a specific error.
        # To prevent errors if running without GPU and models didn't load:
        placeholder_image = Image.new('RGB', (width, height), color = 'grey')
        draw = ImageDraw.Draw(placeholder_image)
        draw.text((10, 10), "GPU models not loaded. Cannot generate image.", fill=(255,0,0))
        images.append(placeholder_image)


    image_paths = [save_image(img) for img in images]
    return image_paths, seed

examples = [
    "3d image, cute girl, in the style of Pixar --ar 1:2 --stylize 750, 4K resolution highlights, Sharp focus, octane render, ray tracing, Ultra-High-Definition, 8k, UHD, HDR, (Masterpiece:1.5), (best quality:1.5)",
    "A glass cup of cold coffee placed on a rustic wooden table, surrounded by soft morning light. The coffee is rich, dark, and topped with a light layer of creamy froth, droplets of condensation sliding down the glass.",
    "Vector illustration of a horse, vector graphic design with flat colors on an brown background in the style of vector art, using simple shapes and graphics with simple details, professionally designed as a tshirt logo ready for print on a white background. --ar 89:82 --v 6.0 --style raw",
    "Man in brown leather jacket posing for camera, in the style of sleek and stylized, clockpunk, subtle shades, exacting precision, ferrania p30  --ar 67:101 --v 5",
    "Commercial photography, giant burger, white lighting, studio light, 8k octane rendering, high resolution photography, insanely detailed, fine details, on white isolated plain, 8k, commercial photography, stock photo, professional color grading, --v 4 --ar 9:16 "
]

css = '''
.gradio-container {
    max-width: 590px !important; /* Existing style */
    margin: 0 auto !important; /* Existing style */
}
h1 {
    text-align: center; /* Existing style */
}
footer {
    visibility: hidden; /* Existing style */
}
'''
with gr.Blocks(css=css, theme="bethecloud/storj_theme") as demo:
    gr.Markdown(DESCRIPTION)
    with gr.Group():
        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.Gallery(label="Result", columns=1, preview=True) # columns=1 for single image below each other if multiple

    with gr.Accordion("Advanced options", open=False):
        style_selection = gr.Dropdown( # MODIFIED: Was gr.Radio, moved into accordion
            label="Image Style",
            choices=STYLE_NAMES,
            value=DEFAULT_STYLE_NAME,
            interactive=True,
            show_label=True,
            container=True,
        )
        use_negative_prompt = gr.Checkbox(label="Use negative prompt", value=True, visible=True)
        negative_prompt = gr.Text(
            label="Negative prompt",
            max_lines=1,
            placeholder="Enter a negative prompt (appended to style's negative)",
            value="(deformed iris, deformed pupils, semi-realistic, cgi, 3d, render, sketch, cartoon, drawing, anime:1.4), text, close up, cropped, out of frame, worst quality, low quality, jpeg artifacts, ugly, duplicate, morbid, mutilated, extra fingers, mutated hands, poorly drawn hands, poorly drawn face, mutation, deformed, blurry, dehydrated, bad anatomy, bad proportions, extra limbs, cloned face, disfigured, gross proportions, malformed limbs, missing arms, missing legs, extra arms, extra legs, fused fingers, too many fingers, long neck",
            visible=True,
        )
        # Note: num_inference_steps and num_images_per_prompt sliders are defined in UI
        # but not wired to the generate function's parameters that control these aspects.
        # Keeping them as is, per "Don't alter the remaining functionality".
        with gr.Row():
            num_inference_steps = gr.Slider( # This UI element is not connected to the backend
                label="Steps (Not Connected)",
                minimum=10,
                maximum=60,
                step=1,
                value=20, # Default value in UI
            )
        with gr.Row():
            num_images_per_prompt = gr.Slider( # This UI element is not connected to the backend
                label="Images (Not Connected)",
                minimum=1,
                maximum=4,
                step=1,
                value=2, # Default value in UI (backend NUM_IMAGES_PER_PROMPT is 1, resulting in 2 total)
            )
        seed = gr.Slider(
            label="Seed",
            minimum=0,
            maximum=MAX_SEED,
            step=1,
            value=0,
            visible=True
        )
        randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
        with gr.Row(visible=True):
            width = gr.Slider(
                label="Width",
                minimum=512,
                maximum=MAX_IMAGE_SIZE, # Use MAX_IMAGE_SIZE
                step=8,
                value=1024,
            )
            height = gr.Slider(
                label="Height",
                minimum=512,
                maximum=MAX_IMAGE_SIZE, # Use MAX_IMAGE_SIZE
                step=8,
                value=1024,
            )
        with gr.Row():
            guidance_scale = gr.Slider(
                label="Guidance Scale",
                minimum=0.1,
                maximum=20.0,
                step=0.1,
                value=3.0,
            )

    # Original style_selection gr.Row has been removed from here.

    gr.Examples(
        examples=examples,
        inputs=prompt,
        outputs=[result, seed], # seed output is good for reproducibility
        fn=generate,
        cache_examples=CACHE_EXAMPLES,
    )

    use_negative_prompt.change(
        fn=lambda x: gr.update(visible=x),
        inputs=use_negative_prompt,
        outputs=negative_prompt,
        api_name=False,
    )

    gr.on(
        triggers=[
            prompt.submit,
            negative_prompt.submit, # Allow submitting negative prompt to trigger run
            run_button.click,
        ],
        fn=generate,
        inputs=[
            prompt,
            negative_prompt,
            use_negative_prompt,
            style_selection, # style_selection is correctly in inputs
            seed,
            width,
            height,
            guidance_scale,
            randomize_seed,
        ],
        outputs=[result, seed],
        api_name="run",
    )

if __name__ == "__main__":
    # For CPU execution, model loading might take time or fail if not handled.
    # The `if torch.cuda.is_available():` block handles model loading for GPU.
    # A CPU fallback for inference would require a CPU-compatible model or different handling in `generate`.
    # The provided code primarily targets GPU.
    # Added a basic placeholder image generation in `generate` if pipes are not loaded.
    # Also need `ImageDraw` for that.
    from PIL import ImageDraw # Add ImageDraw import for CPU placeholder

    demo.queue(max_size=20).launch(ssr_mode=True, show_error=True, share=True)