Spaces:
Sleeping
Sleeping
File size: 6,256 Bytes
f4aaf73 | 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 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 | import gradio as gr
import torch
from diffusers import StableDiffusionXLPipeline, EulerAncestralDiscreteScheduler
from PIL import Image
import random
import os
# Model selection based on art style
MODEL_MAP = {
"default": "stabilityai/stable-diffusion-xl-base-1.0",
"anime": "cagliostrolab/animagine-xl-3.1",
"pixel": "wavymulder/pixel-art-diffusion",
"cinematic": "stabilityai/stable-diffusion-3.5-medium",
}
class PerchanceStyleGenerator:
def __init__(self):
self.current_model = None
self.pipe = None
def load_model(self, model_key="default"):
if self.current_model != model_key:
model_id = MODEL_MAP.get(model_key, MODEL_MAP["default"])
print(f"Loading {model_id}...")
self.pipe = StableDiffusionXLPipeline.from_pretrained(
model_id,
torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
use_safetensors=True
)
self.pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(self.pipe.scheduler.config)
if torch.cuda.is_available():
self.pipe = self.pipe.to("cuda")
self.current_model = model_key
def generate(self, prompt, negative, style, shape, num_images):
# Map Perchance style to model
style_lower = style.lower()
if any(anime in style_lower for anime in ["anime", "manga", "waifu"]):
model_key = "anime"
elif "pixel" in style_lower:
model_key = "pixel"
elif any(cine in style_lower for cine in ["cinematic", "photo"]):
model_key = "cinematic"
else:
model_key = "default"
self.load_model(model_key)
# Calculate dimensions based on shape
if shape == "Square":
width, height = 1024, 1024
elif shape == "Portrait":
width, height = 832, 1216
else: # Landscape
width, height = 1216, 832
# Enhance prompt with style
style_prompt = f"{prompt}, {style.lower()} style, high quality, detailed"
# Generate images
images = []
for i in range(num_images):
generator = torch.Generator().manual_seed(random.randint(1, 999999))
result = self.pipe(
prompt=style_prompt,
negative_prompt=negative,
width=width,
height=height,
num_inference_steps=25,
guidance_scale=7.5,
generator=generator
)
images.append(result.images[0])
if len(images) == 1:
return images[0]
else:
return self.make_grid(images, cols=min(4, num_images))
def make_grid(self, images, cols=4):
rows = (len(images) + cols - 1) // cols
w, h = images[0].size
grid = Image.new('RGB', (cols * w, rows * h))
for i, img in enumerate(images):
grid.paste(img, (i % cols * w, i // cols * h))
return grid
generator = PerchanceStyleGenerator()
ART_STYLES = [
"Painted", "Anime", "Casual Photo", "Cinematic", "Digital Painting",
"Concept Art", "No style", "3D", "Disney Character", "2D Disney Character",
"Disney Sketch", "Concept Sketch", "Painterly", "Oil Painting",
"Oil Painting - Realism", "Oil Painting - Old", "Oil Painting - 70s Pulp",
"Professional Photo", "Anime Drawn", "Anime Anime", "Anime Screencap",
"Cute Anime", "Soft Anime", "Fantasy Painting", "Fantasy Landscape",
"Fantasy Portrait", "Studio Ghibli", "50s Enamel Sign", "Vintage Comic",
"Franco-Belgian Comic", "Tintin Comic", "Medieval", "Pixel Art",
"Furry - Oil", "Furry - Cinematic", "Furry - Painted", "Furry - Drawn",
"Cute Figurine", "3D Emoji", "Illustration", "Cute Illustration",
"Flat Illustration", "Watercolor", "1990s Photo", "1980s Photo",
"1970s Photo", "1960s Photo", "1950s Photo", "1940s Photo", "1930s Photo",
"1920s Photo", "Vintage Pulp Art", "50s Infomercial", "Anime 3D",
"Pokemon Painted", "2D Pokemon", "Vintage Anime", "Neon Vintage Anime",
"Manga", "Fantasy World Map", "Fantasy City Map", "Old World Map",
"3D Isometric", "Icon", "Flat Style Icon", "Flat Style Logo",
"Game Art Icon", "Digital Painting Icon", "Concept Art Icon",
"Cute 3D Icon", "Cute 3D Icon Set", "Crayon Drawing", "Pencil",
"Tattoo Design", "Waifu", "YuGiOh Art", "Traditional Japanese",
"Nihonga Painting", "Claymation", "Cartoon", "Cursed Photo", "MTG Card"
]
with gr.Blocks(theme=gr.themes.Soft(), title="ImageMagic") as demo:
gr.Markdown("# π ImageMagic")
gr.Markdown("Like Perchance, but self-hosted on Hugging Face!")
with gr.Row():
with gr.Column(scale=1):
prompt = gr.Textbox(
label="π Description",
placeholder="Describe what you want to see...",
lines=3
)
negative = gr.Textbox(
label="π« Anti-Description (optional)",
placeholder="What you DON'T want...",
lines=2
)
style = gr.Dropdown(
label="π¨ Art Style",
choices=ART_STYLES,
value="Cinematic"
)
with gr.Row():
shape = gr.Radio(
label="πΌοΈ Shape",
choices=["Square", "Portrait", "Landscape"],
value="Square"
)
num_images = gr.Dropdown(
label="π’ How many?",
choices=[1, 2, 4, 6, 8],
value=1
)
generate_btn = gr.Button("β¨ Generate", variant="primary", size="lg")
with gr.Column(scale=1):
output = gr.Image(label="Result", height=500)
generate_btn.click(
fn=generator.generate,
inputs=[prompt, negative, style, shape, num_images],
outputs=output
)
if __name__ == "__main__":
demo.launch()
|