File size: 14,113 Bytes
f064a5b
 
 
 
 
 
 
3fd4792
f064a5b
 
 
 
9e8de34
3df04cd
f064a5b
 
c09ce28
3df04cd
f064a5b
45ddef4
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
e95ed8a
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
f064a5b
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
3df04cd
f064a5b
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
12a9a98
 
 
 
 
eb17cd6
12a9a98
 
 
f064a5b
 
 
96b7653
 
 
 
 
 
 
 
 
f064a5b
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
45ddef4
 
 
 
 
 
 
 
 
 
f064a5b
 
 
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
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
import os
import random
import sys
from typing import Sequence, Mapping, Any, Union
import torch
import gradio as gr
from PIL import Image
import numpy as np
from huggingface_hub import hf_hub_download
import spaces
from comfy import model_management

CHROMA_VERSION = "chroma-unlocked-v33.safetensors"

# Download required models
t5_path = hf_hub_download(repo_id="comfyanonymous/flux_text_encoders", filename="t5xxl_fp8_e4m3fn.safetensors", local_dir="models/text_encoders/")
vae_path = hf_hub_download(repo_id="lodestones/Chroma", filename="ae.safetensors", local_dir="models/vae")
unet_path = hf_hub_download(repo_id="lodestones/Chroma", filename=CHROMA_VERSION, local_dir="models/unet")

# Example prompts with their parameters
EXAMPLES = [
    [
        "A high-fashion close-up portrait of a blonde woman in clear sunglasses. The image uses a bold teal and red color split for dramatic lighting. The background is a simple teal-green. The photo is sharp and well-composed, and is designed for viewing with anaglyph 3D glasses for optimal effect. It looks professionally done.",
        "low quality, ugly, unfinished, out of focus, deformed, disfigure, blurry, smudged, restricted palette, flat colors",
        1024, 1024, 36, 3.0, 229
    ],
    [
        "A young man smiles broadly while wearing a blue captain's hat and blue mirrored aviator sunglasses. The man is the sharp focus of the image, set against a nicely blurred waterfront background featuring boats and buildings. The shot is cheerful, playful, and amateur looking.",
        "low quality, ugly, unfinished, out of focus, deformed, disfigure, blurry, smudged, restricted palette, flat colors",
        1024, 1024, 40, 2.1, 424
    ],
    [
        "A Shiba Inu smiles broadly while wearing a blue captain's hat and blue mirrored aviator sunglasses. The dog is the sharp focus of the image, set against a nicely blurred waterfront background featuring boats and buildings. The shot is cheerful, playful, and amateur looking.",
        "low quality, ugly, unfinished, out of focus, deformed, disfigure, blurry, smudged, restricted palette, flat colors",
        1024, 1024, 28, 2.9, 413
    ],
    [
        "A painting shows a bearded man in a beige shirt gripping a rope on a sailboat, battling choppy seas under a dramatic sunset. Another figure appears in the distant background, and the style is impressionistic with clear brushstrokes and high skill. The overall effect is one of adventure and freedom.",
        "low quality, ugly, unfinished, out of focus, deformed, disfigure, blurry, smudged, restricted palette, flat colors",
        1024, 1024, 33, 4.6, 281
    ],
    [
        "A painting shows an anthro male wolf sitting on a rock in a forest, anthro, male, solo, canine, canid, werewolf, muscular anthro, black body, black fur, reaching at viewer with his paws, detailed background, outdoors, outside. The scene is detailed, with a shallow depth of field focusing on the anthro wolf. The overall effect is high fantasy and cinematic.",
        "low quality, ugly, unfinished, out of focus, deformed, disfigure, blurry, smudged, restricted palette, flat colors",
        1024, 1024, 26, 3.8, 326
    ],
    [
        "A cute cartoon illustration of A massive black dragon with cosmic galaxy-patterned wings sits atop a castle overlooking a medieval port town at sunset. The style is bright and pastel-toned with soft lines. The artwork is clean and fairly well-executed.",
        "low quality, ugly, unfinished, out of focus, deformed, disfigure, blurry, smudged, restricted palette, flat colors",
        1024, 1024, 41, 5.7, 244
    ],
    [
        "A soft longhaired male anthro dog, with a white chest patch and bright yellow-green eyes, sits on a light grey kitchen island in a sun-drenched NYC apartment. The city view is blurred, showcasing a shallow depth of field. The cat has a dark collar with a small bell. The photo's aesthetic is clean, modern, and slightly dramatic due to the lighting.",
        "low quality, ugly, unfinished, out of focus, deformed, disfigure, blurry, smudged, restricted palette, flat colors",
        1024, 1024, 26, 4.7, 242
    ],
    [
        "A cheerful collage showcases 4 individual leopard geckos, each in a separate square. The geckos vary in color and pattern, but all are shown close-up and looking directly at the camera. Backgrounds differ, illustrating various terrarium elements and simple surfaces. High-quality shots, bright lighting, amateur photo quality. The overall aesthetic is heartwarming and simple.",
        "low quality, ugly, unfinished, out of focus, deformed, disfigure, blurry, smudged, restricted palette, flat colors",
        1024, 1024, 26, 4.7, 249
    ],
    [
        "A cheerful collage showcases 4 cute anthro wolf nick wilde plushie, each in a separate square. Backgrounds differ, illustrating various bed and simple surfaces. High-quality shots, bright lighting, amateur photo quality. The overall aesthetic is heartwarming and simple.",
        "low quality, ugly, unfinished, out of focus, deformed, disfigure, blurry, smudged, restricted palette, flat colors",
        1024, 1024, 26, 5.4, 263
    ],
    [
        "this picture depics a telegram sticker with thick white outlines of a male anthro fox",
        "low quality, ugly, unfinished, out of focus, deformed, disfigure, blurry, smudged, restricted palette, flat colors",
        1024, 1024, 28, 3.9, 399
    ],
    [
        "Extreme close-up of a single tiger eye, direct frontal view. Detailed iris and pupil. Sharp focus on eye texture and color. Natural lighting to capture authentic eye shine and depth. The word \"Chroma\" is painted over it in big, white brush strokes with visible texture.",
        "low quality, ugly, unfinished, out of focus, deformed, disfigure, blurry, smudged, restricted palette, flat colors",
        1024, 1024, 26, 4.0, 9
    ]
]

