import gradio as gr import numpy as np import random import spaces from diffusers import ChromaPipeline, FlowMatchEulerDiscreteScheduler, EulerDiscreteScheduler, EulerAncestralDiscreteScheduler, DPMSolverMultistepScheduler, DDIMScheduler, PNDMScheduler, LMSDiscreteScheduler import torch device = "cuda" if torch.cuda.is_available() else "cpu" model_repo_id = "SG161222/SPARK.Chroma_preview" if torch.cuda.is_available(): torch_dtype = torch.bfloat16 else: torch_dtype = torch.float32 # Load the base pipeline first pipe = ChromaPipeline.from_pretrained(model_repo_id, torch_dtype=torch_dtype) pipe = pipe.to(device) # Map scheduler names to their classes SCHEDULERS = { "Flow Match Euler (Default)": FlowMatchEulerDiscreteScheduler, "Euler": EulerDiscreteScheduler, "Euler Ancestral": EulerAncestralDiscreteScheduler, "DPM Solver++": DPMSolverMultistepScheduler, "DDIM": DDIMScheduler, "PNDM": PNDMScheduler, "LMS": LMSDiscreteScheduler, } MAX_SEED = np.iinfo(np.int32).max MAX_IMAGE_SIZE = 1024 def set_scheduler(pipe, scheduler_name): """Set the scheduler for the pipeline.""" scheduler_class = SCHEDULERS.get(scheduler_name) if scheduler_class: try: scheduler = scheduler_class.from_pipe(pipe.scheduler) pipe.scheduler = scheduler except Exception: # For schedulers that don't have from_pipe method, create fresh instance try: scheduler = scheduler_class() pipe.scheduler = scheduler except Exception: pass # Keep current scheduler if it fails return pipe @spaces.GPU() def infer( prompt, negative_prompt="low quality, ugly, unfinished, out of focus, deformed, disfigured, blurry, smudged, restricted palette, flat colors", seed=42, randomize_seed=False, width=1024, height=1024, guidance_scale=3.0, num_inference_steps=40, scheduler_name="Flow Match Euler (Default)", progress=gr.Progress(track_tqdm=True) ): """Generate an image based on the prompt and settings.""" if randomize_seed: seed = random.randint(0, MAX_SEED) # Set the scheduler global pipe pipe = set_scheduler(pipe, scheduler_name) generator = torch.Generator(device).manual_seed(seed) image = pipe( prompt=prompt, negative_prompt=negative_prompt, guidance_scale=guidance_scale, num_inference_steps=num_inference_steps, width=width, height=height, generator=generator, num_images_per_prompt=1 ).images[0] return image, seed 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.", "A dog eating pizza", "The spirit of a tamagotchi wandering in San Francisco", ] # Custom CSS for styling css = """ #col-container { margin: 0 auto; max-width: 760px; } #run-button { align-self: stretch; font-weight: 600; } """ # Create custom theme custom_theme = gr.themes.Soft( primary_hue="indigo", secondary_hue="blue", neutral_hue="slate", font=gr.themes.GoogleFont("Inter"), text_size="lg", spacing_size="lg", radius_size="md" ).set( button_primary_background_fill="*primary_600", button_primary_background_fill_hover="*primary_700", block_title_text_weight="600", ) with gr.Blocks() as demo: gr.Markdown(f""" # 🎨 SPARK.Chroma [SPARK.Chroma](https://huggingface.co/SG161222/SPARK.Chroma_preview) is an 8.9B parameter text-to-image fine-tuned model based on FLUX.1-schnell *Built with [anycoder](https://huggingface.co/spaces/akhaliq/anycoder)* """) with gr.Column(elem_id="col-container"): with gr.Row(): prompt = gr.Text( label="Prompt", max_lines=1, placeholder="Enter your prompt", scale=2, ) negative_prompt = gr.Text( label="Negative prompt", max_lines=1, placeholder="Enter a negative prompt", value="low quality, ugly, unfinished, out of focus, deformed, disfigured, blurry, smudged, restricted palette, flat colors", scale=1, ) with gr.Row(): run_button = gr.Button("Run", variant="primary", scale=1, elem_id="run-button") result = gr.Image(label="Result", show_label=False, type="pil") with gr.Accordion("⚙️ Advanced Settings", open=False): with gr.Row(): scheduler_name = gr.Dropdown( label="Sampler / Scheduler", choices=list(SCHEDULERS.keys()), value="Flow Match Euler (Default)", info="Choose the sampling method. Flow Match Euler is recommended for FLUX models.", ) guidance_scale = gr.Slider( label="Guidance Scale", minimum=1.0, maximum=10.0, step=0.1, value=3.0, info="Higher values make the model follow the prompt more closely", ) with gr.Row(): seed = gr.Slider( label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=433, ) randomize_seed = gr.Checkbox( label="Randomize seed", value=True, info="Generate a random seed each time", ) 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="Inference Steps", minimum=1, maximum=100, step=1, value=40, info="More steps can produce higher quality but takes longer", ) gr.Markdown("### 💡 Example Prompts") gr.Examples( examples=examples, inputs=[prompt], outputs=[result, seed], fn=infer, cache_examples="lazy" ) # Event listeners for running inference gr.on( triggers=[run_button.click, prompt.submit, negative_prompt.submit], fn=infer, inputs=[prompt, negative_prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps, scheduler_name], outputs=[result, seed], api_visibility="public" ) # Launch with Gradio 6 syntax - theme and all app parameters go in launch() demo.queue().launch( theme=custom_theme, css=css, footer_links=[ {"label": "SPARK.Chroma Model", "url": "https://huggingface.co/SG161222/SPARK.Chroma_preview"}, {"label": "Built with anycoder", "url": "https://huggingface.co/spaces/akhaliq/anycoder"} ] )