Initial commit

2026-09-27 15:20:26 -05:00
commit 72f40c7792
14 changed files with 466 additions and 0 deletions
@@ -0,0 +1,2 @@
+[Desktop Entry]
+Icon=folder-green
@@ -0,0 +1,25 @@
+# Uses the exact same verified environment as your Dev Container
+FROM nvcr.io/nvidia/pytorch:25.02-py3
+
+# Install API and Flux-specific dependencies
+# Note: We skip the torch versions in your list because the NVIDIA base 
+# already includes the optimized sm_120 version we need.
+RUN pip3 install --no-cache-dir \
+    fastapi \
+    uvicorn \
+    pydantic \
+    transformers==4.46.1 \
+    diffusers==0.35.2 \
+    accelerate==0.30.1 \
+    sentencepiece \
+    protobuf
+
+# Set up the application
+WORKDIR /app
+COPY api.py /app/api.py
+
+# Ensure GPU visibility (matching your devcontainer.json)
+ENV NVIDIA_VISIBLE_DEVICES=all
+ENV NVIDIA_DRIVER_CAPABILITIES=compute,utility
+
+CMD ["python3", "api.py"]
@@ -0,0 +1,57 @@
+import os
+import base64
+from io import BytesIO
+from fastapi import FastAPI, HTTPException
+from fastapi.middleware.cors import CORSMiddleware
+from pydantic import BaseModel
+import torch
+from diffusers import FluxPipeline
+
+app = FastAPI()
+
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+print("Initializing Flux-Schnell on RTX 5060 Ti (sm_120)...")
+hf_token = os.environ.get("HF_TOKEN")
+
+pipe = FluxPipeline.from_pretrained(
+    "black-forest-labs/FLUX.1-schnell",
+    torch_dtype=torch.bfloat16,
+    token=hf_token
+)
+
+# Offload execution stages sequentially to maintain 16GB headroom
+pipe.enable_sequential_cpu_offload()
+pipe.vae.enable_tiling()
+
+class ImageRequest(BaseModel):
+    prompt: str
+    model: str = "flux-schnell"
+
+@app.post("/v1/images/generations")
+async def generate(request: ImageRequest):
+    try:
+        result = pipe(
+            request.prompt,
+            guidance_scale=0.0,
+            num_inference_steps=4,
+            max_sequence_length=256,
+        ).images[0]
+
+        buffered = BytesIO()
+        result.save(buffered, format="PNG")
+        img_str = base64.b64encode(buffered.getvalue()).decode()
+
+        return {"data": [{"b64_json": img_str}]}
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
+
+if __name__ == "__main__":
+    import uvicorn
+    uvicorn.run(app, host="0.0.0.0", port=5005)
@@ -0,0 +1,59 @@
+from fastapi import FastAPI, HTTPException
+from fastapi.middleware.cors import CORSMiddleware
+from pydantic import BaseModel
+import torch
+from diffusers import FluxPipeline
+import base64
+from io import BytesIO
+
+app = FastAPI()
+
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+# Load FLUX.1-schnell (Best for 16GB VRAM + Fast Generation)
+print("Initializing Flux-Schnell on RTX 5060 Ti (sm_120)...")
+try:
+    pipe = FluxPipeline.from_pretrained(
+        "black-forest-labs/FLUX.1-schnell", 
+        torch_dtype=torch.bfloat16
+    )
+    # Optimized for 16GB: Moves parts of the model to CPU when not in use
+    # pipe.enable_model_cpu_offload()
+    pipe.enable_sequential_cpu_offload() 
+    pipe.vae.enable_tiling()
+except Exception as e:
+    print(f"Error loading model: {e}")
+
+class ImageRequest(BaseModel):
+    prompt: str
+    model: str = "flux-schnell" # Matches the error fix from earlier
+
+@app.post("/v1/images/generations")
+async def generate(request: ImageRequest):
+    try:
+        # Flux-Schnell generates high quality in just 4 steps
+        result = pipe(
+            request.prompt,
+            guidance_scale=0.0,
+            num_inference_steps=4,
+            max_sequence_length=256,
+        ).images[0]
+        
+        # Convert to Base64 for the API response
+        buffered = BytesIO()
+        result.save(buffered, format="PNG")
+        img_str = base64.b64encode(buffered.getvalue()).decode()
+        
+        return {"data": [{"b64_json": img_str}]}
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
+
+if __name__ == "__main__":
+    import uvicorn
+    uvicorn.run(app, host="0.0.0.0", port=5005)
@@ -0,0 +1,28 @@
+#!/usr/bin/env bash
+set -e
+
+PROMPT="$1"
+SERVER_IP="192.168.2.154" 
+API_URL="http://${SERVER_IP}:5005/v1/images/generations"
+
+# Note: Using "flux-schnell" to match our api.py logic
+RESPONSE=$(curl -s -X POST "$API_URL" \
+  -H "Content-Type: application/json" \
+  -d "{
+    \"prompt\": \"$PROMPT\",
+    \"model\": \"flux-schnell\"
+  }")
+
+BASE64=$(echo "$RESPONSE" | jq -r '.data[0].b64_json')
+
+if [ "$BASE64" == "null" ]; then
+  echo "Error: $(echo "$RESPONSE" | jq -r .detail)"
+  exit 1
+fi
+
+OUTDIR="downloads/flux_images"
+mkdir -p "$OUTDIR"
+OUTFILE="${OUTDIR}/$(date +%s).png"
+
+echo "$BASE64" | base64 -d > "$OUTFILE"
+echo "Image generated: $OUTFILE"
@@ -0,0 +1,295 @@
+<!DOCTYPE html>
+<html lang="en">
+<head>
+  <meta charset="UTF-8">
+  <meta name="viewport" content="width=device-width, initial-scale=1.0">
+  <title>Flux Image Generator</title>
+  <style>
+    * { margin: 0; padding: 0; box-sizing: border-box; }
+
+    body {
+      font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', Roboto, sans-serif;
+      background: #0f0f13;
+      color: #e0e0e0;
+      min-height: 100vh;
+      display: flex;
+      flex-direction: column;
+      align-items: center;
+      padding: 2rem 1rem;
+    }
+
+    h1 {
+      font-size: 1.8rem;
+      font-weight: 600;
+      margin-bottom: 0.3rem;
+      background: linear-gradient(135deg, #a78bfa, #60a5fa);
+      -webkit-background-clip: text;
+      -webkit-text-fill-color: transparent;
+    }
+
+    .subtitle {
+      color: #666;
+      font-size: 0.85rem;
+      margin-bottom: 2rem;
+    }
+
+    .input-area {
+      width: 100%;
+      max-width: 640px;
+      display: flex;
+      gap: 0.5rem;
+      margin-bottom: 2rem;
+    }
+
+    #prompt {
+      flex: 1;
+      padding: 0.85rem 1rem;
+      font-size: 1rem;
+      border: 1px solid #2a2a35;
+      border-radius: 10px;
+      background: #1a1a24;
+      color: #e0e0e0;
+      outline: none;
+      transition: border-color 0.2s;
+    }
+
+    #prompt:focus {
+      border-color: #a78bfa;
+    }
+
+    #prompt::placeholder {
+      color: #555;
+    }
+
+    #generate-btn {
+      padding: 0.85rem 1.5rem;
+      font-size: 1rem;
+      font-weight: 600;
+      border: none;
+      border-radius: 10px;
+      background: linear-gradient(135deg, #7c3aed, #3b82f6);
+      color: white;
+      cursor: pointer;
+      transition: opacity 0.2s, transform 0.1s;
+      white-space: nowrap;
+    }
+
+    #generate-btn:hover:not(:disabled) {
+      opacity: 0.9;
+    }
+
+    #generate-btn:active:not(:disabled) {
+      transform: scale(0.97);
+    }
+
+    #generate-btn:disabled {
+      opacity: 0.5;
+      cursor: not-allowed;
+    }
+
+    .result-area {
+      width: 100%;
+      max-width: 640px;
+      display: flex;
+      flex-direction: column;
+      align-items: center;
+      gap: 1rem;
+    }
+
+    #image-container {
+      display: none;
+      flex-direction: column;
+      align-items: center;
+      gap: 1rem;
+      width: 100%;
+    }
+
+    #image-container.visible {
+      display: flex;
+    }
+
+    #result-image {
+      max-width: 100%;
+      border-radius: 12px;
+      border: 1px solid #2a2a35;
+      background: #1a1a24;
+    }
+
+    #download-btn {
+      padding: 0.7rem 1.4rem;
+      font-size: 0.9rem;
+      font-weight: 600;
+      border: 1px solid #2a2a35;
+      border-radius: 8px;
+      background: #1a1a24;
+      color: #a78bfa;
+      cursor: pointer;
+      transition: background 0.2s;
+      text-decoration: none;
+    }
+
+    #download-btn:hover {
+      background: #22222e;
+    }
+
+    .status {
+      font-size: 0.9rem;
+      color: #888;
+      text-align: center;
+      min-height: 1.4em;
+    }
+
+    .status.error {
+      color: #f87171;
+    }
+
+    .spinner {
+      display: inline-block;
+      width: 1em;
+      height: 1em;
+      margin-right: 0.5em;
+      vertical-align: -0.15em;
+      border: 2px solid #a78bfa;
+      border-right-color: transparent;
+      border-radius: 50%;
+      animation: spin 0.6s linear infinite;
+    }
+
+    @keyframes spin {
+      to { transform: rotate(360deg); }
+    }
+
+    .config {
+      width: 100%;
+      max-width: 640px;
+      margin-bottom: 1.5rem;
+      padding: 0.8rem 1rem;
+      background: #1a1a24;
+      border: 1px solid #2a2a35;
+      border-radius: 10px;
+      display: flex;
+      align-items: center;
+      gap: 0.5rem;
+      font-size: 0.85rem;
+      color: #888;
+    }
+
+    .config label {
+      white-space: nowrap;
+    }
+
+    .config input {
+      flex: 1;
+      padding: 0.35rem 0.5rem;
+      font-size: 0.85rem;
+      border: 1px solid #2a2a35;
+      border-radius: 6px;
+      background: #0f0f13;
+      color: #e0e0e0;
+      outline: none;
+    }
+
+    .config input:focus {
+      border-color: #a78bfa;
+    }
+  </style>
+</head>
+<body>
+  <h1>Flux Image Generator</h1>
+  <p class="subtitle">Powered by flux-schnell</p>
+
+  <div class="config">
+    <label for="api-url">API URL:</label>
+    <input id="api-url" value="http://192.168.2.154:5005/v1/images/generations" />
+  </div>
+
+  <div class="input-area">
+    <input id="prompt" type="text" placeholder="Describe the image you want to create..." autofocus />
+    <button id="generate-btn">Generate</button>
+  </div>
+
+  <p class="status" id="status"></p>
+
+  <div class="result-area">
+    <div id="image-container">
+      <img id="result-image" alt="Generated image" />
+      <a id="download-btn" download="flux-image.png">⬇ Download PNG</a>
+    </div>
+  </div>
+
+  <script>
+    const promptInput = document.getElementById('prompt');
+    const generateBtn = document.getElementById('generate-btn');
+    const apiInput = document.getElementById('api-url');
+    const statusEl = document.getElementById('status');
+    const imageContainer = document.getElementById('image-container');
+    const resultImage = document.getElementById('result-image');
+    const downloadBtn = document.getElementById('download-btn');
+
+    function setStatus(msg, isError = false) {
+      statusEl.innerHTML = msg;
+      statusEl.className = 'status' + (isError ? ' error' : '');
+    }
+
+    async function generate() {
+      const prompt = promptInput.value.trim();
+      if (!prompt) {
+        setStatus('Please enter a prompt.', true);
+        return;
+      }
+
+      const apiUrl = apiInput.value.trim();
+      if (!apiUrl) {
+        setStatus('Please enter an API URL.', true);
+        return;
+      }
+
+      generateBtn.disabled = true;
+      generateBtn.textContent = 'Generating...';
+      setStatus('<span class="spinner"></span> Generating image...');
+      imageContainer.classList.remove('visible');
+
+      try {
+        const response = await fetch(apiUrl, {
+          method: 'POST',
+          headers: { 'Content-Type': 'application/json' },
+          body: JSON.stringify({ prompt, model: 'flux-schnell' })
+        });
+
+        if (!response.ok) {
+          throw new Error(`HTTP ${response.status}: ${response.statusText}`);
+        }
+
+        const data = await response.json();
+        const base64 = data.data?.[0]?.b64_json;
+
+        if (!base64 || base64 === 'null') {
+          throw new Error(data.detail || 'No image data returned.');
+        }
+
+        const blob = await fetch(`data:image/png;base64,${base64}`).then(r => r.blob());
+        const url = URL.createObjectURL(blob);
+
+        resultImage.src = url;
+        downloadBtn.href = url;
+
+        const ts = new Date().toISOString().replace(/[:.]/g, '-');
+        downloadBtn.download = `flux-image-${ts}.png`;
+
+        imageContainer.classList.add('visible');
+        setStatus('');
+      } catch (err) {
+        setStatus(`Error: ${err.message}`, true);
+      } finally {
+        generateBtn.disabled = false;
+        generateBtn.textContent = 'Generate';
+      }
+    }
+
+    generateBtn.addEventListener('click', generate);
+    promptInput.addEventListener('keydown', (e) => {
+      if (e.key === 'Enter') generate();
+    });
+  </script>
+</body>
+</html>