import spaces
import os
from peft import PeftModel
import gradio as gr
from diffusers import StableDiffusionImg2ImgPipeline
from diffusers import AutoPipelineForImage2Image
from diffusers import DiffusionPipeline
import torch
from PIL import Image
from diffusers import StableDiffusionPipeline
# Load the model
# model_id = "nitrosocke/Ghibli-Diffusion"
# pipe = StableDiffusionImg2ImgPipeline.from_pretrained(model_id, torch_dtype=torch.float32)
tk=os.getenv('ghtoken')
print("ttttt",tk)
model_id = "black-forest-labs/FLUX.1-dev"
# pipe =DiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.bfloat16,token=tk)
pipe =AutoPipelineForImage2Image.from_pretrained(model_id, torch_dtype=torch.bfloat16,token=tk)


# # 1. 选择一个基础模型，例如 SD 1.5
# base_model_id = "runwayml/stable-diffusion-v1-5"

# # 2. 加载基础模型
# pipe = StableDiffusionPipeline.from_pretrained(
#     base_model_id,
#     torch_dtype=torch.float32
# )

# # 3. 加载 LoRA 权重
# lora_model_id = "openfree/flux-chatgpt-ghibli-lora"
# pipe.load_lora_weights(lora_model_id)

# pipe = AutoPipelineForImage2Image.from_pretrained('black-forest-labs/FLUX.1-dev', torch_dtype=torch.bfloat16,token=True)
# pipe.load_lora_weights('openfree/flux-chatgpt-ghibli-lora', weight_name='flux-chatgpt-ghibli-lora.safetensors')
pipe.load_lora_weights("alvarobartt/ghibli-characters-flux-lora")


# Move pipeline to GPU if available
device = "cuda" if torch.cuda.is_available() else "cpu"
pipe = pipe.to(device)

# Define the inference function
@spaces.GPU
def ghibli_transform(input_image, prompt="GHBLI anime style photo", guidance_scale=3.5, num_steps=30):
    print('canshu_guidance_scale',guidance_scale)
    print('canshu_num_steps',num_steps)
    if input_image is None:
        raise gr.Error("No image uploaded! Please upload an image before clicking Transform.")
    
    # Process the input image (keep it as PIL Image)
    try:
        init_image = input_image.convert("RGB").resize((1024, 768))
    except Exception as e:
        raise gr.Error(f"Failed to process image: {str(e)}")
    
    # Generate the Ghibli-style image
    try:
        output = pipe(
            prompt=prompt,
            image=init_image,
            # strength=strength,
            # guidance_scale=guidance_scale,
            # num_inference_steps=num_steps  # Use the UI-provided value
            ######
            guidance_scale=guidance_scale,
            num_inference_steps=num_steps
            
            ######
        ).images[0]
    except Exception as e:
        raise gr.Error(f"Pipeline error: {str(e)}")
    
    return output

# Create the Gradio interface
with gr.Blocks(title="Transformer") as demo:
    gr.Markdown("# Transformer")
    gr.Markdown("Upload an image and transform it! [Website:](http://imagetoghibli.online/)")
    
    with gr.Row():
        with gr.Column():
            input_img = gr.Image(label="Upload Image", type="pil")
            prompt = gr.Textbox(label="Prompt", value="GHBLI anime style photo")
           
            guidance = gr.Slider(1, 20, value=3.5, step=0.5, label="Guidance Scale")
            num_steps = gr.Slider(10, 100, value=30, step=5, label="Inference Steps (Higher = Better Quality, Slower)")
            submit_btn = gr.Button("Transform")
        with gr.Column():
            output_img = gr.Image(label="Ghibli-Style Output")
    
    # Connect the button to the function
    submit_btn.click(
        fn=ghibli_transform,
        inputs=[input_img, prompt, guidance, num_steps],
        outputs=output_img
    )

# Launch the Space with share=True for public link
demo.launch(share=True)