import gradio as gr
import numpy as np
import random
import spaces
import torch
from diffusers import DiffusionPipeline, FluxPipeline, AutoencoderTiny, AutoencoderKL
from transformers import CLIPTextModel, CLIPTokenizer, T5EncoderModel, T5TokenizerFast
from live_preview_helpers import calculate_shift, retrieve_timesteps, flux_pipe_call_that_returns_an_iterable_of_images
from download import download_all_models, models, download_vaes

# Call the download function at the start of the app
download_all_models()
download_vaes()

dtype = torch.bfloat16
device = "cuda" if torch.cuda.is_available() else "cpu"

# Load VAEs - these can be reused across models
taef1 = AutoencoderTiny.from_pretrained("madebyollin/taef1", torch_dtype=dtype).to(device)
good_vae = AutoencoderKL.from_pretrained("black-forest-labs/FLUX.1-dev", subfolder="vae", torch_dtype=dtype).to(device)

def load_model(model_name):
    model_info = models[model_name]
    pipeline_class = model_info["pipeline_class"]
    model_id = model_info["model_id"]
    config = model_info["config"]

    pipeline = pipeline_class.from_pretrained(model_id, **config, vae=taef1).to(device)
    
    # Assign the custom function for live preview if it's a FluxPipeline or DiffusionPipeline
    if pipeline_class in (FluxPipeline, DiffusionPipeline):
        pipeline.flux_pipe_call_that_returns_an_iterable_of_images = flux_pipe_call_that_returns_an_iterable_of_images.__get__(pipeline)

    return pipeline, model_info.get("description", "No description available.")

# Initialize with default model
current_model_name = "FLUX.1-dev"
pipe, model_description = load_model(current_model_name)
torch.cuda.empty_cache()

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

@spaces.GPU(duration=75)
def infer(model_name, prompt, seed=42, randomize_seed=False, width=1024, height=1024, guidance_scale=3.5, num_inference_steps=28, progress=gr.Progress(track_tqdm=True)):
    global pipe, current_model_name  # Access the global pipe and model name

    if randomize_seed:
        seed = random.randint(0, MAX_SEED)
    generator = torch.Generator().manual_seed(seed)

    # Only reload the model if a different one is selected
    if model_name != current_model_name:
        pipe, _ = load_model(model_name)
        current_model_name = model_name

    for img in pipe.flux_pipe_call_that_returns_an_iterable_of_images(
            prompt=prompt,
            guidance_scale=guidance_scale,
            num_inference_steps=num_inference_steps,
            width=width,
            height=height,
            generator=generator,
            output_type="pil",
            good_vae=good_vae,
        ):
        yield img, seed

examples = [
    ["FLUX.1-dev", "a tiny astronaut hatching from an egg on the moon"],
    ["FLUX.1-schnell", "a cat holding a sign that says hello world"],
    ["Flux.1-lite-8B-alpha", "an anime illustration of a wiener schnitzel"],
]

css="""
#col-container {
    margin: 0 auto;
    max-width: 520px;
}
"""

with gr.Blocks(css=css) as demo:
    with gr.Column(elem_id="col-container"):
        gr.Markdown(f"""# FLUX.1 Model Selector
Select a model to generate images using the FLUX pipeline.
        """)
        model_selector = gr.Dropdown(
            choices=list(models.keys()),
            value=current_model_name,
            label="Select Model",
        )
        model_description_box = gr.Markdown(model_description)
        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=0,
            )
            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=1024,
                )
            with gr.Row():
                guidance_scale = gr.Slider(
                    label="Guidance Scale",
                    minimum=1,
                    maximum=15,
                    step=0.1,
                    value=3.5,
                )
                num_inference_steps = gr.Slider(
                    label="Number of inference steps",
                    minimum=1,
                    maximum=50,
                    step=1,
                    value=28,
                )
        
    def update_description(selected_model):
        return models[selected_model]["description"]

    model_selector.change(
        fn=update_description,
        inputs=[model_selector],
        outputs=[model_description_box],
    )

    gr.Examples(
        examples=examples,
        fn=infer,
        inputs=[model_selector, prompt],  # Correct order of inputs
        outputs=[result, seed],
        cache_examples=False,
    )

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

