app.py

import random
import os
import uuid
from datetime import datetime
import gradio as gr
import numpy as np
import spaces
import torch
from diffusers import DiffusionPipeline
from PIL import Image

# Create permanent storage directory
SAVE_DIR = "saved_images"  # Gradio will handle the persistence
if not os.path.exists(SAVE_DIR):
    os.makedirs(SAVE_DIR, exist_ok=True)

device = "cuda" if torch.cuda.is_available() else "cpu"
repo_id = "black-forest-labs/FLUX.1-dev"
adapter_id = "flux-lora-korea-palace"

pipeline = DiffusionPipeline.from_pretrained(repo_id, torch_dtype=torch.bfloat16)
pipeline.load_lora_weights(adapter_id)
pipeline = pipeline.to(device)

MAX_SEED = np.iinfo(np.int32).max
MAX_IMAGE_SIZE = 1024

def save_generated_image(image, prompt):
    # Generate unique filename with timestamp
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    unique_id = str(uuid.uuid4())[:8]
    filename = f"{timestamp}_{unique_id}.png"
    filepath = os.path.join(SAVE_DIR, filename)
    
    # Save the image
    image.save(filepath)
    
    # Save metadata
    metadata_file = os.path.join(SAVE_DIR, "metadata.txt")
    with open(metadata_file, "a", encoding="utf-8") as f:
        f.write(f"{filename}|{prompt}|{timestamp}\n")
    
    return filepath

def load_generated_images():
    if not os.path.exists(SAVE_DIR):
        return []
    
    # Load all images from the directory
    image_files = [os.path.join(SAVE_DIR, f) for f in os.listdir(SAVE_DIR) 
                  if f.endswith(('.png', '.jpg', '.jpeg', '.webp'))]
    # Sort by creation time (newest first)
    image_files.sort(key=lambda x: os.path.getctime(x), reverse=True)
    return image_files

def load_predefined_images():
    predefined_images = [
        "assets/cm1.webp",
        "assets/cm2.webp",
        "assets/cm3.webp",
        "assets/cm4.webp",
        "assets/cm5.webp",
        "assets/cm6.webp",
    ]
    return predefined_images

@spaces.GPU(duration=120)
def inference(
    prompt: str,
    seed: int,
    randomize_seed: bool,
    width: int,
    height: int,
    guidance_scale: float,
    num_inference_steps: int,
    lora_scale: float,
    progress: gr.Progress = gr.Progress(track_tqdm=True),
):
    if randomize_seed:
        seed = random.randint(0, MAX_SEED)
    generator = torch.Generator(device=device).manual_seed(seed)
    
    image = pipeline(
        prompt=prompt,
        guidance_scale=guidance_scale,
        num_inference_steps=num_inference_steps,
        width=width,
        height=height,
        generator=generator,
        joint_attention_kwargs={"scale": lora_scale},
    ).images[0]
    
    # Save the generated image
    filepath = save_generated_image(image, prompt)
    
    # Return the image, seed, and updated gallery
    return image, seed, load_generated_images()

