Update ssh

model/model_manager.py

@@ -26,7 +26,30 @@ class ModelManager:
     @spaces.GPU(duration=120)
     def generate_image_ig(self, prompt, model_name):
         pipe = self.load_model_pipe(model_name)
-        result = pipe(prompt=prompt).images[0]
+        if 'cascade' not in name:
+            result = pipe(prompt=prompt).images[0]
+        else:
+            prior, decoder = pipe
+            prior.enable_model_cpu_offload()
+            prior_output = prior(
+                prompt=prompt,
+                height=512,
+                width=512,
+                negative_prompt='',
+                guidance_scale=4.0,
+                num_images_per_prompt=1,
+                num_inference_steps=20
+            )
+
+            decoder.enable_model_cpu_offload()
+            result = decoder(
+                image_embeddings=prior_output.image_embeddings.to(torch.float16),
+                prompt=prompt,
+                negative_prompt='',
+                guidance_scale=0.0,
+                output_type="pil",
+                num_inference_steps=10
+            ).images[0]
         return result
 
     def generate_image_ig_api(self, prompt, model_name):

model/models/huggingface_models.py

@@ -1,5 +1,6 @@
 from diffusers import DiffusionPipeline
 from diffusers import AutoPipelineForText2Image
+from diffusers import StableCascadeDecoderPipeline, StableCascadePriorPipeline
 import torch
 
 
@@ -8,11 +9,16 @@ import torch
 def load_huggingface_model(model_name, model_type):
     if model_name == "SD-turbo":
         pipe = AutoPipelineForText2Image.from_pretrained("stabilityai/sd-turbo", torch_dtype=torch.float16, variant="fp16")
+        pipe = pipe.to("cuda")
     elif model_name == "SDXL-turbo":
         pipe = AutoPipelineForText2Image.from_pretrained("stabilityai/sdxl-turbo", torch_dtype=torch.float16, variant="fp16")
+        pipe = pipe.to("cuda")
+    elif model_name == "Stable-cascade":
+        prior = StableCascadePriorPipeline.from_pretrained("stabilityai/stable-cascade-prior", variant="bf16", torch_dtype=torch.bfloat16)
+        decoder = StableCascadeDecoderPipeline.from_pretrained("stabilityai/stable-cascade", variant="bf16", torch_dtype=torch.float16)
+        pipe = [prior, decoder]
     else:
         raise NotImplementedError
-    pipe = pipe.to("cuda")
     # if model_name == "SD-turbo":
     #     pipe = AutoPipelineForText2Image.from_pretrained("stabilityai/sd-turbo")
     # elif model_name == "SDXL-turbo":
@@ -24,8 +30,9 @@ def load_huggingface_model(model_name, model_type):
 
 
 if __name__ == "__main__":
-    for name in ["SD-turbo", "SDXL-turbo"]:
-        load_huggingface_model(name, "text2image")
+    for name in ["SD-turbo", "SDXL-turbo"]: #"SD-turbo", "SDXL-turbo"
+        pipe = load_huggingface_model(name, "text2image")
+
 
     # for name in ["IF-I-XL-v1.0"]:
     #     pipe = load_huggingface_model(name, 'text2image')

serve/upload.py

@@ -79,18 +79,18 @@ def upload_ssh_all(states, output_dir, data, data_path):
         output_file_list.append(output_file)
         image_list.append(states[i].output)
 
-# with sftp_client as sftp:
-    for i in range(len(output_file_list)):
-        if isinstance(image_list[i], str):
-            print("get url image")
-            image_list[i] = get_image_from_url(image_list[i])
-        with io.BytesIO() as image_byte_stream:
-            image_list[i].save(image_byte_stream, format='JPEG')
-            image_byte_stream.seek(0)  
-            sftp_client.putfo(image_byte_stream, output_file_list[i])
-        print(f"Successfully uploaded image to {output_file_list[i]}")
-    json_data = json.dumps(data, indent=4)
-    with io.BytesIO(json_data.encode('utf-8')) as json_byte_stream:
-        sftp_client.putfo(json_byte_stream, data_path)
-    print(f"Successfully uploaded JSON data to {data_path}")
+    with sftp_client as sftp:
+        for i in range(len(output_file_list)):
+            if isinstance(image_list[i], str):
+                print("get url image")
+                image_list[i] = get_image_from_url(image_list[i])
+            with io.BytesIO() as image_byte_stream:
+                image_list[i].save(image_byte_stream, format='JPEG')
+                image_byte_stream.seek(0)  
+                sftp.putfo(image_byte_stream, output_file_list[i])
+            print(f"Successfully uploaded image to {output_file_list[i]}")
+        json_data = json.dumps(data, indent=4)
+        with io.BytesIO(json_data.encode('utf-8')) as json_byte_stream:
+            sftp.putfo(json_byte_stream, data_path)
+        print(f"Successfully uploaded JSON data to {data_path}")
     # create_ssh_client(SSH_SERVER, SSH_PORT, SSH_USER, SSH_PASSWORD)