Update

app.py

@@ -1,153 +1,471 @@
+import os
+import gc
 import gradio as gr
 import numpy as np
-import random
-
-# import spaces #[uncomment to use ZeroGPU]
-from diffusers import DiffusionPipeline
 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
+from config import (
+    MODEL,
+    MIN_IMAGE_SIZE,
+    MAX_IMAGE_SIZE,
+    USE_TORCH_COMPILE,
+    ENABLE_CPU_OFFLOAD,
+    OUTPUT_DIR,
+    DEFAULT_NEGATIVE_PROMPT,
+    DEFAULT_ASPECT_RATIO,
+    examples,
+    sampler_list,
+    aspect_ratios,
+    style_list,
+)
+import time
+from typing import List, Dict, Tuple, Optional
 
-device = "cuda" if torch.cuda.is_available() else "cpu"
-model_repo_id = "huchukato/pimp-my-pony"  # Replace to the model you would like to use
+# Enhanced logging configuration
+logging.basicConfig(
+    level=logging.INFO,
+    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
+    datefmt='%Y-%m-%d %H:%M:%S'
+)
+logger = logging.getLogger(__name__)
 
-if torch.cuda.is_available():
-    torch_dtype = torch.float16
-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
-
-# @spaces.GPU #[uncomment to use ZeroGPU]
-def infer(
-    prompt,
-    negative_prompt,
-    seed,
-    randomize_seed,
-    width,
-    height,
-    guidance_scale,
-    num_inference_steps,
-    progress=gr.Progress(track_tqdm=True),
-):
-    if randomize_seed:
-        seed = random.randint(0, MAX_SEED)
-
-    generator = torch.Generator().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,
-    ).images[0]
-
-    return image, seed
-
-
-examples = [
-    "source_9, source_8_up, source_7_up, source_6_up, in a loft, dramatic lighting, hard shadows,  1girl,  2b_\(nier:automata\), hairband, white hair, blindfold, covered eyes, black blindfold sleeveless dress, large breasts",
-    "source_9, source_8_up, source_7_up, source_6_up, in a japanese house, at night, dramatic lighting, hard shadows, 1girl, gareth \(fate\), fate/grand order large breasts, solo, cleavage, covered nipples, arm under breast",
-    "source_9, source_8_up, source_7_up, source_6_up, in a greek temple, at night, dramatic lighting, hard shadows, 1girl, kido saori, light purple hair, very long hair, bangs, aqua eyes, sleeveless dress, long dress, puffy dress, large breasts",
-]
-
-css = """
-#col-container {
-    margin: 0 auto;
-    max-width: 640px;
-}
-"""
-
-with gr.Blocks(css=css) as demo:
-    with gr.Column(elem_id="col-container"):
-        gr.Markdown(" # Text-to-Image PimpMyPony")
-
-        with gr.Row():
-            prompt = gr.Text(
-                label="Prompt",
-                show_label=False,
-                max_lines=2,
-                placeholder="Enter your prompt",
-                container=False,
-            )
+# Constants
+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" 
 
-            run_button = gr.Button("Run", scale=0, variant="primary")
+# PyTorch settings for better performance and determinism
+torch.backends.cudnn.deterministic = True
+torch.backends.cudnn.benchmark = False
+torch.backends.cuda.matmul.allow_tf32 = True
 
-        result = gr.Image(label="Result", show_label=False)
+device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+logger.info(f"Using device: {device}")
 
-        with gr.Accordion("Advanced Settings", open=False):
-            negative_prompt = gr.Text(
-                label="Negative prompt",
-                max_lines=1,
-                placeholder="Enter a negative prompt",
-                visible=False,
-            )
+class GenerationError(Exception):
+    """Custom exception for generation errors"""
+    pass
 
-            seed = gr.Slider(
-                label="Seed",
-                minimum=0,
-                maximum=MAX_SEED,
-                step=1,
-                value=0,
-            )
+def validate_prompt(prompt: str) -> str:
+    """Validate and clean up the input prompt."""
+    if not isinstance(prompt, str):
+        raise GenerationError("Prompt must be a string")
+    try:
+        # Ensure proper UTF-8 encoding/decoding
+        prompt = prompt.encode('utf-8').decode('utf-8')
+        # Add space between ! and ,
+        prompt = prompt.replace("!,", "! ,")
+    except UnicodeError:
+        raise GenerationError("Invalid characters in prompt")
+    
+    # Only check if the prompt is completely empty or only whitespace
+    if not prompt or prompt.isspace():
+        raise GenerationError("Prompt cannot be empty")
+    return prompt.strip()
 
