Clean up

nbs/finetune.ipynb

@@ -164,7 +164,7 @@
         "\n",
         "clear_output(wait=True)\n",
         "status_message = \"completed\" if model.status == FineTuneStatus.COMPLETED else \"failed\"\n",
-        "print(f\"\\rModel {model.name} ({model.id}) {status_message} after {elapsed:.0f} seconds\")"
+        "print(f\"Model {model.name} ({model.id}) {status_message} after {elapsed:.0f} seconds\")"
       ]
     },
     {
@@ -177,7 +177,7 @@
       "source": [
         "# If fine-tuning fails for some reason, you can resubmit the model\n",
         "if model.status == FineTuneStatus.FAILED:\n",
-        "    print(f\"Training failed, resubmitting\")\n",
+        "    print(\"Training failed, resubmitting\")\n",
         "    model = resubmit_model(context, model.id)"
       ]
     },

src/stability_sdk/finetune.py

@@ -1,5 +1,6 @@
 import logging
 import mimetypes
+
 from enum import Enum
 from google.protobuf.struct_pb2 import Struct
 from PIL import Image
@@ -134,13 +135,14 @@ def create_model(
                     use=generation.ASSET_USE_INPUT
                 )
             )
-            responses = context._stub_generation.Generate(request)
+
             success = False
             for response in context._stub_generation.Generate(request):
                 for artifact in response.artifacts:
                     if artifact.type == generation.ARTIFACT_TEXT:
                         logging.info(f"Uploaded image {i}: {artifact.text}")
                         success = True
+                        break
             if not success:
                 raise RuntimeError(f"Failed to upload image {image_paths[i]}")