app.py

import gradio as gr
import numpy as np
import random
from diffusers import StableDiffusionXLPipeline, KDPM2AncestralDiscreteScheduler, AutoencoderKL
import torch

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

# Load VAE component
vae = AutoencoderKL.from_pretrained(
    "madebyollin/sdxl-vae-fp16-fix", 
    torch_dtype=torch.float16 if device == "cuda" else torch.float32
)

# Configure the pipeline
pipe = StableDiffusionXLPipeline.from_pretrained(
    "dataautogpt3/ProteusV0.5", 
    vae=vae,
    torch_dtype=torch.float16 if device == "cuda" else torch.float32
)
pipe.scheduler = KDPM2AncestralDiscreteScheduler.from_config(pipe.scheduler.config)
pipe.to(device)

MAX_SEED = np.iinfo(np.int32).max
MAX_IMAGE_SIZE = 1024

def infer(prompt, negative_prompt="", seed=0, randomize_seed=True, width=1024, height=1024, guidance_scale=7.0, num_inference_steps=50, clip_skip=2):
    if randomize_seed:
        seed = random.randint(0, MAX_SEED)
    
    generator = torch.Generator(device=device).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,
        clip_skip=clip_skip
    ).images[0] 
    
    return image

examples = [
    ["Astronaut in a jungle, cold color palette, muted colors, detailed, 8k"],
    ["An astronaut riding a green horse"],
    ["A delicious ceviche cheesecake slice"],
    ["black fluffy gorgeous dangerous cat animal creature, large orange eyes, big fluffy ears, piercing gaze, full moon, dark ambiance, best quality, extremely detailed"],
    ["high quality pixel art, a pixel art silhouette of an anime space-themed girl in a space-punk steampunk style, lying in her bed by the window of a spaceship, smoking, with a rustic feel. The image should embody epic portraiture and double exposure, featuring an isolated landscape visible through the window. The colors should primarily be dynamic and action-packed, with a strong use of negative space. The entire artwork should be in pixel art style, emphasizing the characters shape and set against a white background. Silhouette"]
]

css = """
#col-container { max-width: 700px; margin: 0 auto; }
"""

demo = gr.Blocks(css=css)

with demo:
    gr.Markdown(f"# ProteusV0.5 Demo")
    gr.Markdown(f"Currently running on {'GPU' if torch.cuda.is_available() else 'CPU'}.")
    
    gr.Markdown("""
    ProteusV0.5 is a state-of-the-art text-to-image model that leverages the power of Stable Diffusion to generate high-quality images from text prompts. This model has been fine-tuned on a large dataset of images and has been trained to understand a wide range of prompts and styles.
    
    With ProteusV0.5, you can generate images in a variety of styles, from realistic to abstract, and from simple to complex. The model is also capable of understanding and responding to complex prompts, making it a powerful tool for artists, designers, and anyone looking to generate high-quality images.
    
    To use this demo, simply enter a text prompt in the input field below, and the model will generate an image based on your prompt. You can also adjust the settings to control the quality and style of the generated image.
    """)
    
    with gr.Column(elem_id="col-container"):
        prompt = gr.Text(label="Prompt", placeholder="Enter your prompt")
        run_button = gr.Button("Generate Image")
        result = gr.Image(label="Generated Image")

        with gr.Accordion("Advanced Settings", open=False):
            negative_prompt = gr.Text(label="Negative prompt", placeholder="Enter a negative prompt")
            seed = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=0)
            randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
            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)
            guidance_scale = gr.Slider(label="Guidance scale", minimum=0.0, maximum=20.0, step=0.1, value=7.0)
            num_inference_steps = gr.Slider(label="Number of inference steps", minimum=1, maximum=100, step=1, value=50)
            clip_skip = gr.Slider(label="Clip skip", minimum=1, maximum=12, step=1, value=2)

        gr.Examples(
            examples=examples,
            inputs=prompt,
            outputs=result,
            fn=infer,
            cache_examples=True
        )

    run_button.click(
        fn=infer,
        inputs=[prompt, negative_prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps, clip_skip],
        outputs=result
    )

demo.queue().launch()