Update app.py

app.py

@@ -116,25 +116,49 @@ def initialize_models():
         ae = ae.to(device)
         
         print("Loading Flux model...")
-        # Use the standard Flux model instead of quantized version
-        # This will use more memory but avoid compatibility issues
         from huggingface_hub import hf_hub_download
         from safetensors.torch import load_file
         
         try:
-            # Try to load from the standard Flux checkpoint
-            print("Loading standard Flux model (this may take a while)...")
+            # Try to load from a standard Flux checkpoint
+            # First, let's try the schnell version which might be smaller
+            print("Attempting to load Flux model weights...")
             model = Flux()
-            model = model.to(dtype=torch.bfloat16, device=device)
             
-            # You would need to download the standard Flux weights
-            # For now, let's create a randomly initialized model for testing
-            print("Warning: Using randomly initialized Flux model for testing")
-            print("To use a pretrained model, you need to load proper Flux weights")
+            # Try loading from black-forest-labs directly
+            try:
+                # Note: You might need to authenticate with HuggingFace for this
+                sd = load_file(hf_hub_download(repo_id="black-forest-labs/FLUX.1-schnell", filename="flux1-schnell.safetensors"))
+                # Adjust state dict keys if needed
+                model.load_state_dict(sd, strict=False)
+                print("Loaded Flux schnell model successfully!")
+            except Exception as e1:
+                print(f"Could not load Flux schnell: {e1}")
+                
+                # Try the dev version
+                try:
+                    sd = load_file(hf_hub_download(repo_id="black-forest-labs/FLUX.1-dev", filename="flux1-dev.safetensors"))
+                    model.load_state_dict(sd, strict=False)
+                    print("Loaded Flux dev model successfully!")
+                except Exception as e2:
+                    print(f"Could not load Flux dev: {e2}")
+                    
+                    # If no pretrained weights are available, warn the user
+                    print("\n" + "="*50)
+                    print("WARNING: Could not load pretrained Flux weights!")
+                    print("The model will use random initialization.")
+                    print("For proper results, you need to:")
+                    print("1. Authenticate with HuggingFace: huggingface-cli login")
+                    print("2. Accept the Flux model license agreement")
+                    print("3. Or use a publicly available Flux checkpoint")
+                    print("="*50 + "\n")
+            
+            model = model.to(dtype=torch.bfloat16, device=device)
             
         except Exception as e:
             print(f"Error initializing Flux model: {e}")
-            raise
+            # Continue with random initialization for now
+            model = Flux().to(dtype=torch.bfloat16, device=device)
         
         model_initialized = True
         print("Models initialized successfully!")