import torch
from diffusers import DiffusionPipeline
from pipeline_sdxs import SdxsPipeline

device = "cuda"
pipe_id = "/workspace/sdxs-1b"
pipeline = SdxsPipeline.from_pretrained(
    pipe_id,
).to(device=device, dtype=torch.float16) 
print(pipeline)
prompts = [
"A digital illustration of a young male character with white hair, blue eyes, and a dark complexion. He is wearing a white jacket, brown pants, and a brown belt. His hands are in his pockets, and he is wearing a wristwatch. He has a pouch on his belt and is wearing glasses. The background is a simple"
,"A young woman with striking blue eyes and pointed ears, adorned with a floral kimono and a tattoo. Her hair is styled in a braid, and she wears a pair of ears"
,"In the center of a futuristic space station, a highly advanced N7 soldier stands at attention. The soldier, clad in a suit of power armor, is equipped with a helmet that features a visor and a rifle slung over the shoulder. The armor is a striking combination of gray and red"
,"A muscular, topless male with tiger-like ears and a tail stands in a forest, holding a sword and wearing a blue outfit, gazing directly at the viewer"
,"A young woman with striking features, a mix of black and gold hair floating around her head, vibrant yellow eyes half-closed, and a blue sweater contrasting the black background, framed in a simple yet striking composition."
,'A fluffy domestic cat with piercing green eyes sits attentively in a sunlit room filled with natural light, its soft fur reflecting warm hues of orange through golden windows.'
,"A fierce woman in black-and-white, wearing a spiked iron mask with sharp metallic spikes, sporting a high ponytail, her skin marked with battle-dirt, and eyes that reflect determination and strength. The mask's menacing aura contrasts with her intense expression, capturing her mysterious and intimidating nature."
,"A close-up of an astronaut's helmet with frosted, opaque visor, reflecting space's cold, frozen texture, resting on the visor a butterfly with vibrant, intricately patterned wings, and distant stars' faint glow."
,"A striking character with red eyes and black uniform wields a sword in a defensive stance, poised for battle amidst a stark white background with vibrant red accents."
,"A bald woman with a tattooed upper body stands against a vibrant red background. She is adorned with a pair of dangling earrings and a dress with a dragon design"
,"white cyborg knight riding cyber horse with wings, long white gown, holding scythe, skeleton horse, zombies cyberpunk armor, feathers"
,"A radiant voluptuous woman in Arizona's Grand Canyon at twilight, her fiery ginger hair cascading over a halter top with intricate lace, plunging neckline, black thongs, silver rings, and stacked bangles, surrounded by deep purples and burnt oranges in the sky, her light-blue eyes glowing with wonder as she gazes at the luminous beauty of the night."
,"A man with hair, a white suit and black scarf, anime-style elements, glowing pink horns, and cyberpunk manga art style, standing in a shadowy smoke-filled background, with glowing pink eyes and a side view portrait of his face, white skin, and ear piercings."
,"There is a young male character standing against a vibrant, colorful graffiti wall. he is wearing a straw hat, a black jacket adorned with gold accents, and black shorts."
,"A black BMW M3 sports car with black and yellow rims"
,"A young girl in a flowing, vibrant dress, her glowing eyes capturing the warmth of the day, sits on a grassy field, surrounded by anime-style elements."
,"A young woman with a messy bun and vibrant red hoodie stands ready to strike with a large axe, her intense gaze and dynamic pose evoking energy and anticipation."
,"A small bluebird with a white breast and black stripe perches on a branch, facing right with a blurred background of trees and bushes."
,"A stylized, anime-style illustration of a young male character with white hair, wearing a white suit and holding two guns. He is standing with his arms outstretched and a confident smile on his face."
,"In the center of a white background, a fantasy creature with blue eyes and a white body stands on its hind legs. Its front legs are adorned with armor, and its tail is adorned with scales. The creature has a large, spiky crest on its head and sharp claws"
,"In a serene garden, two young girls stand side by side, their youthful energy palpable. The girl on the left, adorned with a blue dress and a matching blue flower in her hair, gazes directly at the viewer, her eyes sparkling with curiosity."
]
prompts2 = ["Космонавт на мотоцикле"]

negative_prompt="low quality, bad quality, photo, low details"
negative_prompt2 = "worst quality, low quality, photo, low details, blurry, jpeg artifacts, unfinished, sketch, sepia, missing limb, text, bad anatomy, bad proportions, bad hands, missing fingers"
negative_prompt = "bad quality grainy image with low details, incomplete text, despite numerous technical flaws and distorted figures"

all_images = []
# Обработка батчей с прогресс-баром
for i, prompt in enumerate(prompts):
    image = pipeline(
        prompt = prompt,
        negative_prompt = negative_prompt,
        guidance_scale = 5,
        width = 1024,
        height = 1408,
        seed = 0,
        batch_size = 1,
    )[0]
    all_images.extend(image)
import matplotlib.pyplot as plt
import math
import numpy as np
import textwrap

def display_image_grid(images, prompts, cols=3, save_path=None):
    n = len(images)
    if n == 0:
        return
        
    rows = math.ceil(n / cols)

    # constrained_layout вместо tight_layout лучше держит одинаковую высоту ячеек
    fig, axes = plt.subplots(rows, cols, figsize=(cols * 4, rows * 4.5), constrained_layout=True)

    if isinstance(axes, np.ndarray):
        axes = axes.flatten()
    elif isinstance(axes, plt.Axes):
        axes = [axes]
        
    axes = list(axes)[:n]

    for i, (img, prompt) in enumerate(zip(images, prompts)):
        ax = axes[i]
        ax.imshow(img)
        ax.axis("off")
        ax.set_aspect("equal")

        text = (prompt[:200] + "…") if len(prompt) > 200 else prompt
        
        # Разбиваем на строки
        lines = textwrap.wrap(text, width=35)
        
        while len(lines) < 4:
            lines.append("")
            
        ax.set_title("\n".join(lines), fontsize=9, pad=8)
        # -----------------------------------

    if save_path:
        # Для сохранения constrained_layout работает автоматически
        plt.savefig(save_path, bbox_inches="tight", dpi=600, format="jpeg")

    #plt.show()


display_image_grid(all_images, prompts, save_path="media/result_grid.jpg")