DAVINCIDIFFtextgardio.py

import torch
from diffusers import StableDiffusionXLPipeline, AutoencoderTiny
import os
import json
import google.generativeai as genai
import gradio as gr
import random

art_styles = [
    "Impressionism", "Surrealism", "Cubism", "Abstract", "Realism", "Expressionism",
    "Pop Art", "Futurism", "Minimalism", "Conceptual Art", "Baroque", "Renaissance",
    "Gothic", "Neoclassicism", "Romanticism", "Art Nouveau", "Art Deco", "Post-Impressionism"
]

artists = [
    "Pablo Picasso", "Vincent van Gogh", "Leonardo da Vinci", "Claude Monet", "Salvador Dalí",
    "Frida Kahlo", "Rembrandt", "Michelangelo", "Andy Warhol", "Jackson Pollock",
    "Georgia O'Keeffe", "Edvard Munch", "Henri Matisse", "Paul Cézanne", "Gustav Klimt",
    "Caravaggio", "Jean-Michel Basquiat", "Raphael", "Auguste Rodin", "Yayoi Kusama"
]

light_effects = [
    "soft lighting", "dramatic lighting", "rim lighting", "ambient lighting", "studio lighting",
    "natural light", "backlighting", "side lighting", "low-key lighting", "high-key lighting"
]

depth_effects = [
    "shallow depth of field", "deep depth of field", "foreground focus", "background blur",
    "bokeh effect", "layered depth", "aerial perspective", "overlapping elements", "linear perspective",
    "vanishing point"
]

# Function to save the Gemini API key
def save_gemini_api_key(api_key):
    with open("gemini_api_key.json", "w") as f:
        json.dump({"api_key": api_key}, f)

# Function to load the Gemini API key
def load_gemini_api_key():
    try:
        with open("gemini_api_key.json", "r") as f:
            data = json.load(f)
            return data.get("api_key", "")
    except FileNotFoundError:
        return ""

# Function to delete the Gemini API key
def delete_gemini_api_key():
    try:
        if os.path.exists("gemini_api_key.json"):
            os.remove("gemini_api_key.json")
    except Exception as e:
        print(f"Error deleting Gemini API key: {e}")

# Function to enhance prompt using Gemini's API with specific instructions
def enhance_prompt_with_gemini(prompt, api_key):
    try:
        if api_key:
            genai.configure(api_key=api_key)

            # Construct the specific instruction prompt
            instruction = f"You are Professor of creating a prompt for Stable Diffusion to generate an image. Write the phrases for art styles, art movements, artists names, light and depth effects for this {prompt}. Make this prompt the best ever, only response with prompt itself in one paragraph."

            # Generate content using Gemini API with the instruction prompt
            model = genai.GenerativeModel('gemini-1.5-pro')
            response = model.generate_content(instruction)
            enhanced_prompt = response.text.strip()
            return enhanced_prompt
        else:
            return prompt
    except Exception as e:
        print(f"Failed to enhance prompt with Gemini API: {e}")
        return prompt  # Fallback to original prompt on error

# Function to generate images based on user input
def generate_image(prompt, negative_prompt, cfg_scale, steps, width, height, seed, art_style, artist, light_effect, depth_effect, api_key):
    queue = []
    additional_prompts = f"{art_style}, {artist}, {light_effect}, {depth_effect}"
    full_prompt = f"{prompt}, {additional_prompts}"
    enhanced_prompt = enhance_prompt_with_gemini(full_prompt, api_key)

    # Generate a random seed if not provided
    if seed is None or seed == "":
        seed = random.randint(0, 2**32 - 1)
    
    print("\nStarting image generation with the following parameters:")
    print(f"Enhanced Prompt: {enhanced_prompt}")
    print(f"Negative Prompt: {negative_prompt}")
    print(f"CFG Scale: {cfg_scale}")
    print(f"Steps: {steps}")
    print(f"Width: {width}")
    print(f"Height: {height}")
    print(f"Seed: {seed}")
    
    queue.append({
        'prompt': enhanced_prompt,
        'negative_prompt': negative_prompt,
        'cfg_scale': cfg_scale,
        'steps': steps,
        'width': width,
        'height': height,
        'seed': seed,
    })

    vae = AutoencoderTiny.from_pretrained(
      'madebyollin/taesdxl',
      use_safetensors=True,
      torch_dtype=torch.float16,
    ).to('cpu')

    pipe = StableDiffusionXLPipeline.from_pretrained(
        'ABDALLALSWAITI/DAVINCI-DIFF',
        torch_dtype=torch.float16,
        use_safetensors=True,
        vae=vae
    ).to('cpu')

    generator = torch.manual_seed(seed)

    image = pipe(
        prompt=enhanced_prompt,
        negative_prompt=negative_prompt,
        height=height,
        width=width,
        num_inference_steps=steps,
        guidance_scale=cfg_scale,
        generator=generator,
        output_type="pil",
    ).images[0]

    return image

# Gradio Blocks interface
api_key = load_gemini_api_key()

def on_submit(prompt, negative_prompt, cfg_scale, steps, width, height, seed, art_style, artist, light_effect, depth_effect, api_key):
    return generate_image(prompt, negative_prompt, cfg_scale, steps, width, height, seed, art_style, artist, light_effect, depth_effect, api_key)

with gr.Blocks() as interface:
    gr.Markdown("# Image Generation with Stable Diffusion")
    
    with gr.Row():
        prompt = gr.Textbox(label="Prompt", value="3/4 shot, candid photograph of a beautiful 30 year old redhead woman with messy dark hair, peacefully sleeping in her bed, night, dark, light from window, dark shadows, masterpiece, uhd, moody")
        negative_prompt = gr.Textbox(label="Negative Prompt", value="")
    
    with gr.Row():
        cfg_scale = gr.Number(label="CFG Scale", value=7.5)
        steps = gr.Number(label="Steps", value=50)
    
    with gr.Row():
        width = gr.Number(label="Image Width", value=1024)
        height = gr.Number(label="Image Height", value=1024)
    
    seed = gr.Number(label="Seed (leave empty for random seed)", value=None)
    
    with gr.Row():
        art_style = gr.Dropdown(label="Art Style", choices=art_styles, value="Impressionism")
        artist = gr.Dropdown(label="Artist", choices=artists, value="Vincent van Gogh")
    
    with gr.Row():
        light_effect = gr.Dropdown(label="Light Effect", choices=light_effects, value="soft lighting")
        depth_effect = gr.Dropdown(label="Depth Effect", choices=depth_effects, value="shallow depth of field")
    
    api_key = gr.Textbox(label="Gemini API Key", type="password", placeholder="Enter Gemini API Key (optional)")
    
    generate_btn = gr.Button("Generate Image")
    output_image = gr.Image(label="Generated Image")
    
    generate_btn.click(
        on_submit,
        inputs=[prompt, negative_prompt, cfg_scale, steps, width, height, seed, art_style, artist, light_effect, depth_effect, api_key],
        outputs=output_image
    )

interface.launch(share=True)

def on_session_end(session):
    delete_gemini_api_key()