demo.launch()

#working1 
''''
import gradio as gr
import numpy as np
import random
import spaces
import torch
from diffusers import DiffusionPipeline, FluxPipeline, AutoencoderTiny, AutoencoderKL
from transformers import CLIPTextModel, CLIPTokenizer, T5EncoderModel, T5TokenizerFast
from live_preview_helpers import calculate_shift, retrieve_timesteps, flux_pipe_call_that_returns_an_iterable_of_images

# Define models and their configurations
models = {
    "FLUX.1-schnell": {
        "pipeline_class": FluxPipeline,
        "model_id": "black-forest-labs/FLUX.1-schnell",
        "config": {"torch_dtype": torch.bfloat16},
        "description": "**FLUX.1-schnell** is a fast and efficient model designed for quick image generation. It excels at producing high-quality images rapidly, making it ideal for applications where speed is crucial. However, its rapid generation may slightly compromise on the level of detail compared to slower, more meticulous models.",
    },
    "FLUX.1-dev": {
        "pipeline_class": DiffusionPipeline,
        "model_id": "black-forest-labs/FLUX.1-dev",
        "config": {"torch_dtype": torch.bfloat16},
        "description": "**FLUX.1-dev** is a development model that focuses on delivering highly detailed and artistically rich images.",
    },
    "Flux.1-lite-8B-alpha": {
        "pipeline_class": FluxPipeline,
        "model_id": "Freepik/flux.1-lite-8B-alpha",
        "config": {"torch_dtype": torch.bfloat16},
        "description": "**Flux.1-lite-8B-alpha** is a lightweight model optimized for efficiency and ease of use.",
    },
}

dtype = torch.bfloat16
device = "cuda" if torch.cuda.is_available() else "cpu"

# Load VAEs - these can be reused across models
taef1 = AutoencoderTiny.from_pretrained("madebyollin/taef1", torch_dtype=dtype).to(device)
good_vae = AutoencoderKL.from_pretrained("black-forest-labs/FLUX.1-dev", subfolder="vae", torch_dtype=dtype).to(device)

def load_model(model_name):
    model_info = models[model_name]
    pipeline_class = model_info["pipeline_class"]
    model_id = model_info["model_id"]
    config = model_info["config"]

    pipeline = pipeline_class.from_pretrained(model_id, **config, vae=taef1).to(device)
    
    # Assign the custom function for live preview if it's a FluxPipeline or DiffusionPipeline
    if pipeline_class in (FluxPipeline, DiffusionPipeline):
        pipeline.flux_pipe_call_that_returns_an_iterable_of_images = flux_pipe_call_that_returns_an_iterable_of_images.__get__(pipeline)

    return pipeline, model_info.get("description", "No description available.")

# Initialize with default model
current_model_name = "FLUX.1-dev"
pipe, model_description = load_model(current_model_name)
torch.cuda.empty_cache()

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

@spaces.GPU(duration=75)
def infer(model_name, prompt, seed=42, randomize_seed=False, width=1024, height=1024, guidance_scale=3.5, num_inference_steps=28, progress=gr.Progress(track_tqdm=True)):
    global pipe, current_model_name  # Access the global pipe and model name

    if randomize_seed:
        seed = random.randint(0, MAX_SEED)
    generator = torch.Generator().manual_seed(seed)

    # Only reload the model if a different one is selected
    if model_name != current_model_name:
        pipe, _ = load_model(model_name)
        current_model_name = model_name

    for img in pipe.flux_pipe_call_that_returns_an_iterable_of_images(
            prompt=prompt,
            guidance_scale=guidance_scale,
            num_inference_steps=num_inference_steps,
            width=width,
            height=height,
            generator=generator,
            output_type="pil",
            good_vae=good_vae,
        ):
        yield img, seed

examples = [
    ["FLUX.1-dev", "a tiny astronaut hatching from an egg on the moon"],
    ["FLUX.1-schnell", "a cat holding a sign that says hello world"],
    ["Flux.1-lite-8B-alpha", "an anime illustration of a wiener schnitzel"],
]

css="""
#col-container {
    margin: 0 auto;
    max-width: 520px;
}
"""

with gr.Blocks(css=css) as demo:
    with gr.Column(elem_id="col-container"):
        gr.Markdown(f"""# FLUX.1 Model Selector
Select a model to generate images using the FLUX pipeline.
        """)
        model_selector = gr.Dropdown(
            choices=list(models.keys()),
            value=current_model_name,
            label="Select Model",
        )
        model_description_box = gr.Markdown(model_description)
        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=0,
            )
            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=1024,
                )
            with gr.Row():
                guidance_scale = gr.Slider(
                    label="Guidance Scale",
                    minimum=1,
                    maximum=15,
                    step=0.1,
                    value=3.5,
                )
                num_inference_steps = gr.Slider(
                    label="Number of inference steps",
                    minimum=1,
                    maximum=50,
                    step=1,
                    value=28,
                )
        
    def update_description(selected_model):
        return models[selected_model]["description"]

    model_selector.change(
        fn=update_description,
        inputs=[model_selector],
        outputs=[model_description_box],
    )

    gr.Examples(
        examples=examples,
        fn=infer,
        inputs=[model_selector, prompt],  # Correct order of inputs
        outputs=[result, seed],
        cache_examples=False,
    )

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

demo.launch()

'''
#orginal

