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 import gc # 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": "" }, "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>" } } # Pre-generate example images example_images = { "glitch": "examples/glitch_example.jpg", "anime80s": "examples/anime80s_example.jpg", "night": "examples/night_example.jpg" } 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 # Use smaller model for CPU model_id = "runwayml/stable-diffusion-v1-5" if device == "cuda" else "CompVis/stable-diffusion-v1-4" pipe = StableDiffusionPipeline.from_pretrained( model_id, torch_dtype=dtype, low_cpu_mem_usage=True ).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""" # Check if this is one of our examples with pre-generated images if prompt == "A serene mountain landscape with a lake at sunset" and style == "glitch" and seed == 42: if os.path.exists(example_images["glitch"]): image = Image.open(example_images["glitch"]) if apply_guidance: image = apply_purple_guidance(image, guidance_strength) return image if prompt == "A magical forest at twilight" and style == "anime80s" and seed == 789: if os.path.exists(example_images["anime80s"]): image = Image.open(example_images["anime80s"]) if apply_guidance: image = apply_purple_guidance(image, guidance_strength) return image if prompt == "A cyberpunk city at night" and style == "night" and seed == 456: if os.path.exists(example_images["night"]): image = Image.open(example_images["night"]) if apply_guidance: image = apply_purple_guidance(image, guidance_strength) return image 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 with reduced settings for CPU if device == "cpu": # Use much smaller image size and fewer steps on CPU image = pipe( styled_prompt, generator=generator, guidance_scale=7.5, num_inference_steps=10, # Reduced steps height=256, # Smaller height width=256 # Smaller width ).images[0] else: 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) # Clean up memory gc.collect() if device == "cuda": torch.cuda.empty_cache() return image # Create examples directory os.makedirs("examples", exist_ok=True) # 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. Note: Generation may take a few minutes on CPU.""", 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, allow_flagging="never" # Disable flagging to reduce overhead ) if __name__ == "__main__": # Generate and save example images if they don't exist if not all(os.path.exists(path) for path in example_images.values()): print("Pre-generating example images...") # Example 1 if not os.path.exists(example_images["glitch"]): img = generate_image("A serene mountain landscape with a lake at sunset", "glitch", 42, False, 0.5) img.save(example_images["glitch"]) # Example 2 if not os.path.exists(example_images["anime80s"]): img = generate_image("A magical forest at twilight", "anime80s", 789, False, 0.7) img.save(example_images["anime80s"]) # Example 3 if not os.path.exists(example_images["night"]): img = generate_image("A cyberpunk city at night", "night", 456, False, 0.5) img.save(example_images["night"]) # Launch the app demo.launch(share=False, show_error=True)