-            randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
+def validate_dimensions(width: int, height: int) -> None:
+    """Validate image dimensions."""
+    if not MIN_IMAGE_SIZE <= width <= MAX_IMAGE_SIZE:
+        raise GenerationError(f"Width must be between {MIN_IMAGE_SIZE} and {MAX_IMAGE_SIZE}")
+        
+    if not MIN_IMAGE_SIZE <= height <= MAX_IMAGE_SIZE:
+        raise GenerationError(f"Height must be between {MIN_IMAGE_SIZE} and {MAX_IMAGE_SIZE}")
 
-            with gr.Row():
-                width = gr.Slider(
-                    label="Width",
-                    minimum=256,
-                    maximum=MAX_IMAGE_SIZE,
-                    step=32,
-                    value=832,  # Replace with defaults that work for your model
-                )
+@spaces.GPU
+def generate(
+    prompt: str,
+    negative_prompt: str = DEFAULT_NEGATIVE_PROMPT,
+    seed: int = 0,
+    custom_width: int = 1024,
+    custom_height: int = 1024,
+    guidance_scale: float = 6.0,
+    num_inference_steps: int = 25,
+    sampler: str = "Euler a",
+    aspect_ratio_selector: str = DEFAULT_ASPECT_RATIO,
+    style_selector: str = "(None)",
+    use_upscaler: bool = False,
+    upscaler_strength: float = 0.55,
+    upscale_by: float = 1.5,
+    add_quality_tags: bool = True,
+    progress: gr.Progress = gr.Progress(track_tqdm=True),
+) -> Tuple[List[str], Dict]:
+    """Generate images based on the given parameters."""
+    start_time = time.time()
+    upscaler_pipe = None
+    backup_scheduler = None
+    
+    try:
+        # Memory management
+        torch.cuda.empty_cache()
+        gc.collect()
 
-                height = gr.Slider(
-                    label="Height",
-                    minimum=256,
-                    maximum=MAX_IMAGE_SIZE,
-                    step=32,
-                    value=1216,  # Replace with defaults that work for your model
-                )
+        # Input validation
+        prompt = validate_prompt(prompt)
+        if negative_prompt:
+            negative_prompt = negative_prompt.encode('utf-8').decode('utf-8')
+        
+        validate_dimensions(custom_width, custom_height)
+        
+        # Set up generation
+        generator = utils.seed_everything(seed)
+        width, height = utils.aspect_ratio_handler(
+            aspect_ratio_selector,
+            custom_width,
+            custom_height,
+        )
 
-            with gr.Row():
-                guidance_scale = gr.Slider(
-                    label="Guidance scale",
-                    minimum=0.0,
-                    maximum=10.0,
-                    step=0.1,
-                    value=6,  # Replace with defaults that work for your model
-                )
+        # Process prompts
+        if add_quality_tags:
+            prompt = "masterpiece, high score, great score, absurdres, {prompt}".format(prompt=prompt)
+
+        prompt, negative_prompt = utils.preprocess_prompt(
+            styles, style_selector, prompt, negative_prompt
+        )    
+
+        width, height = utils.preprocess_image_dimensions(width, height)
+
+        # Set up pipeline
+        backup_scheduler = pipe.scheduler
+        pipe.scheduler = utils.get_scheduler(pipe.scheduler.config, sampler)
+
+        if use_upscaler:
+            upscaler_pipe = StableDiffusionXLImg2ImgPipeline(**pipe.components)
+            
+        # Prepare metadata
+        metadata = {
+            "prompt": prompt,
+            "negative_prompt": negative_prompt,
+            "resolution": f"{width} x {height}",
+            "guidance_scale": guidance_scale,
+            "num_inference_steps": num_inference_steps,
+            "style_preset": style_selector,
+            "seed": seed,
+            "sampler": sampler,
+            "Model": "PimpMyPony",
+            "Model hash": "e3c47aedb0",
+        }
+
+        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(f"Starting generation with parameters: {json.dumps(metadata, indent=4)}")
+
+        # Generate images
+        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
+
+        # Save images
+        if images:
+            total = len(images)
+            image_paths = []
+            for idx, image in enumerate(images, 1):
+                progress(idx/total, desc="Saving images...")
+                path = utils.save_image(image, metadata, OUTPUT_DIR, IS_COLAB)
+                image_paths.append(path)
+                logger.info(f"Image {idx}/{total} saved as {path}")
 