'''
import gradio as gr
import numpy as np
import random
import spaces
import torch
from diffusers import  DiffusionPipeline, FlowMatchEulerDiscreteScheduler, AutoencoderTiny, AutoencoderKL
from transformers import CLIPTextModel, CLIPTokenizer,T5EncoderModel, T5TokenizerFast
from live_preview_helpers import calculate_shift, retrieve_timesteps, flux_pipe_call_that_returns_an_iterable_of_images

dtype = torch.bfloat16
device = "cuda" if torch.cuda.is_available() else "cpu"

taef1 = AutoencoderTiny.from_pretrained("madebyollin/taef1", torch_dtype=dtype).to(device)
good_vae = AutoencoderKL.from_pretrained("black-forest-labs/FLUX.1-dev", subfolder="vae", torch_dtype=dtype).to(device)
pipe = DiffusionPipeline.from_pretrained("black-forest-labs/FLUX.1-dev", torch_dtype=dtype, vae=taef1).to(device)
torch.cuda.empty_cache()

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

pipe.flux_pipe_call_that_returns_an_iterable_of_images = flux_pipe_call_that_returns_an_iterable_of_images.__get__(pipe)

@spaces.GPU(duration=75)
def infer(prompt, seed=42, randomize_seed=False, width=1024, height=1024, guidance_scale=3.5, num_inference_steps=28, progress=gr.Progress(track_tqdm=True)):
    if randomize_seed:
        seed = random.randint(0, MAX_SEED)
    generator = torch.Generator().manual_seed(seed)
    
    for img in pipe.flux_pipe_call_that_returns_an_iterable_of_images(
            prompt=prompt,
            guidance_scale=guidance_scale,
            num_inference_steps=num_inference_steps,
            width=width,
            height=height,
            generator=generator,
            output_type="pil",
            good_vae=good_vae,
        ):
            yield img, seed
    
examples = [
    "a tiny astronaut hatching from an egg on the moon",
    "a cat holding a sign that says hello world",
    "an anime illustration of a wiener schnitzel",
]

css="""
#col-container {
    margin: 0 auto;
    max-width: 520px;
}
"""

with gr.Blocks(css=css) as demo:
    
    with gr.Column(elem_id="col-container"):
        gr.Markdown(f"""# FLUX.1 [dev]
12B param rectified flow transformer guidance-distilled from [FLUX.1 [pro]](https://blackforestlabs.ai/)  
[[non-commercial license](https://huggingface.co./black-forest-labs/FLUX.1-dev/blob/main/LICENSE.md)] [[blog](https://blackforestlabs.ai/announcing-black-forest-labs/)] [[model](https://huggingface.co./black-forest-labs/FLUX.1-dev)]
        """)
        
        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=0,
            )
            
            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=1024,
                )
            
            with gr.Row():

                guidance_scale = gr.Slider(
                    label="Guidance Scale",
                    minimum=1,
                    maximum=15,
                    step=0.1,
                    value=3.5,
                )
  
                num_inference_steps = gr.Slider(
                    label="Number of inference steps",
                    minimum=1,
                    maximum=50,
                    step=1,
                    value=28,
                )
        
        gr.Examples(
            examples = examples,
            fn = infer,
            inputs = [prompt],
            outputs = [result, seed],
            cache_examples="lazy"
        )

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

demo.launch()

'''




