from huggingface_hub import model_info # /// script # dependencies = [ # "diffusers==0.33.0", # "transformers==4.44.0", # "accelerate", # "bitsandbytes", # "boto3", # "python-dotenv", # "torch==2.4.0", # "torchvision==0.19.0", # "sentencepiece", # "protobuf", # "peft" # ] # /// import json import argparse import sys import os import time import logging from pathlib import Path from datetime import datetime # Disable TQDM progress bars os.environ["TQDM_DISABLE"] = "1" os.environ["HF_HUB_DISABLE_PROGRESS_BARS"] = "1" try: import torch from diffusers import FluxPipeline from transformers import BitsAndBytesConfig from huggingface_hub import model_info except ImportError: pass logging.basicConfig(level=logging.INFO, format='[%(asctime)s] %(message)s') logger = logging.getLogger("hf_a100_worker") def upload_to_s3(file_path, s3_bucket, s3_key): try: import boto3 import mimetypes s3_client = boto3.client('s3', aws_access_key_id=os.environ.get("AWS_ACCESS_KEY"), aws_secret_access_key=os.environ.get("AWS_SECRET_ACCESS_KEY"), region_name=os.environ.get("AWS_REGION", "us-east-2")) content_type, _ = mimetypes.guess_type(str(file_path)) s3_client.upload_file(str(file_path), s3_bucket, s3_key, ExtraArgs={'ContentType': content_type or 'image/webp'}) logger.info(f"S3 Upload successful: {s3_key}") return True except Exception as e: logger.error(f"S3 upload failed: {e}") return False def check_s3_file(s3_bucket, s3_key): try: import boto3 s3 = boto3.client('s3', aws_access_key_id=os.environ.get("AWS_ACCESS_KEY"), aws_secret_access_key=os.environ.get("AWS_SECRET_ACCESS_KEY"), region_name=os.environ.get("AWS_REGION", "us-east-2")) s3.head_object(Bucket=s3_bucket, Key=s3_key) return True except: return False def download_file_from_s3(s3_key, local_path): try: import boto3 s3 = boto3.client('s3', aws_access_key_id=os.environ.get("AWS_ACCESS_KEY"), aws_secret_access_key=os.environ.get("AWS_SECRET_ACCESS_KEY"), region_name=os.environ.get("AWS_REGION", "us-east-2")) s3_bucket = os.environ.get("AWS_BUCKET", "shadow-vault") s3.download_file(s3_bucket, s3_key, local_path) return True except Exception as e: logger.error(f"S3 download failed for {s3_key}: {e}") return False def main(): parser = argparse.ArgumentParser() parser.add_argument("--dry-run", action="store_true", help="Audit connections and exit") args = parser.parse_args() job_json = os.environ.get("JOB_DATA") if not job_json: s3_config_key = os.environ.get("CONFIG_S3_KEY") if s3_config_key: logger.info(f"Downloading configuration from S3: {s3_config_key}...") if download_file_from_s3(s3_config_key, "job_config.json"): with open("job_config.json", "r") as f: job_json = f.read() if not job_json and os.path.exists("job_config.json"): logger.info("Reading configuration from job_config.json...") with open("job_config.json", "r") as f: job_json = f.read() if not job_json: logger.error("Error: No configuration source found (JOB_DATA, CONFIG_S3_KEY, or job_config.json).") sys.exit(1) job_data = json.loads(job_json) job_name = job_data.get("job_name", "hf_masterpiece_a100") tasks = job_data.get("tasks", []) model_id = job_data.get("base_model", "black-forest-labs/FLUX.1-dev") hf_token = os.environ.get("HF_TOKEN") s3_bucket = os.environ.get("AWS_BUCKET", "shadow-vault") logger.info(f"--- ATOMIC RELIABILITY BOOT: {job_name} ---") # 1. Verify S3 Connection if not check_s3_file(s3_bucket, "jobs/cama/base_image.png"): logger.warning("Warning: S3 test file check failed. Verify AWS credentials.") else: logger.info("S3 connectivity verified.") # 2. Verify Model/LoRA Accessibility try: model_info(model_id, token=hf_token) for task in tasks: for lora in task.get("loras", []): model_info(lora["path"], token=hf_token) logger.info("All Model/LoRA assets verified.") except Exception as e: logger.error(f"Asset access failed: {e}") if not args.dry_run: sys.exit(1) if args.dry_run: logger.info("--- Dry Run Successful. All logic verified. ---") return # 3. Production Run logic if not torch.cuda.is_available(): logger.error("CRITICAL: CUDA is not available. Aborting.") sys.exit(1) output_dir = Path("output") output_dir.mkdir(exist_ok=True) logger.info("Loading Pipeline in FULL BFLOAT16 (No Quantization for perfect pixels)...") pipe = FluxPipeline.from_pretrained( model_id, torch_dtype=torch.bfloat16, token=hf_token, device_map="balanced" # Changed from "auto" ) # Enable tiling for memory efficiency during VAE pipe.vae.enable_tiling() # Load LoRAs active_loras = {} for task in tasks: for lora in task.get("loras", []): path = lora["path"] if path not in active_loras: name = f"adapter_{len(active_loras)}" logger.info(f"Loading LoRA: {path}") pipe.load_lora_weights(path, weight_name="rochelle_flux_v1.safetensors" if "rochelle-flux-lora" in path and "-v2" not in path else None, adapter_name=name, token=hf_token) active_loras[path] = name # Process Tasks for idx, task in enumerate(tasks): filename = f"a100_t{idx+1}_{task.get('seed', 42)}.webp" s3_key = f"jobs/{job_name}/{filename}" if check_s3_file(s3_bucket, s3_key): logger.info(f"Task {idx+1} already exists on S3. Skipping.") continue logger.info(f"Generating Task {idx+1}/{len(tasks)}: {task.get('name')}") task_adapters, task_weights = [], [] for lora in task.get("loras", []): if lora["path"] in active_loras: task_adapters.append(active_loras[lora["path"]]) task_weights.append(lora.get("scale", 1.0)) if task_adapters: pipe.set_adapters(task_adapters, adapter_weights=task_weights) else: pipe.disable_lora() with torch.inference_mode(): # CLEAN GENERATION: No manual casting needed on A100 BFloat16 image = pipe( prompt=task.get("prompt", ""), height=task.get("height", 1024), width=task.get("width", 720), guidance_scale=task.get("guidance", 3.5), num_inference_steps=50, generator=torch.Generator("cpu").manual_seed(task.get("seed", 42)) ).images[0] local_path = output_dir / filename image.save(local_path, format="WEBP", quality=95) upload_to_s3(local_path, s3_bucket, s3_key) local_path.unlink() logger.info("Job Complete.") if __name__ == "__main__": main()