-                num_inference_steps = gr.Slider(
-                    label="Number of inference steps",
-                    minimum=1,
-                    maximum=50,
-                    step=1,
-                    value=20,  # Replace with defaults that work for your model
+        generation_time = time.time() - start_time
+        logger.info(f"Generation completed successfully in {generation_time:.2f} seconds")
+        metadata["generation_time"] = f"{generation_time:.2f}s"
+        
+        return image_paths, metadata
+
+    except GenerationError as e:
+        logger.warning(f"Generation validation error: {str(e)}")
+        raise gr.Error(str(e))
+    except Exception as e:
+        logger.exception("Unexpected error during generation")
+        raise gr.Error(f"Generation failed: {str(e)}")
+    finally:
+        # Cleanup
+        torch.cuda.empty_cache()
+        gc.collect()
+        
+        if upscaler_pipe is not None:
+            del upscaler_pipe
+        
+        if backup_scheduler is not None and pipe is not None:
+            pipe.scheduler = backup_scheduler
+            
+        utils.free_memory()
+
+# Model initialization
+if torch.cuda.is_available():
+    try:
+        logger.info("Loading VAE and pipeline...")
+        vae = AutoencoderKL.from_pretrained(
+            "madebyollin/sdxl-vae-fp16-fix",
+            torch_dtype=torch.float16,
+        )
+        pipe = utils.load_pipeline(MODEL, device, vae=vae)
+        logger.info("Pipeline loaded successfully on GPU!")
+    except Exception as e:
+        logger.error(f"Error loading VAE, falling back to default: {e}")
+        pipe = utils.load_pipeline(MODEL, device)
+else:
+    logger.warning("CUDA not available, running on CPU")
+    pipe = None
+
+# Process styles
+styles = {k["name"]: (k["prompt"], k["negative_prompt"]) for k in style_list}
+
+with gr.Blocks(css="style.css", theme="Nymbo/Nymbo_Theme_5") as demo:
+    gr.HTML(
+        """
+        <div class="header">
+            <div class="title">Pimp My Pony</div>
+        </div>
+        """,
+    )
+    
+    with gr.Row():
+        with gr.Column(scale=2):
+            with gr.Group():
+                prompt = gr.Text(
+                    label="Prompt",
+                    max_lines=5,
+                    placeholder="Describe what you want to generate",
+                    info="Enter your image generation prompt here. Be specific and descriptive for better results.",
+                )
+                negative_prompt = gr.Text(
+                    label="Negative Prompt",
+                    max_lines=5,
+                    placeholder="Describe what you want to avoid",
+                    value=DEFAULT_NEGATIVE_PROMPT,
+                    info="Specify elements you don't want in the image.",
                 )
+                add_quality_tags = gr.Checkbox(
+                    label="Quality Tags",
+                    value=True,
+                    info="Add quality-enhancing tags to your prompt automatically.",
+                )
+            with gr.Accordion(label="More Settings", open=False):
+                with gr.Group():
+                    aspect_ratio_selector = gr.Radio(
+                        label="Aspect Ratio",
+                        choices=aspect_ratios,
+                        value=DEFAULT_ASPECT_RATIO,
+                        container=True,
+                        info="Choose the dimensions of your image.",
+                    )
+                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,
+                            info=f"Image width (must be between {MIN_IMAGE_SIZE} and {MAX_IMAGE_SIZE})",
+                        )
+                        custom_height = gr.Slider(
+                            label="Height",
+                            minimum=MIN_IMAGE_SIZE,
+                            maximum=MAX_IMAGE_SIZE,
+                            step=8,
+                            value=1024,
+                            info=f"Image height (must be between {MIN_IMAGE_SIZE} and {MAX_IMAGE_SIZE})",
+                        )
+                with gr.Group():
+                    use_upscaler = gr.Checkbox(
+                        label="Use Upscaler",
+                        value=False,
+                        info="Enable high-resolution upscaling.",
+                    )
+                    with gr.Row() as upscaler_row:
+                        upscaler_strength = gr.Slider(
+                            label="Strength",
+                            minimum=0,
+                            maximum=1,
+                            step=0.05,
+                            value=0.55,
+                            visible=False,
+                            info="Control how much the upscaler affects the final image.",
+                        )
+                        upscale_by = gr.Slider(
+                            label="Upscale by",
+                            minimum=1,
+                            maximum=1.5,
+                            step=0.1,
+                            value=1.5,
+                            visible=False,
+                            info="Multiplier for the final image resolution.",
+                        )
+                with gr.Accordion(label="Advanced Parameters", open=False):
+                    with gr.Group():
+                        style_selector = gr.Dropdown(
+                            label="Style Preset",
+                            interactive=True,
+                            choices=list(styles.keys()),
+                            value="(None)",
+                            info="Apply a predefined style to your generation.",
+                        )
+                    with gr.Group():
+                        sampler = gr.Dropdown(
+                            label="Sampler",
+                            choices=sampler_list,
+                            interactive=True,
+                            value="Euler a",
+                            info="Different samplers can produce varying results.",
+                        )
+                    with gr.Group():
+                        seed = gr.Slider(
+                            label="Seed",
+                            minimum=0,
+                            maximum=utils.MAX_SEED,
+                            step=1,
+                            value=0,
+                            info="Set a specific seed for reproducible results.",
+                        )
+                        randomize_seed = gr.Checkbox(
+                            label="Randomize seed",
+                            value=True,
+                            info="Generate a new random seed for each image.",
+                        )
+                    with gr.Group():
+                        with gr.Row():
+                            guidance_scale = gr.Slider(
+                                label="Guidance scale",
+                                minimum=1,
+                                maximum=12,
+                                step=0.1,
+                                value=6.0,
+                                info="Higher values make the image more closely match your prompt.",
+                            )
+                            num_inference_steps = gr.Slider(
+                                label="Number of inference steps",
+                                minimum=1,
+                                maximum=50,
+                                step=1,
+                                value=25,
+                                info="More steps generally mean higher quality but slower generation.",
+                            )
+        
+        with gr.Column(scale=3):
+            with gr.Blocks():
+                run_button = gr.Button("Generate", variant="primary", elem_id="generate-button")
+            result = gr.Gallery(
+                label="Generated Images",
+                columns=1,
+                height='768px',
+                preview=True,
+                show_label=True,
+            )
+            with gr.Accordion(label="Generation Parameters", open=False):
+                gr_metadata = gr.JSON(
+                    label="Image Metadata",
+                    show_label=True,
+                )
+            gr.Examples(
+                examples=examples,
+                inputs=prompt,
+                outputs=[result, gr_metadata],
+                fn=lambda *args, **kwargs: generate(*args, use_upscaler=True, **kwargs),
+                cache_examples=CACHE_EXAMPLES,
+            )
+    
 
