app.py

import spaces
import gradio as gr
import torch
from diffusers import FluxPipeline, FluxTransformer2DModel, FlowMatchEulerDiscreteScheduler
from huggingface_hub import hf_hub_download
from PIL import Image
import numpy as np
import random
import os

hf_token = os.environ.get('HF_TOKEN')

# Constants
model = "black-forest-labs/FLUX.1-dev"
MAX_SEED = np.iinfo(np.int32).max
MAX_IMAGE_SIZE = 2048

@spaces.GPU()
def infer(prompt, width, height, num_inference_steps, guidance_scale, nums, seed=42, randomize_seed=True, progress=gr.Progress(track_tqdm=True)):
    device = "cuda" if torch.cuda.is_available() else "cpu"
    
    # Initialize model inside the GPU-enabled function
    try:
        transformer = FluxTransformer2DModel.from_single_file(
            "https://huggingface.co/lodestones/Chroma/resolve/main/chroma-unlocked-v27.safetensors",
            torch_dtype=torch.bfloat16,
            token=hf_token
        )
    except KeyError as e:
        print(f"Error loading chroma-unlocked-v27.safetensors: {e}. Falling back to pretrained model.")
        transformer = FluxTransformer2DModel.from_pretrained(
            "lodestones/Chroma",
            subfolder="transformer",
            torch_dtype=torch.bfloat16,
            token=hf_token
        )
    
    pipe = FluxPipeline.from_pretrained(
        model,
        transformer=transformer,
        torch_dtype=torch.bfloat16,
        token=hf_token
    )
    pipe.scheduler = FlowMatchEulerDiscreteScheduler.from_config(
        pipe.scheduler.config, use_beta_sigmas=True
    )
    pipe.to(device)

    # Generate images
    if randomize_seed:
        seed = random.randint(0, MAX_SEED)
    generator = torch.Generator(device=device).manual_seed(seed)
    
    images = pipe(
        prompt=prompt,
        width=width,
        height=height,
        num_inference_steps=num_inference_steps,
        guidance_scale=guidance_scale,
        num_images_per_prompt=nums,
        generator=generator
    ).images

    return images, seed

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

with gr.Blocks(css=css) as demo:
    with gr.Column(elem_id="col-container"):
        gr.HTML("<h1><center>Model Testing</center></h1><p><center>Chroma</center></p>")
        
        with gr.Row():
            prompt = gr.Text(
                label="Prompt",
                show_label=False,
                max_lines=2,
                placeholder="Enter your prompt",
                container=False,
            )
            run_button = gr.Button("Run", scale=0)
        
        result = gr.Gallery(label="Gallery", format="png", columns=1, preview=True, height=400)
        
        with gr.Accordion("Advanced Settings", open=False):
            with gr.Row():
                width = gr.Slider(
                    label="Width",
                    minimum=256,
                    maximum=MAX_IMAGE_SIZE,
                    step=32,
                    value=1024,
                )
                height = gr.Slider(
                    label="Height",
                    minimum=256,
                    maximum=MAX_IMAGE_SIZE,
                    step=32,
                    value=1024,
                )
            
            with gr.Row():
                num_inference_steps = gr.Slider(
                    label="Number of inference steps",
                    minimum=1,
                    maximum=50,
                    step=1,
                    value=30,
                )
                guidance_scale = gr.Slider(
                    label="Guidance Scale",
                    minimum=0,
                    maximum=10,
                    step=0.1,
                    value=3.5,
                )

            with gr.Row():
                nums = gr.Slider(
                    label="Number of Images",
                    minimum=1,
                    maximum=2,
                    step=1,
                    value=1,
                    scale=1,
                )
                seed = gr.Slider(
                    label="Seed",
                    minimum=0,
                    maximum=MAX_SEED,
                    step=1,
                    value=-1,
                )
                randomize_seed = gr.Checkbox(label="Randomize seed", value=True)

    gr.on(
        triggers=[run_button.click, prompt.submit],
        fn=infer,
        inputs=[prompt, width, height, num_inference_steps, guidance_scale, nums, seed, randomize_seed],
        outputs=[result, seed]
    )

demo.launch()