#!/usr/bin/env python from __future__ import annotations import os, random, glob, re, json, base64 from datetime import datetime import gradio as gr import numpy as np import PIL.Image import spaces import torch import pandas as pd from diffusers import AutoencoderKL, DiffusionPipeline DESCRIPTION = """ # šØ ArtForge: OpenDALLE AI Masterpiece Arena š¼ļøš """ if not torch.cuda.is_available(): DESCRIPTION += "\n
Running on CPU š„¶ This demo does not work on CPU. Please use the online demo instead.
" MAX_SEED = np.iinfo(np.int32).max CACHE_EXAMPLES = torch.cuda.is_available() and os.getenv("CACHE_EXAMPLES", "0") == "1" MAX_IMAGE_SIZE = int(os.getenv("MAX_IMAGE_SIZE", "1024")) USE_TORCH_COMPILE = os.getenv("USE_TORCH_COMPILE") == "1" ENABLE_CPU_OFFLOAD = os.getenv("ENABLE_CPU_OFFLOAD") == "1" ENABLE_REFINER = os.getenv("ENABLE_REFINER", "0") == "1" # Global variables for metadata and likes cache image_metadata = pd.DataFrame(columns=['Filename', 'Prompt', 'Likes', 'Dislikes', 'Hearts', 'Created']) LIKES_CACHE_FILE = "likes_cache.json" def load_likes_cache(): if os.path.exists(LIKES_CACHE_FILE): with open(LIKES_CACHE_FILE, 'r') as f: return json.load(f) return {} def save_likes_cache(cache): with open(LIKES_CACHE_FILE, 'w') as f: json.dump(cache, f) likes_cache = load_likes_cache() device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") if torch.cuda.is_available(): vae = AutoencoderKL.from_pretrained("madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float16) pipe = DiffusionPipeline.from_pretrained("dataautogpt3/OpenDalleV1.1", vae=vae, torch_dtype=torch.float16, use_safetensors=True, variant="fp16") if ENABLE_REFINER: refiner = DiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-xl-refiner-1.0", vae=vae, torch_dtype=torch.float16, use_safetensors=True, variant="fp16") if ENABLE_CPU_OFFLOAD: pipe.enable_model_cpu_offload() if ENABLE_REFINER: refiner.enable_model_cpu_offload() else: pipe.to(device) if ENABLE_REFINER: refiner.to(device) if USE_TORCH_COMPILE: pipe.unet = torch.compile(pipe.unet, mode="reduce-overhead", fullgraph=True) if ENABLE_REFINER: refiner.unet = torch.compile(refiner.unet, mode="reduce-overhead", fullgraph=True) def randomize_seed_fn(seed: int, randomize_seed: bool) -> int: return random.randint(0, MAX_SEED) if randomize_seed else seed def create_download_link(filename): with open(filename, "rb") as file: encoded_string = base64.b64encode(file.read()).decode('utf-8') download_link = f'Download Image' return download_link def save_image(image: PIL.Image.Image, prompt: str) -> str: global image_metadata, likes_cache timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") safe_prompt = re.sub(r'[^\w\s-]', '', prompt.lower())[:50] safe_prompt = re.sub(r'[-\s]+', '-', safe_prompt).strip('-') filename = f"{timestamp}_{safe_prompt}.png" image.save(filename) new_row = pd.DataFrame({ 'Filename': [filename], 'Prompt': [prompt], 'Likes': [0], 'Dislikes': [0], 'Hearts': [0], 'Created': [datetime.now()] }) image_metadata = pd.concat([image_metadata, new_row], ignore_index=True) likes_cache[filename] = {'likes': 0, 'dislikes': 0, 'hearts': 0} save_likes_cache(likes_cache) return filename def get_image_gallery(): global image_metadata image_files = image_metadata['Filename'].tolist() return [(file, get_image_caption(file)) for file in image_files if os.path.exists(file)] def get_image_caption(filename): global likes_cache, image_metadata if filename in likes_cache: likes = likes_cache[filename]['likes'] dislikes = likes_cache[filename]['dislikes'] hearts = likes_cache[filename]['hearts'] prompt = image_metadata[image_metadata['Filename'] == filename]['Prompt'].values[0] return f"{filename}\nPrompt: {prompt}\nš {likes} š {dislikes} ā¤ļø {hearts}" return filename def delete_all_images(): global image_metadata, likes_cache for file in image_metadata['Filename']: if os.path.exists(file): os.remove(file) image_metadata = pd.DataFrame(columns=['Filename', 'Prompt', 'Likes', 'Dislikes', 'Hearts', 'Created']) likes_cache = {} save_likes_cache(likes_cache) return get_image_gallery(), image_metadata.values.tolist() def delete_image(filename): global image_metadata, likes_cache if filename and os.path.exists(filename): os.remove(filename) image_metadata = image_metadata[image_metadata['Filename'] != filename] if filename in likes_cache: del likes_cache[filename] save_likes_cache(likes_cache) return get_image_gallery(), image_metadata.values.tolist() def vote(filename, vote_type): global likes_cache if filename in likes_cache: likes_cache[filename][vote_type.lower()] += 1 save_likes_cache(likes_cache) return get_image_gallery(), image_metadata.values.tolist() @spaces.GPU(enable_queue=True) def generate(prompt: str, negative_prompt: str = "", prompt_2: str = "", negative_prompt_2: str = "", use_negative_prompt: bool = False, use_prompt_2: bool = False, use_negative_prompt_2: bool = False, seed: int = 0, width: int = 1024, height: int = 1024, guidance_scale_base: float = 5.0, guidance_scale_refiner: float = 5.0, num_inference_steps_base: int = 25, num_inference_steps_refiner: int = 25, apply_refiner: bool = False, progress=gr.Progress(track_tqdm=True)) -> PIL.Image.Image: print(f"** Generating image for: \"{prompt}\" **") generator = torch.Generator().manual_seed(seed) if not use_negative_prompt: negative_prompt = None if not use_prompt_2: prompt_2 = None if not use_negative_prompt_2: negative_prompt_2 = None if not apply_refiner: image = pipe(prompt=prompt, negative_prompt=negative_prompt, prompt_2=prompt_2, negative_prompt_2=negative_prompt_2, width=width, height=height, guidance_scale=guidance_scale_base, num_inference_steps=num_inference_steps_base, generator=generator, output_type="pil").images[0] else: latents = pipe(prompt=prompt, negative_prompt=negative_prompt, prompt_2=prompt_2, negative_prompt_2=negative_prompt_2, width=width, height=height, guidance_scale=guidance_scale_base, num_inference_steps=num_inference_steps_base, generator=generator, output_type="latent").images image = refiner(prompt=prompt, negative_prompt=negative_prompt, prompt_2=prompt_2, negative_prompt_2=negative_prompt_2, guidance_scale=guidance_scale_refiner, num_inference_steps=num_inference_steps_refiner, image=latents, generator=generator).images[0] filename = save_image(image, prompt) download_link = create_download_link(filename) return image, get_image_gallery(), download_link, image_metadata.values.tolist() examples = [ f"{random.choice(['Impressionist', 'Cubist', 'Surrealist', 'Abstract Expressionist', 'Pop Art', 'Minimalist', 'Baroque', 'Art Nouveau', 'Pointillist', 'Fauvism'])} painting of a majestic lighthouse on a rocky coast. Use bold brushstrokes and a vibrant color palette to capture the interplay of light and shadow as the lighthouse beam cuts through a stormy night sky.", f"{random.choice(['Impressionist', 'Cubist', 'Surrealist', 'Abstract Expressionist', 'Pop Art', 'Minimalist', 'Baroque', 'Art Nouveau', 'Pointillist', 'Fauvism'])} still life featuring a pair of vintage eyeglasses. Focus on the intricate details of the frames and lenses, using a warm color scheme to evoke a sense of nostalgia and wisdom.", f"{random.choice(['Impressionist', 'Cubist', 'Surrealist', 'Abstract Expressionist', 'Pop Art', 'Minimalist', 'Baroque', 'Art Nouveau', 'Pointillist', 'Fauvism'])} depiction of a rustic wooden stool in a sunlit artist's studio. Emphasize the texture of the wood and the interplay of light and shadow, using a mix of earthy tones and highlights.", f"{random.choice(['Impressionist', 'Cubist', 'Surrealist', 'Abstract Expressionist', 'Pop Art', 'Minimalist', 'Baroque', 'Art Nouveau', 'Pointillist', 'Fauvism'])} scene viewed through an ornate window frame. Contrast the intricate details of the window with a dreamy, soft-focus landscape beyond, using a palette that transitions from cool interior tones to warm exterior hues.", f"{random.choice(['Impressionist', 'Cubist', 'Surrealist', 'Abstract Expressionist', 'Pop Art', 'Minimalist', 'Baroque', 'Art Nouveau', 'Pointillist', 'Fauvism'])} close-up study of interlaced fingers. Use a monochromatic color scheme to emphasize the form and texture of the hands, with dramatic lighting to create depth and emotion.", f"{random.choice(['Impressionist', 'Cubist', 'Surrealist', 'Abstract Expressionist', 'Pop Art', 'Minimalist', 'Baroque', 'Art Nouveau', 'Pointillist', 'Fauvism'])} composition featuring a set of dice in motion. Capture the energy and randomness of the throw, using a dynamic color palette and blurred lines to convey movement.", f"{random.choice(['Impressionist', 'Cubist', 'Surrealist', 'Abstract Expressionist', 'Pop Art', 'Minimalist', 'Baroque', 'Art Nouveau', 'Pointillist', 'Fauvism'])} interpretation of heaven. Create an ethereal atmosphere with soft, billowing clouds and radiant light, using a palette of celestial blues, golds, and whites.", f"{random.choice(['Impressionist', 'Cubist', 'Surrealist', 'Abstract Expressionist', 'Pop Art', 'Minimalist', 'Baroque', 'Art Nouveau', 'Pointillist', 'Fauvism'])} portrayal of an ancient, mystical gate. Combine architectural details with elements of fantasy, using a rich, jewel-toned palette to create an air of mystery and magic.", f"{random.choice(['Impressionist', 'Cubist', 'Surrealist', 'Abstract Expressionist', 'Pop Art', 'Minimalist', 'Baroque', 'Art Nouveau', 'Pointillist', 'Fauvism'])} portrait of a curious cat. Focus on capturing the feline's expressive eyes and sleek form, using a mix of bold and subtle colors to bring out the cat's personality.", f"{random.choice(['Impressionist', 'Cubist', 'Surrealist', 'Abstract Expressionist', 'Pop Art', 'Minimalist', 'Baroque', 'Art Nouveau', 'Pointillist', 'Fauvism'])} abstract representation of toes in sand. Use textured brushstrokes to convey the feeling of warm sand, with a palette inspired by a sun-drenched beach." ] css = ''' .gradio-container {max-width: 1024px !important} h1 {text-align: center} footer {visibility: hidden} ''' theme = gr.themes.Soft() with gr.Blocks(css=css, theme=theme) as demo: gr.Markdown(DESCRIPTION) gr.DuplicateButton(value="Duplicate Space for private use", elem_id="duplicate-button", visible=os.getenv("SHOW_DUPLICATE_BUTTON") == "1") with gr.Tab("Generate Images"): with gr.Group(): with gr.Row(): prompt = gr.Text(label="Prompt", show_label=False, max_lines=1, placeholder="Enter your prompt", container=False) run_button = gr.Button("Generate", scale=0) result = gr.Image(label="Result", show_label=False) download_link = gr.HTML(label="Download", show_label=False) with gr.Accordion("Advanced options", open=False): use_negative_prompt = gr.Checkbox(label="Use negative prompt", value=False) use_prompt_2 = gr.Checkbox(label="Use prompt 2", value=False) use_negative_prompt_2 = gr.Checkbox(label="Use negative prompt 2", value=False) negative_prompt = gr.Text(label="Negative prompt", max_lines=1, placeholder="Enter a negative prompt", visible=False) prompt_2 = gr.Text(label="Prompt 2", max_lines=1, placeholder="Enter your second prompt", visible=False) negative_prompt_2 = gr.Text(label="Negative prompt 2", max_lines=1, placeholder="Enter a second negative prompt", visible=False) seed = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=0) randomize_seed = gr.Checkbox(label="Randomize seed", value=True) with gr.Row(): width = gr.Slider(label="Width", minimum=256, maximum=MAX_IMAGE_SIZE, step=32, value=1024) height = gr.Slider(label="Height", minimum=256, maximum=MAX_IMAGE_SIZE, step=32, value=1024) apply_refiner = gr.Checkbox(label="Apply refiner", value=False, visible=ENABLE_REFINER) with gr.Row(): guidance_scale_base = gr.Slider(label="Guidance scale for base", minimum=1, maximum=20, step=0.1, value=5.0) num_inference_steps_base = gr.Slider(label="Number of inference steps for base", minimum=10, maximum=100, step=1, value=25) with gr.Row(visible=False) as refiner_params: guidance_scale_refiner = gr.Slider(label="Guidance scale for refiner", minimum=1, maximum=20, step=0.1, value=5.0) num_inference_steps_refiner = gr.Slider(label="Number of inference steps for refiner", minimum=10, maximum=100, step=1, value=25) with gr.Tab("Gallery and Voting"): image_gallery = gr.Gallery(label="Generated Images", show_label=True, columns=4, height="auto") with gr.Row(): like_button = gr.Button("š Like") dislike_button = gr.Button("š Dislike") heart_button = gr.Button("ā¤ļø Heart") delete_image_button = gr.Button("šļø Delete Selected Image") selected_image = gr.State(None) with gr.Tab("Metadata and Management"): metadata_df = gr.Dataframe( label="Image Metadata", headers=["Filename", "Prompt", "Likes", "Dislikes", "Hearts", "Created"], interactive=False ) delete_all_button = gr.Button("šļø Delete All Images") gr.Examples(examples=examples, inputs=prompt, outputs=[result, image_gallery, download_link, metadata_df], fn=generate, cache_examples=CACHE_EXAMPLES) use_negative_prompt.change(fn=lambda x: gr.update(visible=x), inputs=use_negative_prompt, outputs=negative_prompt, queue=False, api_name=False) use_prompt_2.change(fn=lambda x: gr.update(visible=x), inputs=use_prompt_2, outputs=prompt_2, queue=False, api_name=False) use_negative_prompt_2.change(fn=lambda x: gr.update(visible=x), inputs=use_negative_prompt_2, outputs=negative_prompt_2, queue=False, api_name=False) apply_refiner.change(fn=lambda x: gr.update(visible=x), inputs=apply_refiner, outputs=refiner_params, queue=False, api_name=False) prompt.submit(fn=randomize_seed_fn, inputs=[seed, randomize_seed], outputs=seed, queue=False, api_name=False).then( fn=generate, inputs=[prompt, negative_prompt, prompt_2, negative_prompt_2, use_negative_prompt, use_prompt_2, use_negative_prompt_2, seed, width, height, guidance_scale_base, guidance_scale_refiner, num_inference_steps_base, num_inference_steps_refiner, apply_refiner], outputs=[result, image_gallery, download_link, metadata_df] ) run_button.click(fn=randomize_seed_fn, inputs=[seed, randomize_seed], outputs=seed, queue=False, api_name=False).then( fn=generate, inputs=[prompt, negative_prompt, prompt_2, negative_prompt_2, use_negative_prompt, use_prompt_2, use_negative_prompt_2, seed, width, height, guidance_scale_base, guidance_scale_refiner, num_inference_steps_base, num_inference_steps_refiner, apply_refiner], outputs=[result, image_gallery, download_link, metadata_df] ) image_gallery.select(fn=lambda evt: evt, inputs=[], outputs=[selected_image]) like_button.click(fn=lambda x: vote(x, 'likes'), inputs=[selected_image], outputs=[image_gallery, metadata_df]) dislike_button.click(fn=lambda x: vote(x, 'dislikes'), inputs=[selected_image], outputs=[image_gallery, metadata_df]) heart_button.click(fn=lambda x: vote(x, 'hearts'), inputs=[selected_image], outputs=[image_gallery, metadata_df]) delete_image_button.click(fn=delete_image, inputs=[selected_image], outputs=[image_gallery, metadata_df]) delete_all_button.click(fn=delete_all_images, inputs=[], outputs=[image_gallery, metadata_df]) demo.load(fn=lambda: (get_image_gallery(), image_metadata.values.tolist()), outputs=[image_gallery, metadata_df]) if __name__ == "__main__": demo.queue(max_size=20).launch(share=True, debug=False)