Update app.py

app.py

@@ -6,25 +6,6 @@ from PIL import Image
 import random
 import os
 
-model_id = 'aipicasso/picasso-diffusion-1-0'
-token=os.environ.get("ACCESS_TOKEN")
-
-scheduler = EulerAncestralDiscreteScheduler.from_pretrained(model_id, subfolder="scheduler", use_auth_token=token)
-feature_extractor = CLIPFeatureExtractor.from_pretrained(model_id, use_auth_token=token)
-
-pipe_org = StableDiffusionPipeline.from_pretrained(
-  model_id,
-  torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
-  scheduler=scheduler, use_auth_token=token)
-
-pipe_i2i_org = StableDiffusionImg2ImgPipeline.from_pretrained(
-  model_id,
-  torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
-  scheduler=scheduler,
-  requires_safety_checker=False,
-  safety_checker=None,
-  feature_extractor=feature_extractor, use_auth_token=token
-)
 
 model_id = 'aipicasso/picasso-diffusion-1-0-demo'
 token=os.environ.get("ACCESS_TOKEN")
@@ -48,6 +29,8 @@ pipe_i2i_merged = StableDiffusionImg2ImgPipeline.from_pretrained(
 
 pipe=pipe_merged.to("cuda")
 pipe_i2i=pipe_i2i_merged.to("cuda")
+pipe.enable_xformers_memory_efficient_attention()
+pipe_i2i.enable_xformers_memory_efficient_attention()
 
 def error_str(error, title="Error"):
     return f"""#### {title}
@@ -55,18 +38,6 @@ def error_str(error, title="Error"):
 
 def inference(prompt, guidance, steps, image_size="Square", seed=0, img=None, strength=0.5, neg_prompt="", disable_auto_prompt_correction=False, original_model=False):
   global pipe,pipe_i2i
-
-  if torch.cuda.is_available():
-      pipe=pipe.to("cpu")
-      pipe_i2i=pipe_i2i.to("cpu")
-      if original_model:
-          pipe = pipe_org.to("cuda")
-          pipe_i2i = pipe_i2i_org.to("cuda")
-      else:
-          pipe = pipe_merged.to("cuda")
-          pipe_i2i = pipe_merged.to("cuda")
-      pipe.enable_xformers_memory_efficient_attention()
-      pipe_i2i.enable_xformers_memory_efficient_attention()
   
   generator = torch.Generator('cuda').manual_seed(seed) if seed != 0 else None