app.py

import gradio as gr
import numpy as np
import random
import os
import spaces
import torch
from diffusers import DiffusionPipeline
from transformers import pipeline, AutoTokenizer
from huggingface_hub import login

hf_token = os.getenv("hf_token")
login(token=hf_token)

MAX_SEED = np.iinfo(np.int32).max
MAX_IMAGE_SIZE = 2048
PRELOAD_MODELS = False  # Easy switch for preloading

_text_gen_pipeline = None
_image_gen_pipeline = None

@spaces.GPU()
def get_image_gen_pipeline():
    global _image_gen_pipeline
    if _image_gen_pipeline is None:
        try:
            device = "cuda" if torch.cuda.is_available() else "cpu"
            dtype = torch.bfloat16
            _image_gen_pipeline = DiffusionPipeline.from_pretrained(
                # "black-forest-labs/FLUX.1-schnell",
                "black-forest-labs/FLUX.1-dev",
                torch_dtype=dtype,
            ).to(device)
        except Exception as e:
            print(f"Error loading image generation model: {e}")
            return None
    return _image_gen_pipeline

@spaces.GPU()
def get_text_gen_pipeline():
    global _text_gen_pipeline
    if _text_gen_pipeline is None:
         try:
              device = "cuda" if torch.cuda.is_available() else "cpu"
              tokenizer = AutoTokenizer.from_pretrained(
                  "mistralai/Mistral-7B-Instruct-v0.3",
                  use_fast=True  
              )
              # Set pad_token_id to eos_token_id if pad_token is not set
              if tokenizer.pad_token is None:
                  tokenizer.pad_token = tokenizer.eos_token
                  
              _text_gen_pipeline = pipeline(
                   "text-generation",
                   model="mistralai/Mistral-7B-Instruct-v0.3",
                   tokenizer=tokenizer,
                   max_new_tokens=2048,
                   device=device,
                   pad_token_id=tokenizer.pad_token_id  # Explicitly set pad_token_id
              )
         except Exception as e:
              print(f"Error loading text generation model: {e}")
              return None
    return _text_gen_pipeline

@spaces.GPU()
def refine_prompt(prompt):
    text_gen = get_text_gen_pipeline()
    if text_gen is None:
        return "Text generation model is unavailable."
    try:
        messages = [
            {"role": "system", "content": "Vous êtes un designer produit avec de solides connaissances dans la génération de texte en image. Vous recevrez une demande de produit sous forme de description succincte, et votre mission sera d'imaginer un nouveau design de produit répondant à ce besoin.\n\nLe livrable (réponse générée) sera exclusivement un texte de prompt pour l'IA de texte to image FLUX.1-schnell.\n\nCe prompt devra inclure une description visuelle de l'objet mentionnant explicitement les aspects indispensables de sa fonction.\nA coté de ça vous devez aussi explicitement mentionner dans ce prompt les caractéristiques esthétiques/photo du rendu image (ex : photoréaliste, haute qualité, focale, grain, etc.), sachant que l'image sera l'image principale de cet objet dans le catalogue produit. Le fond de l'image générée doit être entièrement blanc.\nLe prompt doit être sans narration, peut être long mais ne doit pas dépasser 512 jetons."},            {"role": "user", "content": prompt},
        ]
        refined_prompt = text_gen(messages)
        
        # Extract just the assistant's content from the response
        try:
            messages = refined_prompt[0]['generated_text']
            # Find the last message with role 'assistant'
            assistant_messages = [msg for msg in messages if msg['role'] == 'assistant']
            if not assistant_messages:
                return "Error: No assistant response found"
            assistant_content = assistant_messages[-1]['content']
            return assistant_content
        except (KeyError, IndexError):
            return "Error: Unexpected response format from the model"
    except Exception as e:
        return f"Error refining prompt: {str(e)}"

def validate_dimensions(width, height):
    if width * height > MAX_IMAGE_SIZE * MAX_IMAGE_SIZE:
        return False, "Image dimensions too large"
    return True, None

@spaces.GPU()
def infer(prompt, seed=42, randomize_seed=False, width=1024, height=1024, num_inference_steps=4): # , progress=gr.Progress(track_tqdm=True)
    try:
        # progress(0, desc="Starting generation...")
        
        # Validate that prompt is not empty
        if not prompt or prompt.strip() == "":
            return None, "Please provide a valid prompt."

        # progress(0.1, desc="Loading image generation model...")
        pipe = get_image_gen_pipeline()
        if pipe is None:
            return None, "Image generation model is unavailable."
        
        # progress(0.2, desc="Validating dimensions...")
        is_valid, error_msg = validate_dimensions(width, height)
        if not is_valid:
            return None, error_msg

        if randomize_seed:
            seed = random.randint(0, MAX_SEED)
        
        # progress(0.3, desc="Setting up generator...")
        generator = torch.Generator("cuda").manual_seed(seed)  # Explicitly use CUDA generator
        
        # progress(0.4, desc="Generating image...")
        with torch.autocast('cuda'):
            image = pipe(
                prompt=prompt, 
                width=width,
                height=height,
                num_inference_steps=num_inference_steps, 
                generator=generator,
                guidance_scale=0.0,  # Increased guidance scale
                # max_sequence_length=512 
            ).images[0]

        torch.cuda.empty_cache()  # Clean up GPU memory after generation
        # progress(1.0, desc="Done!")
        return image, seed
    except Exception as e:
        print(f"Error in infer: {str(e)}")  # Add detailed error logging
        return None, f"Error generating image: {str(e)}"
 