'''
import os
import torch
import gradio as gr
from diffusers import FluxPipeline, DiffusionPipeline
import spaces
# Helper function to get the Hugging Face token securely
def get_hf_token():
    try:
        from google.colab import userdata
        hf_token = userdata.get('HF_TOKEN')
        if hf_token:
            return hf_token
        else:
            raise RuntimeError("HF_TOKEN not found in Colab secrets.")
    except ImportError:
        return os.getenv("HF_TOKEN", None)

# Securely get the token
_HF_TOKEN = get_hf_token()
if not _HF_TOKEN:
    raise ValueError("HF_TOKEN is not available. Please set it in Colab secrets or environment variables.")

# Define models and their configurations
models = {
    "FLUX.1-schnell": {
        "pipeline_class": FluxPipeline,
        "model_id": "black-forest-labs/FLUX.1-schnell",
        "config": {"torch_dtype": torch.bfloat16},
        "description": "**FLUX.1-schnell** is a fast and efficient model designed for quick image generation. It excels at producing high-quality images rapidly, making it ideal for applications where speed is crucial. However, its rapid generation may slightly compromise on the level of detail compared to slower, more meticulous models.",
    },
    "FLUX.1-dev": {
        "pipeline_class": DiffusionPipeline,
        "model_id": "black-forest-labs/FLUX.1-dev",
        "lora": {
            "repo": "strangerzonehf/Flux-Enrich-Art-LoRA",
            "trigger_word": "enrich art",
        },
        "config": {"torch_dtype": torch.bfloat16},
        "description": "**FLUX.1-dev** is a development model that focuses on delivering highly detailed and artistically rich images.",
    },
    "Flux.1-lite-8B-alpha": {
        "pipeline_class": FluxPipeline,
        "model_id": "Freepik/flux.1-lite-8B-alpha",
        "config": {"torch_dtype": torch.bfloat16},
        "description": "**Flux.1-lite-8B-alpha** is a lightweight model optimized for efficiency and ease of use.",
    },
}

models = {
    "FLUX.1-schnell": {
        "pipeline_class": FluxPipeline,
        "model_id": "black-forest-labs/FLUX.1-schnell",
        "config": {"torch_dtype": torch.bfloat16},
        "description": "**FLUX.1-schnell** is a fast and efficient model designed for quick image generation. It excels at producing high-quality images rapidly, making it ideal for applications where speed is crucial. However, its rapid generation may slightly compromise on the level of detail compared to slower, more meticulous models.",
    },
}



# Function to pre-download models
def download_all_models():
    print("Downloading all models...")
    for model_key, config in models.items():
        try:
            pipeline_class = config["pipeline_class"]
            model_id = config["model_id"]
            # Attempt to download the pipeline without loading it into memory
            pipeline_class.from_pretrained(model_id, token=_HF_TOKEN, **config.get("config", {}))
            if "lora" in config:
                pipeline_class.download_lora_weights(config["lora"]["repo"],token=_HF_TOKEN,)
            print(f"Model '{model_key}' downloaded successfully.")
        except Exception as e:
            print(f"Error downloading model '{model_key}': {e}")
    print("Model download process complete.")

loaded_models = {}
model_load_status = {}  # Dictionary to track model load status

def clear_gpu_memory():
    """Clears GPU memory. Keeps model status information."""
    global loaded_models
    try:
        for model_key in list(loaded_models.keys()):  # Iterate over a copy to allow deletion
            del loaded_models[model_key]
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
            torch.cuda.ipc_collect()
        print("GPU memory cleared.")
        return "GPU memory cleared."
    except Exception as e:
        print(f"Error clearing GPU memory: {e}")
        return f"Error clearing GPU memory: {e}"


def load_model(model_key):
    """Loads a model, clearing GPU memory first if a different model was loaded."""
    global model_load_status

    if model_key not in models:
        model_load_status[model_key] = "Model not found."
        return f"Model '{model_key}' not found in the available models."

    # Clear GPU memory only if a different model is already loaded
    if loaded_models and list(loaded_models.keys())[0] != model_key:
        clear_gpu_memory()

    try:
        config = models[model_key]
        pipeline_class = config["pipeline_class"]
        model_id = config["model_id"]
        
        pipe = pipeline_class.from_pretrained(model_id, token=_HF_TOKEN, **config.get("config", {}))

        if "lora" in config:
            lora_config = config["lora"]
            pipe.load_lora_weights(lora_config["repo"], token=_HF_TOKEN)

        if torch.cuda.is_available():
            pipe.to("cuda")

        loaded_models[model_key] = pipe
        model_load_status[model_key] = "Loaded" # Update load status
        return f"Model '{model_key}' loaded successfully."
    except Exception as e:
        model_load_status[model_key] = "Failed" # Update load status on error
        return f"Error loading model '{model_key}': {e}"
@spaces.GPU()
def generate_image(model, prompt, seed=-1):
    
    generator = torch.Generator(device="cuda" if torch.cuda.is_available() else "cpu")
    if seed != -1:
        generator = generator.manual_seed(seed)

    with torch.no_grad():
        image = model(prompt=prompt, generator=generator).images[0]
    return image

def gradio_generate(selected_model, prompt, seed):
    
    if selected_model not in loaded_models:
      if selected_model in model_load_status and model_load_status[selected_model] == "Loaded":
          # Model should be loaded but isn't in loaded_models, clear it from the status
          del model_load_status[selected_model]
      if selected_model not in model_load_status or model_load_status[selected_model] != "Loaded":
          # Attempt to load the model if not already attempted or failed
          load_model(selected_model)
      
      if selected_model not in loaded_models:
          # If still not loaded after attempt, return an error
          return f"Model not loaded. Load status: {model_load_status.get(selected_model, 'Not attempted')}.", None

    model = loaded_models[selected_model]
    image = generate_image(model, prompt, seed)

    runtime_info = f"Model: {selected_model}\nSeed: {seed}"
    output_path = "generated_image.png"
    image.save(output_path)
    return output_path, runtime_info

def gradio_load_model(selected_model):
    if not selected_model:
        return "No model selected. Please select a model to load."
    return load_model(selected_model)



import gradio as gr

with gr.Blocks(
    css="""
    .container {
        max-width: 800px;
        margin: auto;
        padding: 20px;
        background-color: #ffffff; /* Pure white for a clean background */
        border-radius: 10px; /* Smooth rounded corners */
        box-shadow: 0 4px 8px rgba(0, 0, 0, 0.05); /* Subtle shadow for a lighter feel */
        color: #333333; /* Dark gray text for good contrast and readability */
    }



    """
) as interface:
    with gr.Tab("Image Generator"):
        with gr.Column():
            gr.Markdown("# Text-to-Image Generator")
            model_dropdown = gr.Dropdown(choices=list(models.keys()), label="Select Model")
            prompt_textbox = gr.Textbox(label="Enter Text Prompt")
            seed_slider = gr.Slider(minimum=-1, maximum=1000, step=1, value=-1, label="Random Seed (-1 for random)")
            generate_button = gr.Button("Generate Image")
            output_image = gr.Image(label="Generated Image")
            runtime_info_textbox = gr.Textbox(label="Runtime Information", lines=2, interactive=False)

            # Add example prompts at the bottom
            gr.Markdown("### Example Prompts")
            examples = gr.Examples(
                examples=[
                    [list(models.keys())[0], "Sexy Woman", "sample2.png"],
                   # [list(models.keys())[2], "Sexy girl", "sample3.png"],
                   # [list(models.keys())[1], "Future City", "sample1.png"]
                ],
                inputs=[model_dropdown, prompt_textbox, output_image],
            )

            generate_button.click(gradio_generate, inputs=[model_dropdown, prompt_textbox, seed_slider], outputs=[output_image, runtime_info_textbox])

    with gr.Tab("Model Information"):
        for model_key, model_info in models.items():
            gr.Markdown(f"## {model_key}")
            gr.Markdown(model_info["description"])

    gr.Markdown("""---
    **Credits**: Created by Ruslan Magana Vsevolodovna. For more information, visit [https://ruslanmv.com/](https://ruslanmv.com/).""")

if __name__ == "__main__":

    # Pre-download all models at startup
    download_all_models()      
    interface.launch(debug=True)

'''



