Initial Commit

app.py

@@ -8,6 +8,18 @@ from flux.transformer_flux import FluxTransformer2DModel
 from flux.pipeline_flux_chameleon import FluxPipeline
 import torch.nn as nn
 import math
+import logging
+import sys
+
+# 设置日志
+logging.basicConfig(
+    level=logging.INFO,
+    format='%(asctime)s - %(levelname)s - %(message)s',
+    handlers=[
+        logging.StreamHandler(sys.stdout)
+    ]
+)
+logger = logging.getLogger(__name__)
 
 MODEL_ID = "Djrango/Qwen2vl-Flux"
 
@@ -177,40 +189,40 @@ class FluxInterface:
 
     def generate(self, input_image, prompt="", guidance_scale=3.5, num_inference_steps=28, num_images=2, seed=None, aspect_ratio="1:1"):
         try:
-            print(f"Starting generation with prompt: {prompt}, guidance_scale: {guidance_scale}, steps: {num_inference_steps}")
+            logger.info(f"Starting generation with prompt: {prompt}, guidance_scale: {guidance_scale}, steps: {num_inference_steps}")
             
             if input_image is None:
                 raise ValueError("No input image provided")
                 
             if seed is not None:
                 torch.manual_seed(seed)
-                print(f"Set random seed to: {seed}")
+                logger.info(f"Set random seed to: {seed}")
                 
             self.load_models()
-            print("Models loaded successfully")
+            logger.info("Models loaded successfully")
             
             # Get dimensions from aspect ratio
             if aspect_ratio not in ASPECT_RATIOS:
                 raise ValueError(f"Invalid aspect ratio. Choose from {list(ASPECT_RATIOS.keys())}")
             width, height = ASPECT_RATIOS[aspect_ratio]
-            print(f"Using dimensions: {width}x{height}")
+            logger.info(f"Using dimensions: {width}x{height}")
             
             # Process input image
             try:
                 input_image = self.resize_image(input_image)
-                print(f"Input image resized to: {input_image.size}")
+                logger.info(f"Input image resized to: {input_image.size}")
                 qwen2_hidden_state, image_grid_thw = self.process_image(input_image)
-                print("Input image processed successfully")
+                logger.info("Input image processed successfully")
             except Exception as e:
                 raise RuntimeError(f"Error processing input image: {str(e)}")
             
             try:
                 pooled_prompt_embeds = self.compute_text_embeddings("")
-                print("Base text embeddings computed")
+                logger.info("Base text embeddings computed")
                 
                 # Get T5 embeddings if prompt is provided
                 t5_prompt_embeds = self.compute_t5_text_embeddings(prompt)
-                print("T5 prompt embeddings computed")
+                logger.info("T5 prompt embeddings computed")
             except Exception as e:
                 raise RuntimeError(f"Error computing embeddings: {str(e)}")
             
@@ -226,12 +238,12 @@ class FluxInterface:
                     width=width,
                 ).images
                 
-                print("Images generated successfully")
+                logger.info("Images generated successfully")
                 return output_images
             except Exception as e:
                 raise RuntimeError(f"Error generating images: {str(e)}")
         except Exception as e:
-            print(f"Error during generation: {str(e)}")
+            logger.error(f"Error during generation: {str(e)}")
             raise gr.Error(f"Generation failed: {str(e)}")
 
 # Initialize the interface
@@ -364,13 +376,25 @@ with gr.Blocks(
     # Set up the generation function
     def generate_with_error_handling(*args):
         try:
+            logger.info("Starting image generation with args: %s", str(args))
+            
+            # 输入参数验证
+            input_image, prompt, guidance, steps, num_images, seed, aspect_ratio = args
+            logger.info(f"Input validation - Image: {type(input_image)}, Prompt: '{prompt}', "
+                       f"Guidance: {guidance}, Steps: {steps}, Num Images: {num_images}, "
+                       f"Seed: {seed}, Aspect Ratio: {aspect_ratio}")
+            
+            if input_image is None:
+                raise ValueError("No input image provided")
+            
             gr.Info("Starting image generation...")
             results = interface.generate(*args)
+            logger.info("Generation completed successfully")
             gr.Info("Generation complete!")
             return [results, None]
         except Exception as e:
             error_msg = str(e)
-            print(f"Error in generate_with_error_handling: {error_msg}")
+            logger.error(f"Error in generate_with_error_handling: {error_msg}", exc_info=True)
             return [None, error_msg]
 
     submit_btn.click(