import gradio as gr
import numpy as np
import random
from diffusers import DiffusionPipeline
import torch
import time
import psutil
from huggingface_hub import snapshot_download # New import

# ... (Other imports and variables remain the same)

def load_pipeline(model_id):
    if model_id in PIPELINES:
        return PIPELINES[model_id]
    else:
        # Download model using snapshot_download to handle LFS files
        model_path = snapshot_download(repo_id=model_id, local_dir="./models")  # Download and cache models
        pipe = DiffusionPipeline.from_pretrained(model_path, torch_dtype=torch.float32) 
        pipe.to(DEVICE)
        PIPELINES[model_id] = pipe
        return pipe


# ... (Rest of your code remains the same)


# Get the number of physical CPU cores (excluding hyperthreads)
NUM_CPU_CORES = psutil.cpu_count(logical=True)

# Cap the number of threads to the available physical cores
MAX_THREADS = min(8, NUM_CPU_CORES)

# Device and hardware configuration
DEVICE = "cpu"

# Model Options (optimized for CPU and memory constraints)
MODEL_OPTIONS = {
    "Модель штучного інтелекту середня (Довше-краще якість)": "CompVis/stable-diffusion-v1-4",
    "Модель штучного інтелекту мала (Швидко-гірша якість)": "hf-internal-testing/tiny-stable-diffusion-pipe",
}

# Default to fastest model and lower image size for limited resources
DEFAULT_MODEL_ID = MODEL_OPTIONS["Модель штучного інтелекту мала (Швидко-гірша якість)"]
DEFAULT_IMAGE_SIZE = 512  # Lower default resolution

# Cache models to avoid reloading
PIPELINES = {}

def load_pipeline(model_id):
    if model_id in PIPELINES:
        return PIPELINES[model_id]
    else:
        pipe = DiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float32)
        pipe.to(DEVICE)
        PIPELINES[model_id] = pipe
        return pipe

def generate_image(prompt, negative_prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps, num_images, model_choice):
    if not prompt:
        raise gr.Error("Будь ласка, введіть опис для зображення.")

    torch.set_num_threads(MAX_THREADS) 

    pipe = load_pipeline(MODEL_OPTIONS[model_choice])

    # Adjust memory usage based on available RAM
    torch.cuda.empty_cache()  

    generator = torch.Generator(device=DEVICE)
    if not randomize_seed:
        generator = generator.manual_seed(seed)

    start_time = time.time()
    images = pipe(
        prompt,
        negative_prompt=negative_prompt,
        width=width,
        height=height,
        guidance_scale=guidance_scale,
        num_inference_steps=num_inference_steps,
        num_images_per_prompt=num_images,
        generator=generator,
    ).images

    end_time = time.time()
    generation_time = end_time - start_time

    return images, f"Час генерації: {generation_time:.2f} секунд"

with gr.Blocks() as demo:
    with gr.Row():
        with gr.Column(scale=5):
            prompt = gr.Textbox(label="Опис зображення")
            negative_prompt = gr.Textbox(label="Не показувати", value="")
        with gr.Column(scale=1):
            model_choice = gr.Radio(
                choices=list(MODEL_OPTIONS.keys()),
                label="Якість моделі",
                value=list(MODEL_OPTIONS.keys())[0],
            )
    with gr.Row():
        seed = gr.Slider(label="Seed", minimum=0, maximum=1000000, step=1, value=42)
        randomize_seed = gr.Checkbox(label="Випадковий Seed", value=True)

    with gr.Row():
        width = gr.Slider(label="Ширина", minimum=512, maximum=1024, step=64, value=DEFAULT_IMAGE_SIZE)
        height = gr.Slider(label="Висота", minimum=512, maximum=1024, step=64, value=DEFAULT_IMAGE_SIZE)
    with gr.Row():
        guidance_scale = gr.Slider(label="Guidance Scale", minimum=0, maximum=20, step=0.5, value=7.5)
        num_inference_steps = gr.Slider(label="Кроки інференсу", minimum=20, maximum=50, step=10, value=20)

    with gr.Row():
        num_images = gr.Slider(label="Кількість зображень", minimum=1, maximum=4, step=1, value=1)

    run_button = gr.Button("Створити")
    gallery = gr.Gallery(label="Створені зображення")
    status_text = gr.Textbox(label="Виконання") 

    run_button.click(
        fn=generate_image,
        inputs=[prompt, negative_prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps, num_images, model_choice],
        outputs=[gallery, status_text],
    )
        

demo.launch(share=True)