app.py

import spaces
import torch
from custom_pipeline import FluxPipeline
import gradio as gr

# Load the pre-trained diffusion model
pipe = FluxPipeline.from_pretrained('terminusresearch/flux-booru-v0.2', torch_dtype=torch.bfloat16)
pipe.to('cuda')
import re

def extract_resolution(resolution_str):
 match = re.match(r'(\d+)x(\d+)', resolution_str)
 if match:
 width = int(match.group(1))
 height = int(match.group(2))
 return (width, height)
 else:
 return None

# Define the image generation function with adjustable parameters and a progress bar
@spaces.GPU(duration=115)
def generate(prompt, guidance_scale, guidance_scale_real, num_inference_steps, resolution, negative_prompt):
 width, height = extract_resolution(resolution) or (1024, 1024)
 return pipe(
 prompt,
 negative_prompt=negative_prompt,
 guidance_scale=guidance_scale,
 guidance_scale_real=guidance_scale_real,
 num_inference_steps=num_inference_steps,
 width=width, height=height
 ).images

# Example prompts to demonstrate the model's capabilities
example_prompts = [
 ["A futuristic cityscape at night under a starry sky", 5.5, 35, "blurry, overexposed"],
 ["A serene landscape with a flowing river and autumn trees", 6, 35, "crowded, noisy"],
 ["An abstract painting of joy and energy in bright colors", 5.0, 35, "dark, dull"]
]

# Create a Gradio interface, 1024x1024,1152x960,896x1152
iface = gr.Interface(
 fn=generate,
 inputs=[
 gr.Text(label="Enter your prompt"),
 gr.Slider(1, 4, step=0.1, label="Guidance Scale (Flux)", value=4.0),
 gr.Slider(1, 12, step=0.1, label="Guidance Scale (CFG)", value=5.4),
 gr.Slider(1, 35, step=1, label="Number of Inference Steps", value=20),
 gr.Radio(["1024x1024", "1152x960", "896x1152", "1280x768", "1920x1080", "576x1024", "576x1344"], label="Resolution", value="1024x1024"),
 gr.Text(value="underexposed, blurry, ugly, washed-out", label="Negative Prompt")
 ],
 outputs=gr.Gallery(height=1024, min_width=1024, columns=2),
 examples=example_prompts,
 title="FluxBooru 12B",
 description=(
 "This is a Flux-based 12B parameter model, fully trained across 8xH100 GPUs to reintroduce classifier-free guidance (CFG) sampling."
 "<br />This model is being <strong>actively trained</strong> on 3.5M samples across a wide distribution of photos, synthetic data, cinema, anime, and safe-for-work furry art."
 "<br />"
 "<br />&nbsp;The datasets been filtered for extremist and illegal content, but it is possible to produce toxic outputs. <strong>This model has not been safety-aligned or fine-tuned</strong>."
 " You may receive non-aesthetic results, or prompts might be partially or wholly ignored."
 "<br />Although celebrity names and artist styles haven't been scrubbed from the datasets, the low volume of these samples in the training set result in a lack of representation for public figures."
 "<br />"
 "<br />Be mindful when using this demo space that you do not inadvertently share images without adequate preparation and informing the receivers that these images are AI generated."
 "<br />"
 "<br />This model was developed and funded by <strong>Terminus Research Group</strong>."
 "<br />"
 "<br />"
 "<ul>"
 "<li>Lead trainer: @pseudoterminalx (bghira@GitHub)</li>"
 "<li>Datasets: @jimmycarter, @pseudoterminalx</li>"
 "</ul>"
 )
).launch()