Spaces:
Running
on
Zero
Running
on
Zero
File size: 6,666 Bytes
aa7f465 867f594 aa7f465 867f594 aa7f465 867f594 aa7f465 867f594 aa7f465 867f594 aa7f465 867f594 aa7f465 867f594 aa7f465 867f594 aa7f465 867f594 aa7f465 |
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 |
import gradio as gr
from diffusers import StableDiffusionPipeline, EulerAncestralDiscreteScheduler, DDIMScheduler, PNDMScheduler
import torch
from PIL import ImageEnhance, Image
import numpy as np
# Load Stable Diffusion pipeline
model_id = "CompVis/stable-diffusion-v1-4"
default_scheduler = DDIMScheduler.from_pretrained(model_id, subfolder="scheduler")
pipe = StableDiffusionPipeline.from_pretrained(model_id, scheduler=default_scheduler, torch_dtype=torch.float32, low_cpu_mem_usage=True)
pipe = pipe.to("cpu") # Use CPU
# Scheduler options
schedulers = {
"High-Definition & Fast (DDIM) - Good quality with fastest speed": DDIMScheduler,
"Artistic & Imaginative (Euler Ancestral) - Recommended for creative scenes, moderate speed": EulerAncestralDiscreteScheduler,
"Photo-Realistic (PNDM) - Best for realistic details, moderate speed": PNDMScheduler,
}
# Main image generation function with dynamic scheduling and size option
def generate_image(prompt, use_categories, genre, style, theme, lighting, scheduler_choice, quality, size):
# Check if additional categories should be added to the prompt
if use_categories:
prompt_text = (
f"{prompt.strip()} in a {genre.lower()} wallpaper style, "
f"with {style.lower()} visuals, focusing on a {theme.lower()} theme "
f"and {lighting.lower()} lighting."
)
else:
prompt_text = prompt.strip()
# Set the scheduler based on user choice
scheduler = schedulers[scheduler_choice].from_pretrained(model_id, subfolder="scheduler")
pipe.scheduler = scheduler
# Set output size based on selection
image_size = (256, 256) if size == "Profile Picture" else (384, 384)
# Generate image with specified quality and size
with torch.no_grad():
image = pipe(prompt_text, num_inference_steps=quality, guidance_scale=7.5).images[0]
image = image.resize(image_size) # Resize image to fit selected dimensions
return np.array(image) # Return as NumPy array for Gradio
# Post-processing function for brightness and contrast
def adjust_brightness_contrast(image, brightness, contrast):
image = Image.fromarray(image.astype('uint8'), 'RGB')
image = ImageEnhance.Brightness(image).enhance(brightness)
image = ImageEnhance.Contrast(image).enhance(contrast)
return np.array(image)
# Warning function to show a message if the user selects a high value for quality
def show_warning(quality):
if quality > 30:
return "⚠️ High Quality: This setting may slow down generation. Consider using 10-30 steps for best results."
return ""
# Build Gradio Interface
with gr.Blocks() as demo:
gr.Markdown("# ✨ AI-Powered Wallpaper/Profile Picture Generator\n🖼️ A tool to generate and fine-tune AI-created wallpapers and profile pictures with adjustable styles and effects.")
gr.Markdown("⚠️ **Live effects and advanced prompt engineering coming soon!** Results may vary, so try experimenting with settings to achieve the best results.")
# Image Generation Section
with gr.Tab("Image Generator"):
gr.Markdown("## Generate an Image")
with gr.Row():
with gr.Column():
custom_prompt = gr.Textbox(label="Custom Prompt", placeholder="Describe your image (e.g., 'A forest at sunset')")
# Toggle for using additional categories
use_categories = gr.Checkbox(label="Enable Advanced Settings (Genre, Style, Theme, Lighting)", value=False)
# Additional categories, hidden by default and shown only if use_categories is checked
with gr.Accordion("Advanced Settings", open=False):
genre = gr.Dropdown(["Futuristic", "Nature", "Abstract", "Fantasy", "Sci-Fi", "Cyberpunk"], label="Genre")
style = gr.Dropdown(["Realistic", "Surreal", "Digital Art", "Cartoon", "Photorealistic"], label="Style")
theme = gr.Dropdown(["Landscape", "Portrait", "Abstract Patterns", "Architecture"], label="Theme")
lighting = gr.Dropdown(["Warm", "Cool", "Cinematic", "Soft", "Neon"], label="Lighting")
quality = gr.Slider(10, 30, value=20, step=5, label="Image Quality", info="Higher values yield more detail but take longer to generate.")
warning_message = gr.Markdown("")
# Scheduler selection with default option
scheduler_choice = gr.Dropdown(
[
"High-Definition & Fast (DDIM) - Good quality with fastest speed",
"Artistic & Imaginative (Euler Ancestral) - Recommended for creative scenes, moderate speed",
"Photo-Realistic (PNDM) - Best for realistic details, moderate speed"
],
label="Artistic Style & Speed",
value="High-Definition & Fast (DDIM) - Good quality with fastest speed"
)
size = gr.Dropdown(["Profile Picture", "Wallpaper"], label="Image Size", value="Profile Picture")
generate_button = gr.Button("Generate Image")
with gr.Column():
generated_image = gr.Image(label="Generated Image", interactive=False)
# Display warning message for high-quality settings
quality.change(show_warning, inputs=[quality], outputs=warning_message)
# Bind the generate function to the generate button
generate_button.click(
fn=generate_image,
inputs=[custom_prompt, use_categories, genre, style, theme, lighting, scheduler_choice, quality, size],
outputs=generated_image
)
# Post-Processing Section
with gr.Tab("Edit Generated Image"):
gr.Markdown("## Adjust Brightness & Contrast")
with gr.Row():
with gr.Column():
brightness_slider = gr.Slider(0.5, 2.0, value=1.0, label="Brightness")
contrast_slider = gr.Slider(0.5, 2.0, value=1.0, label="Contrast")
apply_adjustments = gr.Button("Apply Adjustments")
with gr.Column():
output_image = gr.Image(label="Adjusted Image", interactive=False)
# Bind the brightness and contrast adjustment function to the Apply Adjustments button
apply_adjustments.click(
fn=adjust_brightness_contrast,
inputs=[generated_image, brightness_slider, contrast_slider],
outputs=output_image
)
# Launch with a public shareable link
demo.launch(share=True) |