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/winslow-homer"  # Changed to Winslow Homer model

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 for Winslow Homer
examples = [
    "Winslow Homer'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 Homer's characteristic realism and attention to natural light, highlighting the fashionable attire of the period. [trigger]",
    "Winslow Homer'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, naturalistic features characteristic of Homer's portraiture, set against a coastal backdrop with dramatic light and shadow. [trigger]",
    "Winslow Homer's painting of two young girls in 1880s attire seated near a window. One plays with a toy 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 natural sunlight streaming in, creating Homer's distinctive interplay of light and shadow. [trigger]",
    "Winslow Homer'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 near the water, with Homer's masterful depiction of reflections and maritime atmosphere. [trigger]",
    "Winslow Homer's painting of children playing in an 1880s seaside 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 Homer's dramatic seascapes with crashing waves and atmospheric skies. [trigger]",
    "Winslow Homer's depiction of figures in 1880s attire by a rugged coastline, showing the modest clothing of the period. Women wear full-coverage dark dresses with white details, while men are in work clothes or formal attire. The figures are arranged against a backdrop of dramatic ocean waves and weather, showcasing Homer's skill at capturing the relationship between humans and nature. [trigger]"
]

# Improved custom CSS with better visuals - updated colors for Winslow Homer theme
custom_css = """
:root {
    --color-primary: #2B5F75;
    --color-secondary: #D98542;
    --background-fill-primary: linear-gradient(to right, #e8f4f8, #f4e8d8);
}

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: 'Georgia', 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;
    border: 1px solid #d0e0e3;
}
"""

with gr.Blocks(css=custom_css, analytics_enabled=False) as demo:
    gr.HTML('<div class="title">Winslow Homer 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 Winslow Homer, the renowned American landscape painter and printmaker. Homer is known for his marine subjects, dramatic seascapes, and realistic depictions of rural and coastal life. 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()