import os import gradio as gr import numpy as np import random import spaces import torch import json import logging from diffusers import DiffusionPipeline, AutoencoderTiny, AutoencoderKL, AutoPipelineForImage2Image from huggingface_hub import login from diffusers.utils import load_image import time from datetime import datetime from io import BytesIO import torch.nn.functional as F from PIL import Image, ImageFilter import time import boto3 from io import BytesIO import re import json import random import string from diffusers import FluxPipeline from huggingface_hub import hf_hub_download from diffusers.quantizers import PipelineQuantizationConfig from diffusers import (FluxPriorReduxPipeline, FluxInpaintPipeline, FluxFillPipeline, FluxPipeline) # Login Hugging Face Hub HF_TOKEN = os.environ.get("HF_TOKEN") login(token=HF_TOKEN) import diffusers # init dtype = torch.bfloat16 device = "cuda:0" base_model = "black-forest-labs/FLUX.1-dev" txt2img_pipe = FluxPipeline.from_pretrained(base_model, torch_dtype=dtype) txt2img_pipe = txt2img_pipe.to(device) MAX_SEED = 2**32 - 1 class calculateDuration: def __init__(self, activity_name=""): self.activity_name = activity_name def __enter__(self): self.start_time = time.time() self.start_time_formatted = time.strftime("%Y-%m-%d %H:%M:%S", time.localtime(self.start_time)) print(f"Activity: {self.activity_name}, Start time: {self.start_time_formatted}") return self def __exit__(self, exc_type, exc_value, traceback): self.end_time = time.time() self.elapsed_time = self.end_time - self.start_time self.end_time_formatted = time.strftime("%Y-%m-%d %H:%M:%S", time.localtime(self.end_time)) if self.activity_name: print(f"Elapsed time for {self.activity_name}: {self.elapsed_time:.6f} seconds") else: print(f"Elapsed time: {self.elapsed_time:.6f} seconds") def safe_trim_for_clip(text: str, max_words: int = 77) -> str: tokens = re.split(r"\s+", text.strip()) if len(tokens) <= max_words: return text return " ".join(tokens[:max_words]) def upload_image_to_r2(image, account_id, access_key, secret_key, bucket_name): with calculateDuration("Upload images"): connectionUrl = f"https://{account_id}.r2.cloudflarestorage.com" s3 = boto3.client( 's3', endpoint_url=connectionUrl, region_name='auto', aws_access_key_id=access_key, aws_secret_access_key=secret_key ) current_time = datetime.now().strftime("%Y/%m/%d/%H/%M/%S") image_file = f"generated_images/{current_time}/{random.randint(0, MAX_SEED)}.png" buffer = BytesIO() image.save(buffer, "PNG") buffer.seek(0) s3.upload_fileobj(buffer, bucket_name, image_file) print("upload finish", image_file) # Generate thumbnail thumbnail = image.copy() thumbnail_width = 256 aspect_ratio = image.height / image.width thumbnail_height = int(thumbnail_width * aspect_ratio) thumbnail = thumbnail.resize((thumbnail_width, thumbnail_height), Image.LANCZOS) thumbnail_file = image_file.replace(".png", "_thumbnail.png") thumbnail_buffer = BytesIO() thumbnail.save(thumbnail_buffer, "PNG") thumbnail_buffer.seek(0) s3.upload_fileobj(thumbnail_buffer, bucket_name, thumbnail_file) print("upload thumbnail finish", thumbnail_file) return image_file def generate_random_4_digit_string(): return ''.join(random.choices(string.digits, k=4)) @spaces.GPU(duration=120) def run_lora( prompt, image_url, lora_strings_json, image_strength, cfg_scale, steps, randomize_seed, seed, width, height, upload_to_r2, account_id, access_key, secret_key, bucket, progress=gr.Progress(track_tqdm=True) ): print("run_lora", prompt, lora_strings_json, cfg_scale, steps, width, height) gr.Info("Starting process") pipe = txt2img_pipe device = pipe.device print(device) # ========== Seed ========== if randomize_seed: with calculateDuration("Set random seed"): seed = random.randint(0, MAX_SEED) generator = torch.Generator(device=device).manual_seed(seed) # ========== LoRA ========== gr.Info("Start to load LoRA ...") with calculateDuration("Unloading LoRA"): try: pipe.unload_lora_weights() except Exception as _: pass adapter_names = [] adapter_weights = [] if lora_strings_json: try: lora_configs = json.loads(lora_strings_json) except Exception as _: lora_configs = None gr.Warning("Parse lora config json failed") print("parse lora config json failed") if lora_configs: with calculateDuration("Loading LoRA weights"): for lora_info in lora_configs: repo = lora_info.get("repo") weights = lora_info.get("weights") adapter_name = lora_info.get("adapter_name") or f"adp_{generate_random_4_digit_string()}" weight = float(lora_info.get("adapter_weight", 1.0)) if not (repo and weights): print(f"skip invalid lora entry: {lora_info}") continue try: weight_path = hf_hub_download(repo_id=repo, filename=weights) pipe.load_lora_weights(weight_path, adapter_name=adapter_name, prefix=None) adapter_names.append(adapter_name) adapter_weights.append(weight) except Exception as e: print(f"load lora error for {repo}/{weights}: {e}") if adapter_names: pipe.set_adapters(adapter_names, adapter_weights=adapter_weights) try: active = pipe.get_active_adapters() if hasattr(pipe, "get_active_adapters") else [] print("Active adapters:", active) except Exception as e: print("Active adapters query failed:", e) lora_layer_count = 0 for name, module in pipe.transformer.named_modules(): attrs = dir(module) if any(a.startswith("lora_") for a in attrs) or "lora" in module.__class__.__name__.lower(): lora_layer_count += 1 print(f"[DEBUG] transformer LoRA layers: {lora_layer_count}") if lora_layer_count == 0: gr.Warning("LoRA seems not injected (0 layers on transformer). Check whether the LoRA is trained for FLUX and `prefix=None` is set.") pipe.enable_vae_slicing() clip_side_prompt = safe_trim_for_clip(prompt, max_words=77) init_image = None error_message = "" image = None try: gr.Info("Start to generate images ...") joint_attention_kwargs = {"scale": 1} image = pipe( prompt=prompt, num_inference_steps=int(steps), guidance_scale=float(cfg_scale), width=int(width), height=int(height), max_sequence_length=512, generator=generator, joint_attention_kwargs=joint_attention_kwargs ).images[0] except Exception as e: error_message = str(e) gr.Error(error_message) print("fatal error", e) image = None result = {"status": "failed", "message": error_message} if image is None else {"status": "success", "message": "Image generated but not uploaded", "seed": seed} if image is not None and upload_to_r2: try: url = upload_image_to_r2(image, account_id, access_key, secret_key, bucket) result = {"status": "success", "message": "upload image success", "url": url, "seed": seed} except Exception as e: err = f"Upload failed: {e}" gr.Warning(err) print(err) result = {"status": "success", "message": "generated but upload failed", "seed": seed} gr.Info("Completed！") progress(100, "Completed!") # CHANGED: Return both image AND json return image, json.dumps(result) # Gradio interface with gr.Blocks() as demo: gr.Markdown("# flux-dev-multi-lora") with gr.Row(): with gr.Column(): prompt = gr.Text(label="Prompt", placeholder="Enter prompt", lines=10) lora_strings_json = gr.Text( label="LoRA Configs (JSON List String)", placeholder='[{"repo": "lora_repo1", "weights": "weights1.safetensors", "adapter_name": "adapter1", "adapter_weight": 1.0}]', lines=5 ) image_url = gr.Text(label="Image url", placeholder="Enter image url to enable image to image model", lines=1) run_button = gr.Button("Run", scale=0) with gr.Accordion("Advanced Settings", open=False): with gr.Row(): seed = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=0, randomize=True) randomize_seed = gr.Checkbox(label="Randomize seed", value=True) with gr.Row(): width = gr.Slider(label="Width", minimum=256, maximum=1536, step=64, value=1024) height = gr.Slider(label="Height", minimum=256, maximum=1536, step=64, value=1024) with gr.Row(): image_strength = gr.Slider(label="Denoise Strength", info="Lower means more image influence", minimum=0.1, maximum=1.0, step=0.01, value=0.75) cfg_scale = gr.Slider(label="CFG Scale", minimum=1, maximum=20, step=0.5, value=3.5) steps = gr.Slider(label="Steps", minimum=1, maximum=50, step=1, value=28) upload_to_r2 = gr.Checkbox(label="Upload to R2", value=False) account_id = gr.Textbox(label="Account Id", placeholder="Enter R2 account id") access_key = gr.Textbox(label="Access Key", placeholder="Enter R2 access key here") secret_key = gr.Textbox(label="Secret Key", placeholder="Enter R2 secret key here") bucket = gr.Textbox(label="Bucket Name", placeholder="Enter R2 bucket name here") with gr.Column(): # CHANGED: Add image output output_image = gr.Image(label="Generated Image", type="pil") json_text = gr.Text(label="Result JSON") gr.Markdown("**Disclaimer:**") gr.Markdown( "This demo is only for research purpose. This space owner cannot be held responsible for the generation of NSFW (Not Safe For Work) content through the use of this demo. Users are solely responsible for any content they create, and it is their obligation to ensure that it adheres to appropriate and ethical standards. This space owner provides the tools, but the responsibility for their use lies with the individual user." ) inputs = [ prompt, image_url, lora_strings_json, image_strength, cfg_scale, steps, randomize_seed, seed, width, height, upload_to_r2, account_id, access_key, secret_key, bucket ] # CHANGED: Two outputs now outputs = [output_image, json_text] run_button.click( fn=run_lora, inputs=inputs, outputs=outputs ) try: demo.queue().launch() except: print("demo exception ...")