import gradio as gr import torch from diffusers import StableDiffusionPipeline, StableDiffusionXLImg2ImgPipeline, AutoPipelineForText2Image from diffusers.utils import load_image from PIL import Image import time import random import os import gc # Garbage collector import logging # --- Configuration --- # Setup basic logging logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s') logger = logging.getLogger(__name__) # Ensure CPU is used DEVICE = "cpu" TORCH_DTYPE = torch.float32 # float16/bfloat16 not practical on CPU # Model definitions # We need to know the base model for LoRAs and compatible IP-Adapters MODEL_CONFIG = { "BlaireSilver13/youtube-thumbnail": { "repo_id": "BlaireSilver13/youtube-thumbnail", "is_lora": True, "lora_filename": "FLUX-youtube-thumbnails.safetensors", "base_model": "black-forest-labs/FLUX.1-dev", "pipeline_class": AutoPipelineForText2Image, "ip_adapter_repo": "h94/IP-Adapter", "ip_adapter_weights": "ip-adapter_sd15.bin", "ip_adapter_image_encoder": "laion/CLIP-ViT-H-14-laion2B-s32B-b79K" }, "itzzdeep/youtube-thumbnails-sdxl-lora": { "repo_id": "itzzdeep/youtube-thumbnails-sdxl-lora", "is_lora": True, "lora_filename": "youtube-thumbnails-sdxl-lora.safetensors", "base_model": "stabilityai/stable-diffusion-xl-base-1.0", "pipeline_class": AutoPipelineForText2Image, # Handles SDXL loading better "ip_adapter_repo": "h94/IP-Adapter", # SDXL IP-Adapter repo "ip_adapter_weights": "ip-adapter-plus_sdxl_vit-h.bin", # SDXL weights "ip_adapter_image_encoder": "laion/CLIP-ViT-H-14-laion2B-s32B-b79K" # Usually the same encoder repo }, "justmalhar/flux-thumbnails-v3": { "repo_id": "justmalhar/flux-thumbnails-v3", "is_lora": False, # Assuming this is a full SD 1.5 fine-tune based on common practice "base_model": None, "pipeline_class": StableDiffusionPipeline, "ip_adapter_repo": "h94/IP-Adapter", "ip_adapter_weights": "ip-adapter_sd15.bin", "ip_adapter_image_encoder": "laion/CLIP-ViT-H-14-laion2B-s32B-b79K" }, "saq1b/mrbeast-thumbnail-style": { "repo_id": "saq1b/mrbeast-thumbnail-style", "is_lora": True, # This is typically a LoRA "lora_filename": None, # Auto-detect or specify e.g., "pytorch_lora_weights.safetensors" "base_model": "runwayml/stable-diffusion-v1-5", # Common base for SD 1.5 LoRAs "pipeline_class": StableDiffusionPipeline, "ip_adapter_repo": "h94/IP-Adapter", "ip_adapter_weights": "ip-adapter_sd15.bin", "ip_adapter_image_encoder": "laion/CLIP-ViT-H-14-laion2B-s32B-b79K" } } AVAILABLE_MODELS = list(MODEL_CONFIG.keys()) # Global variable to potentially hold the pipeline to avoid reloading *if memory allows* # NOTE: On restricted CPU environments, it's SAFER to load inside the function. # Set to None initially. Let's load dynamically inside the function for safety. # current_pipeline = None # current_model_key = None # --- Helper Functions --- def cleanup_memory(): """Attempts to free GPU memory (less relevant for CPU but good practice).""" logger.info("Attempting to clean up memory...") try: # If a pipeline exists globally (if we change strategy), unload it # global current_pipeline, current_model_key # if current_pipeline is not None: # logger.info(f"Unloading model {current_model_key} from memory.") # del current_pipeline # current_pipeline = None # current_model_key = None gc.collect() if torch.cuda.is_available(): # Only run cuda cache empty if cuda is present torch.cuda.empty_cache() logger.info("Memory cleanup potentially done.") except Exception as e: logger.error(f"Error during memory cleanup: {e}") # --- Main Generation Function --- def generate_thumbnail( model_key: str, prompt: str, negative_prompt: str, reference_image_pil: Image.Image | None, # Gradio provides PIL image num_inference_steps: int, guidance_scale: float, seed: int, ip_adapter_scale: float, progress=gr.Progress() ): """Generates an image using the selected model, IP-Adapter, and settings.""" start_time = time.time() debug_log = f"--- Generation Log ({time.strftime('%Y-%m-%d %H:%M:%S')}) ---\n" debug_log += f"Selected Model Key: {model_key}\n" debug_log += f"Prompt: {prompt}\n" debug_log += f"Negative Prompt: {negative_prompt}\n" debug_log += f"Steps: {num_inference_steps}, CFG Scale: {guidance_scale}\n" debug_log += f"Seed: {seed}\n" debug_log += f"Reference Image Provided: {'Yes' if reference_image_pil else 'No'}\n" debug_log += f"IP Adapter Scale: {ip_adapter_scale}\n" debug_log += f"Device: {DEVICE}, Dtype: {TORCH_DTYPE}\n\n" pipeline = None # Ensure pipeline is defined in this scope try: if not model_key: raise ValueError("No model selected.") config = MODEL_CONFIG[model_key] repo_id = config["repo_id"] is_lora = config["is_lora"] base_model = config["base_model"] pipeline_class = config["pipeline_class"] ip_adapter_repo = config["ip_adapter_repo"] ip_adapter_weights = config["ip_adapter_weights"] # ip_adapter_image_encoder = config["ip_adapter_image_encoder"] # Encoder loaded via IP-Adapter itself usually # --- Model Loading --- load_start_time = time.time() debug_log += f"[{time.time() - start_time:.2f}s] Cleaning up memory before loading...\n" progress(0.1, desc="Cleaning up memory...") cleanup_memory() # Attempt cleanup before loading new model debug_log += f"[{time.time() - start_time:.2f}s] Loading model: {'LoRA ' + repo_id if is_lora else repo_id}...\n" progress(0.2, desc=f"Loading {'LoRA ' + repo_id if is_lora else repo_id}...") model_load_id = base_model if is_lora else repo_id debug_log += f"[{time.time() - start_time:.2f}s] Base/Model ID for pipeline: {model_load_id}\n" pipeline = pipeline_class.from_pretrained( model_load_id, torch_dtype=TORCH_DTYPE, # Add any specific args needed for the pipeline class if necessary # safety_checker=None, # Disable safety checker if needed/causes issues on CPU # requires_safety_checker=False, ) pipeline.to(DEVICE) debug_log += f"[{time.time() - start_time:.2f}s] Base pipeline loaded onto {DEVICE}.\n" if is_lora: lora_load_start = time.time() debug_log += f"[{time.time() - start_time:.2f}s] Loading LoRA weights from {repo_id}...\n" progress(0.4, desc=f"Loading LoRA {repo_id}...") try: lora_filename = config.get("lora_filename") # Get specific filename if provided if lora_filename: debug_log += f"[{time.time() - start_time:.2f}s] Using specified LoRA filename: {lora_filename}\n" pipeline.load_lora_weights(repo_id, weight_name=lora_filename, torch_dtype=TORCH_DTYPE) else: # Let diffusers try to auto-detect standard names like .safetensors or .bin debug_log += f"[{time.time() - start_time:.2f}s] Attempting auto-detection of LoRA filename.\n" pipeline.load_lora_weights(repo_id, torch_dtype=TORCH_DTYPE) # When using LoRA with diffusers >= 0.22, explicitly fuse *or* set adapters # pipeline.fuse_lora() # Fuse creates a new pipeline state (might use more memory) pipeline.set_adapters(pipeline.get_active_adapters(), adapter_weights=1.0) # Recommended for flexibility debug_log += f"[{time.time() - start_time:.2f}s] LoRA weights loaded and adapters set in {time.time() - lora_load_start:.2f}s.\n" except Exception as e: debug_log += f"[{time.time() - start_time:.2f}s] ERROR loading LoRA: {e}. Check LoRA repo structure/filename.\n" # Decide whether to continue without LoRA or raise error raise ValueError(f"Failed to load LoRA weights for {repo_id}: {e}") # --- IP Adapter Loading --- if reference_image_pil and ip_adapter_scale > 0: ip_load_start = time.time() debug_log += f"[{time.time() - start_time:.2f}s] Loading IP-Adapter: {ip_adapter_repo} ({ip_adapter_weights})...\n" progress(0.6, desc="Loading IP-Adapter...") try: # Ensure the pipeline has the load_ip_adapter method if not hasattr(pipeline, "load_ip_adapter"): raise AttributeError("The current pipeline class does not support load_ip_adapter. Check diffusers version or pipeline type.") pipeline.load_ip_adapter( ip_adapter_repo, subfolder="models", # Common subfolder, adjust if needed weight_name=ip_adapter_weights, # image_encoder_folder=ip_adapter_image_encoder # Let diffusers handle encoder loading usually ) pipeline.set_ip_adapter_scale(ip_adapter_scale) debug_log += f"[{time.time() - start_time:.2f}s] IP-Adapter loaded and scale set ({ip_adapter_scale}) in {time.time() - ip_load_start:.2f}s.\n" # Prepare the image for IP-Adapter (often just needs to be a PIL image) ip_image = reference_image_pil.convert("RGB") debug_log += f"[{time.time() - start_time:.2f}s] Reference image prepared for IP-Adapter.\n" except Exception as e: debug_log += f"[{time.time() - start_time:.2f}s] WARNING: Failed to load IP-Adapter: {e}. Proceeding without image guidance.\n" ip_image = None ip_adapter_scale = 0 # Effectively disable it if loading failed pipeline.set_ip_adapter_scale(0) # Ensure scale is 0 else: ip_image = None if hasattr(pipeline, "set_ip_adapter_scale"): pipeline.set_ip_adapter_scale(0) # Ensure scale is 0 if no image/scale=0 debug_log += f"[{time.time() - start_time:.2f}s] No reference image provided or IP Adapter scale is 0. Skipping IP-Adapter loading.\n" debug_log += f"[{time.time() - start_time:.2f}s] Total Model & IP-Adapter Loading time: {time.time() - load_start_time:.2f}s\n" # --- Generation --- gen_start_time = time.time() debug_log += f"[{time.time() - start_time:.2f}s] Starting generation...\n" progress(0.7, desc="Generating image...") # Handle seed if seed == -1: seed = random.randint(0, 2**32 - 1) debug_log += f"[{time.time() - start_time:.2f}s] Using random seed: {seed}\n" generator = torch.Generator(device=DEVICE).manual_seed(seed) # Prepare arguments for pipeline call pipeline_args = { "prompt": prompt, "negative_prompt": negative_prompt, "num_inference_steps": num_inference_steps, "guidance_scale": guidance_scale, "generator": generator, } # Add IP-Adapter image if it's loaded and ready if ip_image is not None and hasattr(pipeline, "set_ip_adapter_scale") and ip_adapter_scale > 0: pipeline_args["ip_adapter_image"] = ip_image # Scale was set earlier with set_ip_adapter_scale debug_log += f"[{time.time() - start_time:.2f}s] Passing reference image to pipeline with IP scale {ip_adapter_scale}.\n" else: debug_log += f"[{time.time() - start_time:.2f}s] Not passing reference image to pipeline.\n" # Run inference with torch.inference_mode(): # More modern than no_grad for inference output_image = pipeline(**pipeline_args).images[0] gen_end_time = time.time() debug_log += f"[{time.time() - start_time:.2f}s] Generation finished in {gen_end_time - gen_start_time:.2f}s.\n" # --- Cleanup --- debug_log += f"[{time.time() - start_time:.2f}s] Unloading model from memory (CPU strategy)...\n" progress(0.95, desc="Cleaning up...") del pipeline # Explicitly delete pipeline cleanup_memory() # Call garbage collection total_time = time.time() - start_time debug_log += f"\n--- Total time: {total_time:.2f} seconds ---\n" return output_image, debug_log except Exception as e: logger.exception(f"Error during generation for model {model_key}") # Log full traceback error_time = time.time() - start_time debug_log += f"\n\n!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!\n" debug_log += f"ERROR occurred after {error_time:.2f}s:\n{e}\n" debug_log += f"!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!\n" # Try cleanup even on error if 'pipeline' in locals() and pipeline is not None: del pipeline cleanup_memory() # Return None for image, and the log containing the error return None, debug_log # --- Gradio Interface --- css = """ #warning { background-color: #FFCCCB; /* Light red */ padding: 10px; border-radius: 5px; text-align: center; font-weight: bold; } #debug_log_area textarea { font-family: monospace; font-size: 10px; /* Smaller font for logs */ white-space: pre-wrap; /* Wrap long lines */ word-wrap: break-word; /* Break words if necessary */ } """ with gr.Blocks(css=css) as demo: gr.Markdown("# YouTube Thumbnail Generator with IP-Adapter") gr.Markdown( "Select a thumbnail model, provide a text prompt, and optionally upload a reference image " "to guide the generation using IP-Adapter." ) gr.HTML("
⚠️ Warning: Inference on CPU is VERY SLOW (minutes per image, especially SDXL models). Please be patient.
") with gr.Row(): with gr.Column(scale=1): model_dropdown = gr.Dropdown( label="Select Thumbnail Model", choices=AVAILABLE_MODELS, value=AVAILABLE_MODELS[0] if AVAILABLE_MODELS else None, ) prompt_input = gr.Textbox(label="Prompt", lines=3, placeholder="e.g., Epic landscape, dramatic lighting, YouTube thumbnail style") negative_prompt_input = gr.Textbox(label="Negative Prompt", lines=2, placeholder="e.g., blurry, low quality, text, signature, watermark") reference_image_input = gr.Image(label="Reference Image (for IP-Adapter)", type="pil", sources=["upload"]) with gr.Accordion("Advanced Settings", open=False): steps_slider = gr.Slider(label="Inference Steps", minimum=10, maximum=100, value=30, step=1) cfg_slider = gr.Slider(label="Guidance Scale (CFG)", minimum=1.0, maximum=20.0, value=7.0, step=0.5) ip_adapter_scale_slider = gr.Slider(label="IP-Adapter Scale", minimum=0.0, maximum=1.5, value=0.6, step=0.05, info="Strength of the reference image influence (0 = disabled).") seed_input = gr.Number(label="Seed", value=-1, precision=0, info="-1 for random seed") generate_button = gr.Button("Generate Thumbnail", variant="primary") with gr.Column(scale=1): output_image = gr.Image(label="Generated Thumbnail", type="pil") debug_output = gr.Textbox(label="Debug Log", lines=20, interactive=False, elem_id="debug_log_area") generate_button.click( fn=generate_thumbnail, inputs=[ model_dropdown, prompt_input, negative_prompt_input, reference_image_input, steps_slider, cfg_slider, seed_input, ip_adapter_scale_slider ], outputs=[output_image, debug_output] ) # --- Launch --- if __name__ == "__main__": logger.info("Starting Gradio App...") # Queueing is important for handling multiple users on Spaces, even if slow demo.queue().launch(debug=True) # debug=True provides Gradio debug info in console