'''

import os
import torch
import gradio as gr
from diffusers import FluxPipeline, DiffusionPipeline
import spaces

# Helper function to get the Hugging Face token securely
def get_hf_token():
    try:
        from google.colab import userdata
        hf_token = userdata.get('HF_TOKEN')
        if hf_token:
            return hf_token
        else:
            raise RuntimeError("HF_TOKEN not found in Colab secrets.")
    except ImportError:
        return os.getenv("HF_TOKEN", None)

# Securely get the token
_HF_TOKEN = get_hf_token()
if not _HF_TOKEN:
    raise ValueError("HF_TOKEN is not available. Please set it in Colab secrets or environment variables.")

# Define models and their configurations
models = {
    "FLUX.1-schnell": {
        "pipeline_class": FluxPipeline,
        "model_id": "black-forest-labs/FLUX.1-schnell",
        "config": {"torch_dtype": torch.bfloat16},
        "description": "**FLUX.1-schnell** is a fast and efficient model designed for quick image generation. It excels at producing high-quality images rapidly, making it ideal for applications where speed is crucial. However, its rapid generation may slightly compromise on the level of detail compared to slower, more meticulous models.",
    },
    "FLUX.1-dev": {
        "pipeline_class": DiffusionPipeline,
        "model_id": "black-forest-labs/FLUX.1-dev",
        "lora": {
            "repo": "strangerzonehf/Flux-Enrich-Art-LoRA",
            "trigger_word": "enrich art",
        },
        "config": {"torch_dtype": torch.bfloat16},
        "description": "**FLUX.1-dev** is a development model that focuses on delivering highly detailed and artistically rich images.",
    },
    "Flux.1-lite-8B-alpha": {
        "pipeline_class": FluxPipeline,
        "model_id": "Freepik/flux.1-lite-8B-alpha",
        "config": {"torch_dtype": torch.bfloat16},
        "description": "**Flux.1-lite-8B-alpha** is a lightweight model optimized for efficiency and ease of use.",
    },
}


# Function to pre-download models
def download_all_models():
    print("Downloading all models...")
    for model_key, config in models.items():
        try:
            pipeline_class = config["pipeline_class"]
            model_id = config["model_id"]
            # Attempt to download the pipeline without loading it into memory
            pipeline_class.from_pretrained(model_id, token=_HF_TOKEN, **config.get("config", {}))
            if "lora" in config:
                pipeline_class.download_lora_weights(config["lora"]["repo"],token=_HF_TOKEN,)
            print(f"Model '{model_key}' downloaded successfully.")
        except Exception as e:
            print(f"Error downloading model '{model_key}': {e}")
    print("Model download process complete.")


loaded_models = {}
model_load_status = {}  # Dictionary to track model load status

def clear_gpu_memory():
    """Clears GPU memory. Keeps model status information."""
    global loaded_models
    try:
        for model_key in list(loaded_models.keys()):  # Iterate over a copy to allow deletion
            del loaded_models[model_key]
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
            torch.cuda.ipc_collect()
        print("GPU memory cleared.")
        return "GPU memory cleared."
    except Exception as e:
        print(f"Error clearing GPU memory: {e}")
        return f"Error clearing GPU memory: {e}"


def load_model(model_key):
    """Loads a model, clearing GPU memory first if a different model was loaded."""
    global model_load_status

    if model_key not in models:
        model_load_status[model_key] = "Model not found."
        return f"Model '{model_key}' not found in the available models."

    # Clear GPU memory only if a different model is already loaded
    if loaded_models and list(loaded_models.keys())[0] != model_key:
        clear_gpu_memory()

    try:
        config = models[model_key]
        pipeline_class = config["pipeline_class"]
        model_id = config["model_id"]
        
        pipe = pipeline_class.from_pretrained(model_id, token=_HF_TOKEN, **config.get("config", {}))

        if "lora" in config:
            lora_config = config["lora"]
            pipe.load_lora_weights(lora_config["repo"], token=_HF_TOKEN)

        if torch.cuda.is_available():
            pipe.to("cuda")

        loaded_models[model_key] = pipe
        model_load_status[model_key] = "Loaded" # Update load status
        return f"Model '{model_key}' loaded successfully."
    except Exception as e:
        model_load_status[model_key] = "Failed" # Update load status on error
        return f"Error loading model '{model_key}': {e}"

@spaces.GPU()
def generate_image(model, prompt, seed=-1):
    
    generator = torch.Generator(device="cuda" if torch.cuda.is_available() else "cpu")
    if seed != -1:
        generator = generator.manual_seed(seed)

    with torch.no_grad():
        image = model(prompt=prompt, generator=generator).images[0]
    return image

def gradio_generate(selected_model, prompt, seed):
    
    if selected_model not in loaded_models:
        if selected_model in model_load_status and model_load_status[selected_model] == "Loaded":
            # Model should be loaded but isn't in loaded_models, clear it from the status
            del model_load_status[selected_model]
        if selected_model not in model_load_status or model_load_status[selected_model] != "Loaded":
            # Attempt to load the model if not already attempted or failed
            load_model(selected_model)
    
    if selected_model not in loaded_models:
        # If still not loaded after attempt, return an error
        return f"Model not loaded. Load status: {model_load_status.get(selected_model, 'Not attempted')}.", None

    model = loaded_models[selected_model]
    image = generate_image(model, prompt, seed)

    runtime_info = f"Model: {selected_model}\nSeed: {seed}"
    output_path = "generated_image.png"
    image.save(output_path)
    return output_path, runtime_info

def gradio_load_model(selected_model):
    if not selected_model:
        return "No model selected. Please select a model to load."
    return load_model(selected_model)


with gr.Blocks(
    css="""
    .container {
        max-width: 800px;
        margin: auto;
        padding: 20px;
        background-color: #f8f8f8;
        border-radius: 10px;
        box-shadow: 0 4px 8px rgba(0, 0, 0, 0.1);
    }
    """
) as interface:
    with gr.Tab("Image Generator"):
        with gr.Column():
            gr.Markdown("# Text-to-Image Generator")
            model_dropdown = gr.Dropdown(choices=list(models.keys()), label="Select Model")
            prompt_textbox = gr.Textbox(label="Enter Text Prompt")
            seed_slider = gr.Slider(minimum=-1, maximum=1000, step=1, value=-1, label="Random Seed (-1 for random)")
            generate_button = gr.Button("Generate Image")
            output_image = gr.Image(label="Generated Image")
            runtime_info_textbox = gr.Textbox(label="Runtime Information", lines=2, interactive=False)

            # Add example prompts at the bottom
            gr.Markdown("### Example Prompts")
            examples = gr.Examples(
                examples=[
                    [list(models.keys())[0], "Sexy Woman", "sample2.png"],
                    [list(models.keys())[2], "Sexy girl", "sample3.png"],
                    [list(models.keys())[1], "Future City", "sample1.png"]
                ],
                inputs=[model_dropdown, prompt_textbox, output_image],
            )

            generate_button.click(gradio_generate, inputs=[model_dropdown, prompt_textbox, seed_slider], outputs=[output_image, runtime_info_textbox])

    with gr.Tab("Model Information"):
        for model_key, model_info in models.items():
            gr.Markdown(f"## {model_key}")
            gr.Markdown(model_info["description"])

    gr.Markdown("""---
    **Credits**: Created by Ruslan Magana Vsevolodovna. For more information, visit [https://ruslanmv.com/](https://ruslanmv.com/).""")

if __name__ == "__main__":

    # Pre-download all models at startup
    download_all_models()

    # Load only the first model
    load_model(list(models.keys())[0])
    
    interface.launch(debug=True)
'''

