File size: 1,798 Bytes
650e153
7c960db
650e153
37cc261
7c960db
 
37cc261
9a4fd44
7c960db
0d8cfe0
7c960db
 
 
 
 
 
083c348
7c960db
 
 
 
 
 
4ae11a0
7c960db
 
 
 
4339127
7c960db
 
 
6e6bdf8
7c960db
 
 
650e153
7c960db
 
d4bdd2d
7c960db
 
 
650e153
7c960db
 
6f37e55
7c960db
 
 
6f37e55
7c960db
 
8a1c757
7c960db
 
 
6f37e55
 
7c960db
 
4ae11a0
d4bdd2d
7c960db
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
# ============================
# FIREWORKS AI MODULE (HYBRID MODE)
# ============================
import os
import base64
import requests
import uuid

FIREWORKS_API_KEY = os.getenv("FIREWORKS_API_KEY")

def fireworks_generate(prompt, model_name, log_callback):
    """
    Generate image using Fireworks AI.
    Returns filepath or None.
    """
    url = "https://api.fireworks.ai/inference/v1/image/generate"

    payload = {
        "model": model_name,
        "prompt": prompt,
        "size": "1024x1024",
        "num_images": 1
    }

    headers = {
        "Authorization": f"Bearer {FIREWORKS_API_KEY}",
        "Content-Type": "application/json"
    }

    try:
        response = requests.post(url, json=payload, headers=headers, timeout=40)
        data = response.json()

        if "images" not in data:
            log_callback(f"  → Fireworks Error: {data}")
            return None

        b64 = data["images"][0]["image_base64"]
        img_bytes = base64.b64decode(b64)

        filename = f"fw_{uuid.uuid4().hex}.png"
        os.makedirs("outputs", exist_ok=True)
        filepath = os.path.join("outputs", filename)

        with open(filepath, "wb") as f:
            f.write(img_bytes)

        if os.path.getsize(filepath) < 10000:
            log_callback("  → Fireworks FAILED (image too small)")
            return None

        log_callback("  → Fireworks SUCCESS")
        return filepath

    except Exception as e:
        log_callback(f"  → Fireworks Exception: {str(e)}")
        return None


def fireworks_schnell(prompt, log_callback):
    return fireworks_generate(prompt, "black-forest-labs/FLUX.1-schnell", log_callback)


def fireworks_dev(prompt, log_callback):
    return fireworks_generate(prompt, "black-forest-labs/FLUX.1-dev", log_callback)