Using GANs and VAEs for creative tasks

import torch
import gc
import matplotlib.pyplot as plt
from diffusers import DiffusionPipeline, StableDiffusionImg2ImgPipeline
from transformers import BlipProcessor, BlipForConditionalGeneration

device = "cuda" if torch.cuda.is_available() else "cpu"
dtype = torch.float16 if device == "cuda" else torch.float32


def clear_memory():
    gc.collect()
    if torch.cuda.is_available():
        torch.cuda.empty_cache()
        torch.cuda.ipc_collect()


def text2image(prompts, model_id="segmind/tiny-sd", seed=42, steps=20, guidance=7.5):
    pipe = DiffusionPipeline.from_pretrained(model_id, torch_dtype=dtype)
    pipe = pipe.to(device if device == "cuda" else "cpu")
    if device == "cuda":
        pipe.enable_attention_slicing()

    images = []
    for i, prompt in enumerate(prompts):
        print(f"T2I prompt {i+1}: {prompt}")
        gen = torch.Generator(device=device).manual_seed(seed + i)
        img = pipe(
            prompt=prompt,
            num_inference_steps=steps,
            guidance_scale=guidance,
            generator=gen
        ).images[0]
        images.append(img)

    del pipe
    clear_memory()
    return images


def caption_images(images, model_id="Salesforce/blip-image-captioning-base", max_new_tokens=30):
    processor = BlipProcessor.from_pretrained(model_id)
    model = BlipForConditionalGeneration.from_pretrained(model_id, torch_dtype=dtype)
    model = model.to(device if device == "cuda" else "cpu")
    model.eval()

    captions = []
    for i, img in enumerate(images):
        inputs = processor(images=img, return_tensors="pt")
        inputs = {k: v.to(device) for k, v in inputs.items()}
        with torch.no_grad():
            out_ids = model.generate(**inputs, max_new_tokens=max_new_tokens)
        caption = processor.decode(out_ids[0], skip_special_tokens=True)
        print(f"Caption {i+1}: {caption}")
        captions.append(caption)

    del model
    del processor
    clear_memory()
    return captions


def image2image(images, prompts, model_id="runwayml/stable-diffusion-v1-5",
                seed=42, steps=25, guidance=7.5, strength=0.55):
    pipe = StableDiffusionImg2ImgPipeline.from_pretrained(model_id, torch_dtype=dtype)
    pipe = pipe.to(device if device == "cuda" else "cpu")
    if device == "cuda":
        pipe.enable_attention_slicing()

    out_images = []
    for i, (img, prompt) in enumerate(zip(images, prompts)):
        print(f"I2I prompt {i+1}: {prompt}")
        gen = torch.Generator(device=device).manual_seed(seed + i)
        styled = pipe(
            prompt=prompt,
            image=img,
            strength=strength,
            guidance_scale=guidance,
            num_inference_steps=steps,
            generator=gen
        ).images[0]
        out_images.append(styled)

    del pipe
    clear_memory()
    return out_images


def show_grid(images, titles, rows=3, cols=2):
    fig, axes = plt.subplots(rows, cols, figsize=(14, 16))
    axes = axes.flatten()

    for i, (img, title) in enumerate(zip(images, titles)):
        axes[i].imshow(img)
        axes[i].set_title(title, fontsize=10)
        axes[i].axis("off")

    for j in range(len(images), len(axes)):
        axes[j].axis("off")

    plt.tight_layout()
    plt.show()


# ===== PRZYKŁAD UŻYCIA =====
t2i_prompts = [
    "a futuristic chair made of wood and glass, studio lighting",
    "a retro red car on a street, detailed, realistic, daytime"
]

i2i_prompts = [
    "transformed into a neon cyberpunk poster, vivid colors, high detail",
    "transformed into a neon cyberpunk poster, vivid colors, high detail"
]

base_images = text2image(t2i_prompts)
captions = caption_images(base_images)
styled_images = image2image(base_images, i2i_prompts)

all_images = [
    base_images[0], base_images[1],
    styled_images[0], styled_images[1]
]

all_titles = [
    f"T2I 1\n{t2i_prompts[0]}",
    f"T2I 2\n{t2i_prompts[1]}",
    f"I2I 1\n{captions[0]}\n{i2i_prompts[0]}",
    f"I2I 2\n{captions[1]}\n{i2i_prompts[1]}"
]

show_grid(all_images, all_titles, rows=3, cols=2)