Update app.py

app.py

@@ -1,11 +1,9 @@
 import os
 import random
-import uuid
 import base64
 import gradio as gr
 import numpy as np
 from PIL import Image
-import spaces
 import torch
 import glob
 from datetime import datetime
@@ -13,44 +11,39 @@ import pandas as pd
 import json
 import re
 import logging
+from diffusers import StableDiffusionXLPipeline, EulerAncestralDiscreteScheduler
 
 # Set up logging
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
-from diffusers import StableDiffusionXLPipeline, EulerAncestralDiscreteScheduler
-
 DESCRIPTION = """# 🎨 ArtForge: Community AI Gallery
+Create, curate, and compete with AI-generated art. Join our creative multiplayer experience! 🖼️🏆✨"""
 
-Create, curate, and compete with AI-generated art. Join our creative multiplayer experience! 🖼️🏆✨
-"""
+METADATA_FILE = "image_metadata.json"
+MAX_SEED = np.iinfo(np.int32).max
 
 # Global variables
 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)
+def load_metadata():
+    global image_metadata
+    if os.path.exists(METADATA_FILE):
+        with open(METADATA_FILE, 'r') as f:
+            image_metadata = pd.DataFrame(json.load(f))
+    else:
+        image_metadata = pd.DataFrame(columns=['Filename', 'Prompt', 'Likes', 'Dislikes', 'Hearts', 'Created'])
 
-likes_cache = load_likes_cache()
+def save_metadata():
+    with open(METADATA_FILE, 'w') as f:
+        json.dump(image_metadata.to_dict('records'), f)
 
-def create_download_link(filename):
-    with open(filename, "rb") as file:
-        encoded_string = base64.b64encode(file.read()).decode('utf-8')
-        download_link = f'<a href="data:image/png;base64,{encoded_string}" download="{filename}">Download Image</a>'
-        return download_link
+load_metadata()
 
 def save_image(img, prompt):
-    global image_metadata, likes_cache
+    global image_metadata
     timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
-    safe_prompt = re.sub(r'[^\w\s-]', '', prompt.lower())[:50]  # Limit to 50 characters
+    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"
     img.save(filename)
@@ -60,84 +53,50 @@ def save_image(img, prompt):
         'Likes': [0],
         'Dislikes': [0],
         'Hearts': [0],
-        'Created': [datetime.now()]
+        'Created': [str(datetime.now())]
     })
     image_metadata = pd.concat([image_metadata, new_row], ignore_index=True, sort=False)
-    likes_cache[filename] = {'likes': 0, 'dislikes': 0, 'hearts': 0}
-    save_likes_cache(likes_cache)
+    save_metadata()
     logger.info(f"Saved new image: {filename}")
     return filename
 
-def randomize_seed_fn(seed: int, randomize_seed: bool) -> int:
-    if randomize_seed:
-        seed = random.randint(0, MAX_SEED)
-    return seed
-
 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)]
+    return [(file, get_image_caption(file)) for file in image_metadata['Filename'] 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}"
+    if filename in image_metadata['Filename'].values:
+        row = image_metadata[image_metadata['Filename'] == filename].iloc[0]
+        return f"{filename}\nPrompt: {row['Prompt']}\n👍 {row['Likes']} 👎 {row['Dislikes']} ❤️ {row['Hearts']}"
     return filename
 
 def delete_all_images():
-    global image_metadata, likes_cache
+    global image_metadata
     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)
+    save_metadata()
     logger.info("All images deleted")
     return get_image_gallery(), image_metadata.values.tolist()
 
 def delete_image(filename):
-    global image_metadata, likes_cache
+    global image_metadata
     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)
+        save_metadata()
         logger.info(f"Deleted image: {filename}")
     return get_image_gallery(), image_metadata.values.tolist()
 
 def vote(filename, vote_type):
-    global likes_cache, image_metadata
-    logger.info(f"Voting on {filename} with {vote_type}")
-    if filename in likes_cache:
-        likes_cache[filename][vote_type.lower()] += 1
-        save_likes_cache(likes_cache)
-        # Update image_metadata as well
-        if filename in image_metadata['Filename'].values:
-            image_metadata.loc[image_metadata['Filename'] == filename, vote_type] += 1
-        logger.info(f"Updated vote count for {filename}: {likes_cache[filename]}")
+    global image_metadata
+    if filename in image_metadata['Filename'].values:
+        image_metadata.loc[image_metadata['Filename'] == filename, vote_type] += 1
+        save_metadata()
+        logger.info(f"Updated {vote_type} count for {filename}")
     else:
-        logger.warning(f"File {filename} not found in likes_cache")
-    return get_image_gallery(), image_metadata.values.tolist()
-
-def get_random_style():
-    styles = [
-        "Impressionist", "Cubist", "Surrealist", "Abstract Expressionist",
-        "Pop Art", "Minimalist", "Baroque", "Art Nouveau", "Pointillist", "Fauvism"
-    ]
-    return random.choice(styles)
-
-MAX_SEED = np.iinfo(np.int32).max
-
-if not torch.cuda.is_available():
-    DESCRIPTION += "\n<p>Running on CPU 🥶 This demo may not work on CPU.</p>"
-
-USE_TORCH_COMPILE = 0
-ENABLE_CPU_OFFLOAD = 0
+        logger.warning(f"File {filename} not found in metadata")
+    return image_metadata.values.tolist()
 
 if torch.cuda.is_available():
     pipe = StableDiffusionXLPipeline.from_pretrained(
@@ -146,30 +105,22 @@ if torch.cuda.is_available():
         use_safetensors=True,
     )
     pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config)
-    
     pipe.load_lora_weights("ehristoforu/dalle-3-xl-v2", weight_name="dalle-3-xl-lora-v2.safetensors", adapter_name="dalle")
     pipe.set_adapters("dalle")
-
     pipe.to("cuda")
+else:
+    DESCRIPTION += "\n<p>Running on CPU 🥶 This demo may not work on CPU.</p>"
 
-@spaces.GPU(enable_queue=True)
 def generate(
-    prompt: str,
-    negative_prompt: str = "",
-    use_negative_prompt: bool = False,
-    seed: int = 0,
-    width: int = 1024,
-    height: int = 1024,
-    guidance_scale: float = 3,
-    randomize_seed: bool = False,
-    progress=gr.Progress(track_tqdm=True),
-    **kwargs  # Add this line to accept any additional arguments
+    prompt, negative_prompt="", use_negative_prompt=False,
+    seed=0, width=1024, height=1024, guidance_scale=3, randomize_seed=False,
+    progress=gr.Progress(track_tqdm=True)
 ):
-    seed = int(randomize_seed_fn(seed, randomize_seed))
-
+    if randomize_seed:
+        seed = random.randint(0, MAX_SEED)
     if not use_negative_prompt:
         negative_prompt = ""
-
+    
     images = pipe(
         prompt=prompt,
         negative_prompt=negative_prompt,
@@ -182,29 +133,12 @@ def generate(
         output_type="pil",
     ).images
     image_paths = [save_image(img, prompt) for img in images]
-    download_links = [create_download_link(path) for path in image_paths]
-
-    return image_paths, seed, download_links, get_image_gallery(), image_metadata.values.tolist()
-
-examples = [
-    f"{get_random_style()} 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"{get_random_style()} 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"{get_random_style()} 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"{get_random_style()} 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"{get_random_style()} 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"{get_random_style()} 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"{get_random_style()} interpretation of heaven. Create an ethereal atmosphere with soft, billowing clouds and radiant light, using a palette of celestial blues, golds, and whites.",
-    f"{get_random_style()} 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"{get_random_style()} 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"{get_random_style()} 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."
-]
+    return image_paths, seed, get_image_gallery(), image_metadata.values.tolist()
 
 css = '''
 .gradio-container{max-width: 1024px !important}
 h1{text-align:center}
-footer {
-visibility: hidden
-}
+footer {visibility: hidden}
 '''
 
 with gr.Blocks(theme=gr.themes.Soft(), css=css) as demo:
@@ -212,153 +146,47 @@ with gr.Blocks(theme=gr.themes.Soft(), css=css) as demo:
 
     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("Run", scale=0)
-            result = gr.Gallery(label="Result", columns=1, preview=True, show_label=False)
+            prompt = gr.Text(label="Prompt", placeholder="Enter your prompt")
+            run_button = gr.Button("Generate")
+        result = gr.Gallery(label="Result", columns=1, preview=True)
         with gr.Accordion("Advanced options", open=False):
-            use_negative_prompt = gr.Checkbox(label="Use negative prompt", value=True)
-            negative_prompt = gr.Text(
-                label="Negative prompt",
-                lines=4,
-                max_lines=6,
-                value="""(deformed, distorted, disfigured:1.3), poorly drawn, bad anatomy, wrong anatomy, extra limb, missing limb, floating limbs, (mutated hands and fingers:1.4), disconnected limbs, mutation, mutated, ugly, disgusting, blurry, amputation, (NSFW:1.25)""",
-                placeholder="Enter a negative prompt",
-                visible=True,
-            )
-            seed = gr.Slider(
-                label="Seed",
-                minimum=0,
-                maximum=MAX_SEED,
-                step=1,
-                value=0,
-                visible=True
-            )
+            use_negative_prompt = gr.Checkbox(label="Use negative prompt", value=False)
+            negative_prompt = gr.Text(label="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(visible=True):
-                width = gr.Slider(
-                    label="Width",
-                    minimum=512,
-                    maximum=2048,
-                    step=8,
-                    value=1920,
-                )
-                height = gr.Slider(
-                    label="Height",
-                    minimum=512,
-                    maximum=2048,
-                    step=8,
-                    value=1080,
-                )
-            with gr.Row():
-                guidance_scale = gr.Slider(
-                    label="Guidance Scale",
-                    minimum=0.1,
-                    maximum=20.0,
-                    step=0.1,
-                    value=20.0,
-                )
-        
-        gr.Examples(
-            examples=examples,
-            inputs=prompt,
-            outputs=[result, seed],
-            fn=generate,
-            cache_examples=False,
-        )
+            width = gr.Slider(label="Width", minimum=512, maximum=2048, step=8, value=1024)
+            height = gr.Slider(label="Height", minimum=512, maximum=2048, step=8, value=1024)
+            guidance_scale = gr.Slider(label="Guidance Scale", minimum=0.1, maximum=20.0, step=0.1, value=7.5)
 
-    with gr.Tab("Gallery and Voting"):
-        image_gallery = gr.Gallery(label="Generated Images", show_label=True, columns=4, height="auto")
-        
+    with gr.Tab("Gallery"):
+        image_gallery = gr.Gallery(label="Generated Images", columns=4, height="auto")
+        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"])
         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")
 
-    use_negative_prompt.change(
-        fn=lambda x: gr.update(visible=x),
-        inputs=use_negative_prompt,
-        outputs=negative_prompt,
-        api_name=False,
-    )
-
-    delete_all_button.click(
-        fn=delete_all_images,
-        inputs=[],
-        outputs=[image_gallery, metadata_df],
-    )
-
-    image_gallery.select(
-        fn=lambda evt: evt,
-        inputs=[],
-        outputs=[selected_image],
-    )
+    use_negative_prompt.change(lambda x: gr.update(visible=x), inputs=use_negative_prompt, outputs=negative_prompt)
 
-    like_button.click(
-        fn=lambda x: vote(x, 'Likes') if x else (get_image_gallery(), image_metadata.values.tolist()),
-        inputs=[selected_image],
-        outputs=[image_gallery, metadata_df],
-    )
-
-    dislike_button.click(
-        fn=lambda x: vote(x, 'Dislikes') if x else (get_image_gallery(), image_metadata.values.tolist()),
-        inputs=[selected_image],
-        outputs=[image_gallery, metadata_df],
-    )
-
-    heart_button.click(
-        fn=lambda x: vote(x, 'Hearts') if x else (get_image_gallery(), image_metadata.values.tolist()),
-        inputs=[selected_image],
-        outputs=[image_gallery, metadata_df],
-    )
-
-    delete_image_button.click(
-        fn=delete_image,
-        inputs=[selected_image],
-        outputs=[image_gallery, metadata_df],
+    run_button.click(
+        fn=generate,
+        inputs=[prompt, negative_prompt, use_negative_prompt, seed, width, height, guidance_scale, randomize_seed],
+        outputs=[result, seed, image_gallery, metadata_df]
     )
 
-    def update_gallery_and_metadata():
-        return gr.update(value=get_image_gallery()), gr.update(value=image_metadata.values.tolist())
+    delete_all_button.click(fn=delete_all_images, outputs=[image_gallery, metadata_df])
+    image_gallery.select(fn=lambda evt: evt, outputs=selected_image)
+    delete_image_button.click(fn=delete_image, inputs=[selected_image], outputs=[image_gallery, metadata_df])
 
-    gr.on(
-        triggers=[
-            prompt.submit,
-            negative_prompt.submit,
-            run_button.click,
-        ],
-        fn=generate,
-        inputs=[
-            prompt,
-            negative_prompt,
-            use_negative_prompt,
-            seed,
-            width,
-            height,
-            guidance_scale,
-            randomize_seed],
-        outputs=[result, seed, gr.HTML(visible=False), image_gallery, metadata_df],
-        api_name="run",
-    )
+    for button, vote_type in [(like_button, 'Likes'), (dislike_button, 'Dislikes'), (heart_button, 'Hearts')]:
+        button.click(fn=lambda x, vt=vote_type: vote(x, vt) if x else image_metadata.values.tolist(), inputs=[selected_image], outputs=[metadata_df])
 
-    demo.load(fn=update_gallery_and_metadata, 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)