add intermediate steps

.idea/misc.xml

@@ -3,5 +3,5 @@
   <component name="Black">
     <option name="sdkName" value="Python 3.13" />
   </component>
-  <component name="ProjectRootManager" version="2" project-jdk-name="Python 3.13 (style_transfer_app)" project-jdk-type="Python SDK" />
+  <component name="ProjectRootManager" version="2" project-jdk-name="Python 3.13 virtualenv at ~/Desktop/my_ai_apps/transfer_image_style/.venv" project-jdk-type="Python SDK" />
 </project>

.idea/style_transfer_app.iml

@@ -4,7 +4,7 @@
     <content url="file://$MODULE_DIR$">
       <excludeFolder url="file://$MODULE_DIR$/.venv" />
     </content>
-    <orderEntry type="jdk" jdkName="Python 3.13 (style_transfer_app)" jdkType="Python SDK" />
+    <orderEntry type="jdk" jdkName="Python 3.13 virtualenv at ~/Desktop/my_ai_apps/transfer_image_style/.venv" jdkType="Python SDK" />
     <orderEntry type="sourceFolder" forTests="false" />
   </component>
 </module>

app.py

@@ -5,7 +5,7 @@ from model import generate_image
 from utils import load_model, load_image, im_convert, device
 
 
-vgg = load_model(d=device)
+model = load_model(d=device)
 
 max_image_size = 400
 
@@ -17,8 +17,12 @@ def generate(content: torch.Tensor, style: torch.Tensor, alpha_slider: float):
 
     target_img = content_img.clone().requires_grad_(True).to(
         device)  # Initialize the target image as a clone of the original content image
-    target = generate_image(model=vgg, content=content_img, style=style_img, target=target_img, steps = 2700, content_wt=alpha_slider)
-    return im_convert(target)
+
+    steps = 500
+    for target, status in generate_image(model=model, content=content_img, style=style_img, target=target_img, steps = steps, content_wt=alpha_slider):
+        yield im_convert(target), status
+
+    # return target
 
 
 def check_inputs(img1, img2):
@@ -46,9 +50,10 @@ with gr.Blocks() as demo:
     alpha_slider = gr.Slider(0, 1, value=1, step=0.1, label="Blending Ratio")
     submit_button = gr.Button("Blend Images", "Generate", variant="primary", interactive=False)
 
-    output = gr.Image(label="Blended Image")
+    output = gr.Image(type="pil", label="Blended Image")
+    status = gr.Textbox(label="Progress")
 
-    submit_button.click(generate, inputs=[content_image, style_image, alpha_slider], outputs=output)
+    submit_button.click(generate, inputs=[content_image, style_image, alpha_slider], outputs=[output, status])
 
     # When images change, check if both are uploaded to enable the button
     content_image.change(fn=check_inputs, inputs=[content_image, style_image], outputs=submit_button)
@@ -56,7 +61,7 @@ with gr.Blocks() as demo:
 
 
 # Launch the demo!
-demo.launch()
+# demo.launch()
 
-# if __name__ == "__main__":
-#     demo.launch()
+if __name__ == "__main__":
+    demo.launch()

model.py

@@ -46,7 +46,15 @@ def generate_image(model: nn.Module, content: torch.Tensor, style: torch.Tensor,
         total_loss.backward()
         optimizer.step()
 
-    return target
+        if ii % 10 == 0:
+            status = f"Processing step {ii} of {steps}"
+            if ii == steps:
+                status = "✅ Completed"
+            yield target, status
+
+
+    # return target
+
 
 
 

utils.py

@@ -1,5 +1,4 @@
 import torch
-import torch.optim as optim
 from torchvision import transforms, models
 from PIL import Image
 import numpy as np
@@ -10,14 +9,14 @@ device = torch.device("cuda" if torch.cuda.is_available() else "mps" if torch.mp
 
 def load_model(d=device):
     weights = models.VGG19_Weights.DEFAULT
-    vgg = models.vgg19(weights=weights).features  # only uses the feature layers of the model
+    model = models.vgg19(weights=weights).features  # only uses the feature layers of the model
 
     # https://pytorch.org/docs/stable/generated/torch.Tensor.requires_grad_.html
-    for param in vgg.parameters():
+    for param in model.parameters():
         param.requires_grad_(False)
 
-    vgg.to(device=d)
-    return vgg
+    model.to(device=d)
+    return model
 
 
 # max_size limits the image size to 400 pixel