File size: 1,650 Bytes
bbc89f6
 
 
 
fdd33ad
bbc89f6
a1e077b
 
bbc89f6
a1e077b
 
 
 
 
 
 
 
 
 
 
 
 
fdd33ad
a1e077b
 
 
 
 
 
 
 
fdd33ad
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
a1e077b
bbc89f6
fdd33ad
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
import replicate
from PIL import Image
import requests
from io import BytesIO
from database import get_lora_models_info


def generate_image(model_name, prompt, steps, cfg_scale, width, height, lora_scale, progress, use_default=False, trigger_word='hi'):
    print(f"Generating image for prompt: {prompt}")
    if use_default:
        img_url = replicate.run(
            "black-forest-labs/flux-pro",
            input={
                "steps": steps,
                "prompt": prompt,
                "guidance": cfg_scale,
                "interval": 2,
                "aspect_ratio": "1:1",
                "safety_tolerance": 2
            }
        )
    else:
        input = {
                "model": "dev",
                "steps": steps,
                "prompt": prompt,
                "guidance": cfg_scale,
                "interval": 2,
                "aspect_ratio": "1:1",
                "safety_tolerance": 2
            }

        db_loras = get_lora_models_info()

        for lora in db_loras:
            if lora["lora_name"] == model_name:
                if lora["hf_repo"]:
                    input["hf_lora"] = lora["hf_repo"]
                    model_name = "lucataco/flux-dev-lora:a22c463f11808638ad5e2ebd582e07a469031f48dd567366fb4c6fdab91d614d"
                    
        if ':' not in model_name:
            model_version = replicate.models.get(model_name).latest_version.id
            print(f"Model version: {model_version}")
            model_name = f"{model_name}:{model_version}"
            
        img_url = replicate.run(
            model_name,
            input=input
        )
    return img_url