'''
with gr.Blocks(
    css="""
    .container {
        max-width: 800px;
        margin: auto;
        padding: 20px;
        background-color: #f8f8f8;
        border-radius: 10px;
        box-shadow: 0 4px 8px rgba(0, 0, 0, 0.1);
    }
    """
) as interface:
    with gr.Tab("Image Generator"):
        with gr.Column():
            gr.Markdown("# Text-to-Image Generator")
            model_dropdown = gr.Dropdown(choices=list(models.keys()), label="Select Model")
            prompt_textbox = gr.Textbox(label="Enter Text Prompt")
            seed_slider = gr.Slider(minimum=-1, maximum=1000, step=1, value=-1, label="Random Seed (-1 for random)")
            generate_button = gr.Button("Generate Image")
            output_image = gr.Image(label="Generated Image")
            runtime_info_textbox = gr.Textbox(label="Runtime Information", lines=2, interactive=False)

            # Add example prompts at the bottom
            gr.Markdown("### Example Prompts")
            examples = gr.Examples(
                examples=[
                    [list(models.keys())[0], "Sexy girl"],
                    [list(models.keys())[0], "Beautiful Woman"],
                    [list(models.keys())[0], "Future City"]
                ],
                inputs=[model_dropdown, prompt_textbox],
            )

            generate_button.click(gradio_generate, inputs=[model_dropdown, prompt_textbox, seed_slider], outputs=[output_image, runtime_info_textbox])

    with gr.Tab("Model Information"):
        for model_key, model_info in models.items():
            gr.Markdown(f"## {model_key}")
            gr.Markdown(model_info["description"])

    gr.Markdown("""---
    **Credits**: Created by Ruslan Magana Vsevolodovna. For more information, visit [https://ruslanmv.com/](https://ruslanmv.com/).""")

if __name__ == "__main__":
    interface.launch(debug=True)
'''

