ajsbsd's picture
Update app.py
4858646 verified
raw
history blame
12.8 kB
import gradio as gr
import torch
from diffusers import StableDiffusionXLPipeline, StableDiffusionXLImg2ImgPipeline
from PIL import Image
import os
import gc
import time
import spaces
from typing import Optional, Tuple
from huggingface_hub import hf_hub_download
# Global pipeline variables
txt2img_pipe = None
img2img_pipe = None
device = "cuda" if torch.cuda.is_available() else "cpu"
# Hugging Face model configuration
MODEL_REPO = "ajsbsd/CyberRealistic-Pony"
MODEL_FILENAME = "cyberrealisticPony_v110.safetensors"
def clear_memory():
"""Clear GPU memory"""
if torch.cuda.is_available():
torch.cuda.empty_cache()
gc.collect()
def load_models():
"""Load both text2img and img2img pipelines optimized for Spaces"""
global txt2img_pipe, img2img_pipe
try:
print("Loading CyberRealistic Pony models...")
# Download model file using huggingface_hub
print(f"Downloading model from {MODEL_REPO}...")
model_path = hf_hub_download(
repo_id=MODEL_REPO,
filename=MODEL_FILENAME,
cache_dir="/tmp/hf_cache" # Use tmp for Spaces
)
print(f"Model downloaded to: {model_path}")
# Load Text2Img pipeline
if txt2img_pipe is None:
txt2img_pipe = StableDiffusionXLPipeline.from_single_file(
model_path,
torch_dtype=torch.float16 if device == "cuda" else torch.float32,
use_safetensors=True,
variant="fp16" if device == "cuda" else None
)
# Aggressive memory optimizations for Spaces
txt2img_pipe.enable_attention_slicing()
txt2img_pipe.enable_vae_slicing()
if device == "cuda":
txt2img_pipe.enable_model_cpu_offload()
txt2img_pipe.enable_sequential_cpu_offload()
else:
txt2img_pipe = txt2img_pipe.to(device)
# Share components for Img2Img to save memory
if img2img_pipe is None:
img2img_pipe = StableDiffusionXLImg2ImgPipeline(
vae=txt2img_pipe.vae,
text_encoder=txt2img_pipe.text_encoder,
text_encoder_2=txt2img_pipe.text_encoder_2,
tokenizer=txt2img_pipe.tokenizer,
tokenizer_2=txt2img_pipe.tokenizer_2,
unet=txt2img_pipe.unet,
scheduler=txt2img_pipe.scheduler,
)
# Same optimizations
img2img_pipe.enable_attention_slicing()
img2img_pipe.enable_vae_slicing()
if device == "cuda":
img2img_pipe.enable_model_cpu_offload()
img2img_pipe.enable_sequential_cpu_offload()
print("Models loaded successfully!")
return True
except Exception as e:
print(f"Error loading models: {e}")
return False
def enhance_prompt(prompt: str, add_quality_tags: bool = True) -> str:
"""Enhance prompt with Pony-style tags"""
if not prompt.strip():
return prompt
if prompt.startswith("score_") or not add_quality_tags:
return prompt
quality_tags = "score_9, score_8_up, score_7_up, masterpiece, best quality, highly detailed"
return f"{quality_tags}, {prompt}"
def validate_dimensions(width: int, height: int) -> Tuple[int, int]:
"""Ensure dimensions are valid for SDXL"""
width = ((width + 63) // 64) * 64
height = ((height + 63) // 64) * 64
# More conservative limits for Spaces
width = max(512, min(1024, width))
height = max(512, min(1024, height))
return width, height
@spaces.GPU(duration=60) # GPU decorator for Spaces
def generate_txt2img(prompt, negative_prompt, num_steps, guidance_scale, width, height, seed, add_quality_tags):
"""Generate image from text prompt with Spaces GPU support"""
global txt2img_pipe
if not prompt.strip():
return None, "Please enter a prompt"
# Lazy load models
if txt2img_pipe is None:
if not load_models():
return None, "Failed to load models. Please try again."
try:
clear_memory()
# Validate dimensions
width, height = validate_dimensions(width, height)
# Set seed
generator = None
if seed != -1:
generator = torch.Generator(device=device).manual_seed(int(seed))
# Enhance prompt
enhanced_prompt = enhance_prompt(prompt, add_quality_tags)
print(f"Generating: {enhanced_prompt[:100]}...")
start_time = time.time()
# Generate with lower memory usage
with torch.no_grad():
result = txt2img_pipe(
prompt=enhanced_prompt,
negative_prompt=negative_prompt or "",
num_inference_steps=min(int(num_steps), 30), # Limit steps for Spaces
guidance_scale=float(guidance_scale),
width=width,
height=height,
generator=generator
)
generation_time = time.time() - start_time
status = f"Generated in {generation_time:.1f}s ({width}x{height})"
return result.images[0], status
except Exception as e:
return None, f"Generation failed: {str(e)}"
finally:
clear_memory()
@spaces.GPU(duration=60) # GPU decorator for Spaces
def generate_img2img(input_image, prompt, negative_prompt, num_steps, guidance_scale, strength, seed, add_quality_tags):
"""Generate image from input image + text prompt with Spaces GPU support"""
global img2img_pipe
if input_image is None:
return None, "Please upload an input image"
if not prompt.strip():
return None, "Please enter a prompt"
# Lazy load models
if img2img_pipe is None:
if not load_models():
return None, "Failed to load models. Please try again."
try:
clear_memory()
# Set seed
generator = None
if seed != -1:
generator = torch.Generator(device=device).manual_seed(int(seed))
# Enhance prompt
enhanced_prompt = enhance_prompt(prompt, add_quality_tags)
# Process input image
if isinstance(input_image, Image.Image):
if input_image.mode != 'RGB':
input_image = input_image.convert('RGB')
# Conservative resize for Spaces
max_size = 768
input_image.thumbnail((max_size, max_size), Image.Resampling.LANCZOS)
w, h = input_image.size
w, h = validate_dimensions(w, h)
input_image = input_image.resize((w, h), Image.Resampling.LANCZOS)
print(f"Transforming: {enhanced_prompt[:100]}...")
start_time = time.time()
with torch.no_grad():
result = img2img_pipe(
prompt=enhanced_prompt,
negative_prompt=negative_prompt or "",
image=input_image,
num_inference_steps=min(int(num_steps), 30), # Limit steps
guidance_scale=float(guidance_scale),
strength=float(strength),
generator=generator
)
generation_time = time.time() - start_time
status = f"Transformed in {generation_time:.1f}s (Strength: {strength})"
return result.images[0], status
except Exception as e:
return None, f"Transformation failed: {str(e)}"
finally:
clear_memory()
# Simplified negative prompt for better performance
DEFAULT_NEGATIVE = """
(low quality:1.3), (worst quality:1.3), (bad quality:1.2), blurry, noisy, ugly, deformed,
(text, watermark:1.4), (extra limbs:1.3), (bad hands:1.3), (bad anatomy:1.2)
"""
# Gradio interface optimized for Spaces
with gr.Blocks(
title="CyberRealistic Pony Generator",
theme=gr.themes.Soft()
) as demo:
gr.Markdown("""
# 🎨 CyberRealistic Pony Image Generator
Generate high-quality images using the CyberRealistic Pony SDXL model.
⚠️ **Note**: First generation may take longer as the model loads. GPU time is limited on Spaces.
""")
with gr.Tabs():
with gr.TabItem("🎨 Text to Image"):
with gr.Row():
with gr.Column():
txt2img_prompt = gr.Textbox(
label="Prompt",
placeholder="beautiful landscape, mountains, sunset",
lines=2
)
with gr.Accordion("Advanced Settings", open=False):
txt2img_negative = gr.Textbox(
label="Negative Prompt",
value=DEFAULT_NEGATIVE,
lines=2
)
txt2img_quality_tags = gr.Checkbox(
label="Add Quality Tags",
value=True
)
with gr.Row():
txt2img_steps = gr.Slider(10, 30, 20, step=1, label="Steps")
txt2img_guidance = gr.Slider(1.0, 15.0, 7.5, step=0.5, label="Guidance")
with gr.Row():
txt2img_width = gr.Slider(512, 1024, 768, step=64, label="Width")
txt2img_height = gr.Slider(512, 1024, 768, step=64, label="Height")
txt2img_seed = gr.Number(label="Seed (-1 for random)", value=-1, precision=0)
txt2img_btn = gr.Button("🎨 Generate", variant="primary", size="lg")
with gr.Column():
txt2img_output = gr.Image(label="Generated Image", height=400)
txt2img_status = gr.Textbox(label="Status", interactive=False)
with gr.TabItem("πŸ–ΌοΈ Image to Image"):
with gr.Row():
with gr.Column():
img2img_input = gr.Image(label="Input Image", type="pil", height=250)
img2img_prompt = gr.Textbox(
label="Prompt",
placeholder="digital painting style, vibrant colors",
lines=2
)
with gr.Accordion("Advanced Settings", open=False):
img2img_negative = gr.Textbox(
label="Negative Prompt",
value=DEFAULT_NEGATIVE,
lines=2
)
img2img_quality_tags = gr.Checkbox(
label="Add Quality Tags",
value=True
)
with gr.Row():
img2img_steps = gr.Slider(10, 30, 20, step=1, label="Steps")
img2img_guidance = gr.Slider(1.0, 15.0, 7.5, step=0.5, label="Guidance")
img2img_strength = gr.Slider(
0.1, 1.0, 0.75, step=0.05,
label="Strength (Higher = more creative)"
)
img2img_seed = gr.Number(label="Seed (-1 for random)", value=-1, precision=0)
img2img_btn = gr.Button("πŸ–ΌοΈ Transform", variant="primary", size="lg")
with gr.Column():
img2img_output = gr.Image(label="Generated Image", height=400)
img2img_status = gr.Textbox(label="Status", interactive=False)
# Event handlers
txt2img_btn.click(
fn=generate_txt2img,
inputs=[txt2img_prompt, txt2img_negative, txt2img_steps, txt2img_guidance,
txt2img_width, txt2img_height, txt2img_seed, txt2img_quality_tags],
outputs=[txt2img_output, txt2img_status]
)
img2img_btn.click(
fn=generate_img2img,
inputs=[img2img_input, img2img_prompt, img2img_negative, img2img_steps, img2img_guidance,
img2img_strength, img2img_seed, img2img_quality_tags],
outputs=[img2img_output, img2img_status]
)
print(f"πŸš€ CyberRealistic Pony Generator initialized on {device}")
if __name__ == "__main__":
demo.launch()