examples = [
    "a backpack for kids, flower style",
    "medieval flip flops",
    "cat shaped cake mold",
]

css="""
#col-container {
    margin: 0 auto;
    max-width: 520px;
}
"""

def preload_models():
    print("Préchargement des modèles...")
    try:
        # Préchargement du modèle de génération de texte
        device = "cuda" if torch.cuda.is_available() else "cpu"
        # Explicitly load the fast tokenizer LGR
        tokenizer = AutoTokenizer.from_pretrained(
            "mistralai/Mistral-7B-Instruct-v0.3",
            use_fast=True  # Ensures a fast tokenizer is used
        )
        _text_gen_pipeline = pipeline(
            "text-generation",
            model="mistralai/Mistral-7B-Instruct-v0.3",
            tokenizer=tokenizer,  # Pass the fast tokenizer in LGR
            max_new_tokens=2048,
            device=device,
        )
        
        # Préchargement du modèle de génération d'images
        dtype = torch.bfloat16
        _image_gen_pipeline = DiffusionPipeline.from_pretrained(
            "black-forest-labs/FLUX.1-schnell",
            torch_dtype=dtype
        ).to(device)
        
        print("Modèles préchargés avec succès!")
        return True
    except Exception as e:
        print(f"Erreur lors du préchargement des modèles: {str(e)}")
        return False

def create_interface():
    # Modify the preloading logic
    if PRELOAD_MODELS:
        models_loaded = preload_models()
        model_status = "✅ Modèles chargés avec succès!" if models_loaded else "⚠️ Erreur lors du chargement des modèles"
    else:
        model_status = "ℹ️ Modèles seront chargés à la demande"

    with gr.Blocks(css=css) as demo:
        info = gr.Info(model_status)

        with gr.Column(elem_id="col-container"):
            gr.Markdown(f"""# Text to Product
            Using Mistral-7B-Instruct-v0.3 + FLUX.1-dev + Trellis
            """)
            
            with gr.Row():
                
                prompt = gr.Text(
                    label="Prompt",
                    show_label=False,
                    max_lines=1,
                    placeholder="Enter your prompt",
                    container=False,
                )
                
                prompt_button = gr.Button("Refine prompt", scale=0)
            
            refined_prompt = gr.Text(
                label="Refined Prompt", 
                show_label=False,
                max_lines=10, 
                placeholder="Prompt refined by Mistral", 
                container=False,
                max_length=2048,
                )
            
            
            run_button = gr.Button("Create visual", scale=0)

            generated_image = gr.Image(label="Generated Image", show_label=False)
            
            with gr.Accordion("Advanced Settings Mistral", open=False):
                gr.Slider(
                    label="Temperature",
                    value=0.9,
                    minimum=0.0,
                    maximum=1.0,
                    step=0.05,
                    interactive=True,
                    info="Higher values produce more diverse outputs",
                ),
                gr.Slider(
                    label="Max new tokens",
                    value=256,
                    minimum=0,
                    maximum=1048,
                    step=64,
                    interactive=True,
                    info="The maximum numbers of new tokens",
                ),
                gr.Slider(
                    label="Top-p (nucleus sampling)",
                    value=0.90,
                    minimum=0.0,
                    maximum=1,
                    step=0.05,
                    interactive=True,
                    info="Higher values sample more low-probability tokens",
                ),
                gr.Slider(
                    label="Repetition penalty",
                    value=1.2,
                    minimum=1.0,
                    maximum=2.0,
                    step=0.05,
                    interactive=True,
                    info="Penalize repeated tokens",
                )
            
            with gr.Accordion("Advanced Settings Flux", open=False):
                
                seed = gr.Slider(
                    label="Seed",
                    minimum=0,
                    maximum=MAX_SEED,
                    step=1,
                    value=0,
                )
                
                randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
                
                with gr.Row():
                    
                    width = gr.Slider(
                        label="Width",
                        minimum=256,
                        maximum=MAX_IMAGE_SIZE,
                        step=32,
                        value=512,
                    )
                    
                    height = gr.Slider(
                        label="Height",
                        minimum=256,
                        maximum=MAX_IMAGE_SIZE,
                        step=32,
                        value=512,
                    )
                
                with gr.Row():
                    
      
                    num_inference_steps = gr.Slider(
                        label="Number of inference steps",
                        minimum=1,
                        maximum=50,
                        step=1,
                        value=4,
                    )
            
            gr.Examples(
                examples=examples,
                fn=refine_prompt,
                inputs = [prompt],
                outputs = [refined_prompt],
                cache_examples=True,
                cache_mode='lazy'
            )


        gr.on(
            triggers=[prompt_button.click, prompt.submit],
            fn = refine_prompt,
            inputs = [prompt],
            outputs = [refined_prompt]
        )

        gr.on(
            triggers=[run_button.click],
            fn = infer,
            inputs = [refined_prompt, seed, randomize_seed, width, height, num_inference_steps],
            outputs = [generated_image, prompt]
        )

    return demo

if __name__ == "__main__":
    demo = create_interface()
    demo.launch()