# Utility functions
def get_value_at_index(obj: Union[Sequence, Mapping], index: int) -> Any:
    try:
        return obj[index]
    except KeyError:
        return obj["result"][index]

def find_path(name: str, path: str = None) -> str:
    if path is None:
        path = os.getcwd()
    if name in os.listdir(path):
        path_name = os.path.join(path, name)
        print(f"{name} found: {path_name}")
        return path_name
    parent_directory = os.path.dirname(path)
    if parent_directory == path:
        return None
    return find_path(name, parent_directory)

def add_comfyui_directory_to_sys_path() -> None:
    comfyui_path = find_path("ComfyUI")
    if comfyui_path is not None and os.path.isdir(comfyui_path):
        sys.path.append(comfyui_path)
        print(f"'{comfyui_path}' added to sys.path")

def add_extra_model_paths() -> None:
    try:
        from main import load_extra_path_config
    except ImportError:
        from utils.extra_config import load_extra_path_config
    extra_model_paths = find_path("extra_model_paths.yaml")
    if extra_model_paths is not None:
        load_extra_path_config(extra_model_paths)
    else:
        print("Could not find the extra_model_paths config file.")

def import_custom_nodes() -> None:
    import asyncio
    import execution
    from nodes import init_extra_nodes
    import server
    loop = asyncio.new_event_loop()
    asyncio.set_event_loop(loop)
    server_instance = server.PromptServer(loop)
    execution.PromptQueue(server_instance)
    init_extra_nodes()

# Initialize paths
add_comfyui_directory_to_sys_path()
add_extra_model_paths()
import_custom_nodes()

# Import all necessary nodes
from nodes import (
    NODE_CLASS_MAPPINGS,
    CLIPTextEncode,
    CLIPLoader,
    VAEDecode,
    UNETLoader,
    VAELoader,
    SaveImage,
)

# Initialize all model loaders outside the function
randomnoise = NODE_CLASS_MAPPINGS["RandomNoise"]()
emptysd3latentimage = NODE_CLASS_MAPPINGS["EmptySD3LatentImage"]()
ksamplerselect = NODE_CLASS_MAPPINGS["KSamplerSelect"]()
cliploader = CLIPLoader()
t5tokenizeroptions = NODE_CLASS_MAPPINGS["T5TokenizerOptions"]()
cliptextencode = CLIPTextEncode()
unetloader = UNETLoader()
vaeloader = VAELoader()
cfgguider = NODE_CLASS_MAPPINGS["CFGGuider"]()
basicscheduler = NODE_CLASS_MAPPINGS["BasicScheduler"]()
samplercustomadvanced = NODE_CLASS_MAPPINGS["SamplerCustomAdvanced"]()
vaedecode = VAEDecode()
saveimage = SaveImage()

# Load models
cliploader_78 = cliploader.load_clip(
    clip_name="t5xxl_fp8_e4m3fn.safetensors", type="chroma", device="default"
)
t5tokenizeroptions_82 = t5tokenizeroptions.set_options(
    min_padding=1, min_length=0, clip=get_value_at_index(cliploader_78, 0)
)
unetloader_76 = unetloader.load_unet(
    unet_name=CHROMA_VERSION, weight_dtype="fp8_e4m3fn"
)
vaeloader_80 = vaeloader.load_vae(vae_name="ae.safetensors")

# Add all the models that load a safetensors file
model_loaders = [cliploader_78, unetloader_76, vaeloader_80]