'''
with gr.Blocks(
    css="""
    .container {
        max-width: 800px;
        margin: auto;
        padding: 20px;
        background-color: #f8f8f8;
        border-radius: 10px;
        box-shadow: 0 4px 8px rgba(0, 0, 0, 0.1);
    }
    """
) as interface:
    with gr.Tab("Image Generator"):
        with gr.Column():
            gr.Markdown("# Text-to-Image Generator")
            model_dropdown = gr.Dropdown(choices=list(models.keys()), label="Select Model")
            #with gr.Row():
                #load_button = gr.Button("Load Model")
                # clear_button = gr.Button("Clear GPU Memory") # Removed clear button
            #load_status = gr.Textbox(label="Model Load Status", interactive=False) # Removed load button
            prompt_textbox = gr.Textbox(label="Enter Text Prompt")
            seed_slider = gr.Slider(minimum=-1, maximum=1000, step=1, value=-1, label="Random Seed (-1 for random)")
            generate_button = gr.Button("Generate Image")
            output_image = gr.Image(label="Generated Image")
            runtime_info_textbox = gr.Textbox(label="Runtime Information", lines=2, interactive=False)

            #load_button.click(gradio_load_model, inputs=[model_dropdown], outputs=[load_status])   # Removed load button click action
            # clear_button.click(clear_gpu_memory, outputs=[load_status]) # Removed clear button click action
            generate_button.click(gradio_generate, inputs=[model_dropdown, prompt_textbox, seed_slider], outputs=[output_image, runtime_info_textbox])

    with gr.Tab("Model Information"):
        for model_key, model_info in models.items():
            gr.Markdown(f"## {model_key}")
            gr.Markdown(model_info["description"])

    gr.Markdown("""---
    **Credits**: Created by Ruslan Magana Vsevolodovna. For more information, visit [https://ruslanmv.com/](https://ruslanmv.com/).""")

if __name__ == "__main__":
    interface.launch(debug=True)
'''