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 = "openfree/paul-cezanne"

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

@spaces.GPU(duration=60)
def inference(
    prompt,
    seed=42,
    randomize_seed=True,
    width=1024,
    height=768,
    guidance_scale=3.5,
    num_inference_steps=30,
    lora_scale=1.0,
    progress=None,
):
    if randomize_seed:
        seed = random.randint(0, MAX_SEED)
    generator = torch.Generator(device=device).manual_seed(int(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 just the image and seed
    return image, seed

# Updated examples with 1880s clothing style
examples = [
    "Cézanne's painting of a lively outdoor gathering in the 1880s, with men in formal top hats, frock coats, and women in bustled dresses with elaborate hats, enjoying a summer afternoon. The scene captures the Belle Époque atmosphere with dappled sunlight filtering through trees, highlighting the fashionable attire of the period. [trigger]",
    "Cézanne's intimate portrait of a young woman from the 1880s, with her hair styled in a fashionable updo, wearing a high-necked dress with lace details and leg-of-mutton sleeves. She wears delicate jewelry and has the soft facial features characteristic of Cézanne's portraiture, set against a background of vibrant colors. [trigger]",
    "Cézanne's painting of two young girls in 1880s attire seated at a piano. One plays while the other stands nearby, both dressed in white frocks with ribbon details, sashes, and high collars typical of the period. The interior setting features Victorian furnishings and decor with warm tones and the distinctive brushwork of Cézanne. [trigger]",
    "Cézanne's painting of an elegant 1880s boating party, with gentlemen in striped boating blazers, straw boater hats, and formal trousers, alongside ladies in bustled day dresses with parasols. The scene captures the leisure activities of French society during the Belle Époque era, with a bright palette emphasizing the fashionable clothing of the period. [trigger]",
    "Cézanne's painting of children playing in an 1880s garden scene, dressed in formal period children's wear including sailor suits for boys and pinafores with full skirts for girls. Their Victorian-era clothing contrasts with their playful activities, set against Cézanne's vibrant treatment of nature and sunlight filtering through the foliage. [trigger]",
    "Cézanne's depiction of bathers in 1880s swimming attire, showing the modest bathing costumes of the period. Women wear full-coverage dark bathing dresses with stockings, while men are in knee-length swimming suits. The figures are arranged in Cézanne's distinctive compositional style against a backdrop of water and landscape rendered in his bold color blocks. [trigger]"
]

# Improved custom CSS with better visuals
custom_css = """
:root {
    --color-primary: #0F4C81;
    --color-secondary: #FF9E6C;
    --background-fill-primary: linear-gradient(to right, #f6f9fc, #edf2f7);
}

footer {
    visibility: hidden;
}

.gradio-container {
    background: var(--background-fill-primary);
}

.title {
    color: var(--color-primary) !important;
    font-size: 3rem !important;
    font-weight: 700 !important;
    text-align: center;
    margin: 1rem 0;
    text-shadow: 2px 2px 4px rgba(0,0,0,0.05);
    font-family: 'Playfair Display', serif;
}

.subtitle {
    color: #4A5568 !important;
    font-size: 1.2rem !important;
    text-align: center;
    margin-bottom: 1.5rem;
    font-style: italic;
}

.collection-link {
    text-align: center;
    margin-bottom: 2rem;
    font-size: 1.1rem;
}

.collection-link a {
    color: var(--color-primary);
    text-decoration: underline;
    transition: color 0.3s ease;
}

.collection-link a:hover {
    color: var(--color-secondary);
}

.model-description {
    background-color: rgba(255, 255, 255, 0.8);
    border-radius: 12px;
    padding: 24px;
    margin: 20px 0;
    box-shadow: 0 4px 12px rgba(0, 0, 0, 0.05);
    border-left: 5px solid var(--color-primary);
}

button.primary {
    background-color: var(--color-primary) !important;
    transition: all 0.3s ease;
}

button:hover {
    transform: translateY(-2px);
    box-shadow: 0 5px 15px rgba(0,0,0,0.1);
}

.input-container {
    border-radius: 10px;
    box-shadow: 0 2px 8px rgba(0,0,0,0.05);
}

.advanced-settings {
    margin-top: 1rem;
    padding: 1rem;
    border-radius: 10px;
    background-color: rgba(255, 255, 255, 0.6);
}

.example-region {
    background-color: rgba(255, 255, 255, 0.5);
    border-radius: 10px;
    padding: 1rem;
    margin-top: 1rem;
}
"""

with gr.Blocks(css=custom_css, analytics_enabled=False) as demo:
    gr.HTML('<div class="title">Paul Cézanne STUDIO</div>')
    
    # Add collection link below title
    gr.HTML('<div class="collection-link"><a href="https://huggingface.co/collections/openfree/painting-art-ai-681453484ec15ef5978bbeb1" target="_blank">View the full Painting Art AI Collection</a></div>')
    
    # Model description with the requested content
    with gr.Group(elem_classes="model-description"):
        gr.HTML('<p>Generate beautiful artwork in the style of Paul Cézanne. Add [trigger] at the end of your prompt for best results.</p>')

    # Simplified structure without tabs and gallery
    with gr.Column(elem_id="col-container"):
        with gr.Row(elem_classes="input-container"):
            prompt = gr.Text(
                label="Prompt",
                max_lines=1,
                placeholder="Enter your prompt (add [trigger] at the end)",
                value=examples[0]  # Set default text instead of generating an image
            )
            run_button = gr.Button("Generate", variant="primary", scale=0)

        result = gr.Image(label="Generated Image")
        seed_output = gr.Number(label="Seed", visible=True)

        with gr.Accordion("Advanced Settings", open=False, elem_classes="advanced-settings"):
            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,
                )

        with gr.Group(elem_classes="example-region"):
            gr.Markdown("### Examples")
            gr.Examples(
                examples=examples,
                inputs=prompt,
                outputs=None,  # Don't auto-run examples
                fn=None,  # No function to run for examples - just fill the prompt
                cache_examples=False,  # Disable caching
            )

    # Event handlers
    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_output],
    )
    
    # No preloading or automatic image generation

demo.queue()
demo.launch()