import gradio as gr import numpy as np import random import os from huggingface_hub import login, hf_hub_download import torch import spaces from diffusers import FluxPipeline import logging # Configure logging logging.basicConfig( level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s', handlers=[ logging.FileHandler("app.log"), logging.StreamHandler() ] ) logger = logging.getLogger("moroccan-ghibli-flux-compare") # Constants MAX_SEED = np.iinfo(np.int32).max MAX_IMAGE_SIZE = 1024 DEFAULT_SEED = 42 device = torch.device("cuda" if torch.cuda.is_available() else "cpu") # Global pipelines pipeline_base = None pipeline_lora = None def load_base_pipeline(): api_key = os.getenv("HF_TOKEN") if not api_key: raise ValueError("HF_TOKEN environment variable not set.") login(token=api_key) logger.info("Loading base Flux model (no LoRA)") pipe = FluxPipeline.from_pretrained( "black-forest-labs/FLUX.1-dev", torch_dtype=torch.bfloat16, ) return pipe def load_lora_pipeline(revision: str): api_key = os.getenv("HF_TOKEN") if not api_key: raise ValueError("HF_TOKEN environment variable not set.") login(token=api_key) logger.info(f"Loading Flux model with LoRA revision: {revision}") pipe = FluxPipeline.from_pretrained( "black-forest-labs/FLUX.1-dev", torch_dtype=torch.bfloat16, ) # Extract step number from revision name (e.g., "step_1500" -> "1500") step_num = revision.split('_')[-1] if len(step_num) < 4: step_num = '0' + step_num filename = f"moroccan_ghibli_flux_lora_00000{step_num}.safetensors" logger.info(f"Downloading LoRA weights: {filename}") lora_path = hf_hub_download( repo_id="atlasia/moroccan-ghibli-flux-lora", filename=filename, revision=revision, token=api_key ) logger.info("Loading LoRA weights into the pipeline") pipe.load_lora_weights(lora_path) return pipe def init_pipelines(style_intensity: str = "low"): global pipeline_base, pipeline_lora revision_map = { "low": "step_750", "medium": "step_1500", "high": "step_2500" } revision = revision_map[style_intensity] if pipeline_base is None: pipeline_base = load_base_pipeline() logger.info("Base pipeline loaded") logger.info(f"Initializing LoRA pipeline with style: {style_intensity} ({revision})") pipeline_lora = load_lora_pipeline(revision) logger.info("LoRA pipeline loaded") # Initial pipelines init_pipelines("low") def update_lora_pipeline(style_intensity): logger.info(f"Updating LoRA pipeline to style: {style_intensity}") init_pipelines(style_intensity) def _run_pipeline(pipe: FluxPipeline, prompt: str, seed: int, width: int, height: int, guidance_scale: float): pipe.to(device) generator = torch.Generator(device=device).manual_seed(seed) max_sequence_length = 512 output = pipe( prompt=[prompt], guidance_scale=guidance_scale, num_inference_steps=50, height=height, width=width, max_sequence_length=max_sequence_length, generator=generator, ) return output.images[0] @spaces.GPU(duration=120) def infer(prompt, seed, width, height, guidance_scale, progress=gr.Progress()): logger.info(f"Generating comparison for prompt: '{prompt[:50]}...'") logger.info(f"Parameters: seed={seed}, width={width}, height={height}, guidance={guidance_scale}") progress(0.1, desc="Preparing base model") base_image = _run_pipeline(pipeline_base, prompt, seed, width, height, guidance_scale) progress(0.55, desc="Preparing LoRA model") lora_image = _run_pipeline(pipeline_lora, prompt, seed, width, height, guidance_scale) progress(0.95, desc="Processing results") logger.info("Comparison generation completed successfully") # Return both images and the used seed return base_image, lora_image, seed def randomize_seed(): return random.randint(0, MAX_SEED) with gr.Blocks() as demo: with gr.Row(): with gr.Column(scale=1): gr.Markdown("# Flux: Base vs Moroccan-Ghibli LoRA") gr.Markdown("Generate side-by-side images with identical settings to compare the base model and the Moroccan Ghibli LoRA.") style_intensity = gr.Dropdown( label="LoRA Style Intensity", choices=["low", "medium", "high"], value="low", interactive=True ) prompt = gr.Textbox( label="Prompt", placeholder="Enter your prompt (e.g., 'Moroccan Ghibli studio style portrait')", ) with gr.Row(): seed = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=DEFAULT_SEED, interactive=True) randomize_button = gr.Button("Randomize Seed") with gr.Row(): width = gr.Slider(label="Width", minimum=256, maximum=MAX_IMAGE_SIZE, step=32, value=1024) height = gr.Slider(label="Height", minimum=256, maximum=MAX_IMAGE_SIZE, step=32, value=1024) guidance_scale = gr.Slider(label="Guidance scale", minimum=0.0, maximum=10.0, step=0.1, value=3.5) run_button = gr.Button("Generate", variant="primary") with gr.Column(scale=2): gr.Markdown("## Generated Images") with gr.Row(): base_image = gr.Image(label="Base Model", height=512) lora_image = gr.Image(label="Moroccan Ghibli LoRA", height=512) output_seed = gr.Number(label="Used Seed", precision=0) gr.Markdown("## Example Prompts") examples = [ ["Moroccan Ghibli studio style portrait of a character in a riad courtyard", DEFAULT_SEED, 1024, 1024, 3.5], ["Moroccan Ghibli studio style image of a bustling souk with flying carpets", DEFAULT_SEED, 1024, 1024, 3.5], ["Moroccan Ghibli studio style landscape of a desert oasis under starry skies", DEFAULT_SEED, 1024, 1024, 3.5], ["Moroccan Ghibli studio style depiction of a magical lantern festival", DEFAULT_SEED, 1024, 1024, 3.5], ["Moroccan Ghibli studio style portrait of a medina at sunset", DEFAULT_SEED, 1024, 1024, 3.5] ] gr.Examples( examples=examples, inputs=[prompt, seed, width, height, guidance_scale], outputs=[base_image, lora_image, output_seed], fn=infer, cache_examples=True ) style_intensity.change( fn=update_lora_pipeline, inputs=[style_intensity], outputs=[] ) gr.on( triggers=[run_button.click, prompt.submit], fn=infer, inputs=[prompt, seed, width, height, guidance_scale], outputs=[base_image, lora_image, output_seed] ) randomize_button.click( fn=randomize_seed, inputs=[], outputs=[seed] ) if __name__ == "__main__": logger.info("Starting application") demo.launch() logger.info("Application closed")