File size: 3,697 Bytes
8d691c3
 
 
 
b79e3f7
8d691c3
 
a15d459
166d9fd
 
e221206
a15d459
 
8d691c3
21e5fc0
8d691c3
b79e3f7
 
 
8d691c3
 
 
 
 
 
 
 
b79e3f7
8d691c3
 
 
 
166d9fd
 
 
21e5fc0
8d691c3
b79e3f7
8d691c3
 
 
 
821924e
8d691c3
821924e
8d691c3
 
 
b79e3f7
 
36abcb2
b79e3f7
 
821924e
8d691c3
 
821924e
8d691c3
21e5fc0
 
 
 
 
166d9fd
 
e221206
 
 
166d9fd
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
21e5fc0
8d691c3
 
 
 
 
 
 
 
8749689
21e5fc0
8d691c3
36abcb2
 
 
166d9fd
21e5fc0
8749689
36abcb2
 
8d691c3
 
36abcb2
 
 
 
 
 
8d691c3
 
36abcb2
8d691c3
36abcb2
 
 
 
 
 
890a980
21e5fc0
821924e
21e5fc0
 
 
 
 
 
 
36abcb2
3a091ee
36abcb2
8d691c3
 
 
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
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
import gradio as gr
import numpy as np
import random

import spaces #[uncomment to use ZeroGPU]
from diffusers import DiffusionPipeline
import torch
import subprocess
from groq import Groq
import base64
import os

subprocess.run("rm -rf /data-nvme/zerogpu-offload/*", env={}, shell=True)

# Load FLUX image generator
device = "cuda" if torch.cuda.is_available() else "cpu"
model_repo_id = "black-forest-labs/FLUX.1-schnell"  # Replace to the model you would like to use
lora_path = "matteomarjanovic/flatsketcher"
weigths_file = "lora.safetensors"

if torch.cuda.is_available():
    torch_dtype = torch.float16
else:
    torch_dtype = torch.float32

pipe = DiffusionPipeline.from_pretrained(model_repo_id, torch_dtype=torch_dtype)
pipe = pipe.to(device)
pipe.load_lora_weights(lora_path, weight_name=weigths_file)

MAX_SEED = np.iinfo(np.int32).max
MAX_IMAGE_SIZE = 1024

def encode_image(image_path):
  with open(image_path, "rb") as image_file:
    return base64.b64encode(image_file.read()).decode('utf-8')


@spaces.GPU #[uncomment to use ZeroGPU]
def infer(
    prompt,
    progress=gr.Progress(track_tqdm=True),
):
    # seed = random.randint(0, MAX_SEED)

    # generator = torch.Generator().manual_seed(seed)

    image = pipe(
        prompt=prompt,
        guidance_scale=0.,
        num_inference_steps=4,
        width=1420,
        height=1080,
        max_sequence_length=256,
        
    ).images[0]

    return image

@spaces.GPU #[uncomment to use ZeroGPU]
def generate_description_fn(
    image,
    progress=gr.Progress(track_tqdm=True),
):
    base64_image = encode_image(image)

    client = Groq(
        api_key=os.environ.get("GROQ_API_KEY"),
    )

    chat_completion = client.chat.completions.create(
        messages=[
            {
                "role": "user",
                "content": [
                    {"type": "text", "text": "What's in this image?"},
                    {
                        "type": "image_url",
                        "image_url": {
                            "url": f"data:image/jpeg;base64,{base64_image}",
                        },
                    },
                ],
            }
        ],
        model="llama-3.2-11b-vision-preview",
    )

    return chat_completion.choices[0].message.content


css = """
#col-container {
    margin: 0 auto;
    max-width: 640px;
}
"""

# generated_prompt = ""

with gr.Blocks(css=css) as demo:
    with gr.Row():
        with gr.Column(elem_id="col-input-image"):
            gr.Markdown(" # Drop your image here")
            input_image = gr.Image(type="filepath")
            generate_button = gr.Button("Generate", scale=0, variant="primary")
            generated_prompt = gr.Markdown("")
        with gr.Column(elem_id="col-container"):
            gr.Markdown(" # Text-to-Image Gradio Template")

            with gr.Row():
                prompt = gr.Text(
                    label="Prompt",
                    show_label=False,
                    max_lines=1,
                    placeholder="Enter your prompt",
                    container=False,
                )

                run_button = gr.Button("Run", scale=0, variant="primary")

            result = gr.Image(label="Result", show_label=False)

        gr.on(
            triggers=[run_button.click, prompt.submit],
            fn=infer,
            inputs=[
                prompt,
            ],
            outputs=[result],
        )

        gr.on(
            triggers=[generate_button.click],
            fn=generate_description_fn,
            inputs=[
                input_image,
            ],
            outputs=[generated_prompt],
        )

if __name__ == "__main__":
    demo.launch()