-        gr.Examples(examples=examples, inputs=[prompt])
+    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,
+    )
+
+    # Combine all triggers including keyboard shortcuts
     gr.on(
-        triggers=[run_button.click, prompt.submit],
-        fn=infer,
+        triggers=[
+            prompt.submit,
+            negative_prompt.submit,
+            run_button.click,
+        ],
+        fn=utils.randomize_seed_fn,
+        inputs=[seed, randomize_seed],
+        outputs=seed,
+        queue=False,
+        api_name=False,
+    ).then(
+        fn=lambda: gr.update(interactive=False, value="Generating..."),
+        outputs=run_button,
+    ).then(
+        fn=generate,
         inputs=[
             prompt,
             negative_prompt,
             seed,
-            randomize_seed,
-            width,
-            height,
+            custom_width,
+            custom_height,
             guidance_scale,
             num_inference_steps,
+            sampler,
+            aspect_ratio_selector,
+            style_selector,
+            use_upscaler,
+            upscaler_strength,
+            upscale_by,
+            add_quality_tags,
         ],
-        outputs=[result, seed],
+        outputs=[result, gr_metadata],
+    ).then(
+        fn=lambda: gr.update(interactive=True, value="Generate"),
+        outputs=run_button,
     )
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.queue(max_size=20).launch(debug=IS_COLAB, share=IS_COLAB)

config.py

@@ -0,0 +1,33 @@
+import os
+import tomli
+from typing import Dict, Any
+
+def fix_escaping(text: str) -> str:
+    # When JSON is loaded, \\\\ becomes \\ automatically
+    # So we don't need to do any transformation
+    return text
+
+def load_config() -> Dict[str, Any]:
+    config_path = os.path.join(os.path.dirname(__file__), 'config.toml')
+    with open(config_path, 'rb') as f:
+        config = tomli.load(f)
+        return config
+
+# Load configuration
+config = load_config()
+
+# Export variables for backward compatibility
+MODEL = os.getenv("MODEL", config['model']['path'])
+MIN_IMAGE_SIZE = int(os.getenv("MIN_IMAGE_SIZE", config['model']['min_image_size']))
+MAX_IMAGE_SIZE = int(os.getenv("MAX_IMAGE_SIZE", config['model']['max_image_size']))
+USE_TORCH_COMPILE = os.getenv("USE_TORCH_COMPILE", str(config['model']['use_torch_compile'])).lower() == "true"
+ENABLE_CPU_OFFLOAD = os.getenv("ENABLE_CPU_OFFLOAD", str(config['model']['enable_cpu_offload'])).lower() == "true"
+OUTPUT_DIR = os.getenv("OUTPUT_DIR", config['model']['output_dir'])
+
+DEFAULT_NEGATIVE_PROMPT = config['prompts']['default_negative']
+DEFAULT_ASPECT_RATIO = config['prompts']['default_aspect_ratio']
+
+examples = config['prompts']['examples']
+sampler_list = config['samplers']['list']
+aspect_ratios = config['aspect_ratios']['list']
+style_list = config['styles']

