Update model.py

model.py

@@ -13,6 +13,13 @@ from shap_e.rendering.torch_mesh import TorchMesh
 from shap_e.util.collections import AttrDict
 from shap_e.util.image_util import load_image
 
+caption = gr.load(name="spaces/fffiloni/CoCa-clone")
+
+def create_image_caption(image_init):
+    cap = caption(image_init, "Beam search", 1.2, 0.5, 5, 20, fn_index=0)
+    print("cap: " + cap)
+    return cap
+
 
 # Copied from https://github.com/openai/shap-e/blob/d99cedaea18e0989e340163dbaeb4b109fa9e8ec/shap_e/util/notebooks.py#L15-L42
 def create_pan_cameras(size: int,
@@ -97,13 +104,15 @@ class Model:
         return mesh_path.name
 
     def run_text(self,
-                 prompt: str,
+                 image: str,
                  seed: int = 0,
                  guidance_scale: float = 15.0,
                  num_steps: int = 64) -> str:
         self.load_model('text300M')
         torch.manual_seed(seed)
 
+        prompt = create_image_caption(image)
+
         latents = sample_latents(
             batch_size=1,
             model=self.model_text,