{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "37de3483",
   "metadata": {},
   "source": [
    "# PixelForge with Z-Image Turbo (Colab Test)\n",
    "Test stärkere Modelle mit kostenloser GPU"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f0d963b3",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 1. Install dependencies\n",
    "!pip install -q torch diffusers transformers accelerate numpy pillow requests -U\n",
    "!pip install -q flask flask-cors python-dotenv\n",
    "print(\"✓ Dependencies installed\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9f93fdb6",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 2. Check GPU\n",
    "import torch\n",
    "print(f\"GPU Available: {torch.cuda.is_available()}\")\n",
    "if torch.cuda.is_available():\n",
    "    print(f\"GPU: {torch.cuda.get_device_name(0)}\")\n",
    "    print(f\"VRAM: {torch.cuda.get_device_properties(0).total_memory / 1e9:.2f} GB\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e1d69f6e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 3. Download & Load Model (beispiel: FLUX oder ähnlich)\n",
    "# Optional: Ersetze mit deinem Z-Image Turbo Modell\n",
    "\n",
    "from diffusers import StableDiffusionPipeline\n",
    "import torch\n",
    "\n",
    "MODEL_ID = \"segmind/tiny-sd\"  # oder dein Z-Image Turbo\n",
    "print(f\"Loading {MODEL_ID}...\")\n",
    "\n",
    "pipe = StableDiffusionPipeline.from_pretrained(\n",
    "    MODEL_ID,\n",
    "    torch_dtype=torch.float16,\n",
    "    safety_checker=None,\n",
    "    requires_safety_checker=False\n",
    ")\n",
    "pipe = pipe.to(\"cuda\")\n",
    "print(\"✓ Model loaded\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a7575df6",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 4. Test Generation\n",
    "prompt = \"A beautiful futuristic city at sunset\"\n",
    "print(f\"Generating: {prompt}\")\n",
    "\n",
    "image = pipe(\n",
    "    prompt,\n",
    "    height=512,\n",
    "    width=512,\n",
    "    num_inference_steps=20,\n",
    "    guidance_scale=7.5\n",
    ").images[0]\n",
    "\n",
    "image.save(\"/tmp/test_output.png\")\n",
    "print(\"✓ Image generated: /tmp/test_output.png\")\n",
    "image"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4c84cc93",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 5. Start API Server (optional: expose to PixelForge)\n",
    "from flask import Flask, request, jsonify\n",
    "import json\n",
    "import base64\n",
    "from io import BytesIO\n",
    "\n",
    "app = Flask(__name__)\n",
    "\n",
    "@app.route('/health', methods=['GET'])\n",
    "def health():\n",
    "    return jsonify({\"status\": \"ok\", \"model\": MODEL_ID, \"gpu\": torch.cuda.get_device_name(0)})\n",
    "\n",
    "@app.route('/generate', methods=['POST'])\n",
    "def generate():\n",
    "    data = request.json\n",
    "    prompt = data.get('prompt', 'a test image')\n",
    "    steps = data.get('steps', 20)\n",
    "    guidance = data.get('guidance', 7.5)\n",
    "    \n",
    "    print(f\"Generating: {prompt}\")\n",
    "    image = pipe(\n",
    "        prompt,\n",
    "        height=512,\n",
    "        width=512,\n",
    "        num_inference_steps=steps,\n",
    "        guidance_scale=guidance\n",
    "    ).images[0]\n",
    "    \n",
    "    # Convert to base64\n",
    "    buffered = BytesIO()\n",
    "    image.save(buffered, format=\"PNG\")\n",
    "    img_base64 = base64.b64encode(buffered.getvalue()).decode()\n",
    "    \n",
    "    return jsonify({\n",
    "        \"success\": True,\n",
    "        \"image\": img_base64,\n",
    "        \"prompt\": prompt\n",
    "    })\n",
    "\n",
    "# Ngrok tunnel (kostenlos)\n",
    "from pyngrok import ngrok\n",
    "ngrok.set_auth_token(\"DEIN_NGROK_TOKEN\")  # Gratis auf ngrok.com registrieren\n",
    "\n",
    "print(\"Starting API server...\")\n",
    "public_url = ngrok.connect(5000)\n",
    "print(f\"Public URL: {public_url}\")\n",
    "print(\"Use this URL in PixelForge config!\")\n",
    "\n",
    "app.run(port=5000)"
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}