Spaces:
Running
on
Zero
Running
on
Zero
File size: 2,550 Bytes
aca93c2 def9f12 a7438d2 e25da44 eeb2755 aca93c2 e25da44 a7438d2 e25da44 |
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 |
import spaces
import torch
import gradio as gr
from diffusers import FluxTransformer2DModel, FluxPipeline, BitsAndBytesConfig
from transformers import T5EncoderModel, BitsAndBytesConfig as BitsAndBytesConfigTF
# Initialize model outside the function
device = "cuda" if torch.cuda.is_available() else "cpu"
dtype = torch.bfloat16
single_file_base_model = "camenduru/FLUX.1-dev-diffusers"
file_url = "https://huggingface.co/lodestones/Chroma/resolve/main/chroma-unlocked-v31.safetensors"
quantization_config_tf = BitsAndBytesConfigTF(load_in_8bit=True, bnb_8bit_compute_dtype=torch.bfloat16)
text_encoder_2 = T5EncoderModel.from_pretrained(single_file_base_model, subfolder="text_encoder_2", torch_dtype=dtype, config=single_file_base_model, quantization_config=quantization_config_tf)
quantization_config = BitsAndBytesConfig(load_in_4bit=True, bnb_4bit_quant_type="nf4", bnb_4bit_use_double_quant=True, bnb_4bit_compute_dtype=torch.bfloat16)
transformer = FluxTransformer2DModel.from_single_file(file_url, subfolder="transformer", torch_dtype=dtype, config=single_file_base_model, quantization_config=quantization_config)
flux_pipeline = FluxPipeline.from_pretrained(single_file_base_model, transformer=transformer, text_encoder_2=text_encoder_2, torch_dtype=dtype)
flux_pipeline.to(device)
@spaces.GPU()
def generate_image(prompt, negative_prompt="", num_inference_steps=30, guidance_scale=7.5):
# Generate image
image = flux_pipeline(
prompt=prompt,
negative_prompt=negative_prompt,
num_inference_steps=num_inference_steps,
guidance_scale=guidance_scale
).images[0]
return image
# Create Gradio interface
iface = gr.Interface(
fn=generate_image,
inputs=[
gr.Textbox(label="Prompt", placeholder="Enter your prompt here..."),
gr.Textbox(label="Negative Prompt", placeholder="Enter negative prompt here...", value=""),
gr.Slider(minimum=1, maximum=100, value=30, step=1, label="Number of Inference Steps"),
gr.Slider(minimum=1.0, maximum=20.0, value=7.5, step=0.1, label="Guidance Scale")
],
outputs=gr.Image(label="Generated Image"),
title="Chroma Image Generator",
description="Generate images using the Chroma model with FLUX pipeline",
examples=[
["A beautiful sunset over mountains, photorealistic, 8k", "blurry, low quality, distorted", 30, 7.5],
["A futuristic cityscape at night, neon lights, cyberpunk style", "ugly, deformed, low resolution", 30, 7.5]
]
)
if __name__ == "__main__":
iface.launch() |