import torch

from PIL import Image
from io import BytesIO

from diffusers import StableDiffusionImg2ImgPipeline

import os

device = "cuda"
pipe = StableDiffusionImg2ImgPipeline.from_pretrained(
    "logo-marks-6-1_5", torch_dtype=torch.float16, safety_checker=None
).to(device)


def generate_image(user_id, seed, image_order, prompt=""):

    init_image = Image.open(
        f"logo-marks-6-1_5/reference-images/{image_order}.png"
    ).convert("RGB")
    init_image.thumbnail((768, 768))

    folder_path = f"images/{user_id}"
    os.makedirs(folder_path, exist_ok=True)

    generator = torch.Generator(device="cuda").manual_seed(seed)
    images = pipe(
        prompt=f"{prompt} logo mark in zvxsbj style",
        image=init_image,
        strength=0.95,
        guidance_scale=30,
        generator=generator,
    ).images
    images[0].save(f"images/{user_id}/{seed}-{image_order}.png")
    print(f"Image with seed {seed} and image_order {image_order} saved")