Upload /bck/hugging.py with huggingface_hub

bck/hugging.py

@@ -0,0 +1,523 @@
+#!/bin/env python3.11
+import os
+import sqlite3
+import replicate
+import argparse
+import requests
+from datetime import datetime
+from fastapi import FastAPI, WebSocket, WebSocketDisconnect, HTTPException, Request, Form, Query,Response
+from fastapi.templating import Jinja2Templates
+from fastapi.responses import FileResponse
+from fastapi.staticfiles import StaticFiles
+from pydantic import BaseModel
+from typing import Optional, List
+import uvicorn
+from asyncio import gather, Semaphore, create_task
+from mistralai import Mistral
+from contextlib import contextmanager
+from io import BytesIO
+import zipfile
+
+import sys
+print(f"Arguments: {sys.argv}")
+
+
+token = os.getenv("HF_TOKEN")
+api_key = os.getenv("MISTRAL_API_KEY")
+agent_id = os.getenv("MISTRAL_FLUX_AGENT")
+
+
+HEADER = "\033[38;2;255;255;153m"
+TITLE = "\033[38;2;255;255;153m"
+MENU = "\033[38;2;255;165;0m"
+SUCCESS = "\033[38;2;153;255;153m"
+ERROR = "\033[38;2;255;69;0m"
+MAIN = "\033[38;2;204;204;255m"
+SPEAKER1 = "\033[38;2;173;216;230m"
+SPEAKER2 = "\033[38;2;255;179;102m"
+RESET = "\033[0m"
+
+#os.system("clear")
+
+#print(f"{HEADER}--------------------\nMY FLUX CREATOR v1.0\n--------------------{RESET}\n")
+
+DOWNLOAD_DIR = "/mnt/d/ai/dialog/2/flux-pics"
+DATABASE_PATH = "flux_logs_neu.db"
+TIMEOUT_DURATION = 900  # Timeout-Dauer in Sekunden
+
+IMAGE_STORAGE_PATH = DOWNLOAD_DIR  # Pfad auf flux-pics setzen
+app = FastAPI()
+
+app.mount("/static", StaticFiles(directory="static"), name="static")
+app.mount("/flux-pics", StaticFiles(directory=DOWNLOAD_DIR), name="flux-pics")
+templates = Jinja2Templates(directory="templates")
+
+@contextmanager
+def get_db_connection(db_path=DATABASE_PATH):
+    conn = sqlite3.connect(db_path)
+    try:
+        yield conn
+    finally:
+        conn.close()
+
+def initialize_database(db_path=DATABASE_PATH):
+    with get_db_connection(db_path) as conn:
+        cursor = conn.cursor()
+        cursor.execute("""
+            CREATE TABLE IF NOT EXISTS generation_logs (
+                id INTEGER PRIMARY KEY AUTOINCREMENT,
+                timestamp TEXT,
+                prompt TEXT,
+                optimized_prompt TEXT,
+                hf_lora TEXT,
+                lora_scale REAL,
+                aspect_ratio TEXT,
+                guidance_scale REAL,
+                output_quality INTEGER,
+                prompt_strength REAL,
+                num_inference_steps INTEGER,
+                output_file TEXT,
+                album_id INTEGER,
+                category_id INTEGER
+            )
+        """)
+        cursor.execute("""
+            CREATE TABLE IF NOT EXISTS albums (
+                id INTEGER PRIMARY KEY AUTOINCREMENT,
+                name TEXT NOT NULL
+            )
+        """)
+        cursor.execute("""
+            CREATE TABLE IF NOT EXISTS categories (
+                id INTEGER PRIMARY KEY AUTOINCREMENT,
+                name TEXT NOT NULL
+            )
+        """)
+        cursor.execute("""
+            CREATE TABLE IF NOT EXISTS pictures (
+                id INTEGER PRIMARY KEY AUTOINCREMENT,
+                timestamp TEXT,
+                file_path TEXT,
+                file_name TEXT,
+                album_id INTEGER,
+                FOREIGN KEY (album_id) REFERENCES albums(id)
+            )
+        """)
+        cursor.execute("""
+            CREATE TABLE IF NOT EXISTS picture_categories (
+                picture_id INTEGER,
+                category_id INTEGER,
+                FOREIGN KEY (picture_id) REFERENCES pictures(id),
+                FOREIGN KEY (category_id) REFERENCES categories(id),
+                PRIMARY KEY (picture_id, category_id)
+            )
+        """)
+        conn.commit()
+
+def log_generation(args, optimized_prompt, image_file):
+    file_path, file_name = os.path.split(image_file)
+    try:
+        with get_db_connection() as conn:
+            cursor = conn.cursor()
+            cursor.execute("""
+                INSERT INTO generation_logs (
+                    timestamp, prompt, optimized_prompt, hf_lora, lora_scale, aspect_ratio, guidance_scale,
+                    output_quality, prompt_strength, num_inference_steps, output_file, album_id, category_id
+                ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
+            """, (
+                datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
+                args.prompt,
+                optimized_prompt,
+                args.hf_lora,
+                args.lora_scale,
+                args.aspect_ratio,
+                args.guidance_scale,
+                args.output_quality,
+                args.prompt_strength,
+                args.num_inference_steps,
+                image_file,
+                args.album_id,
+                args.category_id
+            ))
+            picture_id = cursor.lastrowid
+            cursor.execute("""
+                INSERT INTO pictures (
+                    timestamp, file_path, file_name, album_id
+                ) VALUES (?, ?, ?, ?)
+            """, (
+                datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
+                file_path,
+                file_name,
+                args.album_id
+            ))
+            picture_id = cursor.lastrowid
+
+            # Insert multiple categories
+            for category_id in args.category_ids:
+                cursor.execute("""
+                    INSERT INTO picture_categories (picture_id, category_id)
+                    VALUES (?, ?)
+                """, (picture_id, category_id))
+
+            conn.commit()
+    except sqlite3.Error as e:
+        print(f"Error logging generation: {e}")
+
+@app.on_event("startup")
+def startup_event():
+    initialize_database()
+
+@app.get("/")
+def read_root(request: Request):
+    with get_db_connection() as conn:
+        cursor = conn.cursor()
+        cursor.execute("SELECT id, name FROM albums")
+        albums = cursor.fetchall()
+        cursor.execute("SELECT id, name FROM categories")
+        categories = cursor.fetchall()
+    return templates.TemplateResponse("index.html", {"request": request, "albums": albums, "categories": categories})
+
+@app.get("/archive")
+def read_archive(
+    request: Request,
+    album: Optional[str] = Query(None),
+    category: Optional[List[str]] = Query(None),
+    search: Optional[str] = None,
+    items_per_page: int = Query(30),
+    page: int = Query(1)
+):
+    album_id = int(album) if album and album.isdigit() else None
+    category_ids = [int(cat) for cat in category] if category else []
+    offset = (page - 1) * items_per_page
+
+    with get_db_connection() as conn:
+        cursor = conn.cursor()
+        query = """
+            SELECT gl.timestamp, gl.prompt, gl.optimized_prompt, gl.output_file, a.name as album, c.name as category
+            FROM generation_logs gl
+            LEFT JOIN albums a ON gl.album_id = a.id
+            LEFT JOIN categories c ON gl.category_id = c.id
+            WHERE 1=1
+        """
+        params = []
+
+        if album_id is not None:
+            query += " AND gl.album_id = ?"
+            params.append(album_id)
+
+        if category_ids:
+            query += " AND gl.category_id IN ({})".format(','.join('?' for _ in category_ids))
+            params.extend(category_ids)
+
+        if search:
+            query += " AND (gl.prompt LIKE ? OR gl.optimized_prompt LIKE ?)"
+            params.append(f'%{search}%')
+            params.append(f'%{search}%')
+
+        query += " ORDER BY gl.timestamp DESC LIMIT ? OFFSET ?"
+        params.extend([items_per_page, offset])
+        cursor.execute(query, params)
+        logs = cursor.fetchall()
+
+        logs = [{
+            "timestamp": log[0],
+            "prompt": log[1],
+            "optimized_prompt": log[2],
+            "output_file": log[3],
+            "album": log[4],
+            "category": log[5]
+        } for log in logs]
+
+        cursor.execute("SELECT id, name FROM albums")
+        albums = cursor.fetchall()
+
+        cursor.execute("SELECT id, name FROM categories")
+        categories = cursor.fetchall()
+
+    return templates.TemplateResponse("archive.html", {
+        "request": request,
+        "logs": logs,
+        "albums": albums,
+        "categories": categories,
+        "selected_album": album,
+        "selected_categories": category_ids,
+        "search_query": search,
+        "items_per_page": items_per_page,
+        "page": page
+    })
+
+
+@app.get("/backend")
+def read_backend(request: Request):
+    with get_db_connection() as conn:
+        cursor = conn.cursor()
+        cursor.execute("SELECT id, name FROM albums")
+        albums = cursor.fetchall()
+        cursor.execute("SELECT id, name FROM categories")
+        categories = cursor.fetchall()
+    return templates.TemplateResponse("backend.html", {"request": request, "albums": albums, "categories": categories})
+
+@app.post("/create_album")
+def create_album(name: str = Form(...)):
+    try:
+        with get_db_connection() as conn:
+            cursor = conn.cursor()
+            cursor.execute("INSERT INTO albums (name) VALUES (?)", (name,))
+            conn.commit()
+        return {"message": "Album erstellt"}
+    except sqlite3.Error as e:
+        raise HTTPException(status_code=500, detail=f"Error creating album: {e}")
+
+@app.post("/create_category")
+def create_category(name: str = Form(...)):
+    try:
+        with get_db_connection() as conn:
+            cursor = conn.cursor()
+            cursor.execute("INSERT INTO categories (name) VALUES (?)", (name,))
+            conn.commit()
+        return {"message": "Kategorie erstellt"}
+    except sqlite3.Error as e:
+        raise HTTPException(status_code=500, detail=f"Error creating category: {e}")
+
+@app.post("/flux-pics")
+async def download_images(request: Request):
+    try:
+        body = await request.json()
+        print(f"Received request body: {body}")  # Debug log
+        
+        image_files = body.get("selectedImages", [])
+        if not image_files:
+            raise HTTPException(status_code=400, detail="Keine Bilder ausgewählt.")
+            
+        print(f"Processing image files: {image_files}")  # Debug log
+
+        # Überprüfe ob Download-Verzeichnis existiert
+        if not os.path.exists(IMAGE_STORAGE_PATH):
+            print(f"Storage path not found: {IMAGE_STORAGE_PATH}")  # Debug log
+            raise HTTPException(status_code=500, detail="Storage path not found")
+
+        zip_buffer = BytesIO()
+        with zipfile.ZipFile(zip_buffer, "w", zipfile.ZIP_DEFLATED) as zip_file:
+            for image_file in image_files:
+                image_path = os.path.join(IMAGE_STORAGE_PATH, image_file)
+                print(f"Processing file: {image_path}")  # Debug log
+                
+                if os.path.exists(image_path):
+                    zip_file.write(image_path, arcname=image_file)
+                else:
+                    print(f"File not found: {image_path}")  # Debug log
+                    raise HTTPException(status_code=404, detail=f"Bild {image_file} nicht gefunden.")
+
+        zip_buffer.seek(0)
+        
+        # Korrekter Response mit Buffer
+        return Response(
+            content=zip_buffer.getvalue(),
+            media_type="application/zip",
+            headers={
+                "Content-Disposition": f"attachment; filename=images.zip"
+            }
+        )
+        
+    except Exception as e:
+        print(f"Error in download_images: {str(e)}")  # Debug log
+        raise HTTPException(status_code=500, detail=str(e))
+
+@app.post("/flux-pics/single")
+async def download_single_image(request: Request):
+    try:
+        data = await request.json()
+        filename = data.get("filename")
+        print(f"Requested file download: {filename}")  # Debug log
+        
+        if not filename:
+            print("No filename provided")  # Debug log
+            raise HTTPException(status_code=400, detail="Kein Dateiname angegeben")
+            
+        file_path = os.path.join(IMAGE_STORAGE_PATH, filename)
+        print(f"Full file path: {file_path}")  # Debug log
+        
+        if not os.path.exists(file_path):
+            print(f"File not found: {file_path}")  # Debug log
+            raise HTTPException(status_code=404, detail=f"Datei {filename} nicht gefunden")
+
+        # Determine MIME type
+        file_extension = filename.lower().split('.')[-1]
+        mime_types = {
+            'png': 'image/png',
+            'jpg': 'image/jpeg',
+            'jpeg': 'image/jpeg',
+            'gif': 'image/gif',
+            'webp': 'image/webp'
+        }
+        media_type = mime_types.get(file_extension, 'application/octet-stream')
+        print(f"Serving file with media type: {media_type}")  # Debug log
+            
+        return FileResponse(
+            path=file_path,
+            filename=filename,
+            media_type=media_type,
+            headers={
+                "Content-Disposition": f"attachment; filename={filename}"
+            }
+        )
+    except Exception as e:
+        print(f"Error in download_single_image: {str(e)}")  # Debug log
+        raise HTTPException(status_code=500, detail=str(e))
+
+@app.websocket("/ws")
+async def websocket_endpoint(websocket: WebSocket):
+    await websocket.accept()
+    try:
+        data = await websocket.receive_json()
+        prompts = data.get("prompts", [data])
+
+        for prompt_data in prompts:
+            prompt_data['lora_scale'] = float(prompt_data['lora_scale'])
+            prompt_data['guidance_scale'] = float(prompt_data['guidance_scale'])
+            prompt_data['prompt_strength'] = float(prompt_data['prompt_strength'])
+            prompt_data['num_inference_steps'] = int(prompt_data['num_inference_steps'])
+            prompt_data['num_outputs'] = int(prompt_data['num_outputs'])
+            prompt_data['output_quality'] = int(prompt_data['output_quality'])
+
+            # Handle new album and category creation
+            album_name = prompt_data.get('album_id')
+            category_names = prompt_data.get('category_ids', [])
+
+            if album_name and not album_name.isdigit():
+                with get_db_connection() as conn:
+                    cursor = conn.cursor()
+                    cursor.execute("INSERT INTO albums (name) VALUES (?)", (album_name,))
+                    conn.commit()
+                    prompt_data['album_id'] = cursor.lastrowid
+            else:
+                prompt_data['album_id'] = int(album_name) if album_name else None
+
+            category_ids = []
+            for category_name in category_names:
+                if not category_name.isdigit():
+                    with get_db_connection() as conn:
+                        cursor = conn.cursor()
+                        cursor.execute("INSERT INTO categories (name) VALUES (?)", (category_name,))
+                        conn.commit()
+                        category_ids.append(cursor.lastrowid)
+                else:
+                    category_ids.append(int(category_name) if category_name else None)
+            prompt_data['category_ids'] = category_ids
+
+            args = argparse.Namespace(**prompt_data)
+
+            await websocket.send_json({"message": "Optimiere Prompt..."})
+            optimized_prompt = optimize_prompt(args.prompt) if getattr(args, 'agent', False) else args.prompt
+            await websocket.send_json({"optimized_prompt": optimized_prompt})
+
+            if prompt_data.get("optimize_only"):
+                continue
+
+            await generate_and_download_image(websocket, args, optimized_prompt)
+    except WebSocketDisconnect:
+        print("Client disconnected")
+    except Exception as e:
+        await websocket.send_json({"message": str(e)})
+        raise e
+    finally:
+        await websocket.close()
+
+async def fetch_image(item, index, args, filenames, semaphore, websocket, timestamp):
+    async with semaphore:
+        try:
+            response = requests.get(item, timeout=TIMEOUT_DURATION)
+            if response.status_code == 200:
+                filename = f"{DOWNLOAD_DIR}/image_{timestamp}_{index}.{args.output_format}"
+                with open(filename, "wb") as file:
+                    file.write(response.content)
+                filenames.append(f"/flux-pics/image_{timestamp}_{index}.{args.output_format}")
+                progress = int((index + 1) / args.num_outputs * 100)
+                await websocket.send_json({"progress": progress})
+            else:
+                await websocket.send_json({"message": f"Fehler beim Herunterladen des Bildes {index + 1}: {response.status_code}"})
+        except requests.exceptions.Timeout:
+            await websocket.send_json({"message": f"Timeout beim Herunterladen des Bildes {index + 1}"})
+
+async def generate_and_download_image(websocket: WebSocket, args, optimized_prompt):
+    try:
+        input_data = {
+            "prompt": optimized_prompt,
+            "hf_lora": getattr(args, 'hf_lora', None),  # Use getattr to safely access hf_lora
+            "lora_scale": args.lora_scale,
+            "num_outputs": args.num_outputs,
+            "aspect_ratio": args.aspect_ratio,
+            "output_format": args.output_format,
+            "guidance_scale": args.guidance_scale,
+            "output_quality": args.output_quality,
+            "prompt_strength": args.prompt_strength,
+            "num_inference_steps": args.num_inference_steps,
+            "disable_safety_checker": False
+        }
+
+        await websocket.send_json({"message": "Generiere Bilder..."})
+        
+        # Debug: Log the start of the replication process
+        print(f"Starting replication process for {args.num_outputs} outputs with timeout {TIMEOUT_DURATION}")
+
+        output = replicate.run(
+            "lucataco/flux-dev-lora:091495765fa5ef2725a175a57b276ec30dc9d39c22d30410f2ede68a3eab66b3",
+            input=input_data,
+            timeout=TIMEOUT_DURATION
+        )
+        
+        if not os.path.exists(DOWNLOAD_DIR):
+            os.makedirs(DOWNLOAD_DIR)
+
+        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+        filenames = []
+        semaphore = Semaphore(3)  # Limit concurrent downloads
+
+        tasks = [create_task(fetch_image(item, index, args, filenames, semaphore, websocket, timestamp)) for index, item in enumerate(output)]
+        await gather(*tasks)
+
+        for file in filenames:
+            log_generation(args, optimized_prompt, file)
+        
+        await websocket.send_json({"message": "Bilder erfolgreich generiert", "generated_files": filenames})
+    except requests.exceptions.Timeout:
+        await websocket.send_json({"message": "Fehler bei der Bildgenerierung: Timeout überschritten"})
+    except Exception as e:
+        await websocket.send_json({"message": f"Fehler bei der Bildgenerierung: {str(e)}"})
+        raise Exception(f"Fehler bei der Bildgenerierung: {str(e)}")
+
+def optimize_prompt(prompt):
+    api_key = os.environ.get("MISTRAL_API_KEY")
+    agent_id = os.environ.get("MISTRAL_FLUX_AGENT")
+
+    if not api_key or not agent_id:
+        raise ValueError("MISTRAL_API_KEY oder MISTRAL_FLUX_AGENT nicht gesetzt")
+
+    client = Mistral(api_key=api_key)
+    chat_response = client.agents.complete(
+        agent_id=agent_id,
+        messages=[{"role": "user", "content": f"Optimiere folgenden Prompt für Flux Lora: {prompt}"}]
+    )
+    
+    return chat_response.choices[0].message.content
+
+if __name__ == "__main__":
+    # Parse command line arguments
+    parser = argparse.ArgumentParser(description="Beschreibung")
+    parser.add_argument('--hf_lora', default=None, help='HF LoRA Model')
+    args = parser.parse_args()
+
+    # Pass arguments to the FastAPI application
+    app.state.args = args
+
+    # Run the Uvicorn server
+#    uvicorn.run(app, host="0.0.0.0", port=8000, timeout_keep_alive=900)
+
+    # Run server
+    uvicorn.run(
+        "main:app",
+        host="0.0.0.0",
+        port=8000,
+        reload=True,
+        log_level="debug"
+    )