Update app.py

app.py

@@ -2,6 +2,7 @@ import sys
 import os
 from pathlib import Path
 import gc
+import traceback
 
 # Add the StableCascade and CSD directories to the Python path
 app_dir = Path(__file__).parent
@@ -27,6 +28,7 @@ from utils import WurstCoreCRBM
 from gdf.schedulers import CosineSchedule
 from gdf import VPScaler, CosineTNoiseCond, DDPMSampler, P2LossWeight, AdaptiveLossWeight
 from gdf.targets import EpsilonTarget
+import PIL
 
 # Enable mixed precision
 torch.backends.cuda.matmul.allow_tf32 = True
@@ -73,94 +75,69 @@ if low_vram:
         
         clear_gpu_cache()
 
-# Stage C model configuration
+# Load configurations
 config_file = 'third_party/StableCascade/configs/inference/stage_c_3b.yaml'
 with open(config_file, "r", encoding="utf-8") as file:
     loaded_config = yaml.safe_load(file)
 
-core = WurstCoreCRBM(config_dict=loaded_config, device=device, training=False)
-
-# Stage B model configuration
 config_file_b = 'third_party/StableCascade/configs/inference/stage_b_3b.yaml'
 with open(config_file_b, "r", encoding="utf-8") as file:
     config_file_b = yaml.safe_load(file)
+
+def initialize_models():
+    global models_rbm, models_b, extras, extras_b, core, core_b
+    
+    # Clear any existing models from memory
+    models_rbm = None
+    models_b = None
+    extras = None
+    extras_b = None
+    
+    # Clear GPU cache
+    clear_gpu_cache()
+    
+    # Initialize models
+    core = WurstCoreCRBM(config_dict=loaded_config, device=device, training=False)
+    core_b = WurstCoreB(config_dict=config_file_b, device=device, training=False)
+    
+    extras = core.setup_extras_pre()
+    models = core.setup_models(extras)
+    
+    extras_b = core_b.setup_extras_pre()
+    models_b = core_b.setup_models(extras_b, skip_clip=True)
+    models_b = WurstCoreB.Models(
+        **{**models_b.to_dict(), 'tokenizer': models.tokenizer, 'text_model': models.text_model}
+    )
+    
+    # Initialize models_rbm
+    generator_rbm = StageCRBM()
+    for param_name, param in load_or_fail(core.config.generator_checkpoint_path).items():
+        set_module_tensor_to_device(generator_rbm, param_name, "cpu", value=param)
+    
+    generator_rbm = generator_rbm.to(getattr(torch, core.config.dtype)).to(device)
+    generator_rbm = core.load_model(generator_rbm, 'generator')
+    
+    models_rbm = core.Models(
+        effnet=models.effnet,
+        previewer=models.previewer,
+        generator=generator_rbm,
+        generator_ema=models.generator_ema,
+        tokenizer=models.tokenizer,
+        text_model=models.text_model,
+        image_model=models.image_model
+    )
+    
+    # Move models to appropriate devices
+    models_rbm.generator.to(device).eval().requires_grad_(False)
+    models_b.generator.to(device).eval().requires_grad_(False)
     
-core_b = WurstCoreB(config_dict=config_file_b, device=device, training=False)
-
-# Setup extras and models for Stage C
-extras = core.setup_extras_pre()
-
-gdf_rbm = RBM(
-    schedule=CosineSchedule(clamp_range=[0.0001, 0.9999]),
-    input_scaler=VPScaler(), target=EpsilonTarget(),
-    noise_cond=CosineTNoiseCond(),
-    loss_weight=AdaptiveLossWeight(),
-)
-
-sampling_configs = {
-    "cfg": 5,
-    "sampler": DDPMSampler(gdf_rbm),
-    "shift": 1,
-    "timesteps": 20
-}
-
-extras = core.Extras(
-    gdf=gdf_rbm,
-    sampling_configs=sampling_configs,
-    transforms=extras.transforms,
-    effnet_preprocess=extras.effnet_preprocess,
-    clip_preprocess=extras.clip_preprocess
-)
-
-models = core.setup_models(extras)
-models.generator.eval().requires_grad_(False)
-
-# Setup extras and models for Stage B
-extras_b = core_b.setup_extras_pre()
-models_b = core_b.setup_models(extras_b, skip_clip=True)
-models_b = WurstCoreB.Models(
-    **{**models_b.to_dict(), 'tokenizer': models.tokenizer, 'text_model': models.text_model}
-)
-models_b.generator.bfloat16().eval().requires_grad_(False)
-
-# Off-load old generator (low VRAM mode)
-if low_vram:
-    models.generator.to("cpu")
     clear_gpu_cache()
