Blame for app.py

import shutil
from pathlib import Path
from datetime import datetime

import spaces
import gradio as gr

from src.generate import GenerateConfig, generate_one
from src.styles import load_style

PROJECT_ROOT = Path(__file__).resolve().parent
INPUT_DIR = PROJECT_ROOT / "input_images"
OUTPUT_DIR = PROJECT_ROOT / "output_images"

INPUT_DIR.mkdir(exist_ok=True)
OUTPUT_DIR.mkdir(exist_ok=True)


def save_reference_images(style_name: str, files) -> Path:
    style_name = style_name.strip()
    if not style_name:
        raise ValueError("Style folder name is required.")

    style_dir = INPUT_DIR / style_name
    style_dir.mkdir(parents=True, exist_ok=True)

    # Clear old files so each run uses only the newly uploaded references
    for item in style_dir.iterdir():
        if item.is_file():
            item.unlink()

    saved = 0
    for file_path in files or []:
        src = Path(file_path)
        if src.exists():
            shutil.copy2(src, style_dir / src.name)
            saved += 1

    if saved == 0:
        raise ValueError("Please upload at least one reference image.")

    return style_dir


@spaces.GPU(duration=180)
def run_generation(prompt, files, steps, guidance, width, height, seed):
    if not prompt or not prompt.strip():
        raise ValueError("Prompt is required.")

    style_name = f"style_{datetime.now().strftime('%Y%m%d_%H%M%S')}"

    save_reference_images(style_name, files)
    style = load_style(INPUT_DIR, style_name)

    seed_value = None if seed in ("", None) else int(seed)

    cfg = GenerateConfig(
        model_id="black-forest-labs/FLUX.2-klein-4B",
        steps=int(steps),
        guidance=float(guidance),
        width=int(width),
        height=int(height),
        device="cuda",
        cpu_offload=True,
        use_ref_images=True,
        ref_max_images=6,
        ref_sample_mode="random",
        seed=seed_value,
        num_outputs=1,
    )

    paths = generate_one(
        user_prompt=prompt.strip(),
        style=style,
        out_root=OUTPUT_DIR,
        cfg=cfg,
    )

    return str(paths[0])


