import spaces # MUST be first import gradio as gr import torch import gc from diffusers import FluxPipeline, StableDiffusionXLPipeline, StableDiffusionPipeline from transformers import CLIPTextModel from huggingface_hub import hf_hub_download import os HF_TOKEN = os.environ.get("HF_TOKEN", "") token = HF_TOKEN if HF_TOKEN else None PRIVATE_REPO = "ksalokranjan/downloaded-models" MODELS = { # ── HF-HOSTED FULL PIPELINES ────────────────────────────────────────── "FLUX.1 Schnell (Fast)": { "load": "pretrained", "id": "black-forest-labs/FLUX.1-schnell", "cls": FluxPipeline, "is_flux": True, }, "FHDR Uncensored (FLUX)": { "load": "pretrained", "id": "kpsss34/FHDR_Uncensored", "cls": FluxPipeline, "is_flux": True, }, "UnfilteredAI NSFW Gen V2": { "load": "pretrained", "id": "UnfilteredAI/NSFW-gen-v2", "cls": StableDiffusionXLPipeline, "is_flux": False, }, "Illustrious SDXL NSFW": { "load": "pretrained", "id": "John6666/wai-nsfw-illustrious-v80-sdxl", "cls": StableDiffusionXLPipeline, "is_flux": False, }, "UnfilteredAI Anime NSFW": { "load": "pretrained", "id": "UnfilteredAI/NSFW-GEN-ANIME", "cls": StableDiffusionXLPipeline, "is_flux": False, }, "Shuttle 3 Diffusion": { "load": "pretrained", "id": "shuttleai/shuttle-3-diffusion", "cls": FluxPipeline, "is_flux": True, }, "OpenFLUX.1": { "load": "pretrained", "id": "ostris/OpenFLUX.1", "cls": FluxPipeline, "is_flux": True, }, "Fux Capacity NSFW (HF)": { "load": "pretrained", "id": "BKM1804/fuxCapacityNSFWPornFlux_40FP16", "cls": FluxPipeline, "is_flux": True, }, # ── YOUR PRIVATE REPO — FULL CHECKPOINTS ───────────────────────────── "Cat Ear SDXL Uncensored": { "load": "single_file", "filename": "cat-ear-sdxl-uncensored.safetensors", "cls": StableDiffusionXLPipeline, "is_flux": False, }, "FLUX NSFW Unlocked (v3)": { "load": "single_file", "filename": "flux-nsfw-unlocked.safetensors", "cls": FluxPipeline, "is_flux": True, }, "Fluxed Up NSFW": { "load": "single_file", "filename": "fluxed-up-nsfw.safetensors", "cls": FluxPipeline, "is_flux": True, }, "Fux Capacity NSFW (Full)": { "load": "single_file", "filename": "fux-capacity-nsfw.safetensors", "cls": FluxPipeline, "is_flux": True, }, # ── YOUR PRIVATE REPO — SD1.5 PRUNED CHECKPOINTS ───────────────────── "NSFW Master": { "load": "sd15_single_file", "filename": "nsfw-master.safetensors", "cls": StableDiffusionPipeline, "is_flux": False, }, # ── YOUR PRIVATE REPO — LoRAs ───────────────────────────────────────── "Nobody Rope (LoRA on FLUX)": { "load": "lora", "base_id": "black-forest-labs/FLUX.1-dev", "lora_filename": "nobody-rope.safetensors", "cls": FluxPipeline, "is_flux": True, }, "Realistic Z-Image (LoRA on SDXL)": { "load": "lora", "base_id": "stabilityai/stable-diffusion-xl-base-1.0", "lora_filename": "realistic-z-image-uncensored.safetensors", "cls": StableDiffusionXLPipeline, "is_flux": False, }, "FLUX NSFW Unlock (LoRA on FLUX)": { "load": "lora", "base_id": "black-forest-labs/FLUX.1-dev", "lora_filename": "flux-nsfw-unlock.safetensors", "cls": FluxPipeline, "is_flux": True, }, # ── R&D ─────────────────────────────────────────────────────────────── "⚗️ RnD: FLUX Uncensored Merged": { "load": "pretrained", "id": "shauray/FLUX-UNCENSORED-merged", "cls": FluxPipeline, "is_flux": True, }, } current_pipe = None current_model_name = None def load_model_to_cpu(model_name): global current_pipe, current_model_name if current_model_name == model_name and current_pipe is not None: return if current_pipe is not None: del current_pipe current_pipe = None gc.collect() cfg = MODELS[model_name] print(f"Loading {model_name}...") if cfg["load"] == "pretrained": current_pipe = cfg["cls"].from_pretrained( cfg["id"], torch_dtype=torch.bfloat16, token=token ) elif cfg["load"] == "single_file": local_path = hf_hub_download( repo_id=PRIVATE_REPO, filename=cfg["filename"], token=token ) current_pipe = cfg["cls"].from_single_file( local_path, torch_dtype=torch.bfloat16 ) elif cfg["load"] == "sd15_single_file": # Pruned SD1.5 checkpoint — needs text encoder loaded separately local_path = hf_hub_download( repo_id=PRIVATE_REPO, filename=cfg["filename"], token=token ) text_encoder = CLIPTextModel.from_pretrained( "runwayml/stable-diffusion-v1-5", subfolder="text_encoder", torch_dtype=torch.bfloat16 ) current_pipe = StableDiffusionPipeline.from_single_file( local_path, text_encoder=text_encoder, torch_dtype=torch.bfloat16 ) elif cfg["load"] == "lora": current_pipe = cfg["cls"].from_pretrained( cfg["base_id"], torch_dtype=torch.bfloat16, token=token ) current_pipe.load_lora_weights( PRIVATE_REPO, weight_name=cfg["lora_filename"], token=token ) current_model_name = model_name @spaces.GPU(duration=120) def run_on_gpu(prompt, model_name, negative_prompt, steps, guidance, width, height): global current_pipe cfg = MODELS[model_name] current_pipe.to("cuda") if cfg["is_flux"]: image = current_pipe( prompt=prompt, num_inference_steps=int(steps), guidance_scale=float(guidance), max_sequence_length=256, height=int(height), width=int(width), ).images[0] else: image = current_pipe( prompt=prompt, negative_prompt=negative_prompt.strip() or None, num_inference_steps=int(steps), guidance_scale=float(guidance), height=int(height), width=int(width), ).images[0] current_pipe.to("cpu") torch.cuda.empty_cache() return image def generate_image(prompt, model_name, negative_prompt, steps, guidance, width, height): if not prompt.strip(): return None, "⚠️ Please enter a prompt." try: load_model_to_cpu(model_name) image = run_on_gpu(prompt, model_name, negative_prompt, steps, guidance, width, height) return image, f"✅ Generated with {model_name}!" except Exception as e: return None, f"❌ Error: {str(e)}" with gr.Blocks(title="SilaI - Image Generator") as demo: gr.Markdown("# 🎨 SilaI Image Generator") with gr.Row(): with gr.Column(scale=2): prompt = gr.Textbox(label="Prompt", lines=3, placeholder="Describe your image...") negative_prompt = gr.Textbox(label="Negative Prompt (SDXL only)", lines=2, placeholder="ugly, deformed, blurry...") model_name = gr.Dropdown( choices=list(MODELS.keys()), value=list(MODELS.keys())[0], label="Model" ) with gr.Row(): steps = gr.Slider(1, 50, value=20, step=1, label="Steps") guidance = gr.Slider(0, 20, value=7.5, step=0.5, label="Guidance Scale") with gr.Row(): width = gr.Slider(256, 1280, value=576, step=8, label="Width") height = gr.Slider(256, 1280, value=576, step=8, label="Height") generate_btn = gr.Button("🎨 Generate", variant="primary") with gr.Column(scale=2): output_image = gr.Image(label="Generated Image") status = gr.Textbox(label="Status", interactive=False) generate_btn.click( fn=generate_image, inputs=[prompt, model_name, negative_prompt, steps, guidance, width, height], outputs=[output_image, status] ) demo.launch()