"""Hugging Face Gradio Space entry point for Ragnarok SDXL workflows.""" import os import random import threading import traceback import gradio as gr import spaces import torch from huggingface_hub import hf_hub_download from diffusers import ( ControlNetModel, StableDiffusionXLControlNetImg2ImgPipeline, StableDiffusionXLControlNetPipeline, StableDiffusionXLImg2ImgPipeline, StableDiffusionXLPipeline, ) from character_population import DEFAULT_VISUAL_STYLE, PRODUCTION_CLEAN from hf_space import load_styles, prepare_prompts from space_canny import CannyDownload from space_characters import ALL_TEMPLATES, HEADERS as CHARACTER_HEADERS, parse_batch_prompts, plan_characters, population_render_summary, selected_character_prompt from space_images import IMAGE_TO_IMAGE, RESIZE_MODES, TEXT_TO_IMAGE, prepare_images from space_loras import SDXL_LORAS, prepare_lora, requested_loras from space_models import MODEL_CHOICES, MODEL_ID, RAGNAROK, prepare_model from space_themes import SPACE_SHIPS, default_visual_style STYLES = load_styles() pipe = img2img_pipe = canny_pipe = canny_img2img_pipe = None model_load_error = None CONTROLNET_ID = "diffusers/controlnet-canny-sdxl-1.0" def model_status_text(): if model_load_error: return f"Ragnarok failed to load: {model_load_error}" if pipe is None: return "Ragnarok is loading. Image generation will be available when it is ready." return "Ragnarok is ready. Canny is optional and loads when enabled." def canny_cache_dir(): return "/data/ruinedfooocus/huggingface" if os.path.isdir("/data") else None def download_canny_files(): for filename in ("config.json", "diffusion_pytorch_model.fp16.safetensors"): hf_hub_download(repo_id=CONTROLNET_ID, filename=filename, cache_dir=canny_cache_dir()) CANNY_DOWNLOAD = CannyDownload(download_canny_files) def canny_status_text(enabled): return CANNY_DOWNLOAD.status(enabled) def request_canny(enabled): if enabled: CANNY_DOWNLOAD.start() return canny_status_text(enabled) def get_canny_pipelines(): """Create ControlNet inside a ZeroGPU worker after its files are cached.""" global canny_pipe, canny_img2img_pipe if canny_pipe is None: controlnet = ControlNetModel.from_pretrained( CONTROLNET_ID, variant="fp16", use_safetensors=True, torch_dtype=torch.float16, cache_dir=canny_cache_dir(), local_files_only=True, ).to("cuda") canny_pipe = StableDiffusionXLControlNetPipeline.from_pipe(pipe, controlnet=controlnet) canny_img2img_pipe = StableDiffusionXLControlNetImg2ImgPipeline.from_pipe(pipe, controlnet=controlnet) return canny_pipe, canny_img2img_pipe def load_models(): """Register Ragnarok on ZeroGPU without waiting for optional ControlNet.""" global pipe, img2img_pipe, model_load_error try: print("Preparing Ragnarok checkpoint", flush=True) ragnarok_path = prepare_model(RAGNAROK) print("Loading Ragnarok pipeline", flush=True) loaded_pipe = StableDiffusionXLPipeline.from_single_file( ragnarok_path, config=MODEL_ID, torch_dtype=torch.float16, ) print("Placing Ragnarok pipeline on ZeroGPU", flush=True) loaded_pipe.to("cuda") loaded_pipe.enable_vae_slicing() loaded_img2img = StableDiffusionXLImg2ImgPipeline.from_pipe(loaded_pipe) img2img_pipe = loaded_img2img pipe = loaded_pipe print("Ragnarok text-to-image and image-to-image pipelines ready", flush=True) except Exception as error: model_load_error = str(error) traceback.print_exc() def fetch_selected_model(model_name, mode, input_image, use_canny, canny_image, width, height, resize_mode, strength, low_threshold, high_threshold, canny_scale, canny_start, canny_end, *lora_slots, progress=gr.Progress()): """Validate image controls and fetch selected LoRAs before GPU allocation.""" if model_name != RAGNAROK: raise gr.Error(f"Unsupported model: {model_name}") if pipe is None: raise gr.Error(model_status_text()) try: _, edge_preview = prepare_images( mode, input_image, use_canny, canny_image, width, height, resize_mode, low_threshold, high_threshold, strength, canny_scale, canny_start, canny_end, ) if use_canny: CANNY_DOWNLOAD.wait() selected = requested_loras(lora_slots) for name, _ in selected: prepare_lora(name, progress=progress) except Exception as error: raise gr.Error(str(error)) from error return f"{model_name} ready with {len(selected)} LoRA(s)", edge_preview def render_image(prompt, negative_prompt, style, model_name, width, height, steps, guidance, seed, mode, input_image, resize_mode, strength, use_canny, canny_image, low_threshold, high_threshold, canny_scale, canny_start, canny_end, *lora_slots): if pipe is None: raise gr.Error(model_status_text()) if seed < 0: seed = random.SystemRandom().randrange(0, 2**32) seed = int(seed) positive, negative = prepare_prompts(prompt, negative_prompt, style, seed, STYLES) generator = torch.Generator(device="cuda").manual_seed(seed) if model_name != RAGNAROK: raise gr.Error(f"Unsupported model: {model_name}") try: prepared_input, edge_image = prepare_images( mode, input_image, use_canny, canny_image, width, height, resize_mode, low_threshold, high_threshold, strength, canny_scale, canny_start, canny_end, ) except ValueError as error: raise gr.Error(str(error)) from error selected = requested_loras(lora_slots) pipe.unload_lora_weights() adapters = [] weights = [] for index, (name, weight) in enumerate(selected): adapter_name = f"slot_{index}" pipe.load_lora_weights(prepare_lora(name), adapter_name=adapter_name) adapters.append(adapter_name) weights.append(weight) if adapters: pipe.set_adapters(adapters, adapter_weights=weights) if edge_image is not None: guided_pipe, guided_img2img_pipe = get_canny_pipelines() active_pipe = guided_img2img_pipe if mode == IMAGE_TO_IMAGE else guided_pipe else: active_pipe = img2img_pipe if mode == IMAGE_TO_IMAGE else pipe options = dict( prompt=positive, negative_prompt=negative, width=int(width), height=int(height), num_inference_steps=int(steps), guidance_scale=float(guidance), generator=generator, ) if prepared_input is not None: options.update(image=prepared_input, strength=float(strength)) if edge_image is not None: options.update( controlnet_conditioning_scale=float(canny_scale), control_guidance_start=float(canny_start), control_guidance_end=float(canny_end), ) if prepared_input is None: options["image"] = edge_image else: options["control_image"] = edge_image result = active_pipe(**options) return result.images[0], { "seed": seed, "prompt": positive, "negative_prompt": negative, "width": int(width), "height": int(height), "model": model_name, "mode": mode, "image_strength": float(strength) if prepared_input else None, "canny": bool(edge_image), "canny_scale": float(canny_scale) if edge_image else None, "canny_thresholds": [int(low_threshold), int(high_threshold)] if edge_image else None, "loras": selected, } @spaces.GPU(duration=120) def generate(prompt, negative_prompt, style, model_name, width, height, steps, guidance, seed, mode, input_image, resize_mode, strength, use_canny, canny_image, low_threshold, high_threshold, canny_scale, canny_start, canny_end, *lora_slots): return render_image( prompt, negative_prompt, style, model_name, width, height, steps, guidance, seed, mode, input_image, resize_mode, strength, use_canny, canny_image, low_threshold, high_threshold, canny_scale, canny_start, canny_end, *lora_slots, ) def build_character_panel(total, men, women, min_age, max_age, cast_seed, visual_style, template): try: rows, choices, prompts, batch_prompt, headers = plan_characters( total, men, women, min_age, max_age, cast_seed, visual_style, template, ) except ValueError as error: raise gr.Error(str(error)) from error return ( f"Prepared **{len(prompts)} {'ship' if template == SPACE_SHIPS else 'character'} prompts** using **{template}**.", gr.update(headers=headers, value=rows, visible=True), gr.update(choices=choices, value="1"), prompts[0], batch_prompt, prompts, ) def select_population_template(template): ships = template == SPACE_SHIPS description = ( "Ship designs use the total and seed; gender and age controls do not apply." if ships else f"Selected **{template}**. Build prompts to preview this cast." ) return ( default_visual_style(template), gr.update(visible=not ships), gr.update(visible=not ships), gr.update(visible=not ships), gr.update(visible=not ships), gr.update(label="Total ships" if ships else "Total characters"), description, gr.update(value=[], visible=False), gr.update(choices=[], value=None), "", "", [], ) def copy_character_to_generate(selected, width, height, *lora_slots): return selected, width, height, *lora_slots def start_population_generation(batch_text, negative_prompt, style, model_name, width, height, steps, guidance, seed, mode, input_image, resize_mode, strength, use_canny, canny_image, low_threshold, high_threshold, canny_scale, canny_start, canny_end, *lora_slots): try: prompts = parse_batch_prompts(batch_text) model_message, edge_preview = fetch_selected_model( model_name, mode, input_image, use_canny, canny_image, width, height, resize_mode, strength, low_threshold, high_threshold, canny_scale, canny_start, canny_end, *lora_slots, ) except (ValueError, gr.Error) as error: raise gr.Error(str(error)) from error settings = [ negative_prompt, style, model_name, width, height, steps, guidance, seed, mode, input_image, resize_mode, strength, use_canny, canny_image, low_threshold, high_threshold, canny_scale, canny_start, canny_end, *lora_slots, ] state = {"prompts": prompts, "index": 0, "settings": settings, "running": True} return state, [], None, None, f"Starting population: **0 / {len(prompts)}** images complete.", gr.update(active=True), model_message, edge_preview, gr.update(interactive=False) @spaces.GPU(duration=120) def generate_next_character(state, gallery): if not state or not state.get("running"): return state, gallery or [], gr.skip(), gr.skip(), "No population generation is running.", gr.update(active=False), gr.update(interactive=True) prompts = state["prompts"] index = state["index"] if index >= len(prompts): return state, gallery or [], gr.skip(), gr.skip(), f"Population complete: **{len(prompts)} / {len(prompts)}** images.", gr.update(active=False), gr.update(interactive=True) settings = list(state["settings"]) if settings[7] < 0: settings[7] = random.SystemRandom().randrange(0, 2**32) try: picture, details = render_image(prompts[index], *settings) except Exception as error: stopped = {**state, "running": False} return stopped, gallery or [], gr.skip(), gr.skip(), f"Stopped after **{index} / {len(prompts)}** images: {error}", gr.update(active=False), gr.update(interactive=True) completed = list(gallery or []) lora_names = ", ".join(f"{name.removesuffix('.safetensors')} × {weight:g}" for name, weight in details["loras"]) caption = f"#{index + 1} · {picture.width}×{picture.height} · {lora_names or 'no LoRAs'} · seed {details['seed']}" completed.append((picture, caption)) finished = index + 1 next_state = {**state, "index": finished, "running": finished < len(prompts)} status = f"Population complete: **{finished} / {len(prompts)}** images." if finished == len(prompts) else f"Generating population: **{finished} / {len(prompts)}** images complete." return next_state, completed, picture, details, status, gr.update(active=finished < len(prompts)), gr.update(interactive=finished == len(prompts)) def stop_population_after_failure(state): if not state: return state, "Population generation stopped. No images were completed.", gr.update(active=False), gr.update(interactive=True) stopped = {**state, "running": False} return stopped, f"Population generation stopped after **{state['index']} / {len(state['prompts'])}** images. The GPU request failed; completed images remain in the gallery.", gr.update(active=False), gr.update(interactive=True) with gr.Blocks(title="RuinedFooocus · ZeroGPU") as demo: gr.Markdown("# RuinedFooocus · ZeroGPU\nRagnarok SDXL generation and image editing with styles, wildcards, Canny and LoRAs.") with gr.Tab("Generate"): with gr.Row(): with gr.Column(): prompt = gr.Textbox(label="Prompt", lines=4, placeholder="Describe the image you want to create") negative_prompt = gr.Textbox(label="Negative prompt", lines=2) style = gr.Dropdown(choices=list(STYLES), value=[], multiselect=True, label="Style Selection") model = gr.Dropdown(choices=MODEL_CHOICES, value=RAGNAROK, label="Model") model_status = gr.Markdown(model_status_text()) gr.Timer(5).tick(model_status_text, outputs=model_status, queue=False) with gr.Accordion("Image input and Canny guidance", open=True): mode = gr.Radio([TEXT_TO_IMAGE, IMAGE_TO_IMAGE], value=TEXT_TO_IMAGE, label="Workflow") input_image = gr.Image(label="Image to edit", type="pil") resize_mode = gr.Radio(RESIZE_MODES, value=RESIZE_MODES[0], label="Image resize") strength = gr.Slider(0.05, 1.0, value=0.6, step=0.05, label="Image change strength") use_canny = gr.Checkbox(label="Use Canny edges", value=False) canny_status = gr.Markdown(canny_status_text(False)) use_canny.change(request_canny, inputs=use_canny, outputs=canny_status, queue=False) gr.Timer(5).tick(canny_status_text, inputs=use_canny, outputs=canny_status, queue=False) canny_image = gr.Image(label="Separate Canny guide (optional; otherwise uses image to edit)", type="pil") with gr.Row(): low_threshold = gr.Slider(0, 254, value=100, step=1, label="Canny low threshold") high_threshold = gr.Slider(1, 255, value=200, step=1, label="Canny high threshold") with gr.Row(): canny_scale = gr.Slider(0, 2, value=0.7, step=0.05, label="Canny guidance strength") canny_start = gr.Slider(0, 0.95, value=0, step=0.05, label="Canny start") canny_end = gr.Slider(0.05, 1, value=1, step=0.05, label="Canny end") lora_controls = [] with gr.Accordion("LoRAs", open=False): for index in range(5): with gr.Row(): lora_controls.append(gr.Dropdown(choices=["None", *SDXL_LORAS], value="None", label=f"LoRA {index + 1}")) lora_controls.append(gr.Number(value=0.5, label="Weight", minimum=0, maximum=5, step=0.05)) with gr.Row(): width = gr.Slider(512, 1536, value=1024, step=64, label="Width") height = gr.Slider(512, 1536, value=1024, step=64, label="Height") with gr.Row(): steps = gr.Slider(10, 50, value=25, step=1, label="Steps") guidance = gr.Slider(1, 15, value=7, step=0.5, label="Guidance") seed = gr.Number(value=-1, precision=0, label="Seed (-1 for random)") button = gr.Button("Generate", variant="primary") with gr.Column(): image = gr.Image(label="Generated image", type="pil") canny_preview = gr.Image(label="Canny edge preview", type="pil", interactive=False) metadata = gr.JSON(label="Generation details") with gr.Tab("Character Prompts"): gr.Markdown( "Build a balanced SDXL cast using the original templates or choose " "aliens, medieval warriors, 80s street fighters, Viking warriors, " "space ships, or wizards. Character templates request four full-body " "views in one horizontal row; ships request four views in one row. " "Each finished sheet appears " "below. Edit one prompt for the Generate tab, or " "generate the whole prepared set below. Set the sheet size and LoRA " "weights here. Other controls (styles, image input, Canny, steps, " "guidance, and seed) come from Generate. Each image uses one GPU request." ) with gr.Row(): cast_total = gr.Number(label="Total characters", value=8, precision=0, minimum=1, maximum=50) cast_men = gr.Number(label="Men", value=4, precision=0, minimum=0, maximum=50) cast_women = gr.Number(label="Women", value=4, precision=0, minimum=0, maximum=50) with gr.Row(): cast_min_age = gr.Number(label="Minimum age", value=24, precision=0, minimum=18, maximum=100) cast_max_age = gr.Number(label="Maximum age", value=70, precision=0, minimum=18, maximum=100) cast_seed = gr.Number(label="Cast seed", value=0, precision=0, info="Reuse this value to rebuild the same character matrix.") cast_template = gr.Dropdown( label="Prompt template / theme", choices=list(ALL_TEMPLATES), value=PRODUCTION_CLEAN, info="The original three templates plus themed characters and spacecraft.", ) cast_visual_style = gr.Textbox( label="Shared visual language", value=DEFAULT_VISUAL_STYLE, lines=2, info="Editable for themed templates and Production clean. The two original remixes keep their own wording.", ) with gr.Accordion("Sprite sheet output and LoRA influence", open=True): with gr.Row(): cast_width = gr.Slider(512, 1536, value=1536, step=64, label="Whole sheet width (px)") cast_height = gr.Slider(512, 1536, value=640, step=64, label="Whole sheet height (px)") gr.Markdown("These dimensions are for the **single final image** containing every view, not for each character pose.") cast_lora_controls = [] for index in range(5): with gr.Row(): cast_lora_controls.append(gr.Dropdown( choices=["None", *SDXL_LORAS], value="None", label=f"Population LoRA {index + 1}", )) cast_lora_controls.append(gr.Slider( 0, 5, value=1, step=0.05, label="Influence / weight", info="0 = off; 1 = usual full weight; higher values may overpower the image.", )) cast_settings_summary = gr.Markdown(population_render_summary(1536, 640, *(["None", 1] * 5))) for control in [cast_width, cast_height, *cast_lora_controls]: control.change( population_render_summary, inputs=[cast_width, cast_height, *cast_lora_controls], outputs=cast_settings_summary, queue=False, ) with gr.Row(): cast_build = gr.Button("Build Population Prompts") cast_generate = gr.Button("Generate Population", variant="primary") cast_reference = gr.Button("Use example turnaround settings") cast_status = gr.Markdown() cast_matrix = gr.Dataframe( headers=CHARACTER_HEADERS, datatype=["number"] + ["str"] * (len(CHARACTER_HEADERS) - 1), interactive=False, wrap=True, visible=False, label="Planned character matrix", ) cast_prompts = gr.State([]) cast_choice = gr.Dropdown(label="Character", choices=[]) cast_selected_prompt = gr.Textbox(label="Selected character prompt (editable)", lines=8) cast_use = gr.Button("Use selected prompt and sheet settings in Generate") cast_batch = gr.Textbox(label="All character prompts (separated by ---)", lines=10) cast_run_status = gr.Markdown() cast_image = gr.Image(label="Latest sprite sheet", type="pil", interactive=False) cast_details = gr.JSON(label="Latest sheet settings and applied LoRAs") cast_gallery = gr.Gallery(label="All generated sprite sheets", columns=4, type="filepath") cast_run_state = gr.State() cast_timer = gr.Timer(1, active=False) cast_reference.click( lambda: ( PRODUCTION_CLEAN, DEFAULT_VISUAL_STYLE, [], TEXT_TO_IMAGE, False, 1536, 640, 30, 8, 3461201489, "daggerfall-000001.safetensors", 1.0, "pixel-art-xl-v1.1.safetensors", 1.0, "None", 1.0, "None", 1.0, "None", 1.0, "The supplied turnaround's size, steps, guidance, seed, and " "LoRA names are selected. " "Click Build Population Prompts, then Generate Population.", ), outputs=[cast_template, cast_visual_style, style, mode, use_canny, cast_width, cast_height, steps, guidance, seed, *cast_lora_controls, cast_status], queue=False, ) cast_template.input( select_population_template, inputs=cast_template, outputs=[cast_visual_style, cast_men, cast_women, cast_min_age, cast_max_age, cast_total, cast_status, cast_matrix, cast_choice, cast_selected_prompt, cast_batch, cast_prompts], queue=False, ) cast_build.click( build_character_panel, inputs=[cast_total, cast_men, cast_women, cast_min_age, cast_max_age, cast_seed, cast_visual_style, cast_template], outputs=[cast_status, cast_matrix, cast_choice, cast_selected_prompt, cast_batch, cast_prompts], queue=False, ) cast_choice.change( selected_character_prompt, inputs=[cast_choice, cast_prompts], outputs=cast_selected_prompt, queue=False, ) cast_use.click( copy_character_to_generate, inputs=[cast_selected_prompt, cast_width, cast_height, *cast_lora_controls], outputs=[prompt, width, height, *lora_controls], queue=False, ) cast_generate.click( start_population_generation, inputs=[cast_batch, negative_prompt, style, model, cast_width, cast_height, steps, guidance, seed, mode, input_image, resize_mode, strength, use_canny, canny_image, low_threshold, high_threshold, canny_scale, canny_start, canny_end, *cast_lora_controls], outputs=[cast_run_state, cast_gallery, cast_image, cast_details, cast_run_status, cast_timer, model_status, canny_preview, cast_generate], concurrency_limit=1, ) cast_tick = cast_timer.tick( generate_next_character, inputs=[cast_run_state, cast_gallery], outputs=[cast_run_state, cast_gallery, cast_image, cast_details, cast_run_status, cast_timer, cast_generate], trigger_mode="once", concurrency_limit=1, concurrency_id="sdxl_gpu", ) cast_tick.failure( stop_population_after_failure, inputs=cast_run_state, outputs=[cast_run_state, cast_run_status, cast_timer, cast_generate], queue=False, ) button.click( fn=fetch_selected_model, inputs=[model, mode, input_image, use_canny, canny_image, width, height, resize_mode, strength, low_threshold, high_threshold, canny_scale, canny_start, canny_end, *lora_controls], outputs=[model_status, canny_preview], concurrency_limit=1, ).success( fn=generate, inputs=[prompt, negative_prompt, style, model, width, height, steps, guidance, seed, mode, input_image, resize_mode, strength, use_canny, canny_image, low_threshold, high_threshold, canny_scale, canny_start, canny_end, *lora_controls], outputs=[image, metadata], concurrency_limit=1, concurrency_id="sdxl_gpu", ) if __name__ == "__main__": demo.queue().launch( server_name="0.0.0.0", server_port=int(os.getenv("PORT", "7860")), prevent_thread_lock=True, ) load_models() threading.Event().wait()