import gradio as gr
import numpy as np
import random
import os
from PIL import Image

import spaces  # [uncomment to use ZeroGPU]
from diffusers import StableDiffusionPipeline, UNet2DConditionModel
import torch

# 模态映射
legal_modals = ['OCT', 'CXR', 'Fundus', 'BrainMRI', 'BreastMRI', 'ChestCT']
modal_idx = {
    'OCT': 1,
    'CXR': 2, 
    'Fundus': 3,
    'BrainMRI': 4,
    'BreastMRI': 4,
    'ChestCT': 4
}

device = "cuda" if torch.cuda.is_available() else "cpu"
model_repo_id = "CoheeY/MINIM"

if torch.cuda.is_available():
    torch_dtype = torch.float16
else:
    torch_dtype = torch.float32

# 缓存不同模态的pipeline
pipe_cache = {}

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


def get_pipe_for_modal(modal_name):
    """获取或创建指定模态的pipeline"""
    if modal_name not in pipe_cache:
        print(f"Loading pipeline for modal: {modal_name}")
        
        # 获取模态对应的数字ID
        modal_id = modal_idx[modal_name]
        
        # 加载该模态的微调UNet
        unet_subfolder = f"unets/{modal_id}/unet"
        print(f"Loading UNet from subfolder: {unet_subfolder}")
        
        unet = UNet2DConditionModel.from_pretrained(
            model_repo_id,
            subfolder=unet_subfolder,
            torch_dtype=torch_dtype
        )
        
        # 创建pipeline，使用微调UNet替换基础UNet
        pipe = StableDiffusionPipeline.from_pretrained(
            model_repo_id,
            unet=unet,
            safety_checker=None,
            torch_dtype=torch_dtype
        )
        
        pipe = pipe.to(device)
        
        # 启用优化
        if torch.cuda.is_available():
            try:
                pipe.enable_xformers_memory_efficient_attention()
                print("XFormers enabled")
            except:
                print("XFormers not available")
        
        pipe_cache[modal_name] = pipe
        print(f"Pipeline for {modal_name} loaded successfully")
    
    return pipe_cache[modal_name]


@spaces.GPU
def infer(
    modal,
    prompt,
    negative_prompt,
    seed,
    randomize_seed,
    width,
    height,
    guidance_scale,
    num_inference_steps,
    progress=gr.Progress(track_tqdm=True),
):
    # 设置随机种子
    if randomize_seed:
        seed = random.randint(0, MAX_SEED)
    
    generator = torch.Generator(device=device).manual_seed(seed)
    
    # 获取对应模态的pipeline
    pipe = get_pipe_for_modal(modal)
    
    # 构建输入prompt
    modal_index = modal_idx[modal]
    full_prompt = f"{modal_index}:{prompt}"
    
    print(f"Generating image for modal: {modal}, prompt: {full_prompt}")
    
    # 生成图像
    try:
        image = pipe(
            prompt=full_prompt,
            negative_prompt=negative_prompt,
            guidance_scale=guidance_scale,
            num_inference_steps=num_inference_steps,
            width=width,
            height=height,
            generator=generator,
        ).images[0]
        
        print(f"Image generated successfully")
        return image, seed
        
    except Exception as e:
        print(f"Error during inference: {e}")
        import traceback
        traceback.print_exc()
        
        # 返回错误信息
        error_image = Image.new('RGB', (512, 512), color=(255, 200, 200))
        return error_image, seed


# 医学图像示例
medical_examples = [
    ["CXR", "clear lung fields with normal heart size"],
    ["Fundus", "normal retina with healthy optic disc"],
    ["BrainMRI", "normal brain anatomy in axial view"],
]

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

with gr.Blocks(css=css) as demo:
    with gr.Column(elem_id="col-container"):
        gr.Markdown("# MINIM: Medical Image Generation Model")
        gr.Markdown("### Multi-modal Medical Image Generation with Stable Diffusion")
        
        with gr.Row():
            modal = gr.Dropdown(
                label="Imaging Modality",
                choices=legal_modals,
                value="CXR",
                info="Select medical imaging modality"
            )
            
        with gr.Row():
            prompt = gr.Textbox(
                label="Description",
                show_label=False,
                max_lines=2,
                placeholder="Describe the medical image you want to generate...",
                container=False,
                scale=4
            )
            
            run_button = gr.Button("Generate", scale=1, variant="primary")
        
        result = gr.Image(label="Generated Image", show_label=False)
        
        with gr.Accordion("Advanced Settings", open=False):
            negative_prompt = gr.Textbox(
                label="Negative Prompt",
                max_lines=2,
                placeholder="Describe what you don't want in the image...",
                value="blurry, distorted, artifact",
            )
            
            with gr.Row():
                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=64,
                    value=512,
                )
                
                height = gr.Slider(
                    label="Height",
                    minimum=256,
                    maximum=MAX_IMAGE_SIZE,
                    step=64,
                    value=512,
                )
            
            with gr.Row():
                guidance_scale = gr.Slider(
                    label="Guidance Scale",
                    minimum=0.0,
                    maximum=20.0,
                    step=0.5,
                    value=7.5,
                )
                
                num_inference_steps = gr.Slider(
                    label="Inference Steps",
                    minimum=10,
                    maximum=200,
                    step=5,
                    value=100,
                )
        
        gr.Examples(
            examples=medical_examples,
            inputs=[modal, prompt],
            label="Medical Image Examples"
        )
    
    # 绑定事件
    run_button.click(
        fn=infer,
        inputs=[
            modal,
            prompt,
            negative_prompt,
            seed,
            randomize_seed,
            width,
            height,
            guidance_scale,
            num_inference_steps,
        ],
        outputs=[result, seed]
    )
    
    prompt.submit(
        fn=infer,
        inputs=[
            modal,
            prompt,
            negative_prompt,
            seed,
            randomize_seed,
            width,
            height,
            guidance_scale,
            num_inference_steps,
        ],
        outputs=[result, seed]
    )

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