app.py

import gradio as gr
import torch
import numpy as np
from diffusers import DiffusionPipeline
from diffusers.pipelines.stable_diffusion import StableDiffusionSafetyChecker
from functools import lru_cache
from PIL import Image
from transformers import CLIPImageProcessor

@lru_cache(maxsize=1)
def load_pipeline():
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    # Use FP16 when CUDA is available, along with a revision flag if supported.
    torch_dtype = torch.float16 if device.type == "cuda" else torch.float32
    revision = "fp16" if device.type == "cuda" else None

    base_model = "black-forest-labs/FLUX.1-dev"
    pipe = DiffusionPipeline.from_pretrained(
        base_model,
        torch_dtype=torch_dtype,
        low_cpu_mem_usage=True,
        revision=revision,
    )
    
    # Load LoRA weights
    lora_repo = "strangerzonehf/Flux-Super-Realism-LoRA"
    pipe.load_lora_weights(lora_repo)
    
    # Load safety checker and image processor.
    # If memory remains an issue, you can disable the safety checker below.
    safety_checker = StableDiffusionSafetyChecker.from_pretrained(
        "CompVis/stable-diffusion-safety-checker"
    )
    image_processor = CLIPImageProcessor.from_pretrained("openai/clip-vit-base-patch32")
    
    if device.type == "cuda":
        # Use attention slicing for further memory savings.
        pipe.enable_attention_slicing()
        # Offload layers to CPU when not in use.
        pipe.enable_sequential_cpu_offload()
    
    return pipe, safety_checker, image_processor

pipe, safety_checker, image_processor = load_pipeline()

def generate_image(
    prompt,
    seed=42,
    width=512,   # Keep resolution low by default
    height=512,
    guidance_scale=6,
    steps=28,
    progress=gr.Progress()
):
    try:
        progress(0, desc="Initializing...")
        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        generator = torch.Generator(device=device).manual_seed(seed)
        
        # Auto-add trigger words if not present
        if "super realism" not in prompt.lower():
            prompt = f"Super Realism, {prompt}"
        
        with torch.inference_mode():
            result = pipe(
                prompt=prompt,
                width=width,
                height=height,
                guidance_scale=guidance_scale,
                num_inference_steps=steps,
                generator=generator,
            )
            image = result.images[0]
        
        progress(1, desc="Safety checking...")
        # Process image for safety checking
        safety_input = image_processor(image, return_tensors="pt")
        np_image = np.array(image)
        _, nsfw_detected = safety_checker(
            images=[np_image],
            clip_input=safety_input.pixel_values
        )
        
        if nsfw_detected[0]:
            return Image.new("RGB", (width, height)), "NSFW content detected"
        
        # Clear CUDA cache
        if device.type == "cuda":
            torch.cuda.empty_cache()
        return image, "Generation successful"
    
    except Exception as e:
        return Image.new("RGB", (width, height)), f"Error: {str(e)}"

with gr.Blocks() as app:
    gr.Markdown("# Flux Super Realism Generator")
    
    with gr.Row():
        with gr.Column():
            prompt_input = gr.Textbox(label="Prompt", value="A portrait of a person")
            seed_input = gr.Slider(0, 1000, value=42, label="Seed")
            # Limit the resolution sliders to help avoid memory overuse.
            width_input = gr.Slider(256, 1024, value=512, step=64, label="Width")
            height_input = gr.Slider(256, 1024, value=512, step=64, label="Height")
            guidance_input = gr.Slider(1, 20, value=6, label="Guidance Scale")
            steps_input = gr.Slider(10, 100, value=28, label="Steps")
            submit = gr.Button("Generate")
        
        with gr.Column():
            output_image = gr.Image(label="Result", type="pil")
            status = gr.Textbox(label="Status")
    
    submit.click(
        generate_image,
        inputs=[prompt_input, seed_input, width_input, height_input, guidance_input, steps_input],
        outputs=[output_image, status]
    )
    
    # Queue settings to limit concurrent requests
    app.queue(max_size=3).launch()