old/old_live_test/img_gen_modal_LIVE1.py

#img_gen_modal.py
import modal
import random
import io
from config.config import prompts, models  # Indirect import
import os
import gradio as gr
import torch
import sentencepiece
import torch
from huggingface_hub import login
from transformers import AutoTokenizer
import random
from datetime import datetime

####### IMPORTS FOR LIVE PREVIEW
import numpy as np
from diffusers import  DiffusionPipeline, FlowMatchEulerDiscreteScheduler, AutoencoderTiny, AutoencoderKL
from transformers import CLIPTextModel, CLIPTokenizer,T5EncoderModel, T5TokenizerFast
from live_preview_helpers import calculate_shift, retrieve_timesteps, flux_pipe_call_that_returns_an_iterable_of_images




######### 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

    .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


# 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!")


        ########### LIVE PREVIEW 2/3 ##################
        dtype = torch.bfloat16
        device = "cuda" if torch.cuda.is_available() else "cpu"

        taef1 = AutoencoderTiny.from_pretrained("/data/taef1", torch_dtype=dtype).to(device)
        good_vae = AutoencoderKL.from_pretrained("/data/FLUX.1-dev", subfolder="vae", torch_dtype=dtype).to(device)
        pipe_LIVE = DiffusionPipeline.from_pretrained("/data/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_LIVE.flux_pipe_call_that_returns_an_iterable_of_images = flux_pipe_call_that_returns_an_iterable_of_images.__get__(pipe_LIVE)
        #################################################

        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 ####################
            print("-----SENDING IMG GEN TO PIPE LIVE-----")
            print("-----HOLD ON-----")   
            seed = random.randint(0, MAX_SEED)
            generator = torch.Generator().manual_seed(seed)
            
            for image in pipe_LIVE.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 image, "Update"
            ############################################################


            # ########## 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