import torch
import torchvision.transforms.functional as F
from PIL import Image
from src.pix2pix_turbo import Pix2Pix_Turbo

# Initialize the model
model = Pix2Pix_Turbo("sketch_to_image_stochastic")

# Load the sketch image
sketch_path = "sketch.png"
image = Image.open(sketch_path).convert("RGB")

# Define the prompt
prompt = "ethereal fantasy concept art of   . magnificent, celestial, ethereal, painterly, epic, majestic, magical, fantasy art, cover art, dreamy,"

# Convert the image to a tensor
#image_tensor = F.to_tensor(image) > 0.5
image = image.convert("RGB")
image_t = F.to_tensor(image) > 0.5
# Processing parameters
seed = 42
val_r = 0.4

# Run the model inference
with torch.no_grad():
    c_t = image_t.unsqueeze(0).cuda().float()
    torch.manual_seed(seed)
    B, C, H, W = c_t.shape
    noise = torch.randn((1, 4, H // 8, W // 8), device=c_t.device)
    output_image = model(c_t, prompt, deterministic=False, r=val_r, noise_map=noise)

# Convert the output tensor to an image
output_pil = F.to_pil_image(output_image[0].cpu() * 0.5 + 0.5)

# Save the output image
output_path = "output.png"
output_pil.save(output_path)

print(f"Processing complete. The output image is saved as {output_path}.")