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)