full change

main.py

@@ -34,25 +34,77 @@ class Upscale_CaseCade:
         )
         self.models_b.generator.eval().requires_grad_(False)
         print("STAGE B READY")
+        self.caption = "a photo of image"
+        self.cnet_multiplier = 1.0 # 0.8 # 0.3
+        # Stage C Parameters
+        self.extras.sampling_configs['cfg'] = 1
+        self.extras.sampling_configs['shift'] = 2
+        self.extras.sampling_configs['timesteps'] = 20
+        self.extras.sampling_configs['t_start'] = 1.0
+        # Stage B Parameters
+        self.extras_b.sampling_configs['cfg'] = 1.1
+        self.extras_b.sampling_configs['shift'] = 1
+        self.extras_b.sampling_configs['timesteps'] = 10
+        self.extras_b.sampling_configs['t_start'] = 1.0
+        self.models = ControlNetCore.Models(
+           **{**self.models.to_dict(), 'generator': torch.compile(self.models.generator, mode="reduce-overhead", fullgraph=True)}
+        )
+
+        self.models_b = WurstCoreB.Models(
+           **{**self.models_b.to_dict(), 'generator': torch.compile(self.models_b.generator, mode="reduce-overhead", fullgraph=True)}
+           )
+
 
 
-    def upscale_image(self,image_pil,scale_fator):
+    def upscale_image(self,caption,image_pil,scale_fator):
         batch_size = 1
         cnet_override = None
         images = resize_image(image_pil).unsqueeze(0).expand(batch_size, -1, -1, -1)
         
         batch = {'images': images}
 
+
         with torch.no_grad(), torch.cuda.amp.autocast(dtype=torch.bfloat16):
             effnet_latents = self.core.encode_latents(batch, self.models, self.extras)
             effnet_latents_up = torch.nn.functional.interpolate(effnet_latents, scale_factor=scale_fator, mode="nearest")
             cnet = self.models.controlnet(effnet_latents_up)
             cnet_uncond = cnet
             cnet_input = torch.nn.functional.interpolate(images, scale_factor=scale_fator, mode="nearest")
-            # cnet, cnet_input = core.get_cnet(batch, models, extras)
+            # cnet, cnet_input = self.core.get_cnet(batch, self.models, self.extras)
             # cnet_uncond = cnet
-        og=show_images(batch['images'],return_images=True)
-        upsclae=show_images(cnet_input,return_images=True)
-        return og,upsclae
+            height, width = int(cnet[0].size(-2)*32*4/3), int(cnet[0].size(-1)*32*4/3)
+            stage_c_latent_shape, stage_b_latent_shape = calculate_latent_sizes(height, width, batch_size=batch_size)
+            # PREPARE CONDITIONS
+            batch['captions'] = [caption] * batch_size
+            conditions = self.core.get_conditions(batch, self.models, self.extras, is_eval=True, is_unconditional=False, eval_image_embeds=False)
+            unconditions =  self.core.get_conditions(batch, self.models, self.extras, is_eval=True, is_unconditional=True, eval_image_embeds=False)    
+            conditions['cnet'] = [c.clone() * self.cnet_multiplier if c is not None else c for c in cnet]
+            unconditions['cnet'] = [c.clone() * self.cnet_multiplier if c is not None else c for c in cnet_uncond]
+            conditions_b = self.core_b.get_conditions(batch, self.models_b, self.extras_b, is_eval=True, is_unconditional=False)
+            unconditions_b = self.core_b.get_conditions(batch, self.models_b, self.extras_b, is_eval=True, is_unconditional=True)
+            # torch.manual_seed(42)
+            sampling_c = self.extras.gdf.sample(
+                self.models.generator, conditions, stage_c_latent_shape,
+                unconditions, device=device, **self.extras.sampling_configs,
+            )
+            for (sampled_c, _, _) in tqdm(sampling_c, total=self.extras.sampling_configs['timesteps']):
+                sampled_c = sampled_c
+                
+            # preview_c = models.previewer(sampled_c).float()
+            # show_images(preview_c)
+
+            conditions_b['effnet'] = sampled_c
+            unconditions_b['effnet'] = torch.zeros_like(sampled_c)
+
+            sampling_b = self.extras_b.gdf.sample(
+                self.models_b.generator, conditions_b, stage_b_latent_shape,
+                unconditions_b, device=device, **self.extras_b.sampling_configs
+            )
+            for (sampled_b, _, _) in tqdm(sampling_b, total=self.extras_b.sampling_configs['timesteps']):
+                sampled_b = sampled_b
+            sampled = self.models_b.stage_a.decode(sampled_b).float()
+        # og=show_images(batch['images'],return_images=True)
+        upscale=show_images(sampled,return_images=True)
+        return upscale