config.toml

@@ -0,0 +1,114 @@
+[model]
+path = "huchukato/pimp-my-pony"
+min_image_size = 512
+max_image_size = 2048
+use_torch_compile = false
+enable_cpu_offload = false
+output_dir = "./outputs"
+
+[prompts]
+default_negative = "score_5, score_4, score_3, worst quality, low quality, ugly, malformed, bad anatomy, grayscale"
+default_aspect_ratio = "832 x 1216"
+examples = [
+    "source_9, source_8_up, source_7_up, 1girl, souryuu asuka langley, neon genesis evangelion, eyepatch, red plugsuit, sitting, on throne, crossed legs, head tilt, holding weapon, lance of longinus \\\\(evangelion\\\\), cowboy shot, depth of field, faux traditional media, painterly, impressionism, photo background",
+    "source_9, source_8_up, source_7_up, 1boy, vash the stampede, trigun stampede, red jacket, sunglasses, gun, hand on own hip, aiming, standing, looking at viewer, upper body, desert, cliff, cowboy shot",
+    "source_9, source_8_up, source_7_up, 1girl, vertin \\(reverse:1999\\), reverse:1999, black umbrella, headwear, suitcase, looking at viewer, rain, night, city, bridge, from side, dutch angle, upper body",
+    "source_9, source_8_up, source_7_up, 1girl, 1boy, c.c., lelouch vi britannia, year 2024, code geass, ascot, bare shoulders, black choker, black hair, blue flower, blue rose, breasts, bush, choker, closed mouth, collarbone, couple, dress, flower, frills, green hair, hetero, long hair, long sleeves, looking at another, medium breasts, off-shoulder dress, off shoulder, parted lips, purple ascot, purple eyes, rose, short hair, sitting, straight hair, yellow eyes",
+    "source_9, source_8_up, source_7_up, 1girl, hatsune miku, vocaloid, blue eyes, blue hair, bowl, can, chopsticks, collared shirt, detached sleeves, eating, elbow rest, fish \\(food\\), food, holding, holding chopsticks, katsudon \\(food\\), long hair, long sleeves, looking at viewer, meal, nail polish, necktie, noodles, onigiri, plate, ramen, sashimi, shirt, shrimp, shrimp tempura, sleeveless, sleeveless shirt, solo, spring onion, tempura, twintails",
+    "source_9, source_8_up, source_7_up, 4girls, multiple girls, gotoh hitori, ijichi nijika, kita ikuyo, yamada ryo, bocchi the rock!,  ahoge, black shirt, blank eyes, blonde hair, blue eyes, blue hair, brown sweater, collared shirt, cube hair ornament, detached ahoge, empty eyes, green eyes, hair ornament, hairclip, kessoku band, long sleeves, looking at viewer, medium hair, mole, mole under eye, one side up, pink hair, pink track suit, red eyes, red hair, sailor collar, school uniform, serafuku, shirt, shuka high school uniform, side ahoge, side ponytail, sweater, sweater vest, track suit, white shirt, yellow eyes, painterly, impressionism, faux traditional media, v, double v, waving",
+    "source_9, source_8_up, source_7_up, 1other, solo, outdoors, sky, arm up, night, earth, helmet, outstretched arm, star \\(sky\\), night sky, full moon, floating, starry sky, reaching, jumping, space, cowboy shot, ambiguous gender, spacesuit, moonlight, space helmet, astronaut, horror, black and white, monochromatic, high contrast, abstract background, dutch angle, dark, depth of field, chromatic aberration, faux traditional media"
+]
+
+[samplers]
+list = [
+    "DPM++ 2M Karras",
+    "DPM++ SDE Karras",
+    "DPM++ 2M SDE Karras",
+    "Euler",
+    "Euler a",
+    "DDIM"
+]
+
+[aspect_ratios]
+list = [
+    "1024 x 1024",
+    "1152 x 896",
+    "896 x 1152",
+    "1216 x 832",
+    "832 x 1216",
+    "1344 x 768",
+    "768 x 1344",
+    "1536 x 640",
+    "640 x 1536",
+    "Custom"
+]
+
+[[styles]]
+name = "(None)"
+prompt = "{prompt}"
+negative_prompt = ""
+
+[[styles]]
+name = "Anim4gine"
+prompt = "{prompt}, depth of field, faux traditional media, painterly, impressionism, photo background"
+negative_prompt = ""
+
+[[styles]]
+name = "Painting"
+prompt = "{prompt}, painterly, painting (medium)"
+negative_prompt = ""
+
+[[styles]]
+name = "Pixel art"
+prompt = "{prompt}, pixel art"
+negative_prompt = ""
+
+[[styles]]
+name = "1980s"
+prompt = "{prompt}, 1980s (style), retro artstyle"
+negative_prompt = ""
+
+[[styles]]
+name = "1990s"
+prompt = "{prompt}, 1990s (style), retro artstyle"
+negative_prompt = ""
+
+[[styles]]
+name = "2000s"
+prompt = "{prompt}, 2000s (style), retro artstyle"
+negative_prompt = ""
+
+[[styles]]
+name = "Toon"
+prompt = "{prompt}, toon (style)"
+negative_prompt = ""
+
+[[styles]]
+name = "Lineart"
+prompt = "{prompt}, lineart, thick lineart"
+negative_prompt = ""
+
+[[styles]]
+name = "Art Nouveau"
+prompt = "{prompt}, art nouveau"
+negative_prompt = ""
+
+[[styles]]
+name = "Western Comics"
+prompt = "{prompt}, western comics (style)"
+negative_prompt = ""
+
+[[styles]]
+name = "3D"
+prompt = "{prompt}, 3d"
+negative_prompt = ""
+
+[[styles]]
+name = "Realistic"
+prompt = "{prompt}, realistic, photorealistic"
+negative_prompt = ""
+
+[[styles]]
+name = "Neonpunk"
+prompt = "{prompt}, neonpunk"
+negative_prompt = "" 

