app.py

import os
import gc
import gradio as gr
import numpy as np
import torch
import json
import spaces
import config
import utils
import logging
from PIL import Image, PngImagePlugin
from datetime import datetime
from diffusers.models import AutoencoderKL
from diffusers import StableDiffusionXLPipeline, StableDiffusionXLImg2ImgPipeline

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


if not torch.cuda.is_available():
    DESCRIPTION += "\n<p>Running on CPU 🥶 This demo does not work on CPU. </p>"
IS_COLAB = utils.is_google_colab() or os.getenv("IS_COLAB") == "1"
HF_TOKEN = os.getenv("HF_TOKEN")
CACHE_EXAMPLES = torch.cuda.is_available() and os.getenv("CACHE_EXAMPLES") == "1"
MIN_IMAGE_SIZE = int(os.getenv("MIN_IMAGE_SIZE", "512"))
MAX_IMAGE_SIZE = int(os.getenv("MAX_IMAGE_SIZE", "2048"))
USE_TORCH_COMPILE = os.getenv("USE_TORCH_COMPILE") == "1"
ENABLE_CPU_OFFLOAD = os.getenv("ENABLE_CPU_OFFLOAD") == "1"
OUTPUT_DIR = os.getenv("OUTPUT_DIR", "./outputs")

MODEL = os.getenv(
    "MODEL",
    "https://huggingface.co/AstraliteHeart/pony-diffusion-v6/blob/main/v6.safetensors",
)

torch.backends.cudnn.deterministic = True
torch.backends.cudnn.benchmark = False

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

def load_pipeline(model_name):
    vae = AutoencoderKL.from_pretrained(
        "madebyollin/sdxl-vae-fp16-fix",
        torch_dtype=torch.float16,
    )
    pipeline = (
        StableDiffusionXLPipeline.from_single_file
        if MODEL.endswith(".safetensors")
        else StableDiffusionXLPipeline.from_pretrained
    )

    pipe = pipeline(
        model_name,
        vae=vae,
        torch_dtype=torch.float16,
        custom_pipeline="lpw_stable_diffusion_xl",
        use_safetensors=True,
        add_watermarker=False,
        use_auth_token=HF_TOKEN,
        variant="fp16",
    )

    pipe.to(device)
    return pipe

def parse_json_parameters(json_str):
    try:
        params = json.loads(json_str)
        required_keys = ['prompt', 'negative_prompt', 'resolution', 'guidance_scale', 'num_inference_steps', 'seed', 'sampler']
        for key in required_keys:
            if key not in params:
                raise ValueError(f"Missing required key: {key}")
        
        # Parse resolution
        width, height = map(int, params['resolution'].split(' x '))
        
        return {
            'prompt': params['prompt'],
            'negative_prompt': params['negative_prompt'],
            'seed': params['seed'],
            'width': width,
            'height': height,
            'guidance_scale': params['guidance_scale'],
            'num_inference_steps': params['num_inference_steps'],
            'sampler': params['sampler'],
            'use_upscaler': params.get('use_upscaler', False)
        }
    except json.JSONDecodeError:
        raise ValueError("Invalid JSON format")
    except Exception as e:
        raise ValueError(f"Error parsing JSON: {str(e)}")

