{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Pipeline d'entraînement — Générateur de cartes Pokémon\n",
    "\n",
    "Dataset : images PNG avec métadonnées **pokemon_metadata** dans les chunks PNG (tEXt).\n",
    "Modèle : entrée = JSON (même format que pokemon_metadata). En amont, ta pipeline convertit le prompt utilisateur → JSON ; ici on conditionne directement sur ce JSON (sérialisé en texte)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Using device: cuda\n",
      "Found 21783 file(s), validating...\n",
      "Valid images with pokemon_metadata: 21783\n",
      "Train: 19604, Val: 2179\n"
     ]
    }
   ],
   "source": [
    "from pathlib import Path\n",
    "import json\n",
    "from PIL import Image, ImageFile\n",
    "import torch\n",
    "from torch.utils.data import Dataset, DataLoader, random_split\n",
    "from torchvision import transforms\n",
    "\n",
    "ImageFile.LOAD_TRUNCATED_IMAGES = True\n",
    "\n",
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"mps\" if torch.backends.mps.is_available() else \"cpu\")\n",
    "print(f\"Using device: {device}\")\n",
    "\n",
    "# Répertoire des images : dossiers dans cards/ (sous-dossiers autorisés)\n",
    "image_dir = Path(\"/home/do5-ajlp/juicepyter/cards\")\n",
    "image_extensions = {\".png\", \".jpg\", \".jpeg\"}\n",
    "\n",
    "def get_metadata_from_png(path):\n",
    "    \"\"\"Lit le chunk pokemon_metadata depuis les métadonnées PNG (tEXt).\"\"\"\n",
    "    with Image.open(path) as img:\n",
    "        meta_str = img.info.get(\"pokemon_metadata\")\n",
    "    if meta_str is None:\n",
    "        return None\n",
    "    try:\n",
    "        return json.loads(meta_str)\n",
    "    except json.JSONDecodeError:\n",
    "        return None\n",
    "\n",
    "def metadata_to_conditioning(meta):\n",
    "    \"\"\"Sérialise les metadata (dict, format pokemon_metadata) en chaîne JSON pour le conditionnement.\"\"\"\n",
    "    return json.dumps(meta, sort_keys=True, ensure_ascii=False)\n",
    "\n",
    "all_paths = sorted(\n",
    "    p for p in image_dir.rglob(\"*\")\n",
    "    if p.is_file() and p.suffix.lower() in image_extensions\n",
    ")\n",
    "print(f\"Found {len(all_paths)} file(s), validating...\")\n",
    "\n",
    "image_paths = []\n",
    "for p in all_paths:\n",
    "    try:\n",
    "        meta = get_metadata_from_png(p)\n",
    "        if meta is None:\n",
    "            continue\n",
    "        with Image.open(p) as img:\n",
    "            img.verify()\n",
    "        image_paths.append(p)\n",
    "    except Exception:\n",
    "        continue\n",
    "\n",
    "print(f\"Valid images with pokemon_metadata: {len(image_paths)}\")\n",
    "\n",
    "# Taille cible pour Stable Diffusion\n",
    "IMG_SIZE = 512\n",
    "\n",
    "transform = transforms.Compose([\n",
    "    transforms.Resize((IMG_SIZE, IMG_SIZE)),\n",
    "    transforms.CenterCrop(IMG_SIZE),\n",
    "    transforms.ToTensor(),\n",
    "    transforms.Normalize([0.5], [0.5]),\n",
    "])\n",
    "\n",
    "class PokemonCardDataset(Dataset):\n",
    "    def __init__(self, image_paths, transform, tokenizer, max_length=77):\n",
    "        self.image_paths = image_paths\n",
    "        self.transform = transform\n",
    "        self.tokenizer = tokenizer\n",
    "        self.max_length = max_length\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.image_paths)\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        path = self.image_paths[idx]\n",
    "        meta = get_metadata_from_png(path)\n",
    "        conditioning = metadata_to_conditioning(meta)\n",
    "        img = Image.open(path).convert(\"RGB\")\n",
    "        pixel_values = self.transform(img)\n",
    "        tokens = self.tokenizer(\n",
    "            conditioning,\n",
    "            truncation=True,\n",
    "            padding=\"max_length\",\n",
    "            max_length=self.max_length,\n",
    "            return_tensors=\"pt\",\n",
    "        )\n",
    "        return {\n",
    "            \"pixel_values\": pixel_values,\n",
    "            \"input_ids\": tokens.input_ids.squeeze(0),\n",
    "            \"attention_mask\": tokens.attention_mask.squeeze(0),\n",
    "        }\n",
    "\n",
    "\n",
    "# Tokenizer CLIP (compatible SD 1.5)\n",
    "from transformers import CLIPTokenizer\n",
    "tokenizer = CLIPTokenizer.from_pretrained(\"openai/clip-vit-large-patch14\")\n",
    "\n",
    "dataset = PokemonCardDataset(image_paths, transform, tokenizer)\n",
    "train_size = int(0.9 * len(dataset))\n",
    "val_size = len(dataset) - train_size\n",
    "train_dataset, val_dataset = random_split(dataset, [train_size, val_size])\n",
    "\n",
    "BATCH_SIZE = 4\n",
    "train_loader = DataLoader(train_dataset, batch_size=BATCH_SIZE, shuffle=True, num_workers=0)\n",
    "val_loader = DataLoader(val_dataset, batch_size=BATCH_SIZE, shuffle=False, num_workers=0)\n",
    "\n",
    "print(f\"Train: {len(train_dataset)}, Val: {len(val_dataset)}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/do5-ajlp/lib/python3.12/site-packages/huggingface_hub/utils/_validators.py:206: UserWarning: The `local_dir_use_symlinks` argument is deprecated and ignored in `hf_hub_download`. Downloading to a local directory does not use symlinks anymore.\n",
      "  warnings.warn(\n",
      "Warning: You are sending unauthenticated requests to the HF Hub. Please set a HF_TOKEN to enable higher rate limits and faster downloads.\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "a348b58691b9495dac1a220589791e4c",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading weights:   0%|          | 0/196 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\u001b[1mCLIPTextModel LOAD REPORT\u001b[0m from: runwayml/stable-diffusion-v1-5\n",
      "Key                                | Status     |  | \n",
      "-----------------------------------+------------+--+-\n",
      "text_model.embeddings.position_ids | UNEXPECTED |  | \n",
      "\n",
      "\u001b[3mNotes:\n",
      "- UNEXPECTED\u001b[3m\t:can be ignored when loading from different task/architecture; not ok if you expect identical arch.\u001b[0m\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "UNet parameters: 861,115,332, trainable (LoRA): 1,594,368\n"
     ]
    }
   ],
   "source": [
    "from diffusers import AutoencoderKL, UNet2DConditionModel, DDPMScheduler\n",
    "from transformers import CLIPTextModel\n",
    "from peft import LoraConfig, get_peft_model\n",
    "\n",
    "# Modèle de base: Stable Diffusion 1.5\n",
    "model_id = \"runwayml/stable-diffusion-v1-5\"\n",
    "\n",
    "vae = AutoencoderKL.from_pretrained(model_id, subfolder=\"vae\")\n",
    "unet = UNet2DConditionModel.from_pretrained(model_id, subfolder=\"unet\")\n",
    "text_encoder = CLIPTextModel.from_pretrained(model_id, subfolder=\"text_encoder\")\n",
    "noise_scheduler = DDPMScheduler.from_pretrained(model_id, subfolder=\"scheduler\")\n",
    "\n",
    "# Injecter LoRA dans l'UNet\n",
    "lora_config = LoraConfig(\n",
    "    r=8,\n",
    "    lora_alpha=32,\n",
    "    init_lora_weights=\"gaussian\",\n",
    "    target_modules=[\"to_k\", \"to_q\", \"to_v\", \"to_out.0\"],\n",
    ")\n",
    "unet = get_peft_model(unet, lora_config)\n",
    "\n",
    "vae = vae.to(device)\n",
    "unet = unet.to(device)\n",
    "text_encoder = text_encoder.to(device)\n",
    "\n",
    "# Gel des poids du VAE et du text encoder\n",
    "vae.requires_grad_(False)\n",
    "text_encoder.requires_grad_(False)\n",
    "\n",
    "total_params = sum(p.numel() for p in unet.parameters())\n",
    "trainable_params = sum(p.numel() for p in unet.parameters() if p.requires_grad)\n",
    "print(f\"UNet parameters: {total_params:,}, trainable (LoRA): {trainable_params:,}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Resumed from checkpoint at epoch 1 (train_loss=0.1383, val_loss=0.1398)\n",
      "Training epochs 2 to 10...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[2/10] Train:   3%|▉                             | 152/4901 [01:29<45:46,  1.73it/s, loss=0.1157, ema=0.1271, avg=0.1304]"
     ]
    }
   ],
   "source": [
    "from torch.optim import AdamW\n",
    "from tqdm import tqdm\n",
    "import time\n",
    "\n",
    "EPOCHS = 10\n",
    "LR = 1e-5\n",
    "CHECKPOINT_DIR = \"/home/do5-ajlp/juicepyter/pokemon_card_lora\"\n",
    "CHECKPOINT_PATH = f\"{CHECKPOINT_DIR}/checkpoint.pt\"\n",
    "\n",
    "optimizer = AdamW(unet.parameters(), lr=LR)\n",
    "train_losses = []\n",
    "val_losses = []\n",
    "start_epoch = 0\n",
    "\n",
    "# Resume from checkpoint if it exists\n",
    "if Path(CHECKPOINT_PATH).exists():\n",
    "    checkpoint = torch.load(CHECKPOINT_PATH, map_location=device, weights_only=False)\n",
    "    unet.load_state_dict(checkpoint[\"unet_state_dict\"])\n",
    "    optimizer.load_state_dict(checkpoint[\"optimizer_state_dict\"])\n",
    "    start_epoch = checkpoint[\"epoch\"]\n",
    "    train_losses = checkpoint[\"train_losses\"]\n",
    "    val_losses = checkpoint[\"val_losses\"]\n",
    "    print(f\"Resumed from checkpoint at epoch {start_epoch} (train_loss={train_losses[-1]:.4f}, val_loss={val_losses[-1]:.4f})\")\n",
    "else:\n",
    "    print(\"No checkpoint found, training from scratch.\")\n",
    "\n",
    "if start_epoch >= EPOCHS:\n",
    "    print(f\"Already trained {start_epoch} epochs (requested {EPOCHS}). Increase EPOCHS to continue.\")\n",
    "else:\n",
    "    print(f\"Training epochs {start_epoch + 1} to {EPOCHS}...\")\n",
    "\n",
    "def encode_prompt(text_encoder, input_ids, attention_mask):\n",
    "    with torch.no_grad():\n",
    "        emb = text_encoder(input_ids, attention_mask=attention_mask)[0]\n",
    "    return emb\n",
    "\n",
    "for epoch in range(start_epoch, EPOCHS):\n",
    "    start = time.time()\n",
    "\n",
    "    # Train\n",
    "    unet.train()\n",
    "    running_loss = 0.0\n",
    "    smooth_loss = None  # exponential moving average\n",
    "\n",
    "    pbar = tqdm(train_loader, desc=f\"[{epoch+1}/{EPOCHS}] Train\", leave=True,\n",
    "                bar_format=\"{l_bar}{bar:30}{r_bar}\", dynamic_ncols=True)\n",
    "\n",
    "    for step, batch in enumerate(pbar):\n",
    "        pixel_values = batch[\"pixel_values\"].to(device)\n",
    "        input_ids = batch[\"input_ids\"].to(device)\n",
    "        attention_mask = batch[\"attention_mask\"].to(device)\n",
    "\n",
    "        with torch.no_grad():\n",
    "            latents = vae.encode(pixel_values).latent_dist.sample() * 0.18215\n",
    "\n",
    "        encoder_hidden_states = encode_prompt(text_encoder, input_ids, attention_mask)\n",
    "        noise = torch.randn_like(latents)\n",
    "        bsz = latents.shape[0]\n",
    "        timesteps = torch.randint(0, noise_scheduler.config.num_train_timesteps, (bsz,), device=device).long()\n",
    "        noisy_latents = noise_scheduler.add_noise(latents, noise, timesteps)\n",
    "\n",
    "        pred = unet(noisy_latents, timesteps, encoder_hidden_states).sample\n",
    "        loss = torch.nn.functional.mse_loss(pred.float(), noise.float(), reduction=\"mean\")\n",
    "\n",
    "        optimizer.zero_grad()\n",
    "        loss.backward()\n",
    "        torch.nn.utils.clip_grad_norm_(unet.parameters(), 1.0)\n",
    "        optimizer.step()\n",
    "\n",
    "        running_loss += loss.item() * bsz\n",
    "\n",
    "        # Exponential moving average for smoother loss display\n",
    "        smooth_loss = loss.item() if smooth_loss is None else 0.98 * smooth_loss + 0.02 * loss.item()\n",
    "\n",
    "        pbar.set_postfix(ordered_dict={\n",
    "            \"loss\": f\"{loss.item():.4f}\",\n",
    "            \"ema\": f\"{smooth_loss:.4f}\",\n",
    "            \"avg\": f\"{running_loss / ((step + 1) * bsz):.4f}\",\n",
    "        })\n",
    "\n",
    "    train_loss = running_loss / len(train_dataset)\n",
    "    train_losses.append(train_loss)\n",
    "\n",
    "    # Validate\n",
    "    unet.eval()\n",
    "    running_loss = 0.0\n",
    "\n",
    "    with torch.no_grad():\n",
    "        vbar = tqdm(val_loader, desc=f\"[{epoch+1}/{EPOCHS}] Val  \", leave=True,\n",
    "                    bar_format=\"{l_bar}{bar:30}{r_bar}\", dynamic_ncols=True)\n",
    "\n",
    "        for step, batch in enumerate(vbar):\n",
    "            pixel_values = batch[\"pixel_values\"].to(device)\n",
    "            input_ids = batch[\"input_ids\"].to(device)\n",
    "            attention_mask = batch[\"attention_mask\"].to(device)\n",
    "\n",
    "            latents = vae.encode(pixel_values).latent_dist.sample() * 0.18215\n",
    "            encoder_hidden_states = encode_prompt(text_encoder, input_ids, attention_mask)\n",
    "            noise = torch.randn_like(latents)\n",
    "            bsz = latents.shape[0]\n",
    "            timesteps = torch.randint(0, noise_scheduler.config.num_train_timesteps, (bsz,), device=device).long()\n",
    "            noisy_latents = noise_scheduler.add_noise(latents, noise, timesteps)\n",
    "\n",
    "            pred = unet(noisy_latents, timesteps, encoder_hidden_states).sample\n",
    "            loss = torch.nn.functional.mse_loss(pred.float(), noise.float(), reduction=\"mean\")\n",
    "            running_loss += loss.item() * bsz\n",
    "\n",
    "            vbar.set_postfix(ordered_dict={\n",
    "                \"loss\": f\"{loss.item():.4f}\",\n",
    "                \"avg\": f\"{running_loss / ((step + 1) * bsz):.4f}\",\n",
    "            })\n",
    "\n",
    "    val_loss = running_loss / len(val_dataset)\n",
    "    val_losses.append(val_loss)\n",
    "\n",
    "    elapsed = time.time() - start\n",
    "    mins, secs = divmod(int(elapsed), 60)\n",
    "    print(f\"\\n{'='*60}\")\n",
    "    print(f\"  Epoch {epoch+1}/{EPOCHS} complete — {mins}m{secs:02d}s\")\n",
    "    print(f\"  Train loss : {train_loss:.4f}\")\n",
    "    print(f\"  Val   loss : {val_loss:.4f}\")\n",
    "    print(f\"{'='*60}\\n\")\n",
    "\n",
    "    # Save checkpoint after each epoch\n",
    "    Path(CHECKPOINT_DIR).mkdir(exist_ok=True)\n",
    "    torch.save({\n",
    "        \"epoch\": epoch + 1,\n",
    "        \"unet_state_dict\": unet.state_dict(),\n",
    "        \"optimizer_state_dict\": optimizer.state_dict(),\n",
    "        \"train_losses\": train_losses,\n",
    "        \"val_losses\": val_losses,\n",
    "    }, CHECKPOINT_PATH)\n",
    "    print(f\"  Checkpoint saved (epoch {epoch + 1})\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 28732 characters]",
      "text/plain": [
       "<Figure size 800x400 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "\n",
    "plt.figure(figsize=(8, 4))\n",
    "plt.plot(train_losses, label=\"Train\")\n",
    "plt.plot(val_losses, label=\"Val\")\n",
    "plt.xlabel(\"Epoch\")\n",
    "plt.ylabel(\"MSE Loss\")\n",
    "plt.title(\"Training & Validation Loss\")\n",
    "plt.legend()\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "74dd610894ef48b8853583908a4fa0ce",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading pipeline components...:   0%|          | 0/6 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "You have disabled the safety checker for <class 'diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline'> by passing `safety_checker=None`. Ensure that you abide to the conditions of the Stable Diffusion license and do not expose unfiltered results in services or applications open to the public. Both the diffusers team and Hugging Face strongly recommend to keep the safety filter enabled in all public facing circumstances, disabling it only for use-cases that involve analyzing network behavior or auditing its results. For more information, please have a look at https://github.com/huggingface/diffusers/pull/254 .\n",
      "Token indices sequence length is longer than the specified maximum sequence length for this model (351 > 77). Running this sequence through the model will result in indexing errors\n",
      "The following part of your input was truncated because CLIP can only handle sequences up to 77 tokens: ['your deck .\", \" energytype \": null , \" evolvefrom \": null , \" hp \": null , \" id \": \" sv 0 8 . 5 - 1 0 5 \", \" illustrator \": \" gidora \", \" image \": \" https :// assets . tcgdex . net / en / sv / sv 0 8 . 5 / 1 0 5 \", \" item \": null , \" legal \": {\" expanded \": true , \" standard \": true }, \" level \": null , \" localid \": \" 1 0 5 \", \" name \": \" crispin \", \" rarity \": \" uncommon \", \" regulationmark \": \" h \", \" resistances \": null , \" retreat \": null , \" set \": {\" cardcount \": {\" official \": 1 3 1 , \" total \": 1 8 0 }, \" id \": \" sv 0 8 . 5 \", \" logo \": \" https :// assets . tcgdex . net / en / sv / sv 0 8 . 5 / logo \", \" name \": \" prismatic evolutions \", \" symbol \": null }, \" stage \": null , \" suffix \": null , \" trainertype \": \" supporter \", \" types \": null , \" variants \": {\" firstedition \": false , \" holo \": true , \" normal \": true , \" reverse \": true , \" wpromo \": false }, \" weaknesses \": null }']\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "dbfa4d8442f44140b9e0b9cd17e2b64a",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/30 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "The following part of your input was truncated because CLIP can only handle sequences up to 77 tokens: ['\" effect \": \" for each of your { f } pokémon in play , look at that many cards from the top of your deck and put them back in any order .\", \" energytype \": null , \" evolvefrom \": null , \" hp \": null , \" id \": \" a 4 - 2 0 1 \", \" illustrator \": \" hideki ishikawa \", \" image \": \" https :// assets . tcgdex . net / en / tcgp / a 4 / 2 0 1 \", \" item \": null , \" legal \": {\" expanded \": false , \" standard \": false }, \" level \": null , \" localid \": \" 2 0 1 \", \" name \": \" hiker \", \" rarity \": \" two star \", \" regulationmark \": null , \" resistances \": null , \" retreat \": null , \" set \": {\" cardcount \": {\" official \": 1 6 1 , \" total \": 2 4 1 }, \" id \": \" a 4 \", \" logo \": \" https :// assets . tcgdex . net / en / tcgp / a 4 / logo \", \" name \": \" wisdom of sea and sky \", \" symbol \": null }, \" stage \": null , \" suffix \": null , \" trainertype \": \" supporter \", \" types \": null , \" variants \": {\" firstedition \": false , \" holo \": true , \" normal \": true , \" reverse \": true , \" wpromo \": false }, \" weaknesses \": null }']\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "450ec6affe43451fb03a9a1adff3c86f",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/30 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "The following part of your input was truncated because CLIP can only handle sequences up to 77 tokens: ['\" fighting \", \" fighting \", \" colorless \"], \" damage \": 1 2 0 , \" effect \": null , \" name \": \" skull bash \"}], \" boosters \": null , \" category \": \" pokemon \", \" description \": \" its unique style of coiling allows it to blast sand out of its sand sac more efficiently .\", \" dexid \": [ 8 4 4 ], \" effect \": null , \" energytype \": null , \" evolvefrom \": \" silicobra \", \" hp \": 1 3 0 , \" id \": \" swsh 5 - 8 2 \", \" illustrator \": \" hatachu \", \" image \": \" https :// assets . tcgdex . net / en / swsh / swsh 5 / 8 2 \", \" item \": null , \" legal \": {\" expanded \": true , \" standard \": false }, \" level \": null , \" localid \": \" 8 2 \", \" name \": \" sandaconda \", \" rarity \": \" holo rare \", \" regulationmark \": \" e \", \" resistances \": null , \" retreat \": 3 , \" set \": {\" cardcount \": {\" official \": 1 6 3 , \" total \": 1 8 3 }, \" id \": \" swsh 5 \", \" logo \": \" https :// assets . tcgdex . net / en / swsh / swsh 5 / logo \", \" name \": \" battle styles \", \" symbol \": \" https :// assets . tcgdex . net / univ / swsh / swsh 5 / symbol \"}, \" stage \": \" stage 1 \", \" suffix \": null , \" trainertype \": null , \" types \": [\" fighting \"], \" variants \": {\" firstedition \": false , \" holo \": true , \" normal \": false , \" reverse \": true , \" wpromo \": false }, \" weaknesses \": [{\" type \": \" grass \", \" value \": \"× 2 \"}]}']\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "7cdf72400c3c4acb8dd8773b25592616",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/30 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "The following part of your input was truncated because CLIP can only handle sequences up to 77 tokens: ['\", \" colorless \"], \" damage \": 4 0 , \" effect \": \" the defending pokémon is now asleep .\", \" name \": \" hypno splash \"}, {\" cost \": [\" colorless \", \" colorless \", \" water \", \" water \"], \" damage \": \" 4 0 +\", \" effect \": \" does 4 0 damage plus 1 0 more damage for each of your benched stage 1 evolved pokémon .\", \" name \": \" rely on friends \"}], \" boosters \": null , \" category \": \" pokemon \", \" description \": null , \" dexid \": [ 3 2 1 ], \" effect \": null , \" energytype \": null , \" evolvefrom \": \" wailmer \", \" hp \": 1 2 0 , \" id \": \" ex 1 2 - 1 4 \", \" illustrator \": \" kouki saitou \", \" image \": \" https :// assets . tcgdex . net / en / ex / ex 1 2 / 1 4 \", \" item \": null , \" legal \": {\" expanded \": false , \" standard \": false }, \" level \": null , \" localid \": \" 1 4 \", \" name \": \" wailord \", \" rarity \": \" rare \", \" regulationmark \": null , \" resistances \": null , \" retreat \": null , \" set \": {\" cardcount \": {\" official \": 9 2 , \" total \": 9 3 }, \" id \": \" ex 1 2 \", \" logo \": \" https :// assets . tcgdex . net / en / ex / ex 1 2 / logo \", \" name \": \" legend maker \", \" symbol \": \" https :// assets . tcgdex . net / univ / ex / ex 1 2 / symbol \"}, \" stage \": \" stage 1 \", \" suffix \": null , \" trainertype \": null , \" types \": [\" water \"], \" variants \": {\" firstedition \": false , \" holo \": true , \" normal \": false , \" reverse \": false , \" wpromo \": false }, \" weaknesses \": [{\" type \": \" lightning \", \" value \": \"× 2 \"}]}']\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "f7e897cac2924bea83431a2fe8e635b7",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/30 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 2576232 characters]",
      "text/plain": [
       "<Figure size 1600x800 with 8 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Visualiser des générations sur des prompts de validation\n",
    "from diffusers import StableDiffusionPipeline\n",
    "import random\n",
    "import copy\n",
    "\n",
    "unet.eval()\n",
    "# Merge LoRA weights into a COPY so the original unet stays trainable\n",
    "unet_copy = copy.deepcopy(unet)\n",
    "unet_merged = unet_copy.merge_and_unload()\n",
    "\n",
    "pipe = StableDiffusionPipeline.from_pretrained(\n",
    "    model_id,\n",
    "    vae=vae,\n",
    "    unet=unet_merged,\n",
    "    text_encoder=text_encoder,\n",
    "    tokenizer=tokenizer,\n",
    "    safety_checker=None,\n",
    ")\n",
    "pipe = pipe.to(device)\n",
    "del unet_copy  # free memory\n",
    "\n",
    "indices = random.sample(range(len(val_dataset)), 4)\n",
    "\n",
    "def get_val_path(i):\n",
    "    dataset_idx = int(val_dataset.indices[i])  # tensor → int\n",
    "    return val_dataset.dataset.image_paths[dataset_idx]\n",
    "\n",
    "conditionings = [metadata_to_conditioning(get_metadata_from_png(get_val_path(indices[i]))) for i in range(4)]\n",
    "\n",
    "fig, axes = plt.subplots(2, 4, figsize=(16, 8))\n",
    "for i, cond in enumerate(conditionings):\n",
    "    out = pipe(cond, num_inference_steps=30, guidance_scale=7.5).images[0]\n",
    "    axes[0, i].imshow(out)\n",
    "    axes[0, i].axis(\"off\")\n",
    "    axes[0, i].set_title(cond[:50] + \"...\", fontsize=8)\n",
    "\n",
    "    real = Image.open(get_val_path(indices[i])).convert(\"RGB\").resize((512, 512))\n",
    "    axes[1, i].imshow(real)\n",
    "    axes[1, i].axis(\"off\")\n",
    "    axes[1, i].set_title(\"Ground truth\", fontsize=8)\n",
    "\n",
    "plt.suptitle(\"Generated vs Ground Truth\")\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def generate_card_from_metadata(meta, num_inference_steps=30, guidance_scale=7.5, save_path=None):\n",
    "    \"\"\"Génère une carte à partir d'un dict au format pokemon_metadata (sortie de ta pipeline prompt → JSON).\"\"\"\n",
    "    conditioning = metadata_to_conditioning(meta)\n",
    "    image = pipe(conditioning, num_inference_steps=num_inference_steps, guidance_scale=guidance_scale).images[0]\n",
    "    if save_path:\n",
    "        image.save(save_path)\n",
    "    return image\n",
    "\n",
    "# Exemple : meta = sortie de ta pipeline (prompt utilisateur → JSON)\n",
    "example_meta = {\n",
    "    \"name\": \"Charizard\", \"types\": [\"Fire\"], \"hp\": 120, \"stage\": \"Stage2\", \"rarity\": \"Rare\",\n",
    "    \"attacks\": [{\"name\": \"Fire Spin\", \"damage\": \"100\"}, {\"name\": \"Flamethrower\", \"damage\": \"50\"}],\n",
    "}\n",
    "out = generate_card_from_metadata(example_meta)\n",
    "out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "The following part of your input was truncated because CLIP can only handle sequences up to 77 tokens: ['], \" category \": \" pokemon \", \" description \": \" it makes a nest to suit its long and skinny body . the nest is impossible for other pokémon to enter .\", \" evolvefrom \": \" sentret \", \" hp \": 1 1 0 , \" id \": \" swsh 3 - 1 3 6 \", \" illustrator \": \" tetsuya koizumi \", \" image \": \" https :// assets . tcgdex . net / en / swsh / swsh 3 / 1 3 6 \", \" localid \": \" 1 3 6 \", \" name \": \" furret \", \" rarity \": \" uncommon \", \" regulationmark \": \" d \", \" retreat \": 1 , \" set \": {\" cardcount \": {\" official \": 1 8 9 , \" total \": 2 0 1 }, \" id \": \" swsh 3 \", \" logo \": \" https :// assets . tcgdex . net / en / swsh / swsh 3 / logo \", \" name \": \" darkness ablaze \", \" symbol \": \" https :// assets . tcgdex . net / univ / swsh / swsh 3 / symbol \"}, \" stage \": \" stage 1 \", \" types \": [\" colorless \"], \" weaknesses \": [{\" type \": \" fighting \", \" value \": \"× 2 \"}]}']\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "b90f4851a87f4feb83ee481149dd3347",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/75 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/jpeg": "[encoded data, 80768 characters]",
      "image/png": "[encoded data, 678500 characters]",
      "text/plain": [
       "<PIL.Image.Image image mode=RGB size=512x512>"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "def generate_card_from_metadata(meta, num_inference_steps=30, guidance_scale=7.5, save_path=None):\n",
    "    \"\"\"Génère une carte à partir d'un dict au format pokemon_metadata (sortie de ta pipeline prompt → JSON).\"\"\"\n",
    "    conditioning = metadata_to_conditioning(meta)\n",
    "    image = pipe(conditioning, num_inference_steps=num_inference_steps, guidance_scale=guidance_scale).images[0]\n",
    "    if save_path:\n",
    "        image.save(save_path)\n",
    "    return image\n",
    "\n",
    "# Exemple : meta = sortie de ta pipeline (prompt utilisateur → JSON)\n",
    "example_meta = {\n",
    "  \"category\": \"Pokemon\",\n",
    "  \"id\": \"swsh3-136\",\n",
    "  \"illustrator\": \"tetsuya koizumi\",\n",
    "  \"image\": \"https://assets.tcgdex.net/en/swsh/swsh3/136\",\n",
    "  \"localId\": \"136\",\n",
    "  \"name\": \"Furret\",\n",
    "  \"rarity\": \"Uncommon\",\n",
    "  \"set\": {\n",
    "    \"cardCount\": {\n",
    "      \"official\": 189,\n",
    "      \"total\": 201\n",
    "    },\n",
    "    \"id\": \"swsh3\",\n",
    "    \"logo\": \"https://assets.tcgdex.net/en/swsh/swsh3/logo\",\n",
    "    \"name\": \"Darkness Ablaze\",\n",
    "    \"symbol\": \"https://assets.tcgdex.net/univ/swsh/swsh3/symbol\"\n",
    "  },\n",
    "  \"hp\": 110,\n",
    "  \"types\": [\n",
    "    \"Colorless\"\n",
    "  ],\n",
    "  \"evolveFrom\": \"Sentret\",\n",
    "  \"description\": \"It makes a nest to suit its long and skinny body. The nest is impossible for other Pokémon to enter.\",\n",
    "  \"stage\": \"Stage1\",\n",
    "  \"attacks\": [\n",
    "    {\n",
    "    \"cost\": [\n",
    "      \"Colorless\"\n",
    "    ],\n",
    "    \"name\": \"Feelin' Fine\",\n",
    "    \"effect\": \"Draw 3 cards.\"\n",
    "    },\n",
    "    {\n",
    "    \"cost\": [\n",
    "      \"Colorless\"\n",
    "    ],\n",
    "    \"name\": \"Tail Smash\",\n",
    "    \"effect\": \"Flip a coin. If tails, this attack does nothing.\",\n",
    "    \"damage\": 90\n",
    "    }\n",
    "  ],\n",
    "  \"weaknesses\": [\n",
    "    {\n",
    "      \"type\": \"Fighting\",\n",
    "      \"value\": \"×2\"\n",
    "    }\n",
    "  ],\n",
    "  \"retreat\": 1,\n",
    "  \"regulationMark\": \"D\"\n",
    "}\n",
    "out = generate_card_from_metadata(example_meta)\n",
    "out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "LoRA saved to /home/do5-ajlp/juicepyter/pokemon_card_lora (1 epochs total)\n"
     ]
    }
   ],
   "source": [
    "# Sauvegarder le LoRA et l'historique d'entraînement\n",
    "LORA_PATH = \"/home/do5-ajlp/juicepyter/pokemon_card_lora\"\n",
    "unet.save_pretrained(LORA_PATH)\n",
    "torch.save({\n",
    "    \"train_losses\": train_losses,\n",
    "    \"val_losses\": val_losses,\n",
    "    \"epochs\": len(train_losses),\n",
    "}, LORA_PATH + \"/training_history.pt\")\n",
    "print(f\"LoRA saved to {LORA_PATH} ({len(train_losses)} epochs total)\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
