import os import gc import gradio as gr import numpy as np import torch import json import spaces import random 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 import time from typing import List, Dict, Tuple, Optional from config import ( MODEL, MIN_IMAGE_SIZE, MAX_IMAGE_SIZE, DEFAULT_PROMPT, DEFAULT_NEGATIVE_PROMPT, scheduler_list, ) import io MAX_SEED = np.iinfo(np.int32).max # 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__) # PyTorch settings for better performance and determinism torch.backends.cudnn.deterministic = True torch.backends.cudnn.benchmark = False torch.backends.cuda.matmul.allow_tf32 = True device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") logger.info(f"Using device: {device}") # 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 class GenerationError(Exception): """Custom exception for generation errors""" pass 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() 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}") progress=gr.Progress() @spaces.GPU def generate( prompt: str, negative_prompt: str, width: int, height: int, scheduler: str, opt_strength:float, opt_scale:float, seed: int, randomize_seed: bool, guidance_scale: float, num_inference_steps: int ): progress(0,desc="Starting") if randomize_seed: seed = random.randint(0, MAX_SEED) """Generate images based on the given parameters.""" upscaler_pipe = None backup_scheduler = None def callback1(pipe, step, timestep, callback_kwargs): progress_value = 0.1 + ((step+1.0)/num_inference_steps)*(0.5/1.0) progress(progress_value, desc=f"Image generating, {step + 1}/{num_inference_steps} steps") return callback_kwargs optimizing_steps = int(num_inference_steps * opt_strength) def callback2(pipe, step, timestep, callback_kwargs): progress_value = 0.6 + ((step+1.0)/optimizing_steps)*(0.4/1.0) progress(progress_value, desc=f"Image optimizing, {step + 1}/{optimizing_steps} steps") return callback_kwargs try: # Memory management torch.cuda.empty_cache() gc.collect() # Input validation prompt = validate_prompt(prompt) if negative_prompt: negative_prompt = negative_prompt.encode('utf-8').decode('utf-8') validate_dimensions(width, height) # Set up generation generator = utils.seed_everything(seed) width, height = utils.preprocess_image_dimensions(width, height) # Set up pipeline backup_scheduler = pipe.scheduler pipe.scheduler = utils.get_scheduler(pipe.scheduler.config, scheduler) upscaler_pipe = StableDiffusionXLImg2ImgPipeline(**pipe.components) progress(0.1,desc="Image generating") 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", callback_on_step_end=callback1 ).images upscaled_latents = utils.upscale(latents, "nearest-exact", opt_scale) images = upscaler_pipe( prompt=prompt, negative_prompt=negative_prompt, image=upscaled_latents, guidance_scale=guidance_scale, num_inference_steps=num_inference_steps, strength=opt_strength, generator=generator, output_type="pil", callback_on_step_end=callback2 ).images out_img = images[0] path = utils.save_image(out_img, "./outputs") logger.info(f"output path: {path}") progress(1, desc="Complete") return path 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() title = "# Animagine XL 4.0 Demo" custom_css = """ #row-container { align-items: stretch; } #output-image{ flex-grow: 1; } #output-image *{ max-height: none !important; } """ with gr.Blocks(css=custom_css).queue() as demo: gr.Markdown(title) with gr.Row( elem_id="row-container" ): with gr.Column(): gr.Markdown("### Input") with gr.Column(): prompt = gr.Text( label="Prompt", max_lines=5, placeholder="Enter your prompt", value=DEFAULT_PROMPT, ) negative_prompt = gr.Text( label="Negative prompt", max_lines=5, placeholder="Enter a negative prompt", value=DEFAULT_NEGATIVE_PROMPT, ) with gr.Row(): width = gr.Slider( label="Width", minimum=MIN_IMAGE_SIZE, maximum=MAX_IMAGE_SIZE, step=8, value=832, ) height = gr.Slider( label="Height", minimum=MIN_IMAGE_SIZE, maximum=MAX_IMAGE_SIZE, step=8, value=1216, ) with gr.Row(): optimization_strength = gr.Slider( label="Optimization strength", minimum=0, maximum=1, step=0.05, value=0.55, ) optimization_scale = gr.Slider( label="Optimization scale ratio", minimum=1, maximum=1.5, step=0.1, value=1.5, ) with gr.Column(): scheduler = gr.Dropdown( label="scheduler", choices=scheduler_list, interactive=True, value="Euler a", ) with gr.Column(): seed = gr.Slider( label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=0, ) randomize_seed = gr.Checkbox(label="Randomize seed", value=True) with gr.Row(): guidance_scale = gr.Slider( label="Guidance scale", minimum=1.0, maximum=12.0, step=0.1, value=6.0, ) num_inference_steps = gr.Slider( label="Number of inference steps", minimum=1, maximum=50, step=1, value=25, ) run_button = gr.Button("Run", variant="primary") with gr.Column(): gr.Markdown("### Output") result = gr.Image( type="filepath", label="Generated Image", elem_id="output-image" ) run_button.click( fn=generate, inputs=[ prompt, negative_prompt, width, height, scheduler, optimization_strength,optimization_scale, seed,randomize_seed, guidance_scale,num_inference_steps ], outputs=[result], ) if __name__ == "__main__": demo.queue(max_size=20).launch()