@spaces.GPU
def generate(
    prompt: str,
    negative_prompt: str = "",
    seed: int = 0,
    custom_width: int = 1024,
    custom_height: int = 1024,
    guidance_scale: float = 7.0,
    num_inference_steps: int = 30,
    sampler: str = "DPM++ 2M SDE Karras",
    aspect_ratio_selector: str = "1024 x 1024",
    use_upscaler: bool = False,
    upscaler_strength: float = 0.55,
    upscale_by: float = 1.5,
    json_params: str = "",
    progress=gr.Progress(track_tqdm=True),
) -> Image:
    if json_params:
        try:
            params = parse_json_parameters(json_params)
            prompt = params['prompt']
            negative_prompt = params['negative_prompt']
            seed = params['seed']
            custom_width = params['width']
            custom_height = params['height']
            guidance_scale = params['guidance_scale']
            num_inference_steps = params['num_inference_steps']
            sampler = params['sampler']
            use_upscaler = params['use_upscaler']
        except ValueError as e:
            raise gr.Error(str(e))

    generator = utils.seed_everything(seed)

    width, height = utils.aspect_ratio_handler(
        aspect_ratio_selector,
        custom_width,
        custom_height,
    )

    width, height = utils.preprocess_image_dimensions(width, height)

    backup_scheduler = pipe.scheduler
    pipe.scheduler = utils.get_scheduler(pipe.scheduler.config, sampler)

    if use_upscaler:
        upscaler_pipe = StableDiffusionXLImg2ImgPipeline(**pipe.components)
    metadata = {
        "prompt": prompt,
        "negative_prompt": negative_prompt,
        "resolution": f"{width} x {height}",
        "guidance_scale": guidance_scale,
        "num_inference_steps": num_inference_steps,
        "seed": seed,
        "sampler": sampler,
    }

    if use_upscaler:
        new_width = int(width * upscale_by)
        new_height = int(height * upscale_by)
        metadata["use_upscaler"] = {
            "upscale_method": "nearest-exact",
            "upscaler_strength": upscaler_strength,
            "upscale_by": upscale_by,
            "new_resolution": f"{new_width} x {new_height}",
        }
    else:
        metadata["use_upscaler"] = None
    logger.info(json.dumps(metadata, indent=4))

    try:
        if use_upscaler:
            latents = pipe(
                prompt=prompt,
                negative_prompt=negative_prompt,
                width=width,
                height=height,
                guidance_scale=guidance_scale,
                num_inference_steps=num_inference_steps,
                generator=generator,
                output_type="latent",
            ).images
            upscaled_latents = utils.upscale(latents, "nearest-exact", upscale_by)
            images = upscaler_pipe(
                prompt=prompt,
                negative_prompt=negative_prompt,
                image=upscaled_latents,
                guidance_scale=guidance_scale,
                num_inference_steps=num_inference_steps,
                strength=upscaler_strength,
                generator=generator,
                output_type="pil",
            ).images
        else:
            images = pipe(
                prompt=prompt,
                negative_prompt=negative_prompt,
                width=width,
                height=height,
                guidance_scale=guidance_scale,
                num_inference_steps=num_inference_steps,
                generator=generator,
                output_type="pil",
            ).images

        if images and IS_COLAB:
            for image in images:
                filepath = utils.save_image(image, metadata, OUTPUT_DIR)
                logger.info(f"Image saved as {filepath} with metadata")

        return images, metadata
    except Exception as e:
        logger.exception(f"An error occurred: {e}")
        raise
    finally:
        if use_upscaler:
            del upscaler_pipe
        pipe.scheduler = backup_scheduler
        utils.free_memory()

# Initialize an empty list to store the generation history
generation_history = []

# Function to update the history list
def update_history_list():
    return [item["image"] for item in generation_history]

# Function to handle image click in history
def handle_image_click(evt: gr.SelectData):
    selected = generation_history[evt.index]
    return selected["image"], json.dumps(selected["metadata"], indent=2)

# Modify the generate function to add results to the history
def generate_and_update_history(*args, **kwargs):
    images, metadata = generate(*args, **kwargs)
    timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    generation_history.insert(0, {
        "prompt": metadata["prompt"],
        "timestamp": timestamp,
        "image": images[0],
        "metadata": metadata
    })
    if len(generation_history) > 10:  # Limit history to 10 items
        generation_history.pop()
    return images[0], json.dumps(metadata, indent=2), update_history_list()

if torch.cuda.is_available():
    pipe = load_pipeline(MODEL)
    logger.info("Loaded on Device!")
else:
    pipe = None