requirements.txt

@@ -1,6 +1,8 @@
-accelerate
-diffusers
-invisible_watermark
-torch
-transformers
-xformers
+accelerate>=1.2.1
+diffusers>=0.32.1
+gradio==4.44.1
+hf-transfer>=0.1.9
+spaces>=0.32.0
+torch>=2.4.0
+transformers>=4.48.0
+tomli>=2.0.1

style.css

@@ -0,0 +1,212 @@
+:root {
+  --title-font-size: clamp(1.5rem, 6vw, 3rem);
+  --subtitle-font-size: clamp(1rem, 2vw, 1.2rem);
+  --text-color: #fff;
+  --font-family: 'Helvetica Neue', sans-serif;
+  --gradient-primary: linear-gradient(45deg, #4EACEF, #b428b2);
+  --primary-color: #b61383;
+  --primary-hover: #910f63;
+  --discord-color: #910ed2;
+  --discord-hover: #600962;
+  --border-radius: 12px;
+  --box-shadow: 0 2px 6px rgba(0, 0, 0, 0.05);
+}
+
+body {
+  font-family: var(--font-family);
+  color: var(--text-color);
+  margin: 0;
+  padding: 0;
+  min-height: 100vh;
+  background-color: #f5f5f5;
+}
+
+.header {
+  text-align: center;
+  padding: 1rem 0;
+  margin-bottom: 0.5rem;
+}
+
+.title {
+  font-size: var(--title-font-size);
+  font-weight: 700;
+  text-transform: uppercase;
+  margin-bottom: 0.25rem;
+  background-image: var(--gradient-primary);
+  -webkit-text-fill-color: transparent;
+  -webkit-background-clip: text;
+  background-clip: text;
+  display: inline-block;
+}
+
+.subtitle {
+  font-size: var(--subtitle-font-size);
+  color: #999;
+  margin-bottom: 0.5rem;
+}
+
+.status {
+  display: none;
+}
+
+#duplicate-button {
+  margin: 1.5rem auto 2.5rem;
+  color: #fff;
+  background: #1565c0;
+  border-radius: 100vh;
+  padding: 0.75rem 1.5rem;
+  font-weight: 500;
+  box-shadow: 0 2px 8px rgba(21, 101, 192, 0.25);
+  transition: all 0.2s ease;
+  display: block;
+}
+
+#duplicate-button:hover {
+  background: #1976d2;
+  box-shadow: 0 4px 12px rgba(21, 101, 192, 0.35);
+  transform: translateY(-1px);
+}
+
+.contain {
+  max-width: 80%;
+  margin: 2rem auto;
+  padding: 2rem 1.5rem;
+}
+
+/* Component styling */
+.gr-box {
+  border-radius: var(--border-radius);
+  border: 1px solid #e0e0e0;
+  background: #ffffff;
+  box-shadow: var(--box-shadow);
+  transition: box-shadow 0.2s ease;
+}
+
+.gr-box:hover {
+  box-shadow: 0 4px 12px rgba(0, 0, 0, 0.1);
+}
+
+.gr-button.primary {
+  background: var(--primary-color);
+  border-radius: var(--border-radius);
+  padding: 0.8rem 2rem;
+  font-weight: 500;
+  box-shadow: 0 2px 8px rgba(21, 101, 192, 0.25);
+  transition: all 0.2s ease;
+  text-transform: uppercase;
+  letter-spacing: 0.5px;
+}
+
+.gr-button.primary:hover {
+  background: var(--primary-hover);
+  box-shadow: 0 4px 12px rgba(21, 101, 192, 0.35);
+  transform: translateY(-1px);
+}
+
+/* Form elements */
+.gr-form {
+  background: #fff;
+  padding: 1.5rem;
+  border-radius: var(--border-radius);
+  box-shadow: var(--box-shadow);
+}
+
+.gr-input, .gr-textarea {
+  border: 1px solid #e0e0e0;
+  border-radius: 8px;
+  padding: 0.8rem;
+  transition: all 0.2s ease;
+}
+
+.gr-input:focus, .gr-textarea:focus {
+  border-color: var(--primary-color);
+  box-shadow: 0 0 0 2px rgba(21, 101, 192, 0.1);
+}
+
+/* Accordion styling */
+.gr-accordion {
+  border: none;
+  margin: 1rem 0;
+}
+
+.gr-accordion-header {
+  background: #f8f9fa;
+  border-radius: var(--border-radius);
+  padding: 1rem;
+  font-weight: 500;
+}
+
+/* Gallery styling */
+.gr-gallery {
+  background: #fff;
+  padding: 1rem;
+  border-radius: var(--border-radius);
+  box-shadow: var(--box-shadow);
+}
+
+/* Discord button */
+.discord-btn {
+  display: inline-flex;
+  align-items: center;
+  justify-content: center;
+  background-color: var(--discord-color);
+  color: white !important;
+  text-decoration: none;
+  padding: 12px 24px;
+  border-radius: var(--border-radius);
+  transition: all 0.3s ease;
+  margin-top: 1rem;
+  font-size: 16px;
+  font-weight: 500;
+  width: 100%;
+  border: none;
+  cursor: pointer;
+  box-shadow: 0 2px 8px rgba(88, 101, 242, 0.25);
+}
+
+.discord-btn:hover {
+  background-color: var(--discord-hover);
+  transform: translateY(-2px);
+  box-shadow: 0 4px 12px rgba(88, 101, 242, 0.4);
+}
+
+.discord-icon {
+  width: 24px;
+  height: 24px;
+  margin-right: 12px;
+}
+
+.discord-text {
+  letter-spacing: 0.5px;
+}
+
+/* Tooltips */
+.gr-form small {
+  color: #666;
+  font-size: 0.875rem;
+  margin-top: 0.25rem;
+  display: block;
+}
+
+/* Responsive layout */
+@media (max-width: 768px) {
+  .contain {
+    max-width: 90%;
+    padding: 1rem;
+  }
+  
+  .gr-box {
+    margin: 0.5rem 0;
+  }
+  
+  .gr-button.primary {
+    width: 100%;
+  }
+}
+
+@media (min-width: 1200px) {
+  .contain {
+    max-width: 1400px;
+    padding: 2.5rem 2rem;
+  }
+}

