| from huggingface_hub import model_info
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| import json
|
| import argparse
|
| import sys
|
| import os
|
| import time
|
| import logging
|
| from pathlib import Path
|
| from datetime import datetime
|
|
|
|
|
| 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} ---")
|
|
|
|
|
| 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.")
|
|
|
|
|
| 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
|
|
|
|
|
| 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"
|
| )
|
|
|
|
|
| pipe.vae.enable_tiling()
|
|
|
|
|
| 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
|
|
|
|
|
| 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():
|
|
|
| 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()
|
|
|
|
|
|
|