-
-# Load and configure new generator
-generator_rbm = StageCRBM()
-for param_name, param in load_or_fail(core.config.generator_checkpoint_path).items():
-    set_module_tensor_to_device(generator_rbm, param_name, "cpu", value=param)
-
-generator_rbm = generator_rbm.to(getattr(torch, core.config.dtype)).to(device)
-generator_rbm = core.load_model(generator_rbm, 'generator')
-
-# Create models_rbm instance
-models_rbm = core.Models(
-    effnet=models.effnet,
-    previewer=models.previewer,
-    generator=generator_rbm,
-    generator_ema=models.generator_ema,
-    tokenizer=models.tokenizer,
-    text_model=models.text_model,
-    image_model=models.image_model
-)
-models_rbm.generator.eval().requires_grad_(False)
 
 def infer(style_description, ref_style_file, caption):
     try:
-        # Move all model components to the same device and set to the same precision
-        models_rbm.effnet.to(device).bfloat16()
-        models_rbm.previewer.to(device).bfloat16()
-        models_rbm.generator.to(device).bfloat16()
-        models_rbm.text_model.to(device).bfloat16()
+        # Initialize (or reinitialize) models before each inference
+        initialize_models()
         
-        models_b.generator.to(device).bfloat16()
-        models_b.stage_a.to(device).bfloat16()
-        
-        clear_gpu_cache()  # Clear cache before inference
-
         height = 1024
         width = 1024
         batch_size = 1
@@ -178,7 +155,7 @@ def infer(style_description, ref_style_file, caption):
         extras_b.sampling_configs['timesteps'] = 10
         extras_b.sampling_configs['t_start'] = 1.0
 
-        ref_style = resize_image(PIL.Image.open(ref_style_file).convert("RGB")).unsqueeze(0).expand(batch_size, -1, -1, -1).to(device).bfloat16()
+        ref_style = resize_image(PIL.Image.open(ref_style_file).convert("RGB")).unsqueeze(0).expand(batch_size, -1, -1, -1).to(device)
 
         batch = {'captions': [caption] * batch_size}
         batch['style'] = ref_style
@@ -195,7 +172,7 @@ def infer(style_description, ref_style_file, caption):
             models_to(models_rbm, device="cpu", excepts=["generator", "previewer"])
 
         # Stage C reverse process
-        with torch.cuda.amp.autocast(dtype=torch.bfloat16):  # Use mixed precision with bfloat16
+        with torch.cuda.amp.autocast():
             sampling_c = extras.gdf.sample(
                 models_rbm.generator, conditions, stage_c_latent_shape,
                 unconditions, device=device,
@@ -212,9 +189,6 @@ def infer(style_description, ref_style_file, caption):
 
         clear_gpu_cache()  # Clear cache between stages
 
-        # Ensure all models are on the right device again
-        models_b.generator.to(device).bfloat16()
-        
         # Stage B reverse process
         with torch.no_grad(), torch.cuda.amp.autocast(dtype=torch.bfloat16):                
             conditions_b['effnet'] = sampled_c
@@ -243,6 +217,7 @@ def infer(style_description, ref_style_file, caption):
 
     except Exception as e:
         print(f"An error occurred during inference: {str(e)}")
+        traceback.print_exc()  # This will print the full traceback
         return None
 
     finally:
@@ -252,8 +227,11 @@ def infer(style_description, ref_style_file, caption):
 
 import gradio as gr
 
+def gradio_interface(style_description, ref_style_file, caption):
+    return infer(style_description, ref_style_file, caption)
+
 gr.Interface(
-    fn = infer,
+    fn=gradio_interface,
     inputs=[gr.Textbox(label="style description"), gr.Image(label="Ref Style File", type="filepath"), gr.Textbox(label="caption")],
     outputs=[gr.Image()]
 ).launch()