# Check which models are valid and how to best load them
valid_models = [
    getattr(loader[0], 'patcher', loader[0]) 
    for loader in model_loaders
    if not isinstance(loader[0], dict) and not isinstance(getattr(loader[0], 'patcher', None), dict)
]

# Finally loads the models
model_management.load_models_gpu(valid_models)

@spaces.GPU
def generate_image(prompt, negative_prompt, width, height, steps, cfg, seed):
    with torch.inference_mode():
        # Set random seed if provided
        if seed == -1:
            seed = random.randint(1, 2**64)
        random.seed(seed)
        
        randomnoise_68 = randomnoise.get_noise(noise_seed=seed)
        emptysd3latentimage_69 = emptysd3latentimage.generate(
            width=width, height=height, batch_size=1
        )
        ksamplerselect_72 = ksamplerselect.get_sampler(sampler_name="euler")

        cliptextencode_74 = cliptextencode.encode(
            text=prompt,
            clip=get_value_at_index(t5tokenizeroptions_82, 0),
        )

        cliptextencode_75 = cliptextencode.encode(
            text=negative_prompt,
            clip=get_value_at_index(t5tokenizeroptions_82, 0),
        )

        cfgguider_73 = cfgguider.get_guider(
            cfg=cfg,
            model=get_value_at_index(unetloader_76, 0),
            positive=get_value_at_index(cliptextencode_74, 0),
            negative=get_value_at_index(cliptextencode_75, 0),
        )

        basicscheduler_84 = basicscheduler.get_sigmas(
            scheduler="beta",
            steps=steps,
            denoise=1,
            model=get_value_at_index(unetloader_76, 0),
        )

        samplercustomadvanced_67 = samplercustomadvanced.sample(
            noise=get_value_at_index(randomnoise_68, 0),
            guider=get_value_at_index(cfgguider_73, 0),
            sampler=get_value_at_index(ksamplerselect_72, 0),
            sigmas=get_value_at_index(basicscheduler_84, 0),
            latent_image=get_value_at_index(emptysd3latentimage_69, 0),
        )

        vaedecode_79 = vaedecode.decode(
            samples=get_value_at_index(samplercustomadvanced_67, 0),
            vae=get_value_at_index(vaeloader_80, 0),
        )

        # Save image using SaveImage node with simple string prefix
        saved = saveimage.save_images(
            filename_prefix="Chroma_Generated",
            images=get_value_at_index(vaedecode_79, 0),
        )
        
        # Return the path to the saved image
        saved_path = f"output/{saved['ui']['images'][0]['filename']}"
        return saved_path

# Create Gradio interface
with gr.Blocks() as app:
    gr.Markdown("""
# Chroma

Model: [Chroma](https://huggingface.co/lodestones/Chroma) by [lodestones](https://huggingface.co/lodestones)

Run any ComfyUI Workflow on Spaces: [ComfyUI Workflows](https://huggingface.co/blog/run-comfyui-workflows-on-spaces)

Space Author: [GitHub](https://github.com/gokayfem) | [X.com](https://x.com/gokayfem)
    """)
    
    with gr.Row():
        with gr.Column():
            prompt = gr.Textbox(
                label="Prompt",
                placeholder="Enter your prompt here...",
                lines=3
            )
            negative_prompt = gr.Textbox(
                label="Negative Prompt",
                placeholder="Enter negative prompt here...",
                value="low quality, ugly, unfinished, out of focus, deformed, disfigure, blurry, smudged, restricted palette, flat colors",
                lines=2
            )
            
            with gr.Row():
                width = gr.Slider(
                    minimum=512,
                    maximum=2048,
                    value=1024,
                    step=64,
                    label="Width"
                )
                height = gr.Slider(
                    minimum=512,
                    maximum=2048,
                    value=1024,
                    step=64,
                    label="Height"
                )
            
            with gr.Row():
                steps = gr.Slider(
                    minimum=1,
                    maximum=50,
                    value=26,
                    step=1,
                    label="Steps"
                )
                cfg = gr.Slider(
                    minimum=1,
                    maximum=20,
                    value=4,
                    step=0.5,
                    label="CFG Scale"
                )
                seed = gr.Number(
                    value=-1,
                    label="Seed (-1 for random)"
                )
            
            generate_btn = gr.Button("Generate")
        
        with gr.Column():
            output_image = gr.Image(label="Generated Image")
    
    generate_btn.click(
        fn=generate_image,
        inputs=[prompt, negative_prompt, width, height, steps, cfg, seed],
        outputs=[output_image]
    )
    
    # Add examples section
    gr.Examples(
        examples=EXAMPLES,
        inputs=[prompt, negative_prompt, width, height, steps, cfg, seed],
        outputs=[output_image],
        fn=generate_image,
        cache_examples=True,
        label="Example Prompts - Click to try!"
    )

if __name__ == "__main__":
    app.launch(share=True)