with gr.Blocks(css="style.css") as demo:
    gr.HTML("""
        <header id="site-header">
            <nav>
                <a href="https://sergidev.me">
                    <img
                        src="https://myamber.cloud/life/v1/file?query=%7B%22token%22%3A%227f2b2ba3ef1651d299c9d5a683e8b68c%22%2C%22root%22%3A%22cloud%22%2C%22path%22%3A%22%2Fshared-to%2Ff8b4f187-ec556f6f%22%7D"
                        alt="Logo"
                    />
                </a>
                <h1>HDiffusion</h1>
                <div class="mode-switch">
                    <img
                        src="https://icon-library.com/images/window-icon-png/window-icon-png-0.jpg"
                        id="projects-icon"
                        alt="Projects"
                    />
                    <img
                        src="https://creazilla-store.fra1.digitaloceanspaces.com/icons/3206440/terminal-icon-sm.png"
                        id="terminal-icon"
                        alt="Terminal"
                    />
                </div>
            </nav>
        </header>
    """)
    
    gr.Markdown(
        f"""Gradio demo for [Pony Diffusion V6](https://civitai.com/models/257749/pony-diffusion-v6-xl/) with image gallery, json support, and advanced options.""",
        elem_id="subtitle",
    )

    gr.DuplicateButton(
        value="Duplicate Space for private use",
        elem_id="duplicate-button",
        visible=os.getenv("SHOW_DUPLICATE_BUTTON") == "1",
    )
    with gr.Group():
        with gr.Row():
            prompt = gr.Text(
                label="Prompt",
                show_label=False,
                max_lines=5,
                placeholder="Enter your prompt",
                container=False,
            )
            run_button = gr.Button(
                "Generate", 
                variant="primary", 
                scale=0
            )
        result = gr.Image(
            label="Result", 
            show_label=False
        )
    with gr.Accordion(label="Advanced Settings", open=False):
        negative_prompt = gr.Text(
            label="Negative Prompt",
            max_lines=5,
            placeholder="Enter a negative prompt",
            value=""
        )
        aspect_ratio_selector = gr.Radio(
            label="Aspect Ratio",
            choices=config.aspect_ratios,
            value="1024 x 1024",
            container=True,
        )
        with gr.Group(visible=False) as custom_resolution:
            with gr.Row():
                custom_width = gr.Slider(
                    label="Width",
                    minimum=MIN_IMAGE_SIZE,
                    maximum=MAX_IMAGE_SIZE,
                    step=8,
                    value=1024,
                )
                custom_height = gr.Slider(
                    label="Height",
                    minimum=MIN_IMAGE_SIZE,
                    maximum=MAX_IMAGE_SIZE,
                    step=8,
                    value=1024,
                )
        use_upscaler = gr.Checkbox(label="Use Upscaler", value=False)
        with gr.Row() as upscaler_row:
            upscaler_strength = gr.Slider(
                label="Strength",
                minimum=0,
                maximum=1,
                step=0.05,
                value=0.55,
                visible=False,
            )
            upscale_by = gr.Slider(
                label="Upscale by",
                minimum=1,
                maximum=1.5,
                step=0.1,
                value=1.5,
                visible=False,
            )
        sampler = gr.Dropdown(
            label="Sampler",
            choices=config.sampler_list,
            interactive=True,
            value="DPM++ 2M SDE Karras",
        )
        with gr.Row():
            seed = gr.Slider(
                label="Seed", minimum=0, maximum=utils.MAX_SEED, step=1, value=0
            )
            randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
        with gr.Group():
            with gr.Row():
                guidance_scale = gr.Slider(
                    label="Guidance scale",
                    minimum=1,
                    maximum=12,
                    step=0.1,
                    value=7.0,
                )
                num_inference_steps = gr.Slider(
                    label="Number of inference steps",
                    minimum=1,
                    maximum=50,
                    step=1,
                    value=28,
                )
    with gr.Accordion(label="Generation Parameters", open=False):
        gr_metadata = gr.JSON(label="Metadata", show_label=False)
        json_input = gr.TextArea(label="Edit/Paste JSON Parameters", placeholder="Paste or edit JSON parameters here")
        generate_from_json = gr.Button("Generate from JSON")

    with gr.Accordion("Generation History", open=False) as history_accordion:
        history_gallery = gr.Gallery(
            label="History",
            show_label=False,
            elem_id="history_gallery",
            columns=5,
            rows=2,
            height="auto"
        )
        with gr.Row():
            selected_image = gr.Image(label="Selected Image", interactive=False)
            selected_metadata = gr.JSON(label="Selected Metadata", show_label=False)

    gr.Examples(
        examples=config.examples,
        inputs=prompt,
        outputs=[result, gr_metadata],
        fn=lambda *args, **kwargs: generate_and_update_history(*args, use_upscaler=True, **kwargs),
        cache_examples=CACHE_EXAMPLES,
    )

    use_upscaler.change(
        fn=lambda x: [gr.update(visible=x), gr.update(visible=x)],
        inputs=use_upscaler,
        outputs=[upscaler_strength, upscale_by],
        queue=False,
        api_name=False,
    )
    
    aspect_ratio_selector.change(
        fn=lambda x: gr.update(visible=x == "Custom"),
        inputs=aspect_ratio_selector,
        outputs=custom_resolution,
        queue=False,
        api_name=False,
    )

    inputs = [
        prompt,
        negative_prompt,
        seed,
        custom_width,
        custom_height,
        guidance_scale,
        num_inference_steps,
        sampler,
        aspect_ratio_selector,
        use_upscaler,
        upscaler_strength,
        upscale_by,
        json_input,
    ]

    prompt.submit(
        fn=utils.randomize_seed_fn,
        inputs=[seed, randomize_seed],
        outputs=seed,
        queue=False,
        api_name=False,
    ).then(
        fn=generate_and_update_history,
        inputs=inputs,
        outputs=[result, gr_metadata, history_gallery],
    )
    
    negative_prompt.submit(
        fn=utils.randomize_seed_fn,
        inputs=[seed, randomize_seed],
        outputs=seed,
        queue=False,
        api_name=False,
    ).then(
        fn=generate_and_update_history,
        inputs=inputs,
        outputs=[result, gr_metadata, history_gallery],
    )
    
    run_button.click(
        fn=utils.randomize_seed_fn,
        inputs=[seed, randomize_seed],
        outputs=seed,
        queue=False,
        api_name=False,
    ).then(
        fn=generate_and_update_history,
        inputs=inputs,
        outputs=[result, gr_metadata, history_gallery],
    )
    
    generate_from_json.click(
        fn=generate_and_update_history,
        inputs=inputs,
        outputs=[result, gr_metadata, history_gallery],
    )
    
demo.queue(max_size=20).launch(debug=IS_COLAB, share=IS_COLAB)