Spaces:
Runtime error
Runtime error
File size: 4,072 Bytes
7c5ef3c 90ec6ed 7c5ef3c 90ec6ed 7c5ef3c 90ec6ed 7c5ef3c 0aeb023 7c5ef3c 0aeb023 7c5ef3c 0aeb023 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 |
import gradio as gr
import torch
from diffusers import StableDiffusionPipeline
from huggingface_hub import hf_hub_download
import numpy as np
from PIL import Image
import os
# Suppress symlink warnings
os.environ['HF_HUB_DISABLE_SYMLINKS_WARNING'] = "1"
# Define styles
styles = {
"glitch": {
"concept_url": "sd-concepts-library/001glitch-core",
"seed": 42,
"token": "<glitch-core>"
},
"roth": {
"concept_url": "sd-concepts-library/2814-roth",
"seed": 123,
"token": "<2814-roth>"
},
"night": {
"concept_url": "sd-concepts-library/4tnght",
"seed": 456,
"token": "<4tnght>"
},
"anime80s": {
"concept_url": "sd-concepts-library/80s-anime-ai",
"seed": 789,
"token": "<80s-anime>"
},
"animeai": {
"concept_url": "sd-concepts-library/80s-anime-ai-being",
"seed": 1024,
"token": "<80s-anime-being>"
}
}
def load_pipeline():
"""Load and prepare the pipeline with all style embeddings"""
# Check if CUDA is available
device = "cuda" if torch.cuda.is_available() else "cpu"
dtype = torch.float16 if device == "cuda" else torch.float32
pipe = StableDiffusionPipeline.from_pretrained(
"CompVis/stable-diffusion-v1-4",
torch_dtype=dtype
).to(device)
# Load all embeddings
for style_info in styles.values():
embedding_path = hf_hub_download(
repo_id=style_info["concept_url"],
filename="learned_embeds.bin",
repo_type="model"
)
pipe.load_textual_inversion(embedding_path)
return pipe
def apply_purple_guidance(image, strength=0.5):
"""Apply purple guidance to an image"""
img_array = np.array(image).astype(float)
purple_mask = (img_array[:,:,0] > 100) & (img_array[:,:,2] > 100)
img_array[purple_mask] = img_array[purple_mask] * (1 - strength) + np.array([128, 0, 128]) * strength
return Image.fromarray(np.uint8(img_array.clip(0, 255)))
def generate_image(prompt, style, seed, apply_guidance, guidance_strength=0.5):
"""Generate an image with selected style and optional purple guidance"""
if style not in styles:
return None
# Get style info
style_info = styles[style]
# Prepare generator with appropriate device
device = "cuda" if torch.cuda.is_available() else "cpu"
generator = torch.Generator(device).manual_seed(int(seed))
# Create styled prompt
styled_prompt = f"{prompt} {style_info['token']}"
# Generate image
image = pipe(
styled_prompt,
generator=generator,
guidance_scale=7.5,
num_inference_steps=20
).images[0]
# Apply purple guidance if requested
if apply_guidance:
image = apply_purple_guidance(image, guidance_strength)
return image
# Initialize the pipeline globally
print("Loading pipeline and embeddings...")
pipe = load_pipeline()
# Create the Gradio interface
demo = gr.Interface(
fn=generate_image,
inputs=[
gr.Textbox(label="Prompt", value="A serene mountain landscape with a lake at sunset"),
gr.Radio(choices=list(styles.keys()), label="Style", value="glitch"),
gr.Number(label="Seed", value=42),
gr.Checkbox(label="Apply Purple Guidance", value=False),
gr.Slider(minimum=0.0, maximum=1.0, value=0.5, label="Purple Guidance Strength")
],
outputs=gr.Image(label="Generated Image"),
title="Style-Guided Image Generation with Purple Enhancement",
description="""Generate images in different styles with optional purple color guidance.
Choose a style, enter a prompt, and optionally apply purple color enhancement.""",
examples=[
["A serene mountain landscape with a lake at sunset", "glitch", 42, True, 0.5],
["A magical forest at twilight", "anime80s", 789, True, 0.7],
["A cyberpunk city at night", "night", 456, False, 0.5],
],
cache_examples=True
)
if __name__ == "__main__":
demo.launch(share=False) |