Improve image generation with better processing and model configuration

app.py

@@ -7,6 +7,7 @@ import torch
 from diffusers import DiffusionPipeline
 from transformers import pipeline, AutoTokenizer
 from huggingface_hub import login
+from PIL import Image
 
 hf_token = os.getenv("hf_token")
 login(token=hf_token)
@@ -30,6 +31,8 @@ def get_image_gen_pipeline():
                 "black-forest-labs/FLUX.1-dev",
                 torch_dtype=dtype,
             ).to(device)
+            _image_gen_pipeline.enable_model_cpu_offload()
+            _image_gen_pipeline.enable_vae_slicing()
         except Exception as e:
             print(f"Error loading image generation model: {e}")
             return None
@@ -69,7 +72,7 @@ def refine_prompt(prompt):
         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},
+            {"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 77 jetons."},            {"role": "user", "content": prompt},
         ]
         refined_prompt = text_gen(messages)
         
@@ -93,20 +96,17 @@ def validate_dimensions(width, height):
     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)
+def infer(prompt, seed=42, randomize_seed=False, width=1024, height=1024, num_inference_steps=4): 
     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
@@ -114,26 +114,30 @@ def infer(prompt, seed=42, randomize_seed=False, width=1024, height=1024, num_in
         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(
+            output = 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]
+                guidance_scale=7.5,
+                max_sequence_length=512
+            )
+            
+            # Ensure the image is properly normalized and converted
+            image = output.images[0]
+            if isinstance(image, torch.Tensor):
+                image = (image.clamp(-1, 1) + 1) / 2
+                image = (image * 255).round().clamp(0, 255).to(torch.uint8).cpu().numpy()
+                image = Image.fromarray(image)
 
-        torch.cuda.empty_cache()  # Clean up GPU memory after generation
-        # progress(1.0, desc="Done!")
+        torch.cuda.empty_cache()  
         return image, seed
     except Exception as e:
-        print(f"Error in infer: {str(e)}")  # Add detailed error logging
+        print(f"Error in infer: {str(e)}")  
         return None, f"Error generating image: {str(e)}"
  
 examples = [