with gr.Blocks(title="StyleX") as demo:
    gr.Markdown("# StyleX")
    gr.Markdown(
        "Generate a stylistic image from a text prompt and uploaded reference images."
    )

    with gr.Tabs():
        with gr.Tab("Generate"):
            with gr.Row():
                with gr.Column(scale=2):
                    gr.Markdown("**Prompt**")
                    gr.Markdown(
                        "Describe the image you want the model to create. "
                        "Be specific about the subject, setting, mood, or details you want included."
                    )

                    prompt = gr.Textbox(
                        lines=3,
                        placeholder="Example: A portrait of an Indian woman AI scientist working at a futuristic desk"
                    )

                    with gr.Accordion("Example Prompts", open=False):
                        gr.Markdown(
                            """
Try prompts like:

- **A cyberpunk city street at night with neon lights and rain**
- **A fantasy warrior standing in front of a glowing portal**
- **A futuristic AI scientist working in a high-tech laboratory**
- **A peaceful anime-style village surrounded by mountains**
- **A robot sitting in a cozy library reading a book**
- **A dramatic portrait of a character in a cinematic lighting style**

For best results, describe:
- the subject
- the setting
- the mood
- the lighting
- any important visual details
                            """
                        )

                    gr.Markdown("**Reference Images**")
                    gr.Markdown(
                        "Upload **3–10 images** that all share a similar visual style. "
                        "Around **6 images** is usually the sweet spot for best results."
                    )

                    files = gr.File(
                        file_count="multiple",
                        file_types=["image"],
                        label="Reference Images"
                    )

                    with gr.Accordion("Tips for Better Results", open=False):
                        gr.Markdown(
                            """
### Tips for Stronger Style Matching

- Use reference images that share the **same visual style**
- Avoid mixing very different styles, such as anime, realism, and pixel art together
- Around **6 reference images** usually gives a good balance
- Use clear prompts that describe what you want the final image to contain
- Try multiple generations if the first result does not match the style well
- Adjust **steps** and **guidance** to experiment with quality and style control
- Lower resolutions like **512×512** usually generate faster
- Higher width and height settings may take longer and use more GPU resources

### Prompt Tip

Instead of writing:

`girl in city`

Try writing:

`A cinematic portrait of a girl standing in a futuristic neon city at night, rainy atmosphere, glowing signs, detailed lighting`
                            """
                        )

                    gr.Markdown("**Steps**")
                    gr.Markdown(
                        "Controls how many generation steps the model uses. "
                        "More steps can improve quality, but they also take longer."
                    )

                    steps = gr.Slider(
                        1,
                        10,
                        value=6,
                        step=1,
                        label="Steps"
                    )

                    gr.Markdown("**Guidance**")
                    gr.Markdown(
                        "Controls how strongly the model follows the prompt and style conditioning. "
                        "Higher values can make results more controlled, while lower values may be more flexible."
                    )

                    guidance = gr.Slider(
                        1.0,
                        5.0,
                        value=1.0,
                        step=0.1,
                        label="Guidance"
                    )

                    gr.Markdown("**Seed (optional)**")
                    gr.Markdown(
                        "Use a seed if you want more reproducible results. "
                        "Leave it blank for a random result each time."
                    )

                    seed = gr.Textbox(
                        label="Seed",
                        placeholder="Leave blank for random"
                    )

                    with gr.Row():
                        width = gr.Slider(
                            256,
                            1024,
                            value=512,
                            step=64,
                            label="Width"
                        )

                        height = gr.Slider(
                            256,
                            1024,
                            value=512,
                            step=64,
                            label="Height"
                        )

                    run_btn = gr.Button("Generate Image")
                    output = gr.Image(label="Generated Image")

                    run_btn.click(
                        fn=run_generation,
                        inputs=[prompt, files, steps, guidance, width, height, seed],
                        outputs=output,
                    )

                with gr.Column(scale=1):
                    gr.Markdown(
                        """
## About the Project

**StyleX** is an AI-powered image generation system that creates images using:

- a **text prompt**
- a set of **reference images** that represent a desired style

Instead of generating only from text, StyleX uses uploaded reference images to help guide the visual style of the final output.

---

## How Generation Works

1. You enter a prompt describing the image you want.
2. You upload reference images that share a similar visual style.
3. The system loads the uploaded images as style references.
4. The FLUX model uses the prompt and reference images to guide generation.
5. A new image is created and displayed in the interface.

---

## Best Practices

- Use **3–10** reference images
- Around **6 images** usually works best
- Keep reference images visually consistent
- Use detailed prompts
- Start with **512×512** for faster generation

---

## Disclaimer!

StyleX may need multiple generations to produce the best style match.

Because AI image generation is probabilistic, the first result may not always perfectly capture the desired style. You may need to:

- generate more than once
- adjust the prompt
- change the reference images
- change the steps or guidance values

Results can vary depending on the prompt, reference image quality, and selected settings.

---

## Notes

- If no seed is provided, generation will be random
- Results may vary between runs
- Generation may take time depending on settings
- Hugging Face Spaces may limit daily usage depending on available credits
                        """
                    )

        with gr.Tab("Project Info"):
            gr.Markdown(
                """
# StyleX Project Information

## Overview

**StyleX** is an AI image generation project focused on producing new images from a user prompt while conditioning the output on user-provided reference images.

The goal of StyleX is to combine:

- **user intent** from the text prompt
- **visual style** from the uploaded reference images

This allows the system to generate images that are influenced by both what the user describes and the style examples they provide.

---

## Core Idea

Traditional text-to-image systems generate images mainly from a written prompt.

StyleX expands on this by also using reference images so the output can better match a target visual style.

For example, a user can upload images with an anime, cyberpunk, fantasy, or painterly style, then enter a prompt describing a new subject or scene.

---

## System Workflow

1. The user enters a prompt.
2. The user uploads reference images.
3. StyleX creates a temporary style folder.
4. The uploaded images are saved into that folder.
5. The system loads the selected style references.
6. A `GenerateConfig` object stores generation settings.
7. The FLUX model pipeline runs the generation process.
8. The generated image is saved and returned to the UI.

---

## Architecture Overview

StyleX follows a layered structure:

### Presentation Layer

The Gradio interface allows users to enter prompts, upload images, adjust settings, and view generated results.

### Application Layer

The application controls the workflow, validates inputs, creates style folders, and prepares the generation configuration.

### Service Layer

The generation engine handles style loading, reference image conditioning, and image generation.

### Infrastructure Layer

StyleX uses tools such as:

- PyTorch
- Diffusers
- Transformers
- Pillow
- Hugging Face Spaces GPU support

---

## Model Choice

StyleX uses:

**FLUX.2 Klein**

Model ID:

`black-forest-labs/FLUX.2-klein-4B`

FLUX was selected because it provides stronger visual coherence and better style adherence compared to earlier Stable Diffusion-based testing.

---

## Why FLUX?

FLUX was chosen because StyleX focuses heavily on style representation.

Compared to the earlier Stable Diffusion approach, FLUX provides:

- stronger style matching
- better prompt alignment
- more coherent outputs
- improved visual quality
- better handling of complex scenes

---

## Trade-offs

Although FLUX provides better results, it also has some limitations:

- it is heavier than older Stable Diffusion models
- it uses more VRAM
- it can be slower on limited hardware
- CPU generation is not recommended
- cloud GPU access may have usage limits

---

## Performance Recommendations

Recommended settings:

- **Steps:** 6–10
- **Guidance:** 1.0–3.0
- **Resolution:** 512×512
- **Reference Images:** around 6 images
- **Device:** GPU recommended (Hugging Face will handle this)

Higher settings may improve output quality but can increase generation time.

---

## Features

- Prompt-based image generation
- Multi-image style conditioning
- Uploaded reference image support
- Random or seeded generation
- Adjustable steps
- Adjustable guidance
- Adjustable width and height
- Automatic temporary style folder creation
- Simple Gradio interface
- GPU-based execution through Hugging Face Spaces

---

## Limitations

StyleX is still experimental.

Some current limitations include:

- style matching may not be perfect every time
- mixed reference styles can confuse the model
- more reference images may increase generation time
- higher resolutions require more resources
- Hugging Face usage may be limited by available credits
- results may require multiple attempts

---

## Future Improvements

Possible future improvements include:

- LoRA support for stronger style control
- CLIP-based similarity scoring
- multiple candidate outputs
- automatic best-image selection
- more diverse reference image sampling
- better style analysis
- improved UI design
- ability to view previous generation attempts

---

## Purpose

The purpose of StyleX is to explore how AI image generation can combine a user's creative prompt with reference-based visual style conditioning.

This project demonstrates how modern image generation systems can move beyond simple text prompts and allow users to provide examples of the style they want.
                """
            )


if __name__ == "__main__":
    demo.launch(ssr_mode=False)