utils.py

@@ -0,0 +1,187 @@
+import gc
+import os
+import random
+import numpy as np
+import json
+import torch
+import uuid
+from PIL import Image, PngImagePlugin
+from datetime import datetime
+from dataclasses import dataclass
+from typing import Callable, Dict, Optional, Tuple, Any, List
+from diffusers import (
+    DDIMScheduler,
+    DPMSolverMultistepScheduler,
+    DPMSolverSinglestepScheduler,
+    EulerAncestralDiscreteScheduler,
+    EulerDiscreteScheduler,
+    AutoencoderKL,
+    StableDiffusionXLPipeline,
+)
+import logging
+
+MAX_SEED = np.iinfo(np.int32).max
+
+
+@dataclass
+class StyleConfig:
+    prompt: str
+    negative_prompt: str
+
+
+def randomize_seed_fn(seed: int, randomize_seed: bool) -> int:
+    if randomize_seed:
+        seed = random.randint(0, MAX_SEED)
+    return seed
+
+
+def seed_everything(seed: int) -> torch.Generator:
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed_all(seed)
+    np.random.seed(seed)
+    generator = torch.Generator()
+    generator.manual_seed(seed)
+    return generator
+
+
+def parse_aspect_ratio(aspect_ratio: str) -> Optional[Tuple[int, int]]:
+    if aspect_ratio == "Custom":
+        return None
+    width, height = aspect_ratio.split(" x ")
+    return int(width), int(height)
+
+
+def aspect_ratio_handler(
+    aspect_ratio: str, custom_width: int, custom_height: int
+) -> Tuple[int, int]:
+    if aspect_ratio == "Custom":
+        return custom_width, custom_height
+    else:
+        width, height = parse_aspect_ratio(aspect_ratio)
+        return width, height
+
+
+def get_scheduler(scheduler_config: Dict, name: str) -> Optional[Callable]:
+    scheduler_factory_map = {
+        "DPM++ 2M Karras": lambda: DPMSolverMultistepScheduler.from_config(
+            scheduler_config, use_karras_sigmas=True
+        ),
+        "DPM++ SDE Karras": lambda: DPMSolverSinglestepScheduler.from_config(
+            scheduler_config, use_karras_sigmas=True
+        ),
+        "DPM++ 2M SDE Karras": lambda: DPMSolverMultistepScheduler.from_config(
+            scheduler_config, use_karras_sigmas=True, algorithm_type="sde-dpmsolver++"
+        ),
+        "Euler": lambda: EulerDiscreteScheduler.from_config(scheduler_config),
+        "Euler a": lambda: EulerAncestralDiscreteScheduler.from_config(
+            scheduler_config
+        ),
+        "DDIM": lambda: DDIMScheduler.from_config(scheduler_config),
+    }
+    return scheduler_factory_map.get(name, lambda: None)()
+
+
+def free_memory() -> None:
+    """Free up GPU and system memory."""
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
+        torch.cuda.ipc_collect()
+    gc.collect()
+
+
+def preprocess_prompt(
+    style_dict,
+    style_name: str,
+    positive: str,
+    negative: str = "",
+    add_style: bool = True,
+) -> Tuple[str, str]:
+    p, n = style_dict.get(style_name, style_dict["(None)"])
+
+    if add_style and positive.strip():
+        formatted_positive = p.format(prompt=positive)
+    else:
+        formatted_positive = positive
+
+    combined_negative = n
+    if negative.strip():
+        if combined_negative:
+            combined_negative += ", " + negative
+        else:
+            combined_negative = negative
+
+    return formatted_positive, combined_negative
+
+
+def common_upscale(
+    samples: torch.Tensor,
+    width: int,
+    height: int,
+    upscale_method: str,
+) -> torch.Tensor:
+    return torch.nn.functional.interpolate(
+        samples, size=(height, width), mode=upscale_method
+    )
+
+
+def upscale(
+    samples: torch.Tensor, upscale_method: str, scale_by: float
+) -> torch.Tensor:
+    width = round(samples.shape[3] * scale_by)
+    height = round(samples.shape[2] * scale_by)
+    return common_upscale(samples, width, height, upscale_method)
+
+
+def preprocess_image_dimensions(width, height):
+    if width % 8 != 0:
+        width = width - (width % 8)
+    if height % 8 != 0:
+        height = height - (height % 8)
+    return width, height
+
+
+def save_image(image, metadata, output_dir, is_colab):
+    if is_colab:
+        current_time = datetime.now().strftime("%Y%m%d_%H%M%S")
+        filename = f"image_{current_time}.png"   
+    else:
+        filename = str(uuid.uuid4()) + ".png"
+    os.makedirs(output_dir, exist_ok=True)
+    filepath = os.path.join(output_dir, filename)
+    metadata_str = json.dumps(metadata)
+    info = PngImagePlugin.PngInfo()
+    info.add_text("parameters", metadata_str)
+    image.save(filepath, "PNG", pnginfo=info)
+    return filepath
+    
+    
+def is_google_colab():
+    try:
+        import google.colab
+        return True
+    except:
+        return False
+
+
+def load_pipeline(model_name: str, device: torch.device, hf_token: Optional[str] = None, vae: Optional[AutoencoderKL] = None) -> Any:
+    """Load the Stable Diffusion pipeline."""
+    try:
+        pipeline = (
+            StableDiffusionXLPipeline.from_single_file
+            if model_name.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
+        )
+        pipe.to(device)
+        return pipe
+    except Exception as e:
+        logging.error(f"Failed to load pipeline: {str(e)}", exc_info=True)
+        raise