#img_gen_modal.py import modal import random import io from config.config import prompts, models # Indirect import import os import gradio as gr #MOVED FROM IMAGE IMPORT LIST import torch import sentencepiece import torch from huggingface_hub import login from transformers import AutoTokenizer import random from datetime import datetime #import xformers ########## LIVE PREVIEW TEST 1/3 ########## #from live_preview_helpers import flux_pipe_call_that_returns_an_iterable_of_images ########################################### CACHE_DIR = "/model_cache" # Define the Modal image image = ( modal.Image.from_registry("nvidia/cuda:12.2.0-devel-ubuntu22.04", add_python="3.9") .pip_install_from_requirements("requirements.txt") #modal.Image.debian_slim(python_version="3.9") # Base image # .apt_install( # "git", # ) # .pip_install( # "diffusers", # "transformers", # "xformers", # "torch", # "accelerate", # "gradio>=4.44.1", # "safetensors", # "pillow", # "sentencepiece", # "hf_transfer", # "huggingface_hub[hf_transfer]", # "aria2", # aria2 for ultra-fast parallel downloads # f"git+https://github.com/huggingface/transformers.git" # ) .env( { "HF_HUB_ENABLE_HF_TRANSFER": "1", "HF_HOME": "HF_HOME", "HF_HUB_CACHE": CACHE_DIR } ) ) # Create a Modal app app = modal.App("img-gen-modal", image=image) with image.imports(): import os flux_model_vol = modal.Volume.from_name("flux-model-vol", create_if_missing=True) # Reference your volume ############ LIVE PREVIEW 2/3 ################## # dtype = torch.bfloat16 # device = "cuda" if torch.cuda.is_available() else "cpu" # taef1 = AutoencoderTiny.from_pretrained("madebyollin/taef1", torch_dtype=dtype).to(device) # good_vae = AutoencoderKL.from_pretrained("black-forest-labs/FLUX.1-dev", subfolder="vae", torch_dtype=dtype).to(device) # pipe = DiffusionPipeline.from_pretrained("black-forest-labs/FLUX.1-dev", torch_dtype=dtype, vae=taef1).to(device) # torch.cuda.empty_cache() # MAX_SEED = np.iinfo(np.int32).max # MAX_IMAGE_SIZE = 2048 #pipe.flux_pipe_call_that_returns_an_iterable_of_images = flux_pipe_call_that_returns_an_iterable_of_images.__get__(pipe) ################################################# # GPU FUNCTION @app.function(volumes={"/data": flux_model_vol}, secrets=[modal.Secret.from_name("huggingface-token")], gpu="L40S", timeout = 300 ) def generate_image_gpu(prompt_alias, team_color, model_alias, custom_prompt): image = generate_image(prompt_alias, team_color, model_alias, custom_prompt) return image, "Image generated successfully! Call the banners!" # CPU FUNCTION @app.function(volumes={"/data": flux_model_vol}, secrets=[modal.Secret.from_name("huggingface-token")], cpu = 1, timeout = 300 ) def generate_image_cpu(prompt_alias, team_color, model_alias, custom_prompt): image = generate_image(prompt_alias, team_color, model_alias, custom_prompt) return image, "Image generated successfully! Call the banners!" # MAIN GENERATE IMAGE FUNCTION def generate_image( prompt_alias, team_color, model_alias, custom_prompt, height=360, width=640, num_inference_steps=20, guidance_scale=2.0, seed=-1, progress=gr.Progress(track_tqdm=True) # Add progress parameter ): with modal.enable_output(): print("Hello from ctb_modal!") print("Running debug check...") # Debug function to check installed packages def check_dependencies(): packages = [ "diffusers", # For Stable Diffusion "transformers", # For Hugging Face models "torch", # PyTorch "accelerate", # For distributed training/inference "gradio", # For the Gradio interface (updated to latest version) "safetensors", # For safe model loading "pillow", # For image processing "sentencepiece" ] for package in packages: try: import importlib module = importlib.import_module(package) print(f" {package} is installed. Version:") except ImportError: print(f" {package} is NOT installed.") check_dependencies() # Find the selected prompt and model try: prompt = next(p for p in prompts if p["alias"] == prompt_alias)["text"] model_name = next(m for m in models if m["alias"] == model_alias)["name"] except StopIteration: return None, "ERROR: Invalid prompt or model selected." # Determine the enemy color enemy_color = "blue" if team_color.lower() == "red" else "red" # Print the original prompt and dynamic values for debugging print("Original Prompt:") print(prompt) print(f"Enemy Color: {enemy_color}") print(f"Team Color: {team_color.lower()}") prompt = prompt.format(team_color=team_color.lower(), enemy_color=enemy_color) # Print the formatted prompt for debugging print("\nFormatted Prompt:") print(prompt) # Append the custom prompt (if provided) if custom_prompt and len(custom_prompt.strip()) > 0: prompt += " " + custom_prompt.strip() # Randomize the seed if needed if seed == -1: seed = random.randint(0, 1000000) try: from diffusers import FluxPipeline print("Initializing HF TOKEN") hf_token = os.environ["HF_TOKEN"] print(hf_token) print("HF TOKEN:") login(token=hf_token) print("model_name:") print(model_name) # Use absolute path with leading slash local_path = f"/data/{model_name}" # Changed from "data/" to "/data/" print(f"Loading model from local path: {local_path}") # Debug: Check if the directory exists and list its contents if os.path.exists(local_path): print("Directory exists. Contents:") for item in os.listdir(local_path): print(f" - {item}") else: print(f"Directory does not exist: {local_path}") print("Contents of /data:") print(os.listdir("/data")) # CHECK FOR TORCH USING CUDA print("CHECK FOR TORCH USING CUDA") print(f"CUDA available: {torch.cuda.is_available()}") if torch.cuda.is_available(): print("inside if") print(f"CUDA device count: {torch.cuda.device_count()}") print(f"Current device: {torch.cuda.current_device()}") print(f"Device name: {torch.cuda.get_device_name(torch.cuda.current_device())}") ########## INITIALIZING CPU PIPE ########## print("-----INITIALIZING PIPE-----") pipe = FluxPipeline.from_pretrained( local_path, torch_dtype=torch.bfloat16, #torch_dtype=torch.float16, #torch_dtype=torch.float32, local_files_only=True ) if torch.cuda.is_available(): print("CUDA available") print("using gpu") pipe = pipe.to("cuda") pipe_message = "CUDA" else: print("CUDA not available") print("using cpu") pipe = pipe.to("cpu") pipe_message = "CPU" # pipe.enable_model_cpu_offload() # Use official recommended method print(f"-----{pipe_message} PIPE INITIALIZED-----") print(f"Using device: {pipe.device}") except Exception as e: print(f"Detailed error: {str(e)}") return None, f"ERROR: Failed to initialize PIPE2. Details: {e}" try: print("-----SENDING IMG GEN TO PIPE-----") print("-----HOLD ON-----") # ################ LIVE PREVIEW TEST 3/3 #################### # seed = random.randint(0, MAX_SEED) # generator = torch.Generator().manual_seed(seed) # for img in pipe.flux_pipe_call_that_returns_an_iterable_of_images( # prompt=prompt, # guidance_scale=guidance_scale, # num_inference_steps=num_inference_steps, # width=width, # height=height, # generator=generator, # output_type="pil", # good_vae=good_vae, # ): # yield img, seed # ############################################################ ########## SENDING IMG GEN TO PIPE - WORKING CODE ########## image = pipe( prompt, guidance_scale=guidance_scale, num_inference_steps=num_inference_steps, width=width, height=height, max_sequence_length=512, # seed=seed ).images[0] ############################################################# print("-----IMAGE GENERATED SUCCESSFULLY!-----") print(image) except Exception as e: return f"ERROR: Failed to initialize InferenceClient. Details: {e}" try: print("-----SAVING-----") print("-----DONE!-----") print("-----CALL THE BANNERS!-----") # Save the image with a timestamped filename timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") output_filename = f"/data/images/{timestamp}_{model_alias.replace(' ', '_').lower()}_{prompt_alias.replace(' ', '_').lower()}_{team_color.lower()}.png" # Save the image using PIL's save method image.save(output_filename) print(f"File path: {output_filename}") except Exception as e: print(f"ERROR: Failed to save image. Details: {e}") # Return the filename and success message return image