chroma

app.py

@@ -1,26 +1,26 @@
-import os
+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):
-    # Initialize Flux pipeline
-    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)
-    
     # Generate image
     image = flux_pipeline(
         prompt=prompt,