examples = [
    "Geunjeongjeon Hall of Gyeongbokgung Palace in the crisp winter morning, where fresh snow blankets the vast courtyard. The magnificent throne hall stands proud with its vibrant dancheong colors contrasting against the pure white snow. Intricate dragon carvings along the massive red pillars seem to come alive in the golden morning light, while the multi-tiered upturned eaves create dramatic shadows on the pristine snow below. [trigger]",
    
    "Geunjeongjeon Hall of Gyeongbokgung Palace during a spring dawn, when cherry blossoms frame the grand structure. The first light of day illuminates the ornate throne hall's double-tiered stone platform, while delicate pink petals drift gracefully across the courtyard. The meticulously painted wooden columns and beams showcase the rich vermillion and emerald dancheong patterns, their colors especially vivid against the soft morning sky. [trigger]",
    
    "Geunjeongjeon Hall of Gyeongbokgung Palace at summer dusk, as the setting sun bathes the royal throne hall in warm amber light. The imposing structure's curved roof tiles gleam with a golden sheen, while the intricate ceiling designs cast intricate shadows across the wooden floors. Stone guardians stand sentinel at the base of the grand staircase, their forms dramatic against the deepening purple sky. [trigger]",
    
    "Geunjeongjeon Hall of Gyeongbokgung Palace beneath autumn rain, where water droplets trace elegant paths down the carved dragon pillars. The wet stone steps reflect the deep reds and greens of the painted rafters above, while gray clouds create a dramatic backdrop for the hall's distinctive silhouette. Bronze wind chimes at the corners of the sweeping roof sing softly in the gentle rain. [trigger]",
    
    "Geunjeongjeon Hall of Gyeongbokgung Palace under a full moon, its majestic form illuminated by silvery moonlight. The throne hall's imposing structure casts long shadows across the empty courtyard, while moonbeams highlight the intricate details of the painted ceremonies beneath the eaves. The carved stone foundations seem to glow with an ethereal light, creating a mystical atmosphere worthy of this royal sanctuary. [trigger]",
    
    "Geunjeongjeon Hall of Gyeongbokgung Palace during a traditional ceremony, where the morning sun streams through the towering red pillars. Incense smoke curls gracefully around the ornately decorated beams, while the vast courtyard comes alive with the colorful hanbok of court officials. The ancient throne platform stands dignified beneath the soaring roof, its carved details emphasized by shafts of golden light piercing through the traditional wooden lattice windows. [trigger]"
]

css = """
footer {
    visibility: hidden;
}
"""

with gr.Blocks(theme="Yntec/HaleyCH_Theme_Orange", css=css, analytics_enabled=False) as demo:
    gr.HTML('<div class="title"> Claude Monet STUDIO </div>')
    gr.HTML('<div class="title">😄Image to Video Explore: <a href="https://huggingface.co/spaces/ginigen/theater" target="_blank">https://huggingface.co/spaces/ginigen/theater</a></div>')

    with gr.Tabs() as tabs:
        with gr.Tab("Generation"):
            with gr.Column(elem_id="col-container"):
                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.Image(label="Result", show_label=False)

                with gr.Accordion("Advanced Settings", open=False):
                    seed = gr.Slider(
                        label="Seed",
                        minimum=0,
                        maximum=MAX_SEED,
                        step=1,
                        value=42,
                    )
                    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=768,
                        )

                    with gr.Row():
                        guidance_scale = gr.Slider(
                            label="Guidance scale",
                            minimum=0.0,
                            maximum=10.0,
                            step=0.1,
                            value=3.5,
                        )
                        num_inference_steps = gr.Slider(
                            label="Number of inference steps",
                            minimum=1,
                            maximum=50,
                            step=1,
                            value=30,
                        )
                        lora_scale = gr.Slider(
                            label="LoRA scale",
                            minimum=0.0,
                            maximum=1.0,
                            step=0.1,
                            value=1.0,
                        )

                gr.Examples(
                    examples=examples,
                    inputs=[prompt],
                    outputs=[result, seed],
                )

        with gr.Tab("Gallery"):
            gallery_header = gr.Markdown("### Generated Images Gallery")
            generated_gallery = gr.Gallery(
                label="Generated Images",
                columns=6,
                show_label=False,
                value=load_generated_images(),
                elem_id="generated_gallery",
                height="auto"
            )
            refresh_btn = gr.Button("🔄 Refresh Gallery")

    # Add sample gallery section at the bottom
    gr.Markdown("### Claude Monet Style Examples")
    predefined_gallery = gr.Gallery(
        label="Sample Images", 
        columns=3,
        rows=2,
        show_label=False, 
        value=load_predefined_images()
    )

    # Event handlers
    def refresh_gallery():
        return load_generated_images()

    refresh_btn.click(
        fn=refresh_gallery,
        inputs=None,
        outputs=generated_gallery,
    )

    gr.on(
        triggers=[run_button.click, prompt.submit],
        fn=inference,
        inputs=[
            prompt,
            seed,
            randomize_seed,
            width,
            height,
            guidance_scale,
            num_inference_steps,
            lora_scale,
        ],
        outputs=[result, seed, generated_gallery],
    )

demo.queue()
demo.launch()