Upload app.py with huggingface_hub

app.py

@@ -141,7 +141,7 @@ class FluxDinoPascalModel(nn.Module):
         for timestep in timestep_data['timesteps']:
             single_features = timestep_data['features'][timestep]['single_features']
             multi_timestep_features[timestep] = [
-                f.float().permute(0, 2, 1).reshape(1, 3072, patch_h, patch_w) for f in single_features]
+                f.float().to(device).permute(0, 2, 1).reshape(1, 3072, patch_h, patch_w) for f in single_features]
 
         flux_features, alpha_layers = self.hyperfeature_fusion(multi_timestep_features)
         del multi_timestep_features
@@ -234,9 +234,9 @@ class FluxDinoDUTSModel(nn.Module):
                 val = timestep_data['features'][timestep][key]
                 if isinstance(val, list):
                     timestep_data['features'][timestep][key] = [
-                        f.float() if isinstance(f, torch.Tensor) else f for f in val]
+                        f.float().to(device) if isinstance(f, torch.Tensor) else f for f in val]
                 elif isinstance(val, torch.Tensor):
-                    timestep_data['features'][timestep][key] = val.float()
+                    timestep_data['features'][timestep][key] = val.float().to(device)
 
         dino_features = self.dino_extractor(images)
 
@@ -332,7 +332,7 @@ class FluxNYUDepthModel(nn.Module):
         for timestep in timestep_data['timesteps']:
             single_features = timestep_data['features'][timestep]['single_features']
             multi_timestep_features[timestep] = [
-                f.float().permute(0, 2, 1).reshape(1, 3072, patch_h, patch_w) for f in single_features]
+                f.float().to(device).permute(0, 2, 1).reshape(1, 3072, patch_h, patch_w) for f in single_features]
 
         fused_flux_features = self.hyperfeature_fusion(multi_timestep_features)